From 9509fc8ae61a8eb871a984e759a6d228c930df80 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Thu, 3 Sep 2026 15:39:58 -0700 Subject: [PATCH] [Perf][Kimi-K3] Cut MLA decode concat/cache epilogue latency (#54896) Signed-off-by: Yongye Zhu Co-authored-by: Claude Fable 5 --- ..._kimi_k3_mla_key_concat_kv_cache_kernel.cu | 134 +++++++++++++----- .../test_kimi_k3_mla_fused_epilogue.py | 23 +++ 2 files changed, 121 insertions(+), 36 deletions(-) diff --git a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu index 2e2c4f7cf9..05c7bfc6d1 100644 --- a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu +++ b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu @@ -102,6 +102,9 @@ constexpr int kQkHeadDim = kQkNopeHeadDim + kQkRopeHeadDim; // 192 constexpr int kVHeadDim = 128; // V constexpr int kCacheEntry = kKvLoraRank + kQkRopeHeadDim; // 576 constexpr int kVecElems = 8; // 8 bf16 == one uint4 load / one uint2 fp8 store +// Max token count for the decode epilogue's 3-warps-per-row split; larger +// batches have enough rows for one warp per row. +constexpr int kDecodeRowSplitMaxTokens = 64; #if defined(USE_ROCM) && defined(__gfx942__) constexpr float kFp8Max = 224.0f; @@ -299,14 +302,18 @@ __device__ __forceinline__ void writePrefillQuery( // Concat + store a 576-wide latent: dst[e] = [a512 | b64], e in [0, 576). Used // for the decode query mqa_q = [ql_nope | q_pe] and the plain latent cache // entry [kv_c | k_pe]. FP8 packs to E4M3 (dst_elem_size 1); bf16 stores uint4. +// A single warp passes (laneId, 32); a row split across W warps passes +// (part * 32 + laneId, W * 32). __restrict__ lets the compiler batch loads +// ahead of stores past the byte-typed dst. template -__device__ __forceinline__ void writeLatent576(void* dst, const scalar_t* a512, - const scalar_t* b64, int laneId, - int dst_elem_size, - float scale_inv, - const float* cos_sin = nullptr) { +__device__ __forceinline__ void writeLatent576( + void* __restrict__ dst, const scalar_t* __restrict__ a512, + const scalar_t* __restrict__ b64, int lane, int lane_stride, + int dst_elem_size, float scale_inv, + const float* __restrict__ cos_sin = nullptr) { auto* d = reinterpret_cast(dst); - for (int e = laneId * kVecElems; e < kCacheEntry; e += 32 * kVecElems) { + for (int e = lane * kVecElems; e < kCacheEntry; + e += lane_stride * kVecElems) { if (e < kKvLoraRank) { copyChunk8(d + e * dst_elem_size, a512 + e, scale_inv); } else { @@ -428,7 +435,7 @@ __global__ void fusedKimiK3MLAKeyConcatKVCacheInsertKernel( (slot_id % cache_block_size) * cache_token_stride; writeLatent576( row, kv_c + tokenIdx * kv_c_tok_stride, - k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f, + k_pe + tokenIdx * k_pe_tok_stride, laneId, 32, sizeof(scalar_t), 1.0f, rope_cache); } } @@ -512,7 +519,7 @@ __global__ void fusedKimiK3MLAQKVQuantKVCacheFp8Kernel( (slot_id % cache_block_size) * cache_token_stride; writeLatent576( row, kv_c + tokenIdx * kv_c_tok_stride, - k_pe + tokenIdx * k_pe_tok_stride, laneId, 1, ksi, rope_cache); + k_pe + tokenIdx * k_pe_tok_stride, laneId, 32, 1, ksi, rope_cache); } } @@ -660,8 +667,20 @@ __global__ void fusedKimiK3MLAKeyConcatDsMlaInsertKernel( // Decode epilogue: concat mqa_q = [ql_nope | q_pe] (576) + latent cache insert, // run right before forward_mqa. Q_FP8 quantizes mqa_q; KV_FP8 quantizes the // plain per-tensor cache. (ds_mla cache uses the separate kernel below.) +// +// SPLIT warps cooperate on each 576-wide row: decode batches have too few +// rows to hide the load->store latency with one warp per row, so batches up +// to kDecodeRowSplitMaxTokens launch SPLIT=3 (one 8-element chunk per lane). +// +// Only ql_nope / q_pe come from the producing GEMM chain, so everything else +// (slot_mapping, scales, rope table, cache-row inputs) is read before the +// grid-dependency wait, and the cache-slot warps skip the wait entirely. +// The dependent-launch trigger stays at the end of the kernel: it releases +// the consuming FMHA's grid-dependency wait, so it must not fire before this +// kernel's stores to mqa_q are issued. // ──────────────────────────────────────────────────────────────────────────── -template +template __global__ void fusedKimiK3MLADecodeQConcatKVCacheKernel( const scalar_t* __restrict__ ql_nope, int64_t const qn_tok_stride, int64_t const qn_head_stride, const scalar_t* __restrict__ q_pe, @@ -683,13 +702,13 @@ __global__ void fusedKimiK3MLADecodeQConcatKVCacheKernel( int const laneId = threadIdx.x % 32; int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32; int const slotsPerToken = num_heads + 1; - int const tokenIdx = globalWarpIdx / slotsPerToken; - int const slotIdx = globalWarpIdx % slotsPerToken; + int const globalSlotIdx = globalWarpIdx / SPLIT; + int const part = globalWarpIdx % SPLIT; + int const tokenIdx = globalSlotIdx / slotsPerToken; + int const slotIdx = globalSlotIdx % slotsPerToken; if (tokenIdx >= num_tokens) return; - -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - cudaGridDependencySynchronize(); -#endif + int const lane = part * 32 + laneId; + constexpr int kLaneStride = SPLIT * 32; const float* rope_cache = nullptr; if constexpr (APPLY_ROPE) { @@ -698,12 +717,18 @@ __global__ void fusedKimiK3MLADecodeQConcatKVCacheKernel( if (slotIdx < num_heads) { float const qsi = Q_FP8 ? __ldg(q_scale_inv) : 1.0f; - writeLatent576( + uint8_t* const dst = reinterpret_cast(mqa_q) + - (tokenIdx * mq_tok_stride + slotIdx * mq_head_stride) * kMqElem, - ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride, - q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride, laneId, - kMqElem, qsi, rope_cache); + (tokenIdx * mq_tok_stride + slotIdx * mq_head_stride) * kMqElem; + const scalar_t* const nope_src = + ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride; + const scalar_t* const pe_src = + q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + writeLatent576( + dst, nope_src, pe_src, lane, kLaneStride, kMqElem, qsi, rope_cache); } else { int64_t const slot_id = slot_mapping[tokenIdx]; if (slot_id >= 0) { @@ -714,7 +739,7 @@ __global__ void fusedKimiK3MLADecodeQConcatKVCacheKernel( slot_id % cache_block_size * cache_token_stride) * kCacheElem, kv_c + tokenIdx * kv_c_tok_stride, k_pe + tokenIdx * k_pe_tok_stride, - laneId, kCacheElem, ksi, rope_cache); + lane, kLaneStride, kCacheElem, ksi, rope_cache); } } @@ -760,7 +785,7 @@ __global__ void fusedKimiK3MLADecodeQConcatDsMlaKernel( mqa_q + tokenIdx * mq_tok_stride + slotIdx * mq_head_stride, ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride, q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride, laneId, - sizeof(scalar_t), 1.0f, rope_cache); + 32, sizeof(scalar_t), 1.0f, rope_cache); } else { int64_t const slot_id = slot_mapping[tokenIdx]; if (slot_id >= 0) { @@ -1367,9 +1392,9 @@ void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( VLLM_STABLE_DISPATCH_HALF_TYPES( dt, "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", [&] { - auto launch = [&](auto kernel) { - kk3::launchPdl( - kernel, num_tokens, num_heads, stream, + auto launch = [&](auto kernel, int row_split) { + kk3::launchPdlSlots( + kernel, num_tokens, (num_heads + 1) * row_split, 1, 0, stream, reinterpret_cast(ql_nope.const_data_ptr()), ql_nope.stride(0), ql_nope.stride(1), reinterpret_cast(q_pe.const_data_ptr()), @@ -1388,12 +1413,32 @@ void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert( cos_sin_cache.value().const_data_ptr()) : nullptr); }; + // Split decode-sized rows across 3 warps (see the kernel comment). + bool const split_rows = num_tokens <= kk3::kDecodeRowSplitMaxTokens; if (apply_rope) { - launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + if (split_rows) { + launch( + kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel, + 3); + } else { + launch( + kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel, + 1); + } } else { - launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + if (split_rows) { + launch( + kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel, + 3); + } else { + launch( + kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel, + 1); + } } }); } @@ -1442,9 +1487,9 @@ void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( VLLM_STABLE_DISPATCH_HALF_TYPES( dt, "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", [&] { - auto launch = [&](auto kernel) { - kk3::launchPdl( - kernel, num_tokens, num_heads, stream, + auto launch = [&](auto kernel, int row_split) { + kk3::launchPdlSlots( + kernel, num_tokens, (num_heads + 1) * row_split, 1, 0, stream, reinterpret_cast(ql_nope.const_data_ptr()), ql_nope.stride(0), ql_nope.stride(1), reinterpret_cast(q_pe.const_data_ptr()), @@ -1464,12 +1509,29 @@ void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert( cos_sin_cache.value().const_data_ptr()) : nullptr); }; + // Split decode-sized rows across 3 warps (see the kernel comment). + bool const split_rows = num_tokens <= kk3::kDecodeRowSplitMaxTokens; if (apply_rope) { - launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + if (split_rows) { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel, + 3); + } else { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel, + 1); + } } else { - launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel); + if (split_rows) { + launch( + kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel, + 3); + } else { + launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel, + 1); + } } }); } diff --git a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py index d7cb7ca7eb..aa7c836666 100644 --- a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py +++ b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py @@ -285,3 +285,26 @@ def test_decode_epilogue_preserves_nope_path() -> None: torch.testing.assert_close( _cache_rows(cache, slots), torch.cat((kv_c, k_pe), dim=-1) ) + + +@torch.inference_mode() +@pytest.mark.parametrize("num_tokens", [64, 65]) +def test_decode_epilogue_row_split_boundary(num_tokens: int) -> None: + """Both sides of the M-based warps-per-row dispatch produce the same rows.""" + torch.manual_seed(3) + num_blocks = (num_tokens + _BLOCK_SIZE - 1) // _BLOCK_SIZE + slots = torch.randperm(num_blocks * _BLOCK_SIZE, device="cuda")[:num_tokens].long() + ql_nope = _randn(num_tokens, _NUM_HEADS, 512) + q_pe = _randn(num_tokens, _NUM_HEADS, 64) + kv_c = _randn(num_tokens, 512) + k_pe = _randn(num_tokens, 64) + cache = torch.zeros(num_blocks, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE) + + q_actual = fused_mla_decode_q_concat_kv_cache_insert( + ql_nope, q_pe, kv_c, k_pe, cache, slots + ) + + torch.testing.assert_close(q_actual, torch.cat((ql_nope, q_pe), dim=-1)) + torch.testing.assert_close( + _cache_rows(cache, slots), torch.cat((kv_c, k_pe), dim=-1) + )