[ruocheng-nv]

Description

Split out of #21178 at review request so attention owners can review the FlashMLA build change on its own; #21178
keeps the CUTLASS MoE kernel changes.

flash_mla_cpp_tllm was built with CUDA_SEPARABLE_COMPILATION ON (nvcc -rdc=true). Nothing in FlashMLA needs
cross-TU device linking (its own setup.py builds without -rdc), and -rdc makes ptxas spill in the attention
kernels:

• sm90 sparse prefill (d_qk 512): 0 -> 2044 B spill stores, 296 B stack
• sm90 sparse fp8 decode (h64): 0 -> 220 B
• sm100 sparse prefill (h128): 0 -> 44 B
• sm100 sparse decode (h64): 0 -> 12 B

H200, flash_mla_sparse_fwd with 13312 queries x 64 heads x 2176 top-k rows (d_qk 512): 17.07 -> 7.85 ms per call,
outputs bit-identical. Same FlashMLA commit and nvcc; only -rdc differs. The change applies to every FlashMLA kernel
built in this target (SM90 and SM100).

Test Coverage

• Build stages (x86_64 and SBSA) for the FlashMLA sm90/sm100f objects.
• Existing FlashMLA sparse attention tests under unittest/_torch/attention on H100/B200 (outputs bit-identical).

PR Checklist

• PR description clearly explains what and why.
• No API changes.

Dev Engineer Review

flash_mla_cpp_tllm now disables CUDA separable compilation. The change removes the device-linking configuration and follows FlashMLA’s whole-program compilation model. The stated intent is to avoid ptxas register spills in attention kernels. No API changes are shown. Verify that the target builds and runs correctly on supported CUDA architectures.

QA Engineer Review

No test changes.

Per-File QA Perspective

• cpp/tensorrt_llm/flash_mla/CMakeLists.txt: Verify that the extension builds and its attention kernels run correctly with relocatable device code disabled. No test files or test lists changed.

[coderabbitai]

No actionable comments were generated in the recent review.

Walkthrough

The FlashMLA CMake target now disables CUDA separable compilation and no longer enables device-symbol resolution. The adjacent comment about generating fatbinData was removed; the LTO comment remains.

Changes

FlashMLA build configuration

Layer / File(s)	Summary
Update CUDA compilation settings cpp/tensorrt_llm/flash_mla/CMakeLists.txt	The flash_mla_cpp_tllm target disables CUDA separable compilation and no longer enables device-symbol resolution. The comment about generating fatbinData was removed.

Estimated code review effort: 1 (Trivial) | ~3 minutes

Change: Bug fix

Suggested labels: api-compatible

Merge Risk: Minimal · up to 7793b

The extension now builds without separable CUDA compilation. No merge-blocking issue is established.

Architecture risk: Low · up to 7793b

The change affects 1 system.

Changed systems: cpp

Architecture concerns
No architecture-level concerns identified.

Review details

Systems and components

• observed — cpp (service) was modified; 1 changed file maps to changed impact.

Before / after behavior

• observed — Modified behavior in cpp/tensorrt_llm/flash_mla/CMakeLists.txt: The comment about generating fatbinData was removed; the LTO comment remains.
• observed — Modified behavior in cpp/tensorrt_llm/flash_mla/CMakeLists.txt: flash_mla_cpp_tllm now disables CUDA separable compilation and no longer enables device-symbol resolution; the previous setting enabled both.

Pre-merge checks: 5

Check name	Status	Explanation
Title check		The title follows the repository format and clearly states that FlashMLA will build without relocatable device code.
Description check		The description explains the reason for the change, reports performance and correctness results, lists test coverage, and notes that there are no API changes. The checklist is brief but the descriptio…
Docstring Coverage		No eligible functions found in the changed files. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Description check

Explanation

The description explains the reason for the change, reports performance and correctness results, lists test coverage, and notes that there are no API changes. The checklist is brief but the description is mostly complete.

Full details: Docstring Coverage

Explanation

No eligible functions found in the changed files. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 0 functions across 0 files. (1 skipped: 1 unsupported.)

Finishing Touches

Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/TensorRT-LLM
• Branch: main
• File: .coderabbit.yaml
• Commit: a1c22fc
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: d9b07895-c151-4ad6-be7b-2c91d0fb3388

Commits

Reviewing files that changed from the base of the PR and between a1c22fc and 7793b00.

Files selected for processing: 1

• cpp/tensorrt_llm/flash_mla/CMakeLists.txt

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.