From bc5e704572f8b1b819ef946d00cc874ea878bb56 Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Wed, 29 Jul 2026 15:51:06 -0700 Subject: [PATCH] =?UTF-8?q?analytical-viz:=20Physical=20Layout=20=E2=80=94?= =?UTF-8?q?=20TTFT=20&=20TPOT=20full-model=20latency=20section?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New section at the bottom of the Physical Layout tab that sums the per-stage latencies into end-to-end metrics for the current sharding: - **TPOT** = one decode-step time × N_layers (per output token) - **TTFT** = one prefill pass time × N_layers (for a prompt of the input length, computed in prefill mode) - **E2E** = TTFT + (N_out − 1) × TPOT - **Throughput/replica** = B / TPOT (tokens/sec) Two number inputs: prompt length (default = sidebar S_kv) and target N_out tokens (default 100). Four KPI cards + a breakdown table with per-layer and all-layer µs/ms for attention/FFN in both modes. Both TTFT and TPOT are computed by cloning the sidebar TopologyConfig with mode swapped ('decode' for TPOT, 'prefill' for TTFT with T_q = prompt length), then reusing all_stages / all_ffn_stages — no new math. Caveat noted in caption: PP=1 assumption (no microbatch overlap / bubble modeling). Sanity: Llama 3 8B on CP=2/TP=8/default SIP → TPOT ≈ 4.2 ms/token, TTFT ≈ 1.1 s for an 8k prompt. Cross-checkable on the Per-stage latency tab. Smoke test in test_stage_shapes asserts the section title appears exactly once and sits inside tab_layout (before tab_memory). Co-Authored-By: Claude Opus 4.7 --- tests/analytical_visualization/app.py | 93 +++++++++++++++++++ .../test_stage_shapes.py | 13 +++ 2 files changed, 106 insertions(+) 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"