analytical-viz: roofline tab — S_kv knob, regime formulas, good B/L, decomp plot, split table
Bundle of roofline-tab enhancements requested in-thread:
1. **S_kv slider at top of tab** — override the sidebar's S_kv locally
so all four plots update as you drag it. Handy for exploring how
the KV wall arrives without disturbing the rest of the app config.
2. **Good B / Good L KPI cards** — 3 metrics under the main KPI row:
- Good B = 2·B* (Pope's rule of thumb)
- Good L ceiling = L* (compute-friendly context ceiling)
- Utilization @ current S_kv = 1 / (1 + S_kv/L*)
3. **Plot 4 — latency-decomposition per decode step** — five curves
on one axis:
- Compute (dashed, flat)
- Weight fetch (triangles, shrinks 1/B)
- KV fetch (squares, flat — batching doesn't help)
- Memory total (weights + KV, purple)
- Total (compute + memory, black bold)
Makes "which term dominates at this B?" visible at a glance.
4. **Regime formulas table** — one row per cost term (t_com, weight
fetch, KV fetch, bottleneck) × two columns (short-context vs
long-context regime), plus a 'value now' column using the
current S_kv slider.
5. **How to pick B and S_kv — one-paragraph guidance** with the
numeric recommendations plugged in.
6. **Split formula table** — was cramming symbolic + substituted
form into one cell with newlines; now has four proper columns:
Symbol / Formula / With numbers / Meaning / Value.
New pure functions in chip_roofline.py: t_mem_short, t_mem_long,
t_com, good_batch, good_context, utilization_at. All roofline math
stays in the pure module; app.py just calls them and formats.
10 new tests bring the roofline test count to 27 (all 67 tests still
green): t_mem_long(L*) == t_com, doubling B halves t_mem_short,
good_batch scales with sparsity, utilization at 0/L*/2L* returns
1.0/0.5/0.333, monotonic decrease with context.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -57,9 +57,15 @@ from tests.analytical_visualization.chip_roofline import (
|
||||
balance_context,
|
||||
bound_regime,
|
||||
critical_batch,
|
||||
good_batch,
|
||||
good_context,
|
||||
knee_batch,
|
||||
per_token_latency_curve,
|
||||
t_com,
|
||||
t_mem_long,
|
||||
t_mem_short,
|
||||
total_active_params,
|
||||
utilization_at,
|
||||
)
|
||||
from tests.analytical_visualization.memory_layout import (
|
||||
compute_memory,
|
||||
@@ -2171,8 +2177,29 @@ with tab_roofline:
|
||||
_b_star = critical_batch(_default_machine, model)
|
||||
_l_star = balance_context(_default_machine, model)
|
||||
_n_active = total_active_params(model)
|
||||
|
||||
# ── Local S_kv knob — override sidebar for exploring the plots
|
||||
# without changing the rest of the app's config.
|
||||
st.markdown("**Explore: override S_kv for the plots below**")
|
||||
_skv_min = 128
|
||||
_skv_max = max(int(20 * _l_star), 4 * s_kv, 1_000_000)
|
||||
_rf_skv = st.slider(
|
||||
"S_kv (context length) for roofline plots",
|
||||
min_value=_skv_min, max_value=_skv_max,
|
||||
value=int(s_kv),
|
||||
step=max(1, _skv_min),
|
||||
key="_rf_skv_override",
|
||||
help=("Local to this tab. Drag to see how KV-read time grows "
|
||||
"and where the knee disappears past L*."),
|
||||
)
|
||||
st.caption(
|
||||
f"Using S_kv = **{_rf_skv:,}** tokens for all plots on this tab. "
|
||||
f"L* on your chip is **{_l_star:,.0f}** tokens → "
|
||||
f"S_kv / L* = **{_rf_skv/_l_star:.2f}**."
|
||||
)
|
||||
|
||||
_regime_now = bound_regime(_default_machine, model,
|
||||
batch=max(1, b_batch), s_kv=s_kv)
|
||||
batch=max(1, b_batch), s_kv=_rf_skv)
|
||||
|
||||
# ── KPI cards ─────────────────────────────────────────────────
|
||||
_r1, _r2, _r3, _r4 = st.columns(4)
|
||||
@@ -2186,10 +2213,23 @@ with tab_roofline:
|
||||
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:,})",
|
||||
_r4.metric(f"Regime @ (B={b_batch}, S_kv={_rf_skv:,})",
|
||||
_regime_now.replace("-bound", ""),
|
||||
help="Which term dominates cost at your current sidebar "
|
||||
"settings.")
|
||||
help="Which term dominates cost at the current (B, S_kv) "
|
||||
"you've selected.")
|
||||
|
||||
# ── Recommended B and L (heuristics) ──────────────────────────
|
||||
_rec_b = good_batch(_default_machine, model, sparsity=1.0)
|
||||
_rec_l = good_context(_default_machine, model, s_kv=_rf_skv)
|
||||
_g1, _g2, _g3 = st.columns(3)
|
||||
_g1.metric("Good B (dense)", f"{_rec_b.effective:.0f}",
|
||||
help=_rec_b.reason)
|
||||
_g2.metric("Good L ceiling", f"{_rec_l.max_efficient:,.0f} tok",
|
||||
help="Stay at or below L* to keep decode compute-bound.")
|
||||
_g3.metric(f"Utilization @ S_kv={_rf_skv:,}",
|
||||
f"{_rec_l.utilization_at*100:.1f}%",
|
||||
help="Peak compute utilization at this context: "
|
||||
"1 / (1 + S_kv/L*).")
|
||||
|
||||
# Optional MoE B*: only surface if the preset flags MoE.
|
||||
_is_moe = "MoE" in (preset.family + preset.note)
|
||||
@@ -2211,11 +2251,10 @@ with tab_roofline:
|
||||
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:,})")
|
||||
st.markdown(f"**Per-token latency vs batch size** (S_kv = {_rf_skv:,})")
|
||||
_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)
|
||||
s_kv=_rf_skv)
|
||||
_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],
|
||||
@@ -2232,7 +2271,7 @@ with tab_roofline:
|
||||
_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.set_title(f"Cost curve at S_kv = {_rf_skv:,} tokens")
|
||||
_ax1.grid(True, which="both", alpha=0.3)
|
||||
_ax1.legend(fontsize=8, loc="upper right")
|
||||
plt.tight_layout()
|
||||
@@ -2318,70 +2357,158 @@ with tab_roofline:
|
||||
|
||||
st.divider()
|
||||
|
||||
# ── Plot 4: decode-step latency decomposition ─────────────────
|
||||
st.markdown(f"**Latency per decode step — decomposition** "
|
||||
f"(S_kv = {_rf_skv:,})")
|
||||
_fig4, _ax4 = plt.subplots(figsize=(9, 4.5))
|
||||
_weight_vals = [p.weight_s * 1e3 for p in _pts]
|
||||
_compute_vals = [p.compute_s * 1e3 for p in _pts]
|
||||
_kv_vals = [p.kv_s * 1e3 for p in _pts]
|
||||
_mem_total = [w + k for w, k in zip(_weight_vals, _kv_vals)]
|
||||
_total_vals = [p.total_s * 1e3 for p in _pts]
|
||||
|
||||
_ax4.plot(_xs, _compute_vals, "--", color="#3a86ff",
|
||||
label="Compute", linewidth=1.6)
|
||||
_ax4.plot(_xs, _weight_vals, "^-", color="#ffbe0b",
|
||||
label="Weight fetch (shrinks 1/B)", linewidth=1.6)
|
||||
_ax4.plot(_xs, _kv_vals, "s-", color="#d90429",
|
||||
label="KV fetch (flat)", linewidth=1.6)
|
||||
_ax4.plot(_xs, _mem_total, "-", color="#7b1fa2",
|
||||
label="Memory total (weights + KV)", linewidth=2)
|
||||
_ax4.plot(_xs, _total_vals, "-", color="#212529",
|
||||
label="Total (compute + memory)", linewidth=2.5)
|
||||
_ax4.axvline(_b_star, linestyle=":", color="#2e7d32",
|
||||
label=f"B* = {_b_star:.0f}")
|
||||
_ax4.set_xscale("log", base=2)
|
||||
_ax4.set_yscale("log")
|
||||
_ax4.set_xlabel("Batch size (sequences in flight)")
|
||||
_ax4.set_ylabel("Per-token step time (ms)")
|
||||
_ax4.set_title("Where the time goes: compute vs weight fetch vs KV fetch")
|
||||
_ax4.grid(True, which="both", alpha=0.3)
|
||||
_ax4.legend(fontsize=8, loc="upper right")
|
||||
plt.tight_layout()
|
||||
st.pyplot(_fig4, width='stretch')
|
||||
plt.close(_fig4)
|
||||
|
||||
st.caption(
|
||||
"**Compute** is flat in B. **Weight fetch** shrinks like 1/B "
|
||||
"(the amortization win). **KV fetch** is flat in B — each "
|
||||
"sequence reads its own KV, so batching does nothing for it. "
|
||||
"**Memory total** is weight+KV summed. Which of compute or "
|
||||
"memory-total is bigger at your operating B tells you the "
|
||||
"regime; if memory-total's floor (= KV alone at large B) "
|
||||
"already sits above compute, you're past L*."
|
||||
)
|
||||
|
||||
st.divider()
|
||||
|
||||
# ── Regime formulas (short vs long context) ───────────────────
|
||||
st.markdown("**Cost-term formulas by regime**")
|
||||
_t_com_val = t_com(_default_machine, model) * 1e3 # ms
|
||||
_t_mem_s_at_bstar = t_mem_short(_default_machine, model,
|
||||
int(round(_b_star))) * 1e3
|
||||
_t_mem_l_at_skv = t_mem_long(_default_machine, model, _rf_skv) * 1e3
|
||||
_regime_rows = [
|
||||
{"Term": "t_com (compute)",
|
||||
"Short context (S_kv < L*)": "2 · N / C",
|
||||
"Long context (S_kv > L*)": "2 · N / C (same)",
|
||||
"Value now (ms)": f"{_t_com_val:.3f}"},
|
||||
{"Term": "t_mem — weight fetch",
|
||||
"Short context (S_kv < L*)": "N · b / (W · B)",
|
||||
"Long context (S_kv > L*)": "N · b / (W · B) — becomes small",
|
||||
"Value now (ms)":
|
||||
f"{_t_mem_s_at_bstar:.3f} at B = B*"},
|
||||
{"Term": "t_mem — KV fetch",
|
||||
"Short context (S_kv < L*)": "S_kv · kv_bpt / W — small",
|
||||
"Long context (S_kv > L*)": "S_kv · kv_bpt / W — dominates",
|
||||
"Value now (ms)": f"{_t_mem_l_at_skv:.3f}"},
|
||||
{"Term": "Bottleneck",
|
||||
"Short context (S_kv < L*)":
|
||||
"Weights (B < B*), else Compute (B ≥ B*)",
|
||||
"Long context (S_kv > L*)":
|
||||
"KV bandwidth (regardless of B)",
|
||||
"Value now (ms)":
|
||||
_regime_now.replace("-bound", "")},
|
||||
]
|
||||
st.dataframe(pd.DataFrame(_regime_rows), width='stretch',
|
||||
hide_index=True)
|
||||
st.caption(
|
||||
"N = active params, b = bytes/elem, W = HBM BW, C = peak FLOPs, "
|
||||
"kv_bpt = KV bytes per token = 2 · H_kv · d_head · b · layers. "
|
||||
"'Value now' uses the S_kv slider above and B = B\\*."
|
||||
)
|
||||
|
||||
st.divider()
|
||||
|
||||
# ── Recommended B and L (recap) ───────────────────────────────
|
||||
st.markdown("**How to pick a good B and a good S_kv**")
|
||||
st.markdown(
|
||||
f"- **Good batch**: `B_target = 2 · B* = {_rec_b.effective:.0f}`. "
|
||||
"Below B\\*: memory-bound (doubling B halves cost/token). "
|
||||
"Above 3·B\\*: diminishing returns; latency keeps rising. "
|
||||
"Also cap by your HBM budget: `B_max ≤ (HBM − weights_per_PE) / "
|
||||
"(S_kv · kv_bpt_per_PE)` — check the Memory Breakdown tab.\n"
|
||||
f"- **Good S_kv**: stay ≤ **L\\* = {_l_star:,.0f} tokens** to "
|
||||
"keep decode compute-bound (peak utilization). Past L\\*, "
|
||||
f"utilization = 1/(1 + S_kv/L\\*). At S_kv = {_rf_skv:,}, "
|
||||
f"you're at ~**{_rec_l.utilization_at*100:.1f}%** utilization."
|
||||
)
|
||||
|
||||
st.divider()
|
||||
|
||||
# ── Interpretation table ───────────────────────────────────────
|
||||
st.markdown("**Formulas & interpretation**")
|
||||
# Pre-compute substituted-value strings so each row shows both the
|
||||
# symbolic form and the actual numbers plugged in.
|
||||
_C = _default_machine.peak_flops # FLOPs / s
|
||||
_W = _default_machine.bw_hbm # bytes / s
|
||||
_b_elem = model.bytes_per_elem # bytes / weight
|
||||
_kv_bpt_full = 2 * model.h_kv * model.d_head * model.bytes_per_elem * model.layers
|
||||
|
||||
_sub_ai = (
|
||||
f"= {_C:.2e} / {_W:.2e}\n= {_ai:.2f} FLOPs/byte"
|
||||
)
|
||||
_sub_bstar = (
|
||||
f"= {_C:.2e} · {_b_elem} / (2 · {_W:.2e})\n= {_b_star:.2f}"
|
||||
)
|
||||
if _is_moe:
|
||||
_sub_bmoe = f"= {_b_star:.0f} · 8 (k=8 example)\n= {_b_star * 8:.0f}"
|
||||
else:
|
||||
_sub_bmoe = "N/A (dense preset — sparsity = 1)"
|
||||
_sub_lstar = (
|
||||
f"= 2 · {_n_active:.2e} · {_W:.2e} / ({_C:.2e} · {_kv_bpt_full:,})\n"
|
||||
f"= {_l_star:,.0f} tokens"
|
||||
)
|
||||
_knee_now = knee_batch(_default_machine, model, s_kv)
|
||||
_knee_now = knee_batch(_default_machine, model, _rf_skv)
|
||||
if _knee_now is None:
|
||||
_sub_bknee = (
|
||||
f"= {_b_star:.0f} / (1 − {s_kv:,}/{_l_star:,.0f})\n"
|
||||
f"= {_b_star:.0f} / (1 − {s_kv/_l_star:.3f}) [<= 0 → no knee]"
|
||||
f"{_b_star:.0f} / (1 − {_rf_skv:,}/{_l_star:,.0f}) = "
|
||||
f"{_b_star:.0f} / (1 − {_rf_skv/_l_star:.3f}) [≤ 0]"
|
||||
)
|
||||
_bknee_val = "no knee (S_kv ≥ L*)"
|
||||
else:
|
||||
_sub_bknee = (
|
||||
f"= {_b_star:.0f} / (1 − {s_kv:,}/{_l_star:,.0f})\n"
|
||||
f"= {_b_star:.0f} / {1 - s_kv/_l_star:.3f}\n"
|
||||
f"= {_knee_now:.0f}"
|
||||
f"{_b_star:.0f} / (1 − {_rf_skv:,}/{_l_star:,.0f}) = "
|
||||
f"{_b_star:.0f} / {1 - _rf_skv/_l_star:.3f} = "
|
||||
f"{_knee_now:.0f}"
|
||||
)
|
||||
_bknee_val = f"{_knee_now:.0f}"
|
||||
|
||||
_formula_rows = [
|
||||
{"Symbol": "AI",
|
||||
"Formula": "C / W\n" + _sub_ai,
|
||||
"Meaning": "FLOPs per byte of HBM bandwidth. Chip-only.",
|
||||
"Your chip": f"{_ai:.1f}"},
|
||||
"Formula": "C / W",
|
||||
"With numbers": f"{_C:.2e} / {_W:.2e}",
|
||||
"Meaning": "FLOPs per byte of HBM BW. Chip-only.",
|
||||
"Value": f"{_ai:.2f} FLOPs/byte"},
|
||||
{"Symbol": "B*",
|
||||
"Formula": "C · b / (2 · W)\n" + _sub_bstar,
|
||||
"Meaning": ("Batch size where weight-fetch time = compute time. "
|
||||
"Below: memory-bound. Above: amortized."),
|
||||
"Your chip": f"{_b_star:.0f}"},
|
||||
"Formula": "C · b / (2 · W)",
|
||||
"With numbers": f"{_C:.2e} · {_b_elem} / (2 · {_W:.2e})",
|
||||
"Meaning": ("Batch where weight-fetch = compute. Below: "
|
||||
"memory-bound. Above: amortized."),
|
||||
"Value": f"{_b_star:.0f}"},
|
||||
{"Symbol": "B*_moe",
|
||||
"Formula": "B* · (N_total / N_active)\n" + _sub_bmoe,
|
||||
"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)")},
|
||||
"Formula": "B* · (N_total / N_active)",
|
||||
"With numbers": (f"{_b_star:.0f} · 8 (k=8 example)" if _is_moe
|
||||
else "N/A (dense preset)"),
|
||||
"Meaning": "MoE fetches all experts; computes on active only.",
|
||||
"Value": (f"{_b_star * 8:.0f}" if _is_moe else "1× (dense)")},
|
||||
{"Symbol": "L*",
|
||||
"Formula": "2 · N_active · W / (C · kv_bpt)\n" + _sub_lstar,
|
||||
"Meaning": ("Context where KV-read time = compute time. "
|
||||
"Above: KV-bound regardless of B."),
|
||||
"Your chip": f"{_l_star:,.0f} tok"},
|
||||
"Formula": "2 · N_active · W / (C · kv_bpt)",
|
||||
"With numbers": (f"2 · {_n_active:.2e} · {_W:.2e} / "
|
||||
f"({_C:.2e} · {_kv_bpt_full:,})"),
|
||||
"Meaning": ("Context where KV-read = compute. Above: "
|
||||
"KV-bound regardless of B."),
|
||||
"Value": f"{_l_star:,.0f} tokens"},
|
||||
{"Symbol": "B_knee",
|
||||
"Formula": "B* / (1 − S_kv/L*)\n" + _sub_bknee,
|
||||
"Meaning": ("Batch size where cost curve bends onto its floor. "
|
||||
"Formula": "B* / (1 − S_kv/L*)",
|
||||
"With numbers": _sub_bknee,
|
||||
"Meaning": ("Batch where cost curve bends onto its floor. "
|
||||
"Diverges at S_kv = L*."),
|
||||
"Your chip": _bknee_val},
|
||||
"Value": _bknee_val},
|
||||
]
|
||||
st.dataframe(pd.DataFrame(_formula_rows),
|
||||
width='stretch', hide_index=True,
|
||||
row_height=90)
|
||||
width='stretch', hide_index=True)
|
||||
|
||||
@@ -164,3 +164,95 @@ def bound_regime(machine: MachineParams, model: ModelConfig,
|
||||
"kv-bound": p.kv_s,
|
||||
"compute-bound": p.compute_s}
|
||||
return max(parts, key=parts.get)
|
||||
|
||||
|
||||
# ── Regime-dependent cost terms ────────────────────────────────────
|
||||
|
||||
|
||||
def t_mem_short(machine: MachineParams, model: ModelConfig,
|
||||
batch: int) -> float:
|
||||
"""Per-token weight-fetch time (short-context regime term).
|
||||
|
||||
N_active · b / (W · B). Shrinks as B grows — this is what
|
||||
batching amortizes.
|
||||
"""
|
||||
return (total_active_params(model) * model.bytes_per_elem
|
||||
/ (machine.bw_hbm * max(1, batch)))
|
||||
|
||||
|
||||
def t_mem_long(machine: MachineParams, model: ModelConfig,
|
||||
s_kv: int) -> float:
|
||||
"""Per-token KV-read time (long-context regime term).
|
||||
|
||||
S_kv · kv_bpt / W. Independent of B — each sequence reads its
|
||||
own KV cache; batching doesn't help.
|
||||
"""
|
||||
return s_kv * kv_bytes_per_token(model) / machine.bw_hbm
|
||||
|
||||
|
||||
def t_com(machine: MachineParams, model: ModelConfig) -> float:
|
||||
"""Per-token compute time. Same in both regimes: 2·N/C, peak."""
|
||||
return _FLOPS_PER_PARAM_PER_TOKEN * total_active_params(model) / machine.peak_flops
|
||||
|
||||
|
||||
# ── "Good" batch / context recommendations ─────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchRecommendation:
|
||||
target: float # Pope's rule: 2 × B*
|
||||
b_star: float # B* itself
|
||||
effective: float # what we recommend using
|
||||
reason: str # short explanation
|
||||
|
||||
|
||||
@dataclass
|
||||
class ContextRecommendation:
|
||||
l_star: float # balance context length
|
||||
max_efficient: float # same as l_star (compute-friendly ceiling)
|
||||
utilization_at: float # utilization at current s_kv
|
||||
reason: str
|
||||
|
||||
|
||||
def good_batch(machine: MachineParams, model: ModelConfig,
|
||||
sparsity: float = 1.0) -> BatchRecommendation:
|
||||
"""Recommended batch size: 2 × B* (Pope's rule of thumb).
|
||||
|
||||
Below B*: memory-bound, doubling B halves cost/token.
|
||||
At 2×B*: 50% excess over compute floor — the sweet spot.
|
||||
Beyond 3×B*: diminishing returns; latency keeps growing linearly.
|
||||
"""
|
||||
b_star = critical_batch(machine, model, sparsity)
|
||||
target = 2 * b_star
|
||||
return BatchRecommendation(
|
||||
target=target, b_star=b_star, effective=target,
|
||||
reason=(f"2·B* = 2 · {b_star:.0f} = {target:.0f}. "
|
||||
"Below B*: memory-bound (doubling B halves cost/token). "
|
||||
"Beyond 3·B*: diminishing returns."),
|
||||
)
|
||||
|
||||
|
||||
def good_context(machine: MachineParams, model: ModelConfig,
|
||||
s_kv: int) -> ContextRecommendation:
|
||||
"""Recommended max context: L* — the compute-friendly ceiling.
|
||||
|
||||
Below L*: KV read is cheap relative to compute → good utilization.
|
||||
Above L*: KV bandwidth wall → utilization = 1/(1 + S_kv/L*).
|
||||
"""
|
||||
l_star = balance_context(machine, model)
|
||||
util = utilization_at(s_kv, l_star)
|
||||
return ContextRecommendation(
|
||||
l_star=l_star, max_efficient=l_star, utilization_at=util,
|
||||
reason=(f"L* = {l_star:,.0f} tokens. Below L*: compute-bound "
|
||||
f"(good util). At {s_kv:,} tokens: peak utilization ≈ "
|
||||
f"{util*100:.1f}% (1 / (1 + S_kv/L*))."),
|
||||
)
|
||||
|
||||
|
||||
def utilization_at(s_kv: int, l_star: float) -> float:
|
||||
"""Peak compute utilization at context length s_kv, given L*.
|
||||
|
||||
util = compute / (compute + KV_read) = 1 / (1 + S_kv/L*).
|
||||
At S_kv=0: 100%. At S_kv=L*: 50%. At 2·L*: 33.3%. At 5·L*: 16.7%.
|
||||
"""
|
||||
return 1.0 / (1.0 + s_kv / l_star)
|
||||
|
||||
@@ -10,10 +10,16 @@ from tests.analytical_visualization.chip_roofline import (
|
||||
balance_context,
|
||||
bound_regime,
|
||||
critical_batch,
|
||||
good_batch,
|
||||
good_context,
|
||||
knee_batch,
|
||||
kv_bytes_per_token,
|
||||
per_token_latency_curve,
|
||||
t_com,
|
||||
t_mem_long,
|
||||
t_mem_short,
|
||||
total_active_params,
|
||||
utilization_at,
|
||||
)
|
||||
from tests.analytical_visualization.model_config import (
|
||||
MachineParams, ModelConfig,
|
||||
@@ -193,6 +199,88 @@ def test_kv_bytes_per_token_llama3_8b():
|
||||
assert kv_bpt == 2 * 8 * 128 * 2 * 32
|
||||
|
||||
|
||||
# ── Regime terms + good-batch / good-context recommendations ──────
|
||||
|
||||
|
||||
def test_t_mem_short_shrinks_with_batch():
|
||||
"""t_mem_short = N·b/(W·B) — doubling B halves it."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
assert t_mem_short(m, model, 2) == pytest.approx(
|
||||
t_mem_short(m, model, 1) / 2
|
||||
)
|
||||
|
||||
|
||||
def test_t_mem_long_equals_t_com_at_l_star():
|
||||
"""L* is defined by t_mem_long(L*) == t_com. Cross-check."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
l_star = balance_context(m, model)
|
||||
assert t_mem_long(m, model, int(round(l_star))) == pytest.approx(
|
||||
t_com(m, model), rel=0.01,
|
||||
)
|
||||
|
||||
|
||||
def test_t_com_is_batch_independent():
|
||||
"""t_com = 2·N/C. Constant — doesn't depend on B or S_kv."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
tc = t_com(m, model)
|
||||
assert tc == pytest.approx(2 * total_active_params(model) / m.peak_flops)
|
||||
|
||||
|
||||
def test_good_batch_is_two_b_star():
|
||||
"""Recommendation is 2 × B* (Pope's rule)."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
rec = good_batch(m, model)
|
||||
assert rec.effective == pytest.approx(2 * critical_batch(m, model))
|
||||
assert rec.b_star == pytest.approx(critical_batch(m, model))
|
||||
|
||||
|
||||
def test_good_batch_moe_scales_with_sparsity():
|
||||
"""MoE sparsity=8 → recommended B is 8× dense recommendation."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
r_dense = good_batch(m, model, sparsity=1.0)
|
||||
r_moe = good_batch(m, model, sparsity=8.0)
|
||||
assert r_moe.effective == pytest.approx(8 * r_dense.effective)
|
||||
|
||||
|
||||
def test_good_context_returns_l_star_as_ceiling():
|
||||
"""Recommendation is L*; utilization drops past it."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
rec = good_context(m, model, s_kv=100)
|
||||
assert rec.l_star == pytest.approx(balance_context(m, model))
|
||||
assert rec.max_efficient == pytest.approx(rec.l_star)
|
||||
|
||||
|
||||
def test_utilization_at_l_star_is_half():
|
||||
"""At S_kv = L*, compute and KV read take equal time → 50% util."""
|
||||
l_star = 3000.0
|
||||
assert utilization_at(int(l_star), l_star) == pytest.approx(0.5)
|
||||
|
||||
|
||||
def test_utilization_at_2_l_star_is_one_third():
|
||||
"""At S_kv = 2·L*, util = 1/(1+2) = 33.3%."""
|
||||
l_star = 3000.0
|
||||
assert utilization_at(int(2 * l_star), l_star) == pytest.approx(1 / 3, rel=1e-3)
|
||||
|
||||
|
||||
def test_utilization_at_zero_context_is_100_percent():
|
||||
"""Zero context, no KV to read → all time is compute."""
|
||||
assert utilization_at(0, 3000.0) == pytest.approx(1.0)
|
||||
|
||||
|
||||
def test_utilization_drops_monotonically_with_context():
|
||||
"""util(S_kv) is strictly decreasing."""
|
||||
l_star = 3000.0
|
||||
vals = [utilization_at(s, l_star) for s in [100, 1000, 3000, 6000, 30000]]
|
||||
for a, b in zip(vals, vals[1:]):
|
||||
assert b < a
|
||||
|
||||
|
||||
# ── App wiring: tab exists on the Streamlit app ────────────────────
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user