diff --git a/CMakeLists.txt b/CMakeLists.txt index 02e098f2d9..0a5414b4bf 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -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 diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake index ae5bfec7b4..f6bb0b8ace 100644 --- a/cmake/external_projects/fmha_sm100.cmake +++ b/cmake/external_projects/fmha_sm100.cmake @@ -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 diff --git a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu index 8769ba723b..0a45af8ade 100644 --- a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu @@ -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 +__device__ __forceinline__ void storeIndexQElemsFp8( + uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) { +#ifndef USE_ROCM + using Converter = vllm::_typeConvert; + __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) { + 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(dst) = *reinterpret_cast(out2); +#else + storeScaledQElemsFp8(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(__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 +__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; + float x[kElemsPerLane]; + float group_max = 0.0f; + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + x[i] = static_cast(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(codes); + #endif + packed = clearNegZeroNibbles(packed); + + uint8_t* const slot = page + static_cast(2 * head + kv) * + block_size * + (kNvfp4DataBytes + kNvfp4ScaleBytes); + uint8_t* const data = + slot + token * kNvfp4DataBytes + laneId * (kElemsPerLane / 2); + *reinterpret_cast(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 + 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 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(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(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(tokenIdx) * niq * kHeadDim + - (slot - iq_begin) * kHeadDim + dim_base, - elems); - } else { + storeIndexQElemsFp8( + index_q_out + static_cast(tokenIdx) * niq * kHeadDim + + (slot - iq_begin) * kHeadDim + dim_base, + elems); + } else if (store_ptr != nullptr) { storeElems(store_ptr + dim_base, elems); } - } else { + } else if (store_ptr != nullptr) { storeElems(store_ptr + dim_base, elems); } if (isQ && q_fp8_out != nullptr) { @@ -487,6 +759,16 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( } else { storeElems(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( + reinterpret_cast(kv_cache) + + (sm / block_size) * kv_s_block, + isK ? 0 : 1, head, static_cast(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 +template 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, \ - qkv, q_out, q_fp8_out, reinterpret_cast(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(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><<>>( \ + cudaLaunchKernelEx( \ + &config, \ + fusedMiniMaxM3QNormRopeKVInsertKernel, \ qkv, q_out, q_fp8_out, reinterpret_cast(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(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><<>>( \ + qkv, q_out, q_fp8_out, reinterpret_cast(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(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( \ +#define CALL_FUSED_MINIMAX_M3_IMPL(CACHE_T, KV_DTYPE, NVFP4) \ + vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3( \ reinterpret_cast(qkv.data_ptr()), \ q_out.has_value() ? reinterpret_cast(q_out->data_ptr()) : nullptr, \ q_fp8_out.has_value() \ @@ -654,7 +948,10 @@ void launchFusedMiniMaxM3( static_cast(eps), 1.0f / static_cast(q_fp8_scale), \ static_cast(rotary_dim), num_tokens, nq, nkv, niq, \ static_cast(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 q_fp8_out, // [N, nq*128] contiguous E4M3 - double q_fp8_scale) { + double q_fp8_scale, + std::optional kv_k_scale, // [] f32, NVFP4 only + std::optional 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() : nullptr; + float const* kv_v_scale_ptr = + nvfp4_kv ? kv_v_scale->const_data_ptr() : 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 diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 6d45f99881..3266d28b1e 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -381,7 +381,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert( std::optional q_out, std::optional index_q_out, const std::string& kv_cache_dtype, bool skip_index_branch, - std::optional q_fp8_out, double q_fp8_scale); + std::optional q_fp8_out, double q_fp8_scale, + std::optional kv_k_scale, + std::optional kv_v_scale); #ifdef VLLM_ENABLE_FUSED_KDA_DECODE void fused_kda_decode( diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 98e89c39d3..3d4f9ab273 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -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( diff --git a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py index 5c94968ddc..0007f20c5b 100644 --- a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py +++ b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py @@ -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) diff --git a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py index 0bc3c5692a..16a9be307b 100644 --- a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py @@ -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 + ) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 51a810650f..eef1183697 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -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, ) diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index 75aa90100c..61791dc5ff 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -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 diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index 1c0138946f..06c28a7a38 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -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) diff --git a/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py b/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py index 04313af107..c52dee413d 100644 --- a/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py +++ b/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py @@ -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, ) diff --git a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py index d29e1346bd..5aa23b4ac8 100644 --- a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py +++ b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py @@ -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,