Repository navigation
[Perf] DFLASH: route draft bf16 GEMMs through tuned dispatch on GB300 - #45876
Conversation
nvpohanh
left a comment
nvpohanh
left a comment
There was a problem hiding this comment.
[by Claude Code] This review is done by agents. Please feel free to disagree if they do not make sense.
[by Claude Code] Severity: perf | Confidence: Medium
Issue
The accuracy table cannot show whether the new GEMM entries are correct. In the Muse Spark Open run, these shapes run only in the DFLASH draft, and target verification corrects draft errors. A wrong tactic would leave GSM8K and AIME26 unchanged and only lower the acceptance length.
Why it matters
The end-to-end table measures time per decode cycle, so "throughput scales with 1 / cycle" assumes the acceptance length did not change. The same (M, N, K) keys also run in the target forward of Llama-3-70B-class BF16 models at TP4/TP8, where a wrong result changes outputs.
Fix
Add the mean acceptance length for main and this PR from the same GB300 TP4/TP8 runs.
|
/rerun-test test/registered/e2e/models/test_muse_spark_12_b300_pd.py test/registered/cp/test_glm53_flash_cp.py test/registered/e2e/models/test_muse_spark_unified_memory.py test/registered/e2e/models/test_glm53_flash_b200.py test/registered/unit/spec/test_dflash_logits.py test/registered/e2e/models/test_nvidia_nemotron_3_nano.py test/registered/core/test_basic_sanity_dflash.py test/registered/spec/dflash/test_dflash.py test/registered/spec/test_spec_mixed_chunk.py test/registered/e2e/speculative/test_dflash_domino.py test/registered/spec/test_gemma4_dflash_31b_extra.py test/registered/spec/dflash/test_muse_glimmer_dflash_assistant_gsm8k.py test/registered/unit/spec/test_dflash_domino.py test/registered/kernels/ops/speculative/test_dflash_domino.py test/registered/unit/spec/test_dflash_extra_buffer_lazy.py test/registered/unit/spec/test_dflash_overlap_hostsync.py test/registered/unit/spec/test_oot_dflash_hooks.py test/registered/spec/dspark/test_dspark_stacked_ctx_kv_parity.py test/registered/unit/layers/test_sharded_loader_parallel_groups.py test/registered/unit/mem_cache/test_prefix_valid_fp8_pool.py test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py test/registered/unit/model_executor/model_runner_components/test_spec_aux_hidden_state.py test/registered/unit/model_executor/test_spec_aux_hidden_state_resolver.py test/registered/unit/layer_boundary/test_glm5_next_dflash_capture.py test/registered/unit/mem_cache/test_kv_cache_configurator_draft_gate.py test/registered/unit/mem_cache/test_kv_index_translator_draft_disposition.py test/registered/unit/mem_cache/test_fused_draft_placement.py test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py test/registered/unit/spec/test_weight_update_runners.py test/registered/unit/spec/test_block_verification_args.py test/registered/kernels/ops/speculative/test_block_verification.py test/registered/unit/spec/test_decode_bookkeeping_ownership.py test/registered/unit/server_args/test_unified_memory_spec_gate.py test/registered/scheduler/test_min_free_slots_delayer.py test/registered/unit/npu/graph_runner/test_npu_graph_runner.py test/registered/unit/layers/muse_spark/test_muse_spark_runner_dispatch.py test/registered/kernels/ops/attention/test_muse_spark_qk_norm_rope_store.py test/registered/kernels/ops/communication/test_muse_spark_moe_finalize_all_reduce.py test/registered/kernels/ops/layernorm/test_muse_spark_norm_ops.py test/registered/kernels/ops/layernorm/test_muse_spark_rmsnorm_nvfp4_quant.py test/registered/kernels/ops/moe/test_muse_spark_routing.py test/registered/kernels/ops/moe/test_moe_wna16_marlin.py test/registered/unit/function_call/test_function_call_parser.py test/registered/unit/parser/test_reasoning_parser.py test/registered/unit/entrypoints/openai/test_serving_responses.py test/registered/unit/entrypoints/openai/test_serving_responses_stream.py test/registered/unit/multimodal/test_processor_async_call_sites.py test/registered/attention/unittests/dense/test_fa3.py test/registered/attention/unittests/dense/test_fa4.py test/registered/attention/unittests/dense/test_flashinfer.py test/registered/attention/unittests/dense/test_triton.py test/registered/attention/unittests/gdn/test_flashinfer.py test/registered/attention/unittests/gdn/test_triton.py test/registered/attention/unittests/kda/test_triton.py test/registered/attention/unittests/lightning/test_triton.py test/registered/attention/unittests/mamba/test_mamba2.py test/registered/attention/unittests/mla/test_triton.py test/registered/attention/unittests/swa/test_flashinfer.py test/registered/attention/unittests/swa/test_triton.py |
|
🚀 🚀 🚀 🚀 🚀 🚀 ⛔ |
|
@claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Nothing blocking. The comments below are optional suggestions. There is no need to push a fix for them before merging.
Beyond the inline finding, I also checked two things in dflash.py and found no problems. The ReplicatedLinear swap keeps the kernel_projection.weight parameter name and is built without a quant_config, so it stays unquantized. weight[:num_org] slices leading rows of a contiguous weight, so the view passed to bf16_gemm_dispatch is still contiguous.
Extended reasoning...
Findings were reported inline. This note only records the weight-loading and layout checks in dflash.py that were looked at and turned up no problems.
Motivation
At low concurrency the DFLASH draft model runs its bf16 linears at small M: one block of
num_draft_tokenstokens per request.Two of its paths bypass SGLang's tuned bf16 GEMM dispatch and always use cuBLAS:
DFlashGroupedConv.kernel_projectionis a plainnn.Linear;_project_candidate_logitscallstorch.matmuldirectly.Most of the draft's other GEMM shapes have no GB300 tuned entry and fall back to cuBLAS, which is well below the memory bandwidth
limit at these M sizes. The exception is TP4 qkv at M = 8 and 16, which already used the split-K tactics of the "GB300 TP16" block.
Modifications
models/dflash.pykernel_projectionbecomesReplicatedLinear, so it goes throughUnquantizedLinearMethod.applyand the tuned dispatch.bf16_gemm_dispatchwhen the backend is cutedsl/gemv and the weight is bf16 on CUDA; otherwise itkeeps
torch.matmul.layers/quantization/unquant.py: GB300 split-K / warp split-K entries at M <= 32, TP4 and TP8.kernels/ops/gemm/cutedsl_bf16_gemm.py: TGV M ranges for three shapes at M = 17..32.Each entry was kept only if it beats the current kernel in all three timings of the micro-benchmark below.
The tables are keyed by (M, N, K) only, so the entries are not draft-specific:
models at TP4 and TP8. With the default
autobackend on SM10x those models switch kernels at M <= 32 as well. Only GB300 wasmeasured; GB200 was not.
precedence). They are faster in all three timings; see the micro-benchmark.
precompile_splitk_tactics()(only Qwen4-Exp calls it) compiles every table entry, so it now builds 19 more tactics at startup.Accuracy Tests
Speed Tests and Profiling
Micro-benchmark (GB300, CUDA graphs, cold L2; us per call, after-PDL timing)
End-to-end (Muse Spark Open NVFP4, GB300, DFLASH 8 drafts, FlashInfer autotune off on both sides)
Mean decode cycle (target verify + draft step) from nsys, main -> this PR; FlashInfer autotune off on both sides.
*Includes about ±1% MoE routing variation between runs; the GEMMs this PR changes account for -52.6 us/cycle (-0.40%).
No change at c >= 5 (draft GEMM M > 32, no new entries).
Decode throughput scales with 1 / cycle, so the gain is about +1.1% (TP4) and +1.4% (TP8) at c=1.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #40499863108
Latest PR Test (Extra): ❌ Run #40499863678
Latest PR Test (AMD ROCm 10): ❌ Run #40499862483