[Perf][Kimi-K3] Cut MLA decode concat/cache epilogue latency (#54896)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Yongye Zhu committed
Sep 3, 2026 at 15:39 UTC
9509fc8ae61a8eb871a984e759a6d228c930df80
2 files changed
+121
-36
csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu
+98
-36
@@ -102,6 +102,9 @@ constexpr int kQkHeadDim = kQkNopeHeadDim + kQkRopeHeadDim; // 192
102
constexpr int kVHeadDim = 128; // V
103
constexpr int kCacheEntry = kKvLoraRank + kQkRopeHeadDim; // 576
104
constexpr int kVecElems = 8; // 8 bf16 == one uint4 load / one uint2 fp8 store
105
+// Max token count for the decode epilogue's 3-warps-per-row split; larger
106
+// batches have enough rows for one warp per row.
107
+constexpr int kDecodeRowSplitMaxTokens = 64;
108
109
#if defined(USE_ROCM) && defined(__gfx942__)
110
constexpr float kFp8Max = 224.0f;
@@ -299,14 +302,18 @@ __device__ __forceinline__ void writePrefillQuery(
302
// Concat + store a 576-wide latent: dst[e] = [a512 | b64], e in [0, 576). Used
303
// for the decode query mqa_q = [ql_nope | q_pe] and the plain latent cache
304
// entry [kv_c | k_pe]. FP8 packs to E4M3 (dst_elem_size 1); bf16 stores uint4.
305
+// A single warp passes (laneId, 32); a row split across W warps passes
306
+// (part * 32 + laneId, W * 32). __restrict__ lets the compiler batch loads
307
+// ahead of stores past the byte-typed dst.
308
template <typename scalar_t, bool FP8, bool APPLY_ROPE = false>
303
-__device__ __forceinline__ void writeLatent576(void* dst, const scalar_t* a512,
304
- const scalar_t* b64, int laneId,
305
- int dst_elem_size,
306
- float scale_inv,
307
- const float* cos_sin = nullptr) {
309
+__device__ __forceinline__ void writeLatent576(
310
+ void* __restrict__ dst, const scalar_t* __restrict__ a512,
311
+ const scalar_t* __restrict__ b64, int lane, int lane_stride,
312
+ int dst_elem_size, float scale_inv,
313
+ const float* __restrict__ cos_sin = nullptr) {
314
auto* d = reinterpret_cast<uint8_t*>(dst);
309
- for (int e = laneId * kVecElems; e < kCacheEntry; e += 32 * kVecElems) {
315
+ for (int e = lane * kVecElems; e < kCacheEntry;
316
+ e += lane_stride * kVecElems) {
317
if (e < kKvLoraRank) {
318
copyChunk8<scalar_t, FP8>(d + e * dst_elem_size, a512 + e, scale_inv);
319
} else {
@@ -428,7 +435,7 @@ __global__ void fusedKimiK3MLAKeyConcatKVCacheInsertKernel(
435
(slot_id % cache_block_size) * cache_token_stride;
436
writeLatent576<scalar_t, false, APPLY_ROPE>(
437
row, kv_c + tokenIdx * kv_c_tok_stride,
431
- k_pe + tokenIdx * k_pe_tok_stride, laneId, sizeof(scalar_t), 1.0f,
438
+ k_pe + tokenIdx * k_pe_tok_stride, laneId, 32, sizeof(scalar_t), 1.0f,
439
rope_cache);
440
}
441
}
@@ -512,7 +519,7 @@ __global__ void fusedKimiK3MLAQKVQuantKVCacheFp8Kernel(
519
(slot_id % cache_block_size) * cache_token_stride;
520
writeLatent576<scalar_t, true, APPLY_ROPE>(
521
row, kv_c + tokenIdx * kv_c_tok_stride,
515
- k_pe + tokenIdx * k_pe_tok_stride, laneId, 1, ksi, rope_cache);
522
+ k_pe + tokenIdx * k_pe_tok_stride, laneId, 32, 1, ksi, rope_cache);
523
}
524
}
525
@@ -660,8 +667,20 @@ __global__ void fusedKimiK3MLAKeyConcatDsMlaInsertKernel(
667
// Decode epilogue: concat mqa_q = [ql_nope | q_pe] (576) + latent cache insert,
668
// run right before forward_mqa. Q_FP8 quantizes mqa_q; KV_FP8 quantizes the
669
// plain per-tensor cache. (ds_mla cache uses the separate kernel below.)
670
+//
671
+// SPLIT warps cooperate on each 576-wide row: decode batches have too few
672
+// rows to hide the load->store latency with one warp per row, so batches up
673
+// to kDecodeRowSplitMaxTokens launch SPLIT=3 (one 8-element chunk per lane).
674
+//
675
+// Only ql_nope / q_pe come from the producing GEMM chain, so everything else
676
+// (slot_mapping, scales, rope table, cache-row inputs) is read before the
677
+// grid-dependency wait, and the cache-slot warps skip the wait entirely.
678
+// The dependent-launch trigger stays at the end of the kernel: it releases
679
+// the consuming FMHA's grid-dependency wait, so it must not fire before this
680
+// kernel's stores to mqa_q are issued.
681
// ────────────────────────────────────────────────────────────────────────────
664
-template <typename scalar_t, bool Q_FP8, bool KV_FP8, bool APPLY_ROPE>
682
+template <typename scalar_t, bool Q_FP8, bool KV_FP8, bool APPLY_ROPE,
683
+ int SPLIT = 1>
684
__global__ void fusedKimiK3MLADecodeQConcatKVCacheKernel(
685
const scalar_t* __restrict__ ql_nope, int64_t const qn_tok_stride,
686
int64_t const qn_head_stride, const scalar_t* __restrict__ q_pe,
@@ -683,13 +702,13 @@ __global__ void fusedKimiK3MLADecodeQConcatKVCacheKernel(
702
int const laneId = threadIdx.x % 32;
703
int const globalWarpIdx = blockIdx.x * warpsPerBlock + threadIdx.x / 32;
704
int const slotsPerToken = num_heads + 1;
686
- int const tokenIdx = globalWarpIdx / slotsPerToken;
687
- int const slotIdx = globalWarpIdx % slotsPerToken;
705
+ int const globalSlotIdx = globalWarpIdx / SPLIT;
706
+ int const part = globalWarpIdx % SPLIT;
707
+ int const tokenIdx = globalSlotIdx / slotsPerToken;
708
+ int const slotIdx = globalSlotIdx % slotsPerToken;
709
if (tokenIdx >= num_tokens) return;
689
-
690
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
691
- cudaGridDependencySynchronize();
692
-#endif
710
+ int const lane = part * 32 + laneId;
711
+ constexpr int kLaneStride = SPLIT * 32;
712
713
const float* rope_cache = nullptr;
714
if constexpr (APPLY_ROPE) {
@@ -698,12 +717,18 @@ __global__ void fusedKimiK3MLADecodeQConcatKVCacheKernel(
717
718
if (slotIdx < num_heads) {
719
float const qsi = Q_FP8 ? __ldg(q_scale_inv) : 1.0f;
701
- writeLatent576<scalar_t, Q_FP8, APPLY_ROPE>(
720
+ uint8_t* const dst =
721
reinterpret_cast<uint8_t*>(mqa_q) +
703
- (tokenIdx * mq_tok_stride + slotIdx * mq_head_stride) * kMqElem,
704
- ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride,
705
- q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride, laneId,
706
- kMqElem, qsi, rope_cache);
722
+ (tokenIdx * mq_tok_stride + slotIdx * mq_head_stride) * kMqElem;
723
+ const scalar_t* const nope_src =
724
+ ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride;
725
+ const scalar_t* const pe_src =
726
+ q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride;
727
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
728
+ cudaGridDependencySynchronize();
729
+#endif
730
+ writeLatent576<scalar_t, Q_FP8, APPLY_ROPE>(
731
+ dst, nope_src, pe_src, lane, kLaneStride, kMqElem, qsi, rope_cache);
732
} else {
733
int64_t const slot_id = slot_mapping[tokenIdx];
734
if (slot_id >= 0) {
@@ -714,7 +739,7 @@ __global__ void fusedKimiK3MLADecodeQConcatKVCacheKernel(
739
slot_id % cache_block_size * cache_token_stride) *
740
kCacheElem,
741
kv_c + tokenIdx * kv_c_tok_stride, k_pe + tokenIdx * k_pe_tok_stride,
717
- laneId, kCacheElem, ksi, rope_cache);
742
+ lane, kLaneStride, kCacheElem, ksi, rope_cache);
743
}
744
}
745
@@ -760,7 +785,7 @@ __global__ void fusedKimiK3MLADecodeQConcatDsMlaKernel(
785
mqa_q + tokenIdx * mq_tok_stride + slotIdx * mq_head_stride,
786
ql_nope + tokenIdx * qn_tok_stride + slotIdx * qn_head_stride,
787
q_pe + tokenIdx * qpe_tok_stride + slotIdx * qpe_head_stride, laneId,
763
- sizeof(scalar_t), 1.0f, rope_cache);
788
+ 32, sizeof(scalar_t), 1.0f, rope_cache);
789
} else {
790
int64_t const slot_id = slot_mapping[tokenIdx];
791
if (slot_id >= 0) {
@@ -1367,9 +1392,9 @@ void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(
1392
1393
VLLM_STABLE_DISPATCH_HALF_TYPES(
1394
dt, "fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", [&] {
1370
- auto launch = [&](auto kernel) {
1371
- kk3::launchPdl(
1372
- kernel, num_tokens, num_heads, stream,
1395
+ auto launch = [&](auto kernel, int row_split) {
1396
+ kk3::launchPdlSlots(
1397
+ kernel, num_tokens, (num_heads + 1) * row_split, 1, 0, stream,
1398
reinterpret_cast<scalar_t const*>(ql_nope.const_data_ptr()),
1399
ql_nope.stride(0), ql_nope.stride(1),
1400
reinterpret_cast<scalar_t const*>(q_pe.const_data_ptr()),
@@ -1388,12 +1413,32 @@ void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(
1413
cos_sin_cache.value().const_data_ptr())
1414
: nullptr);
1415
};
1416
+ // Split decode-sized rows across 3 warps (see the kernel comment).
1417
+ bool const split_rows = num_tokens <= kk3::kDecodeRowSplitMaxTokens;
1418
if (apply_rope) {
1392
- launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, false,
1393
- false, true>);
1419
+ if (split_rows) {
1420
+ launch(
1421
+ kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, false,
1422
+ false, true, 3>,
1423
+ 3);
1424
+ } else {
1425
+ launch(
1426
+ kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, false,
1427
+ false, true>,
1428
+ 1);
1429
+ }
1430
} else {
1395
- launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, false,
1396
- false, false>);
1431
+ if (split_rows) {
1432
+ launch(
1433
+ kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, false,
1434
+ false, false, 3>,
1435
+ 3);
1436
+ } else {
1437
+ launch(
1438
+ kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, false,
1439
+ false, false>,
1440
+ 1);
1441
+ }
1442
}
1443
});
1444
}
@@ -1442,9 +1487,9 @@ void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(
1487
1488
VLLM_STABLE_DISPATCH_HALF_TYPES(
1489
dt, "fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert", [&] {
1445
- auto launch = [&](auto kernel) {
1446
- kk3::launchPdl(
1447
- kernel, num_tokens, num_heads, stream,
1490
+ auto launch = [&](auto kernel, int row_split) {
1491
+ kk3::launchPdlSlots(
1492
+ kernel, num_tokens, (num_heads + 1) * row_split, 1, 0, stream,
1493
reinterpret_cast<scalar_t const*>(ql_nope.const_data_ptr()),
1494
ql_nope.stride(0), ql_nope.stride(1),
1495
reinterpret_cast<scalar_t const*>(q_pe.const_data_ptr()),
@@ -1464,12 +1509,29 @@ void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(
1509
cos_sin_cache.value().const_data_ptr())
1510
: nullptr);
1511
};
1512
+ // Split decode-sized rows across 3 warps (see the kernel comment).
1513
+ bool const split_rows = num_tokens <= kk3::kDecodeRowSplitMaxTokens;
1514
if (apply_rope) {
1468
- launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, true,
1469
- true, true>);
1515
+ if (split_rows) {
1516
+ launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, true,
1517
+ true, true, 3>,
1518
+ 3);
1519
+ } else {
1520
+ launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, true,
1521
+ true, true>,
1522
+ 1);
1523
+ }
1524
} else {
1471
- launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, true,
1472
- true, false>);
1525
+ if (split_rows) {
1526
+ launch(
1527
+ kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, true,
1528
+ true, false, 3>,
1529
+ 3);
1530
+ } else {
1531
+ launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, true,
1532
+ true, false>,
1533
+ 1);
1534
+ }
1535
}
1536
});
1537
}
tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py
+23
@@ -285,3 +285,26 @@ def test_decode_epilogue_preserves_nope_path() -> None:
285
torch.testing.assert_close(
286
_cache_rows(cache, slots), torch.cat((kv_c, k_pe), dim=-1)
287
)
288
+
289
+
290
+@torch.inference_mode()
291
+@pytest.mark.parametrize("num_tokens", [64, 65])
292
+def test_decode_epilogue_row_split_boundary(num_tokens: int) -> None:
293
+ """Both sides of the M-based warps-per-row dispatch produce the same rows."""
294
+ torch.manual_seed(3)
295
+ num_blocks = (num_tokens + _BLOCK_SIZE - 1) // _BLOCK_SIZE
296
+ slots = torch.randperm(num_blocks * _BLOCK_SIZE, device="cuda")[:num_tokens].long()
297
+ ql_nope = _randn(num_tokens, _NUM_HEADS, 512)
298
+ q_pe = _randn(num_tokens, _NUM_HEADS, 64)
299
+ kv_c = _randn(num_tokens, 512)
300
+ k_pe = _randn(num_tokens, 64)
301
+ cache = torch.zeros(num_blocks, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE)
302
+
303
+ q_actual = fused_mla_decode_q_concat_kv_cache_insert(
304
+ ql_nope, q_pe, kv_c, k_pe, cache, slots
305
+ )
306
+
307
+ torch.testing.assert_close(q_actual, torch.cat((ql_nope, q_pe), dim=-1))
308
+ torch.testing.assert_close(
309
+ _cache_rows(cache, slots), torch.cat((kv_c, k_pe), dim=-1)
310
+ )