[lucifer1004]

Description

Fixes #20379.

fmhaReductionKernel merges the multi-CTA-KV partials of each output row in one thread, serially over all numCtasKv partials. At low batch its grid is only a few CTAs, so the kernel costs a fixed several microseconds per attention call.

This PR adds a KvSplit template parameter:

• The 512 threads form KvSplit groups. Group g runs the existing online-softmax merge over partials g, g + KvSplit, ... of the same rows. Group 0 then combines the per-group (max, sum, O) through shared memory, rescaling each by exp2(scaleLog2 * (groupMax - max)) and skipping empty groups. Sinks, softmax-stats export and the store are unchanged and done by group 0.
• In the unrolled loads, a slot past the end re-reads the group's current partial with weight 0, so the running max never moves without rescaling.
• Each CTA covers 1/KvSplit of the rows (NumBytePerSlice is per group), so the host launches KvSplit times more slices. TLLM_FMHA_REDUCTION_KV_SPLIT=<2|4|8|16> is an upper bound, and the host halves it until every group holds at least one whole row and the slices still fit the existing 2 x SMs / grid CTA budget. Unset or 1 selects KvSplit = 1, which runs the same per-thread loop as before.
• Static shared memory for the merge is 20 KB per CTA, and only for KvSplit > 1.

Measured on SM107 with Kimi K3 MLA decode (grouped Q64 kernels, batch 1): per-call time -1.9 / -7.3 / -7.4 / -6.5 us at 2K / 8K / 32K / 128K with KvSplit 8, and an end-to-end decode step of -0.10 / -0.23 / -0.21 ms at 2K / 32K / 128K (details in #20379). It is opt-in for now because SM100 / SM103 are not measured.

Test Coverage

• No existing unit test calls runFmhaReduction directly. With the variable unset the launch is KvSplit = 1, the same per-thread loop as before.
• Validated on SM107 (equivalent change): output sums and absolute sums for KvSplit 2/4/8/16 match KvSplit 1 to ~1e-5 at batch 1/2/4 in a TllmGenFmhaRunner microbenchmark; end-to-end accuracy unchanged with KvSplit 8.
• This branch: nvcc -arch=sm_100a compile of fmhaReduction.cu and pre-commit clean. Not yet run on a GPU; SM100 / SM103 not run. Happy to add a dedicated reduction test with synthetic partials if reviewers want one.

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.