perf: add fused norm + RoPE + uniform fp8 store for TRT-LLM DSv4 sparse attention

[akhilg-nv]

Motivation

Depends on #30805, improves e2e perf for dsv4 with trtllm sparse attention

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #32541328105
Latest PR Test (Extra): ❌ Run #32541327985
Latest PR Test (AMD ROCm 7.2): ❌ Run #32541328165

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[b8zhong]

Do we already have some fused CUDA kernel for this? We might try to reuse that one if possible

[akhilg-nv]

I didn't see an existing fused CUDA kernel for this in TRT-LLM. This kernel is very simple with elementwise operations, I haven't done the analysis but it's likely memory bound and most of the savings come from fusing the operations to avoid additional MIO. I'm also not sure how much it impact the e2e latency. So maybe, the triton kernel is "good enough."

I can run some basic e2e benchmark + profiling once the main PR is merged and this is ready for closer review, to validate the above conjecture.

[nvpohanh]

I didn't see an existing fused CUDA kernel for this in TRT-LLM

What about in SGLang code base?

[akhilg-nv]

the counterpart for FlashMLA does a fusion for the fp8 nope + bf16 rope KV pool, but nothing already exists for this exact operation (uniform fp8 pool) as far as I'm aware

[akhilg-nv]

Existing path for FlashMLA: see python wrapper and cuda impl, which is called as part of CompressorBackendMixin._forward_compress_all_in_one here. I can look into adapting it for the uniform FP8 pool format.

[akhilg-nv]

Updated the existing norm + rope + store kernel for the uniform fp8 format for trtllm attention backend by adding a third epilogue for this format, and removed the triton kernel.

Performance is flat across sizes 32 to 4096, so most likely bounded by CPU overhead (measured eagerly).

Measurements on B200:
latest (cuda kernel with uniform fp8 store path): 7.7 µs
triton kernel: 16.7 µs
unfused: 34 µs

[akhilg-nv]

Closing, is folded into #36652 which includes additional performance optimizations.