[TRTLLM-11285][feat] Fuse indexer wk + weights_proj into single GEMM in TF32 for DS-V3.2

[peihu-nv]

Prerequisite

• Depends on [TRTLLM-9521][feat] Unfuse indexer.wk from attention GEMM for DS-V3.2 NVFP4 #11989

Summary by CodeRabbit

• Refactor

  • Optimized sparse attention computation through improved weight fusion mechanism for better performance
  • Simplified sparse attention initialization and forward computation flow

• Tests

  • Updated sparse attention tests to align with refactored weight fusion approach

Description

• Fuse the two DS-V3.2 indexer projections (wk and weights_proj) into a single FP32 cuBLAS GEMM (TF32 tensor cores on Ampere+), saving one kernel launch
• Force wk to FP32 with no quantization to match weights_proj, enabling the fusion
• FP32 (TF32) chosen over FP16 based on accuracy analysis in TRTLLM-10283

Test Coverage

• test_sparse_mla_forward.py
• test_short_seq_mha.py
• NVFP4 inference test (TP=4, DS-V3.2 checkpoint)

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)

• 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.

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes refactor the Indexer weight fusion mechanism by introducing a post-load-weights method that fuses wk and weights_proj tensors. The explicit indexer_k parameter is removed from DSATrtllmAttention forward signature and indexer calls, with weight fusion now occurring during post-load initialization. Related model and test code is updated accordingly.

Changes

Cohort / File(s)	Summary
Indexer Weight Fusion tensorrt_llm/_torch/attention_backend/sparse/dsa.py	Introduces post_load_weights() method and _fused_wk_wp_weight attribute for fusing wk and weights_proj weights. Updates Indexer.__init__ to set wk and weights_proj to float32 with no quantization. Modifies forward computation to use fused cuBLAS GEMM, removing explicit indexer_k parameter from DSATrtllmAttention signature.
Attention Module Updates tensorrt_llm/_torch/models/modeling_deepseekv3.py, tensorrt_llm/_torch/modules/attention.py	Removes post_load_weights indexer fusion logic from DeepseekV32Attention. Updates DeepseekV32Attention.kv_a_proj_with_mqa output features calculation to exclude indexer.head_dim. Modifies forward_impl_with_dsa to split kv_a_proj_with_mqa output into three components instead of four, removing indexer_k handling.
Test Updates tests/unittest/_torch/attention/sparse/test_short_seq_mha.py, tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py	Adds post_load_weights() calls on indexer components after initialization. Removes indexer_k argument from indexer method calls, relying on internal default/top-k computation instead.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title clearly and specifically describes the main change: fusing two indexer projections (wk + weights_proj) into a single GEMM in TF32 for DeepSeek-V3.2, matching the core objective.
Description check	✅ Passed	The PR description clearly explains the objective (fuse indexer projections into FP32 GEMM), provides test coverage examples, and includes a reference to a related TRTLLM ticket for accuracy analysis.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

📝 Coding Plan

• Generate coding plan for human review comments

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

Tip

CodeRabbit can use your project's `pylint` configuration to improve the quality of Python code reviews.

Add a pylint configuration file to your project to customize how CodeRabbit runs pylint.

[peihu-nv]

/bot run

[tensorrt-cicd]

PR_Github #39188 [ run ] triggered by Bot. Commit: c7049ce Link to invocation

[tensorrt-cicd]

PR_Github #39188 [ run ] completed with state SUCCESS. Commit: c7049ce
/LLM/main/L0_MergeRequest_PR pipeline #30442 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[peihu-nv]

/bot run --disable-fail-fast

[peihu-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #39282 [ run ] triggered by Bot. Commit: 6f744c1 Link to invocation

[tensorrt-cicd]

PR_Github #39282 [ run ] completed with state SUCCESS. Commit: 6f744c1
/LLM/main/L0_MergeRequest_PR pipeline #30540 completed with status: 'SUCCESS'

CI Report

Link to invocation

[pengbowang-nv]

LGTM. Leaving a minor question.