From e6eb0d120c480aa5116aafc2d39f19c8ebcf3181 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Mon, 14 Sep 2026 20:12:55 -0700 Subject: [PATCH] [Model][DSv4.1] Store the whole KV in MXFP8 (FlashMLA V4.1 record) (#56893) Signed-off-by: Yongye Zhu Signed-off-by: Yongye Zhu Co-authored-by: Yongye Zhu Co-authored-by: Claude Opus 5 (1M context) --- cmake/external_projects/flashmla.cmake | 85 +++++--- ...deepseek_v4_qnorm_rope_kv_insert_kernel.cu | 94 ++++++--- csrc/libtorch_stable/ops.h | 2 +- csrc/libtorch_stable/torch_bindings.cpp | 2 +- tests/kernels/test_compressor_kv_cache.py | 130 ++++++++++++ ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 98 +++++++++ .../nvidia/ops/dequant_gather_k_cutedsl.py | 62 +++--- vllm/models/deepseek_v41/amd/dspark.py | 2 + vllm/models/deepseek_v41/attention.py | 43 +++- .../deepseek_v41/common/ops/cache_utils.py | 194 +++++++++++++++++- .../common/ops/fused_compress_quant_cache.py | 76 ++++++- vllm/models/deepseek_v41/nvidia/dspark.py | 2 + .../attention/backends/mla/flashmla_sparse.py | 11 + vllm/v1/attention/backends/mla/sparse_swa.py | 18 +- 14 files changed, 709 insertions(+), 110 deletions(-) diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake index 382fd0289a..cfa8243f72 100644 --- a/cmake/external_projects/flashmla.cmake +++ b/cmake/external_projects/flashmla.cmake @@ -19,7 +19,7 @@ else() FetchContent_Declare( flashmla GIT_REPOSITORY https://github.com/vllm-project/FlashMLA - GIT_TAG 6bc49418c5ead572ff0339191ddf3b155749e183 + GIT_TAG 0eee43b12f034b657133cf2afca6a72ebb6efccf GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" @@ -76,43 +76,77 @@ if(FLASH_MLA_ARCHS) set(FlashMLA_SOURCES ${flashmla_SOURCE_DIR}/csrc/api/api.cpp + ${flashmla_SOURCE_DIR}/csrc/api/sparse_prefill.cpp + ${flashmla_SOURCE_DIR}/csrc/api/sparse_decode.cpp + ${flashmla_SOURCE_DIR}/csrc/api/dense_decode.cpp + ${flashmla_SOURCE_DIR}/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp # Misc kernels for decoding - ${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu - ${flashmla_SOURCE_DIR}/csrc/smxx/decode/combine/combine.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/combine/combine.cu # sm90 dense decode - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/fp16.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/bf16.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/fp16.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/bf16.cu # sm90 sparse decode - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu # sm90 sparse prefill - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu - # sm100 dense prefill & backward - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu + # sm100 dense prefill (inference-only: dense backward stays out) + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu # sm100 sparse prefill - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu + + # sm100 fused norm + rope + attn + rope + cast + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu # sm100 sparse decode - ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32_nvfp4_fp8rope.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu ) set(FlashMLA_Extension_SOURCES @@ -125,7 +159,6 @@ if(FLASH_MLA_ARCHS) set(FlashMLA_INCLUDES ${flashmla_SOURCE_DIR}/csrc ${flashmla_SOURCE_DIR}/csrc/kerutils/include - ${flashmla_SOURCE_DIR}/csrc/sm90 ${flashmla_SOURCE_DIR}/csrc/cutlass/include ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include ) diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 06fe7108bd..254b9783bc 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -23,9 +23,15 @@ * cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin * second along last dim; each half is rope_dim/2 = 32 values) * - * Cache layout per paged-cache block (block_size tokens): - * [0, bs*576): token data, 448 fp8 + 128 bf16 each - * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token + * Cache layout per paged-cache block (block_size tokens), selected by the + * `kv_mxfp8` argument: + * V4 (584 B/token, kv_mxfp8=false) + * [0, bs*576): token data, 448 fp8 + 128 bf16 each + * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token + * V4.1 (528 B/token, kv_mxfp8=true) -- RoPE is quantized as well, so the + * whole row is MXFP8 (fp8 e4m3 with one UE8M0 scale per 32 dims) + * [0, bs*512): token data, 512 fp8 each + * [bs*512, bs*512 + bs*16): UE8M0 scales, 16 per token */ #include "torch_utils.h" @@ -91,6 +97,10 @@ constexpr int kQuantBlock = 64; constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock; // 7 constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1; // 8 (7 real + 1 pad) constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576 +// V4.1 MXFP8 record: every dim quantized, one UE8M0 scale per 32 dims. +constexpr int kMxfp8QuantBlock = 32; +constexpr int kMxfp8NumQuantBlocks = kHeadDim / kMxfp8QuantBlock; // 16 +constexpr int kMxfp8TokenDataBytes = kHeadDim; // 512 // FNUZ on gfx942 / OCP elsewhere. FNUZ uses 224.0 (not the dtype's raw // 240.0) to match the rest of vLLM's FNUZ pipeline. #if defined(USE_ROCM) && defined(__gfx942__) @@ -141,6 +151,11 @@ __device__ __forceinline__ uint4 packFp8E4M3x16(float const* values, // ──────────────────────────────────────────────────────────────────────────── // Small inline helpers // ──────────────────────────────────────────────────────────────────────────── +__device__ __forceinline__ float warp2MaxAbs(float val) { + // Reduce absolute max across 2 consecutive lanes (32 dims = one MXFP8 block). + return fmaxf(val, __shfl_xor_sync(FINAL_MASK, val, 1)); +} + __device__ __forceinline__ float warp4MaxAbs(float val) { // Reduce absolute max across 4 consecutive lanes (lane id & 3 group). float peer = __shfl_xor_sync(FINAL_MASK, val, 1); @@ -176,6 +191,10 @@ __device__ __forceinline__ float warpSum(float val) { // slot == kNumHeadsQPadded → KV (RoPE + UE8M0 quant // + paged-cache // insert) +// +// `kv_mxfp8` is grid-uniform and only steers the KV branch, so it stays a +// runtime argument rather than doubling every (head count, q-norm) +// instantiation of a kernel whose hot path is the Q branch. template __device__ __forceinline__ void processDeepseekV4Slot( uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx, @@ -184,7 +203,7 @@ __device__ __forceinline__ void processDeepseekV4Slot( uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping, int64_t const* __restrict__ position_ids, float const* __restrict__ cos_sin_cache, int const cache_block_size, - int const kv_block_stride) { + int const kv_block_stride, bool const kv_mxfp8) { using Converter = vllm::_typeConvert; bool const isKV = (slotIdx == kNumHeadsQPadded); bool const isPadQ = !isKV && (slotIdx >= num_heads_q); @@ -302,19 +321,24 @@ __device__ __forceinline__ void processDeepseekV4Slot( *reinterpret_cast(dst) = out0; *reinterpret_cast(dst + 8) = out1; } else { - // ── KV: FP8 quant on NoPE + bf16 store on RoPE + cache insert. + // ── KV: FP8 quant + cache insert. V4 keeps the RoPE dims in bf16; + // V4.1 (kv_mxfp8) quantizes them too. int64_t const slot_id = slot_mapping[tokenIdx]; if (slot_id >= 0) { + int const token_data_bytes = + kv_mxfp8 ? kMxfp8TokenDataBytes : kTokenDataBytes; + int const scale_bytes = + kv_mxfp8 ? kMxfp8NumQuantBlocks : kScaleBytesPerToken; int64_t const block_idx = slot_id / cache_block_size; int64_t const pos_in_block = slot_id % cache_block_size; uint8_t* block_base = k_cache + block_idx * static_cast(kv_block_stride); - uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes; + uint8_t* token_fp8_ptr = block_base + pos_in_block * token_data_bytes; uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim; uint8_t* token_scale_ptr = block_base + - static_cast(cache_block_size) * kTokenDataBytes + - pos_in_block * kScaleBytesPerToken; + static_cast(cache_block_size) * token_data_bytes + + pos_in_block * scale_bytes; #pragma unroll for (int i = 0; i < kElemsPerLane; i++) { @@ -326,11 +350,13 @@ __device__ __forceinline__ void processDeepseekV4Slot( for (int i = 0; i < kElemsPerLane; i++) { local_absmax = fmaxf(local_absmax, fabsf(elements[i])); } - float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f); + float const absmax = fmaxf( + kv_mxfp8 ? warp2MaxAbs(local_absmax) : warp4MaxAbs(local_absmax), + 1e-4f); float const exponent = ceilf(log2f(absmax / kFp8Max)); float const inv_scale = exp2f(-exponent); - if (!is_rope_lane) { + if (kv_mxfp8 || !is_rope_lane) { uint8_t out_bytes[kElemsPerLane]; #pragma unroll for (int i = 0; i < kElemsPerLane; i++) { @@ -347,13 +373,19 @@ __device__ __forceinline__ void processDeepseekV4Slot( *reinterpret_cast(token_fp8_ptr + dim_base) = *reinterpret_cast(out_bytes); - if ((laneId & 3) == 0) { - int const q_block_idx = laneId >> 2; - float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f); - token_scale_ptr[q_block_idx] = static_cast(encoded); - } - if (laneId == 0) { - token_scale_ptr[kNumQuantBlocks] = 0; + float const encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f); + if (kv_mxfp8) { + // One scale per 32 dims = one per lane pair. + if ((laneId & 1) == 0) { + token_scale_ptr[laneId >> 1] = static_cast(encoded); + } + } else { + if ((laneId & 3) == 0) { + token_scale_ptr[laneId >> 2] = static_cast(encoded); + } + if (laneId == 0) { + token_scale_ptr[kNumQuantBlocks] = 0; + } } } else { uint4 out0, out1; @@ -421,7 +453,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( int const num_tokens_insert, // = slot_mapping.size(0), ≤ num_tokens_full int const num_heads_q, // live Q heads (input layout) int const cache_block_size, // tokens per paged-cache block - int const kv_block_stride) { // bytes per paged-cache block + int const kv_block_stride, // bytes per paged-cache block + bool const kv_mxfp8) { // V4.1 all-dims MXFP8 record #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) // BF16 _typeConvert specialization is unavailable on pre-Ampere. The // DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a @@ -477,7 +510,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( processDeepseekV4Slot( v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size, - kv_block_stride); + kv_block_stride, kv_mxfp8); #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) cudaTriggerProgrammaticLaunchCompletion(); @@ -506,7 +539,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( float const* __restrict__ cos_sin_cache, float const eps, int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const cache_block_size, - int const kv_block_stride) { + int const kv_block_stride, bool const kv_mxfp8) { #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) if constexpr (std::is_same_v) { return; @@ -564,7 +597,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( processDeepseekV4Slot( v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId, num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, - cos_sin_cache, cache_block_size, kv_block_stride); + cos_sin_cache, cache_block_size, kv_block_stride, kv_mxfp8); // ── Buffer rotation: hand the prefetched LDGs to the next iter. v0_curr = v0_next; @@ -590,7 +623,7 @@ static void launchFusedDeepseekV4Templated( uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids, float const* cos_sin_cache, float const eps, int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, - int const cache_block_size, int const kv_block_stride, + int const cache_block_size, int const kv_block_stride, bool const kv_mxfp8, cudaStream_t stream) { constexpr int kBlockSize = 256; constexpr int kWarpsPerBlock = kBlockSize / 32; @@ -630,7 +663,8 @@ static void launchFusedDeepseekV4Templated( scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, - num_heads_q, cache_block_size, kv_block_stride); + num_heads_q, cache_block_size, kv_block_stride, + kv_mxfp8); } else { config.gridDim = dim3(num_tokens_full); cudaLaunchKernelEx( @@ -639,7 +673,7 @@ static void launchFusedDeepseekV4Templated( scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, - kv_block_stride); + kv_block_stride, kv_mxfp8); } #else // ROCm: use standard kernel launch syntax (no PDL/stream serialization) @@ -649,7 +683,7 @@ static void launchFusedDeepseekV4Templated( <<>>( q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, - cache_block_size, kv_block_stride); + cache_block_size, kv_block_stride, kv_mxfp8); #endif } @@ -663,19 +697,19 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const num_heads_q_padded, int const cache_block_size, int const kv_block_stride, - bool const apply_q_norm, cudaStream_t stream) { + bool const apply_q_norm, bool const kv_mxfp8, cudaStream_t stream) { #define DISPATCH(N) \ case N: \ if (apply_q_norm) { \ launchFusedDeepseekV4Templated( \ q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \ cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \ - num_heads_q, cache_block_size, kv_block_stride, stream); \ + num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \ } else { \ launchFusedDeepseekV4Templated( \ q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \ cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \ - num_heads_q, cache_block_size, kv_block_stride, stream); \ + num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \ } \ return; @@ -967,7 +1001,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& position_ids, // [N] int64 torch::stable::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16 int64_t q_head_padded, // padded Q head count for output - double eps, int64_t cache_block_size, bool apply_q_norm) { + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8) { STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(), "q_in must be contiguous CUDA"); STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(), @@ -1036,7 +1070,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( cos_sin_cache.const_data_ptr(), static_cast(eps), num_tokens_full, num_tokens_insert, num_heads_q, num_heads_q_padded, cache_block_size_i, kv_block_stride, - apply_q_norm, stream); + apply_q_norm, kv_mxfp8, stream); }); return q_out; } diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 68a24ca434..e19f49d603 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -264,7 +264,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, torch::stable::Tensor const& position_ids, torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, - double eps, int64_t cache_block_size, bool apply_q_norm); + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8); void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( torch::stable::Tensor& q, torch::stable::Tensor const& kv, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 110f8595ee..54548dfd05 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -430,7 +430,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor q_in, Tensor kv, Tensor! k_cache, " "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " "int q_head_padded, float eps, int cache_block_size, " - "bool apply_q_norm=True) -> Tensor"); + "bool apply_q_norm=True, bool kv_mxfp8=False) -> Tensor"); // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 0f1976cf68..7e092229fe 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -215,6 +215,136 @@ def test_v41_fused_save_compress_and_insert( torch.testing.assert_close(cache_backing, expected_cache, rtol=0, atol=0) +def _mxfp8_record_reference( + latent_row: torch.Tensor, cos_sin_row: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference V4.1 record for one token: 512 fp8 bytes + 16 UE8M0 scales. + + Mirrors FlashMLA's ``KVCacheLayout.V41_FP8Sparse`` quantizer: rotate the + RoPE tail first, then scale every 32-dim tile, RoPE tiles included. + """ + row = latent_row.float() + c, s = cos_sin_row.chunk(2) + rotated = row.clone() + even, odd = row[448::2], row[449::2] + rotated[448::2] = even * c - odd * s + rotated[449::2] = odd * c + even * s + quantized, scales = _ue8m0_reference(rotated, 32, 448.0) + return quantized.view(torch.uint8), (scales.log2() + 127).to(torch.uint8) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize("compress_ratio", [1, 2]) +def test_v41_rope_insert_mxfp8_record(compress_ratio: int): + """The 528-byte V4.1 record quantizes the RoPE dims instead of keeping them + in bf16, so a page is [block_size x 512 fp8][block_size x 16 UE8M0]. + + Byte-exact against the reference quantizer, including the rows and page + padding the kernel must leave alone. + """ + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + rope_quant_insert, + ) + + torch.manual_seed(7) + device = "cuda" + cache_block = 64 + cache_stride = math.ceil(cache_block * 528 / 512) * 512 + num_tokens = 12 + + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) + angles = torch.randn(64, 32, device=device) + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + + cache_backing = torch.full((2, cache_stride), 165, dtype=torch.uint8, device=device) + cache = cache_backing.as_strided((2, cache_block, 528), (cache_stride, 528, 1)) + # Leave one token unmapped so the kernel's negative-slot skip is covered. + slots = torch.arange(num_tokens, dtype=torch.int64, device=device) + slots[3] = -1 + + rope_quant_insert(latent, positions, cos_sin, cache, slots, compress_ratio) + + expected = torch.full_like(cache_backing, 165) + for t in range(num_tokens): + slot = slots[t].item() + pos = positions[t].item() + if slot < 0 or (pos + 1) % compress_ratio: + continue + page, row = divmod(slot, cache_block) + values, scales = _mxfp8_record_reference( + latent[t], cos_sin[pos // compress_ratio * compress_ratio] + ) + expected[page, row * 512 : (row + 1) * 512] = values + scale_offset = cache_block * 512 + row * 16 + expected[page, scale_offset : scale_offset + 16] = scales + + torch.testing.assert_close(cache_backing, expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +def test_v41_mxfp8_cache_round_trip(): + """quantize_and_insert -> dequantize_and_gather recovers the V4.1 record. + + The RoPE dims now go through fp8, so they carry quantization error too -- + the check is that every dim survives within one MXFP8 tile step, and that + the bytes on the way match the reference quantizer. + """ + from vllm.models.deepseek_v41.common.ops import ( + dequantize_and_gather_k_cache, + quantize_and_insert_k_cache, + ) + + torch.manual_seed(11) + device = "cuda" + block_size = 64 + num_tokens = 70 + num_blocks = 4 + page_bytes = math.ceil(block_size * 528 / 512) * 512 + + k = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) + cache = torch.zeros(num_blocks, page_bytes, dtype=torch.uint8, device=device) + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + quantize_and_insert_k_cache( + k, cache, slot_mapping, block_size=block_size, bytes_per_token=528 + ) + + for t in (0, 1, block_size, num_tokens - 1): + values, scales = _ue8m0_reference(k[t].float(), 32, 448.0) + page, row = divmod(t, block_size) + torch.testing.assert_close( + cache[page, row * 512 : (row + 1) * 512], + values.view(torch.uint8), + rtol=0, + atol=0, + ) + scale_offset = block_size * 512 + row * 16 + torch.testing.assert_close( + cache[page, scale_offset : scale_offset + 16], + (scales.log2() + 127).to(torch.uint8), + rtol=0, + atol=0, + ) + + out = torch.zeros(1, num_tokens, 512, dtype=torch.bfloat16, device=device) + dequantize_and_gather_k_cache( + out, + cache.view(num_blocks, block_size, 528), + seq_lens=torch.tensor([num_tokens], dtype=torch.int32, device=device), + gather_lens=None, + block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view( + 1, -1 + ), + block_size=block_size, + offset=0, + ) + # Half an e4m3 ULP at the top of the range (16 units) bounds the error. + tile_amax = k.float().abs().view(num_tokens, 16, 32).amax(-1).clamp(min=1e-4) + tolerance = 16.0 * (tile_amax / 448.0).log2().ceil().exp2() + error = (out[0].float() - k.float()).abs() + assert (error <= tolerance.repeat_interleave(32, dim=-1)).all() + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") @pytest.mark.parametrize("compress_ratio", [1, 2]) @pytest.mark.parametrize("store_fp8", [False, True]) diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index 0e25884368..68eabfb94e 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -424,6 +424,104 @@ def test_kv_path_matches_reference(num_tokens: int, block_size: int): ) +# ── Test 2c: DeepSeek V4.1 MXFP8 record (kv_mxfp8=True) ────────────────────── + +V41_HEAD_BYTES = HEAD_DIM + HEAD_DIM // 32 # 512 fp8 + 16 UE8M0 scales = 528 + + +@pytest.mark.parametrize("num_tokens", [1, 17, 64]) +@pytest.mark.parametrize("block_size", [32, 64]) +def test_kv_path_mxfp8_record_matches_reference(num_tokens: int, block_size: int): + """``kv_mxfp8=True`` writes DeepSeek's V4.1 record: all 512 dims as fp8 e4m3 + with one UE8M0 scale per 32 dims, no bf16 RoPE tail. + + Both sides quantize the same bf16-rounded rotated row, so the bytes must + match the Triton encoder exactly apart from the documented 1-ULP bf16 + rounding split on the rotation itself. + """ + from vllm.models.deepseek_v41.common.ops import ( + dequantize_and_gather_k_cache as v41_dequantize_and_gather_k_cache, + ) + from vllm.models.deepseek_v41.common.ops import ( + quantize_and_insert_k_cache as v41_quantize_and_insert_k_cache, + ) + + torch.manual_seed(3) + device = "cuda" + dtype = torch.bfloat16 + max_pos = 4096 + + kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device) + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device) + + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + + k_cache_ref = torch.zeros( + num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device + ) + v41_quantize_and_insert_k_cache( + apply_rope_gptj_last_k(kv, positions, cos_sin_cache), + k_cache_ref, + slot_mapping, + block_size=block_size, + use_fnuz=USE_FNUZ, + bytes_per_token=V41_HEAD_BYTES, + ) + + k_cache_fused = torch.zeros_like(k_cache_ref) + q_dummy = torch.zeros(num_tokens, 1, HEAD_DIM, dtype=dtype, device=device) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q_dummy, + kv, + k_cache_fused, + slot_mapping, + positions, + cos_sin_cache, + 64, + 1e-6, + block_size, + True, + True, # kv_mxfp8 + ) + + # A pre-quantization bf16 tie can land the two encoders on adjacent e4m3 + # codes, so compare after dequantizing and allow one step at the tile scale. + def _dequant(cache_2d): + out = torch.zeros(1, num_tokens, HEAD_DIM, dtype=dtype, device=device) + v41_dequantize_and_gather_k_cache( + out, + cache_2d.view(num_blocks, block_size, V41_HEAD_BYTES), + torch.tensor([num_tokens], dtype=torch.int32, device=device), + None, + torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(0), + block_size, + offset=0, + use_fnuz=USE_FNUZ, + ) + return out[0] + + rec_ref = _dequant(k_cache_ref).float() + rec_fused = _dequant(k_cache_fused).float() + tile_scale = ( + (rec_ref.abs().view(num_tokens, 16, 32).amax(-1).clamp(min=1e-4) / FP8_MAX) + .log2() + .ceil() + .exp2() + ) + # One e4m3 ULP at the top of the range is 32 quantized units. + step = 32.0 * tile_scale.repeat_interleave(32, dim=-1) + assert ((rec_fused - rec_ref).abs() <= step).all() + # The NoPE dims see no rotation at all, so they must be byte-identical. + nope_bytes = slice(0, block_size * HEAD_DIM) + ref_rows = k_cache_ref[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM) + fused_rows = k_cache_fused[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM) + torch.testing.assert_close( + fused_rows[..., :NOPE_DIM], ref_rows[..., :NOPE_DIM], rtol=0, atol=0 + ) + + # ── Test 2b: DP padding (slot_mapping shorter than q/kv) ───────────────────── diff --git a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py index fcb12740f2..5920ed290c 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py @@ -23,10 +23,11 @@ from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( class DequantGatherKCacheKernel( VllmCuTeDSLJitKernel["DequantGatherKCacheKernel.CompileKey"] ): - # Hard-coded for DSv4. + # DSv4 geometry. The record format (how much of the row is quantized and + # how wide a scale group is) comes from the compile key: V4 quantizes the + # 448 NoPE dims in groups of 64 and keeps a bf16 RoPE tail, V4.1 quantizes + # all 512 dims in groups of 32 (MXFP8) and has no tail. head_dim = 512 - fp8_dim = 448 - group_size = 64 # 1 scale per 64 elems num_warps = 4 tb_size = num_warps * 32 num_stages = 4 @@ -35,17 +36,20 @@ class DequantGatherKCacheKernel( class CompileKey: block_size: int has_gather_lens: bool + fp8_dim: int = 448 + group_size: int = 64 # dims per scale @staticmethod def kernel(compile_key: CompileKey) -> Any: head_dim = DequantGatherKCacheKernel.head_dim - fp8_dim = DequantGatherKCacheKernel.fp8_dim - group_size = DequantGatherKCacheKernel.group_size + fp8_dim = compile_key.fp8_dim + group_size = compile_key.group_size num_warps = DequantGatherKCacheKernel.num_warps tb_size = DequantGatherKCacheKernel.tb_size num_stages = DequantGatherKCacheKernel.num_stages bf16_dim = head_dim - fp8_dim data_dim = fp8_dim + bf16_dim * 2 + scale_bytes = head_dim // group_size block_size = compile_key.block_size @cute.jit @@ -62,8 +66,8 @@ class DequantGatherKCacheKernel( ): op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL) cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128) - cp8_atom = cute.make_copy_atom( - cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64 + cp_scale_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=scale_bytes * 8 ) page_id = block_table[req_id, pos // block_size] block_offset = pos % block_size @@ -88,7 +92,7 @@ class DequantGatherKCacheKernel( ) elif idx == cutlass.const_expr(data_dim // 16): cute.copy( - cp8_atom, + cp_scale_atom, k_scale[page_id, block_offset, None], s_kscale[None, stage_id], ) @@ -119,8 +123,8 @@ class DequantGatherKCacheKernel( )[None, warp_id, None] s_kscale = smem.allocate_tensor( Uint8, - cute.make_layout((8, num_warps, num_stages)), - byte_alignment=8, + cute.make_layout((scale_bytes, num_warps, num_stages)), + byte_alignment=scale_bytes, )[None, warp_id, None] # Prepare for 16B cp.async, also for BF16 smem loads later. @@ -239,14 +243,16 @@ class DequantGatherKCacheKernel( # Last 64 elems are BF16 tail, corresponds to dequant1 of last # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 + - # 8x 16B for BF16. - if lane_id + 32 >= fp8_dim // 8: - idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8 - cute.copy( - cp16_atom, - s_kdata_16B_slice[(None, idx), compute_stage], - dequant1, - ) + # 8x 16B for BF16. The V4.1 record is fully quantized and has + # no tail. + if cutlass.const_expr(bf16_dim > 0): # noqa: SIM102 + if lane_id + 32 >= fp8_dim // 8: + idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8 + cute.copy( + cp16_atom, + s_kdata_16B_slice[(None, idx), compute_stage], + dequant1, + ) # Store two 16B BF16 chunks per lane: first half, then second half. dst = out_slice[req_id, offset + i, (None, lane_id)] @@ -269,7 +275,7 @@ class DequantGatherKCacheKernel( ): # Split k_cache into k_data and k_scale. Each [block_size, head_bytes] # block is actually a concat of - # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8]. + # [block_size, fp8_dim + bf16_dim * 2] and [block_size, scale_bytes]. k_data = cute.make_tensor( k_cache.iterator, layout=cute.make_layout( @@ -282,8 +288,8 @@ class DequantGatherKCacheKernel( k_scale = cute.make_tensor( k_cache.iterator + (block_size * data_dim), layout=cute.make_layout( - (k_cache.shape[0], block_size, 8), - stride=(k_cache.stride[0], 8, 1), + (k_cache.shape[0], block_size, scale_bytes), + stride=(k_cache.stride[0], scale_bytes, 1), ), ) @@ -311,10 +317,14 @@ class DequantGatherKCacheKernel( *, block_size: int, has_gather_lens: bool, + fp8_dim: int = 448, + group_size: int = 64, ) -> CompileKey: return self.CompileKey( block_size=block_size, has_gather_lens=has_gather_lens, + fp8_dim=fp8_dim, + group_size=group_size, ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: @@ -345,9 +355,11 @@ class DequantGatherKCacheKernel( def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: num_reqs = cute.sym_int() head_dim = self.head_dim - fp8_dim = self.fp8_dim + fp8_dim = compile_key.fp8_dim block_size = compile_key.block_size - head_bytes = fp8_dim + (head_dim - fp8_dim) * 2 + 8 + head_bytes = ( + fp8_dim + (head_dim - fp8_dim) * 2 + head_dim // compile_key.group_size + ) out = make_fake_tensor(BFloat16, (num_reqs, cute.sym_int(), head_dim), 16) k_cache = cute.runtime.make_fake_tensor( @@ -384,10 +396,14 @@ class DequantGatherKCacheKernel( block_table: torch.Tensor, block_size: int, offset: int, + fp8_dim: int = 448, + group_size: int = 64, ) -> CuTeDSLLaunchSpec["DequantGatherKCacheKernel.CompileKey"]: compile_key = self.dispatch( block_size=block_size, has_gather_lens=gather_lens is not None, + fp8_dim=fp8_dim, + group_size=group_size, ) launch_args = ( out, diff --git a/vllm/models/deepseek_v41/amd/dspark.py b/vllm/models/deepseek_v41/amd/dspark.py index 3a6cef8aa7..78f07b6b11 100644 --- a/vllm/models/deepseek_v41/amd/dspark.py +++ b/vllm/models/deepseek_v41/amd/dspark.py @@ -259,6 +259,8 @@ def _insert_context_kv( attn.padded_heads, attn.eps, block_size, + True, # apply_q_norm (unused: the query is a discarded dummy) + attn.kv_mxfp8, ) elif cache_dtype == torch.bfloat16: swa_3d = swa_cache.view(-1, block_size, attn.head_dim) diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 131cf4d97e..fa26351681 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -51,6 +51,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.models.utils import extract_layer_index from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope from vllm.models.deepseek_v41.compressor import DeepseekCompressor +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.multi_stream_utils import ( execute_in_parallel, @@ -109,6 +110,16 @@ def _fill_short_context_topk_indices( ) +# Which packed fp8_ds_mla record a V4.1 layer writes. FlashMLA decodes +# DeepSeek's V4.1 record -- all 512 dims (RoPE included) as fp8 e4m3 with one +# UE8M0 scale per 32 dims, 512 data bytes + 16 scale bytes per token, pages +# rounded to the kernel's 512 B TMA stride -- only in its SM100 sparse-decode +# kernels. Every other arch keeps the V4 record: 448 fp8 NoPE + 64 bf16 RoPE +# plus 7 UE8M0 scales of 64 dims and a pad byte (584 B, 576 B pages). +def _use_v41_mxfp8_kv_record() -> bool: + return current_platform.is_device_capability_family(100) + + def _resolve_dsv4_kv_cache_dtype( use_fp8_ds_mla_layout: bool, kv_cache_dtype: str, @@ -445,6 +456,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype( self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config ) + self.kv_mxfp8 = _use_v41_mxfp8_kv_record() + self.kv_bytes_per_token = 528 if self.kv_mxfp8 else 584 + self.kv_page_alignment = 512 if self.kv_mxfp8 else 576 self.swa_cache_layer = DeepseekV4SWACache( head_dim=self.head_dim, @@ -454,6 +468,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): cache_config=cache_config, backend_cls=self.swa_backend_cls, block_size=32, + packed_bytes_per_token=self.kv_bytes_per_token, + packed_page_alignment=self.kv_page_alignment, ) # The attention layer itself was already registered with the @@ -879,6 +895,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self.eps, swa_metadata.block_size, False, + self.kv_mxfp8, ) # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV @@ -934,9 +951,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # DeepseekV4SWACache. if not self.is_kv_source: return None - # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B - # alignment; plain bf16 / per-tensor fp8 rows use natural element-size - # pages. + # fp8_ds_mla is a UE8M0 block-scaled uint8 layout whose page rounds up + # to the decode kernel's TMA stride; plain bf16 / per-tensor fp8 rows + # use natural element-size pages. uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla" return MLAAttentionSpec( block_size=vllm_config.cache_config.block_size, @@ -945,12 +962,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype, tokens_per_state=self.compress_ratio, cache_dtype_str=self.kv_cache_dtype, - alignment=576 if uses_fp8_ds_mla_layout else 512, + alignment=self.kv_page_alignment if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), - # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token; - # head_size stays semantic (512). - state_content_bytes=584 if uses_fp8_ds_mla_layout else None, + # Packed record width; head_size stays semantic (512). + state_content_bytes=( + self.kv_bytes_per_token if uses_fp8_ds_mla_layout else None + ), ) def _compressed_kv_cache(self) -> torch.Tensor: @@ -998,8 +1016,15 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): head_size=self.head_dim, dtype=self.dtype, tokens_per_state=self.compress_ratio, - # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). - alignment=576 if uses_fp8_ds_mla_layout else 512, + # The indexer page shares a physical block with the main KV cache + # (the block stride is the sum of every page in the group), so it + # has to carry the same padding: the packed record's TMA stride, + # else 512B for FlashInfer sparse (#44577). + alignment=( + 576 + if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() + else 512 + ), ) def forward(self): ... diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 64462eb73b..51743016bf 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -36,6 +36,18 @@ from vllm.triton_utils import tl, triton from vllm.utils.import_utils import has_cutedsl from vllm.utils.math_utils import next_power_of_2 +# Per-token byte width of the two paged fp8 records. Both put a page's whole +# data region ahead of its whole scale region, so ``k_cache.shape[-1]`` -- the +# width the KV-cache spec publishes -- identifies the record: +# V4 (584 B): 448 fp8 e4m3 NoPE + 64 bf16 RoPE, then 7 UE8M0 scales of 64 +# dims each plus a pad byte. +# V4.1 (528 B): all 512 dims as fp8 e4m3 (RoPE included), then 16 UE8M0 +# scales of 32 dims each. FlashMLA's ``ModelType::V41``. +V4_BYTES_PER_TOKEN = 584 +V41_BYTES_PER_TOKEN = 528 +V41_QUANT_BLOCK = 32 +V41_NUM_SCALES = 512 // V41_QUANT_BLOCK # 16 + @triton.jit def quantize_and_insert_k_kernel( @@ -163,6 +175,67 @@ def quantize_and_insert_k_kernel( tl.store(bf16_out_ptr + chunk_offsets, bf16_vals) +@triton.jit +def _quantize_and_insert_k_mxfp8_kernel( + k_ptr, # [num_tokens, 512] bf16 + slot_mapping_ptr, # [num_tokens] int64 + k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view) + num_tokens, + head_dim: tl.constexpr, # 512, RoPE dims included + scale_dim: tl.constexpr, # 16 + quant_block: tl.constexpr, # 32 dims per UE8M0 scale + cache_block_size: tl.constexpr, # paged cache block size, in tokens + block_stride: tl.constexpr, # total bytes per block (padded) + fp8_max: tl.constexpr, + use_fnuz: tl.constexpr = False, +): + """Quantize a bf16 K row to MXFP8 and insert it into a V4.1 paged cache. + + One program per token. Unlike the V4 record there is no bf16 tail, so every + dim is scaled. + """ + pid = tl.program_id(0) + if pid >= num_tokens: + return + + slot_idx = tl.load(slot_mapping_ptr + pid) + if slot_idx == -1: + return + + block_idx = slot_idx // cache_block_size + pos_in_block = slot_idx % cache_block_size + + # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks. + cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride + token_data_ptr = cache_block_ptr + pos_in_block * head_dim + token_scale_ptr = ( + cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim + ) + + d = tl.arange(0, head_dim) + x = tl.load(k_ptr + pid.to(tl.int64) * head_dim + d).to(tl.float32) + + tiles = tl.reshape(x, (scale_dim, quant_block)) + # Floor the absmax like the CUDA encoder so an all-zero tile still lands on + # a finite scale. + amax = tl.maximum(tl.max(tl.abs(tiles), 1), 1e-4) + exponent = tl.ceil(tl.log2(amax * (1.0 / fp8_max))) + scaled = tl.clamp( + tiles * tl.reshape(tl.exp2(-exponent), (scale_dim, 1)), -fp8_max, fp8_max + ) + if use_fnuz: # noqa: SIM108 + fp8 = scaled.to(tl.float8e4b8) + else: + fp8 = scaled.to(tl.float8e4nv) + tl.store( + token_data_ptr + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (head_dim,)) + ) + + # UE8M0 encoding: stored_value = exponent + 127 (bias). + encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), 255.0) + tl.store(token_scale_ptr + tl.arange(0, scale_dim), encoded.to(tl.uint8)) + + def quantize_and_insert_k_cache( k: torch.Tensor, # [num_tokens, 512] bf16 k_cache: torch.Tensor, # [num_blocks, block_bytes] uint8 @@ -170,16 +243,13 @@ def quantize_and_insert_k_cache( block_size: int = 64, is_ue8m0: bool = True, use_fnuz: bool = False, + bytes_per_token: int = V4_BYTES_PER_TOKEN, ): """ Quantize K tensor and insert into paged K cache. - K Cache block layout (block_size=64 tokens): - - First 64 * 576 = 36864 bytes: Token data - - Each token: 448 bytes (fp8) + 128 bytes (bf16) - - Next 64 * 8 = 512 bytes: Scales - - Each token: 8 bytes (uint8 scales, 7 real + 1 padding) - - Padded to multiple of 576 + ``bytes_per_token`` picks the record (see the module header): the V4 one, + or V4.1's all-dims MXFP8 one. ``use_fnuz=True`` selects FNUZ E4M3 cache encoding and is only valid on platforms whose FP8 format is FNUZ. ``use_fnuz=False`` selects OCP E4M3, @@ -196,6 +266,26 @@ def quantize_and_insert_k_cache( num_tokens = slot_mapping.shape[0] block_stride = k_cache.stride(0) # bytes per block + if bytes_per_token == V41_BYTES_PER_TOKEN: + if use_fnuz and not current_platform.is_fp8_fnuz(): + raise ValueError("use_fnuz=True requires a platform using FNUZ FP8") + _quantize_and_insert_k_mxfp8_kernel[(num_tokens,)]( + k, + slot_mapping, + k_cache, + num_tokens, + head_dim=512, + scale_dim=V41_NUM_SCALES, + quant_block=V41_QUANT_BLOCK, + cache_block_size=block_size, + block_stride=block_stride, + fp8_max=get_fp8_min_max()[1] + if use_fnuz + else torch.finfo(torch.float8_e4m3fn).max, + use_fnuz=use_fnuz, + ) + return + TOKEN_FP8_DIM = 448 TOKEN_BF16_DIM = 64 TOKEN_SCALE_DIM = 8 @@ -343,6 +433,69 @@ def _dequantize_and_gather_k_kernel( tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals) +@triton.jit +def _dequantize_and_gather_k_mxfp8_kernel( + out_ptr, + out_stride0, + out_stride1, + k_cache_ptr, + seq_lens_ptr, + block_table_ptr, + offset, + gather_lens_ptr, + max_blocks_per_seq: tl.constexpr, + head_dim: tl.constexpr, # 512 + scale_dim: tl.constexpr, # 16 + quant_block: tl.constexpr, # 32 + cache_block_size: tl.constexpr, + block_stride: tl.constexpr, + use_fnuz: tl.constexpr = False, +): + """Gather V4.1 MXFP8 rows into a bf16 workspace, dequantizing in place.""" + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) + + seq_len = tl.load(seq_lens_ptr + batch_idx) + if gather_lens_ptr is not None: # noqa: SIM108 + gather_len = tl.load(gather_lens_ptr + batch_idx) + else: + gather_len = seq_len + start_pos = seq_len - gather_len + + d = tl.arange(0, head_dim) + s = tl.arange(0, scale_dim) + for i in range(worker_id, gather_len, num_workers): + pos = start_pos + i + block_in_seq = pos // cache_block_size + pos_in_block = pos % cache_block_size + + block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq + physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32 + + # int64: physical_block_idx * block_stride can exceed 2^31. + cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride + token_data_ptr = cache_block_ptr + pos_in_block * head_dim + token_scale_ptr = ( + cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim + ) + + x_uint8 = tl.load(token_data_ptr + d) + if use_fnuz: + x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True) + else: + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) + tiles = tl.reshape(x_fp8.to(tl.float32), (scale_dim, quant_block)) + + # UE8M0: scale = 2^(stored_value - 127). + encoded = tl.load(token_scale_ptr + s) + scale = tl.exp2(encoded.to(tl.float32) - 127.0) + dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,)) + + output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 + tl.store(output_row_ptr + d, dequant.to(tl.bfloat16)) + + def dequantize_and_gather_k_cache_triton( # [num_reqs, max_num_tokens, head_size] out: torch.Tensor, @@ -358,6 +511,28 @@ def dequantize_and_gather_k_cache_triton( offset: int, use_fnuz: bool = False, ) -> None: + num_reqs = seq_lens.shape[0] + NUM_WORKERS = 128 + if k_cache.shape[-1] == V41_BYTES_PER_TOKEN: + _dequantize_and_gather_k_mxfp8_kernel[(num_reqs, NUM_WORKERS)]( + out, + out.stride(0), + out.stride(1), + k_cache, + seq_lens, + block_table, + offset, + gather_lens, + max_blocks_per_seq=block_table.shape[-1], + head_dim=512, + scale_dim=V41_NUM_SCALES, + quant_block=V41_QUANT_BLOCK, + cache_block_size=block_size, + block_stride=k_cache.stride(0), + use_fnuz=use_fnuz, + ) + return + TOKEN_FP8_DIM = 448 TOKEN_BF16_DIM = 64 TOKEN_SCALE_DIM = 8 @@ -365,8 +540,6 @@ def dequantize_and_gather_k_cache_triton( FP8_MAX = 448.0 TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 - num_reqs = seq_lens.shape[0] - NUM_WORKERS = 128 _dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)]( out, out.stride(0), @@ -408,6 +581,8 @@ def dequantize_and_gather_k_cache( ) -> None: """Dequantize and gather a paged DSv4 K cache. + The record is read off ``k_cache.shape[-1]``; see the module header. + ``use_fnuz`` MUST match the encoder of the specific cache being read: ``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere), ``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder @@ -419,6 +594,7 @@ def dequantize_and_gather_k_cache( _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, ) + mxfp8 = k_cache.shape[-1] == V41_BYTES_PER_TOKEN _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL( out=out, k_cache=k_cache, @@ -427,6 +603,8 @@ def dequantize_and_gather_k_cache( block_table=block_table, block_size=block_size, offset=offset, + fp8_dim=512 if mxfp8 else 448, + group_size=V41_QUANT_BLOCK if mxfp8 else 64, ) return diff --git a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py index 5950c4b6ba..96f15849dd 100644 --- a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py @@ -232,10 +232,13 @@ def rope_quant_insert( The BF16 latent supplies both NoPE quantization and RoPE input. It is read only for valid slots at group boundaries. The cache dtype selects the - layout: ``uint8`` is the fp8_ds_mla paged layout (576 value bytes and eight - segregated UE8M0 scale bytes per token, including one zero padding scale); - ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows - read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``. + layout: ``uint8`` is the fp8_ds_mla paged layout, whose record the + per-token byte width names -- 584 B for V4 (576 value bytes and eight + segregated UE8M0 scale bytes, including one zero padding scale) or 528 B + for V4.1 (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 + scales of 32 dims each). ``bfloat16`` and ``float8_e4m3fn`` are the plain + [448 NoPE | 64 RoPE] rows read by FlashInfer, the latter scaled by the + per-tensor ``fp8_scale``. """ assert compress_ratio in (1, 2) assert latent.shape[1] == 512 and latent.dtype == torch.bfloat16 @@ -246,8 +249,15 @@ def rope_quant_insert( return launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {} if kv_cache.dtype == torch.uint8: - assert kv_cache.shape[-1] == 584 - _rope_quant_insert_kernel[(num_tokens,)]( + assert kv_cache.shape[-1] in (584, 528), ( + f"unsupported paged KV record width {kv_cache.shape[-1]}" + ) + kernel = ( + _rope_quant_insert_mxfp8_kernel + if kv_cache.shape[-1] == 528 + else _rope_quant_insert_kernel + ) + kernel[(num_tokens,)]( latent, positions, cos_sin_cache, @@ -338,6 +348,60 @@ def _rope_quant_insert_kernel( tl.store(rope_dst + d - 448, rotated.to(tl.bfloat16), d >= 448) +@triton.jit +def _rope_quant_insert_mxfp8_kernel( + latent, + positions, + cos_sin, + cache, + cache_slots, + COS_STRIDE: tl.constexpr, + CACHE_STRIDE: tl.constexpr, + CACHE_BLOCK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + SANITIZE_CACHE_NANS: tl.constexpr, +): + """V4.1 record: RoPE first, then MXFP8-quantize all 512 dims. + + The RoPE dims are quantized here too, so unlike the V4 kernel the rotation + has to happen before the scales are picked. + """ + t = tl.program_id(0) + slot = tl.load(cache_slots + t) + if slot < 0: + return + position = tl.load(positions + t) + if (position + 1) % COMPRESS_RATIO != 0: + return + d = tl.arange(0, 512) + normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32) + + # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there. + even, odd = tl.split(tl.reshape(normed, (256, 2))) + pair = tl.arange(0, 256) - 224 + cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE + c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32) + s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32) + rotated = tl.interleave(even * c - odd * s, odd * c + even * s) + if SANITIZE_CACHE_NANS: + rotated = tl.where(rotated == rotated, rotated, 0.0) + + page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE + values = page + (slot % CACHE_BLOCK) * 512 + scales = page + CACHE_BLOCK * 512 + (slot % CACHE_BLOCK) * 16 + + quant = tl.reshape(rotated, (16, 32)) + amax = tl.maximum(tl.max(tl.abs(quant), 1), 1e-4) + exponent = tl.ceil(tl.log2(amax * (1.0 / 448.0))) + scaled = quant * tl.reshape(tl.exp2(-exponent), (16, 1)) + fp8 = tl.clamp(scaled, -448.0, 448.0).to(tl.float8e4nv) + tl.store(values + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (512,))) + + max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0 + encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), max_encoded) + tl.store(scales + tl.arange(0, 16), encoded.to(tl.uint8)) + + @triton.jit def _rope_plain_insert_kernel( latent, diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 3c711681aa..88646404f4 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -264,6 +264,8 @@ def _insert_context_kv( attn.padded_heads, attn.eps, block_size, + True, # apply_q_norm (unused: the query is a discarded dummy) + attn.kv_mxfp8, ) elif cache_dtype == torch.bfloat16: swa_3d = swa_cache.view(-1, block_size, attn.head_dim) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index d8c4368401..9c440776ae 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -91,6 +91,17 @@ Bytes, structured as: The first `ue8m0` is the scale for the first 64 `float8_e4m3` values, the second for the next 64, and so on. +For DeepSeek V4.1, each token's KV cache is 528 Bytes, structured as: +- **First 512 bytes:** all 512 dims as `float8_e4m3`. Unlike V4 the 64 RoPE + dims are quantized too, so there is no `bfloat16` part. +- **Last 16 bytes:** Scale factors, containing 16 `ue8m0` values, one per 32 + consecutive `float8_e4m3` values (i.e. MXFP8). + +The V4 and V4.1 records are not laid out token-by-token within a page: a page +holds all its data rows first and all its scale rows after, so a page is +`block_size * bytes_per_token` bytes rounded up to the decode kernel's TMA +stride (576 B for V4, 512 B for V4.1). + In the "nvfp4_ds_mla" format (SM100 only, DeepSeek V3.2 geometry), each token's KV cache is 352 Bytes, structured as: - **First 256 bytes:** 512 `e2m1` NoPE values packed 2/byte (low nibble = diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 62c3921b5e..c4f3cd8b4b 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -79,6 +79,8 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): cache_config: CacheConfig, backend_cls: "type[AttentionBackend] | None" = None, block_size: int = 64, + packed_bytes_per_token: int = 584, + packed_page_alignment: int = 576, ): super().__init__() self.backend_cls = backend_cls or DeepseekSparseSWABackend @@ -88,6 +90,8 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): self.prefix = prefix self.cache_config = cache_config self.dtype = dtype + self.packed_bytes_per_token = packed_bytes_per_token + self.packed_page_alignment = packed_page_alignment compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") @@ -105,8 +109,9 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): self.kv_cache = kv_cache.squeeze(1) def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: - # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous - # bf16/fp8 cache uses the natural element-size page. + # fp8_ds_mla's UE8M0 paged layout rounds its page up to the decode + # kernel's TMA stride; contiguous bf16/fp8 cache uses the natural + # element-size page. uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla" return SlidingWindowMLASpec( block_size=self.block_size, @@ -115,10 +120,11 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): dtype=self.dtype, sliding_window=self.window_size, cache_dtype_str=self.cache_config.cache_dtype, - # DeepseekV4 fp8_ds_mla: 584B per token (448B NoPE + 128B RoPE + 8B scales) - state_content_bytes=584 if uses_fp8_ds_mla_layout else None, - # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). - alignment=576 if uses_fp8_ds_mla_layout else 512, + state_content_bytes=( + self.packed_bytes_per_token if uses_fp8_ds_mla_layout else None + ), + # FlashMLA packing stride; 512B for FlashInfer sparse (#44577). + alignment=self.packed_page_alignment if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype), )