ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

analytical-viz: 'Chip roofline & B*' tab — AI, B*, L*, cost curves

Browse Source 
New tab surfaces the arithmetic-intensity story from LLM-serving
practice for the current sidebar chip + model. All derived from
MachineParams + ModelConfig; no new configuration.

chip_roofline module (pure functions):
- arithmetic_intensity      = C / W   (FLOPs per byte HBM BW)
- critical_batch (B*)       = C * b / (2 * W) * sparsity
- balance_context (L*)      = 2 * N_active * W / (C * kv_bpt)
- knee_batch (B_knee)       = B* / (1 - S_kv/L*)  (None past L*)
- per_token_latency_curve   = weight/B + compute + KV, per B
- bound_regime              = which term dominates now

Peak roofline convention (no compute_util factor) so weight_s ==
compute_s exactly at B*. Comm + TP/CP sharding intentionally
excluded — stage_latencies is the full latency model; this is the
back-of-envelope chip-vs-model view.

Tab shows:
- 4 KPI cards: AI, B*, L*, current regime at (B, S_kv)
- MoE hint when preset flags MoE ('300 x sparsity' rule)
- Plot 1: cost vs B at current S_kv (weight, compute floor, KV
  floor, total) with B* marker line
- Plot 2: cost curves at S_kv/L* = 0.25, 0.5, 1, 2, 5 — shows the
  no-knee regime past L*
- Plot 3: B_knee vs S_kv — knee slides right, diverges at L*
- Formulas + interpretation table

