[Model][DSv4.1] Store the whole KV in MXFP8 (FlashMLA V4.1 record) (#56893)
Signed-off-by: Yongye Zhu <yongye@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Yongye Zhu <yongye@inferact.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Yongye Zhu committed
Sep 14, 2026 at 20:12 UTC
e6eb0d120c480aa5116aafc2d39f19c8ebcf3181
14 files changed
+709
-110
cmake/external_projects/flashmla.cmake
+59
-26
@@ -19,7 +19,7 @@ else()
19
FetchContent_Declare(
20
flashmla
21
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
22
- GIT_TAG 6bc49418c5ead572ff0339191ddf3b155749e183
22
+ GIT_TAG 0eee43b12f034b657133cf2afca6a72ebb6efccf
23
GIT_PROGRESS TRUE
24
CONFIGURE_COMMAND ""
25
BUILD_COMMAND ""
@@ -76,43 +76,77 @@ if(FLASH_MLA_ARCHS)
76
77
set(FlashMLA_SOURCES
78
${flashmla_SOURCE_DIR}/csrc/api/api.cpp
79
+ ${flashmla_SOURCE_DIR}/csrc/api/sparse_prefill.cpp
80
+ ${flashmla_SOURCE_DIR}/csrc/api/sparse_decode.cpp
81
+ ${flashmla_SOURCE_DIR}/csrc/api/dense_decode.cpp
82
+ ${flashmla_SOURCE_DIR}/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp
83
84
# Misc kernels for decoding
81
- ${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
82
- ${flashmla_SOURCE_DIR}/csrc/smxx/decode/combine/combine.cu
85
+ ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu
86
+ ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/combine/combine.cu
87
88
# sm90 dense decode
85
- ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/fp16.cu
86
- ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/bf16.cu
89
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/fp16.cu
90
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/bf16.cu
91
92
# sm90 sparse decode
89
- ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu
90
- ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu
91
- ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu
92
- ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu
93
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu
94
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu
95
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu
96
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu
97
98
# sm90 sparse prefill
95
- ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu
96
- ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu
97
- ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
98
- ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu
99
- ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
99
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu
100
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu
101
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu
102
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
103
101
- # sm100 dense prefill & backward
102
- ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
104
+ # sm100 dense prefill (inference-only: dense backward stays out)
105
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
106
107
# sm100 sparse prefill
105
- ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu
106
- ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu
107
- ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
108
- ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
109
- ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu
108
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu
109
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu
110
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
111
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
112
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu
113
+
114
+ # sm100 fused norm + rope + attn + rope + cast
115
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu
116
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu
117
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu
118
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu
119
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu
120
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu
121
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu
122
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu
123
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu
124
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu
125
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu
126
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu
127
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu
128
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu
129
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu
130
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu
131
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu
132
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu
133
134
# sm100 sparse decode
112
- ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu
113
- ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu
114
- ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32_nvfp4_fp8rope.cu
115
- ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
135
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu
136
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu
137
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu
138
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu
139
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu
140
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu
141
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu
142
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu
143
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu
144
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
145
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu
146
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu
147
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu
148
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu
149
+ ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu
150
)
151
152
set(FlashMLA_Extension_SOURCES
@@ -125,7 +159,6 @@ if(FLASH_MLA_ARCHS)
159
set(FlashMLA_INCLUDES
160
${flashmla_SOURCE_DIR}/csrc
161
${flashmla_SOURCE_DIR}/csrc/kerutils/include
128
- ${flashmla_SOURCE_DIR}/csrc/sm90
162
${flashmla_SOURCE_DIR}/csrc/cutlass/include
163
${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include
164
)
csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+64
-30
@@ -23,9 +23,15 @@
23
* cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin
24
* second along last dim; each half is rope_dim/2 = 32 values)
25
*
26
- * Cache layout per paged-cache block (block_size tokens):
27
- * [0, bs*576): token data, 448 fp8 + 128 bf16 each
28
- * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token
26
+ * Cache layout per paged-cache block (block_size tokens), selected by the
27
+ * `kv_mxfp8` argument:
28
+ * V4 (584 B/token, kv_mxfp8=false)
29
+ * [0, bs*576): token data, 448 fp8 + 128 bf16 each
30
+ * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token
31
+ * V4.1 (528 B/token, kv_mxfp8=true) -- RoPE is quantized as well, so the
32
+ * whole row is MXFP8 (fp8 e4m3 with one UE8M0 scale per 32 dims)
33
+ * [0, bs*512): token data, 512 fp8 each
34
+ * [bs*512, bs*512 + bs*16): UE8M0 scales, 16 per token
35
*/
36
37
#include "torch_utils.h"
@@ -91,6 +97,10 @@ constexpr int kQuantBlock = 64;
97
constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock; // 7
98
constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1; // 8 (7 real + 1 pad)
99
constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576
100
+// V4.1 MXFP8 record: every dim quantized, one UE8M0 scale per 32 dims.
101
+constexpr int kMxfp8QuantBlock = 32;
102
+constexpr int kMxfp8NumQuantBlocks = kHeadDim / kMxfp8QuantBlock; // 16
103
+constexpr int kMxfp8TokenDataBytes = kHeadDim; // 512
104
// FNUZ on gfx942 / OCP elsewhere. FNUZ uses 224.0 (not the dtype's raw
105
// 240.0) to match the rest of vLLM's FNUZ pipeline.
106
#if defined(USE_ROCM) && defined(__gfx942__)
@@ -141,6 +151,11 @@ __device__ __forceinline__ uint4 packFp8E4M3x16(float const* values,
151
// ────────────────────────────────────────────────────────────────────────────
152
// Small inline helpers
153
// ────────────────────────────────────────────────────────────────────────────
154
+__device__ __forceinline__ float warp2MaxAbs(float val) {
155
+ // Reduce absolute max across 2 consecutive lanes (32 dims = one MXFP8 block).
156
+ return fmaxf(val, __shfl_xor_sync(FINAL_MASK, val, 1));
157
+}
158
+
159
__device__ __forceinline__ float warp4MaxAbs(float val) {
160
// Reduce absolute max across 4 consecutive lanes (lane id & 3 group).
161
float peer = __shfl_xor_sync(FINAL_MASK, val, 1);
@@ -176,6 +191,10 @@ __device__ __forceinline__ float warpSum(float val) {
191
// slot == kNumHeadsQPadded → KV (RoPE + UE8M0 quant
192
// + paged-cache
193
// insert)
194
+//
195
+// `kv_mxfp8` is grid-uniform and only steers the KV branch, so it stays a
196
+// runtime argument rather than doubling every (head count, q-norm)
197
+// instantiation of a kernel whose hot path is the Q branch.
198
template <typename scalar_t_in, int kNumHeadsQPadded, bool APPLY_Q_NORM>
199
__device__ __forceinline__ void processDeepseekV4Slot(
200
uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx,
@@ -184,7 +203,7 @@ __device__ __forceinline__ void processDeepseekV4Slot(
203
uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping,
204
int64_t const* __restrict__ position_ids,
205
float const* __restrict__ cos_sin_cache, int const cache_block_size,
187
- int const kv_block_stride) {
206
+ int const kv_block_stride, bool const kv_mxfp8) {
207
using Converter = vllm::_typeConvert<scalar_t_in>;
208
bool const isKV = (slotIdx == kNumHeadsQPadded);
209
bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
@@ -302,19 +321,24 @@ __device__ __forceinline__ void processDeepseekV4Slot(
321
*reinterpret_cast<uint4*>(dst) = out0;
322
*reinterpret_cast<uint4*>(dst + 8) = out1;
323
} else {
305
- // ── KV: FP8 quant on NoPE + bf16 store on RoPE + cache insert.
324
+ // ── KV: FP8 quant + cache insert. V4 keeps the RoPE dims in bf16;
325
+ // V4.1 (kv_mxfp8) quantizes them too.
326
int64_t const slot_id = slot_mapping[tokenIdx];
327
if (slot_id >= 0) {
328
+ int const token_data_bytes =
329
+ kv_mxfp8 ? kMxfp8TokenDataBytes : kTokenDataBytes;
330
+ int const scale_bytes =
331
+ kv_mxfp8 ? kMxfp8NumQuantBlocks : kScaleBytesPerToken;
332
int64_t const block_idx = slot_id / cache_block_size;
333
int64_t const pos_in_block = slot_id % cache_block_size;
334
uint8_t* block_base =
335
k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
312
- uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
336
+ uint8_t* token_fp8_ptr = block_base + pos_in_block * token_data_bytes;
337
uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
338
uint8_t* token_scale_ptr =
339
block_base +
316
- static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
317
- pos_in_block * kScaleBytesPerToken;
340
+ static_cast<int64_t>(cache_block_size) * token_data_bytes +
341
+ pos_in_block * scale_bytes;
342
343
#pragma unroll
344
for (int i = 0; i < kElemsPerLane; i++) {
@@ -326,11 +350,13 @@ __device__ __forceinline__ void processDeepseekV4Slot(
350
for (int i = 0; i < kElemsPerLane; i++) {
351
local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
352
}
329
- float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
353
+ float const absmax = fmaxf(
354
+ kv_mxfp8 ? warp2MaxAbs(local_absmax) : warp4MaxAbs(local_absmax),
355
+ 1e-4f);
356
float const exponent = ceilf(log2f(absmax / kFp8Max));
357
float const inv_scale = exp2f(-exponent);
358
333
- if (!is_rope_lane) {
359
+ if (kv_mxfp8 || !is_rope_lane) {
360
uint8_t out_bytes[kElemsPerLane];
361
#pragma unroll
362
for (int i = 0; i < kElemsPerLane; i++) {
@@ -347,13 +373,19 @@ __device__ __forceinline__ void processDeepseekV4Slot(
373
*reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
374
*reinterpret_cast<uint4 const*>(out_bytes);
375
350
- if ((laneId & 3) == 0) {
351
- int const q_block_idx = laneId >> 2;
352
- float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
353
- token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
354
- }
355
- if (laneId == 0) {
356
- token_scale_ptr[kNumQuantBlocks] = 0;
376
+ float const encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
377
+ if (kv_mxfp8) {
378
+ // One scale per 32 dims = one per lane pair.
379
+ if ((laneId & 1) == 0) {
380
+ token_scale_ptr[laneId >> 1] = static_cast<uint8_t>(encoded);
381
+ }
382
+ } else {
383
+ if ((laneId & 3) == 0) {
384
+ token_scale_ptr[laneId >> 2] = static_cast<uint8_t>(encoded);
385
+ }
386
+ if (laneId == 0) {
387
+ token_scale_ptr[kNumQuantBlocks] = 0;
388
+ }
389
}
390
} else {
391
uint4 out0, out1;
@@ -421,7 +453,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
453
int const num_tokens_insert, // = slot_mapping.size(0), ≤ num_tokens_full
454
int const num_heads_q, // live Q heads (input layout)
455
int const cache_block_size, // tokens per paged-cache block
424
- int const kv_block_stride) { // bytes per paged-cache block
456
+ int const kv_block_stride, // bytes per paged-cache block
457
+ bool const kv_mxfp8) { // V4.1 all-dims MXFP8 record
458
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
459
// BF16 _typeConvert specialization is unavailable on pre-Ampere. The
460
// DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a
@@ -477,7 +510,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
510
processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>(
511
v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out,
512
k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size,
480
- kv_block_stride);
513
+ kv_block_stride, kv_mxfp8);
514
515
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
516
cudaTriggerProgrammaticLaunchCompletion();
@@ -506,7 +539,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
539
float const* __restrict__ cos_sin_cache, float const eps,
540
int const num_tokens_full, int const num_tokens_insert,
541
int const num_heads_q, int const cache_block_size,
509
- int const kv_block_stride) {
542
+ int const kv_block_stride, bool const kv_mxfp8) {
543
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
544
if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
545
return;
@@ -564,7 +597,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
597
processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>(
598
v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
599
num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids,
567
- cos_sin_cache, cache_block_size, kv_block_stride);
600
+ cos_sin_cache, cache_block_size, kv_block_stride, kv_mxfp8);
601
602
// ── Buffer rotation: hand the prefetched LDGs to the next iter.
603
v0_curr = v0_next;
@@ -590,7 +623,7 @@ static void launchFusedDeepseekV4Templated(
623
uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids,
624
float const* cos_sin_cache, float const eps, int const num_tokens_full,
625
int const num_tokens_insert, int const num_heads_q,
593
- int const cache_block_size, int const kv_block_stride,
626
+ int const cache_block_size, int const kv_block_stride, bool const kv_mxfp8,
627
cudaStream_t stream) {
628
constexpr int kBlockSize = 256;
629
constexpr int kWarpsPerBlock = kBlockSize / 32;
@@ -630,7 +663,8 @@ static void launchFusedDeepseekV4Templated(
663
scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>,
664
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
665
cos_sin_cache, eps, num_tokens_full, num_tokens_insert,
633
- num_heads_q, cache_block_size, kv_block_stride);
666
+ num_heads_q, cache_block_size, kv_block_stride,
667
+ kv_mxfp8);
668
} else {
669
config.gridDim = dim3(num_tokens_full);
670
cudaLaunchKernelEx(
@@ -639,7 +673,7 @@ static void launchFusedDeepseekV4Templated(
673
scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>,
674
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
675
eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
642
- kv_block_stride);
676
+ kv_block_stride, kv_mxfp8);
677
}
678
#else
679
// ROCm: use standard kernel launch syntax (no PDL/stream serialization)
@@ -649,7 +683,7 @@ static void launchFusedDeepseekV4Templated(
683
<<<grid, kBlockSize, 0, stream>>>(
684
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
685
cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q,
652
- cache_block_size, kv_block_stride);
686
+ cache_block_size, kv_block_stride, kv_mxfp8);
687
#endif
688
}
689
@@ -663,19 +697,19 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
697
int const num_tokens_full, int const num_tokens_insert,
698
int const num_heads_q, int const num_heads_q_padded,
699
int const cache_block_size, int const kv_block_stride,
666
- bool const apply_q_norm, cudaStream_t stream) {
700
+ bool const apply_q_norm, bool const kv_mxfp8, cudaStream_t stream) {
701
#define DISPATCH(N) \
702
case N: \
703
if (apply_q_norm) { \
704
launchFusedDeepseekV4Templated<scalar_t_in, N, true>( \
705
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \
706
cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \
673
- num_heads_q, cache_block_size, kv_block_stride, stream); \
707
+ num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \
708
} else { \
709
launchFusedDeepseekV4Templated<scalar_t_in, N, false>( \
710
q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \
711
cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \
678
- num_heads_q, cache_block_size, kv_block_stride, stream); \
712
+ num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \
713
} \
714
return;
715
@@ -967,7 +1001,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
1001
torch::stable::Tensor const& position_ids, // [N] int64
1002
torch::stable::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16
1003
int64_t q_head_padded, // padded Q head count for output
970
- double eps, int64_t cache_block_size, bool apply_q_norm) {
1004
+ double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8) {
1005
STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(),
1006
"q_in must be contiguous CUDA");
1007
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(
1070
cos_sin_cache.const_data_ptr<float>(), static_cast<float>(eps),
1071
num_tokens_full, num_tokens_insert, num_heads_q,
1072
num_heads_q_padded, cache_block_size_i, kv_block_stride,
1039
- apply_q_norm, stream);
1073
+ apply_q_norm, kv_mxfp8, stream);
1074
});
1075
return q_out;
1076
}
csrc/libtorch_stable/ops.h
+1
-1
@@ -264,7 +264,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
264
torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
265
torch::stable::Tensor const& position_ids,
266
torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
267
- double eps, int64_t cache_block_size, bool apply_q_norm);
267
+ double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8);
268
269
void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
270
torch::stable::Tensor& q, torch::stable::Tensor const& kv,
csrc/libtorch_stable/torch_bindings.cpp
+1
-1
@@ -430,7 +430,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
430
"Tensor q_in, Tensor kv, Tensor! k_cache, "
431
"Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
432
"int q_head_padded, float eps, int cache_block_size, "
433
- "bool apply_q_norm=True) -> Tensor");
433
+ "bool apply_q_norm=True, bool kv_mxfp8=False) -> Tensor");
434
435
// FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
436
// FP8 tensor, and KV into a contiguous 512-wide token-strided cache.
tests/kernels/test_compressor_kv_cache.py
+130
@@ -215,6 +215,136 @@ def test_v41_fused_save_compress_and_insert(
215
torch.testing.assert_close(cache_backing, expected_cache, rtol=0, atol=0)
216
217
218
+def _mxfp8_record_reference(
219
+ latent_row: torch.Tensor, cos_sin_row: torch.Tensor
220
+) -> tuple[torch.Tensor, torch.Tensor]:
221
+ """Reference V4.1 record for one token: 512 fp8 bytes + 16 UE8M0 scales.
222
+
223
+ Mirrors FlashMLA's ``KVCacheLayout.V41_FP8Sparse`` quantizer: rotate the
224
+ RoPE tail first, then scale every 32-dim tile, RoPE tiles included.
225
+ """
226
+ row = latent_row.float()
227
+ c, s = cos_sin_row.chunk(2)
228
+ rotated = row.clone()
229
+ even, odd = row[448::2], row[449::2]
230
+ rotated[448::2] = even * c - odd * s
231
+ rotated[449::2] = odd * c + even * s
232
+ quantized, scales = _ue8m0_reference(rotated, 32, 448.0)
233
+ return quantized.view(torch.uint8), (scales.log2() + 127).to(torch.uint8)
234
+
235
+
236
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
237
+@pytest.mark.parametrize("compress_ratio", [1, 2])
238
+def test_v41_rope_insert_mxfp8_record(compress_ratio: int):
239
+ """The 528-byte V4.1 record quantizes the RoPE dims instead of keeping them
240
+ in bf16, so a page is [block_size x 512 fp8][block_size x 16 UE8M0].
241
+
242
+ Byte-exact against the reference quantizer, including the rows and page
243
+ padding the kernel must leave alone.
244
+ """
245
+ from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
246
+ rope_quant_insert,
247
+ )
248
+
249
+ torch.manual_seed(7)
250
+ device = "cuda"
251
+ cache_block = 64
252
+ cache_stride = math.ceil(cache_block * 528 / 512) * 512
253
+ num_tokens = 12
254
+
255
+ positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
256
+ latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device)
257
+ angles = torch.randn(64, 32, device=device)
258
+ cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
259
+
260
+ cache_backing = torch.full((2, cache_stride), 165, dtype=torch.uint8, device=device)
261
+ cache = cache_backing.as_strided((2, cache_block, 528), (cache_stride, 528, 1))
262
+ # Leave one token unmapped so the kernel's negative-slot skip is covered.
263
+ slots = torch.arange(num_tokens, dtype=torch.int64, device=device)
264
+ slots[3] = -1
265
+
266
+ rope_quant_insert(latent, positions, cos_sin, cache, slots, compress_ratio)
267
+
268
+ expected = torch.full_like(cache_backing, 165)
269
+ for t in range(num_tokens):
270
+ slot = slots[t].item()
271
+ pos = positions[t].item()
272
+ if slot < 0 or (pos + 1) % compress_ratio:
273
+ continue
274
+ page, row = divmod(slot, cache_block)
275
+ values, scales = _mxfp8_record_reference(
276
+ latent[t], cos_sin[pos // compress_ratio * compress_ratio]
277
+ )
278
+ expected[page, row * 512 : (row + 1) * 512] = values
279
+ scale_offset = cache_block * 512 + row * 16
280
+ expected[page, scale_offset : scale_offset + 16] = scales
281
+
282
+ torch.testing.assert_close(cache_backing, expected, rtol=0, atol=0)
283
+
284
+
285
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
286
+def test_v41_mxfp8_cache_round_trip():
287
+ """quantize_and_insert -> dequantize_and_gather recovers the V4.1 record.
288
+
289
+ The RoPE dims now go through fp8, so they carry quantization error too --
290
+ the check is that every dim survives within one MXFP8 tile step, and that
291
+ the bytes on the way match the reference quantizer.
292
+ """
293
+ from vllm.models.deepseek_v41.common.ops import (
294
+ dequantize_and_gather_k_cache,
295
+ quantize_and_insert_k_cache,
296
+ )
297
+
298
+ torch.manual_seed(11)
299
+ device = "cuda"
300
+ block_size = 64
301
+ num_tokens = 70
302
+ num_blocks = 4
303
+ page_bytes = math.ceil(block_size * 528 / 512) * 512
304
+
305
+ k = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device)
306
+ cache = torch.zeros(num_blocks, page_bytes, dtype=torch.uint8, device=device)
307
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
308
+ quantize_and_insert_k_cache(
309
+ k, cache, slot_mapping, block_size=block_size, bytes_per_token=528
310
+ )
311
+
312
+ for t in (0, 1, block_size, num_tokens - 1):
313
+ values, scales = _ue8m0_reference(k[t].float(), 32, 448.0)
314
+ page, row = divmod(t, block_size)
315
+ torch.testing.assert_close(
316
+ cache[page, row * 512 : (row + 1) * 512],
317
+ values.view(torch.uint8),
318
+ rtol=0,
319
+ atol=0,
320
+ )
321
+ scale_offset = block_size * 512 + row * 16
322
+ torch.testing.assert_close(
323
+ cache[page, scale_offset : scale_offset + 16],
324
+ (scales.log2() + 127).to(torch.uint8),
325
+ rtol=0,
326
+ atol=0,
327
+ )
328
+
329
+ out = torch.zeros(1, num_tokens, 512, dtype=torch.bfloat16, device=device)
330
+ dequantize_and_gather_k_cache(
331
+ out,
332
+ cache.view(num_blocks, block_size, 528),
333
+ seq_lens=torch.tensor([num_tokens], dtype=torch.int32, device=device),
334
+ gather_lens=None,
335
+ block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view(
336
+ 1, -1
337
+ ),
338
+ block_size=block_size,
339
+ offset=0,
340
+ )
341
+ # Half an e4m3 ULP at the top of the range (16 units) bounds the error.
342
+ tile_amax = k.float().abs().view(num_tokens, 16, 32).amax(-1).clamp(min=1e-4)
343
+ tolerance = 16.0 * (tile_amax / 448.0).log2().ceil().exp2()
344
+ error = (out[0].float() - k.float()).abs()
345
+ assert (error <= tolerance.repeat_interleave(32, dim=-1)).all()
346
+
347
+
348
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
349
@pytest.mark.parametrize("compress_ratio", [1, 2])
350
@pytest.mark.parametrize("store_fp8", [False, True])
tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+98
@@ -424,6 +424,104 @@ def test_kv_path_matches_reference(num_tokens: int, block_size: int):
424
)
425
426
427
+# ── Test 2c: DeepSeek V4.1 MXFP8 record (kv_mxfp8=True) ──────────────────────
428
+
429
+V41_HEAD_BYTES = HEAD_DIM + HEAD_DIM // 32 # 512 fp8 + 16 UE8M0 scales = 528
430
+
431
+
432
+@pytest.mark.parametrize("num_tokens", [1, 17, 64])
433
+@pytest.mark.parametrize("block_size", [32, 64])
434
+def test_kv_path_mxfp8_record_matches_reference(num_tokens: int, block_size: int):
435
+ """``kv_mxfp8=True`` writes DeepSeek's V4.1 record: all 512 dims as fp8 e4m3
436
+ with one UE8M0 scale per 32 dims, no bf16 RoPE tail.
437
+
438
+ Both sides quantize the same bf16-rounded rotated row, so the bytes must
439
+ match the Triton encoder exactly apart from the documented 1-ULP bf16
440
+ rounding split on the rotation itself.
441
+ """
442
+ from vllm.models.deepseek_v41.common.ops import (
443
+ dequantize_and_gather_k_cache as v41_dequantize_and_gather_k_cache,
444
+ )
445
+ from vllm.models.deepseek_v41.common.ops import (
446
+ quantize_and_insert_k_cache as v41_quantize_and_insert_k_cache,
447
+ )
448
+
449
+ torch.manual_seed(3)
450
+ device = "cuda"
451
+ dtype = torch.bfloat16
452
+ max_pos = 4096
453
+
454
+ kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device)
455
+ positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
456
+ cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
457
+
458
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
459
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
460
+
461
+ k_cache_ref = torch.zeros(
462
+ num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device
463
+ )
464
+ v41_quantize_and_insert_k_cache(
465
+ apply_rope_gptj_last_k(kv, positions, cos_sin_cache),
466
+ k_cache_ref,
467
+ slot_mapping,
468
+ block_size=block_size,
469
+ use_fnuz=USE_FNUZ,
470
+ bytes_per_token=V41_HEAD_BYTES,
471
+ )
472
+
473
+ k_cache_fused = torch.zeros_like(k_cache_ref)
474
+ q_dummy = torch.zeros(num_tokens, 1, HEAD_DIM, dtype=dtype, device=device)
475
+ torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
476
+ q_dummy,
477
+ kv,
478
+ k_cache_fused,
479
+ slot_mapping,
480
+ positions,
481
+ cos_sin_cache,
482
+ 64,
483
+ 1e-6,
484
+ block_size,
485
+ True,
486
+ True, # kv_mxfp8
487
+ )
488
+
489
+ # A pre-quantization bf16 tie can land the two encoders on adjacent e4m3
490
+ # codes, so compare after dequantizing and allow one step at the tile scale.
491
+ def _dequant(cache_2d):
492
+ out = torch.zeros(1, num_tokens, HEAD_DIM, dtype=dtype, device=device)
493
+ v41_dequantize_and_gather_k_cache(
494
+ out,
495
+ cache_2d.view(num_blocks, block_size, V41_HEAD_BYTES),
496
+ torch.tensor([num_tokens], dtype=torch.int32, device=device),
497
+ None,
498
+ torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(0),
499
+ block_size,
500
+ offset=0,
501
+ use_fnuz=USE_FNUZ,
502
+ )
503
+ return out[0]
504
+
505
+ rec_ref = _dequant(k_cache_ref).float()
506
+ rec_fused = _dequant(k_cache_fused).float()
507
+ tile_scale = (
508
+ (rec_ref.abs().view(num_tokens, 16, 32).amax(-1).clamp(min=1e-4) / FP8_MAX)
509
+ .log2()
510
+ .ceil()
511
+ .exp2()
512
+ )
513
+ # One e4m3 ULP at the top of the range is 32 quantized units.
514
+ step = 32.0 * tile_scale.repeat_interleave(32, dim=-1)
515
+ assert ((rec_fused - rec_ref).abs() <= step).all()
516
+ # The NoPE dims see no rotation at all, so they must be byte-identical.
517
+ nope_bytes = slice(0, block_size * HEAD_DIM)
518
+ ref_rows = k_cache_ref[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM)
519
+ fused_rows = k_cache_fused[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM)
520
+ torch.testing.assert_close(
521
+ fused_rows[..., :NOPE_DIM], ref_rows[..., :NOPE_DIM], rtol=0, atol=0
522
+ )
523
+
524
+
525
# ── Test 2b: DP padding (slot_mapping shorter than q/kv) ─────────────────────
526
527
vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py
+39
-23
@@ -23,10 +23,11 @@ from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
23
class DequantGatherKCacheKernel(
24
VllmCuTeDSLJitKernel["DequantGatherKCacheKernel.CompileKey"]
25
):
26
- # Hard-coded for DSv4.
26
+ # DSv4 geometry. The record format (how much of the row is quantized and
27
+ # how wide a scale group is) comes from the compile key: V4 quantizes the
28
+ # 448 NoPE dims in groups of 64 and keeps a bf16 RoPE tail, V4.1 quantizes
29
+ # all 512 dims in groups of 32 (MXFP8) and has no tail.
30
head_dim = 512
28
- fp8_dim = 448
29
- group_size = 64 # 1 scale per 64 elems
31
num_warps = 4
32
tb_size = num_warps * 32
33
num_stages = 4
@@ -35,17 +36,20 @@ class DequantGatherKCacheKernel(
36
class CompileKey:
37
block_size: int
38
has_gather_lens: bool
39
+ fp8_dim: int = 448
40
+ group_size: int = 64 # dims per scale
41
42
@staticmethod
43
def kernel(compile_key: CompileKey) -> Any:
44
head_dim = DequantGatherKCacheKernel.head_dim
42
- fp8_dim = DequantGatherKCacheKernel.fp8_dim
43
- group_size = DequantGatherKCacheKernel.group_size
45
+ fp8_dim = compile_key.fp8_dim
46
+ group_size = compile_key.group_size
47
num_warps = DequantGatherKCacheKernel.num_warps
48
tb_size = DequantGatherKCacheKernel.tb_size
49
num_stages = DequantGatherKCacheKernel.num_stages
50
bf16_dim = head_dim - fp8_dim
51
data_dim = fp8_dim + bf16_dim * 2
52
+ scale_bytes = head_dim // group_size
53
block_size = compile_key.block_size
54
55
@cute.jit
@@ -62,8 +66,8 @@ class DequantGatherKCacheKernel(
66
):
67
op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL)
68
cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128)
65
- cp8_atom = cute.make_copy_atom(
66
- cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64
69
+ cp_scale_atom = cute.make_copy_atom(
70
+ cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=scale_bytes * 8
71
)
72
page_id = block_table[req_id, pos // block_size]
73
block_offset = pos % block_size
@@ -88,7 +92,7 @@ class DequantGatherKCacheKernel(
92
)
93
elif idx == cutlass.const_expr(data_dim // 16):
94
cute.copy(
91
- cp8_atom,
95
+ cp_scale_atom,
96
k_scale[page_id, block_offset, None],
97
s_kscale[None, stage_id],
98
)
@@ -119,8 +123,8 @@ class DequantGatherKCacheKernel(
123
)[None, warp_id, None]
124
s_kscale = smem.allocate_tensor(
125
Uint8,
122
- cute.make_layout((8, num_warps, num_stages)),
123
- byte_alignment=8,
126
+ cute.make_layout((scale_bytes, num_warps, num_stages)),
127
+ byte_alignment=scale_bytes,
128
)[None, warp_id, None]
129
130
# Prepare for 16B cp.async, also for BF16 smem loads later.
@@ -239,14 +243,16 @@ class DequantGatherKCacheKernel(
243
244
# Last 64 elems are BF16 tail, corresponds to dequant1 of last
245
# 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 +
242
- # 8x 16B for BF16.
243
- if lane_id + 32 >= fp8_dim // 8:
244
- idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8
245
- cute.copy(
246
- cp16_atom,
247
- s_kdata_16B_slice[(None, idx), compute_stage],
248
- dequant1,
249
- )
246
+ # 8x 16B for BF16. The V4.1 record is fully quantized and has
247
+ # no tail.
248
+ if cutlass.const_expr(bf16_dim > 0): # noqa: SIM102
249
+ if lane_id + 32 >= fp8_dim // 8:
250
+ idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8
251
+ cute.copy(
252
+ cp16_atom,
253
+ s_kdata_16B_slice[(None, idx), compute_stage],
254
+ dequant1,
255
+ )
256
257
# Store two 16B BF16 chunks per lane: first half, then second half.
258
dst = out_slice[req_id, offset + i, (None, lane_id)]
@@ -269,7 +275,7 @@ class DequantGatherKCacheKernel(
275
):
276
# Split k_cache into k_data and k_scale. Each [block_size, head_bytes]
277
# block is actually a concat of
272
- # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8].
278
+ # [block_size, fp8_dim + bf16_dim * 2] and [block_size, scale_bytes].
279
k_data = cute.make_tensor(
280
k_cache.iterator,
281
layout=cute.make_layout(
@@ -282,8 +288,8 @@ class DequantGatherKCacheKernel(
288
k_scale = cute.make_tensor(
289
k_cache.iterator + (block_size * data_dim),
290
layout=cute.make_layout(
285
- (k_cache.shape[0], block_size, 8),
286
- stride=(k_cache.stride[0], 8, 1),
291
+ (k_cache.shape[0], block_size, scale_bytes),
292
+ stride=(k_cache.stride[0], scale_bytes, 1),
293
),
294
)
295
@@ -311,10 +317,14 @@ class DequantGatherKCacheKernel(
317
*,
318
block_size: int,
319
has_gather_lens: bool,
320
+ fp8_dim: int = 448,
321
+ group_size: int = 64,
322
) -> CompileKey:
323
return self.CompileKey(
324
block_size=block_size,
325
has_gather_lens=has_gather_lens,
326
+ fp8_dim=fp8_dim,
327
+ group_size=group_size,
328
)
329
330
def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
@@ -345,9 +355,11 @@ class DequantGatherKCacheKernel(
355
def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
356
num_reqs = cute.sym_int()
357
head_dim = self.head_dim
348
- fp8_dim = self.fp8_dim
358
+ fp8_dim = compile_key.fp8_dim
359
block_size = compile_key.block_size
350
- head_bytes = fp8_dim + (head_dim - fp8_dim) * 2 + 8
360
+ head_bytes = (
361
+ fp8_dim + (head_dim - fp8_dim) * 2 + head_dim // compile_key.group_size
362
+ )
363
364
out = make_fake_tensor(BFloat16, (num_reqs, cute.sym_int(), head_dim), 16)
365
k_cache = cute.runtime.make_fake_tensor(
@@ -384,10 +396,14 @@ class DequantGatherKCacheKernel(
396
block_table: torch.Tensor,
397
block_size: int,
398
offset: int,
399
+ fp8_dim: int = 448,
400
+ group_size: int = 64,
401
) -> CuTeDSLLaunchSpec["DequantGatherKCacheKernel.CompileKey"]:
402
compile_key = self.dispatch(
403
block_size=block_size,
404
has_gather_lens=gather_lens is not None,
405
+ fp8_dim=fp8_dim,
406
+ group_size=group_size,
407
)
408
launch_args = (
409
out,
vllm/models/deepseek_v41/amd/dspark.py
+2
@@ -259,6 +259,8 @@ def _insert_context_kv(
259
attn.padded_heads,
260
attn.eps,
261
block_size,
262
+ True, # apply_q_norm (unused: the query is a discarded dummy)
263
+ attn.kv_mxfp8,
264
)
265
elif cache_dtype == torch.bfloat16:
266
swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
vllm/models/deepseek_v41/attention.py
+34
-9
@@ -51,6 +51,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
51
from vllm.model_executor.models.utils import extract_layer_index
52
from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope
53
from vllm.models.deepseek_v41.compressor import DeepseekCompressor
54
+from vllm.platforms import current_platform
55
from vllm.triton_utils import tl, triton
56
from vllm.utils.multi_stream_utils import (
57
execute_in_parallel,
@@ -109,6 +110,16 @@ def _fill_short_context_topk_indices(
110
)
111
112
113
+# Which packed fp8_ds_mla record a V4.1 layer writes. FlashMLA decodes
114
+# DeepSeek's V4.1 record -- all 512 dims (RoPE included) as fp8 e4m3 with one
115
+# UE8M0 scale per 32 dims, 512 data bytes + 16 scale bytes per token, pages
116
+# rounded to the kernel's 512 B TMA stride -- only in its SM100 sparse-decode
117
+# kernels. Every other arch keeps the V4 record: 448 fp8 NoPE + 64 bf16 RoPE
118
+# plus 7 UE8M0 scales of 64 dims and a pad byte (584 B, 576 B pages).
119
+def _use_v41_mxfp8_kv_record() -> bool:
120
+ return current_platform.is_device_capability_family(100)
121
+
122
+
123
def _resolve_dsv4_kv_cache_dtype(
124
use_fp8_ds_mla_layout: bool,
125
kv_cache_dtype: str,
@@ -445,6 +456,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
456
self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype(
457
self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config
458
)
459
+ self.kv_mxfp8 = _use_v41_mxfp8_kv_record()
460
+ self.kv_bytes_per_token = 528 if self.kv_mxfp8 else 584
461
+ self.kv_page_alignment = 512 if self.kv_mxfp8 else 576
462
463
self.swa_cache_layer = DeepseekV4SWACache(
464
head_dim=self.head_dim,
@@ -454,6 +468,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
468
cache_config=cache_config,
469
backend_cls=self.swa_backend_cls,
470
block_size=32,
471
+ packed_bytes_per_token=self.kv_bytes_per_token,
472
+ packed_page_alignment=self.kv_page_alignment,
473
)
474
475
# The attention layer itself was already registered with the
@@ -879,6 +895,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
895
self.eps,
896
swa_metadata.block_size,
897
False,
898
+ self.kv_mxfp8,
899
)
900
901
# Plain-row path: the [num_blocks, block_size, 512] cache stores the KV
@@ -934,9 +951,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
951
# DeepseekV4SWACache.
952
if not self.is_kv_source:
953
return None
937
- # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B
938
- # alignment; plain bf16 / per-tensor fp8 rows use natural element-size
939
- # pages.
954
+ # fp8_ds_mla is a UE8M0 block-scaled uint8 layout whose page rounds up
955
+ # to the decode kernel's TMA stride; plain bf16 / per-tensor fp8 rows
956
+ # use natural element-size pages.
957
uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla"
958
return MLAAttentionSpec(
959
block_size=vllm_config.cache_config.block_size,
@@ -945,12 +962,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
962
dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
963
tokens_per_state=self.compress_ratio,
964
cache_dtype_str=self.kv_cache_dtype,
948
- alignment=576 if uses_fp8_ds_mla_layout else 512,
965
+ alignment=self.kv_page_alignment if uses_fp8_ds_mla_layout else 512,
966
model_version="deepseek_v4",
967
kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
951
- # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token;
952
- # head_size stays semantic (512).
953
- state_content_bytes=584 if uses_fp8_ds_mla_layout else None,
968
+ # Packed record width; head_size stays semantic (512).
969
+ state_content_bytes=(
970
+ self.kv_bytes_per_token if uses_fp8_ds_mla_layout else None
971
+ ),
972
)
973
974
def _compressed_kv_cache(self) -> torch.Tensor:
@@ -998,8 +1016,15 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
1016
head_size=self.head_dim,
1017
dtype=self.dtype,
1018
tokens_per_state=self.compress_ratio,
1001
- # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
1002
- alignment=576 if uses_fp8_ds_mla_layout else 512,
1019
+ # The indexer page shares a physical block with the main KV cache
1020
+ # (the block stride is the sum of every page in the group), so it
1021
+ # has to carry the same padding: the packed record's TMA stride,
1022
+ # else 512B for FlashInfer sparse (#44577).
1023
+ alignment=(
1024
+ 576
1025
+ if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record()
1026
+ else 512
1027
+ ),
1028
)
1029
1030
def forward(self): ...
vllm/models/deepseek_v41/common/ops/cache_utils.py
+186
-8
@@ -36,6 +36,18 @@ from vllm.triton_utils import tl, triton
36
from vllm.utils.import_utils import has_cutedsl
37
from vllm.utils.math_utils import next_power_of_2
38
39
+# Per-token byte width of the two paged fp8 records. Both put a page's whole
40
+# data region ahead of its whole scale region, so ``k_cache.shape[-1]`` -- the
41
+# width the KV-cache spec publishes -- identifies the record:
42
+# V4 (584 B): 448 fp8 e4m3 NoPE + 64 bf16 RoPE, then 7 UE8M0 scales of 64
43
+# dims each plus a pad byte.
44
+# V4.1 (528 B): all 512 dims as fp8 e4m3 (RoPE included), then 16 UE8M0
45
+# scales of 32 dims each. FlashMLA's ``ModelType::V41``.
46
+V4_BYTES_PER_TOKEN = 584
47
+V41_BYTES_PER_TOKEN = 528
48
+V41_QUANT_BLOCK = 32
49
+V41_NUM_SCALES = 512 // V41_QUANT_BLOCK # 16
50
+
51
52
@triton.jit
53
def quantize_and_insert_k_kernel(
@@ -163,6 +175,67 @@ def quantize_and_insert_k_kernel(
175
tl.store(bf16_out_ptr + chunk_offsets, bf16_vals)
176
177
178
+@triton.jit
179
+def _quantize_and_insert_k_mxfp8_kernel(
180
+ k_ptr, # [num_tokens, 512] bf16
181
+ slot_mapping_ptr, # [num_tokens] int64
182
+ k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view)
183
+ num_tokens,
184
+ head_dim: tl.constexpr, # 512, RoPE dims included
185
+ scale_dim: tl.constexpr, # 16
186
+ quant_block: tl.constexpr, # 32 dims per UE8M0 scale
187
+ cache_block_size: tl.constexpr, # paged cache block size, in tokens
188
+ block_stride: tl.constexpr, # total bytes per block (padded)
189
+ fp8_max: tl.constexpr,
190
+ use_fnuz: tl.constexpr = False,
191
+):
192
+ """Quantize a bf16 K row to MXFP8 and insert it into a V4.1 paged cache.
193
+
194
+ One program per token. Unlike the V4 record there is no bf16 tail, so every
195
+ dim is scaled.
196
+ """
197
+ pid = tl.program_id(0)
198
+ if pid >= num_tokens:
199
+ return
200
+
201
+ slot_idx = tl.load(slot_mapping_ptr + pid)
202
+ if slot_idx == -1:
203
+ return
204
+
205
+ block_idx = slot_idx // cache_block_size
206
+ pos_in_block = slot_idx % cache_block_size
207
+
208
+ # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks.
209
+ cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride
210
+ token_data_ptr = cache_block_ptr + pos_in_block * head_dim
211
+ token_scale_ptr = (
212
+ cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim
213
+ )
214
+
215
+ d = tl.arange(0, head_dim)
216
+ x = tl.load(k_ptr + pid.to(tl.int64) * head_dim + d).to(tl.float32)
217
+
218
+ tiles = tl.reshape(x, (scale_dim, quant_block))
219
+ # Floor the absmax like the CUDA encoder so an all-zero tile still lands on
220
+ # a finite scale.
221
+ amax = tl.maximum(tl.max(tl.abs(tiles), 1), 1e-4)
222
+ exponent = tl.ceil(tl.log2(amax * (1.0 / fp8_max)))
223
+ scaled = tl.clamp(
224
+ tiles * tl.reshape(tl.exp2(-exponent), (scale_dim, 1)), -fp8_max, fp8_max
225
+ )
226
+ if use_fnuz: # noqa: SIM108
227
+ fp8 = scaled.to(tl.float8e4b8)
228
+ else:
229
+ fp8 = scaled.to(tl.float8e4nv)
230
+ tl.store(
231
+ token_data_ptr + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (head_dim,))
232
+ )
233
+
234
+ # UE8M0 encoding: stored_value = exponent + 127 (bias).
235
+ encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), 255.0)
236
+ tl.store(token_scale_ptr + tl.arange(0, scale_dim), encoded.to(tl.uint8))
237
+
238
+
239
def quantize_and_insert_k_cache(
240
k: torch.Tensor, # [num_tokens, 512] bf16
241
k_cache: torch.Tensor, # [num_blocks, block_bytes] uint8
@@ -170,16 +243,13 @@ def quantize_and_insert_k_cache(
243
block_size: int = 64,
244
is_ue8m0: bool = True,
245
use_fnuz: bool = False,
246
+ bytes_per_token: int = V4_BYTES_PER_TOKEN,
247
):
248
"""
249
Quantize K tensor and insert into paged K cache.
250
177
- K Cache block layout (block_size=64 tokens):
178
- - First 64 * 576 = 36864 bytes: Token data
179
- - Each token: 448 bytes (fp8) + 128 bytes (bf16)
180
- - Next 64 * 8 = 512 bytes: Scales
181
- - Each token: 8 bytes (uint8 scales, 7 real + 1 padding)
182
- - Padded to multiple of 576
251
+ ``bytes_per_token`` picks the record (see the module header): the V4 one,
252
+ or V4.1's all-dims MXFP8 one.
253
254
``use_fnuz=True`` selects FNUZ E4M3 cache encoding and is only valid on
255
platforms whose FP8 format is FNUZ. ``use_fnuz=False`` selects OCP E4M3,
@@ -196,6 +266,26 @@ def quantize_and_insert_k_cache(
266
num_tokens = slot_mapping.shape[0]
267
block_stride = k_cache.stride(0) # bytes per block
268
269
+ if bytes_per_token == V41_BYTES_PER_TOKEN:
270
+ if use_fnuz and not current_platform.is_fp8_fnuz():
271
+ raise ValueError("use_fnuz=True requires a platform using FNUZ FP8")
272
+ _quantize_and_insert_k_mxfp8_kernel[(num_tokens,)](
273
+ k,
274
+ slot_mapping,
275
+ k_cache,
276
+ num_tokens,
277
+ head_dim=512,
278
+ scale_dim=V41_NUM_SCALES,
279
+ quant_block=V41_QUANT_BLOCK,
280
+ cache_block_size=block_size,
281
+ block_stride=block_stride,
282
+ fp8_max=get_fp8_min_max()[1]
283
+ if use_fnuz
284
+ else torch.finfo(torch.float8_e4m3fn).max,
285
+ use_fnuz=use_fnuz,
286
+ )
287
+ return
288
+
289
TOKEN_FP8_DIM = 448
290
TOKEN_BF16_DIM = 64
291
TOKEN_SCALE_DIM = 8
@@ -343,6 +433,69 @@ def _dequantize_and_gather_k_kernel(
433
tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals)
434
435
436
+@triton.jit
437
+def _dequantize_and_gather_k_mxfp8_kernel(
438
+ out_ptr,
439
+ out_stride0,
440
+ out_stride1,
441
+ k_cache_ptr,
442
+ seq_lens_ptr,
443
+ block_table_ptr,
444
+ offset,
445
+ gather_lens_ptr,
446
+ max_blocks_per_seq: tl.constexpr,
447
+ head_dim: tl.constexpr, # 512
448
+ scale_dim: tl.constexpr, # 16
449
+ quant_block: tl.constexpr, # 32
450
+ cache_block_size: tl.constexpr,
451
+ block_stride: tl.constexpr,
452
+ use_fnuz: tl.constexpr = False,
453
+):
454
+ """Gather V4.1 MXFP8 rows into a bf16 workspace, dequantizing in place."""
455
+ batch_idx = tl.program_id(0)
456
+ worker_id = tl.program_id(1)
457
+ num_workers = tl.num_programs(1)
458
+
459
+ seq_len = tl.load(seq_lens_ptr + batch_idx)
460
+ if gather_lens_ptr is not None: # noqa: SIM108
461
+ gather_len = tl.load(gather_lens_ptr + batch_idx)
462
+ else:
463
+ gather_len = seq_len
464
+ start_pos = seq_len - gather_len
465
+
466
+ d = tl.arange(0, head_dim)
467
+ s = tl.arange(0, scale_dim)
468
+ for i in range(worker_id, gather_len, num_workers):
469
+ pos = start_pos + i
470
+ block_in_seq = pos // cache_block_size
471
+ pos_in_block = pos % cache_block_size
472
+
473
+ block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq
474
+ physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32
475
+
476
+ # int64: physical_block_idx * block_stride can exceed 2^31.
477
+ cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
478
+ token_data_ptr = cache_block_ptr + pos_in_block * head_dim
479
+ token_scale_ptr = (
480
+ cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim
481
+ )
482
+
483
+ x_uint8 = tl.load(token_data_ptr + d)
484
+ if use_fnuz:
485
+ x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)
486
+ else:
487
+ x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
488
+ tiles = tl.reshape(x_fp8.to(tl.float32), (scale_dim, quant_block))
489
+
490
+ # UE8M0: scale = 2^(stored_value - 127).
491
+ encoded = tl.load(token_scale_ptr + s)
492
+ scale = tl.exp2(encoded.to(tl.float32) - 127.0)
493
+ dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,))
494
+
495
+ output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1
496
+ tl.store(output_row_ptr + d, dequant.to(tl.bfloat16))
497
+
498
+
499
def dequantize_and_gather_k_cache_triton(
500
# [num_reqs, max_num_tokens, head_size]
501
out: torch.Tensor,
@@ -358,6 +511,28 @@ def dequantize_and_gather_k_cache_triton(
511
offset: int,
512
use_fnuz: bool = False,
513
) -> None:
514
+ num_reqs = seq_lens.shape[0]
515
+ NUM_WORKERS = 128
516
+ if k_cache.shape[-1] == V41_BYTES_PER_TOKEN:
517
+ _dequantize_and_gather_k_mxfp8_kernel[(num_reqs, NUM_WORKERS)](
518
+ out,
519
+ out.stride(0),
520
+ out.stride(1),
521
+ k_cache,
522
+ seq_lens,
523
+ block_table,
524
+ offset,
525
+ gather_lens,
526
+ max_blocks_per_seq=block_table.shape[-1],
527
+ head_dim=512,
528
+ scale_dim=V41_NUM_SCALES,
529
+ quant_block=V41_QUANT_BLOCK,
530
+ cache_block_size=block_size,
531
+ block_stride=k_cache.stride(0),
532
+ use_fnuz=use_fnuz,
533
+ )
534
+ return
535
+
536
TOKEN_FP8_DIM = 448
537
TOKEN_BF16_DIM = 64
538
TOKEN_SCALE_DIM = 8
@@ -365,8 +540,6 @@ def dequantize_and_gather_k_cache_triton(
540
FP8_MAX = 448.0
541
TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2
542
368
- num_reqs = seq_lens.shape[0]
369
- NUM_WORKERS = 128
543
_dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)](
544
out,
545
out.stride(0),
@@ -408,6 +581,8 @@ def dequantize_and_gather_k_cache(
581
) -> None:
582
"""Dequantize and gather a paged DSv4 K cache.
583
584
+ The record is read off ``k_cache.shape[-1]``; see the module header.
585
+
586
``use_fnuz`` MUST match the encoder of the specific cache being read:
587
``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere),
588
``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder
@@ -419,6 +594,7 @@ def dequantize_and_gather_k_cache(
594
_DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL,
595
)
596
597
+ mxfp8 = k_cache.shape[-1] == V41_BYTES_PER_TOKEN
598
_DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL(
599
out=out,
600
k_cache=k_cache,
@@ -427,6 +603,8 @@ def dequantize_and_gather_k_cache(
603
block_table=block_table,
604
block_size=block_size,
605
offset=offset,
606
+ fp8_dim=512 if mxfp8 else 448,
607
+ group_size=V41_QUANT_BLOCK if mxfp8 else 64,
608
)
609
return
610
vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py
+70
-6
@@ -232,10 +232,13 @@ def rope_quant_insert(
232
233
The BF16 latent supplies both NoPE quantization and RoPE input. It is read
234
only for valid slots at group boundaries. The cache dtype selects the
235
- layout: ``uint8`` is the fp8_ds_mla paged layout (576 value bytes and eight
236
- segregated UE8M0 scale bytes per token, including one zero padding scale);
237
- ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows
238
- read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``.
235
+ layout: ``uint8`` is the fp8_ds_mla paged layout, whose record the
236
+ per-token byte width names -- 584 B for V4 (576 value bytes and eight
237
+ segregated UE8M0 scale bytes, including one zero padding scale) or 528 B
238
+ for V4.1 (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0
239
+ scales of 32 dims each). ``bfloat16`` and ``float8_e4m3fn`` are the plain
240
+ [448 NoPE | 64 RoPE] rows read by FlashInfer, the latter scaled by the
241
+ per-tensor ``fp8_scale``.
242
"""
243
assert compress_ratio in (1, 2)
244
assert latent.shape[1] == 512 and latent.dtype == torch.bfloat16
@@ -246,8 +249,15 @@ def rope_quant_insert(
249
return
250
launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {}
251
if kv_cache.dtype == torch.uint8:
249
- assert kv_cache.shape[-1] == 584
250
- _rope_quant_insert_kernel[(num_tokens,)](
252
+ assert kv_cache.shape[-1] in (584, 528), (
253
+ f"unsupported paged KV record width {kv_cache.shape[-1]}"
254
+ )
255
+ kernel = (
256
+ _rope_quant_insert_mxfp8_kernel
257
+ if kv_cache.shape[-1] == 528
258
+ else _rope_quant_insert_kernel
259
+ )
260
+ kernel[(num_tokens,)](
261
latent,
262
positions,
263
cos_sin_cache,
@@ -338,6 +348,60 @@ def _rope_quant_insert_kernel(
348
tl.store(rope_dst + d - 448, rotated.to(tl.bfloat16), d >= 448)
349
350
351
+@triton.jit
352
+def _rope_quant_insert_mxfp8_kernel(
353
+ latent,
354
+ positions,
355
+ cos_sin,
356
+ cache,
357
+ cache_slots,
358
+ COS_STRIDE: tl.constexpr,
359
+ CACHE_STRIDE: tl.constexpr,
360
+ CACHE_BLOCK: tl.constexpr,
361
+ COMPRESS_RATIO: tl.constexpr,
362
+ SANITIZE_CACHE_NANS: tl.constexpr,
363
+):
364
+ """V4.1 record: RoPE first, then MXFP8-quantize all 512 dims.
365
+
366
+ The RoPE dims are quantized here too, so unlike the V4 kernel the rotation
367
+ has to happen before the scales are picked.
368
+ """
369
+ t = tl.program_id(0)
370
+ slot = tl.load(cache_slots + t)
371
+ if slot < 0:
372
+ return
373
+ position = tl.load(positions + t)
374
+ if (position + 1) % COMPRESS_RATIO != 0:
375
+ return
376
+ d = tl.arange(0, 512)
377
+ normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32)
378
+
379
+ # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there.
380
+ even, odd = tl.split(tl.reshape(normed, (256, 2)))
381
+ pair = tl.arange(0, 256) - 224
382
+ cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE
383
+ c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32)
384
+ s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32)
385
+ rotated = tl.interleave(even * c - odd * s, odd * c + even * s)
386
+ if SANITIZE_CACHE_NANS:
387
+ rotated = tl.where(rotated == rotated, rotated, 0.0)
388
+
389
+ page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE
390
+ values = page + (slot % CACHE_BLOCK) * 512
391
+ scales = page + CACHE_BLOCK * 512 + (slot % CACHE_BLOCK) * 16
392
+
393
+ quant = tl.reshape(rotated, (16, 32))
394
+ amax = tl.maximum(tl.max(tl.abs(quant), 1), 1e-4)
395
+ exponent = tl.ceil(tl.log2(amax * (1.0 / 448.0)))
396
+ scaled = quant * tl.reshape(tl.exp2(-exponent), (16, 1))
397
+ fp8 = tl.clamp(scaled, -448.0, 448.0).to(tl.float8e4nv)
398
+ tl.store(values + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (512,)))
399
+
400
+ max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0
401
+ encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), max_encoded)
402
+ tl.store(scales + tl.arange(0, 16), encoded.to(tl.uint8))
403
+
404
+
405
@triton.jit
406
def _rope_plain_insert_kernel(
407
latent,
vllm/models/deepseek_v41/nvidia/dspark.py
+2
@@ -264,6 +264,8 @@ def _insert_context_kv(
264
attn.padded_heads,
265
attn.eps,
266
block_size,
267
+ True, # apply_q_norm (unused: the query is a discarded dummy)
268
+ attn.kv_mxfp8,
269
)
270
elif cache_dtype == torch.bfloat16:
271
swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
vllm/v1/attention/backends/mla/flashmla_sparse.py
+11
@@ -91,6 +91,17 @@ Bytes, structured as:
91
The first `ue8m0` is the scale for the first 64 `float8_e4m3` values,
92
the second for the next 64, and so on.
93
94
+For DeepSeek V4.1, each token's KV cache is 528 Bytes, structured as:
95
+- **First 512 bytes:** all 512 dims as `float8_e4m3`. Unlike V4 the 64 RoPE
96
+ dims are quantized too, so there is no `bfloat16` part.
97
+- **Last 16 bytes:** Scale factors, containing 16 `ue8m0` values, one per 32
98
+ consecutive `float8_e4m3` values (i.e. MXFP8).
99
+
100
+The V4 and V4.1 records are not laid out token-by-token within a page: a page
101
+holds all its data rows first and all its scale rows after, so a page is
102
+`block_size * bytes_per_token` bytes rounded up to the decode kernel's TMA
103
+stride (576 B for V4, 512 B for V4.1).
104
+
105
In the "nvfp4_ds_mla" format (SM100 only, DeepSeek V3.2 geometry), each
106
token's KV cache is 352 Bytes, structured as:
107
- **First 256 bytes:** 512 `e2m1` NoPE values packed 2/byte (low nibble =
vllm/v1/attention/backends/mla/sparse_swa.py
+12
-6
@@ -79,6 +79,8 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
79
cache_config: CacheConfig,
80
backend_cls: "type[AttentionBackend] | None" = None,
81
block_size: int = 64,
82
+ packed_bytes_per_token: int = 584,
83
+ packed_page_alignment: int = 576,
84
):
85
super().__init__()
86
self.backend_cls = backend_cls or DeepseekSparseSWABackend
@@ -88,6 +90,8 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
90
self.prefix = prefix
91
self.cache_config = cache_config
92
self.dtype = dtype
93
+ self.packed_bytes_per_token = packed_bytes_per_token
94
+ self.packed_page_alignment = packed_page_alignment
95
compilation_config = get_current_vllm_config().compilation_config
96
if prefix in compilation_config.static_forward_context:
97
raise ValueError(f"Duplicate layer name: {prefix}")
@@ -105,8 +109,9 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
109
self.kv_cache = kv_cache.squeeze(1)
110
111
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
108
- # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous
109
- # bf16/fp8 cache uses the natural element-size page.
112
+ # fp8_ds_mla's UE8M0 paged layout rounds its page up to the decode
113
+ # kernel's TMA stride; contiguous bf16/fp8 cache uses the natural
114
+ # element-size page.
115
uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla"
116
return SlidingWindowMLASpec(
117
block_size=self.block_size,
@@ -115,10 +120,11 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
120
dtype=self.dtype,
121
sliding_window=self.window_size,
122
cache_dtype_str=self.cache_config.cache_dtype,
118
- # DeepseekV4 fp8_ds_mla: 584B per token (448B NoPE + 128B RoPE + 8B scales)
119
- state_content_bytes=584 if uses_fp8_ds_mla_layout else None,
120
- # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
121
- alignment=576 if uses_fp8_ds_mla_layout else 512,
123
+ state_content_bytes=(
124
+ self.packed_bytes_per_token if uses_fp8_ds_mla_layout else None
125
+ ),
126
+ # FlashMLA packing stride; 512B for FlashInfer sparse (#44577).
127
+ alignment=self.packed_page_alignment if uses_fp8_ds_mla_layout else 512,
128
model_version="deepseek_v4",
129
kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype),
130
)