[None][perf] Fuse DeepSeek-V4 Indexer Q projection with CuTe DSL#16532
Conversation
mingyangHao
changed the title
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds a Blackwell CuTe DSL path for FP8 indexer-Q quantization, fused RoPE and FP4 output, including CUDA operators, persistent GEMM kernels, runtime tactic selection, DeepSeek-V4 integration, and bit-exact tests. ChangesIndexer-Q CuTe DSL fusion
Estimated code review effort: 5 (Critical) | ~90+ minutes Sequence Diagram(s)sequenceDiagram
participant DeepseekV4Indexer
participant cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell
participant CuteDSLIndexerQBlackwellRunner
participant fp8_quantize_1x128_cutedsl_ue8m0
participant BlackwellGEMM
DeepseekV4Indexer->>cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell: request fused Q projection
cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell->>CuteDSLIndexerQBlackwellRunner: select tactic and dispatch
CuteDSLIndexerQBlackwellRunner->>fp8_quantize_1x128_cutedsl_ue8m0: create FP8 input and UE8M0 scales
CuteDSLIndexerQBlackwellRunner->>BlackwellGEMM: launch GEMM with RoPE and FP4 epilogue
BlackwellGEMM-->>DeepseekV4Indexer: return packed FP4 Q and scales
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py (1)
734-742: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the fake kernel’s return annotation.
Proposed fix
`@torch.library.register_fake`("trtllm::fp8_quantize_1x128_cutedsl_ue8m0") -def _(input: torch.Tensor): +def _(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:As per coding guidelines, “Annotate every function” and “prefer built-in generic types.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py` around lines 734 - 742, Add a return type annotation to the fake kernel function registered as trtllm::fp8_quantize_1x128_cutedsl_ue8m0, using the appropriate built-in generic tuple type describing its float8 tensor and uint8 tensor results.Source: Coding guidelines
tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py (1)
496-498: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExercise the TVM-FFI branch used in production.
deepseek_v4.pyLine 1170 always passesuse_tvm_ffi=True, but this test forcesFalse; the production stream and pointer invocation path remains untested. Coverage is insufficient—add at least one small-M and one native-M TVM-FFI case in this test file. Please also annotate the new test signature.Proposed parametrization
-@pytest.mark.parametrize("num_tokens", [1, 4, 5, 8, 128]) -def test_cute_dsl_fp8_indexer_q_gemm_rope_fp4_matches_unfused(num_tokens): +@pytest.mark.parametrize( + ("num_tokens", "use_tvm_ffi"), + [ + (1, False), + (4, False), + (5, False), + (8, False), + (128, False), + (1, True), + (128, True), + ], +) +def test_cute_dsl_fp8_indexer_q_gemm_rope_fp4_matches_unfused( + num_tokens: int, use_tvm_ffi: bool +) -> None: ... - runner = CuteDSLIndexerQBlackwellRunner(use_tvm_ffi=False) + runner = CuteDSLIndexerQBlackwellRunner(use_tvm_ffi=use_tvm_ffi) ... - use_tvm_ffi=False, + use_tvm_ffi=use_tvm_ffi,As per path instructions, test reviews must assess coverage and provide a concrete test location; follow up in
tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py.Also applies to: 534-534, 554-562
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py` around lines 496 - 498, Extend test_cute_dsl_fp8_indexer_q_gemm_rope_fp4_matches_unfused to cover the production TVM-FFI path by adding parameterized cases for at least one small-M and one native-M configuration, ensuring use_tvm_ffi=True reaches the stream and pointer invocation logic. Annotate the updated test signature with the required typing, and preserve the existing unfused comparison coverage.Sources: Coding guidelines, Path instructions
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py (1)
3659-3660: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFully type the new runner and custom-op interfaces.
unique_id,_ptr,tactic,dtype,**kwargs, and the fake implementation lack precise annotations, while several signatures use legacyList/Tuple. Please add Python 3.10+ built-in annotations and document the public custom-op arguments.As per coding guidelines, “Annotate every function” and “prefer built-in generic types,” with documented public function arguments.
Also applies to: 3706-3718, 3747-3769, 3909-3917, 3933-3941
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py` around lines 3659 - 3660, Fully annotate _prepare_cutedsl_indexer_q_tuning_inputs and the related runner, custom-op, and fake implementations, including precise types for unique_id, _ptr, tactic, dtype, and **kwargs. Replace legacy List/Tuple annotations with Python 3.10+ built-in generics throughout the affected signatures, and document arguments exposed by public custom-op interfaces.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py`:
- Around line 56-62: Update DEEPSEEK_V4_INDEXER_Q_FULL_FUSION_MIN_TOKENS from
128 to 16 so the fused native path is selected for M=16 and above, while
preserving the existing small-M threshold and maximum-token boundary.
In
`@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_act_fusion.py`:
- Around line 186-207: Update the activation validation in the constructor
around _act_is_identity so ActivationType.Identity is accepted only when
indexer_q_fusion is enabled. Reject identity activation with indexer_q_fusion
disabled before epilogue execution, while preserving the existing
identity-specific tile and scale-factor validation for fused configurations.
---
Nitpick comments:
In `@tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py`:
- Around line 734-742: Add a return type annotation to the fake kernel function
registered as trtllm::fp8_quantize_1x128_cutedsl_ue8m0, using the appropriate
built-in generic tuple type describing its float8 tensor and uint8 tensor
results.
In `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py`:
- Around line 3659-3660: Fully annotate _prepare_cutedsl_indexer_q_tuning_inputs
and the related runner, custom-op, and fake implementations, including precise
types for unique_id, _ptr, tactic, dtype, and **kwargs. Replace legacy
List/Tuple annotations with Python 3.10+ built-in generics throughout the
affected signatures, and document arguments exposed by public custom-op
interfaces.
In `@tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py`:
- Around line 496-498: Extend
test_cute_dsl_fp8_indexer_q_gemm_rope_fp4_matches_unfused to cover the
production TVM-FFI path by adding parameterized cases for at least one small-M
and one native-M configuration, ensuring use_tvm_ffi=True reaches the stream and
pointer invocation logic. Annotate the updated test signature with the required
typing, and preserve the existing unfused comparison coverage.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 61f1fdaa-05a5-4946-98cd-340034cb844d
📒 Files selected for processing (11)
cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cucpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.hcpp/tensorrt_llm/thop/fp8Quantize.cpptensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.pytensorrt_llm/_torch/custom_ops/cpp_custom_ops.pytensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_act_fusion.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.pytensorrt_llm/_torch/modules/linear.pytests/integration/test_lists/test-db/l0_dgx_b200.ymltests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
13b204f to
73ca517
Compare
|
/bot run |
|
PR_Github #60158 [ run ] triggered by Bot. Commit: |
|
PR_Github #60158 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60204 [ run ] triggered by Bot. Commit: |
|
PR_Github #60204 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60236 [ run ] triggered by Bot. Commit: |
|
PR_Github #60236 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60303 [ run ] triggered by Bot. Commit: |
|
PR_Github #60861 [ run ] triggered by Bot. Commit: |
zhaoyangwang-nvidia
left a comment
zhaoyangwang-nvidia
left a comment
There was a problem hiding this comment.
Approve with little comments.
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
9b3726e to
c74f213
Compare
|
/bot kill |
|
PR_Github #60884 [ kill ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #60861 [ run ] completed with state |
|
PR_Github #60884 [ kill ] completed with state |
|
PR_Github #60889 [ run ] triggered by Bot. Commit: |
|
PR_Github #60889 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61146 [ run ] triggered by Bot. Commit: |
|
PR_Github #61146 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61213 [ run ] triggered by Bot. Commit: |
|
PR_Github #61213 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61307 [ run ] triggered by Bot. Commit: |
|
PR_Github #61307 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61463 [ run ] triggered by Bot. Commit: |
|
PR_Github #61463 [ run ] completed with state |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py`:
- Around line 3762-3770: Update _fallback_tactic to validate native
configurations with the same kernel_class.can_implement logic used by
TunableRunner.get_valid_tactics. Iterate through _native_tactics and return the
first feasible native tactic for the given m, n, and k; preserve the existing
small-M selection and ensure a safe fallback is handled if no native tactic is
implementable.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5c13b42b-1dc7-4544-ab1f-31c89ceb165c
📒 Files selected for processing (2)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.pytensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
|
/bot skip --comment 'accidentally update branch' |
|
PR_Github #61512 [ skip ] triggered by Bot. Commit: |
|
PR_Github #61512 [ skip ] completed with state |
5 participants
Description
Paths compared
mlaRoPEInplaceKernel+fusedCatFp4Kernel+ output scale copies.kernel + output scale copies. It uses the swapped small-M kernel for
M<=8and the native CuTe fused kernel for every
M>=16.There is no fallback to the original chain at any measured M. NSYS topology
validation found no standalone RoPE or FP4-concat kernel in any forced-fused
candidate graph.
Results
Times are the average of the forward/reverse capture medians. The percentage
column is the range from the two independent capture orders; positive means
forced-fused CuTe is faster.
Test Coverage
PR 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.Dev Engineer Review
FP8 quantization / scale-layout correctness
fp8_blockscale_quant_packed.cuintroducesOutputCuteDslSfto switch UE8M0 scale output format:false): packs four UE8M0 scale bytes into a 32-bitpackedvalue (with K-bound masking) and writespacked_scale_output.true): bypassespacked; replicates each UE8M0 scale byte into a replicated 32-bit pattern and writes into the SM100 CuTe 128×4 swizzled layout, emitting only lanes wherelane_id % 8 == 0and zero-filling padded rows.OutputCuteDslSf=falsefor the default launcher andtruefor the CuTe launcher, passingpadded_masscale_leading_dim_uint32.New SM100-only FP8 quant op (CuTe DSL / UE8M0)
fp8_blockscale_quant_packed.haddslaunch_fp8_quantize_1x128_cutedsl_bf16_e4m3(...)withpadded_mfor the swizzled scale layout.thop/fp8Quantize.cppaddsfp8_quantize_1x128_cutedsl_ue8m0(Tensor)(SM100-family + input validation, includingk % 128 == 0), returning(fp8 values, ue8m0 scale tensor).cpp_custom_ops.pyadds the corresponding fake/meta registration fortrtllm::fp8_quantize_1x128_cutedsl_ue8m0, computing padded scale shapes fromm(padded to 128) and thek-dependent padded scale dimensions.DeepSeek-V4 indexer-Q fused routing
deepseek_v4.pyrefactors Q handling to unify projection→RoPE→quantization in_project_and_quantize_q, using the CuTe fused optorch.ops.trtllm.cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwellwhen the CuTe/DSv4 constraints are satisfied (MXFP4_BLOCKWISE cache, non-Neox, exact head/rope dims, CuTe cutedsl scale attributes present, SM100f, bf16 path, andqrbf16).CuTe DSL runner + operator (Blackwell path)
cute_dsl_custom_ops.pyadds:CuteDSLIndexerQBlackwellRunner(SM100-family gated) with autotuning that:position_ids = 0during tuning to keep RoPE cache lookups consistent.tactic == -1.cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell(...)plus matching fake/meta registration.wrapper_indexer_q_swap_abvswrapper_indexer_q) and returns packed FP4 bytes (int8) and output scale (int32).Blackwell persistent GEMM epilogue integration
dense_blockscaled_gemm_act_fusion.pyanddense_blockscaled_gemm_persistent.pyadd anindexer_q_fusionmode:indexer_scale_tensor,position_ids_tensor,cos_sin_cache_tensor).wrapper_indexer_q_swap_ab.Weight transform / buffers for fused layout
tensorrt_llm/_torch/modules/linear.pyextends the FP8 resmoothing/layout flow to generate CUTEDSL-specific indexer-Q buffers (indexer_q_weight_scale_cutedsl,indexer_q_alpha_cutedsl) when the new fused layout is enabled.CI stability / failure correlation
QA Engineer Review
Test files touched
tests/integration/test_lists/test-db/l0_dgx_b200.ymlunittest/_torch/attention/sparse/test_cpp_custom_ops.py::test_cute_dsl_fp8_indexer_q_gemm_rope_fp4_matches_unfusedsystem_gpu_count: gte: 4, lte: 4, Ubuntu,*b200*GPUs,stage: pre_merge,backend: pytorch,orchestrator: mpi.tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.pytest_cute_dsl_fp8_indexer_q_gemm_rope_fp4_matches_unfused(num_tokens)num_tokens(and performs bit-for-bit equality vs an unfused reference chain).Coverage / verdict
test-dbintegration list? Yes (single new entry inl0_dgx_b200.yml).