From 903d2efe7eb6d045d03c0a8192b1ac7bec50d168 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Wed, 12 Aug 2026 22:42:57 -0700 Subject: [PATCH] [Attention][MLA] Fuse Kimi-K3 chunked-context K/V packing (#51772) Signed-off-by: Yongye Zhu Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- ..._kimi_k3_mla_key_concat_kv_cache_kernel.cu | 281 ++++++++++++++- csrc/libtorch_stable/ops.h | 10 + csrc/libtorch_stable/torch_bindings.cpp | 11 + .../test_kimi_k3_mla_fused_epilogue.py | 70 +++- .../kimi_k3/test_mla_prefill_context.py | 319 ++++++++++++++++++ .../v1/attention/test_mla_prefill_registry.py | 4 +- .../layers/attention/mla_attention.py | 38 ++- vllm/models/kimi_k3/nvidia/mla.py | 198 ++++++++++- .../ops/fused_mla_key_concat_kv_cache.py | 56 +++ .../backends/mla/prefill/aiter_flash_attn.py | 5 + .../v1/attention/backends/mla/prefill/base.py | 14 +- .../backends/mla/prefill/flash_attn.py | 2 + .../backends/mla/prefill/flashinfer.py | 2 + .../backends/mla/prefill/tokenspeed_mla.py | 5 +- .../backends/mla/prefill/trtllm_ragged.py | 16 +- 15 files changed, 996 insertions(+), 35 deletions(-) create mode 100644 tests/models/kimi_k3/test_mla_prefill_context.py 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 b0d52e8b81..db4805dde0 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 @@ -25,6 +25,14 @@ * matching MLA's _q_scale / _k_scale / _v_scale (the cache latent uses a * separate cache scale). Per-tensor E4M3. * + * K/V pack (fused_kimi_k3_mla_kv_concat{,_quant_fp8}): + * - k_out[t, h] = [k_nope[t, h] | k_pe[t]], cast to fp8 by the _quant_fp8 + * variant, which also writes v_fp8[t, h] = cast(v[t, h]) + * The chunked-context epilogue: no cache insert, no q, no RoPE (the cached + * k_pe is already rotated) and no scaling. `k_nope`/`v` are strided views + * of one kv_b_proj output and `k_pe` may already be fp8 (a plain fp8 cache is + * gathered without dequantizing); the outputs are contiguous. + * * fp8_ds_mla (fused_kimi_k3_mla_key_concat_ds_mla_insert): * - full key concat (bf16), and cache insert in DeepSeek's 656-byte * 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, } } +// Cast 8 source elements (one uint4 of bf16/fp16) to E4M3 with no scaling, +// using the native pairwise converters so the result is bit-identical to +// `src.to(torch.float8_e4m3fn)`. +template +__device__ __forceinline__ void copyChunk8UnitFp8(uint8_t* dst, + const scalar_t* src) { +#ifndef USE_ROCM + uint4 const input = *reinterpret_cast(src); + using Converter = vllm::_typeConvert; + auto const* input2 = + reinterpret_cast(&input); + uint2 output; + auto* output2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&output); + #pragma unroll + for (int i = 0; i < 4; ++i) { + if constexpr (std::is_same_v) { + output2[i] = __nv_cvt_bfloat16raw2_to_fp8x2( + static_cast<__nv_bfloat162_raw>(input2[i]), __NV_SATFINITE, + __NV_E4M3); + } else { + output2[i] = __nv_cvt_halfraw2_to_fp8x2( + static_cast<__half2_raw>(input2[i]), __NV_SATFINITE, __NV_E4M3); + } + } + *reinterpret_cast(dst) = output; +#else + copyChunk8(dst, src, 1.0f); +#endif +} + // Concat + store one head's full key: dst[e] = [k_nope | k_pe], e in [0, 192). // FP8 dst is byte-addressed; bf16 dst is scalar_t-addressed (dst_elem_size). template @@ -202,6 +240,38 @@ __device__ __forceinline__ void writeFullKey(void* dst, const scalar_t* k_nope, } } +// Unscaled full-key writer for the chunked-context pack, driven by a half warp +// (16 lanes x 8 elems covers NoPE 128; 8 of them cover RoPE 64). Unlike +// `writeFullKey` there is no scale and no RoPE (a gathered k_pe is already +// rotated), and KPE_FP8 handles a k_pe that the gather left in the fp8 cache +// layout -- so neither dtype needs an intermediate cast or a broadcast copy. +template +__device__ __forceinline__ void writeFullKeyPack(void* dst, + const scalar_t* k_nope, + const void* k_pe, int laneId) { + static_assert(!KPE_FP8 || FP8, "an fp8 k_pe requires an fp8 key output"); + constexpr int kElemSize = FP8 ? 1 : sizeof(scalar_t); + auto* d = reinterpret_cast(dst); + for (int e = laneId * kVecElems; e < kQkHeadDim; e += 16 * kVecElems) { + void* out = d + e * kElemSize; + const scalar_t* src = k_nope + e; + if (e >= kQkNopeHeadDim) { + int const rope_e = e - kQkNopeHeadDim; + if constexpr (KPE_FP8) { + *reinterpret_cast(out) = *reinterpret_cast( + reinterpret_cast(k_pe) + rope_e); + continue; + } + src = reinterpret_cast(k_pe) + rope_e; + } + if constexpr (FP8) { + copyChunk8UnitFp8(reinterpret_cast(out), src); + } else { + *reinterpret_cast(out) = *reinterpret_cast(src); + } + } +} + // Store a prefill query, rotating only q[..., 128:192]. For bf16 dst may alias // q (in-place); fp8 writes the quantized query directly to its output. template @@ -457,6 +527,72 @@ __global__ void fusedKimiK3MLAQKVQuantKVCacheFp8Kernel( #endif } +// ──────────────────────────────────────────────────────────────────────────── +// K/V pack variant (chunked context): concatenate the strided kv_b_proj output +// with k_pe into a contiguous row-major key, casting to E4M3 (and casting V +// into its own contiguous output) when FP8. Unlike the other kernels this one +// has no per-token cache slot, so a warp handles two (token, head) rows -- 16 +// lanes each -- and the grid is capped so a long context becomes a grid-stride +// loop instead of a huge launch. +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLAKVConcatPackKernel( + const scalar_t* __restrict__ k_nope, int64_t const kn_tok_stride, + int64_t const kn_head_stride, const void* __restrict__ k_pe, + int64_t const k_pe_tok_stride, const scalar_t* __restrict__ v, + int64_t const v_tok_stride, int64_t const v_head_stride, + void* __restrict__ k_out, int64_t const ko_tok_stride, + int64_t const ko_head_stride, uint8_t* __restrict__ v_fp8, + int64_t const vo_tok_stride, int64_t const vo_head_stride, + int const num_tokens, int const num_heads) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 16; + int const rowInWarp = (threadIdx.x % 32) / 16; + int64_t const globalWarpIdx = + static_cast(blockIdx.x) * warpsPerBlock + threadIdx.x / 32; + int64_t const totalRows = static_cast(num_tokens) * num_heads; + if (globalWarpIdx * 2 >= totalRows) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + constexpr int kOutElemSize = FP8 ? 1 : sizeof(scalar_t); + int64_t const gridStride = + static_cast(gridDim.x) * warpsPerBlock * 2; + for (int64_t row = globalWarpIdx * 2 + rowInWarp; row < totalRows; + row += gridStride) { + int const tokenIdx = static_cast(row / num_heads); + int const headIdx = static_cast(row % num_heads); + const void* pe; + if constexpr (KPE_FP8) { + pe = reinterpret_cast(k_pe) + tokenIdx * k_pe_tok_stride; + } else { + pe = reinterpret_cast(k_pe) + tokenIdx * k_pe_tok_stride; + } + writeFullKeyPack( + reinterpret_cast(k_out) + + (tokenIdx * ko_tok_stride + headIdx * ko_head_stride) * + kOutElemSize, + k_nope + tokenIdx * kn_tok_stride + headIdx * kn_head_stride, pe, + laneId); + + // bf16 V needs no cast, so it stays a strided view of the kv_b_proj output. + if constexpr (FP8) { + const scalar_t* vh = + v + tokenIdx * v_tok_stride + headIdx * v_head_stride; + uint8_t* vo = v_fp8 + tokenIdx * vo_tok_stride + headIdx * vo_head_stride; + for (int e = laneId * kVecElems; e < kVHeadDim; e += 16 * kVecElems) { + copyChunk8UnitFp8(vo + e, vh + e); + } + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + // ──────────────────────────────────────────────────────────────────────────── // ds_mla variant: concat full key (bf16) + fp8_ds_mla latent cache insert // @@ -648,16 +784,20 @@ __global__ void fusedKimiK3MLADecodeQConcatDsMlaKernel( #endif } -// PDL-aware launch of a (token, num_heads + 1)-warp grid. +// PDL-aware launch of a (token, slots_per_token)-row grid, `rows_per_warp` rows +// to a warp. `max_blocks > 0` caps the grid, leaving the kernel's grid-stride +// loop to cover the remaining rows. template -static void launchPdl(KernelT kernel, int num_tokens, int num_heads, - cudaStream_t stream, Args... args) { +static void launchPdlSlots(KernelT kernel, int num_tokens, int slots_per_token, + int rows_per_warp, int max_blocks, + cudaStream_t stream, Args... args) { constexpr int kBlockSize = 256; constexpr int kWarpsPerBlock = kBlockSize / 32; - int64_t const total_warps = - static_cast(num_tokens) * (num_heads + 1); - int const grid = + int64_t const total_rows = static_cast(num_tokens) * slots_per_token; + int64_t const total_warps = (total_rows + rows_per_warp - 1) / rows_per_warp; + int grid = static_cast((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock); + if (max_blocks > 0 && grid > max_blocks) grid = max_blocks; #ifndef USE_ROCM static int const sm_version = getSMVersion(); cudaLaunchConfig_t config; @@ -681,6 +821,13 @@ static void launchPdl(KernelT kernel, int num_tokens, int num_heads, #endif } +// Cache-inserting kernels use one extra warp per token for the cache slot. +template +static void launchPdl(KernelT kernel, int num_tokens, int num_heads, + cudaStream_t stream, Args... args) { + launchPdlSlots(kernel, num_tokens, num_heads + 1, 1, 0, stream, args...); +} + void checkBfloat16Support(torch::headeronly::ScalarType dtype) { #ifndef USE_ROCM if (dtype == torch::headeronly::ScalarType::BFloat16) { @@ -913,6 +1060,128 @@ void fused_kimi_k3_mla_key_concat_ds_mla_insert( }); } +namespace { +// Shared checks for the chunked-context K/V pack ops. `k_nope` / `v` are +// strided views of one kv_b_proj output and `k_pe` is a strided view of the +// gathered context workspace, so only a unit last-dim stride (not full +// contiguity) is required of the inputs. Returns true when `k_pe` is still in +// the fp8 cache layout. +bool check_kv_concat_inputs(torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor const& k_out, + torch::headeronly::ScalarType out_dtype) { + using torch::headeronly::ScalarType; + ScalarType const dt = k_nope.scalar_type(); + STD_TORCH_CHECK(k_nope.device().is_cuda() && k_nope.dim() == 3 && + k_nope.size(2) == 128 && k_nope.stride(2) == 1, + "k_nope must be CUDA [T, H, 128] with unit last stride"); + bool const k_pe_is_fp8 = k_pe.scalar_type() == ScalarType::Float8_e4m3fn; + STD_TORCH_CHECK( + k_pe.device() == k_nope.device() && k_pe.dim() == 2 && + k_pe.size(1) == 64 && k_pe.stride(1) == 1 && + (k_pe.scalar_type() == dt || + (k_pe_is_fp8 && out_dtype == ScalarType::Float8_e4m3fn)), + "k_pe must be CUDA [T, 64] with unit last-dim stride and use the input " + "dtype, or float8_e4m3fn when the key output is fp8"); + STD_TORCH_CHECK(k_out.device() == k_nope.device() && k_out.is_contiguous() && + k_out.scalar_type() == out_dtype && k_out.dim() == 3 && + k_out.size(2) == 192 && k_out.size(0) == k_nope.size(0) && + k_out.size(1) == k_nope.size(1) && + k_pe.size(0) == k_nope.size(0), + "k_out must be a contiguous CUDA [T, H, 192] tensor of the " + "output dtype matching k_nope's token and head dimensions"); + vllm::kimi_k3_fused_ops::checkBfloat16Support(dt); + return k_pe_is_fp8; +} +} // namespace + +void fused_kimi_k3_mla_kv_concat( + torch::stable::Tensor const& k_nope, // [T, H, 128] bf16/fp16 + torch::stable::Tensor const& k_pe, // [T, 64] same dtype as k_nope + torch::stable::Tensor& k_out) { // [T, H, 192] same dtype, written + namespace kk3 = vllm::kimi_k3_fused_ops; + torch::headeronly::ScalarType const dt = k_nope.scalar_type(); + check_kv_concat_inputs(k_nope, k_pe, k_out, dt); + + int const num_tokens = static_cast(k_nope.size(0)); + int const num_heads = static_cast(k_nope.size(1)); + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + k_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(k_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES(dt, "fused_kimi_k3_mla_kv_concat", [&] { + kk3::launchPdlSlots( + kk3::fusedKimiK3MLAKVConcatPackKernel, + num_tokens, num_heads, 2, get_device_prop()->multiProcessorCount * 8, + stream, reinterpret_cast(k_nope.const_data_ptr()), + k_nope.stride(0), k_nope.stride(1), k_pe.const_data_ptr(), + k_pe.stride(0), static_cast(nullptr), int64_t{0}, + int64_t{0}, k_out.mutable_data_ptr(), k_out.stride(0), k_out.stride(1), + static_cast(nullptr), int64_t{0}, int64_t{0}, num_tokens, + num_heads); + }); +} + +void fused_kimi_k3_mla_kv_concat_quant_fp8( + torch::stable::Tensor const& k_nope, // [T, H, 128] bf16/fp16 + torch::stable::Tensor const& k_pe, // [T, 64] bf16/fp16/fp8 + torch::stable::Tensor const& v, // [T, H, 128] bf16/fp16 + torch::stable::Tensor& k_fp8, // [T, H, 192] fp8, written + torch::stable::Tensor& v_fp8) { // [T, H, 128] fp8, written + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + ScalarType const dt = k_nope.scalar_type(); + bool const k_pe_is_fp8 = + check_kv_concat_inputs(k_nope, k_pe, k_fp8, ScalarType::Float8_e4m3fn); + STD_TORCH_CHECK(v.device() == k_nope.device() && v.scalar_type() == dt && + v.dim() == 3 && v.size(2) == 128 && v.stride(2) == 1 && + v.size(0) == k_nope.size(0) && + v.size(1) == k_nope.size(1), + "v must match k_nope's dtype and token/head dimensions and " + "have shape [T, H, 128] with unit last stride"); + STD_TORCH_CHECK(v_fp8.device() == k_nope.device() && v_fp8.is_contiguous() && + v_fp8.scalar_type() == ScalarType::Float8_e4m3fn && + v_fp8.dim() == 3 && v_fp8.size(2) == 128 && + v_fp8.size(0) == k_nope.size(0) && + v_fp8.size(1) == k_nope.size(1), + "v_fp8 must be a contiguous CUDA [T, H, 128] fp8 tensor " + "matching k_nope's token and head dimensions"); + + int const num_tokens = static_cast(k_nope.size(0)); + int const num_heads = static_cast(k_nope.size(1)); + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + k_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(k_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_kv_concat_quant_fp8", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdlSlots( + kernel, num_tokens, num_heads, 2, + get_device_prop()->multiProcessorCount * 8, stream, + reinterpret_cast(k_nope.const_data_ptr()), + k_nope.stride(0), k_nope.stride(1), k_pe.const_data_ptr(), + k_pe.stride(0), + reinterpret_cast(v.const_data_ptr()), + v.stride(0), v.stride(1), k_fp8.mutable_data_ptr(), + k_fp8.stride(0), k_fp8.stride(1), + reinterpret_cast(v_fp8.mutable_data_ptr()), + v_fp8.stride(0), v_fp8.stride(1), num_tokens, num_heads); + }; + if (k_pe_is_fp8) { + launch(kk3::fusedKimiK3MLAKVConcatPackKernel); + } else { + launch(kk3::fusedKimiK3MLAKVConcatPackKernel); + } + }); +} + void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( torch::stable::Tensor const& q, // [Tp, H, 192] bf16 torch::stable::Tensor const& k_nope, // [Tp, H, 128] bf16 diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 7dfe0da78c..85c875402c 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -300,6 +300,16 @@ void fused_kimi_k3_mla_key_concat_ds_mla_insert( std::optional position_ids, std::optional cos_sin_cache); +void fused_kimi_k3_mla_kv_concat(torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor& k_out); + +void fused_kimi_k3_mla_kv_concat_quant_fp8(torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor const& v, + torch::stable::Tensor& k_fp8, + torch::stable::Tensor& v_fp8); + void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope, torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index fd4e8a848c..801aed03d9 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -466,6 +466,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " "int cache_block_size, Tensor? position_ids=None, " "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_kv_concat(Tensor k_nope, Tensor k_pe, Tensor! k_out) " + "-> ()"); + ops.def( + "fused_kimi_k3_mla_kv_concat_quant_fp8(" + "Tensor k_nope, Tensor k_pe, Tensor v, Tensor! k_fp8, Tensor! v_fp8) " + "-> ()"); ops.def( "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(" "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) { TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert)); ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert", TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert)); + ops.impl("fused_kimi_k3_mla_kv_concat", + TORCH_BOX(&fused_kimi_k3_mla_kv_concat)); + ops.impl("fused_kimi_k3_mla_kv_concat_quant_fp8", + TORCH_BOX(&fused_kimi_k3_mla_kv_concat_quant_fp8)); ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert)); ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", 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 0e66e0da1e..d7cb7ca7eb 100644 --- a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py +++ b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py @@ -9,6 +9,8 @@ from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import ( fused_mla_decode_q_concat_kv_cache_insert, fused_mla_key_concat_ds_mla_insert, fused_mla_key_concat_kv_cache_insert, + fused_mla_kv_concat, + fused_mla_kv_concat_quant_fp8, fused_mla_qkv_quant_kv_cache_fp8_insert, ) from vllm.platforms import current_platform @@ -25,8 +27,8 @@ _POSITIONS = (1, 7, 13) _SLOTS = (0, 3, 9) -def _randn(*shape: int) -> torch.Tensor: - return torch.randn(*shape, device="cuda", dtype=_DTYPE) * 0.2 +def _randn(*shape: int, dtype: torch.dtype = _DTYPE) -> torch.Tensor: + return torch.randn(*shape, device="cuda", dtype=dtype) * 0.2 def _rope_cache(max_position: int = 32) -> torch.Tensor: @@ -66,6 +68,70 @@ def _assert_fp8_close(actual: torch.Tensor, expected: torch.Tensor) -> None: ) +def _strided_context_inputs( + num_tokens: int, num_heads: int, dtype: torch.dtype, k_pe_fp8: bool +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The real prefill-context layouts: ``k_nope``/``v`` as the two strided + halves of one kv_b_proj output and ``k_pe`` as a column slice of the gather + workspace (left in the fp8 cache layout for a plain fp8 cache).""" + kv_nope = _randn(num_tokens, num_heads, 256, dtype=dtype) + k_nope, v = kv_nope.split((128, 128), dim=-1) + workspace = _randn(num_tokens, 576, dtype=dtype) + if k_pe_fp8: + workspace = workspace.to(torch.float8_e4m3fn) + return k_nope, workspace[:, 512:].unsqueeze(1), v + + +@pytest.mark.parametrize("num_tokens", [0, _NUM_TOKENS]) +@pytest.mark.parametrize("num_heads", [3, _NUM_HEADS]) +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16]) +@torch.inference_mode() +def test_context_kv_concat_accepts_strided_inputs( + num_tokens: int, num_heads: int, input_dtype: torch.dtype +) -> None: + k_nope, k_pe, _ = _strided_context_inputs( + num_tokens, num_heads, input_dtype, k_pe_fp8=False + ) + + k = fused_mla_kv_concat(k_nope, k_pe) + + assert k.is_contiguous() + assert k.dtype == input_dtype + torch.testing.assert_close(k[..., :128], k_nope, atol=0, rtol=0) + torch.testing.assert_close( + k[..., 128:], k_pe.expand(-1, num_heads, -1), atol=0, rtol=0 + ) + + +@pytest.mark.parametrize("k_pe_dtype", ["input", "fp8"]) +@pytest.mark.parametrize("num_tokens", [0, _NUM_TOKENS]) +@pytest.mark.parametrize("num_heads", [3, _NUM_HEADS]) +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16]) +@torch.inference_mode() +def test_context_kv_pack_quantizes_strided_inputs( + k_pe_dtype: str, + num_tokens: int, + num_heads: int, + input_dtype: torch.dtype, +) -> None: + k_nope, k_pe, v = _strided_context_inputs( + num_tokens, num_heads, input_dtype, k_pe_fp8=k_pe_dtype == "fp8" + ) + + k_actual, v_actual = fused_mla_kv_concat_quant_fp8(k_nope, k_pe, v) + + fp8 = torch.float8_e4m3fn + k_expected = torch.empty_like(k_actual) + k_expected[..., :128] = k_nope.to(fp8) + k_expected[..., 128:] = k_pe.to(fp8) + assert k_actual.is_contiguous() + assert v_actual.is_contiguous() + # The pack casts with the native pairwise converters, so it must be + # bit-identical to torch's `.to(fp8)` rather than merely close. + torch.testing.assert_close(k_actual.float(), k_expected.float(), atol=0, rtol=0) + torch.testing.assert_close(v_actual.float(), v.to(fp8).float(), atol=0, rtol=0) + + @pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"]) @torch.inference_mode() def test_prefill_epilogue_fuses_gptj_rope(cache_kind: str) -> None: diff --git a/tests/models/kimi_k3/test_mla_prefill_context.py b/tests/models/kimi_k3/test_mla_prefill_context.py new file mode 100644 index 0000000000..6c6e4bf20e --- /dev/null +++ b/tests/models/kimi_k3/test_mla_prefill_context.py @@ -0,0 +1,319 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K3's fused chunked-context prefill must match the generic MLA impl. + +The layer owns its own context loop so it can fuse the per-chunk K/V pack and +skip re-quantizing an already-quantized query. That is only safe if it feeds the +prefill backend exactly what ``MLACommonBaseImpl._compute_prefill_context`` +would, chunk for chunk. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonBaseImpl, + MLACommonPrefillMetadata, + build_mla_chunked_context_metadata, +) +from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="Kimi-K3 fused MLA requires CUDA" +) + +_KV_LORA_RANK = 512 +_QK_NOPE = 128 +_QK_ROPE = 64 +_V_HEAD_DIM = 128 +_ENTRY = _KV_LORA_RANK + _QK_ROPE +_NUM_HEADS = 2 +_BLOCK_SIZE = 16 +_WORKSPACE_TOKENS = 128 +# Splits one long request across chunks, packs short ones together, and leaves +# the last request without any context. +_CONTEXT_LENS = [200, 48, 32, 0] +_QUERY_LENS = [8, 4, 6, 5] + + +class _RecordingPrefillBackend: + """Records what each chunk is asked to attend over. + + ``honors_out`` mimics the two backend families: one that writes into a + caller-provided ``out`` (trtllm_ragged, flashinfer, tokenspeed) and one that + always returns its own buffer (flash_attn with a padded V, aiter). + """ + + def __init__(self, honors_out: bool = False) -> None: + self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + self.out_destinations: list[torch.Tensor | None] = [] + self._honors_out = honors_out + + @staticmethod + def get_name() -> str: + return "recording" + + def supports_out(self) -> bool: + return self._honors_out + + def run_prefill_context_chunk(self, *, chunk, q, k, v, out=None): + self.calls.append((q.float().clone(), k.float().clone(), v.float().clone())) + self.out_destinations.append(out) + assert out is None or self._honors_out + # Fold K/V into the partial so any packing difference shows up in the + # merged context output, not just in the recorded calls. + digest = (k.float().mean() + v.float().mean()).item() + num_q = q.shape[0] + if out is None: + out = torch.empty( + (num_q, _NUM_HEADS, _V_HEAD_DIM), device=q.device, dtype=torch.bfloat16 + ) + else: + assert out.shape == (num_q, _NUM_HEADS, _V_HEAD_DIM) + out.fill_(digest) + lse = torch.full( + (_NUM_HEADS, num_q), + 1.0 + chunk.index, + device=q.device, + dtype=torch.float32, + ) + return out, lse + + +class _KVBProj(torch.nn.Module): + """Stand-in for the layer's ``kv_b_proj`` (returns an (out, bias) tuple). + + Enforces the same input contract as the real linear methods, because that is + what decides whether the gathered latent needs a cast: + + * an fp8 weight consumes the fp8 latent directly and dequantizes internally, + so it takes fp8 or bf16; + * a bf16 weight -- what a stock K3 checkpoint carries -- is a plain + ``F.linear`` and rejects anything but bf16, exactly as torch does. + + Either weight dtype produces a bf16 output. + """ + + def __init__(self, device: torch.device, weight_dtype: torch.dtype) -> None: + super().__init__() + weight = ( + torch.randn( + _NUM_HEADS * (_QK_NOPE + _V_HEAD_DIM), + _KV_LORA_RANK, + device=device, + dtype=torch.bfloat16, + ) + * 0.05 + ) + self.register_buffer("weight", weight.to(weight_dtype)) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]: + if self.weight.dtype == torch.bfloat16 and x.dtype != torch.bfloat16: + raise RuntimeError( + "a bfloat16 kv_b_proj cannot consume the gathered latent as " + f"{x.dtype}; it must be cast first" + ) + return torch.nn.functional.linear( + x.to(torch.bfloat16), self.weight.to(torch.bfloat16) + ), None + + +class _FusedLayer: + """Only the attributes K3's context loop reads.""" + + _compute_prefill_context = MultiHeadLatentAttention._compute_prefill_context + _gather_context_latent = MultiHeadLatentAttention._gather_context_latent + _attn_read_kv_cache = MultiHeadLatentAttention._attn_read_kv_cache + + def __init__(self, kv_b_proj, kv_cache, kv_cache_dtype, k_scale) -> None: + self.kv_b_proj = kv_b_proj + self.kv_cache = kv_cache + self.kv_cache_dtype = kv_cache_dtype + self._k_scale = k_scale + self.kv_lora_rank = _KV_LORA_RANK + self.num_local_heads = _NUM_HEADS + self.qk_nope_head_dim = _QK_NOPE + self.v_head_dim = _V_HEAD_DIM + + +class _ReferenceImpl: + """Only the attributes the generic context loop reads.""" + + _compute_prefill_context = MLACommonBaseImpl._compute_prefill_context + _concat_k_nope_k_pe = MLACommonBaseImpl._concat_k_nope_k_pe + _use_flashinfer_concat_mla_k = False + + def __init__(self, kv_b_proj, kv_cache_dtype) -> None: + self.kv_b_proj = kv_b_proj + self.kv_cache_dtype = kv_cache_dtype + self.kv_lora_rank = _KV_LORA_RANK + self.num_heads = _NUM_HEADS + self.qk_nope_head_dim = _QK_NOPE + self.qk_rope_head_dim = _QK_ROPE + self.v_head_dim = _V_HEAD_DIM + + +def _build_prefill_metadata( + device: torch.device, + workspace_dtype: torch.dtype, + q_data_type: torch.dtype, + backend: _RecordingPrefillBackend, +) -> MLACommonPrefillMetadata: + query_start_loc_cpu = torch.zeros(len(_QUERY_LENS) + 1, dtype=torch.int32) + query_start_loc_cpu[1:] = torch.tensor(_QUERY_LENS, dtype=torch.int32).cumsum(0) + workspace = torch.empty( + (_WORKSPACE_TOKENS, _ENTRY), dtype=workspace_dtype, device=device + ) + chunked_context = build_mla_chunked_context_metadata( + context_lens_cpu=torch.tensor(_CONTEXT_LENS, dtype=torch.int32), + prefill_query_start_loc_cpu=query_start_loc_cpu, + chunked_prefill_workspace=workspace, + chunked_prefill_workspace_size=_WORKSPACE_TOKENS, + block_size=_BLOCK_SIZE, + align_chunk_to_block=True, + device=device, + dcp_world_size=1, + dcp_local_block_size=1, + dcp_virtual_block_size=1, + ) + assert chunked_context is not None + assert len(chunked_context.chunks) > 1, "the batch must exercise accumulation" + + max_blocks = (max(_CONTEXT_LENS) + max(_QUERY_LENS)) // _BLOCK_SIZE + 1 + num_blocks = max_blocks * len(_CONTEXT_LENS) + block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).view( + len(_CONTEXT_LENS), max_blocks + ) + return MLACommonPrefillMetadata( + block_table=block_table, + query_start_loc=query_start_loc_cpu.to(device), + max_query_len=max(_QUERY_LENS), + chunked_context=chunked_context, + q_data_type=q_data_type, + output_dtype=torch.bfloat16, + prefill_backend=backend, + ), num_blocks + + +@pytest.mark.parametrize("honors_out", [False, True], ids=["copy_out", "writes_out"]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize( + "kv_b_proj_quantized", [True, False], ids=["fp8_kv_b_proj", "bf16_kv_b_proj"] +) +@torch.inference_mode() +def test_fused_context_matches_generic_impl( + kv_b_proj_quantized: bool, kv_cache_dtype: str, honors_out: bool +) -> None: + """Cache dtype and ``kv_b_proj`` dtype vary independently. + + A stock K3 checkpoint pairs a bf16 ``kv_b_proj`` with an fp8 cache, so the + fused loop cannot assume the gathered latent is already in the dtype + ``kv_b_proj`` accepts. + """ + torch.manual_seed(0) + device = torch.device("cuda") + fp8 = current_platform.fp8_dtype() + quantized = kv_cache_dtype == "fp8" + # An fp8 cache is read by an fp8 query, so the workspace keeps the fp8 + # layout; a bf16 cache dequantizes into a bf16 workspace. + q_data_type = fp8 if quantized else torch.bfloat16 + workspace_dtype = q_data_type + + kv_b_proj = _KVBProj( + device, weight_dtype=fp8 if kv_b_proj_quantized else torch.bfloat16 + ) + k_scale = torch.ones(1, dtype=torch.float32, device=device) + + backend_fused = _RecordingPrefillBackend(honors_out=honors_out) + # The reference impl never passes `out`, so it always allocates its own. + backend_ref = _RecordingPrefillBackend() + prefill_fused, num_blocks = _build_prefill_metadata( + device, workspace_dtype, q_data_type, backend_fused + ) + prefill_ref, _ = _build_prefill_metadata( + device, workspace_dtype, q_data_type, backend_ref + ) + + cache = torch.randn( + (num_blocks, _BLOCK_SIZE, _ENTRY), device=device, dtype=torch.bfloat16 + ) + kv_cache = cache.to(fp8) if quantized else cache + q = ( + torch.randn( + (sum(_QUERY_LENS), _NUM_HEADS, _QK_NOPE + _QK_ROPE), + device=device, + dtype=torch.bfloat16, + ) + * 0.2 + ).to(q_data_type) + + layer = _FusedLayer(kv_b_proj, kv_cache, kv_cache_dtype, k_scale) + fused_out, fused_lse = layer._compute_prefill_context( + q, SimpleNamespace(prefill=prefill_fused) + ) + + impl = _ReferenceImpl(kv_b_proj, kv_cache_dtype) + ref_out, ref_lse = impl._compute_prefill_context( + q, kv_cache, SimpleNamespace(prefill=prefill_ref), k_scale + ) + + assert len(backend_fused.calls) == len(backend_ref.calls) + for chunk_idx, (fused_call, ref_call) in enumerate( + zip(backend_fused.calls, backend_ref.calls, strict=True) + ): + for name, fused_t, ref_t in zip( + ("q", "k", "v"), fused_call, ref_call, strict=True + ): + torch.testing.assert_close( + fused_t, + ref_t, + atol=0, + rtol=0, + msg=lambda m, n=name, i=chunk_idx: f"chunk {i} {n} differs: {m}", + ) + torch.testing.assert_close(fused_out, ref_out, atol=0, rtol=0) + torch.testing.assert_close(fused_lse, ref_lse, atol=0, rtol=0) + + # Every chunk but the continuation should have been written in place, i.e. + # straight into the returned accumulator, with no intermediate copy. + wrote_in_place = [ + out is not None and out.data_ptr() == fused_out[chunk.token_slice].data_ptr() + for out, chunk in zip( + backend_fused.out_destinations, + prefill_fused.chunked_context.chunks, + strict=True, + ) + ] + continuations = [c.is_continuation for c in prefill_fused.chunked_context.chunks] + assert any(continuations), "the batch must exercise a continuation chunk" + if honors_out: + assert wrote_in_place == [not c for c in continuations] + else: + assert not any(wrote_in_place) + + +@torch.inference_mode() +def test_fused_context_rejects_an_unquantized_query() -> None: + """The fp8 query is produced by the new-token epilogue, not re-cast here.""" + torch.manual_seed(0) + device = torch.device("cuda") + fp8 = current_platform.fp8_dtype() + prefill, num_blocks = _build_prefill_metadata( + device, fp8, fp8, _RecordingPrefillBackend() + ) + layer = _FusedLayer( + _KVBProj(device, weight_dtype=fp8), + torch.zeros((num_blocks, _BLOCK_SIZE, _ENTRY), device=device, dtype=fp8), + "fp8", + torch.ones(1, dtype=torch.float32, device=device), + ) + q = torch.zeros( + (sum(_QUERY_LENS), _NUM_HEADS, _QK_NOPE + _QK_ROPE), + device=device, + dtype=torch.bfloat16, + ) + with pytest.raises(AssertionError, match="new-token epilogue"): + layer._compute_prefill_context(q, SimpleNamespace(prefill=prefill)) diff --git a/tests/v1/attention/test_mla_prefill_registry.py b/tests/v1/attention/test_mla_prefill_registry.py index 05dcb77838..fc5efe3a5b 100644 --- a/tests/v1/attention/test_mla_prefill_registry.py +++ b/tests/v1/attention/test_mla_prefill_registry.py @@ -24,7 +24,7 @@ class CustomMLAPrefillBackend(MLAPrefillBackend): def run_prefill_new_tokens(self, q, k, v, return_softmax_lse): raise NotImplementedError - def run_prefill_context_chunk(self, chunk, q, k, v): + def run_prefill_context_chunk(self, chunk, q, k, v, out=None): raise NotImplementedError @@ -112,7 +112,7 @@ def test_register_custom_backend_as_decorator(): def run_prefill_new_tokens(self, q, k, v, return_softmax_lse): raise NotImplementedError - def run_prefill_context_chunk(self, chunk, q, k, v): + def run_prefill_context_chunk(self, chunk, q, k, v, out=None): raise NotImplementedError assert MLAPrefillBackendEnum.CUSTOM.is_overridden() diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 33bcbf8aaf..c57bd0ea31 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -2414,6 +2414,22 @@ def reorg_kvcache( return reorganized_kv_c_normed, reorganized_k_pe +def neutralize_empty_context_partials( + chunked_context: "MLACommonPrefillMetadata.ChunkedContextMetadata", + output: torch.Tensor, + output_lse: torch.Tensor, +) -> None: + """Neutralize the partial of every prefill that no chunk covers. + + A prefill without context is never gathered, so nothing would write its rows; + a zero output with an ``-inf`` lse carries no weight into the final merge + against the suffix partial. + """ + for token_slice in chunked_context.empty_token_slices: + output[token_slice].zero_() + output_lse[:, token_slice].fill_(float("-inf")) + + def init_mla_context_partial( chunked_context: "MLACommonPrefillMetadata.ChunkedContextMetadata", attn_output: torch.Tensor, @@ -2423,7 +2439,9 @@ def init_mla_context_partial( """Allocate the running context partial over all prefill tokens. Laid out like the chunk partials so the final whole-batch merge against the - suffix partial sees matching head strides. + suffix partial sees matching head strides. Callers whose backend honors an + ``out`` tensor already know that layout and can allocate directly, pairing it + with ``neutralize_empty_context_partials``. """ output = torch.empty( (num_tokens, *attn_output.shape[1:]), @@ -2435,10 +2453,7 @@ def init_mla_context_partial( dtype=attn_softmax_lse.dtype, device=attn_softmax_lse.device, ) - # No chunk covers a prefill without context, so neutralize its partial. - for token_slice in chunked_context.empty_token_slices: - output[token_slice].zero_() - output_lse[:, token_slice].fill_(float("-inf")) + neutralize_empty_context_partials(chunked_context, output, output_lse) return output, output_lse @@ -2448,16 +2463,26 @@ def accumulate_mla_context_chunk( attn_softmax_lse: torch.Tensor, output: torch.Tensor, output_lse: torch.Tensor, + output_written: bool = False, ) -> None: """Fold one chunk's partial into the running context partial. Only the first request may be a continuation; its tokens are merged and the remaining token range is initialized. + + Args: + output_written: The chunk's attention output already landed in + ``output[chunk.token_slice]`` because the backend was handed it as + ``out``, leaving only the lse to fold. Invalid for a continuation + chunk, whose leading tokens must be merged rather than overwritten. """ token_start = chunk.token_slice.start token_end = chunk.token_slice.stop init_start = token_start if chunk.is_continuation: + assert not output_written, ( + "a continuation chunk must not write over the partial it merges with" + ) init_start = chunk.continuation_token_end num_merged = init_start - token_start merge_attn_states( @@ -2470,7 +2495,8 @@ def accumulate_mla_context_chunk( ) if init_start < token_end: written = init_start - token_start - output[init_start:token_end].copy_(attn_output[written:]) + if not output_written: + output[init_start:token_end].copy_(attn_output[written:]) output_lse[:, init_start:token_end].copy_(attn_softmax_lse[:, written:]) diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 8f70574142..708104b073 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -8,8 +8,9 @@ This is a self-contained MLA layer that owns the full attention path: -> fused pre-attention ops (fused_qkv_a_proj / norms / q_b_proj) -> explicit prefill / decode split prefill: fused key-concat + cache-insert kernel -> run_prefill_new_tokens - (+ chunked-context merge); dispatched by cache dtype - (bf16 / plain fp8 / fp8_ds_mla) + (+ chunked-context merge, whose per-chunk gather -> kv_b_proj + -> fused K/V pack loop this layer owns); dispatched by cache + dtype (bf16 / plain fp8 / fp8_ds_mla) decode : W_UK absorb (BMM1) -> fused q-concat + cache-insert kernel -> impl.forward_mqa -> W_UV up-proj (MQA) -> optional output gate @@ -34,6 +35,7 @@ from typing import TYPE_CHECKING, cast import torch from torch import nn +from vllm import _custom_ops as ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import ( CacheConfig, @@ -48,6 +50,12 @@ from vllm.model_executor.layers.attention.attention import ( set_default_quant_scales, should_load_quant_weights, ) +from vllm.model_executor.layers.attention.mla_attention import ( + _get_kv_b_proj_input_dtype, + accumulate_mla_context_chunk, + init_mla_context_partial, + neutralize_empty_context_partials, +) from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm 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 ( fused_mla_decode_q_concat_kv_cache_insert, fused_mla_key_concat_ds_mla_insert, fused_mla_key_concat_kv_cache_insert, + fused_mla_kv_concat, + fused_mla_kv_concat_quant_fp8, fused_mla_qkv_quant_kv_cache_fp8_insert, ) from vllm.platforms import current_platform @@ -696,6 +706,181 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): cos_sin_cache=cos_sin_cache, ) + def _compute_prefill_context( + self, + q: torch.Tensor, + attn_metadata: "MLACommonMetadata", + ) -> tuple[torch.Tensor, torch.Tensor]: + """Chunked-context prefill, K3-fused. Replaces the impl's version. + + Per chunk the impl gathers the paged latent, up-projects it, then casts + and concatenates K (and casts V) in two or three more launches. Here that + tail is one fused kernel per chunk -- ``fused_mla_kv_concat`` for a bf16 + query, ``fused_mla_kv_concat_quant_fp8`` when the query is fp8 -- reading + the strided ``kv_b_proj`` output in place and writing a contiguous key, so + only the gather and ``kv_b_proj`` remain. + + The impl's query cast is gone as well: ``q`` already carries + ``prefill.q_data_type`` because the new-token epilogue quantized it. The + gathered latent still gets the impl's cast to whatever ``kv_b_proj`` + consumes -- free (a no-op ``.to``) for a checkpoint whose ``kv_b_proj`` + takes the fp8 latent directly, and required for a bf16 one, which is + what a stock K3 checkpoint carries. Its output is bf16 either way. + + The gathered ``k_pe`` is likewise used as-is (fp8 for a plain fp8 cache) + and needs no RoPE: it was rotated on the way in. + + Chunk partials are written straight into the accumulating context partial + when the prefill backend honors ``out``, so only the (64x smaller) lse is + copied per chunk. + + Decode context parallelism keeps using + ``impl._context_parallel_compute_prefill_context``; its extra allgather + and reorg are not fused here. + """ + prefill = attn_metadata.prefill + assert prefill is not None + prefill_backend = prefill.prefill_backend + assert prefill_backend is not None + chunked_context = prefill.chunked_context + assert chunked_context is not None + assert q.dtype == prefill.q_data_type, ( + "Kimi-K3 chunked context expects the new-token epilogue to have " + f"produced a {prefill.q_data_type} query; got {q.dtype}." + ) + + fp8_prefill = q.dtype == current_platform.fp8_dtype() + workspace = chunked_context.workspace + kv_cache = self._attn_read_kv_cache() + kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(self.kv_b_proj, fp8_prefill) + + def run_chunk( + chunk, out: torch.Tensor | None = None + ) -> tuple[torch.Tensor, torch.Tensor]: + self._gather_context_latent(chunk, kv_cache, prefill, fp8_prefill) + gathered = workspace[: chunk.num_context_tokens] + kv_c_normed = gathered[..., : self.kv_lora_rank] + if kv_b_proj_input_dtype is not None: + kv_c_normed = kv_c_normed.to(kv_b_proj_input_dtype) + kv_nope = self.kv_b_proj(kv_c_normed)[0].view( + -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim + ) + k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) + k_pe = gathered[..., self.kv_lora_rank :] + if fp8_prefill: + k, v = fused_mla_kv_concat_quant_fp8(k_nope, k_pe, v) + else: + k = fused_mla_kv_concat(k_nope, k_pe) + attn_output, attn_lse = prefill_backend.run_prefill_context_chunk( + chunk=chunk, q=q[chunk.token_slice], k=k, v=v, out=out + ) + assert out is None or attn_output.data_ptr() == out.data_ptr(), ( + f"{prefill_backend.get_name()} reports supports_out() but did not " + "write the context chunk into the `out` it was given." + ) + return attn_output, attn_lse + + chunks = chunked_context.chunks + if len(chunks) == 1 and not chunked_context.empty_token_slices: + # One chunk covering every prefill token: its partial *is* the context + # partial, so it needs neither an accumulator nor a copy. + return run_chunk(chunks[0]) + + # A backend honoring `out` writes each chunk's partial straight into the + # accumulator, so the per-chunk output copy disappears -- and because that + # contract fixes the trailing shape, the accumulator can be sized before + # any chunk runs. Otherwise the shape is only knowable from a real partial, + # so the first chunk runs ahead of the loop and is copied in. + writes_out = prefill_backend.supports_out() + if writes_out: + assert prefill.output_dtype is not None + output = torch.empty( + (q.shape[0], self.num_local_heads, self.v_head_dim), + dtype=prefill.output_dtype, + device=q.device, + ) + output_lse = torch.empty( + (self.num_local_heads, q.shape[0]), + dtype=torch.float32, + device=q.device, + ) + neutralize_empty_context_partials(chunked_context, output, output_lse) + else: + attn_output, attn_lse = run_chunk(chunks[0]) + output, output_lse = init_mla_context_partial( + chunked_context, attn_output, attn_lse, num_tokens=q.shape[0] + ) + accumulate_mla_context_chunk( + chunks[0], attn_output, attn_lse, output, output_lse + ) + chunks = chunks[1:] + + for chunk in chunks: + # A continuation chunk's leading tokens have to be merged with the + # partial already sitting there, so it cannot write in place. + out = ( + output[chunk.token_slice] + if writes_out and not chunk.is_continuation + else None + ) + attn_output, attn_lse = run_chunk(chunk, out=out) + accumulate_mla_context_chunk( + chunk, + attn_output, + attn_lse, + output, + output_lse, + output_written=out is not None, + ) + return output, output_lse + + def _gather_context_latent( + self, + chunk, + kv_cache: torch.Tensor, + prefill, + fp8_prefill: bool, + ) -> None: + """Gather one chunk's paged context latent into the workspace. + + Dispatched exactly as in ``impl._compute_prefill_context``: an fp8 query + reads the plain fp8 cache in its stored layout, anything else lands in the + workspace as the model dtype. + """ + workspace = prefill.chunked_context.workspace + toks = chunk.num_context_tokens + block_table = prefill.block_table[chunk.request_slice] + if self.kv_cache_dtype == "fp8_ds_mla": + ops.cp_gather_and_upconvert_fp8_kv_cache( + src_cache=kv_cache, + dst=workspace[:toks], + block_table=block_table, + workspace_starts=chunk.cu_seq_lens, + batch_size=chunk.num_requests, + seq_starts=chunk.starts, + ) + elif not fp8_prefill: + ops.gather_and_maybe_dequant_cache( + src_cache=kv_cache, + dst=workspace, + block_table=block_table, + cu_seq_lens=chunk.cu_seq_lens, + token_to_seq=chunk.token_to_seq, + num_tokens=toks, + kv_cache_dtype=self.kv_cache_dtype, + scale=self._k_scale, + seq_starts=chunk.starts, + ) + else: + ops.cp_gather_cache( + src_cache=kv_cache, + dst=workspace[:toks], + block_table=block_table, + cu_seq_lens=chunk.cu_seq_lens, + batch_size=chunk.num_requests, + seq_starts=chunk.starts, + ) + def _forward_prefill_fused( self, q: torch.Tensor, @@ -710,8 +895,9 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): """Prefill using the fused key-concat + cache-insert kernel. Replaces ``_concat_k_nope_k_pe`` and the prefill cache write with one - fused kernel launch, dispatched by cache dtype. The chunked context - gather + online-softmax merge are delegated to the impl. + fused kernel launch, dispatched by cache dtype. Chunked context runs + through this layer's ``_compute_prefill_context``, except under DCP where + it is delegated to the impl. Supported configs (K3 fp8 policy): - bf16 cache -> bf16 prefill query @@ -812,8 +998,8 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): ) ) else: - context_output, context_lse = self.impl._compute_prefill_context( # type: ignore[attr-defined] - q, self._attn_read_kv_cache(), attn_metadata, self._k_scale + context_output, context_lse = self._compute_prefill_context( + q, attn_metadata ) suffix_output, suffix_lse = output_prefill out = out.view(-1, self.num_local_heads, self.v_head_dim) diff --git a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py index 816da37153..24ae1a1a82 100644 --- a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py +++ b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py @@ -12,6 +12,9 @@ mirror ``fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_{bf16,fp8}_insert``. - ``fused_mla_qkv_quant_kv_cache_fp8_insert`` (fp8): additionally quantize ``q``/``k``/``v`` to E4M3 with ``q_scale`` / ``k_scale`` / ``v_scale`` (the cache shares ``k_scale``, as in ``concat_and_cache_mla``). +- ``fused_mla_kv_concat`` / ``fused_mla_kv_concat_quant_fp8`` (chunked context): + the same K concat (plus the fp8 K/V cast) without the cache insert, for context + chunks whose latent was gathered back out of the paged cache. The optional ``positions`` / ``cos_sin_cache`` pair enables GPT-J-style RoPE 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( return q_fp8, k_fp8, v_fp8 +def _empty_full_key( + k_nope: torch.Tensor, k_pe: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: + num_tokens, num_heads, qk_nope_head_dim = k_nope.shape + return torch.empty( + (num_tokens, num_heads, qk_nope_head_dim + k_pe.shape[1]), + dtype=dtype, + device=k_nope.device, + ) + + +def fused_mla_kv_concat( + k_nope: torch.Tensor, # [T, H, qk_nope_head_dim], may be strided + k_pe: torch.Tensor, # [T, rope] or [T, 1, rope], same dtype as k_nope +) -> torch.Tensor: + """Concat ``k = [k_nope | k_pe]`` into a contiguous key, in one launch. + + The chunked-context counterpart of ``fused_mla_key_concat_kv_cache_insert``: + no query, no cache insert and no RoPE (the gathered ``k_pe`` is already + rotated). ``k_nope`` is a strided half of one ``kv_b_proj`` output and + ``k_pe`` a strided view of the gather workspace, so neither has to be made + contiguous first. + """ + k_pe = k_pe.reshape(k_pe.shape[0], k_pe.shape[-1]) + k = _empty_full_key(k_nope, k_pe, k_nope.dtype) + if k.shape[0]: + torch.ops._C.fused_kimi_k3_mla_kv_concat(k_nope, k_pe, k) + return k + + +def fused_mla_kv_concat_quant_fp8( + k_nope: torch.Tensor, # [T, H, qk_nope_head_dim], may be strided + k_pe: torch.Tensor, # [T, rope] or [T, 1, rope], k_nope's dtype or fp8 + v: torch.Tensor, # [T, H, v_head_dim], may be strided +) -> tuple[torch.Tensor, torch.Tensor]: + """``fused_mla_kv_concat`` plus an fp8 cast of the key and of ``v``. + + ``k_pe`` may already be fp8: a plain fp8 cache is gathered without + dequantizing, and those bytes are copied through as-is. + + Returns contiguous ``(k_fp8, v_fp8)``. + """ + k_pe = k_pe.reshape(k_pe.shape[0], k_pe.shape[-1]) + fp8 = torch.float8_e4m3fn + k_fp8 = _empty_full_key(k_nope, k_pe, fp8) + v_fp8 = torch.empty(v.shape, dtype=fp8, device=v.device) + if k_fp8.shape[0]: + torch.ops._C.fused_kimi_k3_mla_kv_concat_quant_fp8( + k_nope, k_pe, v, k_fp8, v_fp8 + ) + return k_fp8, v_fp8 + + def fused_mla_decode_q_concat_kv_cache_insert( ql_nope: torch.Tensor, # [B, H, kv_lora_rank] (BMM1 output, absorbed q) q_pe: torch.Tensor, # [B, H, qk_rope_head_dim] diff --git a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py index 09726ec8cc..3ba5a5a5fa 100644 --- a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py @@ -106,7 +106,12 @@ class AiterFlashAttnPrefillBackend(MLAPrefillBackend): q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: + assert out is None, ( + "AiterFlashAttnPrefillBackend does not report supports_out(), so it " + "is never given a context-chunk `out` to write into." + ) out, lse = self.flash_attn_varlen_func( q=q, k=k, diff --git a/vllm/v1/attention/backends/mla/prefill/base.py b/vllm/v1/attention/backends/mla/prefill/base.py index 986d0ca347..1b589cb295 100644 --- a/vllm/v1/attention/backends/mla/prefill/base.py +++ b/vllm/v1/attention/backends/mla/prefill/base.py @@ -74,13 +74,16 @@ class MLAPrefillBackend(ABC): return False def supports_out(self) -> bool: - """Whether `run_prefill_new_tokens` honors a caller-provided `out` - tensor of shape `[num_tokens, num_heads, v_head_dim]`, writing the - final result into it in place. + """Whether `run_prefill_new_tokens` and `run_prefill_context_chunk` honor + a caller-provided `out` tensor of shape + `[num_tokens, num_heads, v_head_dim]`, writing the result into it in place + and returning it. When True, callers may pass `out` and skip the post-hoc - slice/flatten/copy. False for backends that ignore `out` or emit a - padded (`qk_head_dim`) output. Overridden by backends that support it. + slice/flatten/copy -- and, for context chunks, size the accumulating + partial before running any chunk. False for backends that ignore `out` or + emit a padded (`qk_head_dim`) output. Overridden by backends that support + it. """ return False @@ -179,5 +182,6 @@ class MLAPrefillBackend(ABC): q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: raise NotImplementedError diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index ce81be07f1..eace6f5a10 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -477,6 +477,7 @@ class FlashAttnPrefillBackend(MLAPrefillBackend): q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: return self._flash_attn_varlen_diff_headdims( q=q, @@ -489,6 +490,7 @@ class FlashAttnPrefillBackend(MLAPrefillBackend): softmax_scale=self.scale, causal=False, # Context is unmasked return_softmax_lse=True, + out=out, ) diff --git a/vllm/v1/attention/backends/mla/prefill/flashinfer.py b/vllm/v1/attention/backends/mla/prefill/flashinfer.py index 902aebbdde..f2707aa970 100644 --- a/vllm/v1/attention/backends/mla/prefill/flashinfer.py +++ b/vllm/v1/attention/backends/mla/prefill/flashinfer.py @@ -228,11 +228,13 @@ class FlashInferPrefillBackend(MLAPrefillBackend): q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: attn_out, lse = self._prefill_chunks[chunk.index].run( q=q, k=k, v=v, + out=out, return_lse=True, ) diff --git a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py index 7452300a03..6985b0a70c 100644 --- a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py @@ -166,11 +166,13 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend): q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: from tokenspeed_mla import tokenspeed_mla_prefill # See note in run_prefill_new_tokens — `v` is a split-view of `kv_nope` - # in `_compute_prefill_context` and arrives non-contiguous. + # in the generic `_compute_prefill_context` and arrives non-contiguous. + # (K3's fused fp8 K/V pack already hands over a contiguous `v`.) v = v.contiguous() attn_out, lse = tokenspeed_mla_prefill( @@ -187,6 +189,7 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend): cum_seq_lens_q=chunk.query_start_loc, max_seq_len_q=chunk.max_query_len, enable_pdl=False, + out=out, ) # Convert from (q_len, num_heads) to (num_heads, q_len) diff --git a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py index b4f4181bdb..144a75aded 100644 --- a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py +++ b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py @@ -147,16 +147,18 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend): q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: from flashinfer.prefill import trtllm_ragged_attention_deepseek - out = torch.empty( - q.shape[0], - q.shape[1], - v.shape[2], - device=q.device, - dtype=self._prefill_metadata.output_dtype, - ) + if out is None: + out = torch.empty( + q.shape[0], + q.shape[1], + v.shape[2], + device=q.device, + dtype=self._prefill_metadata.output_dtype, + ) attn_out, lse = trtllm_ragged_attention_deepseek( query=q,