[lucifer1004]

Description

Fixes #20545.

In Kimi K3 decode, every MoE layer ends with prefix_sum + (routed_out + shared_out): two elementwise kernels whose result is read first by the next attention-residual + RMSNorm. This PR lets that consumer do both additions inside the kernel it already runs. Opt-in via KIMI_K3_MOE_TAIL_ADD_FOLD=1.

Op (attnResFwd.cu/.h, attnResOp.cpp, cpp_custom_ops.py)

• trtllm::attn_res_add_rmsnorm_fwd gets an optional trailing Tensor? layer_residual_add2=None. With it, the updated prefix sum is bf16(layer_residual + bf16(layer_residual_add + layer_residual_add2)). That is the result of the two BF16 tensor additions, bit for bit, subnormals included: both additions use add.rn.f32, so --use_fast_math does not flush them.
• This is a new FUSE_LAYER_ADD2 template flag on the existing single-CTA (N <= 4) and split-K (N >= 5) kernels, not a copy of them. Without the argument, the op and the existing instantiations are unchanged.
• The persistent kernel does not take the second addend.

Model (modeling_kimi_linear.py)

• KimiK3MoERuntime.forward(..., return_parts=True) returns (routed_out, shared_out) unsummed.
• With the env set, KimiLinearModel.forward passes defer_moe_tail_add=True to MoE layers. Such a layer returns a _MoeTailPending(prefix, routed, shared) when its token count is inside the fused window (KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS, H = 7168, BF16 CUDA tensors).
• The next layer's pre-attention attention-residual, or the model's output one, consumes it through _apply_attn_res_add_and_rmsnorm(..., addend2=shared).
• The addition is materialized first, and the existing path then runs unchanged, wherever the fused per-token kernel is not the path that would be taken anyway:
  • a layer with a speculative-decoding hidden-state capture (the tap needs the pre-norm mixture);
  • no snapshots yet;
  • the persistent kernel's topology;
  • N > 12;
  • the final-layer capture.
• The prefix sum that is written to the snapshot bank or carried on is therefore bitwise the same as without the env.

Performance, Rubin, TP8 across 2 nodes (MNNVL), Kimi K3 with speculative decoding, CUDA graphs, single request, decode step time. Measured with the same kernel arithmetic and the same hand-off on our integration branch, as a separate op. There, the two additions were already a single kernel, so on main, where they are two, the saving should be at least this large.

	2K	32K	64K	96K	128K	393K
fold vs. separate additions	-0.106 ms	-0.117 ms	-0.121 ms	-0.126 ms	-0.107 ms	-0.107 ms

Accuracy and needle checks unchanged. In a single-GPU microbenchmark (Rubin single GPU, 200-SM engineering sample), the fused op costs 0.58-0.76 µs more than the attention-residual alone with the single-CTA kernel and 0.10-0.17 µs more with split-K, against a median of 1.6 µs for the standalone (single-kernel) addition it replaces in production.

The fold applies only where the fused per-token epilogue is used, i.e. up to KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS tokens per rank (default 1; #20372 proposes 32). The numbers above were taken with a window covering the decode batch.

Relation to other work:

• [None][feat] Kimi k3 fuse decode attention residual tail #17194 (draft) added the attention-residual + RMSNorm fusions that are now on main. This PR builds on those ops and does not overlap with it.
• [#20540][perf] Kimi K3 MoE decode: router GEMM and routing kernel on a router stream #20541 also edits KimiK3MoERuntime.forward. The two are independent but will conflict textually; whichever lands second gets rebased.

Test Coverage

• tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py (in l0_b200.yml and l0_gb300_multi_gpus.yml):
  • two addends give bitwise the same updated prefix sum and output as one addend holding their BF16 sum, and the same prefix sum as prefix + (routed + shared) in PyTorch (N = 1 to 12, 1 / 8 / 16 / 32 tokens);
  • subnormal sums are kept;
  • the model helper dispatches the fused add2 path and matches the pre-summed one bitwise;
  • with the persistent topology the helper declines addend2, so the caller adds first.
• tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py: a 25-layer decoder chain with deferred MoE tails is bitwise equal to the eager chain, including the snapshot bank and a capture tap. This test runs on CPU through the fallback, so it covers the hand-off rather than the kernel.
• GPU results at this commit (Rubin single GPU (200-SM engineering sample, SM107); C++ built from main plus the C++ changes of this PR, [#20540][perf] Kimi K3 MoE decode: router GEMM and routing kernel on a router stream #20541 and [#20545][perf] Kimi K3: fold the MoE-tail additions into the next attention-residual #20546 (they touch disjoint files)):
  • test_attn_res_rmsnorm_op.py: 56 passed, 1 skipped (on main the 18 new cases fail), including the bitwise pre-summed comparison and the subnormal (non-FTZ) cases;
  • test_attn_res_op.py: 13 passed;
  • test_custom_ops.py -k test_register_fake failed only on trtllm::mxe2m1_moe_routing_workspace, an op of [#20540][perf] Kimi K3 MoE decode: router GEMM and routing kernel on a router stream #20541 that was compiled into the shared library while this branch's Python does not register [#20540][perf] Kimi K3 MoE decode: router GEMM and routing kernel on a router stream #20541's fakes. This PR's own ops all have fakes; with this PR's C++ alone the test is expected to pass (not run separately).
• Locally (CPU): the test above, tests/unittest/_torch/modeling/test_kimi_linear_modeling.py and tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dflash_scaffold.py pass; a mutation in the hand-off makes the new CPU test fail.
• Earlier validation of the same arithmetic on Rubin: 420 op-level cases bitwise equal (including subnormals), CUDA-graph replay, and the end-to-end A/B above.
• Not measured: SM100 / SM103, attention DP.
• 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.