Repository navigation
[Fix] dsv4: Build the DSpark draft block metadata without host sequence lengths - #46891
Conversation
…ce lengths Since sgl-project#45925 the DSpark draft batch carries no host sequence lengths when the draft backend opts out of them, which the DeepSeek-V4 backend does on CUDA. Its draft-block prep then read them back with seq_lens.tolist(): a device sync on every draft step, and inside the metadata glue capture an error that turned the glue graph off. The prep needs only the batch size and the device tensors, so derive the extend lengths from those. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
/rerun-test test_ragged_verify.py test_dsv41_draft_metadata_glue.py test_dsv41_target_verify_metadata_glue.py test_forked_draft_metadata_prep.py test_dspark_draft_block.py test_fused_draft_placement.py test_dspark_block_accept_estimator.py test_dspark_confidence_metrics.py test_dspark_dp_tier.py test_dspark_draft_path_default.py test_dspark_info_dumper.py test_dspark_scheduler.py test_dspark_sps_table.py test_dspark_sts.py test_dspark_draft_graph.py test_dspark_kernel_parity.py test_dspark_stacked_ctx_kv_parity.py test_ragged_verify_backend_capability.py test_dspark_step_inputs_once.py test_dspark_sampling_accept_graph.py test_basic_sanity_dspark.py registered/spec/dspark/test_dspark_dp_spec_prefill_coordination.py test_dspark_minimax_mxfp8_stacked_kv.py test_deepseek_v41_flash_gb300_dp2ep2.py test_deepseek_v41_flash_gb300_tp2ep2.py test_deepseek_v41_flash_gb300_encoder_bcg.py test_deepseek_v41_flash_gb300_trtllm_bounded_bcg.py test_deepseek_v41_pd_cp_dspark.py test_deepseek_v4_flash_fp4_b200.py test_deepseek_v4_flash_fp4_megamoe_b200.py |
|
🚀 🚀 🚀 |
|
🚀 🚀 🚀 🚀 🚀 🚀 🚀 |
|
The four
None of the four logs has The |
Motivation
Since #45925, the DSpark draft batch carries no host sequence lengths when the draft attention backend sets
needs_cpu_seq_lens = False. That is right fortrtllm_mha, but the DeepSeek-V4 backend also sets it to False under DSpark on CUDA. Its draft-block prep,init_forward_metadata_dspark_draft_block, then read the lengths back withseq_lens.tolist(), which caused two problems:Cannot copy between CPU and CUDA tensors during CUDA graph capture, and the runner falls back to eager draft metadata prep for the rest of the run:seq_lens.tolist()again.On CUDA the draft-block prep does not need host lengths at all:
[block_size] * bs.Modifications
compute_uniform_extend_lengthsacceptsseq_lens_cpu=None. It then takes the batch size fromseq_lensand returns no host copy of the extended lengths.Nonethrough instead of callingseq_lens.tolist().Accuracy Tests
The draft metadata comes from the same device tensors as before. Accept length is unchanged with real drafts: DEP2 with real acceptance gives a median of 2.0 to 2.2 on both arms at every batch size.
Speed Tests and Profiling
GB300, DeepSeek-V4.1-Flash, DSpark block 5, one node per arm, one run each. Main is
cb55d77d37, and the fix is the same commit plus this change. sa-bench uses input length 1024. The output length is 4096 for TP4 and 1024 for DEP2. Forward occupancy comes fromSGLANG_ENABLE_METRICS_DEVICE_TIMER=1.TP4, EP1, simulated acceptance 3.51:
DEP2 with DP attention and MegaMoE, real acceptance:
The gain shows where the step is short and host-bound: at TP4 batch 1, output throughput is 17% higher and the median ITL is 15% lower. DEP2 is within run-to-run noise.
Checklist
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #40592800125
Latest PR Test (Extra): ❌ Run #40592800066
Latest PR Test (AMD ROCm 10): ❌ Run #40592799993