perf(kimi-k2): fuse decode-path small kernels to reduce launch overhead

[xiaguan]

Summary

• Fusion 1: split_qkv_a + 2× rms_norm → single kimi_mla_split_qkv_a_norm kernel (3→1 launch, ×61 layers, 9.5→3.5μs on H20)
• Fusion 2: add_into + scaled_add → single kimi_residual_add_scaled_f32 kernel (2→1 launch, ×60 MoE layers, 6.3→2.8μs on H20)
• Both fused kernels are bitwise equal to the original path
• Total: 182 fewer kernel launches per decode step, ~0.58ms TPOT reduction at bs=1

E2E Results (H20, TP1 DP8 EP8 PPLX, bs=64, output_len=128)

Metric	Baseline	Fused	Delta
TPOT p50	40.10ms	39.64ms	−0.46ms (−1.1%)
TPOT p99	43.72ms	43.36ms	−0.36ms

Test plan

• Standalone C++ bench: bitwise correctness at B∈{1,4,8,16,32,64}
• ncu profile: confirmed launch reduction and kernel duration improvement
• E2E bench_serving dp8 tp1 ep8 bs64 on H20 cluster

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces optimized CUDA kernels to fuse operations and improve performance. Specifically, it adds kimi_residual_add_scaled_f32_kernel to combine residual addition and scaling, and split_qkv_a_norm_kernel to fuse the splitting and RMS normalization of QKV projections. These are integrated into the Rust runner and worker pipelines. The review feedback highlights a return type mismatch between the CUDA implementation (cudaError_t) and Rust FFI bindings (CUresult) for the new MLA split-norm function, as well as an optimization opportunity in the CUDA kernel to restrict shared memory writes to a single thread within the warp.

[gemini-code-assist]

The function kimi_mla_split_qkv_a_norm_cuda is defined as returning cudaError_t in C++, but is declared as returning CUresult in the Rust FFI bindings (ffi.rs). While both are integer types and success (0) maps correctly, their non-success error codes do not align 1:1. This mismatch can lead to incorrect error propagation or silent failures. Consider returning CUresult and mapping the error code using cudaGetLastError() similar to how it is done in kimi_experts.cu (e.g., err == cudaSuccess ? CUDA_SUCCESS : CUDA_ERROR_LAUNCH_FAILED).

[gemini-code-assist]

In split_qkv_a_norm_kernel, all 32 threads of warp 0 (ty == 0) write the same reduced total_q and total_ckv values to smem_q[0] and smem_ckv[0]. To avoid redundant writes and potential bank conflicts/coherency overhead, it is recommended to restrict these writes to thread 0 of the warp (tx == 0).

    if (tx == 0) {
      smem_q[0] = total_q;
    }

    float total_ckv = (tx < CKV_WARPS) ? smem_ckv[tx] : 0.f;
#pragma unroll
    for (int offset = kWarpSize / 2; offset > 0; offset /= 2)
      total_ckv += fi_shfl_xor(total_ckv, offset);
    if (tx == 0) {
      smem_ckv[0] = total_ckv;
    }