[lucifer1004]

Description

Fixes #20634.

Depends on #20512: the new path is taken only when the CSA2 projections run in BF16 on SM107, which is what #20512 enables. Without #20512 the DeepSeek-V4.1 model always requests MXFP8 projections, which SM107 rejects and SM100/SM103 keep, so nothing changes for the model.

At decode sizes the two hidden-state projections of DeepseekV41Attention (wq_a: 5120 -> 1280, wkv: 5120 -> 512) are N-small, K = hidden_size GEMMs. In BF16 on Vera Rubin (SM107) cuBLAS reaches 0.9-2 TB/s on them; trtllm::tinygemm2 streams the same weights 1.6-1.8x faster. This PR uses it for those two projections in the BF16 projection mode on SM107:

• post_load_weights copies the wq_a and wkv weights into one row-concatenated [1792, 5120] tensor and rebinds both Linear weights as views into it (no extra memory). Checkpoint reloads copy into the views, so the aliasing survives them.
• One token: both projections run as a single tinygemm2 over the concatenated weight. The two halves are returned with dense [1, N] strides.
• 2-8 tokens: one tinygemm2 per projection.
• Anything else uses the existing Linear path. This includes more tokens, other dtypes, unaligned or non-contiguous inputs, a weight that is no longer a view of the concatenated tensor (for example a replaced Parameter), MXFP8 projections, and other GPUs.

tinygemm2 calls cudaTriggerProgrammaticLaunchCompletion() before its output stores and is always launched with programmatic stream serialization (the concern in #15729 / #15790, both closed). The consumers here, the q_norm / kv_norm RMSNorms (FlashInfer, PDL enabled by default), wait with griddepcontrol.wait, which covers completion and memory visibility of the preceding grid. We also stress-tested this exact pattern (see Test Coverage). Moving the trigger after the stores in the kernel would still be the more robust fix and is left to a separate change.

Measured on Vera Rubin (TP4/EP4, DeepSeek-V4.1-Flash, batch 1, 32K context, CUDA graphs, BF16 CSA2 projections):

• op level, CUDA graph replay, weights rotated past L2, M = 1: fused call 4.13 µs vs. two tinygemm2 calls 7.23 µs vs. two cuBLAS calls 12.52 µs;
• end to end, same server configuration with only this change: decode step period 5.97 ms -> 5.70 ms (-0.27 ms, -4.5%); GSM8K (100 questions) 99 -> 99, passkey 8K/32K/128K 5/5 in both.
• together with [#20635][perf] DeepSeek-V4 MoE gate: Triton kernel for decode-size batches #20637 (Triton DeepSeek-V4 MoE gate; the two PRs are independent and can land in either order), on another node against a baseline server on that node: 5.95 ms -> 5.50 ms (-0.45 ms, -7.5%); GSM8K 99 -> 98, passkey 5/5 in both.

Test Coverage

• tests/unittest/_torch/attention/sparse/csa2/test_module.py::test_tinygemm_hidden_state_projections (new; SM107 only). It builds the native test module (BF16 projections), calls post_load_weights, and checks:
  • the aliasing of both weights;
  • which path runs: a counting wrapper around torch.ops.trtllm.tinygemm2 sees one fused call for 1 token, two calls for 2 and 8 tokens, and none for 9 and 32;
  • shapes and dense strides;
  • closeness to an FP32 F.linear reference;
  • CUDA graph capture / replay of the fused call, bitwise equal to eager over three new inputs;
  • a checkpoint reload, which writes through the views and keeps the fused call;
  • a replaced wkv parameter, which falls back to separate calls.
• _native_model (test helper) now skips non-SM100/103 GPUs only for MXFP8 projections, so the BF16 module can be built on SM107. test_module_sends_whole_context_and_generation_phases builds its module without __init__ and now sets the two new attributes.
• Stress test of the PDL ordering (not part of this PR), on one SM107 GPU: in a CUDA graph, a PDL RMSNorm writes x, then tinygemm2 runs (fused 1792 rows at M = 1, or 1280 + 512 at M = 1/4/8), then two PDL RMSNorms. Each replay gets a new input and is compared bitwise with the same ops run eagerly, without PDL and with a device sync between ops. Result: 0 mismatches in 20,000 replays for each of the 4 cases. The test has no positive control.
• Results on Vera Rubin (compute capability 10.7, torch 2.14.0+cu132):
  • with the change: the new test passes; test_module.py 73 passed, 17 skipped, 1 failed;
  • without the change: the new test fails (no post_load_weights); test_module.py 72 passed, 17 skipped, 2 failed.
  • The failure common to both runs is test_module_composed_flash_helper[flash_mla]: the FlashMLA module is not installed in our image.
• pre-commit on the changed files: clean.

PR Checklist

• PR description clearly explains what and why.
• PR follows the TRT-LLM coding guidelines to the best of my knowledge.
• Test cases are provided for new code paths.
• No API changes, no new dependencies, no CODEOWNERS changes.
• Please check this after reviewing the above items as appropriate for this PR.