tp>1 prefill is host-bound: per-token fp4 GEMM + per-layer H2D activation staging
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 423
- Forks
- 53
- Avg merge
- 20h 26m
- Merged PRs (30d)
- 310
Description
Problem
TP>1 prefill on the 27B NVFP4 checkpoint is host-bound, not GPU-bound. The fp4 device GEMM (vt_cuda_mlp_shard_runT_fp4 / MlpGuActFp4 / MlpDownFp4, src/vt/cuda/nccl_communicator.cu) launches once per token instead of once per batch, and every fp4 layer stages its activations through host memory (Sm70Fp8PackQpn). Decode hides this (CUDA-graph replay saturates the GPU); prefill exposes it.
Measured with nsys on tp=2 (fresh binary, exclusive $HOME/gpu.lock, trace /tmp/prof_tp2.nsys-rep):
MlpGuActFp4launches with identical duration (~845 us) and identical gridX (8704) at T=5 and T=1 — per-token, NOT batched. 5 tokens cost ~5x a decode-step MLP.cudaMemcpyAsyncH2D is 82.6% of CUDA-API time (15.9 s; 33.6 GB / 9,195 copies; 19.6 GB >10 MB during the request window in ~4.1 s).Sm70Fp8PackQpnis 25.4% of kernel time (1.74 s, 2,560 launches) — staging activations through host.- Warm prefill: GPU0 2.33 s kernel work over 4.37 s wall = ~53% busy; 459 >2 ms inter-launch stalls +
cudaStreamSynchronize~140 ms.
Not a defect in the shipped fp8w row (#2058, BACKEND-DISTRIBUTED-TP-CLEAN); it is the next wave in the same mirror line.
Fix directions (to be scoped in the spec)
- Batched T-token fp4 GEMM: gridX scales with T so a prefill batch costs one launch set, not T of them.
- Device-resident activations across the fp4 layers: kill the per-layer host round-trip.
- Fold
Sm70Fp8PackQpnstaging on-device.
Axis
Prefill wall (TTFT / prefill-only throughput at tp=1 vs tp=2), with the decode line held flat. Established token-exact baseline ' Paris.\nThe capital of Germany is Berlin.' (10 tokens) stays the correctness gate.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in src/vt/cuda/nccl_communicator.cu, tracing vt_cuda_mlp_shard_runT_fp4, MlpGuActFp4, MlpDownFp4, and Sm70Fp8PackQpn while reproducing the tp=2 prefill profile. Compare T=1 and T=5 launches and the token-exact baseline; done means batched fp4 work and device-resident activation handling improve prefill wall time without changing the decode line or correctness.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- distributed-systems, performance
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100