analytical-viz: Physical Layout — TTFT & TPOT full-model latency section

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 <noreply@anthropic.com>
2026-07-29 15:51:06 -07:00
parent c39bbb16aa
commit bc5e704572
2 changed files with 106 additions and 0 deletions
@@ -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:
@@ -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"