diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index 56cc12d..11dfac7 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -2178,14 +2178,13 @@ with tab_roofline: "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) - # ── Local knobs: S_kv, B, FLOPs multiplier, BW multiplier ───── + # ── Local knobs: S_kv, B, FLOPs, HBM BW ─────────────────────── # Override sidebar for exploring the plots without changing the - # rest of the app's config. + # rest of the app's config. FLOPs/BW knobs drive ALL roofline + # math and every plot/formula on this tab (memory-budget stacks + # excluded — those are HBM-capacity, not BW). st.markdown("**Explore: local overrides for the plots below**") _kn_a, _kn_b, _kn_c, _kn_d = st.columns(4) with _kn_a: @@ -2210,44 +2209,56 @@ with tab_roofline: help="Local to this tab. Marker on plots + PE-memory KV.", ) with _kn_c: - _rf_flops_mult = st.slider( - "FLOPs × (chip)", - min_value=0.25, max_value=8.0, - value=1.0, step=0.25, - key="_rf_flops_mult", - help=("Scale peak_tflops for the AI-sensitivity plot. " - "Doubling FLOPs doubles AI and B*."), + _rf_flops_tflops = st.slider( + "FLOPs (TFLOPS)", + min_value=1.0, max_value=32.0, + value=float(_default_machine.peak_tflops_f16), + step=0.5, + key="_rf_flops_tflops", + help=("Chip peak TFLOPS for the AI-sensitivity plot. " + "Higher FLOPs → higher AI and higher B*."), ) with _kn_d: - _rf_bw_mult = st.slider( - "HBM BW × (chip)", - min_value=0.25, max_value=8.0, - value=1.0, step=0.25, - key="_rf_bw_mult", - help=("Scale bw_hbm for the AI-sensitivity plot. " - "Doubling BW halves AI and B*."), + _rf_bw_gbs = st.slider( + "HBM BW (GB/s)", + min_value=128.0, max_value=1024.0, + value=float(_default_machine.bw_hbm_gbs), + step=16.0, + key="_rf_bw_gbs", + help=("Chip HBM bandwidth for the AI-sensitivity plot. " + "Higher BW → lower AI (more bytes to feed per FLOP)."), ) - # Build a scaled machine for the AI-sensitivity view. + # Build the effective chip from the knobs; every roofline number + # on this tab is derived from it. _scaled_machine = dataclasses.replace( _default_machine, - peak_tflops_f16=_default_machine.peak_tflops_f16 * _rf_flops_mult, - bw_hbm_gbs=_default_machine.bw_hbm_gbs * _rf_bw_mult, + peak_tflops_f16=_rf_flops_tflops, + bw_hbm_gbs=_rf_bw_gbs, ) - _ai_scaled = arithmetic_intensity(_scaled_machine) - _b_star_scaled = critical_batch(_scaled_machine, model) + _ai = arithmetic_intensity(_scaled_machine) + _b_star = critical_batch(_scaled_machine, model) + _l_star = balance_context(_scaled_machine, model) + _flops_mult_now = _rf_flops_tflops / _default_machine.peak_tflops_f16 + _bw_mult_now = _rf_bw_gbs / _default_machine.bw_hbm_gbs + _ai_base = arithmetic_intensity(_default_machine) + _b_star_base = critical_batch(_default_machine, model) st.caption( f"**S_kv = {_rf_skv:,}**, **B = {_rf_b}**, " - f"**FLOPs × {_rf_flops_mult:.2f}**, **BW × {_rf_bw_mult:.2f}**. " - f"L* = **{_l_star:,.0f}** tokens (base chip) → S_kv/L* = " - f"**{_rf_skv/_l_star:.2f}**. Base B* = **{_b_star:.0f}**; " - f"scaled B* = **{_b_star_scaled:.0f}**." + f"**FLOPs = {_rf_flops_tflops:.1f} TFLOPS** " + f"(× {_flops_mult_now:.2f} of base), " + f"**BW = {_rf_bw_gbs:.0f} GB/s** " + f"(× {_bw_mult_now:.2f} of base). " + f"Effective **AI = {_ai:.1f}** (base {_ai_base:.1f}), " + f"**B\\* = {_b_star:.0f}** (base {_b_star_base:.0f}), " + f"**L\\* = {_l_star:,.0f}** tokens → S_kv/L\\* = " + f"**{_rf_skv/_l_star:.2f}**." ) # ── Headline plots: step latency and cost per token (=÷B) ───── _b_range = [1, 2, 4, 8, 16, 32, 64, 128, 256] - _step_pts = step_latency_curve(_default_machine, model, _b_range, + _step_pts = step_latency_curve(_scaled_machine, model, _b_range, s_kv=_rf_skv) - _tok_pts = per_token_latency_curve(_default_machine, model, _b_range, + _tok_pts = per_token_latency_curve(_scaled_machine, model, _b_range, s_kv=_rf_skv) _xs_head = [p.batch for p in _step_pts] @@ -2463,7 +2474,7 @@ with tab_roofline: st.divider() - _regime_now = bound_regime(_default_machine, model, + _regime_now = bound_regime(_scaled_machine, model, batch=max(1, _rf_b), s_kv=_rf_skv) # ── KPI cards ───────────────────────────────────────────────── @@ -2484,8 +2495,8 @@ with tab_roofline: "you've selected via the sliders above.") # ── 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) + _rec_b = good_batch(_scaled_machine, model, sparsity=1.0) + _rec_l = good_context(_scaled_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) @@ -2508,19 +2519,22 @@ with tab_roofline: 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." + f"**Effective chip (from knobs):** " + f"{_scaled_machine.peak_tflops_f16:.1f} TFLOPs BF16, " + f"{_scaled_machine.bw_hbm_gbs:.0f} GB/s HBM. " + f"Base chip from sidebar Hardware: " + f"{_default_machine.peak_tflops_f16:.0f} TFLOPs, " + f"{_default_machine.bw_hbm_gbs:.0f} GB/s." ) 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, + _t_com_val = t_com(_scaled_machine, model) * 1e3 # ms + _t_mem_s_at_bstar = t_mem_short(_scaled_machine, model, int(round(_b_star))) * 1e3 - _t_mem_l_at_skv = t_mem_long(_default_machine, model, _rf_skv) * 1e3 + _t_mem_l_at_skv = t_mem_long(_scaled_machine, model, _rf_skv) * 1e3 _regime_rows = [ {"Term": "t_com (compute)", "Short context (S_kv < L*)": "2 · N / C", @@ -2571,12 +2585,12 @@ with tab_roofline: # ── Interpretation table ─────────────────────────────────────── st.markdown("**Formulas & interpretation**") - _C = _default_machine.peak_flops # FLOPs / s - _W = _default_machine.bw_hbm # bytes / s + _C = _scaled_machine.peak_flops # FLOPs / s (from knobs) + _W = _scaled_machine.bw_hbm # bytes / s (from knobs) _b_elem = model.bytes_per_elem # bytes / weight _kv_bpt_full = 2 * model.h_kv * model.d_head * model.bytes_per_elem * model.layers - _knee_now = knee_batch(_default_machine, model, _rf_skv) + _knee_now = knee_batch(_scaled_machine, model, _rf_skv) if _knee_now is None: _sub_bknee = ( f"{_b_star:.0f} / (1 − {_rf_skv:,}/{_l_star:,.0f}) = "