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>
This commit is contained in:
@@ -1130,6 +1130,99 @@ with tab_layout:
|
|||||||
st.dataframe(pd.DataFrame(ffn_stage_shape_rows(cfg)),
|
st.dataframe(pd.DataFrame(ffn_stage_shape_rows(cfg)),
|
||||||
width='stretch', hide_index=True)
|
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 ─────────────────────────────────────
|
# ── TAB 2: Memory breakdown ─────────────────────────────────────
|
||||||
with tab_memory:
|
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, (
|
assert "ffn_stage_shape_rows(cfg)" not in stages_block, (
|
||||||
"per-stage FFN shape table must not remain on Per-stage latency tab"
|
"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"
|
||||||
|
|||||||
Reference in New Issue
Block a user