diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index 46140be..dbd0ad3 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -1130,6 +1130,99 @@ with tab_layout: st.dataframe(pd.DataFrame(ffn_stage_shape_rows(cfg)), width='stretch', hide_index=True) + # Row H: Full-model latency (TTFT & TPOT) ────────────────────── + st.divider() + st.subheader("Full-model latency (TTFT & TPOT)") + st.caption( + "End-to-end per-token latency (TPOT) is one decode step summed " + "across all layers. TTFT is one prefill pass (all layers) for a " + "prompt of the length below. **Assumes PP=1** (no PP-bubble " + "modeling); comm is the same as your Per-stage latency tab." + ) + + _tt_cols = st.columns([1, 1, 2]) + with _tt_cols[0]: + _prompt_len = st.number_input( + "Prompt length (for TTFT)", min_value=1, max_value=2_000_000, + value=int(s_kv), step=128, key="_tt_prompt_len", + help="T_q used for the prefill pass. Defaults to sidebar S_kv.", + ) + with _tt_cols[1]: + _tt_n_out = st.number_input( + "Output tokens (for end-to-end)", min_value=1, max_value=100_000, + value=100, step=1, key="_tt_n_out_tokens", + help="Number of tokens to generate. E2E = TTFT + (N-1) × TPOT.", + ) + + # Build two temporary FullConfigs: one for decode (TPOT), one for + # prefill (TTFT). Everything else — sharding, B, machine — stays + # the same as the sidebar cfg. + _tpot_topo = dataclasses.replace(topo, mode="decode", s_kv=int(s_kv)) + _ttft_topo = dataclasses.replace(topo, mode="prefill", + s_kv=int(_prompt_len)) + _tpot_cfg = FullConfig(model=model, topo=_tpot_topo, + machine=_default_machine) + _ttft_cfg = FullConfig(model=model, topo=_ttft_topo, + machine=_default_machine) + + _decode_attn_s = sum(s.visible_s for s in all_stages(_tpot_cfg)) + _decode_ffn_s = sum(s.visible_s for s in all_ffn_stages(_tpot_cfg)) + _prefill_attn_s = sum(s.visible_s for s in all_stages(_ttft_cfg)) + _prefill_ffn_s = sum(s.visible_s for s in all_ffn_stages(_ttft_cfg)) + + _tpot_per_layer_s = _decode_attn_s + _decode_ffn_s + _ttft_per_layer_s = _prefill_attn_s + _prefill_ffn_s + _tpot_total_s = _tpot_per_layer_s * model.layers + _ttft_total_s = _ttft_per_layer_s * model.layers + _e2e_total_s = _ttft_total_s + (int(_tt_n_out) - 1) * _tpot_total_s + _throughput_tokens_per_s = (max(1, int(b_batch)) / _tpot_total_s + if _tpot_total_s > 0 else 0.0) + + _kpi = st.columns(4) + _kpi[0].metric("TPOT (per token)", + f"{_tpot_total_s * 1e3:.2f} ms", + help="One decode-step time × all layers.") + _kpi[1].metric(f"TTFT (prompt={_prompt_len:,})", + f"{_ttft_total_s * 1e3:.2f} ms", + help="One prefill pass × all layers.") + _kpi[2].metric(f"E2E for {_tt_n_out} tokens", + f"{_e2e_total_s * 1e3:.2f} ms", + help="TTFT + (N-1) × TPOT.") + _kpi[3].metric(f"Throughput / replica (B={b_batch})", + f"{_throughput_tokens_per_s:.1f} tok/s", + help="B ÷ TPOT — decode-side throughput this replica.") + + _tt_rows = [ + {"Component": "Attention (decode, TPOT contribution)", + "Per layer (µs)": round(_decode_attn_s * 1e6, 3), + "All layers (ms)": round(_decode_attn_s * model.layers * 1e3, 3)}, + {"Component": "FFN (decode, TPOT contribution)", + "Per layer (µs)": round(_decode_ffn_s * 1e6, 3), + "All layers (ms)": round(_decode_ffn_s * model.layers * 1e3, 3)}, + {"Component": "Attention (prefill, TTFT contribution)", + "Per layer (µs)": round(_prefill_attn_s * 1e6, 3), + "All layers (ms)": round(_prefill_attn_s * model.layers * 1e3, 3)}, + {"Component": "FFN (prefill, TTFT contribution)", + "Per layer (µs)": round(_prefill_ffn_s * 1e6, 3), + "All layers (ms)": round(_prefill_ffn_s * model.layers * 1e3, 3)}, + {"Component": "TPOT total (per token)", + "Per layer (µs)": round(_tpot_per_layer_s * 1e6, 3), + "All layers (ms)": round(_tpot_total_s * 1e3, 3)}, + {"Component": "TTFT total (prefill)", + "Per layer (µs)": round(_ttft_per_layer_s * 1e6, 3), + "All layers (ms)": round(_ttft_total_s * 1e3, 3)}, + ] + st.dataframe(pd.DataFrame(_tt_rows), width='stretch', hide_index=True) + + st.caption( + f"Sidebar params in use: **{model.name}** on **{topo.cp}**×" + f"**{topo.tp}**×**{topo.pp}** (CP·TP·PP), B=**{b_batch}**, " + f"layers=**{model.layers}**. TTFT uses prefill-mode formulas " + f"(compute-bound, T_q={_prompt_len:,}); TPOT uses decode-mode " + f"(memory-bound, T_q=1). Cross-check per-stage breakdown on " + f"the *Per-stage latency* tab." + ) + # ── TAB 2: Memory breakdown ───────────────────────────────────── with tab_memory: diff --git a/tests/analytical_visualization/test_stage_shapes.py b/tests/analytical_visualization/test_stage_shapes.py index 338e1c5..d6062ac 100644 --- a/tests/analytical_visualization/test_stage_shapes.py +++ b/tests/analytical_visualization/test_stage_shapes.py @@ -152,3 +152,16 @@ def test_shapes_section_is_on_physical_layout_tab(): assert "ffn_stage_shape_rows(cfg)" not in stages_block, ( "per-stage FFN shape table must not remain on Per-stage latency tab" ) + + +def test_ttft_tpot_section_wired_on_layout_tab(): + """The 'Full-model latency (TTFT & TPOT)' subheader must live on + tab_layout, before tab_memory. Guard the section presence + placement.""" + from pathlib import Path + src = ( + Path(__file__).parent / "app.py" + ).resolve().read_text(encoding="utf-8") + assert src.count('"Full-model latency (TTFT & TPOT)"') == 1 + pos = src.index('"Full-model latency (TTFT & TPOT)"') + tab_memory_pos = src.index("with tab_memory:") + assert pos < tab_memory_pos, "TTFT section must live inside tab_layout"