From f8460ef23659afd69046a81529ec2295e4c3f686 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Thu, 15 Oct 2026 13:58:25 -0700 Subject: [PATCH] [Kernel][DSv4.1] Bump FlashMLA to the 2026.09.30 release and its interleaved V4.1 KV cache (#61371) Signed-off-by: juntianl Co-authored-by: Claude Opus 5.5 Co-authored-by: Kimi Code Co-authored-by: Yifan Qiao --- cmake/external_projects/flashmla.cmake | 151 ++++++++++-------- ...deepseek_v4_qnorm_rope_kv_insert_kernel.cu | 27 ++-- tests/kernels/test_compressor_kv_cache.py | 99 +++++++++--- ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 5 +- .../layers/test_mla_short_prefill_indexer.py | 23 ++- .../attention/dsa/sparse_mqa_logits.py | 6 +- .../nvidia/ops/dequant_gather_k_cutedsl.py | 17 +- vllm/models/deepseek_v41/attention.py | 34 ++-- .../deepseek_v41/common/ops/cache_utils.py | 31 ++-- .../common/ops/fused_compress_quant_cache.py | 16 +- vllm/utils/deep_gemm.py | 2 +- .../attention/backends/mla/flashmla_sparse.py | 9 +- vllm/v1/attention/backends/mla/sparse_swa.py | 7 +- 13 files changed, 262 insertions(+), 165 deletions(-) diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake index cfa8243f72..daec85069c 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 0eee43b12f034b657133cf2afca6a72ebb6efccf + GIT_TAG 5dceadbd9eabc9acfc547ef18ec8f7b14b980682 GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" @@ -81,72 +81,83 @@ if(FLASH_MLA_ARCHS) ${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/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/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/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/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 + # Current kernels (DeepSeek V4.1, token-interleaved KV cache scales) + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/smxx/decode/combine/combine.cu # sm100 dense prefill (inference-only: dense backward stays out) - ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu # sm100 sparse prefill - ${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 + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu # sm100 sparse decode - ${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 + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu + + # sm100 fused norm + rope + attn + rope + cast + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_prefill_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_prefill_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_prefill_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_prefill_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu + + # Legacy kernels (upstream keeps them under cuda_kernels/legacy/): + # V4, V3.2 and V3.2-NVFP4, and everything on Hopper + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/smxx/decode/combine/combine.cu + + # sm90 dense decode + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/dense/instantiations/fp16.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/dense/instantiations/bf16.cu + + # sm90 sparse decode + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/sparse/instantiations/v4_persistent_h64.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/sparse/instantiations/v4_persistent_h128.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/sparse/instantiations/v32_persistent_h64.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/decode/sparse/instantiations/v32_persistent_h128.cu + + # sm90 sparse prefill + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/prefill/sparse/instantiations/phase1_k512.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/prefill/sparse/instantiations/phase1_k576.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu + + # sm100 sparse prefill, d_qk = 576 + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu + + # sm100 sparse decode + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/head64/instantiations/v32_h64.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/head64/instantiations/v4_h64.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu + + # sm100 fused norm + rope + attn + rope + cast, V4 decode + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/cuda_kernels/legacy/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu ) set(FlashMLA_Extension_SOURCES @@ -158,17 +169,17 @@ if(FLASH_MLA_ARCHS) set(FlashMLA_INCLUDES ${flashmla_SOURCE_DIR}/csrc - ${flashmla_SOURCE_DIR}/csrc/kerutils/include - ${flashmla_SOURCE_DIR}/csrc/cutlass/include - ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include + ${flashmla_SOURCE_DIR}/csrc/3rdparty/kerutils/include + ${flashmla_SOURCE_DIR}/csrc/3rdparty/cutlass/include + ${flashmla_SOURCE_DIR}/csrc/3rdparty/cutlass/tools/util/include ) set(FlashMLA_Extension_INCLUDES ${flashmla_SOURCE_DIR}/csrc - ${flashmla_SOURCE_DIR}/csrc/kerutils/include + ${flashmla_SOURCE_DIR}/csrc/3rdparty/kerutils/include ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/ - ${flashmla_SOURCE_DIR}/csrc/cutlass/include - ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include + ${flashmla_SOURCE_DIR}/csrc/3rdparty/cutlass/include + ${flashmla_SOURCE_DIR}/csrc/3rdparty/cutlass/tools/util/include ) set_gencode_flags_for_srcs( @@ -202,6 +213,12 @@ if(FLASH_MLA_ARCHS) target_compile_definitions(_flashmla_C PRIVATE USE_CUDA) endif() + # Select the CUDA branch of FlashMLA's params.h and kerutils; kerutils + # detects CUDA by itself in .cu files, so only host files need the flag. + target_compile_definitions(_flashmla_C PRIVATE + FLASH_MLA_IS_BUILD_ON_CUDA + $<$:KERUTILS_IS_BUILD_ON_CUDA>) + define_extension_target( _flashmla_extension_C DESTINATION vllm diff --git a/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 4eb8d7756b..c90c48d1ba 100644 --- a/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -32,9 +32,10 @@ * [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 + * whole row is MXFP8 (fp8 e4m3 with one UE8M0 scale per 32 dims), and each + * token's scales follow its own data + * [i*528, i*528 + 512): token i data, 512 fp8 + * [i*528 + 512, (i+1)*528): token i UE8M0 scales, 16 */ #include "torch_utils.h" @@ -104,6 +105,8 @@ constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576 constexpr int kMxfp8QuantBlock = 32; constexpr int kMxfp8NumQuantBlocks = kHeadDim / kMxfp8QuantBlock; // 16 constexpr int kMxfp8TokenDataBytes = kHeadDim; // 512 +constexpr int kMxfp8BytesPerToken = + kMxfp8TokenDataBytes + kMxfp8NumQuantBlocks; // 528 // 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__) @@ -356,20 +359,20 @@ __device__ __forceinline__ void processDeepseekV4Slot( // 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 * token_data_bytes; + uint8_t* token_fp8_ptr = + block_base + + pos_in_block * (kv_mxfp8 ? kMxfp8BytesPerToken : kTokenDataBytes); uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim; uint8_t* token_scale_ptr = - block_base + - static_cast(cache_block_size) * token_data_bytes + - pos_in_block * scale_bytes; + kv_mxfp8 + ? token_fp8_ptr + kMxfp8TokenDataBytes + : block_base + + static_cast(cache_block_size) * kTokenDataBytes + + pos_in_block * kScaleBytesPerToken; #pragma unroll for (int i = 0; i < kElemsPerLane; i++) { @@ -790,7 +793,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( // backend. Differences from the legacy path: // * No Q head padding — output Q layout matches the input num_heads_q. // * KV is written as a *contiguous* 512-wide row per token (token-strided), -// not the legacy UE8M0 paged layout with a separate scale tail. +// not the V4 paged record with a separate scale tail. // * Q/KV are stored either as bf16 or as per-tensor E4M3 FP8 (one global // scale), selected by the STORE_Q_FP8 / STORE_KV_FP8 template flags. // diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 237c552fc1..23cc4f8bd3 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -360,7 +360,7 @@ def _mxfp8_record_reference( @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]. + in bf16, so a token is [512 fp8][16 UE8M0]. Byte-exact against the reference quantizer, including the rows and page padding the kernel must leave alone. @@ -398,9 +398,8 @@ def test_v41_rope_insert_mxfp8_record(compress_ratio: int): 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 + expected[page, row * 528 : row * 528 + 512] = values + expected[page, row * 528 + 512 : (row + 1) * 528] = scales torch.testing.assert_close(cache_backing, expected, rtol=0, atol=0) @@ -510,14 +509,13 @@ def test_v41_mxfp8_cache_round_trip(): 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], + cache[page, row * 528 : row * 528 + 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], + cache[page, row * 528 + 512 : (row + 1) * 528], (scales.log2() + 127).to(torch.uint8), rtol=0, atol=0, @@ -564,7 +562,8 @@ def _decode_nvfp4_row(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tenso @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") @pytest.mark.parametrize("compress_ratio", [1, 2]) def test_v41_rope_insert_nvfp4_record(compress_ratio: int): - """The 288-byte V4.1 NVFP4 record: 256 B of e2m1 pairs then 32 e4m3 scales. + """The 288-byte V4.1 NVFP4 record: 256 B of e2m1 pairs then 32 e4m3 scales + per token. Scales are byte-exact against the reference (``amax / 6`` clamped to the e4m3 range); values are checked as round-to-nearest on the e2m1 grid, whose @@ -604,17 +603,13 @@ def test_v41_rope_insert_nvfp4_record(compress_ratio: int): amax = rotated.view(32, 16).abs().amax(-1) scale = (amax / 6.0).clamp(2.0**-9, 448.0).to(torch.float8_e4m3fn) - got_scales = cache_backing[page, cache_block * 256 + row * 32 :][:32] + record = cache_backing[page, row * 288 : (row + 1) * 288] + got_scales = record[256:] torch.testing.assert_close(got_scales, scale.view(torch.uint8), rtol=0, atol=0) - decoded = _decode_nvfp4_row( - cache_backing[page, row * 256 : (row + 1) * 256], got_scales - ) + decoded = _decode_nvfp4_row(record[:256], got_scales) step = scale.float().repeat_interleave(16) assert ((decoded - rotated).abs() <= step).all() - untouched[page, row * 256 : (row + 1) * 256] = cache_backing[ - page, row * 256 : (row + 1) * 256 - ] - untouched[page, cache_block * 256 + row * 32 :][:32] = got_scales + untouched[page, row * 288 : (row + 1) * 288] = record # Rows the kernel must not have touched, page padding included. torch.testing.assert_close(cache_backing, untouched, rtol=0, atol=0) @@ -670,10 +665,8 @@ def test_v41_nvfp4_gather_matches_insert(num_reqs, num_tokens, partial): offset=offset, ) - expected = _decode_nvfp4_row( - backing[:, : block_size * 256].reshape(-1, 256), - backing[:, block_size * 256 : block_size * 288].reshape(-1, 32), - ) + records = backing[:, : block_size * 288].reshape(-1, 288) + expected = _decode_nvfp4_row(records[:, :256], records[:, 256:]) def check(shorter_by): assert (storage[:, num_tokens + offset :] == -123).all() @@ -702,6 +695,72 @@ def test_v41_nvfp4_gather_matches_insert(num_reqs, num_tokens, partial): check(1) +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="The V4.1 records are SM100-only", +) +@pytest.mark.parametrize("record", ["swa_fp8", "compressed_fp8", "compressed_fp4"]) +def test_v41_cache_writers_match_flashmla_decode(record: str): + """FlashMLA and vLLM's gathers read back the token a V4.1 writer put mid + page in a padded block, so all agree on the interleaved record layout.""" + import vllm.v1.attention.ops.flashmla as fm + from vllm.models.deepseek_v41.common.ops.cache_utils import ( + dequantize_and_gather_k_cache, + dequantize_and_gather_k_cache_triton, + ) + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + rope_quant_insert, + ) + + device, block_size, topk = "cuda", 64, 64 + slot = block_size + 37 + bytes_per_token = 288 if record == "compressed_fp4" else 528 + torch.manual_seed(0) + kv = torch.randn(1, 512, dtype=torch.bfloat16, device=device) + positions = torch.tensor([slot], device=device) + # cos = 1, sin = 0: RoPE is the identity, so the cache should hold `kv`. + cos_sin = torch.zeros(slot + 1, 64, device=device) + cos_sin[:, :32] = 1 + # Padded like a packed block: a multiple of the record and of 512 B. + cache = torch.zeros(2, block_size + 32, bytes_per_token, dtype=torch.uint8) + cache = cache.to(device)[:, :block_size] + if record == "swa_fp8": + torch.ops._C.fused_deepseek_v4_kv_rope_insert( + kv, cache, positions, positions, cos_sin, block_size, None, True + ) + else: + rope_quant_insert(kv, positions, cos_sin, cache, positions, 1) + + indices = torch.full((1, 1, topk), -1, dtype=torch.int32, device=device) + indices[..., 0] = slot + extra = None + if record == "compressed_fp4": + cache, extra = torch.zeros(2, block_size, 528).to(cache), cache + out, _ = fm.flash_mla_with_kvcache( + q=torch.randn(1, 1, 64, 512, dtype=torch.bfloat16, device=device), + k_cache=cache.unsqueeze(2), + block_table=None, + cache_seqlens=None, + head_dim_v=512, + tile_scheduler_metadata=fm.get_mla_metadata()[0], + is_fp8_kvcache=True, + indices=indices if extra is None else torch.full_like(indices, -1), + extra_k_cache=None if extra is None else extra.unsqueeze(2), + extra_indices_in_kvcache=None if extra is None else indices, + ) + for gather in (dequantize_and_gather_k_cache, dequantize_and_gather_k_cache_triton): + rows = torch.zeros(1, slot + 1, 512, dtype=torch.bfloat16, device=device) + block_table = torch.arange(2, dtype=torch.int32, device=device)[None] + seq_lens = torch.tensor([slot + 1], dtype=torch.int32, device=device) + k_cache = cache if extra is None else extra + gather(rows, k_cache, seq_lens, None, block_table, block_size, offset=0) + torch.testing.assert_close( + out[0, 0], rows[0, slot].expand(64, -1), rtol=0, atol=0 + ) + cos = torch.nn.functional.cosine_similarity(rows[0, slot].float(), kv[0].float(), 0) + assert cos > (0.98 if extra is not None else 0.999) + + @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 c41af1e4b9..a14c889caa 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 @@ -525,9 +525,8 @@ def test_kv_path_mxfp8_record_matches_reference(num_tokens: int, block_size: int 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) + ref_rows = k_cache_ref.view(num_blocks, block_size, V41_HEAD_BYTES) + fused_rows = k_cache_fused.view(num_blocks, block_size, V41_HEAD_BYTES) torch.testing.assert_close( fused_rows[..., :NOPE_DIM], ref_rows[..., :NOPE_DIM], rtol=0, atol=0 ) diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index 7ca4cff404..8c071b6730 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -830,15 +830,19 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): _assert_same_selection(dense, out, ref) -def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): +def _build_fp4_paged_cache( + req_lens: torch.Tensor, page_kv: int, head_dim: int, page_pad: int = 0 +): """MXFP4 paged cache in DeepGEMM's fused layout: per page, all tokens' - fp4 data first, then their packed UE8M0 scales.""" + fp4 data first, then their packed UE8M0 scales; ``page_pad`` bytes pad + each page like the other layers' pages in a packed block.""" device = req_lens.device pages_per = (req_lens + page_kv - 1) // page_kv num_pages, max_pages = int(pages_per.sum()), int(pages_per.max()) + head_bytes = head_dim // 2 + 4 kv_cache = torch.zeros( - num_pages, page_kv, 1, head_dim // 2 + 4, dtype=torch.uint8, device=device - ) + num_pages, page_kv * head_bytes + page_pad, dtype=torch.uint8, device=device + )[:, : page_kv * head_bytes].view(num_pages, page_kv, 1, head_bytes) block_table = torch.zeros( len(req_lens), max_pages, dtype=torch.int32, device=device ) @@ -864,7 +868,10 @@ def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): @pytest.mark.parametrize("next_n", [1, 2]) @pytest.mark.parametrize("capture", [False, True]) -def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool): +@pytest.mark.parametrize("page_pad", [0, 256]) +def test_sparse_mqa_logits_paged_matches_dense_masked( + next_n: int, capture: bool, page_pad: int +): """Paged sparse-MQA decode selects exactly the values ``torch.topk`` picks from the dense bf16 masked logits, incl. MTP rows (next_n > 1 flattened to one query row each), CUDA graph capture/replay and schedule reuse by a @@ -886,8 +893,10 @@ def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool .clamp(min=0) .to(torch.int32) ) - kv_cache, block_table = _build_fp4_paged_cache(req_lens, page_kv, head_dim) - assert kv_cache.stride(0) % 512 == 0 # the paged sparse kernel requires it + kv_cache, block_table = _build_fp4_paged_cache( + req_lens, page_kv, head_dim, page_pad + ) + assert kv_cache.stride(0) % 512 == page_pad # padded pages are not 512B q_fp, q_sf = _quant_fp4( torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) diff --git a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py index aeeb2f63b4..0e3ae398f6 100644 --- a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py +++ b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py @@ -427,7 +427,7 @@ def sparse_mqa_logits_paged_decode( q: [rows, 1, H, D] packed Q values (MXFP4 viewed as int8). q_scale: [rows, 1, H] int32 packed UE8M0 Q scales. kv_cache: 4D [num_pages, page_kv, 1, head_bytes] uint8 fused cache - view; the page stride must be 512B-aligned. + view; the page stride must be 256B-aligned. weights: [rows, H] bf16 per-head weights. context_lens: [rows] int32 per-row (compressed) context lengths. block_table: [rows, P] int32 per-row page table, ``stride(-1) == 1``. @@ -451,8 +451,8 @@ def sparse_mqa_logits_paged_decode( rows = q.shape[0] page_kv = kv_cache.shape[1] assert weights.dtype == torch.bfloat16, weights.dtype - assert kv_cache.stride(0) % 512 == 0, ( - "paged cache page stride must be 512B-aligned: " + assert kv_cache.stride(0) % 256 == 0, ( + "paged cache page stride must be 256B-aligned: " f"shape={tuple(kv_cache.shape)} stride={kv_cache.stride()}" ) assert page_kv % sparse_block_kv == 0, (page_kv, sparse_block_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 d97c54a011..a714aa96c5 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 @@ -52,6 +52,12 @@ class DequantGatherKCacheKernel( data_dim = fp8_dim + bf16_dim * 2 scale_bytes = head_dim // group_size block_size = compile_key.block_size + # V4.1 stores each token's scales right after its data; V4 puts all of + # a block's data rows ahead of its scale rows. + interleaved = bf16_dim == 0 + data_stride = data_dim + scale_bytes if interleaved else data_dim + scale_offset = data_dim if interleaved else block_size * data_dim + scale_stride = data_stride if interleaved else scale_bytes @cute.jit def load_g2s( @@ -274,23 +280,22 @@ class DequantGatherKCacheKernel( offset: Int32, stream: CUstream, ): - # 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, scale_bytes]. + # Split k_cache into k_data [block_size, fp8_dim + bf16_dim * 2] and + # k_scale [block_size, scale_bytes] views of each block. k_data = cute.make_tensor( k_cache.iterator, layout=cute.make_layout( (k_cache.shape[0], block_size, data_dim), - stride=(k_cache.stride[0], data_dim, 1), + stride=(k_cache.stride[0], data_stride, 1), ), ) # k_data_slice: [num_blocks, block_size, (16, data_dim/16)] # s_kdata_slice: [(4, data_dim/16), num_stages] k_scale = cute.make_tensor( - k_cache.iterator + (block_size * data_dim), + k_cache.iterator + scale_offset, layout=cute.make_layout( (k_cache.shape[0], block_size, scale_bytes), - stride=(k_cache.stride[0], scale_bytes, 1), + stride=(k_cache.stride[0], scale_stride, 1), ), ) diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index b68879c5c8..78114f0161 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -122,10 +122,10 @@ 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). +# UE8M0 scale per 32 dims, each token's 512 data bytes followed by its 16 scale +# bytes -- 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) @@ -515,8 +515,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): 288 if self.kv_cache_dtype == "nvfp4_ds_mla" else self.swa_bytes_per_token ) # One alignment for every page in the block: the block stride is their - # sum, and 512 satisfies both TMA strides in play (512 for the V4.1 - # fp8 record, 256 for NVFP4). + # sum. V4 needs its 576 B TMA stride; V4.1's is on the SWA block + # stride, and 512 is free (its pages are already 512 B multiples). self.kv_page_alignment = 512 if self.kv_mxfp8 else 576 if self.kv_cache_dtype == "nvfp4_ds_mla" and not self.kv_mxfp8: raise ValueError( @@ -562,6 +562,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): block_size=32, packed_bytes_per_token=self.swa_bytes_per_token, packed_page_alignment=self.kv_page_alignment, + # FlashMLA reads the interleaved V4.1 records as rows of the whole + # cache, so the packed block has to hold whole rows of both. + packed_block_stride_alignment=( + math.lcm(self.swa_bytes_per_token, self.compressed_bytes_per_token) + if self.kv_mxfp8 + else None + ), bounded_replay=swa_bounded_replay, ) @@ -1222,15 +1229,18 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): tokens_per_state=self.compress_ratio, # 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, + # has to carry the same padding: the packed KV pages' alignment, # else 512B for FlashInfer sparse (#44577). alignment=page_alignment, - # DeepGEMM's paged sparse MQA logits address pages by the block - # stride, which packs every layer's page in this layout, and need - # it 512B-aligned; the main KV record read from the same blocks - # needs its own TMA stride, so keep both. + # DeepGEMM's paged sparse logits need a 256B-aligned block stride; + # V4 records keep their TMA stride (V4.1's is on the SWA spec). block_stride_alignment=( - math.lcm(512, page_alignment) + math.lcm( + 256, + 1 + if uses_fp8_ds_mla_layout and _use_v41_mxfp8_kv_record() + else page_alignment, + ) if self.sparse_logits and not self.rocm_mxfp4 else None ), diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 000fb03195..b7c8a42fd2 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -36,16 +36,17 @@ from vllm.triton_utils import tl, triton from vllm.utils.import_utils import has_cutedsl from vllm.utils.math_utils import cdiv, 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 +# Per-token byte width of the paged records; ``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. +# dims each plus a pad byte. A page puts its whole data region +# ahead of its whole scale region. # 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.1 NVFP4 (288 B): 512 e2m1 values packed two per byte (even element in # the low nibble), then 32 e4m3 scales of 16 dims each. # FlashMLA's ``ModelType::V41_FP4``, compressed cache only. +# The V4.1 records are interleaved: each token's scales follow its own data. V4_BYTES_PER_TOKEN = 584 V41_BYTES_PER_TOKEN = 528 V41_QUANT_BLOCK = 32 @@ -212,10 +213,8 @@ def _quantize_and_insert_k_mxfp8_kernel( # 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 - ) + token_data_ptr = cache_block_ptr + pos_in_block * (head_dim + scale_dim) + token_scale_ptr = token_data_ptr + head_dim d = tl.arange(0, head_dim) x = tl.load(k_ptr + pid.to(tl.int64) * head_dim + d).to(tl.float32) @@ -568,10 +567,8 @@ def _dequantize_and_gather_k_mxfp8_kernel( # 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 - ) + token_data_ptr = cache_block_ptr + pos_in_block * (head_dim + scale_dim) + token_scale_ptr = token_data_ptr + head_dim x_uint8 = tl.load(token_data_ptr + d) if use_fnuz: @@ -629,10 +626,9 @@ def _dequantize_and_gather_k_nvfp4_kernel( block_table_ptr + batch_idx * max_blocks_per_seq + block_in_seq ) page = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride + token_ptr = page + pos_in_block * (packed_bytes + scale_dim) - packed = tl.load( - page + pos_in_block * packed_bytes + tl.arange(0, packed_bytes) - ) + packed = tl.load(token_ptr + tl.arange(0, packed_bytes)) # Even element in the low nibble, odd in the high nibble. codes = tl.interleave((packed & 0xF).to(tl.int32), (packed >> 4).to(tl.int32)) # e2m1 magnitudes: 0, 0.5, 1, 1.5, 2, 3, 4, 6. @@ -642,12 +638,7 @@ def _dequantize_and_gather_k_nvfp4_kernel( mag = tl.where(mag_code < 2, m * 0.5, (1.0 + m * 0.5) * tl.exp2(e - 1.0)) vals = tl.where(codes >= 8, -mag, mag) - sf = tl.load( - page - + cache_block_size * packed_bytes - + pos_in_block * scale_dim - + tl.arange(0, scale_dim) - ) + sf = tl.load(token_ptr + packed_bytes + tl.arange(0, scale_dim)) scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32) tiles = tl.reshape(vals, (scale_dim, quant_block)) dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,)) 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 2ba9dbea3d..92ade5fef9 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 @@ -239,7 +239,8 @@ def rope_quant_insert( UE8M0 scale bytes, including one zero padding scale), 528 B for V4.1 (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 scales of 32 dims each), or 288 B for V4.1 NVFP4 (256 bytes of e2m1 pairs then 32 - e4m3 scales of 16 dims each), which only the compressed cache uses. + e4m3 scales of 16 dims each), which only the compressed cache uses. The + two V4.1 records keep each token's scales right after its own values. ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``. """ @@ -390,8 +391,8 @@ def _rope_quant_insert_mxfp8_kernel( 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 + values = page + (slot % CACHE_BLOCK) * 528 + scales = values + 512 quant = tl.reshape(rotated, (16, 32)) amax = tl.maximum(tl.max(tl.abs(quant), 1), 1e-4) @@ -453,12 +454,9 @@ def _rope_quant_insert_nvfp4_kernel( lo, hi = tl.split(tl.reshape(tl.reshape(scaled, (512,)), (256, 2))) page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE - tl.store( - page + (slot % CACHE_BLOCK) * 256 + tl.arange(0, 256), - _fp32x2_to_fp4x2(lo, hi), - ) - scales = page + CACHE_BLOCK * 256 + (slot % CACHE_BLOCK) * 32 - tl.store(scales + tl.arange(0, 32), scale.to(tl.uint8, bitcast=True)) + values = page + (slot % CACHE_BLOCK) * 288 + tl.store(values + tl.arange(0, 256), _fp32x2_to_fp4x2(lo, hi)) + tl.store(values + 256 + tl.arange(0, 32), scale.to(tl.uint8, bitcast=True)) @triton.jit diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 5badccf6b5..b995d19d89 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -948,7 +948,7 @@ def fp8_fp4_paged_sparse_mqa_logits( Args: q: ``(q_values, q_scale)``; q_values is [num_q_tokens, 1, H, D]. kv_cache: Fused paged cache [num_pages, page_kv, 1, D' + 4] uint8, - page stride 512B-aligned. + page stride 256B-aligned. weights: [num_q_tokens, H] ``torch.bfloat16``. metadata: From `get_paged_sparse_mqa_logits_metadata`. num_max_sparse_blocks: Candidate blocks per row. diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index fd4b562f53..bbbb111d8b 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -103,10 +103,11 @@ For DeepSeek V4.1, each token's KV cache is 528 Bytes, structured as: - **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). +The V4 record is 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 576 B +TMA stride. The V4.1 records are interleaved: each token's scales directly +follow its own data, and the decode kernel's TMA stride is the record width. In the "nvfp4_ds_mla" format (SM100 only, DeepSeek V3.2 geometry), each token's KV cache is 352 Bytes, structured as: diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 65df2fcee4..15daace007 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -85,6 +85,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): block_size: int = 64, packed_bytes_per_token: int = 584, packed_page_alignment: int = 576, + packed_block_stride_alignment: int | None = None, bounded_replay: bool = False, ): super().__init__() @@ -99,6 +100,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): self.dtype = dtype self.packed_bytes_per_token = packed_bytes_per_token self.packed_page_alignment = packed_page_alignment + self.packed_block_stride_alignment = packed_block_stride_alignment compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") @@ -135,8 +137,11 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): state_content_bytes=( self.packed_bytes_per_token if uses_fp8_ds_mla_layout else None ), - # FlashMLA packing stride; 512B for FlashInfer sparse (#44577). + # Packed page alignment; 512B for FlashInfer sparse (#44577). alignment=self.packed_page_alignment if uses_fp8_ds_mla_layout else 512, + block_stride_alignment=( + self.packed_block_stride_alignment if uses_fp8_ds_mla_layout else None + ), model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype), )