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