[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>
This commit is contained in:
Yongye Zhu
co-authored by
Claude Fable 5
Claude Fable 5
parent
c7e6e36fa9
commit
9509fc8ae6
@@ -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 <typename scalar_t, bool FP8, bool APPLY_ROPE = false>
|
||||
__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<uint8_t*>(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<scalar_t, FP8>(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<scalar_t, false, APPLY_ROPE>(
|
||||
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<scalar_t, true, APPLY_ROPE>(
|
||||
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 <typename scalar_t, bool Q_FP8, bool KV_FP8, bool APPLY_ROPE>
|
||||
template <typename scalar_t, bool Q_FP8, bool KV_FP8, bool APPLY_ROPE,
|
||||
int SPLIT = 1>
|
||||
__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<scalar_t, Q_FP8, APPLY_ROPE>(
|
||||
uint8_t* const dst =
|
||||
reinterpret_cast<uint8_t*>(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<scalar_t, Q_FP8, APPLY_ROPE>(
|
||||
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<scalar_t const*>(ql_nope.const_data_ptr()),
|
||||
ql_nope.stride(0), ql_nope.stride(1),
|
||||
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(
|
||||
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<scalar_t, false,
|
||||
false, true>);
|
||||
if (split_rows) {
|
||||
launch(
|
||||
kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, false,
|
||||
false, true, 3>,
|
||||
3);
|
||||
} else {
|
||||
launch(
|
||||
kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, false,
|
||||
false, true>,
|
||||
1);
|
||||
}
|
||||
} else {
|
||||
launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, false,
|
||||
false, false>);
|
||||
if (split_rows) {
|
||||
launch(
|
||||
kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, false,
|
||||
false, false, 3>,
|
||||
3);
|
||||
} else {
|
||||
launch(
|
||||
kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, false,
|
||||
false, false>,
|
||||
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<scalar_t const*>(ql_nope.const_data_ptr()),
|
||||
ql_nope.stride(0), ql_nope.stride(1),
|
||||
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(
|
||||
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<scalar_t, true,
|
||||
true, true>);
|
||||
if (split_rows) {
|
||||
launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, true,
|
||||
true, true, 3>,
|
||||
3);
|
||||
} else {
|
||||
launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, true,
|
||||
true, true>,
|
||||
1);
|
||||
}
|
||||
} else {
|
||||
launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, true,
|
||||
true, false>);
|
||||
if (split_rows) {
|
||||
launch(
|
||||
kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, true,
|
||||
true, false, 3>,
|
||||
3);
|
||||
} else {
|
||||
launch(kk3::fusedKimiK3MLADecodeQConcatKVCacheKernel<scalar_t, true,
|
||||
true, false>,
|
||||
1);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user