test_chip_roofline covers B* against H100 reference (295), sparsity
scaling, L* scaling with HBM BW, knee divergence at L*, per-token
curve monotonicity + asymptote to compute+KV floor, regime
classification. Smoke test guards the tab is wired in app.py.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-07-29 11:59:19 -07:00
parent 5e92b89821
commit 1ddd1baa7a
3 changed files with 582 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/analytical_visualization/app.py
+209 -1
View File
@@ -30,6 +30,7 @@ for _mod_name in (
"tests.analytical_visualization.memory_layout",
"tests.analytical_visualization.stage_latencies",
"tests.analytical_visualization.stage_shapes",
"tests.analytical_visualization.chip_roofline",
"tests.analytical_visualization.auto_explore",
"tests.analytical_visualization.auto_hardware",
"tests.analytical_visualization.pe_weight_layout",
@@ -51,6 +52,15 @@ from tests.analytical_visualization.stage_shapes import (
attn_stage_shape_rows,
ffn_stage_shape_rows,
)
from tests.analytical_visualization.chip_roofline import (
arithmetic_intensity,
balance_context,
bound_regime,
critical_batch,
knee_batch,
per_token_latency_curve,
total_active_params,
)
from tests.analytical_visualization.memory_layout import (
compute_memory,
attention_weight_rows,
@@ -531,11 +541,13 @@ if _warnings:
# Attn + FFN/MoE) so the user can toggle scope without switching tabs.
# Auto Suggest is renamed "Parallelism" since it only varies parallelism
# knobs (hardware is held fixed at the sidebar values).
tab_layout, tab_auto, tab_memory, tab_stages, tab_compare, tab_hw = st.tabs([
(tab_layout, tab_auto, tab_memory, tab_stages, tab_compare, tab_hw,
tab_roofline) = st.tabs([
"Physical layout",
"Auto Suggest Parallelism",
"Memory breakdown", "Per-stage latency",
"Save & compare", "Auto Hardware",
"Chip roofline & B*",
])
@@ -2142,3 +2154,199 @@ def _render_auto_hardware_tab():
with tab_hw:
_render_auto_hardware_tab()
# ── TAB 7: Chip roofline & B* ────────────────────────────────────
with tab_roofline:
st.subheader("Chip roofline & critical batch size (B*)")
st.caption(
"Back-of-envelope answer to 'is my chip well-matched to this "
"model?' — see how big the batch must be before decode becomes "
"compute-bound, and how big the context can grow before the KV "
"bandwidth wall kills your utilization. All numbers per PE, "
"peak roofline (no compute-util factor), no comm."
)
_ai = arithmetic_intensity(_default_machine)
_b_star = critical_batch(_default_machine, model)
_l_star = balance_context(_default_machine, model)
_n_active = total_active_params(model)
_regime_now = bound_regime(_default_machine, model,
batch=max(1, b_batch), s_kv=s_kv)
# ── KPI cards ─────────────────────────────────────────────────
_r1, _r2, _r3, _r4 = st.columns(4)
_r1.metric("AI (FLOPs/byte)", f"{_ai:.1f}",
help="C/W. Peak FLOPs per byte of HBM bandwidth. "
"Chip-only — no model-dependent factor.")
_r2.metric("B* (dense)", f"{_b_star:.0f}",
help="C*b/(2*W). Batch size where weight-fetch time "
"equals compute time. Below this you're memory-bound.")
_r3.metric("L* (tokens)", f"{_l_star:,.0f}",
help="Context length where KV-read time equals compute "
"time. Above this, no batch size gets you back to "
"compute-bound.")
_r4.metric(f"Regime @ (B={b_batch}, S_kv={s_kv:,})",
_regime_now.replace("-bound", ""),
help="Which term dominates cost at your current sidebar "
"settings.")
# Optional MoE B*: only surface if the preset flags MoE.
_is_moe = "MoE" in (preset.family + preset.note)
if _is_moe:
st.info(
"This preset is dense-approximated as an MoE. Reiner Pope's "
"'300 × sparsity' rule: for real sparsity k = N_total/N_active, "
f"B*_moe ≈ {_b_star:.0f} × k. E.g. DeepSeek 32-of-256 experts "
f"(k=8) → B* ≈ {_b_star*8:.0f}."
)
st.caption(
f"**Full-model attention+FFN params:** {_n_active/1e9:.2f} B. "
f"**Chip:** {_default_machine.peak_tflops_f16:.0f} TFLOPs BF16, "
f"{_default_machine.bw_hbm_gbs:.0f} GB/s HBM. "
f"Change these in the sidebar Hardware panel to see the roofline shift."
)
st.divider()
# ── Plot 1: cost vs B at current S_kv ─────────────────────────
st.markdown("**Per-token latency vs batch size** (current S_kv = "
f"{s_kv:,})")
_b_range = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]
_pts = per_token_latency_curve(_default_machine, model, _b_range,
s_kv=s_kv)
_fig1, _ax1 = plt.subplots(figsize=(9, 4.5))
_xs = [p.batch for p in _pts]
_ax1.plot(_xs, [p.weight_s * 1e3 for p in _pts],
"o-", label="Weight fetch (1/B)", color="#ffbe0b")
_ax1.axhline(_pts[0].compute_s * 1e3, linestyle="--",
color="#3a86ff", label="Compute floor (peak)")
_ax1.axhline(_pts[0].kv_s * 1e3, linestyle=":",
color="#d90429", label="KV read floor")
_ax1.plot(_xs, [p.total_s * 1e3 for p in _pts],
"-", color="#212529", linewidth=2, label="Total")
_ax1.axvline(_b_star, linestyle=":", color="#2e7d32",
label=f"B* = {_b_star:.0f}")
_ax1.set_xscale("log", base=2)
_ax1.set_yscale("log")
_ax1.set_xlabel("Batch size (sequences in flight)")
_ax1.set_ylabel("Per-token step time (ms)")
_ax1.set_title(f"Cost curve at S_kv = {s_kv:,} tokens")
_ax1.grid(True, which="both", alpha=0.3)
_ax1.legend(fontsize=8, loc="upper right")
plt.tight_layout()
st.pyplot(_fig1, width='stretch')
plt.close(_fig1)
st.caption(
"**Below B\\*:** cost is dominated by weight fetch; doubling B "
"halves cost per token. **At B\\*:** weight fetch = compute. "
"**Above 2-3× B\\*:** you've captured most amortization; further "
"batching mostly buys latency you don't want."
)
st.divider()
# ── Plot 2: cost curves at multiple context lengths ───────────
st.markdown("**The 'no-knee' phenomenon — cost vs B at different S_kv**")
_fig2, _ax2 = plt.subplots(figsize=(9, 4.5))
_l_ratios = [0.25, 0.5, 1.0, 2.0, 5.0]
_colors = ["#3a86ff", "#2e7d32", "#ffbe0b", "#ef6c00", "#d90429"]
for _r, _col in zip(_l_ratios, _colors):
_skv_at = max(1, int(_r * _l_star))
_pts_r = per_token_latency_curve(_default_machine, model, _b_range,
s_kv=_skv_at)
_label = f"S_kv = {_r:.2g} × L* ({_skv_at:,} tok)"
_ax2.plot(_xs, [p.total_s * 1e3 for p in _pts_r], "-",
color=_col, label=_label, linewidth=1.6)
_ax2.axhline(_pts[0].compute_s * 1e3, linestyle="--",
color="#888", label="Compute floor")
_ax2.axvline(_b_star, linestyle=":", color="#2e7d32",
label=f"B* = {_b_star:.0f}")
_ax2.set_xscale("log", base=2)
_ax2.set_yscale("log")
_ax2.set_xlabel("Batch size")
_ax2.set_ylabel("Per-token step time (ms)")
_ax2.set_title("Cost curves at several context lengths "
f"(L* = {_l_star:,.0f} tokens)")
_ax2.grid(True, which="both", alpha=0.3)
_ax2.legend(fontsize=8, loc="upper right")
plt.tight_layout()
st.pyplot(_fig2, width='stretch')
plt.close(_fig2)
st.caption(
"**S_kv < L\\*:** curve bends at B_knee = B\\* / (1 − S_kv/L\\*), "
"then flattens onto the compute floor. **S_kv = L\\*:** curves "
"become parallel — you approach the floor but never touch it. "
"**S_kv > L\\*:** KV read is above the compute floor at all B; "
"no batch size recovers utilization. This is the algebraic core "
"of 'at 1M context, there's no B that gets you compute-bound'."
)
st.divider()
# ── Plot 3: knee vs context ───────────────────────────────────
st.markdown("**Effective knee B_knee(S_kv) = B\\* / (1 − S_kv/L\\*)**")
_skv_range = [int(_r * _l_star)
for _r in (0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6,
0.7, 0.8, 0.85, 0.9, 0.95, 0.98)]
_knees = [knee_batch(_default_machine, model, s) for s in _skv_range]
_fig3, _ax3 = plt.subplots(figsize=(9, 3.5))
_ax3.plot(_skv_range, _knees, "o-", color="#7b1fa2", linewidth=2)
_ax3.axhline(_b_star, linestyle=":", color="#2e7d32",
label=f"B* = {_b_star:.0f}")
_ax3.axvline(_l_star, linestyle=":", color="#d90429",
label=f"L* = {_l_star:,.0f}")
_ax3.set_yscale("log")
_ax3.set_xlabel("Context length S_kv (tokens)")
_ax3.set_ylabel("Effective knee batch size")
_ax3.set_title("How far right the cost-curve knee slides as context grows")
_ax3.grid(True, which="both", alpha=0.3)
_ax3.legend(fontsize=8, loc="upper left")
plt.tight_layout()
st.pyplot(_fig3, width='stretch')
plt.close(_fig3)
st.caption(
"The knee starts at B\\* for short context and slides right as "
"S_kv/L\\* grows. It diverges at S_kv = L\\* and disappears "
"beyond that. Sparse attention (e.g. DeepSeek-style) replaces "
"S_kv with ~√S_kv in the KV term — bending this curve back down."
)
st.divider()
# ── Interpretation table ───────────────────────────────────────
st.markdown("**Formulas & interpretation**")
_formula_rows = [
{"Symbol": "AI",
"Formula": "C / W",
"Meaning": "FLOPs per byte of HBM bandwidth. Chip-only.",
"Your chip": f"{_ai:.1f}"},
{"Symbol": "B*",
"Formula": "C · b / (2 · W)",
"Meaning": ("Batch size where weight-fetch time = compute time. "
"Below: memory-bound. Above: amortized."),
"Your chip": f"{_b_star:.0f}"},
{"Symbol": "B*_moe",
"Formula": "B* · (N_total / N_active)",
"Meaning": "MoE fetches all experts but computes on active only.",
"Your chip": ("N/A (dense preset)" if not _is_moe
else f"{_b_star * 8:.0f} (k=8 example)")},
{"Symbol": "L*",
"Formula": "2 · N_active · W / (C · kv_bpt)",
"Meaning": ("Context where KV-read time = compute time. "
"Above: KV-bound regardless of B."),
"Your chip": f"{_l_star:,.0f} tok"},
{"Symbol": "B_knee",
"Formula": "B* / (1 − S_kv/L*)",
"Meaning": ("Batch size where cost curve bends onto its floor. "
"Diverges at S_kv = L*."),
"Your chip": (f"{knee_batch(_default_machine, model, s_kv):.0f}"
if knee_batch(_default_machine, model, s_kv) is not None
else "no knee (S_kv >= L*)")},
]
st.dataframe(pd.DataFrame(_formula_rows),
width='stretch', hide_index=True)
tests/analytical_visualization/chip_roofline.py
+166
View File
@@ -0,0 +1,166 @@
"""Chip-level roofline math: AI, B*, L*, per-token latency curves.
Surfaces the arithmetic-intensity story from LLM-serving practice:
- **AI** = C / W (peak FLOPs per byte of HBM bandwidth).
- **B\*** = C * b / (2 * W) * sparsity — the critical batch size at
which weight-fetch time equals compute time for one decode step.
Sparsity = N_total / N_active (MoE factor; 1 for dense).
- **L\*** = 2 * N_active / (AI * kv_bytes_per_token) — the balance
context length at which KV-read time equals compute time.
- **B_knee(S_kv)** = B* / (1 - S_kv/L*) — the batch size where the
cost curve bends (weight-fetch drops below the compute+KV floor).
Diverges at S_kv = L* and no knee exists past it.
Per-token decode-step latency, per PE (dense-approx, no comm):
t(B, S_kv) = N_active * b / (W * B) <-- weight fetch, 1/B
+ 2 * N_active / C <-- compute (peak), flat
+ S_kv * kv_bpt / W <-- KV read, flat
All numbers per PE / per one forward pass. **Peak roofline — no
utilization factor.** Comm cost and TP/CP sharding are intentionally
NOT in the roofline — this is the back-of-envelope chip-vs-model view
the transcript talks about, not the full latency model that
stage_latencies.py builds. With this convention weight_s == compute_s
exactly at B*.
"""
from __future__ import annotations
from dataclasses import dataclass
from .model_config import MachineParams, ModelConfig
# Per-token bytes of BF16 MAC arithmetic: one multiply + one add = 2 FLOPs.
_FLOPS_PER_PARAM_PER_TOKEN = 2
# ── Chip / model derived quantities ────────────────────────────────
def arithmetic_intensity(machine: MachineParams) -> float:
"""FLOPs per byte of HBM bandwidth. Peak roofline; utilization
is applied only in the compute-time formula, not here."""
return machine.peak_flops / machine.bw_hbm
def total_active_params(model: ModelConfig) -> int:
"""Full-model parameter count (attention + FFN, all layers).
Attention: 4 projections × hidden × H_q * d_head effective per layer.
(W_Q hidden×H_q*d_h, W_O H_q*d_h×hidden, W_K/W_V hidden×H_kv*d_h.)
FFN: 3 × hidden × ffn_dim per layer (gate, up, down).
"""
m = model
attn = (
m.hidden * m.h_q * m.d_head # W_Q
+ m.hidden * m.h_kv * m.d_head * 2 # W_K + W_V
+ m.h_q * m.d_head * m.hidden # W_O
)
ffn = 3 * m.hidden * m.ffn_dim
return (attn + ffn) * m.layers
def kv_bytes_per_token(model: ModelConfig) -> int:
"""Bytes of KV cache one new token adds across ALL layers, per one
sequence, un-sharded (K + V, H_kv heads * d_h * bytes)."""
m = model
return 2 * m.h_kv * m.d_head * m.bytes_per_elem * m.layers
def critical_batch(machine: MachineParams, model: ModelConfig,
sparsity: float = 1.0) -> float:
"""B* = C * b / (2 * W) * sparsity.
Sparsity = N_total / N_active (>= 1). Dense = 1. MoE 8-of-256 = 8.
"""
b = model.bytes_per_elem
ai = arithmetic_intensity(machine)
return ai * b / _FLOPS_PER_PARAM_PER_TOKEN * sparsity
def balance_context(machine: MachineParams, model: ModelConfig) -> float:
"""L* = 2 * N_active / (AI * kv_bpt).
Context length (in tokens) at which per-step KV read matches the
per-step compute cost. Beyond L*, the KV term is dominant and no
batch size gets you compute-bound.
"""
n_active = total_active_params(model)
ai = arithmetic_intensity(machine)
kv_bpt = kv_bytes_per_token(model)
return _FLOPS_PER_PARAM_PER_TOKEN * n_active / (ai * kv_bpt)
def knee_batch(machine: MachineParams, model: ModelConfig,
s_kv: int) -> float | None:
"""B_knee(S_kv) = B* / (1 - S_kv/L*).
Returns None when S_kv >= L* (no knee exists — the total-cost
curve never touches the compute floor).
"""
b_star = critical_batch(machine, model)
l_star = balance_context(machine, model)
r = s_kv / l_star
if r >= 1.0:
return None
return b_star / (1.0 - r)
# ── Per-token latency curves ───────────────────────────────────────
@dataclass
class RooflinePoint:
batch: int
weight_s: float # weight fetch time, 1/B
compute_s: float # compute time, flat
kv_s: float # KV read time, flat
total_s: float
def per_token_latency_curve(machine: MachineParams, model: ModelConfig,
batch_range: list[int],
s_kv: int) -> list[RooflinePoint]:
"""Per-token decode-step latency curve across a range of batch sizes.
Returns one point per batch. All times per PE, dense-approx,
utilization from machine.compute_util. Comm and TP/CP sharding are
excluded — this is the roofline model.
"""
n_active = total_active_params(model)
b = model.bytes_per_elem
weight_bytes = n_active * b
compute_flops = _FLOPS_PER_PARAM_PER_TOKEN * n_active
kv_read_bytes = s_kv * kv_bytes_per_token(model)
compute_s = compute_flops / machine.peak_flops
kv_s = kv_read_bytes / machine.bw_hbm
points: list[RooflinePoint] = []
for bs in batch_range:
weight_s = weight_bytes / machine.bw_hbm / max(1, bs)
points.append(RooflinePoint(
batch=bs,
weight_s=weight_s,
compute_s=compute_s,
kv_s=kv_s,
total_s=weight_s + compute_s + kv_s,
))
return points
def bound_regime(machine: MachineParams, model: ModelConfig,
batch: int, s_kv: int) -> str:
"""Which term dominates at the current (batch, S_kv) point.
Returns 'memory-bound' if weight_fetch is the largest term,
'kv-bound' if KV read is largest, 'compute-bound' if compute.
"""
pts = per_token_latency_curve(machine, model, [batch], s_kv)
p = pts[0]
parts = {"memory-bound": p.weight_s,
"kv-bound": p.kv_s,
"compute-bound": p.compute_s}
return max(parts, key=parts.get)
tests/analytical_visualization/test_chip_roofline.py
+207
View File
@@ -0,0 +1,207 @@
"""Tests for chip_roofline: AI, B*, L*, per-token latency curves."""
from __future__ import annotations
import math
import pytest
from tests.analytical_visualization.chip_roofline import (
arithmetic_intensity,
balance_context,
bound_regime,
critical_batch,
knee_batch,
kv_bytes_per_token,
per_token_latency_curve,
total_active_params,
)
from tests.analytical_visualization.model_config import (
MachineParams, ModelConfig,
)
from tests.analytical_visualization.model_presets import PRESETS
# ── Arithmetic intensity ────────────────────────────────────────────
def test_arithmetic_intensity_matches_ratio():
"""AI = peak_flops / bw_hbm — pure ratio, no util factor."""
m = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
assert arithmetic_intensity(m) == pytest.approx(8e12 / 256e9)
def test_ai_scales_with_flops_and_bandwidth():
m1 = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
m2 = MachineParams(peak_tflops_f16=16.0, bw_hbm_gbs=256.0)
m3 = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=128.0)
assert arithmetic_intensity(m2) == pytest.approx(2 * arithmetic_intensity(m1))
assert arithmetic_intensity(m3) == pytest.approx(2 * arithmetic_intensity(m1))
# ── Critical batch B* ───────────────────────────────────────────────
def test_critical_batch_h100_reference():
"""H100-class: 989 TFLOPs BF16, 3.35 TB/s HBM3 → B* ≈ 295."""
m = MachineParams(peak_tflops_f16=989.0, bw_hbm_gbs=3350.0)
model = PRESETS["Llama 3 8B"].model # bf16 (b=2)
b_star = critical_batch(m, model)
assert b_star == pytest.approx(295, rel=0.05), b_star
def test_critical_batch_scales_with_sparsity():
"""MoE 8× sparsity gives 8× B*."""
m = MachineParams()
model = PRESETS["Llama 3 8B"].model
b_dense = critical_batch(m, model, sparsity=1)
b_moe8 = critical_batch(m, model, sparsity=8)
assert b_moe8 == pytest.approx(8 * b_dense)
def test_critical_batch_bf16_formula():
"""For BF16 (b=2), B* = AI in flops-per-byte units → numerically."""
m = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
model = ModelConfig(bytes_per_elem=2)
ai = arithmetic_intensity(m)
assert critical_batch(m, model) == pytest.approx(ai), (
critical_batch(m, model), ai,
)
# ── Balance context L* ──────────────────────────────────────────────
def test_balance_context_positive_finite():
"""L* for a real model on real machine is a positive finite number."""
m = MachineParams()
model = PRESETS["Llama 3 8B"].model
l_star = balance_context(m, model)
assert 0 < l_star < 1e9, l_star
def test_balance_context_scales_with_bandwidth():
"""L* = 2N*W/(C*kv_bpt): doubling HBM BW doubles L*.
Slower memory hits the KV wall at a shorter context. Machine-only
change so N_active and kv_bpt stay fixed."""
fast = MachineParams(bw_hbm_gbs=512.0)
slow = MachineParams(bw_hbm_gbs=256.0)
model = PRESETS["Llama 3 8B"].model
assert (balance_context(fast, model)
== pytest.approx(2 * balance_context(slow, model)))
# ── Knee ────────────────────────────────────────────────────────────
def test_knee_equals_b_star_at_short_context():
"""S_kv → 0: B_knee → B*."""
m = MachineParams()
model = PRESETS["Llama 3 8B"].model
b_star = critical_batch(m, model)
assert knee_batch(m, model, s_kv=1) == pytest.approx(b_star, rel=1e-3)
def test_knee_diverges_at_balance_context():
"""S_kv = L*: knee is infinite; S_kv > L*: no knee (None)."""
m = MachineParams()
model = PRESETS["Llama 3 8B"].model
l_star = balance_context(m, model)
assert knee_batch(m, model, s_kv=int(2 * l_star)) is None
# At 0.9 * L*, knee ~= 10 * B*; assert at least 5x to survive rounding.
below = knee_batch(m, model, s_kv=int(0.9 * l_star))
assert below is not None and below > 5 * critical_batch(m, model)
def test_knee_slides_right_with_context():
"""S_kv = L*/2 → B_knee = 2 * B*."""
m = MachineParams()
model = PRESETS["Llama 3 8B"].model
b_star = critical_batch(m, model)
l_star = balance_context(m, model)
got = knee_batch(m, model, s_kv=int(l_star / 2))
assert got == pytest.approx(2 * b_star, rel=0.01), (got, 2 * b_star)
# ── Per-token latency curve ─────────────────────────────────────────
def test_latency_curve_monotonically_decreasing():
"""total_s must strictly decrease as batch increases (weight/B shrinks)."""
m = MachineParams()
model = PRESETS["Llama 3 8B"].model
pts = per_token_latency_curve(m, model, [1, 2, 4, 8, 16, 32, 64],
s_kv=1024)
totals = [p.total_s for p in pts]
for a, b in zip(totals, totals[1:]):
assert b < a, (a, b)
def test_latency_curve_asymptotes_to_compute_plus_kv():
"""At very large B, total → compute + KV (weight/B → 0)."""
m = MachineParams()
model = PRESETS["Llama 3 8B"].model
pts = per_token_latency_curve(m, model, [10_000_000], s_kv=1024)
p = pts[0]
assert p.total_s == pytest.approx(p.compute_s + p.kv_s, rel=1e-3)
def test_latency_at_b_star_is_roughly_2x_floor():
"""At B*, weight = compute (dense, S_kv small), so total ≈ 2 * compute."""
m = MachineParams()
model = PRESETS["Llama 3 8B"].model
b_star = int(round(critical_batch(m, model)))
pts = per_token_latency_curve(m, model, [b_star], s_kv=1)
p = pts[0]
# weight_s ≈ compute_s at B*
assert p.weight_s == pytest.approx(p.compute_s, rel=0.05), (
p.weight_s, p.compute_s,
)
# ── Regime classification ───────────────────────────────────────────
def test_regime_memory_bound_at_b1():
m = MachineParams()
model = PRESETS["Llama 3 8B"].model
assert bound_regime(m, model, batch=1, s_kv=1024) == "memory-bound"
def test_regime_kv_bound_past_balance_context():
"""S_kv > L* with reasonable batch → KV dominates."""
m = MachineParams()
model = PRESETS["Llama 3 8B"].model
l_star = balance_context(m, model)
b_star = int(round(critical_batch(m, model)))
assert bound_regime(m, model, batch=b_star * 4,
s_kv=int(3 * l_star)) == "kv-bound"
# ── Component sanity ────────────────────────────────────────────────
def test_total_active_params_llama3_8b_ballpark():
"""Llama 3 8B — attn+FFN alone (no embeddings/LM head) is ~6.9B,
HF-reported total 8.03B. Guard the ballpark."""
n = total_active_params(PRESETS["Llama 3 8B"].model)
assert 6.5e9 < n < 8e9, n
def test_kv_bytes_per_token_llama3_8b():
"""Llama 3 8B: 2*8*128*2 bytes/layer/token × 32 layers = 131072 bytes."""
kv_bpt = kv_bytes_per_token(PRESETS["Llama 3 8B"].model)
assert kv_bpt == 2 * 8 * 128 * 2 * 32
# ── App wiring: tab exists on the Streamlit app ────────────────────
def test_roofline_tab_registered_in_app():
"""The new 'Chip roofline & B*' tab is listed in st.tabs and gets a
corresponding `with tab_roofline:` block."""
from pathlib import Path
src = (Path(__file__).parent / "app.py").resolve().read_text(
encoding="utf-8"
)
assert src.count('"Chip roofline & B*"') == 1
assert "with tab_roofline:" in src