diff --git a/tests/analytical_visualization/app.py b/tests/analytical_visualization/app.py index 4cf4be4..6f0257a 100644 --- a/tests/analytical_visualization/app.py +++ b/tests/analytical_visualization/app.py @@ -3498,3 +3498,133 @@ with tab_prules: "(DeepSeek-V3, vLLM PP-on-non-NVLink) show that no single " "recipe fits all workloads." ) + + st.divider() + + # ── Section 11: TP vs CP — the shared-budget tradeoff ───────── + st.markdown("### 11. TP vs CP: the shared-budget tradeoff") + st.markdown( + "**Fixing `TP × CP` doesn't equalize the two configurations.** " + "Larger TP additionally divides *model-weight* traffic; larger " + "CP only divides the *sequence* dimension. That's why (TP=8, " + "CP=2) often beats (TP=2, CP=8) during decode — not because " + "CP comm is heavier, but because CP doesn't touch weights." + ) + + st.markdown("#### 11a. Same TP·CP = 16 → different weight sharding") + _cptp_config = [ + {"Config": "TP=8, CP=2", + "Weight shard / GPU": "W / 8", + "Context shard": "L / 2", + "KV shard / GPU": "KV / 16", + "Weight bytes / PE (Llama-like layer)": "~214 MB", + "KV read / PE @ 1M ctx": "~256 MB"}, + {"Config": "TP=2, CP=8", + "Weight shard / GPU": "W / 2", + "Context shard": "L / 8", + "KV shard / GPU": "KV / 16", + "Weight bytes / PE (Llama-like layer)": "~856 MB (4× more)", + "KV read / PE @ 1M ctx": "~256 MB"}, + ] + st.dataframe(pd.DataFrame(_cptp_config), width='stretch', + hide_index=True) + st.caption( + "Concrete Llama-like layer: `d_model=8192`, `d_ff=28672`, " + "`H_kv=8`, `d_head=128`, BF16 → weights per layer ≈ 1.7 GB; " + "KV read at 1M context ≈ 4.1 GB per layer. **Both configs " + "give the same KV sharding**, but TP=8 reads 4× fewer weight " + "bytes per PE per token." + ) + + st.markdown("#### 11b. What each dimension actually shards") + _cptp_shard = [ + {"Axis": "TP", + "Shards weights?": "✓ (by head)", + "Shards KV?": "✓ (by KV head)", + "Shards context (attention math)?": "✗"}, + {"Axis": "CP", + "Shards weights?": "✗", + "Shards KV?": "✓ (by position)", + "Shards context (attention math)?": "✓"}, + {"Axis": "PP", + "Shards weights?": "✓ (by layer)", + "Shards KV?": "✓ (by layer)", + "Shards context (attention math)?": "✗"}, + {"Axis": "DP", + "Shards weights?": "✗ (or optimizer only in training w/ ZeRO)", + "Shards KV?": "✓ (across independent sequences)", + "Shards context (attention math)?": "✗"}, + ] + st.dataframe(pd.DataFrame(_cptp_shard), width='stretch', + hide_index=True) + st.caption( + "Key asymmetry: **CP doesn't touch the FFN weight footprint " + "at all**. FFN is a huge chunk of decode weight-fetch cost, " + "so if the FFN block is dominant, TP is the only axis that " + "helps." + ) + + st.markdown("#### 11c. When to prefer TP vs CP") + _cptp_prefer = [ + {"Situation": "Weight fetch + FFN dominates step time", + "Prefer": "Higher TP", + "Why": "TP divides weight bandwidth linearly; CP doesn't touch it"}, + {"Situation": "KV read + attention dominates", + "Prefer": "Higher CP", + "Why": "CP shards positions; attention scales O(S_kv)"}, + {"Situation": "Moderate context, FFN-heavy model", + "Prefer": "Higher TP", + "Why": "Attention is a small fraction; FFN weight matters more"}, + {"Situation": "Very long context (≫ L*), attention-heavy", + "Prefer": "Higher CP", + "Why": "KV read is the wall; CP lets multiple PEs share it"}, + {"Situation": "TPOT-tight decode", + "Prefer": "Higher TP (up to ~n_kv_heads)", + "Why": "Divides the big cost (weights); collectives small at T_q=1"}, + {"Situation": "Long-prompt prefill under TTFT SLO", + "Prefer": "Higher CP with `kv` ring variant", + "Why": "CP hops overlap with attention compute; bytes stay bounded"}, + ] + st.dataframe(pd.DataFrame(_cptp_prefer), width='stretch', + hide_index=True) + + st.markdown("#### 11d. CP communication profile (prefill vs decode)") + _cptp_comm = [ + {"Mode": "Prefill (T_q = prompt length, large)", + "CP comm bytes/hop": + "~O(S_local · H_kv · d_h) with kv-variant (bounded)", + "Overlap potential": "Good — hides in attention compute", + "Real cost": "Moderate", + "Notes": "kv-variant: bytes independent of T_q → good scaling; " + "qoml-variant grows with T_q, worst case"}, + {"Mode": "Decode (T_q = 1)", + "CP comm bytes/hop": "O(d_model) for merge (small)", + "Overlap potential": "Poor — only 1 query token", + "Real cost": "Sync-latency-dominated", + "Notes": "~2·(CP−1)·α per layer; α × hops beats bytes at T_q=1"}, + ] + st.dataframe(pd.DataFrame(_cptp_comm), width='stretch', + hide_index=True) + st.caption( + "So CP comm isn't necessarily *heavier* than TP comm — it's a " + "different profile. In decode, CP suffers from per-hop α " + "latency (small bytes, many syncs); in prefill, CP with the " + "right ring variant can be very cheap and overlappable." + ) + + st.divider() + st.markdown("#### 11e. Decode cost decomposition (mental model)") + st.markdown( + "> `T_decode ≈ T_weight_read(TP) + T_KV_read(TP, CP) + " + "T_TP_comm(TP) + T_CP_merge(CP)`\n\n" + "- **Larger TP**: cuts `T_weight_read`, cuts `T_KV_read` (by " + "KV head), grows `T_TP_comm`, may make GEMMs too small.\n" + "- **Larger CP**: cuts `T_KV_read` (by position) and speeds up " + "attention math, **doesn't touch weights or FFN**, grows " + "`T_CP_merge` sync overhead.\n\n" + "So the real question isn't 'which comm is heavier' — it's " + "**'which cost is dominating T_decode right now, and which " + "axis attacks it?'** For weight-fetch-bound decode with " + "moderate context, larger TP is the right lever. For " + "KV-bandwidth-bound decode with extreme context, larger CP is." + ) diff --git a/tests/analytical_visualization/test_stage_shapes.py b/tests/analytical_visualization/test_stage_shapes.py index 3f9d3d1..34f1f9d 100644 --- a/tests/analytical_visualization/test_stage_shapes.py +++ b/tests/analytical_visualization/test_stage_shapes.py @@ -176,3 +176,12 @@ def test_parallelism_rules_tab_registered(): ).resolve().read_text(encoding="utf-8") assert src.count('"Parallelism rules"') == 1 assert "with tab_prules:" in src + + +def test_tp_vs_cp_section_present(): + """Section 11 'TP vs CP: the shared-budget tradeoff' must appear.""" + from pathlib import Path + src = ( + Path(__file__).parent / "app.py" + ).resolve().read_text(encoding="utf-8") + assert src.count("TP vs CP: the shared-budget tradeoff") == 1