[Kernel][DSv4.1] Bump FlashMLA to the 2026.09.30 release and its interleaved V4.1 KV cache (#61371)
Signed-off-by: juntianl <juntianl@inferact.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Kimi Code <noreply@moonshot.ai> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
This commit is contained in:
Juntian Liu
co-authored by
Claude Opus 5.5
Kimi Code
Yifan Qiao
Claude Opus 5.5
Kimi Code
Yifan Qiao
parent
33212db2c2
commit
f8460ef236
@@ -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
|
||||
$<$<COMPILE_LANGUAGE:CXX>:KERUTILS_IS_BUILD_ON_CUDA>)
|
||||
|
||||
define_extension_target(
|
||||
_flashmla_extension_C
|
||||
DESTINATION vllm
|
||||
|
||||
@@ -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<int64_t>(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<int64_t>(cache_block_size) * token_data_bytes +
|
||||
pos_in_block * scale_bytes;
|
||||
kv_mxfp8
|
||||
? token_fp8_ptr + kMxfp8TokenDataBytes
|
||||
: block_base +
|
||||
static_cast<int64_t>(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.
|
||||
//
|
||||
|
||||
@@ -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])
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
),
|
||||
|
||||
@@ -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,))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user