Repository navigation
[DSpark] Fix draft CUDA graph stream explosion - #40658
Conversation
|
/rerun-test test_broadcast_cuda_graph.py test_kimi_k3_b300.py test_kimi_k3_b300_low_latency.py |
|
Results for 🚀 🚀 |
|
/rerun-test test_kimi_k3_b300.py test_kimi_k3_b300_low_latency.py test_gemma4_dflash_31b_extra.py test_dflash_domino.py test_qwen35_fp4_mtp.py |
|
Results for 🚀 🚀 🚀 ⛔
Please provide the full path, e.g. |
|
/rerun-test test/registered/e2e/speculative/test_dflash_domino.py |
|
Results for 🚀 |
(cherry picked from commit bc22e1d)
Motivation
DSpark draft CUDA graph replay expands to seven streams on B300 with TP8/DCP1. Its seven dependent token broadcasts go through ProcessGroupNCCL even while the graph capture context has enabled PyNccl.
Modifications
Use the enabled PyNccl communicator for CUDA broadcasts, keeping broadcasts on the current capture stream. Preserve the torch.distributed fallback when PyNccl is disabled or unavailable. Keep the existing typed PyNccl broadcast unchanged. Audited SRT callers use supported integer or floating-point payloads; no bool/complex broadcast support is added.
Accuracy Tests
Existing Kimi K3 B300 correctness tests are being retriggered on the updated head. Dummy weights and simulated acceptance do not validate model accuracy. Pre-commit checks passed for the runtime change.
Speed Tests and Profiling
8× B300, TP8, 60k input / 1k output, DSpark block 7, simulated acceptance 5.5, dummy MXFP4 target and dummy draft weights, FP8 E4M3 KV, FP32 SSM, CuTe DSL MLA, FlashInfer MXFP4 MoE, CUDA device connections 8.
Earlier-revision measurements: the following N=5 table was collected at
ff17b2e04b, which included byte broadcasting. These timings have not been rerun after restoring the existing typed broadcast ate91dd97ff0. The earlier DCP8 traces showed one draft stream and four target streams before and after.N=5 unprofiled samples per variant, topology, and batch size:
Decode throughput is effectively unchanged in these runs; full-batch steady-state throughput gives the same conclusion. Mean total latency changes are +0.10%, +0.09%, +0.12%, and −0.26%, respectively. The last mean includes one slower baseline prefill, retained in the results; no end-to-end speedup is established.
Baseline is main
c4d3770a68; patched isff17b2e04b. Each variant/topology uses one fresh server, an excluded full-shape BS1/BS16 warmup pair, then five alternating measured pairs. Order: baseline DCP1, patched DCP1, patched DCP8, baseline DCP8. Client/server seeds are 42, corresponding acceptance lengths match, and all requests complete the requested output length. Samples share a server launch, so this is limited regression evidence rather than a general performance guarantee. DCP8 usesfi_a2awith replicated Q projection. The client capacity check is skipped because it reads physical capacity while DCP exposes logical capacity; server memory settings are unchanged.Checklist
CI States
Latest PR Test (Base): ❌ Run #35692480924
Latest PR Test (Extra): ❌ Run #35692480469
Latest PR Test (AMD ROCm 10): ❌ Run #35692480624