[#13061][bugfix] Add rotate_activation to DSA fused_cat_fp8 path

[waynehacking8]

Summary

Fixes #13061 — the DSA sparse attention indexer's _prep_q_or_k method calls fused_cat_fp8 to concatenate [qk_pe, qk_nope] and quantize to FP8 in a single fused kernel, but skips rotate_activation (Hadamard rotation) which DSA requires between concatenation and quantization.

• Root cause: fused_cat_fp8 is a single CUDA kernel that does cat+quant atomically — there's no way to inject the rotation in between.
• Fix: Unfuse the FP8 path into torch.cat → rotate_activation → fp8_quantize_1x128. The scale transpose logic follows the reference implementation in test_fused_cat_fp8.py.
• Backward-compatible: rotate_activation gracefully returns identity when fast-hadamard-transform is not installed (existing warning).
• FP4 gap: The FP4 path (fused_cat_fp4) has the same issue but needs a standalone FP4 quantize-only kernel to unfuse — noted in the docstring as a follow-up.

Code flow (before → after)

Before (buggy):

qk_pe, qk_nope → fused_cat_fp8 (cat + quant) → fp8_out, scale
                   rotate_activation skipped

After (fixed):

qk_pe, qk_nope → torch.cat → rotate_activation → fp8_quantize_1x128 → fp8_out, scale
                               Hadamard rotation applied

Test plan

• rotate_activation fallback: returns identity when fast-hadamard-transform is missing — no regression for existing setups

Not run in my environment (no GPU build / no DeepSeek V3.2 weights) — flagged for CI / reviewer validation:

• test_fused_cat_fp8.py: the fused_cat_fp8 CUDA kernel is unchanged by this PR (the FP8 path now reaches the quantizer via fp8_quantize_1x128_sf_transpose after the Hadamard rotation), so the kernel test is unaffected.
• DSA sparse attention end-to-end with DeepSeek V3.2 (requires model weights + fast-hadamard-transform).

Closes #13061

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-compatible or api-breaking. For api-breaking, include BREAKING in 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

• Refactor
  • Optimized sparse attention backend with improved quantization implementation and enhanced tensor processing operations.

[coderabbitai]

📝 Walkthrough

Walkthrough

This change fixes the FP8 quantization path in DSA (Dense Sparse Attention) by inserting Hadamard rotation before quantization. The method now concatenates Q/K tensors, applies rotation, quantizes with fp8_quantize_1x128, and normalizes the scale tensor shape before returning.

Changes

FP8 Quantization with Hadamard Rotation

Layer / File(s)	Summary
FP8 Q/K preprocessing with rotation before quantization tensorrt_llm/_torch/attention_backend/sparse/dsa.py	Indexer._prep_q_or_k FP8 path now applies rotate_activation after concatenating qk_pe/qk_nope, switches from fused_cat_fp8 to torch.ops.trtllm.fp8_quantize_1x128, and post-processes scale tensor with padding, reshape, and transpose when 1D; FP4 path unchanged.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main change: adding rotate_activation to the DSA fused_cat_fp8 path, with a bugfix label and reference to issue #13061.
Linked Issues check	✅ Passed	The code changes directly address issue #13061 by unfusing the FP8 path to include rotate_activation between concatenation and quantization, exactly as required.
Out of Scope Changes check	✅ Passed	All changes are focused on fixing the DSA FP8 path; no unrelated modifications are present beyond the scope of issue #13061.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Description check	✅ Passed	The PR description comprehensively explains the bug, root cause, solution, code flow before/after, and test plan, following the template structure with all required sections completed.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (1)

tensorrt_llm/_torch/attention_backend/sparse/dsa.py (1)

2345-2350: 💤 Low value

Consider extracting scale normalization to a shared helper.

The scale normalization logic here duplicates fp8_quantize_1x128_sf_transpose from tensorrt_llm/quantization/utils/fp8_utils.py. To reduce duplication, consider either:

1. Adding a use_ue8m0 parameter to the existing fp8_quantize_1x128_sf_transpose function, or
2. Extracting just the scale normalization into a small helper that both call sites can use.

Not blocking since the current implementation is correct and matches the reference pattern from context snippet 2.

🤖 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/dsa.py` around lines 2345 -
2350, The scale normalization block in dsa.py duplicates logic from
fp8_quantize_1x128_sf_transpose in tensorrt_llm/quantization/utils/fp8_utils.py;
refactor by extracting the normalization into a shared helper (e.g.,
normalize_fp8_scale(scale, combined_2d, head_dim, padded_block=128)) or by
adding a use_ue8m0 parameter to fp8_quantize_1x128_sf_transpose and reusing it
from this module; update the code in
tensorrt_llm/_torch/attention_backend/sparse/dsa.py to call the shared helper
(or the extended fp8_quantize_1x128_sf_transpose) instead of reimplementing the
slicing, padding and transpose logic so both call sites use the same
implementation.

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

Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py`:
- Around line 2345-2350: The scale normalization block in dsa.py duplicates
logic from fp8_quantize_1x128_sf_transpose in
tensorrt_llm/quantization/utils/fp8_utils.py; refactor by extracting the
normalization into a shared helper (e.g., normalize_fp8_scale(scale,
combined_2d, head_dim, padded_block=128)) or by adding a use_ue8m0 parameter to
fp8_quantize_1x128_sf_transpose and reusing it from this module; update the code
in tensorrt_llm/_torch/attention_backend/sparse/dsa.py to call the shared helper
(or the extended fp8_quantize_1x128_sf_transpose) instead of reimplementing the
slicing, padding and transpose logic so both call sites use the same
implementation.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c1e815b4-dcc0-46d7-a462-b9b9228512b2

📥 Commits

Reviewing files that changed from the base of the PR and between 45e2523 and 2a39f63.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/attention_backend/sparse/dsa.py

[waynehacking8]

Status check: this PR has been mergeable and review-ready since 06-09 (coderabbit pass addressed at open time). Could a maintainer trigger CI (/bot run) when convenient? Happy to rebase or adjust if anything in the recent refactors moved — the fix still applies cleanly to current main.

[karljang]

@waynehacking8,
Thank you for your contributions! Before we proceed, could you please address the pre-commit check failures?

[waynehacking8]

Thanks @karljang — addressed in e570480: the _prep_q_or_k import and fp8_quantize_1x128_sf_transpose call now match the repo's isort/yapf style (backslash-continued single import; closing paren on the last-argument line), which was the pre-commit failure. I also resolved the two open items in the PR checklist (the kernel-test one is unaffected since fused_cat_fp8 is unchanged; the DeepSeek V3.2 end-to-end run needs model weights + a GPU I don't have locally, so I've flagged it for CI/reviewer validation). Ready for /bot run whenever you'd like.

[PerkzZheng]

Add @lfr-0531 to review as well. thanks

[yuxianq]

This PR basically revert #11899, which removes Hadamard rotation.
@kaiyux as the author, could you tell us why we can remove Hadamard rotation here?

[waynehacking8]

Replied in detail in a top-level comment — TL;DR: this isn't a revert. #11899 was a [perf] fusion, and fusedCatFp8.cu only does concat + FP8 quantize (no Hadamard), so the rotation was dropped unintentionally (that's #13061). I've restored the pre-#11899 cat → rotate → quantize path verbatim (same rotate_activation, same fp8_utils.fp8_quantize_1x128_sf_transpose). Happy to fold the rotation into the kernel as a perf follow-up.

[yuxianq]

I find a comment about it: #11899 (comment), which says:

Yes, I intentionally removed it. hadamard_transform was not used, and previous verification shows that it's safe to remove with no accuracy drop.

@waynehacking8 Can you show us a case that accuracy drop for a benchmark without Hadamard rotation?

[waynehacking8]

@yuxianq @PerkzZheng Thanks for the careful review. I looked into the history — this isn't an intentional revert of a design choice, it's restoring a regression:

#11899 was [perf]-only: it fused cat+quantize into fusedCatFp8.cu. That kernel does only concat (Stage 1) + FP8 quantize (Stage 2) — no Hadamard rotation — and its dsa.py diff deleted q_or_k = rotate_activation(q_or_k) without re-adding it anywhere. So the rotation was dropped as a side effect of the fusion, not by design; #13061 is exactly that silent regression.

This PR restores the pre-#11899 fp8 path. I've aligned the fp8 branch to mirror the original implementation verbatim — maybe_compiled_cat → rotate_activation → fp8_utils.fp8_quantize_1x128_sf_transpose — so it's byte-for-byte the code #11899 replaced (only fused_cat_fp4 is kept for the fp4 case). Same Hadamard, same quantizer ⇒ numerically the original intended behavior. Rotate-before-FP8 matters for quant accuracy (Hadamard smooths outliers); DSV3.2 is robust enough to mask it, but it's still a correctness drop.

The cost is that un-fusing gives up #11899's kernel speedup. Happy to follow up by folding the rotation into fusedCatFp8.cu (and the FP4 path, which still lacks it) to recover perf — separate PR or in this one, whichever you and @kaiyux prefer.

[waynehacking8]

@yuxianq Thanks — kaiyux's #11899 comment is the decisive context. I worked through it and his call holds.

The indexer logit is out_ij = Σ_h ReLU(q[i,h]·k[j]) · weights[i,h], and _prep_q_or_k applies the same orthonormal Hadamard H to both q and k (dsa.py:2376-2377). Since (Hq)·(Hk) = q·k, the per-head dot product — and therefore the ReLU, the weights, and the summed logit — are all unchanged in exact arithmetic. The Hadamard's only real effect is lowering FP8 quantization error on that dot product before it's computed (the QuaRot rationale), and kaiyux's verification shows that doesn't move the indexer's top-k selection.

I can't produce a config with an actual accuracy drop (no GPU / DSV3.2 weights on my side), and given the cancellation I don't expect one for the FP8 path. So I'll defer to the verification and close this — #13061 is really a question, not a regression. If FP4's larger quant error ever turns out to move top-k, that'd be the place to revisit.

Thanks all for the quick reviews.