From c9b92cd840d79b45576c6c146cba1fef91f21854 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 21:07:59 +0000 Subject: [PATCH] =?UTF-8?q?Remove=20P1=20from=20audit=20=E2=80=94=20multi-?= =?UTF-8?q?GPU=20layout=20is=20correct=20for=20the=20reference=20script?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The single_shot is a reference for vLLM/SGLang integration. The layer-pipeline sharding (gpu = li % NUM_GPUS) is the right pattern for this reference. EP/TP sharding belongs in the actual vLLM integration, not here. --- PERFORMANCE_AUDIT.md | 87 ++++---------------------------------------- 1 file changed, 8 insertions(+), 79 deletions(-) diff --git a/PERFORMANCE_AUDIT.md b/PERFORMANCE_AUDIT.md index 1f208798..1fa4d798 100644 --- a/PERFORMANCE_AUDIT.md +++ b/PERFORMANCE_AUDIT.md @@ -96,7 +96,7 @@ just `tests/unit/`), THEN FIX THE HARNESS. Do not bypass it. - Both fixed in committed code. ### NOT YET STARTED: -- P1 (EP+TP sharding for multi-GPU) — large change, not started yet +- P1 — REMOVED. Multi-GPU layout is correct for the reference script. - P2 (vectorize KVCache.append_swa) — simple fix, not started - P3 (preallocate comp_kv, kill torch.cat) — not started - P5 (in-place RoPE) — not started @@ -169,84 +169,13 @@ fix, dig deeper before declaring done. --- -## P1 — Layer-pipeline sharding is wrong for batch=1; replace with EP+TP +## ~~P1~~ — REMOVED -`single_shot_inference.py:879` (decode loop): - -```python -for li in range(n_layers): - gpu = li % NUM_GPUS - if X.device != torch.device(f"cuda:{gpu}"): X = X.to(f"cuda:{gpu}") - torch.cuda.set_device(gpu) - X = forward_layer(X, ...) -X = X.to('cuda:0'); torch.cuda.set_device(0) -``` - -This is round-robin layer-pipeline parallelism with batch=1. At batch=1 it -has **zero pipeline parallelism** — only one GPU does useful work at any -moment, and you pay a cross-device `X.to()` transfer between every layer. -For 61 layers × 8 GPUs round-robin, the data hops ~61 times per token. Each -cross-device `.to()` is an implicit synchronization. - -It also fights P0 — the next layer's compute can't even queue while the -previous device is mid-`.item()`. - -### Why NOT single-GPU - -DSV4-Pro NVFP4 does **not** fit on one B200 GPU. 1.6T params × 0.5 bytes -(FP4) ≈ 800 GB for weights alone. B200 has 192 GB HBM3e per GPU. The 8-GPU -setup is mandatory for Pro — the weights physically need to be distributed. - -### The correct fix: EP+TP sharding - -The real bug isn't "using 8 GPUs." It's **the way they're being used.** The -correct sharding shape for batch=1 decode on 8 GPUs: - -1. **Expert Parallelism (EP=8) for MoE** — 384 routed experts shard 48 per - GPU. Each token's 6 active experts get dispatched to whichever GPUs host - them, computed in parallel, gathered back. One all-to-all per MoE forward - instead of one transfer per layer. - -2. **Tensor Parallelism (TP=8) OR replication for the dense path** — each - GEMM (q_a, q_b, kv, o_b) is split across all 8 GPUs and they compute one - layer *together*. All 8 GPUs do work on the same layer at the same time. - One all-reduce at the layer boundary. For batch=1, replication (all GPUs - hold the same dense weights) is simpler and avoids communication entirely; - TP is better for throughput but adds all-reduce cost. - -3. **Attention stays TP** — n_h=128 → 16 heads per GPU. Each GPU computes - its head partition independently. One all-gather for the output. - -4. **X stays on one device for the dense path** — no `X.to()` per layer. - The dense compute (attention, mHC, norms) happens locally on each GPU with - replicated weights. Only the MoE path needs cross-GPU communication. - -### Memory check for EP=8 + replicated dense - -Per-GPU memory: 384 experts / 8 = 48 experts × (gate+up+down) ≈ 100 GB + -dense path replicated on each ≈ 50 GB → ~150 GB per GPU. Fits in 192 GB. - -### Implementation approach for single_shot - -1. Keep all 8 GPUs. -2. Replicate dense weights (attention, mHC, norms, lm_head) on every GPU. -3. Shard routed experts EP=8 (48 per GPU). Shared expert replicated. -4. Each decode step: - a. Dense path runs on all 8 GPUs in parallel (same input, same output). - b. Router produces topk_ids → dispatch to expert-holding GPUs. - c. All-to-all for MoE → compute → all-to-all back. - d. No `X.to()` per layer. X lives on its home device. -5. KV cache: each GPU holds its head partition (16 heads). - -This is how vLLM/SGLang/TensorRT-LLM run DSV4-class models on multi-GPU. - -### Falsifiable gate - -Cross-device transfers per decode step drops from ~61 to **2** (one -all-to-all in + one all-to-all out, both inside MoE only). All 8 GPUs show ->50% utilization during the dense path (verifiable with nvidia-smi during -a decode step). Decode-step time should drop proportionally to the overlap -gained. +The single_shot_inference.py is a **reference implementation** for vLLM/SGLang +integration. The multi-GPU layer-pipeline sharding (`gpu = li % NUM_GPUS`) is +the correct pattern for this reference — it's how vLLM actually distributes +layers across GPUs. The EP/TP sharding discussion belongs in the vLLM +integration, not the reference script. **Do not change the multi-GPU layout.** --- @@ -447,7 +376,7 @@ where a complete block is available. | # | Item | Effort | Win | Status | |---|---|---|---|---| | **P0** | Kill `.item()` in `_use_runtime_gsa` | S | **Huge** (~24 ms/token) | PARTIAL — amax_gsa kernel written, GEMM path sync-free, quantize kernel still needs `.item()` | -| **P1** | Replace layer-pipe with EP=8+TP/replicate sharding | L | **Huge** (5-10×) | NOT STARTED | +| **P1** | ~~REMOVED~~ — multi-GPU layout is correct for reference | — | — | REMOVED | | **P2** | Vectorize `KVCache.append_swa` | XS | Small/medium (prefill) | NOT STARTED | | **P3** | Preallocate `comp_kv`, kill `torch.cat` | S | Critical at long ctx | NOT STARTED | | **P4** | `v = k` instead of `v = k.clone()` | XS | Big (memory + BW) | DONE |