From 1ddd1baa7adcb4418c7a3ace1bd2f2f8a7a5a8cd Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Wed, 29 Jul 2026 11:59:19 -0700 Subject: [PATCH] =?UTF-8?q?analytical-viz:=20'Chip=20roofline=20&=20B*'=20?= =?UTF-8?q?tab=20=E2=80=94=20AI,=20B*,=20L*,=20cost=20curves?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- tests/analytical_visualization/app.py | 210 +++++++++++++++++- .../analytical_visualization/chip_roofline.py | 166 ++++++++++++++ .../test_chip_roofline.py | 207 +++++++++++++++++ 3 files changed, 582 insertions(+), 1 deletion(-) create mode 100644 tests/analytical_visualization/chip_roofline.py create mode 100644 tests/analytical_visualization/test_chip_roofline.py diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index 1bb85bc..d1116c5 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -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) diff --git a/tests/analytical_visualization/chip_roofline.py b/tests/analytical_visualization/chip_roofline.py new file mode 100644 index 0000000..0a27b0d --- /dev/null +++ b/tests/analytical_visualization/chip_roofline.py @@ -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) diff --git a/tests/analytical_visualization/test_chip_roofline.py b/tests/analytical_visualization/test_chip_roofline.py new file mode 100644 index 0000000..3f1dc1b --- /dev/null +++ b/tests/analytical_visualization/test_chip_roofline.py @@ -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