[None][perf] Fuse DeepSeek-V4 Indexer Q projection with CuTe DSL#16532
Open
mingyangHao wants to merge 1 commit into
Open
[None][perf] Fuse DeepSeek-V4 Indexer Q projection with CuTe DSL#16532mingyangHao wants to merge 1 commit into
mingyangHao wants to merge 1 commit into
Conversation
mingyangHao
requested review from
JunyiXu-nv,
joyang-nv,
kris1025,
longlee0622,
yunruis and
zhaoyangwang-nvidia
July 17, 2026 07:01
mingyangHao
changed the title
Contributor
📝 WalkthroughWalkthroughAdds 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 fp8_quantize_1x128_cutedsl_ue8m0
participant CuteDSLIndexerQBlackwellRunner
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 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}
Contributor
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
Comment on lines
+56
to
+62
| # The swapped-A/B CuTe DSL fusion wins through M=8 with the production main-Q | ||
| # sibling present, while the native CuTe fusion wins from M=128. Keep the | ||
| # production chain in the measured gap and above the validated sweep range; | ||
| # the TRT-LLM autotuner selects the concrete CuTe tactic within each route. | ||
| DEEPSEEK_V4_INDEXER_Q_SMALL_M_FUSION_MAX_TOKENS = 8 | ||
| DEEPSEEK_V4_INDEXER_Q_FULL_FUSION_MIN_TOKENS = 128 | ||
| DEEPSEEK_V4_INDEXER_Q_FULL_FUSION_MAX_TOKENS = 16384 |
Contributor
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
Enable the fused native path from M=16 as specified.
The 128 lower bound routes M=16–127 through the unfused chain, contradicting the stated native-fusion range and excluding benchmarked cases.
Proposed fix
-DEEPSEEK_V4_INDEXER_Q_FULL_FUSION_MIN_TOKENS = 128
+DEEPSEEK_V4_INDEXER_Q_FULL_FUSION_MIN_TOKENS = 16📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| # The swapped-A/B CuTe DSL fusion wins through M=8 with the production main-Q | |
| # sibling present, while the native CuTe fusion wins from M=128. Keep the | |
| # production chain in the measured gap and above the validated sweep range; | |
| # the TRT-LLM autotuner selects the concrete CuTe tactic within each route. | |
| DEEPSEEK_V4_INDEXER_Q_SMALL_M_FUSION_MAX_TOKENS = 8 | |
| DEEPSEEK_V4_INDEXER_Q_FULL_FUSION_MIN_TOKENS = 128 | |
| DEEPSEEK_V4_INDEXER_Q_FULL_FUSION_MAX_TOKENS = 16384 | |
| # The swapped-A/B CuTe DSL fusion wins through M=8 with the production main-Q | |
| # sibling present, while the native CuTe fusion wins from M=128. Keep the | |
| # production chain in the measured gap and above the validated sweep range; | |
| # the TRT-LLM autotuner selects the concrete CuTe tactic within each route. | |
| DEEPSEEK_V4_INDEXER_Q_SMALL_M_FUSION_MAX_TOKENS = 8 | |
| DEEPSEEK_V4_INDEXER_Q_FULL_FUSION_MIN_TOKENS = 16 | |
| DEEPSEEK_V4_INDEXER_Q_FULL_FUSION_MAX_TOKENS = 16384 |
🤖 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/attention_backend/sparse/deepseek_v4/deepseek_v4.py`
around lines 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.
Comment on lines
+186
to
+207
| self.indexer_q_fusion = indexer_q_fusion | ||
| self.is_gated = is_gated_activation(activation_type) | ||
| # Precompute per-activation flags as plain Python bools so the epilogue | ||
| # dispatch can use cutlass.const_expr(...) on them (like is_gated). An | ||
| # inline enum comparison inside @cute.jit is not folded as a constant. | ||
| self._act_is_swiglu = activation_type == ActivationType.Swiglu | ||
| self._act_is_gelu = activation_type == ActivationType.Gelu | ||
| self._act_is_identity = activation_type == ActivationType.Identity | ||
| # Host-side guard (raise is not allowed inside the @cute.kernel epilogue): | ||
| # only SwiGLU and GELU(tanh) are wired in the fused epilogue. | ||
| if not (self._act_is_swiglu or self._act_is_gelu): | ||
| if not (self._act_is_swiglu or self._act_is_gelu or self._act_is_identity): | ||
| raise NotImplementedError( | ||
| f"Fused epilogue activation {activation_type} not implemented " | ||
| f"(only Swiglu and Gelu are wired)." | ||
| f"(only Swiglu, Gelu, and the indexer-Q identity epilogue are wired)." | ||
| ) | ||
| if self.indexer_q_fusion: | ||
| if activation_type != ActivationType.Identity: | ||
| raise ValueError("indexer_q_fusion requires ActivationType.Identity") | ||
| if sf_vec_size != 32 or mma_tiler_mn[1] != 128: | ||
| raise ValueError( | ||
| "indexer_q_fusion requires MXF8 sf_vec_size=32 and a 128-column MMA tile" | ||
| ) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Reject identity activation when indexer-Q fusion is disabled.
ActivationType.Identity is accepted unconditionally, but its epilogue stores raw acc_vec_up without applying alpha. A non-fused identity configuration therefore silently produces incorrect output.
Proposed fix
if self.indexer_q_fusion:
if activation_type != ActivationType.Identity:
raise ValueError("indexer_q_fusion requires ActivationType.Identity")
if sf_vec_size != 32 or mma_tiler_mn[1] != 128:
raise ValueError(
"indexer_q_fusion requires MXF8 sf_vec_size=32 and a 128-column MMA tile"
)
+elif self._act_is_identity:
+ raise ValueError(
+ "ActivationType.Identity is only supported with indexer_q_fusion"
+ )Also applies to: 1606-1613
🤖 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/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_act_fusion.py`
around lines 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.
Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
mingyangHao
force-pushed
the
user/mingyangh/dsv4-indexer-q-fusion
branch
from
July 17, 2026 07:23
13b204f to
73ca517
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.Summary by CodeRabbit
New Features
Tests