iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Attention][MiniMax-M3] NVFP4 KV cache on the MSA sparse attention path (#59300)

Browse Source 
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yongye Zhu
2026-10-01 14:39:55 -07:00
committed by GitHub
co-authored by Claude Opus 5.5
parent ea84051819
commit 7566d83bd3
12 changed files with 1297 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+21
View File
@@ -590,9 +590,30 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
endif()
# The fused MiniMax-M3 KV insert quantizes NVFP4 with SM100-family-only
# cvt instructions, so its SM10x/11x code is built for the family targets.
set(MINIMAX_M3_FUSED_SRC
"csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu")
list(REMOVE_ITEM VLLM_STABLE_EXT_SRC "${MINIMAX_M3_FUSED_SRC}")
set_gencode_flags_for_srcs(
SRCS "${VLLM_STABLE_EXT_SRC}"
CUDA_ARCHS "${CUDA_ARCHS}")
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
cuda_archs_loose_intersection(MINIMAX_M3_FUSED_ARCHS
"10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
else()
cuda_archs_loose_intersection(MINIMAX_M3_FUSED_ARCHS
"10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
endif()
foreach(_ARCH ${CUDA_ARCHS})
if(_ARCH VERSION_LESS 10.0 OR _ARCH VERSION_GREATER_EQUAL 12.0)
list(APPEND MINIMAX_M3_FUSED_ARCHS "${_ARCH}")
endif()
endforeach()
set_gencode_flags_for_srcs(
SRCS "${MINIMAX_M3_FUSED_SRC}"
CUDA_ARCHS "${MINIMAX_M3_FUSED_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${MINIMAX_M3_FUSED_SRC}")
if(COOPERATIVE_TOPK_ARCHS)
list(APPEND VLLM_STABLE_EXT_SRC
cmake/external_projects/fmha_sm100.cmake
+15 -1
View File
@@ -17,7 +17,7 @@ else()
FetchContent_Declare(
fmha_sm100
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
GIT_TAG f355c37eb4e1413f21ee2ad8bbad25079e6bef9d
GIT_TAG be4a52aea420b6bef73be9af1a649859964790cb
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
@@ -39,6 +39,8 @@ install(FILES
"${FMHA_SM100_PY_ROOT}/api.py"
"${FMHA_SM100_PY_ROOT}/bench_utils.py"
"${FMHA_SM100_PY_ROOT}/jit.py"
"${FMHA_SM100_PY_ROOT}/nvfp4_kv.py"
"${FMHA_SM100_PY_ROOT}/q8kv4_decode_adapter.py"
"${FMHA_SM100_PY_ROOT}/sparse.py"
"${FMHA_SM100_PY_ROOT}/sparse_fmha_adapter.py"
DESTINATION vllm/third_party/fmha_sm100
@@ -51,6 +53,18 @@ install(DIRECTORY "${FMHA_SM100_PY_ROOT}/csrc/"
PATTERN "*.pyc" EXCLUDE
PATTERN ".git*" EXCLUDE)
install(DIRECTORY "${FMHA_SM100_PY_ROOT}/decode_q8kv4/"
DESTINATION vllm/third_party/fmha_sm100/decode_q8kv4
COMPONENT fmha_sm100
PATTERN "__pycache__" EXCLUDE
PATTERN "*.pyc" EXCLUDE)
install(DIRECTORY "${FMHA_SM100_PY_ROOT}/kvouter/"
DESTINATION vllm/third_party/fmha_sm100/kvouter
COMPONENT fmha_sm100
PATTERN "__pycache__" EXCLUDE
PATTERN "*.pyc" EXCLUDE)
install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/"
DESTINATION vllm/third_party/fmha_sm100/cute
COMPONENT fmha_sm100
csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
+391 -42
View File
@@ -285,6 +285,262 @@ __device__ __forceinline__ void storeScaledQElemsFp8(
storeElemsFp8(dst, scaled);
}
// index_q: E4M3(scalar_t(v)) with unit scale, i.e. q is rounded to the model
// dtype before the fp8 cast. On CUDA each element pair is rounded to scalar_t
// and then converted to E4M3, with one packed cvt each.
template <typename scalar_t>
__device__ __forceinline__ void storeIndexQElemsFp8(
uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
#ifndef USE_ROCM
using Converter = vllm::_typeConvert<scalar_t>;
__nv_fp8x2_storage_t out2[kElemsPerLane / 2];
#pragma unroll
for (int i = 0; i < kElemsPerLane / 2; i++) {
auto const v =
Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1]));
if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
out2[i] = __nv_cvt_bfloat16raw2_to_fp8x2(v, __NV_SATFINITE, __NV_E4M3);
} else {
out2[i] = __nv_cvt_halfraw2_to_fp8x2(v, __NV_SATFINITE, __NV_E4M3);
}
}
*reinterpret_cast<uint32_t*>(dst) = *reinterpret_cast<uint32_t const*>(out2);
#else
storeScaledQElemsFp8<scalar_t>(dst, elems, 1.0f);
#endif
}
// ────────────────────────────────────────────────────────────────────────────
// NVFP4 main K/V cache store
// ────────────────────────────────────────────────────────────────────────────
// The block encoding of reshape_and_cache_nvfp4_kernel
// (nvfp4_kv_cache_kernels.cu) in the MSA readers' head-slot page layout. The
// quantizer differs from that kernel's reciprocal-multiply math: per 16-value
// block, scale = min(448, max(E4M3(max(amax, 1e-12) / 6), 1/512)) and
// code = sign | RNE-E2M1(min(|x / scale|, 6)), with zero stored as +0, all on
// x / the global scale.
// An HND page is 2*nkv slots of [bs, 72] bytes; slot 2*head + side (K = 0,
// V = 1) holds that head's packed E2M1 data [bs, 64] followed by its E4M3 block
// scales [bs, 8]. Every head is one contiguous 2*bs*72-byte run, so a page
// split across TP ranks is a contiguous byte slice per rank. K scales are
// linear, V scales use the SM100 trtllm-gen 4x4 token-quad swizzle. One
// 16-element scale group is four lanes. SM100-family builds of this file
// (CMake gives it family-specific SM10x/11x targets) encode E2M1 with
// cvt.rn.satfinite.e2m1x2; other targets use a bit-exact software encoder.
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
defined(__CUDA_ARCH_FAMILY_SPECIFIC__)
#define MINIMAX_M3_NVFP4_NATIVE_CVT 1
#else
#define MINIMAX_M3_NVFP4_NATIVE_CVT 0
#endif
#ifndef USE_ROCM
constexpr int kNvfp4DataBytes = kHeadDim / 2; // 64
constexpr int kNvfp4ScaleBytes = kHeadDim / 16; // 8
constexpr int kNvfp4LanesPerGroup = 16 / kElemsPerLane; // 4
// Round-to-nearest-even, saturating E2M1 code of ``x`` (matches
// cvt.rn.satfinite.e2m1x2.f32). Below 1 the grid is the 0.5-step subnormals;
// from 1 up it is fp32 rounded to one mantissa bit, whose (exponent, mantissa)
// bits are the code offset by 252 once the exponent is rebiased.
__device__ __forceinline__ uint32_t e2m1Code(float x) {
float const ax = fabsf(x);
uint32_t mag;
if (ax < 1.0f) {
mag = static_cast<uint32_t>(__float2int_rn(2.0f * ax));
} else {
uint32_t const bits = __float_as_uint(ax);
uint32_t const rounded = bits + 0x1FFFFFu + ((bits >> 22) & 1u);
mag = min((rounded >> 22) - 252u, 7u);
}
return mag | ((__float_as_uint(x) >> 28) & 0x8u);
}
__device__ __forceinline__ float rcpApproxFtz(float a) {
float b;
asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(b) : "f"(a));
return b;
}
// a / b rounded to nearest even from r = RN(1 / b): Markstein's correction of
// a * r, which is __fdiv_rn's fast path without its range-check branch. Exact
// unless a * r overflows or the residual a - q * b underflows (a within a few
// binades of the subnormals).
__device__ __forceinline__ float divRnByRcp(float a, float b, float r) {
float const q = __fmul_rn(a, r);
return __fmaf_rn(__fmaf_rn(-q, b, a), r, q);
}
// Block scale E4M3(max(amax, 1e-12) / 6) clamped to [1/512, 448], with
// sf_scale applied to it as to the values, and its reciprocal RN(1 / sf_value).
__device__ __forceinline__ void nvfp4GroupScale(float sf_scale, float group_max,
uint8_t& sf_byte,
float& sf_value,
float& sf_rcp) {
constexpr float kRcp6 = 0x1.555556p-3f; // RN(1 / 6)
float const amax = fmaxf(group_max, 1e-12f);
// amax / 6 (RNE). amax * kRcp6 bounds it from above, and fminf picks it for
// amax = inf, where the correction is NaN.
float const amax_6 =
fminf(divRnByRcp(amax, 6.0f, kRcp6), __fmul_rn(amax, kRcp6));
// Flooring before the rounding gives the 1/512 minimum; satfinite caps it
// at 448.
__nv_fp8_e4m3 const sf(fmaxf(sf_scale * amax_6, 0x1p-9f));
sf_byte = sf.__x;
sf_value = float(sf);
// rcp.approx is correctly rounded on every positive E4M3 value.
sf_rcp = rcpApproxFtz(sf_value);
}
// nvfp4GroupScale for sf_scale = 1 and a bf16/fp16 group_max. Uncorrected,
// amax * RN(1 / 6) still rounds to the E4M3 byte of amax / 6 (checked for every
// such amax: amax / 6 is too far from an E4M3 rounding boundary to cross it,
// and exact ties round back), and the 1e-12 floor lies below the 1/512 one.
__device__ __forceinline__ void nvfp4UnitGroupScale(float group_max,
uint8_t& sf_byte,
float& sf_rcp) {
constexpr float kRcp6 = 0x1.555556p-3f; // RN(1 / 6)
__nv_fp8_e4m3 const sf(fmaxf(__fmul_rn(group_max, kRcp6), 0x1p-9f));
sf_byte = sf.__x;
sf_rcp = rcpApproxFtz(float(sf));
}
// (x * sf_scale) / sf_value with the E2M1 code of the IEEE quotient. The
// quotient itself is the IEEE one except below 1e-35 (code 0 either way) and
// for magnitudes over 4096, which are clamped (NaN kept) so that the division
// stays finite: 4096 / 448 > 5, so their code stays the saturated 6.
__device__ __forceinline__ float nvfp4Quotient(float x, float sf_scale,
float sf_value, float sf_rcp) {
constexpr float kMaxMag = 4096.0f;
float n = __fmul_rn(x, sf_scale);
n = fabsf(n) > kMaxMag ? copysignf(kMaxMag, n) : n;
return divRnByRcp(n, sf_value, sf_rcp);
}
// Zero is stored as +0: clear the sign of every nibble whose magnitude is 0.
// Adding 7 to a 3-bit magnitude carries into the sign bit exactly when it is
// nonzero.
__device__ __forceinline__ uint16_t clearNegZeroNibbles(uint16_t packed) {
return packed & (((packed & 0x7777u) + 0x7777u) | 0x7777u);
}
// Quantize one K or V head row of one token into its NVFP4 page, where
// ``scale`` is the dequantization (global) scale. Every lane of the warp must
// call this (the scale-group reductions shuffle).
template <typename scalar_t>
__device__ __forceinline__ void storeNvfp4CacheElems(
uint8_t* __restrict__ page, int const kv, int const head, int const token,
int const block_size, int const laneId, float const (&elems)[kElemsPerLane],
float const scale) {
// Quantize the model-dtype value that the unfused path reads back from qkv.
using Converter = vllm::_typeConvert<scalar_t>;
float x[kElemsPerLane];
float group_max = 0.0f;
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
x[i] = static_cast<float>(Converter::convert(elems[i]));
group_max = fmaxf(group_max, fabsf(x[i]));
}
group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 1, 32));
group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 2, 32));
uint8_t sf_byte;
float sf_rcp;
float q[kElemsPerLane]; // (x / scale) / sf_value, E2M1-encoded below
if (scale == 1.0f) {
static_assert(sizeof(scalar_t) == 2, "the unit-scale path needs bf16/fp16");
nvfp4UnitGroupScale(group_max, sf_byte, sf_rcp);
// x / sf_value truncated from x * RN(1 / sf_value). RN(1 / s) >= 1 / s for
// every E4M3 s, so a bf16/fp16 x on an E2M1 midpoint lands exactly on it
// (and the cvt rounds it to even), while any other bf16/fp16 x is too far
// from a midpoint to reach one. No clamp either: a finite product never
// truncates to inf, and inf saturates in the cvt.
#if MINIMAX_M3_NVFP4_NATIVE_CVT
float2 const r2 = make_float2(sf_rcp, sf_rcp);
#pragma unroll
for (int i = 0; i < kElemsPerLane; i += 2) {
float2 const p = __fmul2_rz(make_float2(x[i], x[i + 1]), r2);
q[i] = p.x;
q[i + 1] = p.y;
}
#else
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) q[i] = __fmul_rz(x[i], sf_rcp);
#endif
} else {
float const sf_scale = __frcp_rn(scale); // exactly 1.0f / scale
float sf_value;
nvfp4GroupScale(sf_scale, group_max, sf_byte, sf_value, sf_rcp);
// A correctly rounded division, not a reciprocal multiply, so that exact
// E2M1 midpoints round to even.
#if MINIMAX_M3_NVFP4_NATIVE_CVT
// nvfp4Quotient on element pairs, with SM100's packed FP32 instructions.
float2 const s2 = make_float2(sf_scale, sf_scale);
float2 const neg_v2 = make_float2(-sf_value, -sf_value);
float2 const r2 = make_float2(sf_rcp, sf_rcp);
#pragma unroll
for (int i = 0; i < kElemsPerLane; i += 2) {
float2 n = __fmul2_rn(make_float2(x[i], x[i + 1]), s2);
// |n| <= 4096 as in nvfp4Quotient, NaN kept.
asm("max.NaN.f32 %0, %0, 0fC5800000;\n\t"
"min.NaN.f32 %0, %0, 0f45800000;\n\t"
"max.NaN.f32 %1, %1, 0fC5800000;\n\t"
"min.NaN.f32 %1, %1, 0f45800000;"
: "+f"(n.x), "+f"(n.y));
float2 const q0 = __fmul2_rn(n, r2);
float2 const p = __ffma2_rn(__ffma2_rn(q0, neg_v2, n), r2, q0);
q[i] = p.x;
q[i + 1] = p.y;
}
#else
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
q[i] = nvfp4Quotient(x[i], sf_scale, sf_value, sf_rcp);
}
#endif
}
uint16_t packed;
#if MINIMAX_M3_NVFP4_NATIVE_CVT
// Low nibble = even element; cvt packs its first source in the high nibble.
asm("{\n"
" .reg .b8 lo, hi;\n"
" cvt.rn.satfinite.e2m1x2.f32 lo, %2, %1;\n"
" cvt.rn.satfinite.e2m1x2.f32 hi, %4, %3;\n"
" mov.b16 %0, {lo, hi};\n"
"}"
: "=h"(packed)
: "f"(q[0]), "f"(q[1]), "f"(q[2]), "f"(q[3]));
#else
uint32_t codes = 0;
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) codes |= e2m1Code(q[i]) << (4 * i);
packed = static_cast<uint16_t>(codes);
#endif
packed = clearNegZeroNibbles(packed);
uint8_t* const slot = page + static_cast<int64_t>(2 * head + kv) *
block_size *
(kNvfp4DataBytes + kNvfp4ScaleBytes);
uint8_t* const data =
slot + token * kNvfp4DataBytes + laneId * (kElemsPerLane / 2);
*reinterpret_cast<uint16_t*>(data) = packed;
if (laneId % kNvfp4LanesPerGroup == 0) {
int const group = laneId / kNvfp4LanesPerGroup;
int offset;
if (kv == 0) {
offset = token * kNvfp4ScaleBytes + group;
} else {
constexpr int kGroupsPerQuad = kNvfp4ScaleBytes / 4;
int const swizzled_t = (token / 4) * 4 + group / kGroupsPerQuad;
int const swizzled_s = (group % kGroupsPerQuad) * 4 + token % 4;
offset = swizzled_t * kNvfp4ScaleBytes + swizzled_s;
}
slot[block_size * kNvfp4DataBytes + offset] = sf_byte;
}
}
#endif // USE_ROCM
// ────────────────────────────────────────────────────────────────────────────
// Kernel
// ────────────────────────────────────────────────────────────────────────────
@@ -307,9 +563,8 @@ __device__ __forceinline__ void storeScaledQElemsFp8(
// branch and writes index_q/index_k outputs. Skip-index-topk reuse layers keep
// kHasIndex=true but set kProcessIndex=false.
template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
typename out_idx_t, bool kHasIndex, bool kInsertKV,
bool kProcessIndex,
bool kFp8Idx>
bool kNvfp4, typename out_idx_t, bool kHasIndex, bool kInsertKV,
bool kProcessIndex, bool kFp8Idx>
__global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse)
scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
@@ -332,7 +587,10 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
// The content (last) dim is always innermost-contiguous (stride 1), so the
// NHD/HND layout choice is captured by the head/token strides.
int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
int64_t const kv_s_dim) {
int64_t const kv_s_dim,
// NVFP4 main cache only: per-layer K/V dequantization (global) scales.
float const* __restrict__ kv_k_scale,
float const* __restrict__ kv_v_scale) {
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
// _typeConvert<BFloat16> is unavailable on pre-Ampere; the M3 kernel only
// runs with bf16/fp16 inputs in practice. Discard the bf16 body there.
@@ -352,8 +610,11 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
int const idx_slots = kProcessIndex ? niq + 1 : 0;
int const slots_per_token = nq + nkv + v_slots + idx_slots;
int const tokenIdx = globalWarpIdx / slots_per_token;
int const slot = globalWarpIdx % slots_per_token;
// Unsigned: both are non-negative, and the signed division's sign handling
// would cost every warp.
unsigned const warp_u = globalWarpIdx, slots_u = slots_per_token;
int const tokenIdx = warp_u / slots_u;
int const slot = warp_u % slots_u;
if (tokenIdx >= num_tokens) return;
// Slot boundaries.
@@ -420,10 +681,13 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
// flat TMA descriptor can address them as [tokens*heads, head_dim]; this
// folds the de-interleaving into the store the kernel already does, instead
// of a separate q.contiguous() copy. Everything else stays in place.
// Given q_fp8_out but no q_out, q is emitted only in fp8 (null store_ptr).
scalar_t* store_ptr = row_ptr;
if (isQ && q_out != nullptr) {
store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
slot * kHeadDim;
} else if (isQ && q_fp8_out != nullptr) {
store_ptr = nullptr;
} else if (isIQ && index_q_out != nullptr) {
// bf16 index_q_out: gather here. fp8: written by the explicit fp8 store.
if constexpr (!kFp8Idx) {
@@ -433,6 +697,13 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
}
}
// NVFP4: the per-layer global scale is not produced by the predecessor,
// so load it before the PDL wait.
float nvfp4_scale = 0.0f;
if constexpr (kNvfp4) {
if (isK || isV) nvfp4_scale = isK ? *kv_k_scale : *kv_v_scale;
}
// PDL: wait for the predecessor kernel (the qkv-projection GEMM that
// produces ``qkv``) to finish before touching any global memory. No-op
// when PDL is not enabled on the launch. The CUDA runtime wrapper emits
@@ -451,16 +722,17 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
cos_ptr, /*apply_norm=*/norm_w != nullptr);
if constexpr (kFp8Idx) {
// index_q is e4m3 bytes; Q/K (and in-place index_k) stay scalar_t.
// index_q is E4M3(scalar_t(q)) bytes; Q/K (and in-place index_k) stay
// scalar_t.
if (isIQ && index_q_out != nullptr) {
storeElemsFp8(index_q_out +
static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
(slot - iq_begin) * kHeadDim + dim_base,
elems);
} else {
storeIndexQElemsFp8<scalar_t>(
index_q_out + static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
(slot - iq_begin) * kHeadDim + dim_base,
elems);
} else if (store_ptr != nullptr) {
storeElems<scalar_t>(store_ptr + dim_base, elems);
}
} else {
} else if (store_ptr != nullptr) {
storeElems<scalar_t>(store_ptr + dim_base, elems);
}
if (isQ && q_fp8_out != nullptr) {
@@ -487,6 +759,16 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
} else {
storeElems<scalar_t>(index_cache + sm * kHeadDim + dim_base, elems);
}
} else if constexpr (kNvfp4) {
#ifndef USE_ROCM
// kv_cache is [num_blocks, 2*nkv, block_size, 72] bytes (HND), slot
// 2*head + side; kv_s_block is its page stride.
storeNvfp4CacheElems<scalar_t>(
reinterpret_cast<uint8_t*>(kv_cache) +
(sm / block_size) * kv_s_block,
isK ? 0 : 1, head, static_cast<int>(sm % block_size), block_size,
laneId, elems, nvfp4_scale);
#endif
} else if (isK || isV) {
// kv_cache logical shape [num_blocks, nkv, block_size, 2*head_dim].
// Paging is logical (block = sm/block_size, token = sm%block_size);
@@ -515,7 +797,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
// ────────────────────────────────────────────────────────────────────────────
// Launch wrapper
// ────────────────────────────────────────────────────────────────────────────
template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
bool kNvfp4 = false>
void launchFusedMiniMaxM3(
scalar_t* qkv, scalar_t* q_out, uint8_t* q_fp8_out, void* index_q_out,
scalar_t const* q_norm_w, scalar_t const* k_norm_w,
@@ -527,7 +810,8 @@ void launchFusedMiniMaxM3(
int const nq, int const nkv, int const niq, int const block_size,
int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
int64_t const kv_s_dim, bool const has_index, bool const insert_kv,
bool const process_index, bool const fp8_idx, cudaStream_t stream) {
bool const process_index, bool const fp8_idx, float const* kv_k_scale,
float const* kv_v_scale, cudaStream_t stream) {
// Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the
// void* pointers per instantiation in the LAUNCH macro.
// Slot count must match the kernel's compile-time gating.
@@ -559,35 +843,44 @@ void launchFusedMiniMaxM3(
config.attrs = attrs;
config.numAttrs = (sm_version >= 90) ? 1 : 0;
#define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \
cudaLaunchKernelEx( \
&config, \
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
HAS_INDEX, INSERT, \
PROCESS_INDEX, FP8>, \
qkv, q_out, q_fp8_out, reinterpret_cast<OUT_T*>(index_q_out), \
q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
slot_mapping, index_slot_mapping, kv_cache, \
reinterpret_cast<OUT_T*>(index_cache), eps, q_fp8_inv_scale, \
rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
kv_s_head, kv_s_token, kv_s_dim)
#else
// ROCm: standard kernel launch syntax (no PDL/stream serialization).
// clang-format off
#define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \
fusedMiniMaxM3QNormRopeKVInsertKernel< \
scalar_t, cache_t, kv_dt, OUT_T, HAS_INDEX, INSERT, PROCESS_INDEX, \
FP8><<<grid, kBlockSize, 0, stream>>>( \
cudaLaunchKernelEx( \
&config, \
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, \
kNvfp4, OUT_T, HAS_INDEX, \
INSERT, PROCESS_INDEX, FP8>, \
qkv, q_out, q_fp8_out, reinterpret_cast<OUT_T*>(index_q_out), \
q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
slot_mapping, index_slot_mapping, kv_cache, \
reinterpret_cast<OUT_T*>(index_cache), eps, q_fp8_inv_scale, \
rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
kv_s_head, kv_s_token, kv_s_dim)
kv_s_head, kv_s_token, kv_s_dim, kv_k_scale, kv_v_scale)
#else
// ROCm: standard kernel launch syntax (no PDL/stream serialization).
// clang-format off
#define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \
fusedMiniMaxM3QNormRopeKVInsertKernel< \
scalar_t, cache_t, kv_dt, kNvfp4, OUT_T, HAS_INDEX, INSERT, \
PROCESS_INDEX, FP8><<<grid, kBlockSize, 0, stream>>>( \
qkv, q_out, q_fp8_out, reinterpret_cast<OUT_T*>(index_q_out), \
q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
slot_mapping, index_slot_mapping, kv_cache, \
reinterpret_cast<OUT_T*>(index_cache), eps, q_fp8_inv_scale, \
rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
kv_s_head, kv_s_token, kv_s_dim, kv_k_scale, kv_v_scale)
// clang-format on
#endif
if (has_index) {
if constexpr (kNvfp4) {
// NVFP4 main cache: sparse serving only (the host checks insert_kv).
if (!process_index) {
LAUNCH(true, true, false, false, scalar_t);
} else if (fp8_idx) {
LAUNCH(true, true, true, true, uint8_t);
} else {
LAUNCH(true, true, true, false, scalar_t);
}
} else if (has_index) {
if (!process_index) {
if (insert_kv) {
LAUNCH(true, true, false, false, scalar_t);
@@ -621,8 +914,9 @@ void launchFusedMiniMaxM3(
} // namespace vllm
// clang-format off
#define CALL_FUSED_MINIMAX_M3(_RAW_T, CACHE_T, KV_DTYPE) \
vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>( \
#define CALL_FUSED_MINIMAX_M3_IMPL(CACHE_T, KV_DTYPE, NVFP4) \
vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE, \
NVFP4>( \
reinterpret_cast<st*>(qkv.data_ptr()), \
q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr()) : nullptr, \
q_fp8_out.has_value() \
@@ -654,7 +948,10 @@ void launchFusedMiniMaxM3(
static_cast<float>(eps), 1.0f / static_cast<float>(q_fp8_scale), \
static_cast<int>(rotary_dim), num_tokens, nq, nkv, niq, \
static_cast<int>(block_size), kv_s_block, kv_s_head, kv_s_token, \
kv_s_dim, has_index, insert_kv, process_index, fp8_idx, stream)
kv_s_dim, has_index, insert_kv, process_index, fp8_idx, kv_k_scale_ptr, \
kv_v_scale_ptr, stream)
#define CALL_FUSED_MINIMAX_M3(_RAW_T, CACHE_T, KV_DTYPE) \
CALL_FUSED_MINIMAX_M3_IMPL(CACHE_T, KV_DTYPE, false)
// clang-format on
// ────────────────────────────────────────────────────────────────────────────
@@ -681,7 +978,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
const std::string& kv_cache_dtype, bool skip_index_branch,
std::optional<torch::stable::Tensor>
q_fp8_out, // [N, nq*128] contiguous E4M3
double q_fp8_scale) {
double q_fp8_scale,
std::optional<torch::stable::Tensor> kv_k_scale, // [] f32, NVFP4 only
std::optional<torch::stable::Tensor> kv_v_scale) {
STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
"qkv must be contiguous CUDA");
STD_TORCH_CHECK(
@@ -721,8 +1020,10 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
bool const has_index = niq > 0;
bool const insert_kv = kv_cache.has_value();
bool const process_index = has_index && !skip_index_branch;
bool const nvfp4_kv = kv_cache_dtype == "nvfp4";
vllm::Fp8KVCacheDataType const kv_dt =
vllm::get_fp8_kv_cache_data_type(kv_cache_dtype);
nvfp4_kv ? vllm::Fp8KVCacheDataType::kAuto
: vllm::get_fp8_kv_cache_data_type(kv_cache_dtype);
int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
int const expected_row =
(nq + 2 * nkv + (has_index ? niq + 1 : 0)) * kHeadDim;
@@ -769,8 +1070,13 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
index_slot_mapping->numel() == slot_mapping->numel()),
"index_slot_mapping must be int64 CUDA with slot_mapping length");
}
// Main attention KV cache: auto matches qkv, fp8 uses uint8 storage.
if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
// Main attention KV cache: auto matches qkv, fp8 and nvfp4 use uint8
// storage.
if (nvfp4_kv) {
STD_TORCH_CHECK(
kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte,
"nvfp4 kv_cache must use uint8 storage");
} else if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
"auto kv_cache dtype must match qkv");
} else {
@@ -790,6 +1096,27 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
STD_TORCH_CHECK(kv_cache->dim() == 4 && kv_cache->stride(3) == 1,
"kv_cache must be [nb,nkv,bs,2*head_dim] with contiguous "
"content dim (stride(3)==1)");
if (nvfp4_kv) {
// Packed HND pages [nb, 2*nkv, bs, 72]: slot 2*head + side is one
// contiguous [bs, 72] block, which the store addresses directly.
int64_t const full_dim = kHeadDim / 2 + kHeadDim / 16;
STD_TORCH_CHECK(
kv_cache->size(1) == 2 * nkv && kv_cache->size(2) == block_size &&
kv_cache->size(3) == full_dim &&
kv_cache->stride(2) == full_dim &&
kv_cache->stride(1) == block_size * full_dim &&
block_size % 4 == 0,
"nvfp4 kv_cache must be HND [nb, 2*num_kv_heads, block_size, 72] "
"with contiguous pages and block_size % 4 == 0");
for (auto const* scale : {&kv_k_scale, &kv_v_scale}) {
STD_TORCH_CHECK(
scale->has_value() && (*scale)->is_cuda() &&
(*scale)->scalar_type() ==
torch::headeronly::ScalarType::Float &&
(*scale)->numel() >= 1,
"nvfp4 kv_cache requires float32 CUDA kv_k_scale and kv_v_scale");
}
}
kv_s_block = kv_cache->stride(0);
kv_s_head = kv_cache->stride(1);
kv_s_token = kv_cache->stride(2);
@@ -854,10 +1181,31 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
"fp8 index path: index_q_out must be fp8 e4m3");
}
STD_TORCH_CHECK(!nvfp4_kv || insert_kv,
"nvfp4 kv_cache_dtype requires kv_cache (insert mode)");
float const* kv_k_scale_ptr =
nvfp4_kv ? kv_k_scale->const_data_ptr<float>() : nullptr;
float const* kv_v_scale_ptr =
nvfp4_kv ? kv_v_scale->const_data_ptr<float>() : nullptr;
const torch::stable::accelerator::DeviceGuard device_guard(
qkv.get_device_index());
auto stream = get_current_cuda_stream(qkv.get_device_index());
if (nvfp4_kv) {
#ifndef USE_ROCM
VLLM_STABLE_DISPATCH_HALF_TYPES(
qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert_nvfp4", [&] {
using st = scalar_t;
CALL_FUSED_MINIMAX_M3_IMPL(uint8_t, vllm::Fp8KVCacheDataType::kAuto,
true);
});
#else
STD_TORCH_CHECK(false, "nvfp4 kv_cache is not supported on ROCm");
#endif
return;
}
VLLM_STABLE_DISPATCH_HALF_TYPES(
qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert", [&] {
using st = scalar_t;
@@ -867,3 +1215,4 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
}
#undef CALL_FUSED_MINIMAX_M3
#undef CALL_FUSED_MINIMAX_M3_IMPL
csrc/libtorch_stable/ops.h
+3 -1
View File
@@ -381,7 +381,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
std::optional<torch::stable::Tensor> q_out,
std::optional<torch::stable::Tensor> index_q_out,
const std::string& kv_cache_dtype, bool skip_index_branch,
std::optional<torch::stable::Tensor> q_fp8_out, double q_fp8_scale);
std::optional<torch::stable::Tensor> q_fp8_out, double q_fp8_scale,
std::optional<torch::stable::Tensor> kv_k_scale,
std::optional<torch::stable::Tensor> kv_v_scale);
#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
void fused_kda_decode(
csrc/libtorch_stable/torch_bindings.cpp
+2 -1
View File
@@ -513,7 +513,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"Tensor!? kv_cache, Tensor!? index_cache, "
"int block_size, Tensor!? q_out, Tensor!? index_q_out, "
"str kv_cache_dtype, bool skip_index_branch=False, "
"Tensor!? q_fp8_out=None, float q_fp8_scale=1.0) -> ()");
"Tensor!? q_fp8_out=None, float q_fp8_scale=1.0, "
"Tensor? kv_k_scale=None, Tensor? kv_v_scale=None) -> ()");
#ifdef VLLM_ENABLE_FUSED_KDA_DECODE
ops.def(
tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
+441 -1
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Correctness tests for MiniMax M3 CUTLASS sparse decode."""
"""Correctness tests for MiniMax M3 MSA sparse attention (CUTLASS decode and
NVFP4 KV cache)."""
import math
from types import SimpleNamespace
@@ -8,6 +9,7 @@ from types import SimpleNamespace
import pytest
import torch
from tests.kernels.quantization.nvfp4_utils import dequant_nvfp4_kv_cache
from vllm import _custom_ops as ops
from vllm.config import AttentionConfig, CUDAGraphMode
from vllm.forward_context import ForwardContext, override_forward_context
@@ -18,6 +20,7 @@ from vllm.models.minimax_m3.common.sparse_attention import (
MiniMaxM3SparseBackend,
MiniMaxM3SparseMetadata,
MiniMaxM3SparseMetadataBuilder,
MiniMaxM3SparsePrefillMetadata,
MiniMaxM3SparseTritonImpl,
select_main_backend_and_impl_cls,
)
@@ -28,6 +31,7 @@ from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention
from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import (
MSACutlassDecodePlanCache,
msa_cutlass_sparse_decode,
nvfp4_kv_cache_views,
prepare_decode_metadata,
should_prepare_decode_metadata,
)
@@ -36,6 +40,7 @@ from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
MiniMaxM3SparseMSADecodeMetadata,
MiniMaxM3SparseMSAImpl,
MiniMaxM3SparseMSAMetadataBuilder,
MiniMaxM3SparseMSANvfp4Backend,
)
from vllm.platforms import current_platform
from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -164,6 +169,53 @@ def test_msa_cutlass_decode_static_dispatch_requires_sm100(
)
@pytest.mark.parametrize("decode_backend", ["cutlass", "triton"])
def test_msa_cutlass_decode_static_dispatch_nvfp4_has_no_fallback(
decode_backend: str,
) -> None:
"""NVFP4 has no Triton decode, so it plans CUTLASS for any batch size."""
assert should_prepare_decode_metadata(
1,
DEFAULT_QUERY_LEN,
decode_backend=decode_backend, # type: ignore[arg-type]
num_q_heads=16,
num_kv_heads=1,
kv_cache_dtype="nvfp4",
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
)
def test_nvfp4_selects_msa_nvfp4_backend_and_requires_sm100(
monkeypatch: pytest.MonkeyPatch,
) -> None:
backend_cls, impl_cls = select_main_backend_and_impl_cls(
topk_blocks=TOPK,
kv_cache_dtype="nvfp4",
num_kv_heads=1,
)
assert backend_cls is MiniMaxM3SparseMSANvfp4Backend
assert impl_cls is MiniMaxM3SparseMSAImpl
with pytest.raises(ValueError, match="nvfp4_4over6"):
select_main_backend_and_impl_cls(
topk_blocks=TOPK,
kv_cache_dtype="nvfp4_4over6",
num_kv_heads=1,
)
monkeypatch.setattr(
current_platform,
"is_device_capability_family",
lambda _: False,
)
with pytest.raises(ValueError, match="NVFP4"):
select_main_backend_and_impl_cls(
topk_blocks=TOPK,
kv_cache_dtype="nvfp4",
num_kv_heads=1,
)
@pytest.mark.parametrize(
("backend", "expected"),
[
@@ -639,3 +691,391 @@ def test_msa_cutlass_decode_matches_triton_with_interleaved_cache(
graph.replay()
current_platform.synchronize()
torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02)
def _make_paged_layout(
seq_lens_list: list[int],
) -> tuple[torch.Tensor, int]:
"""Block table over a shuffled physical page pool."""
pages_per_request = [math.ceil(seq_len / BLOCK_SIZE) for seq_len in seq_lens_list]
num_pages = sum(pages_per_request)
block_table = torch.zeros(
len(seq_lens_list), max(pages_per_request), dtype=torch.int32, device="cuda"
)
physical_pages = torch.randperm(num_pages, dtype=torch.int32, device="cuda")
offset = 0
for request, request_pages in enumerate(pages_per_request):
block_table[request, :request_pages] = physical_pages[
offset : offset + request_pages
]
offset += request_pages
return block_table, num_pages
def _make_nvfp4_kv_cache(
num_pages: int,
num_kv_heads: int,
kv_cache_dtype: str,
k_scale: torch.Tensor,
v_scale: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Write random K/V through the NVFP4 ``reshape_and_cache_flash`` writer,
whose pages hold every head's data then every head's scales per side, and
repack them into head slots.
Returns the packed HND cache ``[pages, 2 * Hkv, 128, 72]`` and its BF16
dequantization ``[pages, Hkv, 128, 2 * 128]`` for the Triton reference.
"""
num_slots = num_pages * BLOCK_SIZE
key = torch.randn(
num_slots, num_kv_heads, HEAD_DIM, dtype=torch.bfloat16, device="cuda"
)
value = torch.randn_like(key)
written = torch.zeros(
num_pages,
2 * num_kv_heads,
BLOCK_SIZE,
HEAD_DIM // 2 + HEAD_DIM // 16,
dtype=torch.uint8,
device="cuda",
)
k_cache, v_cache = written.transpose(1, 2).split(num_kv_heads, dim=-2)
ops.reshape_and_cache_flash(
key,
value,
k_cache,
v_cache,
torch.arange(num_slots, dtype=torch.int64, device="cuda"),
kv_cache_dtype,
k_scale,
v_scale,
)
kv_cache = torch.empty_like(written)
head_slot_views = iter(nvfp4_kv_cache_views(kv_cache))
dequantized = []
for side, global_scale, swizzled in (
(written[:, :num_kv_heads], k_scale, False),
(written[:, num_kv_heads:], v_scale, True),
):
data = side.as_strided(
(num_pages, num_kv_heads, BLOCK_SIZE, HEAD_DIM // 2),
(side.stride(0), BLOCK_SIZE * HEAD_DIM // 2, HEAD_DIM // 2, 1),
)
scale = side.as_strided(
(num_pages, num_kv_heads, BLOCK_SIZE, HEAD_DIM // 16),
(side.stride(0), BLOCK_SIZE * HEAD_DIM // 16, HEAD_DIM // 16, 1),
side.storage_offset() + num_kv_heads * BLOCK_SIZE * HEAD_DIM // 2,
)
next(head_slot_views).copy_(data)
next(head_slot_views).copy_(scale)
dequantized.append(
dequant_nvfp4_kv_cache(
data,
scale,
global_scale.item(),
HEAD_DIM,
BLOCK_SIZE,
swizzled_scales=swizzled,
)
)
return kv_cache, torch.cat(dequantized, dim=-1).to(torch.bfloat16)
def test_nvfp4_kv_cache_views_address_head_slots() -> None:
"""Slot 2h (K) and 2h + 1 (V) hold head h's data then scales, so each head
is one contiguous run of the page: a TP rank's page is a byte slice."""
num_kv_heads, slot_bytes = 4, BLOCK_SIZE * (HEAD_DIM // 2 + HEAD_DIM // 16)
kv_cache = torch.randint(
0,
256,
(3, 2 * num_kv_heads, BLOCK_SIZE, HEAD_DIM // 2 + HEAD_DIM // 16),
dtype=torch.uint8,
)
k_data, k_sf, v_data, v_sf = nvfp4_kv_cache_views(kv_cache)
runs = kv_cache.view(3, num_kv_heads, 2 * slot_bytes)
for head in range(num_kv_heads):
expected = torch.cat(
[v[:, head].flatten(1) for v in (k_data, k_sf, v_data, v_sf)], dim=1
)
assert torch.equal(runs[:, head], expected)
@pytest.mark.parametrize(
("num_q_heads", "num_kv_heads", "num_request_pairs", "query_len", "capture"),
[
pytest.param(64, 4, 1, 1, True, id="tp1-small-batch-query-len-1"),
pytest.param(64, 4, 8, 8, False, id="tp1-query-len-8"),
pytest.param(16, 1, 16, 4, True, id="tp4-query-len-4"),
],
)
def test_msa_cutlass_decode_nvfp4_matches_triton_on_dequantized_cache(
num_q_heads: int,
num_kv_heads: int,
num_request_pairs: int,
query_len: int,
capture: bool,
) -> None:
"""CUTLASS decode reads vLLM's packed NVFP4 pages, including swizzled V
scales and device global scales, as their exact dequantized values."""
torch.manual_seed(0)
seq_lens_list = [257, 513] * num_request_pairs
seq_lens_cpu = torch.tensor(seq_lens_list, dtype=torch.int32)
seq_lens = seq_lens_cpu.cuda()
block_table, num_pages = _make_paged_layout(seq_lens_list)
k_scale = torch.tensor(0.5, dtype=torch.float32, device="cuda")
v_scale = torch.tensor(2.0, dtype=torch.float32, device="cuda")
kv_cache, kv_reference = _make_nvfp4_kv_cache(
num_pages, num_kv_heads, "nvfp4", k_scale, v_scale
)
num_query_tokens = len(seq_lens_list) * query_len
query = torch.randn(
num_query_tokens, num_q_heads, HEAD_DIM, dtype=torch.bfloat16, device="cuda"
)
query_fp8 = query.to(torch.float8_e4m3fn)
topk = _make_topk(seq_lens_list, num_kv_heads, query_len)
expected = torch.empty_like(query)
minimax_m3_sparse_attn_decode(
query_fp8.to(torch.bfloat16),
kv_reference,
topk.transpose(0, 1),
block_table,
seq_lens,
num_kv_heads,
SM_SCALE,
expected,
query_len,
)
assert should_prepare_decode_metadata(
len(seq_lens_list),
query_len,
decode_backend="triton",
num_q_heads=num_q_heads,
num_kv_heads=num_kv_heads,
kv_cache_dtype="nvfp4",
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
)
metadata = prepare_decode_metadata(
block_table,
seq_lens,
seq_lens_cpu,
query_len,
num_q_heads=num_q_heads,
num_kv_heads=num_kv_heads,
page_size=BLOCK_SIZE,
topk_blocks=TOPK,
)
actual = torch.zeros_like(query)
def run() -> None:
msa_cutlass_sparse_decode(
query_fp8,
kv_cache,
topk,
actual,
metadata,
scale=SM_SCALE,
q_scale_float=1.0,
k_scale_float=1.0,
v_scale_float=1.0,
k_scale=k_scale,
v_scale=v_scale,
)
if capture:
run()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
run()
actual.zero_()
graph.replay()
else:
run()
current_platform.synchronize()
torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.05)
@pytest.mark.parametrize(
("num_q_heads", "num_kv_heads"),
[pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")],
)
def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache(
num_q_heads: int,
num_kv_heads: int,
) -> None:
"""MSA prefill reads NVFP4 pages as their exact dequantized values,
including chunked prefill over cached context, and takes q from its fp8
copy (the only one the fused insert writes on this path)."""
torch.manual_seed(0)
q_lens = [300, 128, 77]
seq_lens_list = [300, 640, 1100]
block_table, num_pages = _make_paged_layout(seq_lens_list)
k_scale = torch.tensor(0.5, dtype=torch.float32, device="cuda")
v_scale = torch.tensor(2.0, dtype=torch.float32, device="cuda")
kv_cache, kv_reference = _make_nvfp4_kv_cache(
num_pages, num_kv_heads, "nvfp4", k_scale, v_scale
)
total_q = sum(q_lens)
seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda")
context_lens = seq_lens - torch.tensor(q_lens, dtype=torch.int32, device="cuda")
cu_seqlens_q = torch.zeros(len(q_lens) + 1, dtype=torch.int32, device="cuda")
cu_seqlens_q[1:] = torch.tensor(q_lens, device="cuda").cumsum(0)
cu_seqlens_k = torch.zeros_like(cu_seqlens_q)
cu_seqlens_k[1:] = seq_lens.cumsum(0)
prefill = MiniMaxM3SparsePrefillMetadata(
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
seq_lens=seq_lens,
context_lens=context_lens,
block_table=block_table,
max_query_len=max(q_lens),
max_seq_len=max(seq_lens_list),
total_kv_blocks=sum(math.ceil(n / BLOCK_SIZE) for n in seq_lens_list),
)
metadata = MiniMaxM3SparseMetadata(
seq_lens=seq_lens,
max_seq_len=max(seq_lens_list),
slot_mapping=torch.empty(total_q, dtype=torch.int64, device="cuda"),
num_actual_tokens=total_q,
num_decodes=0,
num_decode_tokens=0,
num_prefills=len(q_lens),
num_prefill_tokens=total_q,
prefill=prefill,
)
# Every causally visible page (<= TOPK per query), ending with the local one.
topk = torch.full(
(total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device="cuda"
)
token = 0
for q_len, seq_len in zip(q_lens, seq_lens_list):
for local_q in range(q_len):
visible = (seq_len - q_len + local_q) // BLOCK_SIZE + 1
assert visible <= TOPK
topk[token, :, :visible] = torch.arange(visible, device="cuda")
token += 1
layer_name = "model.layers.0.self_attn.attn"
q_scale = 0.5
layer = SimpleNamespace(
layer_name=layer_name,
topk_indices_buffer=topk,
_q_scale_float=q_scale,
_k_scale=k_scale,
_v_scale=v_scale,
)
query_fp8 = (
torch.randn(
total_q, num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda"
)
/ q_scale
).to(torch.float8_e4m3fn)
query = query_fp8.to(torch.bfloat16) * q_scale
common = dict(
num_heads=num_q_heads,
head_size=HEAD_DIM,
scale=SM_SCALE,
num_kv_heads=num_kv_heads,
topk_blocks=TOPK,
sparse_block_size=BLOCK_SIZE,
)
forward_context = ForwardContext(
no_compile_layers={},
attn_metadata={layer_name: metadata},
slot_mapping={},
)
with override_forward_context(forward_context):
expected = MiniMaxM3SparseTritonImpl(**common).forward(
layer, query, kv_reference, torch.empty_like(query)
)
actual = MiniMaxM3SparseMSAImpl(**common, kv_cache_dtype="nvfp4").forward(
layer,
torch.full_like(query, float("nan")),
kv_cache,
torch.zeros_like(query),
query_fp8=query_fp8,
)
current_platform.synchronize()
torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.05)
def test_msa_triton_decode_fallback_reads_dequantized_query_fp8() -> None:
"""With CUTLASS decode on but not planned for the batch, the Triton decode
fallback attends with the dequantized fp8 q, the only q written."""
torch.manual_seed(0)
num_q_heads, num_kv_heads, q_scale = 16, 1, 0.5
seq_lens_list = [257, 513, 129]
seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda")
block_table, num_pages = _make_paged_layout(seq_lens_list)
kv_cache = (
torch.randn(
num_pages,
num_kv_heads,
BLOCK_SIZE,
2 * HEAD_DIM,
dtype=torch.bfloat16,
device="cuda",
)
* 0.25
).to(torch.float8_e4m3fn)
num_tokens = len(seq_lens_list)
metadata = MiniMaxM3SparseMetadata(
seq_lens=seq_lens,
max_seq_len=max(seq_lens_list),
slot_mapping=torch.empty(num_tokens, dtype=torch.int64, device="cuda"),
num_actual_tokens=num_tokens,
num_decodes=num_tokens,
num_decode_tokens=num_tokens,
num_prefills=0,
num_prefill_tokens=0,
decode=MiniMaxM3SparseMSADecodeMetadata(
seq_lens=seq_lens,
block_table=block_table,
decode_query_len=1,
msa_cutlass=None,
),
)
layer_name = "model.layers.0.self_attn.attn"
layer = SimpleNamespace(
layer_name=layer_name,
topk_indices_buffer=_make_topk(seq_lens_list, num_kv_heads, 1),
_q_scale_float=q_scale,
)
query_fp8 = torch.randn(
num_tokens, num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda"
).to(torch.float8_e4m3fn)
query = query_fp8.to(torch.bfloat16) * q_scale
common = dict(
num_heads=num_q_heads,
head_size=HEAD_DIM,
scale=SM_SCALE,
num_kv_heads=num_kv_heads,
kv_cache_dtype="fp8",
topk_blocks=TOPK,
sparse_block_size=BLOCK_SIZE,
)
impl = MiniMaxM3SparseMSAImpl(**common, msa_decode_backend="cutlass")
assert impl.use_cutlass_decode
forward_context = ForwardContext(
no_compile_layers={},
attn_metadata={layer_name: metadata},
slot_mapping={},
)
with override_forward_context(forward_context):
expected = MiniMaxM3SparseTritonImpl(**common).forward(
layer, query, kv_cache, torch.empty_like(query)
)
actual = impl.forward(
layer,
torch.full_like(query, float("nan")),
kv_cache,
torch.zeros_like(query),
query_fp8=query_fp8,
)
current_platform.synchronize()
torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02)
tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
+226 -4
View File
@@ -320,6 +320,38 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0
)
# Without q_out, q is emitted only in fp8: the same bytes, while qkv keeps
# its raw q and every other output matches the run above.
outputs = [qkv_orig.clone(), torch.empty_like(q_fp8), torch.empty_like(index_q)]
caches = [torch.zeros_like(kv_cache), torch.zeros_like(index_cache)]
ops.fused_minimax_m3_qknorm_rope_kv_insert(
outputs[0],
q_w,
k_w,
cos_sin,
positions,
num_heads,
num_kv_heads,
ROTARY_DIM,
eps,
iq_w,
ik_w,
num_idx_heads,
slot_mapping,
index_slot_mapping,
*caches,
block_size,
None,
outputs[2],
kv_cache_dtype,
q_fp8_out=outputs[1],
q_fp8_scale=0.5,
)
for actual, expected in zip(
outputs + caches, [qkv, q_fp8, index_q, kv_cache, index_cache]
):
assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
@pytest.mark.parametrize("num_tokens", [1, 64, 513])
@pytest.mark.parametrize("block_size", [16, 64])
@@ -441,8 +473,8 @@ def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype):
# ── Test 4: fp8 (e4m3) index outputs ─────────────────────────────────────────
# The fp8 score path stores index_q and the index-K cache as e4m3 while q/k/v +
# q_out stay bf16. Asserts: (1) q/k/v/q_out are bit-identical to the bf16 run
# (the index dtype must not perturb the main branch), and (2) the e4m3 index
# outputs dequantize close to the bf16 reference.
# (the index dtype must not perturb the main branch), and (2) index_q is the
# e4m3 cast of the bf16 run's index_q, and the index-K cache is within fp8 ulp.
@pytest.mark.skipif(
@@ -528,6 +560,196 @@ def test_sparse_full_fp8_index(num_tokens, block_size):
torch.testing.assert_close(qkv_fp, qkv_bf, rtol=0, atol=0)
torch.testing.assert_close(kvc_fp, kvc_bf, rtol=0, atol=0)
# (2) Dequantized e4m3 index outputs match the bf16 reference within fp8 ulp.
torch.testing.assert_close(iq_fp.float(), iq_bf.float(), rtol=0.13, atol=0.05)
# (2) index_q is e4m3(bf16(q)); the index-K cache is within fp8 ulp.
assert torch.equal(
iq_fp.view(torch.uint8), iq_bf.to(torch.float8_e4m3fn).view(torch.uint8)
)
torch.testing.assert_close(idxc_fp.float(), idxc_bf.float(), rtol=0.13, atol=0.05)
# ── Test 5: sparse mode with an NVFP4 main KV cache ──────────────────────────
# The NVFP4 inserts quantize x * sf_scale, where sf_scale = 1 / the global
# scale: per 16-value block,
# scale = min(448, max(E4M3(max(amax, 1e-12) / 6), 1/512)) and
# code = sign | RNE E2M1 of min(|x / scale|, 6), with zero stored as +0.
E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)
def _div(a, b):
# Division by a Python scalar would become a reciprocal multiply.
return a / torch.as_tensor(b, dtype=a.dtype, device=a.device)
def nvfp4_ref(x, sf_scale):
"""Returns the E2M1 codes [..., 128] and E4M3 scale bytes [..., 8] of x."""
blocks = x.float().unflatten(-1, (-1, 16))
amax = blocks.abs().amax(-1, keepdim=True).clamp_min(1e-12)
scale = (sf_scale * _div(amax, 6.0)).clamp(max=448.0)
scale = scale.to(torch.float8_e4m3fn).float().clamp(min=2.0**-9)
q = _div(blocks * sf_scale, scale)
# Alternating > / >= rounds exact midpoints to the even code.
mag = sum(
(q.abs() >= m) if i % 2 else (q.abs() > m) for i, m in enumerate(E2M1_MIDPOINTS)
).to(torch.uint8)
codes = torch.where(mag > 0, mag | (torch.signbit(q).to(torch.uint8) << 3), 0)
scale_bytes = scale.squeeze(-1).to(torch.float8_e4m3fn).view(torch.uint8)
return codes.flatten(-2), scale_bytes
def nvfp4_pages(sides, slot_mapping, num_blocks, block_size):
"""Expected NVFP4 pages [nb, len(sides) * nkv, bs, 72] for the K (then V)
side given as (x [N, nkv, 128], sf_scale). Slot h * len(sides) + side is
head h's packed E2M1 data [bs, 64] (low nibble = even element), then its
E4M3 scales [bs, 8]: linear for K, 4x4 token-quad swizzled for V."""
nkv, ns, device = sides[0][0].shape[1], HEAD_DIM // 16, slot_mapping.device
row_bytes, data_bytes = HEAD_DIM // 2 + ns, block_size * HEAD_DIM // 2
pages = torch.zeros(
num_blocks,
nkv,
len(sides),
block_size * row_bytes,
dtype=torch.uint8,
device=device,
)
valid = slot_mapping >= 0
blk, tok = slot_mapping[valid] // block_size, slot_mapping[valid] % block_size
g, t = torch.arange(ns, device=device), tok[:, None]
scale_offsets = (t * ns + g, ((t // 4) * 4 + g // 2) * ns + (g % 2) * 4 + t % 4)
for side, (x, sf_scale) in enumerate(sides):
codes, scales = nvfp4_ref(x[valid], sf_scale)
data = pages[:, :, side, :data_bytes].view(num_blocks, nkv, block_size, -1)
data[blk, :, tok] = codes[..., 0::2] | (codes[..., 1::2] << 4)
scale_region = pages[:, :, side, data_bytes:]
for h in range(nkv):
scale_region[blk[:, None], h, scale_offsets[side]] = scales[:, h]
return pages.view(num_blocks, -1, block_size, row_bytes)
def nvfp4_edge_values(num_rows, global_scale):
"""[num_rows, 128] bf16 whose 16-value blocks, divided by global_scale, hit
the quantizer's corner cases: exact E2M1 midpoints under non-power-of-2 scales
(blocks 0-1), amax under the 1/512 scale floor (2), all -0 (3), and
saturation with |x| > 4096 up to bf16 max (4) or with inf (5)."""
g = torch.Generator().manual_seed(0)
x = torch.randn(num_rows, 8, 16, generator=g)
sign = torch.where(torch.rand(x.shape, generator=g) < 0.5, -1.0, 1.0)
steps = torch.tensor(E2M1_MIDPOINTS + (0.5, 1.0, 3.0))
mant = 1 + torch.randint(1, 8, (num_rows, 2, 1), generator=g) / 8
scale = mant * 2.0 ** torch.randint(-6, 3, (num_rows, 2, 1), generator=g)
x[:, :2] = steps[torch.randint(len(steps), (num_rows, 2, 16), generator=g)] * scale
x[:, :2, :1] = 6 * scale # amax = 6 * scale, so the block scale is `scale`
x[:, 2] = torch.rand(num_rows, 16, generator=g) * 2.0**-8
x[:, 4] *= 5000
x = x * sign * global_scale
x[:, 3] = -0.0
x[:, 4, 0] = sign[:, 4, 0] * torch.finfo(torch.bfloat16).max
x[:, 5, 0] = sign[:, 5, 0] * float("inf")
return x.flatten(1).to(torch.bfloat16)
@pytest.mark.skipif(
not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
reason="NVFP4 KV caches need SM100",
)
@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
@pytest.mark.parametrize("num_kv_heads", [1, 4])
@pytest.mark.parametrize("block_size", [16, 128])
@pytest.mark.parametrize("kv_scales", [(0.5, 2.0), (1.0, 1.0)])
def test_sparse_full_nvfp4(num_tokens, num_kv_heads, block_size, kv_scales):
"""The fused NVFP4 insert writes the reference quantizer's bytes (packed
E2M1, linear K / swizzled V E4M3 block scales) for the kernel's own
post-RoPE k and for a raw v that hits the quantizer's corner cases, and
skips padded slots. Unit global scales take the kernel's separate fast
path."""
kv_cache_dtype = "nvfp4"
torch.manual_seed(5)
device, dtype, eps = "cuda", torch.bfloat16, 1e-6
base, max_pos = 5_000_000.0, 4096
num_heads, num_idx_heads = 4 * num_kv_heads, 1
weights = [
torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1 for _ in range(4)
]
cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
positions = torch.randint(
0, max_pos, (num_tokens,), dtype=torch.int64, device=device
)
qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
iqsz = num_idx_heads * HEAD_DIM
qkv = torch.randn(
num_tokens, qsz + 2 * kvsz + iqsz + HEAD_DIM, dtype=dtype, device=device
)
k_scale, v_scale = (
torch.tensor(s, dtype=torch.float32, device=device) for s in kv_scales
)
qkv[:, qsz + kvsz : qsz + 2 * kvsz] = nvfp4_edge_values(
num_tokens * num_kv_heads, v_scale.item()
).view(num_tokens, kvsz)
qkv_orig = qkv.clone()
num_blocks = (num_tokens + block_size - 1) // block_size + 1
kv_cache = torch.zeros(
num_blocks,
2 * num_kv_heads,
block_size,
HEAD_DIM // 2 + HEAD_DIM // 16,
dtype=torch.uint8,
device=device,
)
index_cache = torch.zeros(
num_blocks, block_size, HEAD_DIM, dtype=dtype, device=device
)
slot_mapping = torch.randperm(
num_blocks * block_size, dtype=torch.int64, device=device
)[:num_tokens]
slot_mapping[::5] = -1 # padded / unscheduled tokens are not written
index_slot_mapping = torch.randperm(
num_blocks * block_size, dtype=torch.int64, device=device
)[:num_tokens]
ops.fused_minimax_m3_qknorm_rope_kv_insert(
qkv,
*weights[:2],
cos_sin,
positions,
num_heads,
num_kv_heads,
ROTARY_DIM,
eps,
*weights[2:],
num_idx_heads,
slot_mapping,
index_slot_mapping,
kv_cache,
index_cache,
block_size,
torch.empty(num_tokens, qsz, dtype=dtype, device=device),
torch.empty(num_tokens, iqsz, dtype=dtype, device=device),
kv_cache_dtype,
kv_k_scale=k_scale,
kv_v_scale=v_scale,
)
_, k_out, v_out, _, index_k = qkv.split([qsz, kvsz, kvsz, iqsz, HEAD_DIM], dim=-1)
torch.testing.assert_close(
v_out, qkv_orig.split([qsz, kvsz, kvsz, iqsz, HEAD_DIM], dim=-1)[2]
)
expected_kv_cache = nvfp4_pages(
[
(k_out.reshape(num_tokens, num_kv_heads, HEAD_DIM), 1 / k_scale),
(v_out.reshape(num_tokens, num_kv_heads, HEAD_DIM), 1 / v_scale),
],
slot_mapping,
num_blocks,
block_size,
)
assert torch.equal(kv_cache, expected_kv_cache), (
f"{int((kv_cache != expected_kv_cache).sum())} NVFP4 cache bytes differ"
)
expected_index_cache = torch.zeros_like(index_cache).view(-1, HEAD_DIM)
expected_index_cache[index_slot_mapping] = index_k
torch.testing.assert_close(
index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0
)
vllm/_custom_ops.py
+12 -1
View File
@@ -2587,6 +2587,8 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
skip_index_branch: bool = False,
q_fp8_out: torch.Tensor | None = None,
q_fp8_scale: float = 1.0,
kv_k_scale: torch.Tensor | None = None,
kv_v_scale: torch.Tensor | None = None,
) -> None:
"""Fused MiniMax-M3 attention pre-processing (in-place).
@@ -2610,7 +2612,14 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
attention's flat TMA descriptor.
If ``q_fp8_out`` is given, the same normalized q is also written in FP8
E4M3 using ``q_fp8_scale`` as its dequantization scale.
E4M3 using ``q_fp8_scale`` as its dequantization scale. Without ``q_out``,
q is then written only in FP8 and the q slice of ``qkv`` is left as is.
``kv_cache_dtype="nvfp4"`` quantizes k/v into the packed HND
``[num_blocks, 2 * num_kv_heads, block_size, 72]`` uint8 cache using the
device dequantization scales ``kv_k_scale``/``kv_v_scale``. Slot
``2 * head + side`` (K = 0, V = 1) holds that head's E2M1 data followed by
its E4M3 block scales, so every head is one contiguous run of the page.
When ``skip_index_branch`` is true, sparse rows still keep their packed
``[index_q | index_k]`` tail, but the kernel only processes the main q/k/v
@@ -2641,6 +2650,8 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
skip_index_branch,
q_fp8_out,
q_fp8_scale,
kv_k_scale,
kv_v_scale,
)
vllm/models/minimax_m3/common/sparse_attention.py
+19 -3
View File
@@ -56,6 +56,7 @@ from vllm.v1.attention.backends.utils import (
from vllm.v1.kv_cache_interface import (
AttentionSpec,
KVCacheLayout,
get_kv_quant_mode,
is_quantized_kv_cache,
)
@@ -531,10 +532,10 @@ def select_main_backend_and_impl_cls(
"""Pick the main attention backend and implementation.
Blackwell (SM100) uses the MSA attend for supported top-k block counts
when the KV cache is BF16 or FP8 E4M3; MI355 uses AITER sparse PA
when the KV cache is BF16, FP8 E4M3 or NVFP4; MI355 uses AITER sparse PA
with shuffle KV cache layout; Other platforms and FP8 E5M2 fall
back to Triton. The MSA modules are imported lazily to avoid import errors
on unsupported platforms.
back to Triton. NVFP4 is MSA-only. The MSA modules are imported lazily to
avoid import errors on unsupported platforms.
"""
use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(
num_kv_heads, emits_sparse_block_table=emits_sparse_block_table
@@ -545,6 +546,18 @@ def select_main_backend_and_impl_cls(
and topk_blocks in (4, 8, 16, 32)
and kv_cache_dtype != "fp8_e5m2"
)
use_nvfp4 = get_kv_quant_mode(kv_cache_dtype).is_nvfp4
if use_nvfp4 and kv_cache_dtype != "nvfp4":
raise ValueError(
"MiniMax M3 sparse attention supports only the 'nvfp4' NVFP4 KV "
f"cache dtype, got {kv_cache_dtype}"
)
if use_nvfp4 and (use_aiter_sparse_pa or not use_msa):
raise ValueError(
"MiniMax M3 sparse attention supports an NVFP4 KV cache only with "
f"the SM100 MSA backend (kv_cache_dtype={kv_cache_dtype}, "
f"topk_blocks={topk_blocks})"
)
selected = (
"AITER_SPARSE_PA" if use_aiter_sparse_pa else ("MSA" if use_msa else "Triton")
)
@@ -565,8 +578,11 @@ def select_main_backend_and_impl_cls(
from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
MiniMaxM3SparseMSABackend,
MiniMaxM3SparseMSAImpl,
MiniMaxM3SparseMSANvfp4Backend,
)
if use_nvfp4:
return MiniMaxM3SparseMSANvfp4Backend, MiniMaxM3SparseMSAImpl
return MiniMaxM3SparseMSABackend, MiniMaxM3SparseMSAImpl
return MiniMaxM3SparseBackend, MiniMaxM3SparseTritonImpl
vllm/models/minimax_m3/nvidia/model.py
+7 -1
View File
@@ -500,6 +500,8 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
# of wrapping the generic Attention module. Keep the same runtime scale
# attributes so FP8 KV reads can honor vLLM's per-layer descale contract.
set_default_quant_scales(self, register_buffer=True)
# The fused insert quantizes k/v into NVFP4 pages with these scales.
self.use_nvfp4_kv = get_kv_quant_mode(self.kv_cache_dtype).is_nvfp4
# Indexer side-cache dtype, mirroring --kv-cache-dtype for the main
# cache (--attention-config '{"indexer_kv_dtype": ...}').
self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
@@ -611,6 +613,8 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
main_slot_mapping = fwd_slot_mapping[self.layer_name]
index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
q = qkv.new_empty((num_tokens, self.q_size))
# With query_fp8, q is emitted only in fp8; the attend dequantizes it
# into ``q`` for the tokens that run on bf16-query kernels.
query_fp8 = self._allocate_query_fp8(qkv)
# index_q matches the index-K cache dtype (e4m3 for the fp8 score path);
# the fused kernel emits fp8 directly when this buffer is e4m3.
@@ -636,11 +640,13 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
self.kv_cache,
self.indexer.index_cache.kv_cache,
self.kv_cache.size(2), # paged-cache block size
q,
q if query_fp8 is None else None,
index_q,
self.kv_cache_dtype,
q_fp8_out=query_fp8,
q_fp8_scale=self._q_scale_float,
kv_k_scale=self._k_scale if self.use_nvfp4_kv else None,
kv_v_scale=self._v_scale if self.use_nvfp4_kv else None,
)
output = torch.empty_like(q)
vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py
+52 -7
View File
@@ -12,6 +12,7 @@ import torch
from vllm.config.attention import MiniMaxM3MSADecodeBackend
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.v1.kv_cache_interface import get_kv_quant_mode
_MAX_NUM_Q_HEADS = 64
_MAX_NUM_KV_HEADS = 4
@@ -23,9 +24,36 @@ _TOPK = 16
_MAX_QUERY_HEAD_ROWS = 65536
_MAX_DECODE_QUERY_LEN = 32
# Kernel benchmarks put the CUTLASS crossover at 16 requests for TP1 and TP4.
# NVFP4 has no Triton fallback, so it takes CUTLASS at every batch size.
_MIN_CUTLASS_BATCH_SIZE = 16
def is_nvfp4_kv_cache(kv_cache_dtype: str) -> bool:
return get_kv_quant_mode(kv_cache_dtype).is_nvfp4
def nvfp4_kv_cache_views(
kv_cache: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Return the K data, K scale, V data and V scale views of an NVFP4 cache.
``kv_cache`` is uint8 ``[pages, 2 * Hkv, page_size, 72]``; slot
``2 * head + side`` (K = 0, V = 1) is that head's E2M1 data
``[page_size, 64]`` followed by its E4M3 block scales ``[page_size, 8]``,
the one page layout ``fmha_sm100`` reads.
"""
from vllm.third_party.fmha_sm100.nvfp4_kv import nvfp4_head_slot_views
return nvfp4_head_slot_views(*nvfp4_kv_cache_slots(kv_cache))
def nvfp4_kv_cache_slots(
kv_cache: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Return the per-head K and V slots of an NVFP4 cache."""
return kv_cache[:, 0::2], kv_cache[:, 1::2]
@dataclass
class MSACutlassDecodeMetadata:
plan: Any
@@ -91,6 +119,9 @@ class MSACutlassDecodePlanCache:
use_fp8_kvcache=True,
split_prefill_decode=False,
device=device,
# The Q8KV4 route keeps plan-time lengths and its own page offsets,
# which this cached plan does not refresh.
decode_backend="kv_mode3",
)
plan_info = plan[3]
@@ -214,11 +245,12 @@ def supports_cutlass_sparse_decode(
topk_blocks: int,
) -> bool:
"""Return whether static model geometry supports CUTLASS sparse decode."""
nvfp4 = is_nvfp4_kv_cache(kv_cache_dtype)
return (
decode_backend == "cutlass"
(decode_backend == "cutlass" or nvfp4)
and current_platform.is_cuda()
and current_platform.is_device_capability_family(100)
and kv_cache_dtype in ("fp8", "fp8_e4m3")
and (kv_cache_dtype in ("fp8", "fp8_e4m3") or nvfp4)
and _supported_head_geometry(num_q_heads, num_kv_heads)
and page_size == _PAGE_SIZE
and topk_blocks == _TOPK
@@ -248,7 +280,7 @@ def should_prepare_decode_metadata(
topk_blocks=topk_blocks,
)
and 1 <= decode_query_len <= _MAX_DECODE_QUERY_LEN
and batch_size >= _MIN_CUTLASS_BATCH_SIZE
and (batch_size >= _MIN_CUTLASS_BATCH_SIZE or is_nvfp4_kv_cache(kv_cache_dtype))
and total_q * num_q_heads <= _MAX_QUERY_HEAD_ROWS
)
@@ -292,9 +324,22 @@ def msa_cutlass_sparse_decode(
q_scale_float: float,
k_scale_float: float,
v_scale_float: float,
k_scale: torch.Tensor | None = None,
v_scale: torch.Tensor | None = None,
) -> None:
"""Run CUTLASS sparse decode with metadata prepared by the MSA builder."""
key, value = kv_cache.split(_HEAD_DIM, dim=-1)
"""Run CUTLASS sparse decode with metadata prepared by the MSA builder.
NVFP4 caches are uint8 head-slot pages (see ``nvfp4_kv_cache_views``) and
take the device global scales ``k_scale``/``v_scale``.
"""
if kv_cache.dtype == torch.uint8:
assert k_scale is not None and v_scale is not None
key, value = nvfp4_kv_cache_slots(kv_cache)
k_scale_arg: float | torch.Tensor = k_scale
v_scale_arg: float | torch.Tensor = v_scale
else:
key, value = kv_cache.split(_HEAD_DIM, dim=-1)
k_scale_arg, v_scale_arg = k_scale_float, v_scale_float
from vllm.third_party.fmha_sm100.api import fmha_sm100
@@ -310,7 +355,7 @@ def msa_cutlass_sparse_decode(
output_o=True,
sm_scale=scale,
q_scale=q_scale_float,
k_scale=k_scale_float,
v_scale=v_scale_float,
k_scale=k_scale_arg,
v_scale=v_scale_arg,
o_scale=1.0,
)
vllm/models/minimax_m3/nvidia/sparse_attention_msa.py
+108 -11
View File
@@ -4,15 +4,18 @@
Prefill attends with ``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``).
Decode uses Triton split-K by default, with an opt-in CUTLASS ``fmha_sm100``
path for regular decode and speculative verification.
path for regular decode and speculative verification. NVFP4 KV caches read the
vLLM packed pages directly in both prefill and CUTLASS decode.
"""
from dataclasses import dataclass
from dataclasses import dataclass, replace
from typing import ClassVar
import torch
from vllm.config import VllmConfig
from vllm.config.attention import MiniMaxM3MSADecodeBackend
from vllm.config.cache import CacheDType
from vllm.forward_context import get_forward_context
from vllm.logger import init_logger
from vllm.models.minimax_m3.common.ops.sparse_attn import (
@@ -29,20 +32,33 @@ from vllm.models.minimax_m3.common.sparse_attention import (
from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import (
MSACutlassDecodeMetadata,
MSACutlassDecodePlanCache,
is_nvfp4_kv_cache,
msa_cutlass_sparse_decode,
nvfp4_kv_cache_views,
prepare_decode_metadata,
should_prepare_decode_metadata,
supports_cutlass_sparse_decode,
)
from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim
from vllm.v1.attention.backend import (
AttentionLayer,
CommonAttentionMetadata,
)
from vllm.v1.kv_cache_interface import AttentionSpec
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
logger = init_logger(__name__)
def _dequantize_query(
query_fp8: torch.Tensor, q_scale: float, out: torch.Tensor
) -> torch.Tensor:
"""Write ``query_fp8 * q_scale`` into ``out`` (E4M3 is exact in BF16)."""
out.copy_(query_fp8)
if q_scale != 1.0:
out.mul_(q_scale)
return out
class MiniMaxM3SparseMSABackend(MiniMaxM3SparseBackend):
"""MiniMax M3 backend with NVIDIA MSA-specific decode metadata."""
@@ -51,6 +67,30 @@ class MiniMaxM3SparseMSABackend(MiniMaxM3SparseBackend):
return MiniMaxM3SparseMSAMetadataBuilder
class MiniMaxM3SparseMSANvfp4Backend(MiniMaxM3SparseMSABackend):
"""MSA backend over HND NVFP4 KV cache pages in per-head K/V slots."""
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["nvfp4"]
@classmethod
def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
"""Store each head's K and V as two slots of packed data + block scales
(see ``nvfp4_kv_cache_views``)."""
if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_nvfp4:
return spec
return replace(
spec,
num_head_slots=2 * spec.num_kv_heads,
state_content_bytes=nvfp4_kv_cache_full_dim(spec.head_size)
* get_dtype_size(spec.dtype),
)
@classmethod
def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None:
# Every head's K/V slots are one contiguous run of the page.
return (KVCacheLayout.LBHNC,)
class MiniMaxM3SparseCutlassBackend(MiniMaxM3SparseMSABackend):
"""Attention-backend alias selecting CUTLASS MSA sparse decode."""
@@ -166,6 +206,7 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
topk_blocks=topk_blocks,
sparse_block_size=sparse_block_size,
)
self.use_nvfp4_kv = is_nvfp4_kv_cache(kv_cache_dtype)
self.use_cutlass_decode = supports_cutlass_sparse_decode(
decode_backend=msa_decode_backend,
num_q_heads=self.num_heads,
@@ -174,6 +215,12 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
page_size=self.block_size,
topk_blocks=self.topk_blocks,
)
if self.use_nvfp4_kv and not self.use_cutlass_decode:
raise ValueError(
"MiniMax M3 NVFP4 KV cache requires CUTLASS MSA sparse decode "
f"(num_heads={self.num_heads}, num_kv_heads={self.num_kv_heads}, "
f"block_size={self.block_size}, topk_blocks={self.topk_blocks})"
)
logger.info_once(
"MiniMax M3 MSA sparse decode selected %s",
"CUTLASS" if self.use_cutlass_decode else "Triton",
@@ -218,9 +265,17 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
hd = self.head_size
q = query[:num_tokens].view(-1, self.num_heads, hd)
out = output[:num_tokens].view(-1, self.num_heads, hd)
kv_cache = (
kv_cache.view(self.kv_cache_fp8_dtype) if self.use_fp8_kv else kv_cache
# Given query_fp8, it is the only q the fused insert wrote: CUTLASS
# decode reads it directly, and the bf16-query kernels read it
# dequantized into ``query``.
q_fp8 = (
None
if query_fp8 is None
else query_fp8[:num_tokens].view(-1, self.num_heads, hd)
)
q_scale_float = getattr(layer, "_q_scale_float", 1.0)
if self.use_fp8_kv and not self.use_nvfp4_kv:
kv_cache = kv_cache.view(self.kv_cache_fp8_dtype)
k_scale = getattr(layer, "_k_scale", None) if self.use_fp8_kv else None
v_scale = getattr(layer, "_v_scale", None) if self.use_fp8_kv else None
@@ -234,21 +289,31 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
else None
)
if self.use_cutlass_decode and msa_metadata is not None:
assert query_fp8 is not None
assert q_fp8 is not None
msa_cutlass_sparse_decode(
query_fp8[:nd].view(-1, self.num_heads, hd),
q_fp8[:nd],
kv_cache,
topk[:nd],
out[:nd],
msa_metadata,
scale=self.scale,
q_scale_float=getattr(layer, "_q_scale_float", 1.0),
q_scale_float=q_scale_float,
k_scale_float=getattr(layer, "_k_scale_float", 1.0),
v_scale_float=getattr(layer, "_v_scale_float", 1.0),
k_scale=k_scale if self.use_nvfp4_kv else None,
v_scale=v_scale if self.use_nvfp4_kv else None,
)
elif self.use_nvfp4_kv:
raise RuntimeError(
"MiniMax M3 NVFP4 KV cache has no Triton decode fallback; "
f"CUTLASS decode was not planned for {main_md.num_decodes} "
f"requests with query length {d.decode_query_len}"
)
else:
minimax_m3_sparse_attn_decode(
q[:nd],
q[:nd]
if q_fp8 is None
else _dequantize_query(q_fp8[:nd], q_scale_float, q[:nd]),
kv_cache,
topk[:nd].transpose(0, 1),
d.block_table,
@@ -266,6 +331,7 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
from vllm.third_party.fmha_sm100.sparse import (
build_k2q_csr,
sparse_atten_func,
sparse_atten_nvfp4_kv_func,
)
p = main_md.prefill
@@ -273,8 +339,11 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
# [H, prefill, MK] transposed view; build_k2q_csr consumes the
# strided view directly (topK stays innermost-contiguous).
prefill_topk = topk[nd:num_tokens].transpose(0, 1)
qp = q[nd:]
k_cache, v_cache = kv_cache.split(self.head_size, dim=-1)
qp = (
q[nd:]
if q_fp8 is None
else _dequantize_query(q_fp8[nd:], q_scale_float, q[nd:])
)
k2q_row_ptr, k2q_q_indices, schedule = build_k2q_csr(
prefill_topk,
p.cu_seqlens_q,
@@ -287,6 +356,34 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
qhead_per_kv=qp.shape[1] // self.num_kv_heads,
return_schedule=True,
)
if self.use_nvfp4_kv:
k_data, k_sf, v_data, v_sf = nvfp4_kv_cache_views(kv_cache)
sparse_atten_nvfp4_kv_func(
qp,
k_data,
v_data,
k_sf,
v_sf,
k_scale,
v_scale,
k2q_row_ptr,
k2q_q_indices,
self.topk_blocks,
cu_seqlens_q=p.cu_seqlens_q,
cu_seqlens_k=p.cu_seqlens_k,
max_seqlen_q=p.max_query_len,
max_seqlen_k=p.max_seq_len,
blk_kv=SPARSE_BLOCK_SIZE,
causal=True,
softmax_scale=self.scale,
page_table=p.block_table,
seqused_k=p.seq_lens,
schedule=schedule,
out=out[nd:],
kv_layout="vllm",
)
return output
k_cache, v_cache = kv_cache.split(self.head_size, dim=-1)
sparse_atten_func(
qp,
k_cache,