@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 903d2efe7e

[Attention][MLA] Fuse Kimi-K3 chunked-context K/V packing (#51772)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Yongye Zhu committed Aug 12, 2026 at 22:42 UTC 903d2efe7eb6d045d03c0a8192b1ac7bec50d168
15 files changed +996 -35
csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu
+275 -6
@@ -25,6 +25,14 @@
25 * matching MLA's _q_scale / _k_scale / _v_scale (the cache latent uses a
26 * separate cache scale). Per-tensor E4M3.
27 *
28 + * K/V pack (fused_kimi_k3_mla_kv_concat{,_quant_fp8}):
29 + * - k_out[t, h] = [k_nope[t, h] | k_pe[t]], cast to fp8 by the _quant_fp8
30 + * variant, which also writes v_fp8[t, h] = cast(v[t, h])
31 + * The chunked-context epilogue: no cache insert, no q, no RoPE (the cached
32 + * k_pe is already rotated) and no scaling. `k_nope`/`v` are strided views
33 + * of one kv_b_proj output and `k_pe` may already be fp8 (a plain fp8 cache is
34 + * gathered without dequantizing); the outputs are contiguous.
35 + *
36 * fp8_ds_mla (fused_kimi_k3_mla_key_concat_ds_mla_insert):
37 * - full key concat (bf16), and cache insert in DeepSeek's 656-byte
38 * block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile dynamic
@@ -183,6 +191,36 @@ __device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src,
191 }
192 }
193
194 +// Cast 8 source elements (one uint4 of bf16/fp16) to E4M3 with no scaling,
195 +// using the native pairwise converters so the result is bit-identical to
196 +// `src.to(torch.float8_e4m3fn)`.
197 +template <typename scalar_t>
198 +__device__ __forceinline__ void copyChunk8UnitFp8(uint8_t* dst,
199 + const scalar_t* src) {
200 +#ifndef USE_ROCM
201 + uint4 const input = *reinterpret_cast<const uint4*>(src);
202 + using Converter = vllm::_typeConvert<scalar_t>;
203 + auto const* input2 =
204 + reinterpret_cast<typename Converter::packed_hip_type const*>(&input);
205 + uint2 output;
206 + auto* output2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&output);
207 + #pragma unroll
208 + for (int i = 0; i < 4; ++i) {
209 + if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
210 + output2[i] = __nv_cvt_bfloat16raw2_to_fp8x2(
211 + static_cast<__nv_bfloat162_raw>(input2[i]), __NV_SATFINITE,
212 + __NV_E4M3);
213 + } else {
214 + output2[i] = __nv_cvt_halfraw2_to_fp8x2(
215 + static_cast<__half2_raw>(input2[i]), __NV_SATFINITE, __NV_E4M3);
216 + }
217 + }
218 + *reinterpret_cast<uint2*>(dst) = output;
219 +#else
220 + copyChunk8<scalar_t, true>(dst, src, 1.0f);
221 +#endif
222 +}
223 +
224 // Concat + store one head's full key: dst[e] = [k_nope | k_pe], e in [0, 192).
225 // FP8 dst is byte-addressed; bf16 dst is scalar_t-addressed (dst_elem_size).
226 template <typename scalar_t, bool FP8, bool APPLY_ROPE = false>
@@ -202,6 +240,38 @@ __device__ __forceinline__ void writeFullKey(void* dst, const scalar_t* k_nope,
240 }
241 }
242
243 +// Unscaled full-key writer for the chunked-context pack, driven by a half warp
244 +// (16 lanes x 8 elems covers NoPE 128; 8 of them cover RoPE 64). Unlike
245 +// `writeFullKey` there is no scale and no RoPE (a gathered k_pe is already
246 +// rotated), and KPE_FP8 handles a k_pe that the gather left in the fp8 cache
247 +// layout -- so neither dtype needs an intermediate cast or a broadcast copy.
248 +template <typename scalar_t, bool FP8, bool KPE_FP8>
249 +__device__ __forceinline__ void writeFullKeyPack(void* dst,
250 + const scalar_t* k_nope,
251 + const void* k_pe, int laneId) {
252 + static_assert(!KPE_FP8 || FP8, "an fp8 k_pe requires an fp8 key output");
253 + constexpr int kElemSize = FP8 ? 1 : sizeof(scalar_t);
254 + auto* d = reinterpret_cast<uint8_t*>(dst);
255 + for (int e = laneId * kVecElems; e < kQkHeadDim; e += 16 * kVecElems) {
256 + void* out = d + e * kElemSize;
257 + const scalar_t* src = k_nope + e;
258 + if (e >= kQkNopeHeadDim) {
259 + int const rope_e = e - kQkNopeHeadDim;
260 + if constexpr (KPE_FP8) {
261 + *reinterpret_cast<uint2*>(out) = *reinterpret_cast<const uint2*>(
262 + reinterpret_cast<const uint8_t*>(k_pe) + rope_e);
263 + continue;
264 + }
265 + src = reinterpret_cast<const scalar_t*>(k_pe) + rope_e;
266 + }
267 + if constexpr (FP8) {
268 + copyChunk8UnitFp8(reinterpret_cast<uint8_t*>(out), src);
269 + } else {
270 + *reinterpret_cast<uint4*>(out) = *reinterpret_cast<const uint4*>(src);
271 + }
272 + }
273 +}
274 +
275 // Store a prefill query, rotating only q[..., 128:192]. For bf16 dst may alias
276 // q (in-place); fp8 writes the quantized query directly to its output.
277 template <typename scalar_t, bool FP8, bool APPLY_ROPE = false>
@@ -457,6 +527,72 @@ __global__ void fusedKimiK3MLAQKVQuantKVCacheFp8Kernel(
527 #endif
528 }
529
530 +// ────────────────────────────────────────────────────────────────────────────
531 +// K/V pack variant (chunked context): concatenate the strided kv_b_proj output
532 +// with k_pe into a contiguous row-major key, casting to E4M3 (and casting V
533 +// into its own contiguous output) when FP8. Unlike the other kernels this one
534 +// has no per-token cache slot, so a warp handles two (token, head) rows -- 16
535 +// lanes each -- and the grid is capped so a long context becomes a grid-stride
536 +// loop instead of a huge launch.
537 +// ────────────────────────────────────────────────────────────────────────────
538 +template <typename scalar_t, bool FP8, bool KPE_FP8>
539 +__global__ void fusedKimiK3MLAKVConcatPackKernel(
540 + const scalar_t* __restrict__ k_nope, int64_t const kn_tok_stride,
541 + int64_t const kn_head_stride, const void* __restrict__ k_pe,
542 + int64_t const k_pe_tok_stride, const scalar_t* __restrict__ v,
543 + int64_t const v_tok_stride, int64_t const v_head_stride,
544 + void* __restrict__ k_out, int64_t const ko_tok_stride,
545 + int64_t const ko_head_stride, uint8_t* __restrict__ v_fp8,
546 + int64_t const vo_tok_stride, int64_t const vo_head_stride,
547 + int const num_tokens, int const num_heads) {
548 + int const warpsPerBlock = blockDim.x / 32;
549 + int const laneId = threadIdx.x % 16;
550 + int const rowInWarp = (threadIdx.x % 32) / 16;
551 + int64_t const globalWarpIdx =
552 + static_cast<int64_t>(blockIdx.x) * warpsPerBlock + threadIdx.x / 32;
553 + int64_t const totalRows = static_cast<int64_t>(num_tokens) * num_heads;
554 + if (globalWarpIdx * 2 >= totalRows) return;
555 +
556 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
557 + cudaGridDependencySynchronize();
558 +#endif
559 +
560 + constexpr int kOutElemSize = FP8 ? 1 : sizeof(scalar_t);
561 + int64_t const gridStride =
562 + static_cast<int64_t>(gridDim.x) * warpsPerBlock * 2;
563 + for (int64_t row = globalWarpIdx * 2 + rowInWarp; row < totalRows;
564 + row += gridStride) {
565 + int const tokenIdx = static_cast<int>(row / num_heads);
566 + int const headIdx = static_cast<int>(row % num_heads);
567 + const void* pe;
568 + if constexpr (KPE_FP8) {
569 + pe = reinterpret_cast<const uint8_t*>(k_pe) + tokenIdx * k_pe_tok_stride;
570 + } else {
571 + pe = reinterpret_cast<const scalar_t*>(k_pe) + tokenIdx * k_pe_tok_stride;
572 + }
573 + writeFullKeyPack<scalar_t, FP8, KPE_FP8>(
574 + reinterpret_cast<uint8_t*>(k_out) +
575 + (tokenIdx * ko_tok_stride + headIdx * ko_head_stride) *
576 + kOutElemSize,
577 + k_nope + tokenIdx * kn_tok_stride + headIdx * kn_head_stride, pe,
578 + laneId);
579 +
580 + // bf16 V needs no cast, so it stays a strided view of the kv_b_proj output.
581 + if constexpr (FP8) {
582 + const scalar_t* vh =
583 + v + tokenIdx * v_tok_stride + headIdx * v_head_stride;
584 + uint8_t* vo = v_fp8 + tokenIdx * vo_tok_stride + headIdx * vo_head_stride;
585 + for (int e = laneId * kVecElems; e < kVHeadDim; e += 16 * kVecElems) {
586 + copyChunk8UnitFp8(vo + e, vh + e);
587 + }
588 + }
589 + }
590 +
591 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
592 + cudaTriggerProgrammaticLaunchCompletion();
593 +#endif
594 +}
595 +
596 // ────────────────────────────────────────────────────────────────────────────
597 // ds_mla variant: concat full key (bf16) + fp8_ds_mla latent cache insert
598 //
@@ -648,16 +784,20 @@ __global__ void fusedKimiK3MLADecodeQConcatDsMlaKernel(
784 #endif
785 }
786
651 -// PDL-aware launch of a (token, num_heads + 1)-warp grid.
787 +// PDL-aware launch of a (token, slots_per_token)-row grid, `rows_per_warp` rows
788 +// to a warp. `max_blocks > 0` caps the grid, leaving the kernel's grid-stride
789 +// loop to cover the remaining rows.
790 template <typename KernelT, typename... Args>
653 -static void launchPdl(KernelT kernel, int num_tokens, int num_heads,
654 - cudaStream_t stream, Args... args) {
791 +static void launchPdlSlots(KernelT kernel, int num_tokens, int slots_per_token,
792 + int rows_per_warp, int max_blocks,
793 + cudaStream_t stream, Args... args) {
794 constexpr int kBlockSize = 256;
795 constexpr int kWarpsPerBlock = kBlockSize / 32;
657 - int64_t const total_warps =
658 - static_cast<int64_t>(num_tokens) * (num_heads + 1);
659 - int const grid =
796 + int64_t const total_rows = static_cast<int64_t>(num_tokens) * slots_per_token;
797 + int64_t const total_warps = (total_rows + rows_per_warp - 1) / rows_per_warp;
798 + int grid =
799 static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
800 + if (max_blocks > 0 && grid > max_blocks) grid = max_blocks;
801 #ifndef USE_ROCM
802 static int const sm_version = getSMVersion();
803 cudaLaunchConfig_t config;
@@ -681,6 +821,13 @@ static void launchPdl(KernelT kernel, int num_tokens, int num_heads,
821 #endif
822 }
823
824 +// Cache-inserting kernels use one extra warp per token for the cache slot.
825 +template <typename KernelT, typename... Args>
826 +static void launchPdl(KernelT kernel, int num_tokens, int num_heads,
827 + cudaStream_t stream, Args... args) {
828 + launchPdlSlots(kernel, num_tokens, num_heads + 1, 1, 0, stream, args...);
829 +}
830 +
831 void checkBfloat16Support(torch::headeronly::ScalarType dtype) {
832 #ifndef USE_ROCM
833 if (dtype == torch::headeronly::ScalarType::BFloat16) {
@@ -913,6 +1060,128 @@ void fused_kimi_k3_mla_key_concat_ds_mla_insert(
1060 });
1061 }
1062
1063 +namespace {
1064 +// Shared checks for the chunked-context K/V pack ops. `k_nope` / `v` are
1065 +// strided views of one kv_b_proj output and `k_pe` is a strided view of the
1066 +// gathered context workspace, so only a unit last-dim stride (not full
1067 +// contiguity) is required of the inputs. Returns true when `k_pe` is still in
1068 +// the fp8 cache layout.
1069 +bool check_kv_concat_inputs(torch::stable::Tensor const& k_nope,
1070 + torch::stable::Tensor const& k_pe,
1071 + torch::stable::Tensor const& k_out,
1072 + torch::headeronly::ScalarType out_dtype) {
1073 + using torch::headeronly::ScalarType;
1074 + ScalarType const dt = k_nope.scalar_type();
1075 + STD_TORCH_CHECK(k_nope.device().is_cuda() && k_nope.dim() == 3 &&
1076 + k_nope.size(2) == 128 && k_nope.stride(2) == 1,
1077 + "k_nope must be CUDA [T, H, 128] with unit last stride");
1078 + bool const k_pe_is_fp8 = k_pe.scalar_type() == ScalarType::Float8_e4m3fn;
1079 + STD_TORCH_CHECK(
1080 + k_pe.device() == k_nope.device() && k_pe.dim() == 2 &&
1081 + k_pe.size(1) == 64 && k_pe.stride(1) == 1 &&
1082 + (k_pe.scalar_type() == dt ||
1083 + (k_pe_is_fp8 && out_dtype == ScalarType::Float8_e4m3fn)),
1084 + "k_pe must be CUDA [T, 64] with unit last-dim stride and use the input "
1085 + "dtype, or float8_e4m3fn when the key output is fp8");
1086 + STD_TORCH_CHECK(k_out.device() == k_nope.device() && k_out.is_contiguous() &&
1087 + k_out.scalar_type() == out_dtype && k_out.dim() == 3 &&
1088 + k_out.size(2) == 192 && k_out.size(0) == k_nope.size(0) &&
1089 + k_out.size(1) == k_nope.size(1) &&
1090 + k_pe.size(0) == k_nope.size(0),
1091 + "k_out must be a contiguous CUDA [T, H, 192] tensor of the "
1092 + "output dtype matching k_nope's token and head dimensions");
1093 + vllm::kimi_k3_fused_ops::checkBfloat16Support(dt);
1094 + return k_pe_is_fp8;
1095 +}
1096 +} // namespace
1097 +
1098 +void fused_kimi_k3_mla_kv_concat(
1099 + torch::stable::Tensor const& k_nope, // [T, H, 128] bf16/fp16
1100 + torch::stable::Tensor const& k_pe, // [T, 64] same dtype as k_nope
1101 + torch::stable::Tensor& k_out) { // [T, H, 192] same dtype, written
1102 + namespace kk3 = vllm::kimi_k3_fused_ops;
1103 + torch::headeronly::ScalarType const dt = k_nope.scalar_type();
1104 + check_kv_concat_inputs(k_nope, k_pe, k_out, dt);
1105 +
1106 + int const num_tokens = static_cast<int>(k_nope.size(0));
1107 + int const num_heads = static_cast<int>(k_nope.size(1));
1108 + if (num_tokens == 0) return;
1109 +
1110 + const torch::stable::accelerator::DeviceGuard device_guard(
1111 + k_nope.get_device_index());
1112 + const cudaStream_t stream =
1113 + get_current_cuda_stream(k_nope.get_device_index());
1114 +
1115 + VLLM_STABLE_DISPATCH_HALF_TYPES(dt, "fused_kimi_k3_mla_kv_concat", [&] {
1116 + kk3::launchPdlSlots(
1117 + kk3::fusedKimiK3MLAKVConcatPackKernel<scalar_t, false, false>,
1118 + num_tokens, num_heads, 2, get_device_prop()->multiProcessorCount * 8,
1119 + stream, reinterpret_cast<scalar_t const*>(k_nope.const_data_ptr()),
1120 + k_nope.stride(0), k_nope.stride(1), k_pe.const_data_ptr(),
1121 + k_pe.stride(0), static_cast<scalar_t const*>(nullptr), int64_t{0},
1122 + int64_t{0}, k_out.mutable_data_ptr(), k_out.stride(0), k_out.stride(1),
1123 + static_cast<uint8_t*>(nullptr), int64_t{0}, int64_t{0}, num_tokens,
1124 + num_heads);
1125 + });
1126 +}
1127 +
1128 +void fused_kimi_k3_mla_kv_concat_quant_fp8(
1129 + torch::stable::Tensor const& k_nope, // [T, H, 128] bf16/fp16
1130 + torch::stable::Tensor const& k_pe, // [T, 64] bf16/fp16/fp8
1131 + torch::stable::Tensor const& v, // [T, H, 128] bf16/fp16
1132 + torch::stable::Tensor& k_fp8, // [T, H, 192] fp8, written
1133 + torch::stable::Tensor& v_fp8) { // [T, H, 128] fp8, written
1134 + using torch::headeronly::ScalarType;
1135 + namespace kk3 = vllm::kimi_k3_fused_ops;
1136 + ScalarType const dt = k_nope.scalar_type();
1137 + bool const k_pe_is_fp8 =
1138 + check_kv_concat_inputs(k_nope, k_pe, k_fp8, ScalarType::Float8_e4m3fn);
1139 + STD_TORCH_CHECK(v.device() == k_nope.device() && v.scalar_type() == dt &&
1140 + v.dim() == 3 && v.size(2) == 128 && v.stride(2) == 1 &&
1141 + v.size(0) == k_nope.size(0) &&
1142 + v.size(1) == k_nope.size(1),
1143 + "v must match k_nope's dtype and token/head dimensions and "
1144 + "have shape [T, H, 128] with unit last stride");
1145 + STD_TORCH_CHECK(v_fp8.device() == k_nope.device() && v_fp8.is_contiguous() &&
1146 + v_fp8.scalar_type() == ScalarType::Float8_e4m3fn &&
1147 + v_fp8.dim() == 3 && v_fp8.size(2) == 128 &&
1148 + v_fp8.size(0) == k_nope.size(0) &&
1149 + v_fp8.size(1) == k_nope.size(1),
1150 + "v_fp8 must be a contiguous CUDA [T, H, 128] fp8 tensor "
1151 + "matching k_nope's token and head dimensions");
1152 +
1153 + int const num_tokens = static_cast<int>(k_nope.size(0));
1154 + int const num_heads = static_cast<int>(k_nope.size(1));
1155 + if (num_tokens == 0) return;
1156 +
1157 + const torch::stable::accelerator::DeviceGuard device_guard(
1158 + k_nope.get_device_index());
1159 + const cudaStream_t stream =
1160 + get_current_cuda_stream(k_nope.get_device_index());
1161 +
1162 + VLLM_STABLE_DISPATCH_HALF_TYPES(
1163 + dt, "fused_kimi_k3_mla_kv_concat_quant_fp8", [&] {
1164 + auto launch = [&](auto kernel) {
1165 + kk3::launchPdlSlots(
1166 + kernel, num_tokens, num_heads, 2,
1167 + get_device_prop()->multiProcessorCount * 8, stream,
1168 + reinterpret_cast<scalar_t const*>(k_nope.const_data_ptr()),
1169 + k_nope.stride(0), k_nope.stride(1), k_pe.const_data_ptr(),
1170 + k_pe.stride(0),
1171 + reinterpret_cast<scalar_t const*>(v.const_data_ptr()),
1172 + v.stride(0), v.stride(1), k_fp8.mutable_data_ptr(),
1173 + k_fp8.stride(0), k_fp8.stride(1),
1174 + reinterpret_cast<uint8_t*>(v_fp8.mutable_data_ptr()),
1175 + v_fp8.stride(0), v_fp8.stride(1), num_tokens, num_heads);
1176 + };
1177 + if (k_pe_is_fp8) {
1178 + launch(kk3::fusedKimiK3MLAKVConcatPackKernel<scalar_t, true, true>);
1179 + } else {
1180 + launch(kk3::fusedKimiK3MLAKVConcatPackKernel<scalar_t, true, false>);
1181 + }
1182 + });
1183 +}
1184 +
1185 void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(
1186 torch::stable::Tensor const& q, // [Tp, H, 192] bf16
1187 torch::stable::Tensor const& k_nope, // [Tp, H, 128] bf16
csrc/libtorch_stable/ops.h
+10
@@ -300,6 +300,16 @@ void fused_kimi_k3_mla_key_concat_ds_mla_insert(
300 std::optional<torch::stable::Tensor> position_ids,
301 std::optional<torch::stable::Tensor> cos_sin_cache);
302
303 +void fused_kimi_k3_mla_kv_concat(torch::stable::Tensor const& k_nope,
304 + torch::stable::Tensor const& k_pe,
305 + torch::stable::Tensor& k_out);
306 +
307 +void fused_kimi_k3_mla_kv_concat_quant_fp8(torch::stable::Tensor const& k_nope,
308 + torch::stable::Tensor const& k_pe,
309 + torch::stable::Tensor const& v,
310 + torch::stable::Tensor& k_fp8,
311 + torch::stable::Tensor& v_fp8);
312 +
313 void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(
314 torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope,
315 torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
csrc/libtorch_stable/torch_bindings.cpp
+11
@@ -466,6 +466,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
466 "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, "
467 "int cache_block_size, Tensor? position_ids=None, "
468 "Tensor? cos_sin_cache=None) -> ()");
469 + ops.def(
470 + "fused_kimi_k3_mla_kv_concat(Tensor k_nope, Tensor k_pe, Tensor! k_out) "
471 + "-> ()");
472 + ops.def(
473 + "fused_kimi_k3_mla_kv_concat_quant_fp8("
474 + "Tensor k_nope, Tensor k_pe, Tensor v, Tensor! k_fp8, Tensor! v_fp8) "
475 + "-> ()");
476 ops.def(
477 "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert("
478 "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, "
@@ -769,6 +776,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
776 TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert));
777 ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert",
778 TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert));
779 + ops.impl("fused_kimi_k3_mla_kv_concat",
780 + TORCH_BOX(&fused_kimi_k3_mla_kv_concat));
781 + ops.impl("fused_kimi_k3_mla_kv_concat_quant_fp8",
782 + TORCH_BOX(&fused_kimi_k3_mla_kv_concat_quant_fp8));
783 ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert",
784 TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert));
785 ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert",
tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py
+68 -2
@@ -9,6 +9,8 @@ from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import (
9 fused_mla_decode_q_concat_kv_cache_insert,
10 fused_mla_key_concat_ds_mla_insert,
11 fused_mla_key_concat_kv_cache_insert,
12 + fused_mla_kv_concat,
13 + fused_mla_kv_concat_quant_fp8,
14 fused_mla_qkv_quant_kv_cache_fp8_insert,
15 )
16 from vllm.platforms import current_platform
@@ -25,8 +27,8 @@ _POSITIONS = (1, 7, 13)
27 _SLOTS = (0, 3, 9)
28
29
28 -def _randn(*shape: int) -> torch.Tensor:
29 - return torch.randn(*shape, device="cuda", dtype=_DTYPE) * 0.2
30 +def _randn(*shape: int, dtype: torch.dtype = _DTYPE) -> torch.Tensor:
31 + return torch.randn(*shape, device="cuda", dtype=dtype) * 0.2
32
33
34 def _rope_cache(max_position: int = 32) -> torch.Tensor:
@@ -66,6 +68,70 @@ def _assert_fp8_close(actual: torch.Tensor, expected: torch.Tensor) -> None:
68 )
69
70
71 +def _strided_context_inputs(
72 + num_tokens: int, num_heads: int, dtype: torch.dtype, k_pe_fp8: bool
73 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
74 + """The real prefill-context layouts: ``k_nope``/``v`` as the two strided
75 + halves of one kv_b_proj output and ``k_pe`` as a column slice of the gather
76 + workspace (left in the fp8 cache layout for a plain fp8 cache)."""
77 + kv_nope = _randn(num_tokens, num_heads, 256, dtype=dtype)
78 + k_nope, v = kv_nope.split((128, 128), dim=-1)
79 + workspace = _randn(num_tokens, 576, dtype=dtype)
80 + if k_pe_fp8:
81 + workspace = workspace.to(torch.float8_e4m3fn)
82 + return k_nope, workspace[:, 512:].unsqueeze(1), v
83 +
84 +
85 +@pytest.mark.parametrize("num_tokens", [0, _NUM_TOKENS])
86 +@pytest.mark.parametrize("num_heads", [3, _NUM_HEADS])
87 +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16])
88 +@torch.inference_mode()
89 +def test_context_kv_concat_accepts_strided_inputs(
90 + num_tokens: int, num_heads: int, input_dtype: torch.dtype
91 +) -> None:
92 + k_nope, k_pe, _ = _strided_context_inputs(
93 + num_tokens, num_heads, input_dtype, k_pe_fp8=False
94 + )
95 +
96 + k = fused_mla_kv_concat(k_nope, k_pe)
97 +
98 + assert k.is_contiguous()
99 + assert k.dtype == input_dtype
100 + torch.testing.assert_close(k[..., :128], k_nope, atol=0, rtol=0)
101 + torch.testing.assert_close(
102 + k[..., 128:], k_pe.expand(-1, num_heads, -1), atol=0, rtol=0
103 + )
104 +
105 +
106 +@pytest.mark.parametrize("k_pe_dtype", ["input", "fp8"])
107 +@pytest.mark.parametrize("num_tokens", [0, _NUM_TOKENS])
108 +@pytest.mark.parametrize("num_heads", [3, _NUM_HEADS])
109 +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16])
110 +@torch.inference_mode()
111 +def test_context_kv_pack_quantizes_strided_inputs(
112 + k_pe_dtype: str,
113 + num_tokens: int,
114 + num_heads: int,
115 + input_dtype: torch.dtype,
116 +) -> None:
117 + k_nope, k_pe, v = _strided_context_inputs(
118 + num_tokens, num_heads, input_dtype, k_pe_fp8=k_pe_dtype == "fp8"
119 + )
120 +
121 + k_actual, v_actual = fused_mla_kv_concat_quant_fp8(k_nope, k_pe, v)
122 +
123 + fp8 = torch.float8_e4m3fn
124 + k_expected = torch.empty_like(k_actual)
125 + k_expected[..., :128] = k_nope.to(fp8)
126 + k_expected[..., 128:] = k_pe.to(fp8)
127 + assert k_actual.is_contiguous()
128 + assert v_actual.is_contiguous()
129 + # The pack casts with the native pairwise converters, so it must be
130 + # bit-identical to torch's `.to(fp8)` rather than merely close.
131 + torch.testing.assert_close(k_actual.float(), k_expected.float(), atol=0, rtol=0)
132 + torch.testing.assert_close(v_actual.float(), v.to(fp8).float(), atol=0, rtol=0)
133 +
134 +
135 @pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"])
136 @torch.inference_mode()
137 def test_prefill_epilogue_fuses_gptj_rope(cache_kind: str) -> None:
tests/models/kimi_k3/test_mla_prefill_context.py new
+319
@@ -0,0 +1,319 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""K3's fused chunked-context prefill must match the generic MLA impl.
4 +
5 +The layer owns its own context loop so it can fuse the per-chunk K/V pack and
6 +skip re-quantizing an already-quantized query. That is only safe if it feeds the
7 +prefill backend exactly what ``MLACommonBaseImpl._compute_prefill_context``
8 +would, chunk for chunk.
9 +"""
10 +
11 +from types import SimpleNamespace
12 +
13 +import pytest
14 +import torch
15 +
16 +from vllm.model_executor.layers.attention.mla_attention import (
17 + MLACommonBaseImpl,
18 + MLACommonPrefillMetadata,
19 + build_mla_chunked_context_metadata,
20 +)
21 +from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention
22 +from vllm.platforms import current_platform
23 +
24 +pytestmark = pytest.mark.skipif(
25 + not current_platform.is_cuda(), reason="Kimi-K3 fused MLA requires CUDA"
26 +)
27 +
28 +_KV_LORA_RANK = 512
29 +_QK_NOPE = 128
30 +_QK_ROPE = 64
31 +_V_HEAD_DIM = 128
32 +_ENTRY = _KV_LORA_RANK + _QK_ROPE
33 +_NUM_HEADS = 2
34 +_BLOCK_SIZE = 16
35 +_WORKSPACE_TOKENS = 128
36 +# Splits one long request across chunks, packs short ones together, and leaves
37 +# the last request without any context.
38 +_CONTEXT_LENS = [200, 48, 32, 0]
39 +_QUERY_LENS = [8, 4, 6, 5]
40 +
41 +
42 +class _RecordingPrefillBackend:
43 + """Records what each chunk is asked to attend over.
44 +
45 + ``honors_out`` mimics the two backend families: one that writes into a
46 + caller-provided ``out`` (trtllm_ragged, flashinfer, tokenspeed) and one that
47 + always returns its own buffer (flash_attn with a padded V, aiter).
48 + """
49 +
50 + def __init__(self, honors_out: bool = False) -> None:
51 + self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
52 + self.out_destinations: list[torch.Tensor | None] = []
53 + self._honors_out = honors_out
54 +
55 + @staticmethod
56 + def get_name() -> str:
57 + return "recording"
58 +
59 + def supports_out(self) -> bool:
60 + return self._honors_out
61 +
62 + def run_prefill_context_chunk(self, *, chunk, q, k, v, out=None):
63 + self.calls.append((q.float().clone(), k.float().clone(), v.float().clone()))
64 + self.out_destinations.append(out)
65 + assert out is None or self._honors_out
66 + # Fold K/V into the partial so any packing difference shows up in the
67 + # merged context output, not just in the recorded calls.
68 + digest = (k.float().mean() + v.float().mean()).item()
69 + num_q = q.shape[0]
70 + if out is None:
71 + out = torch.empty(
72 + (num_q, _NUM_HEADS, _V_HEAD_DIM), device=q.device, dtype=torch.bfloat16
73 + )
74 + else:
75 + assert out.shape == (num_q, _NUM_HEADS, _V_HEAD_DIM)
76 + out.fill_(digest)
77 + lse = torch.full(
78 + (_NUM_HEADS, num_q),
79 + 1.0 + chunk.index,
80 + device=q.device,
81 + dtype=torch.float32,
82 + )
83 + return out, lse
84 +
85 +
86 +class _KVBProj(torch.nn.Module):
87 + """Stand-in for the layer's ``kv_b_proj`` (returns an (out, bias) tuple).
88 +
89 + Enforces the same input contract as the real linear methods, because that is
90 + what decides whether the gathered latent needs a cast:
91 +
92 + * an fp8 weight consumes the fp8 latent directly and dequantizes internally,
93 + so it takes fp8 or bf16;
94 + * a bf16 weight -- what a stock K3 checkpoint carries -- is a plain
95 + ``F.linear`` and rejects anything but bf16, exactly as torch does.
96 +
97 + Either weight dtype produces a bf16 output.
98 + """
99 +
100 + def __init__(self, device: torch.device, weight_dtype: torch.dtype) -> None:
101 + super().__init__()
102 + weight = (
103 + torch.randn(
104 + _NUM_HEADS * (_QK_NOPE + _V_HEAD_DIM),
105 + _KV_LORA_RANK,
106 + device=device,
107 + dtype=torch.bfloat16,
108 + )
109 + * 0.05
110 + )
111 + self.register_buffer("weight", weight.to(weight_dtype))
112 +
113 + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]:
114 + if self.weight.dtype == torch.bfloat16 and x.dtype != torch.bfloat16:
115 + raise RuntimeError(
116 + "a bfloat16 kv_b_proj cannot consume the gathered latent as "
117 + f"{x.dtype}; it must be cast first"
118 + )
119 + return torch.nn.functional.linear(
120 + x.to(torch.bfloat16), self.weight.to(torch.bfloat16)
121 + ), None
122 +
123 +
124 +class _FusedLayer:
125 + """Only the attributes K3's context loop reads."""
126 +
127 + _compute_prefill_context = MultiHeadLatentAttention._compute_prefill_context
128 + _gather_context_latent = MultiHeadLatentAttention._gather_context_latent
129 + _attn_read_kv_cache = MultiHeadLatentAttention._attn_read_kv_cache
130 +
131 + def __init__(self, kv_b_proj, kv_cache, kv_cache_dtype, k_scale) -> None:
132 + self.kv_b_proj = kv_b_proj
133 + self.kv_cache = kv_cache
134 + self.kv_cache_dtype = kv_cache_dtype
135 + self._k_scale = k_scale
136 + self.kv_lora_rank = _KV_LORA_RANK
137 + self.num_local_heads = _NUM_HEADS
138 + self.qk_nope_head_dim = _QK_NOPE
139 + self.v_head_dim = _V_HEAD_DIM
140 +
141 +
142 +class _ReferenceImpl:
143 + """Only the attributes the generic context loop reads."""
144 +
145 + _compute_prefill_context = MLACommonBaseImpl._compute_prefill_context
146 + _concat_k_nope_k_pe = MLACommonBaseImpl._concat_k_nope_k_pe
147 + _use_flashinfer_concat_mla_k = False
148 +
149 + def __init__(self, kv_b_proj, kv_cache_dtype) -> None:
150 + self.kv_b_proj = kv_b_proj
151 + self.kv_cache_dtype = kv_cache_dtype
152 + self.kv_lora_rank = _KV_LORA_RANK
153 + self.num_heads = _NUM_HEADS
154 + self.qk_nope_head_dim = _QK_NOPE
155 + self.qk_rope_head_dim = _QK_ROPE
156 + self.v_head_dim = _V_HEAD_DIM
157 +
158 +
159 +def _build_prefill_metadata(
160 + device: torch.device,
161 + workspace_dtype: torch.dtype,
162 + q_data_type: torch.dtype,
163 + backend: _RecordingPrefillBackend,
164 +) -> MLACommonPrefillMetadata:
165 + query_start_loc_cpu = torch.zeros(len(_QUERY_LENS) + 1, dtype=torch.int32)
166 + query_start_loc_cpu[1:] = torch.tensor(_QUERY_LENS, dtype=torch.int32).cumsum(0)
167 + workspace = torch.empty(
168 + (_WORKSPACE_TOKENS, _ENTRY), dtype=workspace_dtype, device=device
169 + )
170 + chunked_context = build_mla_chunked_context_metadata(
171 + context_lens_cpu=torch.tensor(_CONTEXT_LENS, dtype=torch.int32),
172 + prefill_query_start_loc_cpu=query_start_loc_cpu,
173 + chunked_prefill_workspace=workspace,
174 + chunked_prefill_workspace_size=_WORKSPACE_TOKENS,
175 + block_size=_BLOCK_SIZE,
176 + align_chunk_to_block=True,
177 + device=device,
178 + dcp_world_size=1,
179 + dcp_local_block_size=1,
180 + dcp_virtual_block_size=1,
181 + )
182 + assert chunked_context is not None
183 + assert len(chunked_context.chunks) > 1, "the batch must exercise accumulation"
184 +
185 + max_blocks = (max(_CONTEXT_LENS) + max(_QUERY_LENS)) // _BLOCK_SIZE + 1
186 + num_blocks = max_blocks * len(_CONTEXT_LENS)
187 + block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).view(
188 + len(_CONTEXT_LENS), max_blocks
189 + )
190 + return MLACommonPrefillMetadata(
191 + block_table=block_table,
192 + query_start_loc=query_start_loc_cpu.to(device),
193 + max_query_len=max(_QUERY_LENS),
194 + chunked_context=chunked_context,
195 + q_data_type=q_data_type,
196 + output_dtype=torch.bfloat16,
197 + prefill_backend=backend,
198 + ), num_blocks
199 +
200 +
201 +@pytest.mark.parametrize("honors_out", [False, True], ids=["copy_out", "writes_out"])
202 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
203 +@pytest.mark.parametrize(
204 + "kv_b_proj_quantized", [True, False], ids=["fp8_kv_b_proj", "bf16_kv_b_proj"]
205 +)
206 +@torch.inference_mode()
207 +def test_fused_context_matches_generic_impl(
208 + kv_b_proj_quantized: bool, kv_cache_dtype: str, honors_out: bool
209 +) -> None:
210 + """Cache dtype and ``kv_b_proj`` dtype vary independently.
211 +
212 + A stock K3 checkpoint pairs a bf16 ``kv_b_proj`` with an fp8 cache, so the
213 + fused loop cannot assume the gathered latent is already in the dtype
214 + ``kv_b_proj`` accepts.
215 + """
216 + torch.manual_seed(0)
217 + device = torch.device("cuda")
218 + fp8 = current_platform.fp8_dtype()
219 + quantized = kv_cache_dtype == "fp8"
220 + # An fp8 cache is read by an fp8 query, so the workspace keeps the fp8
221 + # layout; a bf16 cache dequantizes into a bf16 workspace.
222 + q_data_type = fp8 if quantized else torch.bfloat16
223 + workspace_dtype = q_data_type
224 +
225 + kv_b_proj = _KVBProj(
226 + device, weight_dtype=fp8 if kv_b_proj_quantized else torch.bfloat16
227 + )
228 + k_scale = torch.ones(1, dtype=torch.float32, device=device)
229 +
230 + backend_fused = _RecordingPrefillBackend(honors_out=honors_out)
231 + # The reference impl never passes `out`, so it always allocates its own.
232 + backend_ref = _RecordingPrefillBackend()
233 + prefill_fused, num_blocks = _build_prefill_metadata(
234 + device, workspace_dtype, q_data_type, backend_fused
235 + )
236 + prefill_ref, _ = _build_prefill_metadata(
237 + device, workspace_dtype, q_data_type, backend_ref
238 + )
239 +
240 + cache = torch.randn(
241 + (num_blocks, _BLOCK_SIZE, _ENTRY), device=device, dtype=torch.bfloat16
242 + )
243 + kv_cache = cache.to(fp8) if quantized else cache
244 + q = (
245 + torch.randn(
246 + (sum(_QUERY_LENS), _NUM_HEADS, _QK_NOPE + _QK_ROPE),
247 + device=device,
248 + dtype=torch.bfloat16,
249 + )
250 + * 0.2
251 + ).to(q_data_type)
252 +
253 + layer = _FusedLayer(kv_b_proj, kv_cache, kv_cache_dtype, k_scale)
254 + fused_out, fused_lse = layer._compute_prefill_context(
255 + q, SimpleNamespace(prefill=prefill_fused)
256 + )
257 +
258 + impl = _ReferenceImpl(kv_b_proj, kv_cache_dtype)
259 + ref_out, ref_lse = impl._compute_prefill_context(
260 + q, kv_cache, SimpleNamespace(prefill=prefill_ref), k_scale
261 + )
262 +
263 + assert len(backend_fused.calls) == len(backend_ref.calls)
264 + for chunk_idx, (fused_call, ref_call) in enumerate(
265 + zip(backend_fused.calls, backend_ref.calls, strict=True)
266 + ):
267 + for name, fused_t, ref_t in zip(
268 + ("q", "k", "v"), fused_call, ref_call, strict=True
269 + ):
270 + torch.testing.assert_close(
271 + fused_t,
272 + ref_t,
273 + atol=0,
274 + rtol=0,
275 + msg=lambda m, n=name, i=chunk_idx: f"chunk {i} {n} differs: {m}",
276 + )
277 + torch.testing.assert_close(fused_out, ref_out, atol=0, rtol=0)
278 + torch.testing.assert_close(fused_lse, ref_lse, atol=0, rtol=0)
279 +
280 + # Every chunk but the continuation should have been written in place, i.e.
281 + # straight into the returned accumulator, with no intermediate copy.
282 + wrote_in_place = [
283 + out is not None and out.data_ptr() == fused_out[chunk.token_slice].data_ptr()
284 + for out, chunk in zip(
285 + backend_fused.out_destinations,
286 + prefill_fused.chunked_context.chunks,
287 + strict=True,
288 + )
289 + ]
290 + continuations = [c.is_continuation for c in prefill_fused.chunked_context.chunks]
291 + assert any(continuations), "the batch must exercise a continuation chunk"
292 + if honors_out:
293 + assert wrote_in_place == [not c for c in continuations]
294 + else:
295 + assert not any(wrote_in_place)
296 +
297 +
298 +@torch.inference_mode()
299 +def test_fused_context_rejects_an_unquantized_query() -> None:
300 + """The fp8 query is produced by the new-token epilogue, not re-cast here."""
301 + torch.manual_seed(0)
302 + device = torch.device("cuda")
303 + fp8 = current_platform.fp8_dtype()
304 + prefill, num_blocks = _build_prefill_metadata(
305 + device, fp8, fp8, _RecordingPrefillBackend()
306 + )
307 + layer = _FusedLayer(
308 + _KVBProj(device, weight_dtype=fp8),
309 + torch.zeros((num_blocks, _BLOCK_SIZE, _ENTRY), device=device, dtype=fp8),
310 + "fp8",
311 + torch.ones(1, dtype=torch.float32, device=device),
312 + )
313 + q = torch.zeros(
314 + (sum(_QUERY_LENS), _NUM_HEADS, _QK_NOPE + _QK_ROPE),
315 + device=device,
316 + dtype=torch.bfloat16,
317 + )
318 + with pytest.raises(AssertionError, match="new-token epilogue"):
319 + layer._compute_prefill_context(q, SimpleNamespace(prefill=prefill))
tests/v1/attention/test_mla_prefill_registry.py
+2 -2
@@ -24,7 +24,7 @@ class CustomMLAPrefillBackend(MLAPrefillBackend):
24 def run_prefill_new_tokens(self, q, k, v, return_softmax_lse):
25 raise NotImplementedError
26
27 - def run_prefill_context_chunk(self, chunk, q, k, v):
27 + def run_prefill_context_chunk(self, chunk, q, k, v, out=None):
28 raise NotImplementedError
29
30
@@ -112,7 +112,7 @@ def test_register_custom_backend_as_decorator():
112 def run_prefill_new_tokens(self, q, k, v, return_softmax_lse):
113 raise NotImplementedError
114
115 - def run_prefill_context_chunk(self, chunk, q, k, v):
115 + def run_prefill_context_chunk(self, chunk, q, k, v, out=None):
116 raise NotImplementedError
117
118 assert MLAPrefillBackendEnum.CUSTOM.is_overridden()
vllm/model_executor/layers/attention/mla_attention.py
+32 -6
@@ -2414,6 +2414,22 @@ def reorg_kvcache(
2414 return reorganized_kv_c_normed, reorganized_k_pe
2415
2416
2417 +def neutralize_empty_context_partials(
2418 + chunked_context: "MLACommonPrefillMetadata.ChunkedContextMetadata",
2419 + output: torch.Tensor,
2420 + output_lse: torch.Tensor,
2421 +) -> None:
2422 + """Neutralize the partial of every prefill that no chunk covers.
2423 +
2424 + A prefill without context is never gathered, so nothing would write its rows;
2425 + a zero output with an ``-inf`` lse carries no weight into the final merge
2426 + against the suffix partial.
2427 + """
2428 + for token_slice in chunked_context.empty_token_slices:
2429 + output[token_slice].zero_()
2430 + output_lse[:, token_slice].fill_(float("-inf"))
2431 +
2432 +
2433 def init_mla_context_partial(
2434 chunked_context: "MLACommonPrefillMetadata.ChunkedContextMetadata",
2435 attn_output: torch.Tensor,
@@ -2423,7 +2439,9 @@ def init_mla_context_partial(
2439 """Allocate the running context partial over all prefill tokens.
2440
2441 Laid out like the chunk partials so the final whole-batch merge against the
2426 - suffix partial sees matching head strides.
2442 + suffix partial sees matching head strides. Callers whose backend honors an
2443 + ``out`` tensor already know that layout and can allocate directly, pairing it
2444 + with ``neutralize_empty_context_partials``.
2445 """
2446 output = torch.empty(
2447 (num_tokens, *attn_output.shape[1:]),
@@ -2435,10 +2453,7 @@ def init_mla_context_partial(
2453 dtype=attn_softmax_lse.dtype,
2454 device=attn_softmax_lse.device,
2455 )
2438 - # No chunk covers a prefill without context, so neutralize its partial.
2439 - for token_slice in chunked_context.empty_token_slices:
2440 - output[token_slice].zero_()
2441 - output_lse[:, token_slice].fill_(float("-inf"))
2456 + neutralize_empty_context_partials(chunked_context, output, output_lse)
2457 return output, output_lse
2458
2459
@@ -2448,16 +2463,26 @@ def accumulate_mla_context_chunk(
2463 attn_softmax_lse: torch.Tensor,
2464 output: torch.Tensor,
2465 output_lse: torch.Tensor,
2466 + output_written: bool = False,
2467 ) -> None:
2468 """Fold one chunk's partial into the running context partial.
2469
2470 Only the first request may be a continuation; its tokens are merged and the
2471 remaining token range is initialized.
2472 +
2473 + Args:
2474 + output_written: The chunk's attention output already landed in
2475 + ``output[chunk.token_slice]`` because the backend was handed it as
2476 + ``out``, leaving only the lse to fold. Invalid for a continuation
2477 + chunk, whose leading tokens must be merged rather than overwritten.
2478 """
2479 token_start = chunk.token_slice.start
2480 token_end = chunk.token_slice.stop
2481 init_start = token_start
2482 if chunk.is_continuation:
2483 + assert not output_written, (
2484 + "a continuation chunk must not write over the partial it merges with"
2485 + )
2486 init_start = chunk.continuation_token_end
2487 num_merged = init_start - token_start
2488 merge_attn_states(
@@ -2470,7 +2495,8 @@ def accumulate_mla_context_chunk(
2495 )
2496 if init_start < token_end:
2497 written = init_start - token_start
2473 - output[init_start:token_end].copy_(attn_output[written:])
2498 + if not output_written:
2499 + output[init_start:token_end].copy_(attn_output[written:])
2500 output_lse[:, init_start:token_end].copy_(attn_softmax_lse[:, written:])
2501
2502
vllm/models/kimi_k3/nvidia/mla.py
+192 -6
@@ -8,8 +8,9 @@ This is a self-contained MLA layer that owns the full attention path:
8 -> fused pre-attention ops (fused_qkv_a_proj / norms / q_b_proj)
9 -> explicit prefill / decode split
10 prefill: fused key-concat + cache-insert kernel -> run_prefill_new_tokens
11 - (+ chunked-context merge); dispatched by cache dtype
12 - (bf16 / plain fp8 / fp8_ds_mla)
11 + (+ chunked-context merge, whose per-chunk gather -> kv_b_proj
12 + -> fused K/V pack loop this layer owns); dispatched by cache
13 + dtype (bf16 / plain fp8 / fp8_ds_mla)
14 decode : W_UK absorb (BMM1) -> fused q-concat + cache-insert kernel
15 -> impl.forward_mqa -> W_UV up-proj (MQA)
16 -> optional output gate
@@ -34,6 +35,7 @@ from typing import TYPE_CHECKING, cast
35 import torch
36 from torch import nn
37
38 +from vllm import _custom_ops as ops
39 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
40 from vllm.config import (
41 CacheConfig,
@@ -48,6 +50,12 @@ from vllm.model_executor.layers.attention.attention import (
50 set_default_quant_scales,
51 should_load_quant_weights,
52 )
53 +from vllm.model_executor.layers.attention.mla_attention import (
54 + _get_kv_b_proj_input_dtype,
55 + accumulate_mla_context_chunk,
56 + init_mla_context_partial,
57 + neutralize_empty_context_partials,
58 +)
59 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
60 from vllm.model_executor.layers.layernorm import RMSNorm
61 from vllm.model_executor.layers.linear import (
@@ -67,6 +75,8 @@ from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import (
75 fused_mla_decode_q_concat_kv_cache_insert,
76 fused_mla_key_concat_ds_mla_insert,
77 fused_mla_key_concat_kv_cache_insert,
78 + fused_mla_kv_concat,
79 + fused_mla_kv_concat_quant_fp8,
80 fused_mla_qkv_quant_kv_cache_fp8_insert,
81 )
82 from vllm.platforms import current_platform
@@ -696,6 +706,181 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
706 cos_sin_cache=cos_sin_cache,
707 )
708
709 + def _compute_prefill_context(
710 + self,
711 + q: torch.Tensor,
712 + attn_metadata: "MLACommonMetadata",
713 + ) -> tuple[torch.Tensor, torch.Tensor]:
714 + """Chunked-context prefill, K3-fused. Replaces the impl's version.
715 +
716 + Per chunk the impl gathers the paged latent, up-projects it, then casts
717 + and concatenates K (and casts V) in two or three more launches. Here that
718 + tail is one fused kernel per chunk -- ``fused_mla_kv_concat`` for a bf16
719 + query, ``fused_mla_kv_concat_quant_fp8`` when the query is fp8 -- reading
720 + the strided ``kv_b_proj`` output in place and writing a contiguous key, so
721 + only the gather and ``kv_b_proj`` remain.
722 +
723 + The impl's query cast is gone as well: ``q`` already carries
724 + ``prefill.q_data_type`` because the new-token epilogue quantized it. The
725 + gathered latent still gets the impl's cast to whatever ``kv_b_proj``
726 + consumes -- free (a no-op ``.to``) for a checkpoint whose ``kv_b_proj``
727 + takes the fp8 latent directly, and required for a bf16 one, which is
728 + what a stock K3 checkpoint carries. Its output is bf16 either way.
729 +
730 + The gathered ``k_pe`` is likewise used as-is (fp8 for a plain fp8 cache)
731 + and needs no RoPE: it was rotated on the way in.
732 +
733 + Chunk partials are written straight into the accumulating context partial
734 + when the prefill backend honors ``out``, so only the (64x smaller) lse is
735 + copied per chunk.
736 +
737 + Decode context parallelism keeps using
738 + ``impl._context_parallel_compute_prefill_context``; its extra allgather
739 + and reorg are not fused here.
740 + """
741 + prefill = attn_metadata.prefill
742 + assert prefill is not None
743 + prefill_backend = prefill.prefill_backend
744 + assert prefill_backend is not None
745 + chunked_context = prefill.chunked_context
746 + assert chunked_context is not None
747 + assert q.dtype == prefill.q_data_type, (
748 + "Kimi-K3 chunked context expects the new-token epilogue to have "
749 + f"produced a {prefill.q_data_type} query; got {q.dtype}."
750 + )
751 +
752 + fp8_prefill = q.dtype == current_platform.fp8_dtype()
753 + workspace = chunked_context.workspace
754 + kv_cache = self._attn_read_kv_cache()
755 + kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(self.kv_b_proj, fp8_prefill)
756 +
757 + def run_chunk(
758 + chunk, out: torch.Tensor | None = None
759 + ) -> tuple[torch.Tensor, torch.Tensor]:
760 + self._gather_context_latent(chunk, kv_cache, prefill, fp8_prefill)
761 + gathered = workspace[: chunk.num_context_tokens]
762 + kv_c_normed = gathered[..., : self.kv_lora_rank]
763 + if kv_b_proj_input_dtype is not None:
764 + kv_c_normed = kv_c_normed.to(kv_b_proj_input_dtype)
765 + kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
766 + -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
767 + )
768 + k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
769 + k_pe = gathered[..., self.kv_lora_rank :]
770 + if fp8_prefill:
771 + k, v = fused_mla_kv_concat_quant_fp8(k_nope, k_pe, v)
772 + else:
773 + k = fused_mla_kv_concat(k_nope, k_pe)
774 + attn_output, attn_lse = prefill_backend.run_prefill_context_chunk(
775 + chunk=chunk, q=q[chunk.token_slice], k=k, v=v, out=out
776 + )
777 + assert out is None or attn_output.data_ptr() == out.data_ptr(), (
778 + f"{prefill_backend.get_name()} reports supports_out() but did not "
779 + "write the context chunk into the `out` it was given."
780 + )
781 + return attn_output, attn_lse
782 +
783 + chunks = chunked_context.chunks
784 + if len(chunks) == 1 and not chunked_context.empty_token_slices:
785 + # One chunk covering every prefill token: its partial *is* the context
786 + # partial, so it needs neither an accumulator nor a copy.
787 + return run_chunk(chunks[0])
788 +
789 + # A backend honoring `out` writes each chunk's partial straight into the
790 + # accumulator, so the per-chunk output copy disappears -- and because that
791 + # contract fixes the trailing shape, the accumulator can be sized before
792 + # any chunk runs. Otherwise the shape is only knowable from a real partial,
793 + # so the first chunk runs ahead of the loop and is copied in.
794 + writes_out = prefill_backend.supports_out()
795 + if writes_out:
796 + assert prefill.output_dtype is not None
797 + output = torch.empty(
798 + (q.shape[0], self.num_local_heads, self.v_head_dim),
799 + dtype=prefill.output_dtype,
800 + device=q.device,
801 + )
802 + output_lse = torch.empty(
803 + (self.num_local_heads, q.shape[0]),
804 + dtype=torch.float32,
805 + device=q.device,
806 + )
807 + neutralize_empty_context_partials(chunked_context, output, output_lse)
808 + else:
809 + attn_output, attn_lse = run_chunk(chunks[0])
810 + output, output_lse = init_mla_context_partial(
811 + chunked_context, attn_output, attn_lse, num_tokens=q.shape[0]
812 + )
813 + accumulate_mla_context_chunk(
814 + chunks[0], attn_output, attn_lse, output, output_lse
815 + )
816 + chunks = chunks[1:]
817 +
818 + for chunk in chunks:
819 + # A continuation chunk's leading tokens have to be merged with the
820 + # partial already sitting there, so it cannot write in place.
821 + out = (
822 + output[chunk.token_slice]
823 + if writes_out and not chunk.is_continuation
824 + else None
825 + )
826 + attn_output, attn_lse = run_chunk(chunk, out=out)
827 + accumulate_mla_context_chunk(
828 + chunk,
829 + attn_output,
830 + attn_lse,
831 + output,
832 + output_lse,
833 + output_written=out is not None,
834 + )
835 + return output, output_lse
836 +
837 + def _gather_context_latent(
838 + self,
839 + chunk,
840 + kv_cache: torch.Tensor,
841 + prefill,
842 + fp8_prefill: bool,
843 + ) -> None:
844 + """Gather one chunk's paged context latent into the workspace.
845 +
846 + Dispatched exactly as in ``impl._compute_prefill_context``: an fp8 query
847 + reads the plain fp8 cache in its stored layout, anything else lands in the
848 + workspace as the model dtype.
849 + """
850 + workspace = prefill.chunked_context.workspace
851 + toks = chunk.num_context_tokens
852 + block_table = prefill.block_table[chunk.request_slice]
853 + if self.kv_cache_dtype == "fp8_ds_mla":
854 + ops.cp_gather_and_upconvert_fp8_kv_cache(
855 + src_cache=kv_cache,
856 + dst=workspace[:toks],
857 + block_table=block_table,
858 + workspace_starts=chunk.cu_seq_lens,
859 + batch_size=chunk.num_requests,
860 + seq_starts=chunk.starts,
861 + )
862 + elif not fp8_prefill:
863 + ops.gather_and_maybe_dequant_cache(
864 + src_cache=kv_cache,
865 + dst=workspace,
866 + block_table=block_table,
867 + cu_seq_lens=chunk.cu_seq_lens,
868 + token_to_seq=chunk.token_to_seq,
869 + num_tokens=toks,
870 + kv_cache_dtype=self.kv_cache_dtype,
871 + scale=self._k_scale,
872 + seq_starts=chunk.starts,
873 + )
874 + else:
875 + ops.cp_gather_cache(
876 + src_cache=kv_cache,
877 + dst=workspace[:toks],
878 + block_table=block_table,
879 + cu_seq_lens=chunk.cu_seq_lens,
880 + batch_size=chunk.num_requests,
881 + seq_starts=chunk.starts,
882 + )
883 +
884 def _forward_prefill_fused(
885 self,
886 q: torch.Tensor,
@@ -710,8 +895,9 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
895 """Prefill using the fused key-concat + cache-insert kernel.
896
897 Replaces ``_concat_k_nope_k_pe`` and the prefill cache write with one
713 - fused kernel launch, dispatched by cache dtype. The chunked context
714 - gather + online-softmax merge are delegated to the impl.
898 + fused kernel launch, dispatched by cache dtype. Chunked context runs
899 + through this layer's ``_compute_prefill_context``, except under DCP where
900 + it is delegated to the impl.
901
902 Supported configs (K3 fp8 policy):
903 - bf16 cache -> bf16 prefill query
@@ -812,8 +998,8 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
998 )
999 )
1000 else:
815 - context_output, context_lse = self.impl._compute_prefill_context( # type: ignore[attr-defined]
816 - q, self._attn_read_kv_cache(), attn_metadata, self._k_scale
1001 + context_output, context_lse = self._compute_prefill_context(
1002 + q, attn_metadata
1003 )
1004 suffix_output, suffix_lse = output_prefill
1005 out = out.view(-1, self.num_local_heads, self.v_head_dim)
vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py
+56
@@ -12,6 +12,9 @@ mirror ``fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_{bf16,fp8}_insert``.
12 - ``fused_mla_qkv_quant_kv_cache_fp8_insert`` (fp8): additionally quantize
13 ``q``/``k``/``v`` to E4M3 with ``q_scale`` / ``k_scale`` / ``v_scale`` (the
14 cache shares ``k_scale``, as in ``concat_and_cache_mla``).
15 +- ``fused_mla_kv_concat`` / ``fused_mla_kv_concat_quant_fp8`` (chunked context):
16 + the same K concat (plus the fp8 K/V cast) without the cache insert, for context
17 + chunks whose latent was gathered back out of the paged cache.
18
19 The optional ``positions`` / ``cos_sin_cache`` pair enables GPT-J-style RoPE
20 inside the epilogue. Omitting both keeps the K3 NoPE fast path. The kernels use
@@ -156,6 +159,59 @@ def fused_mla_qkv_quant_kv_cache_fp8_insert(
159 return q_fp8, k_fp8, v_fp8
160
161
162 +def _empty_full_key(
163 + k_nope: torch.Tensor, k_pe: torch.Tensor, dtype: torch.dtype
164 +) -> torch.Tensor:
165 + num_tokens, num_heads, qk_nope_head_dim = k_nope.shape
166 + return torch.empty(
167 + (num_tokens, num_heads, qk_nope_head_dim + k_pe.shape[1]),
168 + dtype=dtype,
169 + device=k_nope.device,
170 + )
171 +
172 +
173 +def fused_mla_kv_concat(
174 + k_nope: torch.Tensor, # [T, H, qk_nope_head_dim], may be strided
175 + k_pe: torch.Tensor, # [T, rope] or [T, 1, rope], same dtype as k_nope
176 +) -> torch.Tensor:
177 + """Concat ``k = [k_nope | k_pe]`` into a contiguous key, in one launch.
178 +
179 + The chunked-context counterpart of ``fused_mla_key_concat_kv_cache_insert``:
180 + no query, no cache insert and no RoPE (the gathered ``k_pe`` is already
181 + rotated). ``k_nope`` is a strided half of one ``kv_b_proj`` output and
182 + ``k_pe`` a strided view of the gather workspace, so neither has to be made
183 + contiguous first.
184 + """
185 + k_pe = k_pe.reshape(k_pe.shape[0], k_pe.shape[-1])
186 + k = _empty_full_key(k_nope, k_pe, k_nope.dtype)
187 + if k.shape[0]:
188 + torch.ops._C.fused_kimi_k3_mla_kv_concat(k_nope, k_pe, k)
189 + return k
190 +
191 +
192 +def fused_mla_kv_concat_quant_fp8(
193 + k_nope: torch.Tensor, # [T, H, qk_nope_head_dim], may be strided
194 + k_pe: torch.Tensor, # [T, rope] or [T, 1, rope], k_nope's dtype or fp8
195 + v: torch.Tensor, # [T, H, v_head_dim], may be strided
196 +) -> tuple[torch.Tensor, torch.Tensor]:
197 + """``fused_mla_kv_concat`` plus an fp8 cast of the key and of ``v``.
198 +
199 + ``k_pe`` may already be fp8: a plain fp8 cache is gathered without
200 + dequantizing, and those bytes are copied through as-is.
201 +
202 + Returns contiguous ``(k_fp8, v_fp8)``.
203 + """
204 + k_pe = k_pe.reshape(k_pe.shape[0], k_pe.shape[-1])
205 + fp8 = torch.float8_e4m3fn
206 + k_fp8 = _empty_full_key(k_nope, k_pe, fp8)
207 + v_fp8 = torch.empty(v.shape, dtype=fp8, device=v.device)
208 + if k_fp8.shape[0]:
209 + torch.ops._C.fused_kimi_k3_mla_kv_concat_quant_fp8(
210 + k_nope, k_pe, v, k_fp8, v_fp8
211 + )
212 + return k_fp8, v_fp8
213 +
214 +
215 def fused_mla_decode_q_concat_kv_cache_insert(
216 ql_nope: torch.Tensor, # [B, H, kv_lora_rank] (BMM1 output, absorbed q)
217 q_pe: torch.Tensor, # [B, H, qk_rope_head_dim]
vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py
+5
@@ -106,7 +106,12 @@ class AiterFlashAttnPrefillBackend(MLAPrefillBackend):
106 q: torch.Tensor,
107 k: torch.Tensor,
108 v: torch.Tensor,
109 + out: torch.Tensor | None = None,
110 ) -> tuple[torch.Tensor, torch.Tensor]:
111 + assert out is None, (
112 + "AiterFlashAttnPrefillBackend does not report supports_out(), so it "
113 + "is never given a context-chunk `out` to write into."
114 + )
115 out, lse = self.flash_attn_varlen_func(
116 q=q,
117 k=k,
vllm/v1/attention/backends/mla/prefill/base.py
+9 -5
@@ -74,13 +74,16 @@ class MLAPrefillBackend(ABC):
74 return False
75
76 def supports_out(self) -> bool:
77 - """Whether `run_prefill_new_tokens` honors a caller-provided `out`
78 - tensor of shape `[num_tokens, num_heads, v_head_dim]`, writing the
79 - final result into it in place.
77 + """Whether `run_prefill_new_tokens` and `run_prefill_context_chunk` honor
78 + a caller-provided `out` tensor of shape
79 + `[num_tokens, num_heads, v_head_dim]`, writing the result into it in place
80 + and returning it.
81
82 When True, callers may pass `out` and skip the post-hoc
82 - slice/flatten/copy. False for backends that ignore `out` or emit a
83 - padded (`qk_head_dim`) output. Overridden by backends that support it.
83 + slice/flatten/copy -- and, for context chunks, size the accumulating
84 + partial before running any chunk. False for backends that ignore `out` or
85 + emit a padded (`qk_head_dim`) output. Overridden by backends that support
86 + it.
87 """
88 return False
89
@@ -179,5 +182,6 @@ class MLAPrefillBackend(ABC):
182 q: torch.Tensor,
183 k: torch.Tensor,
184 v: torch.Tensor,
185 + out: torch.Tensor | None = None,
186 ) -> tuple[torch.Tensor, torch.Tensor]:
187 raise NotImplementedError
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+2
@@ -477,6 +477,7 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
477 q: torch.Tensor,
478 k: torch.Tensor,
479 v: torch.Tensor,
480 + out: torch.Tensor | None = None,
481 ) -> tuple[torch.Tensor, torch.Tensor]:
482 return self._flash_attn_varlen_diff_headdims(
483 q=q,
@@ -489,6 +490,7 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
490 softmax_scale=self.scale,
491 causal=False, # Context is unmasked
492 return_softmax_lse=True,
493 + out=out,
494 )
495
496
vllm/v1/attention/backends/mla/prefill/flashinfer.py
+2
@@ -228,11 +228,13 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
228 q: torch.Tensor,
229 k: torch.Tensor,
230 v: torch.Tensor,
231 + out: torch.Tensor | None = None,
232 ) -> tuple[torch.Tensor, torch.Tensor]:
233 attn_out, lse = self._prefill_chunks[chunk.index].run(
234 q=q,
235 k=k,
236 v=v,
237 + out=out,
238 return_lse=True,
239 )
240
vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py
+4 -1
@@ -166,11 +166,13 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
166 q: torch.Tensor,
167 k: torch.Tensor,
168 v: torch.Tensor,
169 + out: torch.Tensor | None = None,
170 ) -> tuple[torch.Tensor, torch.Tensor]:
171 from tokenspeed_mla import tokenspeed_mla_prefill
172
173 # See note in run_prefill_new_tokens — `v` is a split-view of `kv_nope`
173 - # in `_compute_prefill_context` and arrives non-contiguous.
174 + # in the generic `_compute_prefill_context` and arrives non-contiguous.
175 + # (K3's fused fp8 K/V pack already hands over a contiguous `v`.)
176 v = v.contiguous()
177
178 attn_out, lse = tokenspeed_mla_prefill(
@@ -187,6 +189,7 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
189 cum_seq_lens_q=chunk.query_start_loc,
190 max_seq_len_q=chunk.max_query_len,
191 enable_pdl=False,
192 + out=out,
193 )
194
195 # Convert from (q_len, num_heads) to (num_heads, q_len)
vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py
+9 -7
@@ -147,16 +147,18 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
147 q: torch.Tensor,
148 k: torch.Tensor,
149 v: torch.Tensor,
150 + out: torch.Tensor | None = None,
151 ) -> tuple[torch.Tensor, torch.Tensor]:
152 from flashinfer.prefill import trtllm_ragged_attention_deepseek
153
153 - out = torch.empty(
154 - q.shape[0],
155 - q.shape[1],
156 - v.shape[2],
157 - device=q.device,
158 - dtype=self._prefill_metadata.output_dtype,
159 - )
154 + if out is None:
155 + out = torch.empty(
156 + q.shape[0],
157 + q.shape[1],
158 + v.shape[2],
159 + device=q.device,
160 + dtype=self._prefill_metadata.output_dtype,
161 + )
162
163 attn_out, lse = trtllm_ragged_attention_deepseek(
164 query=q,