Remove P1 from audit — multi-GPU layout is correct for the reference script
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.
This commit is contained in:
@@ -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 |
|
||||
|
||||
Reference in New Issue
Block a user