Repository navigation
[None][feat] DeepSeek V4 / DSpark optimizations on Rubin - #19961
Conversation
0f66060 to
7cfa431
Compare
e2696e8 to
a9426df
Compare
Barry-Delaney
left a comment
Barry-Delaney
left a comment
There was a problem hiding this comment.
Reviewed the code at a9426df, including the changes relative to the shared merge-back prerequisite. The SM107 Q-projection fusion gates and fallback paths, mixed context/generation handling, single mutable MLA output contract, non-aliasing O-LoRA scratch, Blackwell output consumers, and DFlash dummy-slot reset look sound. No blocking code findings.
Validation: source and call-path review; all 12 changed Python files parse and git diff --check passes. GPU execution was not performed locally; accepting the author-reported functional results. This is a code approval; CI and draft readiness remain with the author.
1a7be57 to
39a3e84
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #78047 [ run ] triggered by Bot. Commit: |
reasonsolo
added
the
api-compatible
39a3e84 to
296540e
Compare
|
/bot run --disable-fail-fast |
CUDA-graph padding and attention-DP dummy requests all map to the DFlash worker's dummy slot, which no request ever frees. If the generation step advanced that slot's context length like a real slot's, the dummies' drafter context would grow every step up to the context cap. Main zeroes it after each update (DFlashWorker._advance_ctx_len, NVIDIA#19988). Add a regression test that drives prepare_1st_drafter_inputs over a real slot and a padding row and checks that only the real slot advances. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
… kv-norm The q_b GEMM + RMSNorm + RoPE fusion writes the complete rotated Q, so the generation half runs mla_rope_generation with kv_only=True. kv_only needs kv_norm_weight, which is only passed when _fused_kv_norm_active is set; the q_b gate did not check that flag, so a generation batch where the fused kv-norm prologue is off (for example when _fused_q_rope_specs cannot build a spec) would trip "kv_only requires kv_norm_weight" in the op. Require _fused_kv_norm_active before taking the q_b fusion, matching the "both or neither" coupling the fused prologue already enforces. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…onfig contexts Runner applied ModelLoader.load_config_and_apply_defaults before entering scaled_from_ctx and vision_config_ctx, so with vision_config="none" the defaults were resolved from the composite checkpoint config (and its *ForConditionalGeneration class) while the model itself loads from the text-only config. Resolve the defaults and build the ModelLoader inside the same contexts that model_loader.load() runs under. Add a CPU unit test that stubs ModelLoader and checks the llm_args it receives: parallelism derived from a multi-rank mapping, a model default applied, and the defaults resolved through the text-only config loader. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…match AccuracyTask.evaluate() used a caller-supplied prepared_evaluator as is. In INTEGRATION_TEST mode the hypothesis test switches to one sample, but a full-size prepared evaluator would still run every sample. Raise a ValueError when the prepared evaluator's num_samples differs from the sample count of this evaluation. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…uest timeout - The generation node released its worker on the first failed health probe of the frontend, so one slow answer during the GSM8K run could tear the generation worker down mid-evaluation. Require ten consecutive failed probes before treating the frontend as shut down. - The OpenAI client timeout is in seconds; 1_800_000 let a stalled request hold a worker thread for weeks. Use 3600 s per attempt, which with the default two retries still ends inside the test timeout. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
The fusion gate takes the Q lengths from attn_metadata.cu_q_seqlens, which DeepSeek V4 metadata does not set, or from cu_seq_lens_cuda for pure generation batches, so only decode batches take the fused path. Correct the docstrings and comments that described context coverage; the context-side handling stays in place but is not exercised yet. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…utes The GB300 list gave test_gsm8k_dep4_megamoe_cutedsl a 240-minute TIMEOUT. Lower it to 30 minutes, which is enough for this test on GB300. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…ding test mla_rope_generation now passes q_rope_applied to the kernel after quant_scale_qkv. Cover it in the forwarded controls and compare the trailing arguments by the controls' length instead of a fixed count. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…ning cache key cu_q_seqlens and kv_cache_lengths are sized by the batch and were left as static dims, so every decode batch whose size differed from the tuning batch missed the autotuner cache and ran the fallback tactic. The kernel takes positions from position_ids and only stages the batch metadata, so constrain both inputs to a single tuning sequence (which maps them to -1 in the cache key) and fill them with valid values in the tuning pre-hook. The SM107 test now checks that a different number of sequences with the same M hits the tuned entry. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…oRA output - The q_b GEMM fusion gate now also requires the precomputed cu_seqlens hook, 128 local heads and the 448 + 64 head layout. Without the hook, generation would assert after the fused kernel had already run. - The fused epilogue output is always a contiguous [tokens, 1, hidden] buffer, so the non-direct path no longer returns a strided slice. O-LoRA always goes to a separate scratch. - Drop the unreachable copy branch in the single-phase shortcut. - Allocate the unit FP8 Q scale as a module buffer at init instead of lazily, so a first call under CUDA graph capture issues no H2D copy. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…d check every q_b tactic - An output_tensor without partition_id is a plain in-place GEMM only on runners that opt in (the MXFP8 in-place runner). Other runners reject it again, as before. - The MXFP8 in-place runner inherits the parent's kernel cache and derives its tuning config from the parent, adding only the output M constraint. - The q_b fusion test checks every registered tactic against the reference, not only the one the autotuner picked. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…ator has it The Evaluator base class has no num_samples, so a correctly built evaluator of another type would fail with a misleading ValueError. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
…V4 on SM107 DeepseekV4ForCausalLM.get_model_defaults only returned KV cache defaults, so on SM107 the FP8 block-scale linears fell back to the non-CuTe DSL GEMM path and prefill ran noticeably slower, which reduced high-concurrency disaggregated throughput. Return use_cute_dsl_blockscaling_mm=True as a model default when running on SM107, keeping the existing KV cache defaults logic. Update the unit tests accordingly and add an SM107 case. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
… KV slice A DSpark generation-only request allocates room for its draft tokens past the prompt. Add a case where only that reserve spills into a new block, and check that the block is dropped from the transfer table while the in-window SWA blocks keep their ordinals. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
… test Main replaced DFlashWorker._batch_to_slot with a DeviceCopyBuffer (_batch_to_slot_buffer). The test still set the old attribute, so prepare_1st_drafter_inputs read .tensor from None. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
83ed552 to
2237b43
Compare
|
/bot run |
|
PR_Github #80307 [ run ] triggered by Bot. Commit: |
|
PR_Github #80307 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #80341 [ run ] triggered by Bot. Commit: |
|
PR_Github #80341 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #80365 [ run ] triggered by Bot. Commit: |
|
PR_Github #80365 [ run ] completed with state Resource maintenance in LLM/main/L0_Test-x86_64-Multi-GPU #2991: B200 multi-GPU stages are paused to relieve the CI queue backlog (109 running / 459 pending on 2026-10-22). Owner: @sharabeshj-nvda. The CI team decides when to lift it.
|
|
/bot run --disable-fail-fast |
|
PR_Github #80386 [ run ] triggered by Bot. Commit: |
|
PR_Github #80386 [ run ] completed with state Resource maintenance in LLM/main/L0_Test-x86_64-Multi-GPU #2992: B200 multi-GPU stages are paused to relieve the CI queue backlog (109 running / 459 pending on 2026-10-22). Owner: @sharabeshj-nvda. The CI team decides when to lift it. |
Dev Engineer Review
The current diff only reformats imports in
cute_dsl_custom_ops.py,accuracy_core.py, andtest_fp8_block_scale_gemm.py. It does not change runtime behavior, APIs, test logic, or test-list configuration. Review finding counts and test results are unavailable.QA Engineer Review
One test file changed. Its edits only reformat imports; test functions, selectors, and test-list entries are unchanged. No integration test-list update applies. Coverage verdict: sufficient for the formatting-only diff. Test execution is not established.
Per-File QA Perspective
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py: Import formatting only. No observable behavior change requires QA verification.tests/integration/defs/accuracy/accuracy_core.py: Import formatting only. No evaluator behavior change requires QA verification.tests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py: Import formatting only. Existing FP8 GEMM and DSv4 fusion tests are unchanged; the file has no integration test-list entry.Description
This PR completes the DeepSeek V4 fast paths on Rubin (SM107) whose kernels and ops are already in main but
not wired in. It also adds the DSpark accuracy coverage used to validate them.
q_b GEMM + RMSNorm + RoPE fusion. On SM107, with FP8 block-scaled
q_b_projon the CuTe DSL path, theDeepSeek V4 MLA module now calls
cute_dsl_dsv4_qb_gemm_fused_rmsnorm_rope_quant.cu_q_seqlens, so contextand mixed batches take the unfused path.
mla_rope_generation(q_rope_applied). The context-side plumbing (the thop sparse-MLA context branchaccepting no
q_pewhenquant_q_bufferalready holds rotated Q) is in place for a later prefill pathbut is not exercised yet.
kv_a_projbefore the aux-stream compressor and indexer prework, which share one startevent.
TRTLLM_DISABLE_DSV4_Q_B_GEMM_NORM_ROPE_FUSION=1.Direct-output o_b projection. With the fused attention epilogue, both O-LoRA projections run inside the
MLA custom op and write the hidden states into its single mutable output, as the attention developer guide
describes.
into the output (new
cute_dsl_mxfp8_gemm_inplace_rubin).Tests.
test_deepseek_v4_q_b_fusion.py.test_mla_registry.py.TestDeepSeekV4ProDSpark::test_gsm8k_dep4_megamoe_cutedsl: NVFP4 export, MegaMoE CuTe DSL, DEP4,GSM8K + acceptance length. Runs on l0_gb300_multi_gpus.
TestDeepSeekV4ProDSparkMultinode::test_gsm8k_1p1d_dep4: two-node 1P1D disaggregated GSM8K. QAmulti-node list, because it needs one pytest task per node.
SM107 fixes.
use_cute_dsl_blockscaling_mmto on for SM107, so its FP8 block-scale linears take the CuTe DSL GEMM path.sm_100fon SM107 again (new3rdparty/patches/deepgemm_sm107_arch_family.patch). After the DeepGEMM bump, they were built assm_107f. The bundled CUTLASS has no SM107 target, so the TMA descriptor prefetch hit a device assert, for example in the DSA indexer'sfp8_paged_mqa_logits.DFlash dummy slot. CUDA-graph padding and attention-DP dummy requests share the DFlash worker's dummy slot,
which is never freed. Main now keeps that slot's drafter context empty (
_advance_ctx_len, #19988); this PR addsa regression test that drives the drafter-input step over a real slot and a padding row.
Not included. Block-scaled FP8 projections for DSpark: main does not run the DSpark CuTe DSL kernels on
SM107 yet, so that work needs a separate enablement PR first. The SM107 block-scale default above covers the DeepSeek V4 main model only.
Test Coverage
tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_q_b_fusion.pytests/unittest/_torch/attention/test_mla_registry.pytests/unittest/_torch/thop/parallel/test_fp8_block_scale_gemm.py(SM107)tests/unittest/_torch/modeling/test_modeling_deepseekv4.py(model defaults, including SM107)tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py::test_generation_step_keeps_the_dummy_slot_context_emptytests/unittest/disaggregated/test_cache_reuse_adapter.py::TestCreateKvSliceBlockSpan::test_dspark_gen_init_draft_reserve_block_is_trimmedaccuracy/test_llm_api_pytorch.py::TestDeepSeekV4ProDSpark::test_gsm8k_dep4_megamoe_cutedslaccuracy/test_disaggregated_multinode.py::TestDeepSeekV4ProDSparkMultinode::test_gsm8k_1p1d_dep4PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.