@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7566d83bd3

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

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Yongye Zhu committed Oct 1, 2026 at 14:39 UTC 7566d83bd35cdfd7cd3f9e521a005d2fc1ed1c33
12 files changed +1293 -69
CMakeLists.txt
+21
@@ -590,9 +590,30 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
590 endif()
591 endif()
592
593 + # The fused MiniMax-M3 KV insert quantizes NVFP4 with SM100-family-only
594 + # cvt instructions, so its SM10x/11x code is built for the family targets.
595 + set(MINIMAX_M3_FUSED_SRC
596 + "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu")
597 + list(REMOVE_ITEM VLLM_STABLE_EXT_SRC "${MINIMAX_M3_FUSED_SRC}")
598 set_gencode_flags_for_srcs(
599 SRCS "${VLLM_STABLE_EXT_SRC}"
600 CUDA_ARCHS "${CUDA_ARCHS}")
601 + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
602 + cuda_archs_loose_intersection(MINIMAX_M3_FUSED_ARCHS
603 + "10.0f;10.7f;11.0f" "${CUDA_ARCHS}")
604 + else()
605 + cuda_archs_loose_intersection(MINIMAX_M3_FUSED_ARCHS
606 + "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
607 + endif()
608 + foreach(_ARCH ${CUDA_ARCHS})
609 + if(_ARCH VERSION_LESS 10.0 OR _ARCH VERSION_GREATER_EQUAL 12.0)
610 + list(APPEND MINIMAX_M3_FUSED_ARCHS "${_ARCH}")
611 + endif()
612 + endforeach()
613 + set_gencode_flags_for_srcs(
614 + SRCS "${MINIMAX_M3_FUSED_SRC}"
615 + CUDA_ARCHS "${MINIMAX_M3_FUSED_ARCHS}")
616 + list(APPEND VLLM_STABLE_EXT_SRC "${MINIMAX_M3_FUSED_SRC}")
617
618 if(COOPERATIVE_TOPK_ARCHS)
619 list(APPEND VLLM_STABLE_EXT_SRC
cmake/external_projects/fmha_sm100.cmake
+15 -1
@@ -17,7 +17,7 @@ else()
17 FetchContent_Declare(
18 fmha_sm100
19 GIT_REPOSITORY https://github.com/vllm-project/MSA.git
20 - GIT_TAG f355c37eb4e1413f21ee2ad8bbad25079e6bef9d
20 + GIT_TAG be4a52aea420b6bef73be9af1a649859964790cb
21 GIT_PROGRESS TRUE
22 CONFIGURE_COMMAND ""
23 BUILD_COMMAND ""
@@ -39,6 +39,8 @@ install(FILES
39 "${FMHA_SM100_PY_ROOT}/api.py"
40 "${FMHA_SM100_PY_ROOT}/bench_utils.py"
41 "${FMHA_SM100_PY_ROOT}/jit.py"
42 + "${FMHA_SM100_PY_ROOT}/nvfp4_kv.py"
43 + "${FMHA_SM100_PY_ROOT}/q8kv4_decode_adapter.py"
44 "${FMHA_SM100_PY_ROOT}/sparse.py"
45 "${FMHA_SM100_PY_ROOT}/sparse_fmha_adapter.py"
46 DESTINATION vllm/third_party/fmha_sm100
@@ -51,6 +53,18 @@ install(DIRECTORY "${FMHA_SM100_PY_ROOT}/csrc/"
53 PATTERN "*.pyc" EXCLUDE
54 PATTERN ".git*" EXCLUDE)
55
56 +install(DIRECTORY "${FMHA_SM100_PY_ROOT}/decode_q8kv4/"
57 + DESTINATION vllm/third_party/fmha_sm100/decode_q8kv4
58 + COMPONENT fmha_sm100
59 + PATTERN "__pycache__" EXCLUDE
60 + PATTERN "*.pyc" EXCLUDE)
61 +
62 +install(DIRECTORY "${FMHA_SM100_PY_ROOT}/kvouter/"
63 + DESTINATION vllm/third_party/fmha_sm100/kvouter
64 + COMPONENT fmha_sm100
65 + PATTERN "__pycache__" EXCLUDE
66 + PATTERN "*.pyc" EXCLUDE)
67 +
68 install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/"
69 DESTINATION vllm/third_party/fmha_sm100/cute
70 COMPONENT fmha_sm100
csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
+387 -38
@@ -285,6 +285,262 @@ __device__ __forceinline__ void storeScaledQElemsFp8(
285 storeElemsFp8(dst, scaled);
286 }
287
288 +// index_q: E4M3(scalar_t(v)) with unit scale, i.e. q is rounded to the model
289 +// dtype before the fp8 cast. On CUDA each element pair is rounded to scalar_t
290 +// and then converted to E4M3, with one packed cvt each.
291 +template <typename scalar_t>
292 +__device__ __forceinline__ void storeIndexQElemsFp8(
293 + uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
294 +#ifndef USE_ROCM
295 + using Converter = vllm::_typeConvert<scalar_t>;
296 + __nv_fp8x2_storage_t out2[kElemsPerLane / 2];
297 + #pragma unroll
298 + for (int i = 0; i < kElemsPerLane / 2; i++) {
299 + auto const v =
300 + Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1]));
301 + if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
302 + out2[i] = __nv_cvt_bfloat16raw2_to_fp8x2(v, __NV_SATFINITE, __NV_E4M3);
303 + } else {
304 + out2[i] = __nv_cvt_halfraw2_to_fp8x2(v, __NV_SATFINITE, __NV_E4M3);
305 + }
306 + }
307 + *reinterpret_cast<uint32_t*>(dst) = *reinterpret_cast<uint32_t const*>(out2);
308 +#else
309 + storeScaledQElemsFp8<scalar_t>(dst, elems, 1.0f);
310 +#endif
311 +}
312 +
313 +// ────────────────────────────────────────────────────────────────────────────
314 +// NVFP4 main K/V cache store
315 +// ────────────────────────────────────────────────────────────────────────────
316 +// The block encoding of reshape_and_cache_nvfp4_kernel
317 +// (nvfp4_kv_cache_kernels.cu) in the MSA readers' head-slot page layout. The
318 +// quantizer differs from that kernel's reciprocal-multiply math: per 16-value
319 +// block, scale = min(448, max(E4M3(max(amax, 1e-12) / 6), 1/512)) and
320 +// code = sign | RNE-E2M1(min(|x / scale|, 6)), with zero stored as +0, all on
321 +// x / the global scale.
322 +// An HND page is 2*nkv slots of [bs, 72] bytes; slot 2*head + side (K = 0,
323 +// V = 1) holds that head's packed E2M1 data [bs, 64] followed by its E4M3 block
324 +// scales [bs, 8]. Every head is one contiguous 2*bs*72-byte run, so a page
325 +// split across TP ranks is a contiguous byte slice per rank. K scales are
326 +// linear, V scales use the SM100 trtllm-gen 4x4 token-quad swizzle. One
327 +// 16-element scale group is four lanes. SM100-family builds of this file
328 +// (CMake gives it family-specific SM10x/11x targets) encode E2M1 with
329 +// cvt.rn.satfinite.e2m1x2; other targets use a bit-exact software encoder.
330 +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
331 + defined(__CUDA_ARCH_FAMILY_SPECIFIC__)
332 + #define MINIMAX_M3_NVFP4_NATIVE_CVT 1
333 +#else
334 + #define MINIMAX_M3_NVFP4_NATIVE_CVT 0
335 +#endif
336 +
337 +#ifndef USE_ROCM
338 +constexpr int kNvfp4DataBytes = kHeadDim / 2; // 64
339 +constexpr int kNvfp4ScaleBytes = kHeadDim / 16; // 8
340 +constexpr int kNvfp4LanesPerGroup = 16 / kElemsPerLane; // 4
341 +
342 +// Round-to-nearest-even, saturating E2M1 code of ``x`` (matches
343 +// cvt.rn.satfinite.e2m1x2.f32). Below 1 the grid is the 0.5-step subnormals;
344 +// from 1 up it is fp32 rounded to one mantissa bit, whose (exponent, mantissa)
345 +// bits are the code offset by 252 once the exponent is rebiased.
346 +__device__ __forceinline__ uint32_t e2m1Code(float x) {
347 + float const ax = fabsf(x);
348 + uint32_t mag;
349 + if (ax < 1.0f) {
350 + mag = static_cast<uint32_t>(__float2int_rn(2.0f * ax));
351 + } else {
352 + uint32_t const bits = __float_as_uint(ax);
353 + uint32_t const rounded = bits + 0x1FFFFFu + ((bits >> 22) & 1u);
354 + mag = min((rounded >> 22) - 252u, 7u);
355 + }
356 + return mag | ((__float_as_uint(x) >> 28) & 0x8u);
357 +}
358 +
359 +__device__ __forceinline__ float rcpApproxFtz(float a) {
360 + float b;
361 + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(b) : "f"(a));
362 + return b;
363 +}
364 +
365 +// a / b rounded to nearest even from r = RN(1 / b): Markstein's correction of
366 +// a * r, which is __fdiv_rn's fast path without its range-check branch. Exact
367 +// unless a * r overflows or the residual a - q * b underflows (a within a few
368 +// binades of the subnormals).
369 +__device__ __forceinline__ float divRnByRcp(float a, float b, float r) {
370 + float const q = __fmul_rn(a, r);
371 + return __fmaf_rn(__fmaf_rn(-q, b, a), r, q);
372 +}
373 +
374 +// Block scale E4M3(max(amax, 1e-12) / 6) clamped to [1/512, 448], with
375 +// sf_scale applied to it as to the values, and its reciprocal RN(1 / sf_value).
376 +__device__ __forceinline__ void nvfp4GroupScale(float sf_scale, float group_max,
377 + uint8_t& sf_byte,
378 + float& sf_value,
379 + float& sf_rcp) {
380 + constexpr float kRcp6 = 0x1.555556p-3f; // RN(1 / 6)
381 + float const amax = fmaxf(group_max, 1e-12f);
382 + // amax / 6 (RNE). amax * kRcp6 bounds it from above, and fminf picks it for
383 + // amax = inf, where the correction is NaN.
384 + float const amax_6 =
385 + fminf(divRnByRcp(amax, 6.0f, kRcp6), __fmul_rn(amax, kRcp6));
386 + // Flooring before the rounding gives the 1/512 minimum; satfinite caps it
387 + // at 448.
388 + __nv_fp8_e4m3 const sf(fmaxf(sf_scale * amax_6, 0x1p-9f));
389 + sf_byte = sf.__x;
390 + sf_value = float(sf);
391 + // rcp.approx is correctly rounded on every positive E4M3 value.
392 + sf_rcp = rcpApproxFtz(sf_value);
393 +}
394 +
395 +// nvfp4GroupScale for sf_scale = 1 and a bf16/fp16 group_max. Uncorrected,
396 +// amax * RN(1 / 6) still rounds to the E4M3 byte of amax / 6 (checked for every
397 +// such amax: amax / 6 is too far from an E4M3 rounding boundary to cross it,
398 +// and exact ties round back), and the 1e-12 floor lies below the 1/512 one.
399 +__device__ __forceinline__ void nvfp4UnitGroupScale(float group_max,
400 + uint8_t& sf_byte,
401 + float& sf_rcp) {
402 + constexpr float kRcp6 = 0x1.555556p-3f; // RN(1 / 6)
403 + __nv_fp8_e4m3 const sf(fmaxf(__fmul_rn(group_max, kRcp6), 0x1p-9f));
404 + sf_byte = sf.__x;
405 + sf_rcp = rcpApproxFtz(float(sf));
406 +}
407 +
408 +// (x * sf_scale) / sf_value with the E2M1 code of the IEEE quotient. The
409 +// quotient itself is the IEEE one except below 1e-35 (code 0 either way) and
410 +// for magnitudes over 4096, which are clamped (NaN kept) so that the division
411 +// stays finite: 4096 / 448 > 5, so their code stays the saturated 6.
412 +__device__ __forceinline__ float nvfp4Quotient(float x, float sf_scale,
413 + float sf_value, float sf_rcp) {
414 + constexpr float kMaxMag = 4096.0f;
415 + float n = __fmul_rn(x, sf_scale);
416 + n = fabsf(n) > kMaxMag ? copysignf(kMaxMag, n) : n;
417 + return divRnByRcp(n, sf_value, sf_rcp);
418 +}
419 +
420 +// Zero is stored as +0: clear the sign of every nibble whose magnitude is 0.
421 +// Adding 7 to a 3-bit magnitude carries into the sign bit exactly when it is
422 +// nonzero.
423 +__device__ __forceinline__ uint16_t clearNegZeroNibbles(uint16_t packed) {
424 + return packed & (((packed & 0x7777u) + 0x7777u) | 0x7777u);
425 +}
426 +
427 +// Quantize one K or V head row of one token into its NVFP4 page, where
428 +// ``scale`` is the dequantization (global) scale. Every lane of the warp must
429 +// call this (the scale-group reductions shuffle).
430 +template <typename scalar_t>
431 +__device__ __forceinline__ void storeNvfp4CacheElems(
432 + uint8_t* __restrict__ page, int const kv, int const head, int const token,
433 + int const block_size, int const laneId, float const (&elems)[kElemsPerLane],
434 + float const scale) {
435 + // Quantize the model-dtype value that the unfused path reads back from qkv.
436 + using Converter = vllm::_typeConvert<scalar_t>;
437 + float x[kElemsPerLane];
438 + float group_max = 0.0f;
439 + #pragma unroll
440 + for (int i = 0; i < kElemsPerLane; i++) {
441 + x[i] = static_cast<float>(Converter::convert(elems[i]));
442 + group_max = fmaxf(group_max, fabsf(x[i]));
443 + }
444 + group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 1, 32));
445 + group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 2, 32));
446 +
447 + uint8_t sf_byte;
448 + float sf_rcp;
449 + float q[kElemsPerLane]; // (x / scale) / sf_value, E2M1-encoded below
450 + if (scale == 1.0f) {
451 + static_assert(sizeof(scalar_t) == 2, "the unit-scale path needs bf16/fp16");
452 + nvfp4UnitGroupScale(group_max, sf_byte, sf_rcp);
453 + // x / sf_value truncated from x * RN(1 / sf_value). RN(1 / s) >= 1 / s for
454 + // every E4M3 s, so a bf16/fp16 x on an E2M1 midpoint lands exactly on it
455 + // (and the cvt rounds it to even), while any other bf16/fp16 x is too far
456 + // from a midpoint to reach one. No clamp either: a finite product never
457 + // truncates to inf, and inf saturates in the cvt.
458 + #if MINIMAX_M3_NVFP4_NATIVE_CVT
459 + float2 const r2 = make_float2(sf_rcp, sf_rcp);
460 + #pragma unroll
461 + for (int i = 0; i < kElemsPerLane; i += 2) {
462 + float2 const p = __fmul2_rz(make_float2(x[i], x[i + 1]), r2);
463 + q[i] = p.x;
464 + q[i + 1] = p.y;
465 + }
466 + #else
467 + #pragma unroll
468 + for (int i = 0; i < kElemsPerLane; i++) q[i] = __fmul_rz(x[i], sf_rcp);
469 + #endif
470 + } else {
471 + float const sf_scale = __frcp_rn(scale); // exactly 1.0f / scale
472 + float sf_value;
473 + nvfp4GroupScale(sf_scale, group_max, sf_byte, sf_value, sf_rcp);
474 + // A correctly rounded division, not a reciprocal multiply, so that exact
475 + // E2M1 midpoints round to even.
476 + #if MINIMAX_M3_NVFP4_NATIVE_CVT
477 + // nvfp4Quotient on element pairs, with SM100's packed FP32 instructions.
478 + float2 const s2 = make_float2(sf_scale, sf_scale);
479 + float2 const neg_v2 = make_float2(-sf_value, -sf_value);
480 + float2 const r2 = make_float2(sf_rcp, sf_rcp);
481 + #pragma unroll
482 + for (int i = 0; i < kElemsPerLane; i += 2) {
483 + float2 n = __fmul2_rn(make_float2(x[i], x[i + 1]), s2);
484 + // |n| <= 4096 as in nvfp4Quotient, NaN kept.
485 + asm("max.NaN.f32 %0, %0, 0fC5800000;\n\t"
486 + "min.NaN.f32 %0, %0, 0f45800000;\n\t"
487 + "max.NaN.f32 %1, %1, 0fC5800000;\n\t"
488 + "min.NaN.f32 %1, %1, 0f45800000;"
489 + : "+f"(n.x), "+f"(n.y));
490 + float2 const q0 = __fmul2_rn(n, r2);
491 + float2 const p = __ffma2_rn(__ffma2_rn(q0, neg_v2, n), r2, q0);
492 + q[i] = p.x;
493 + q[i + 1] = p.y;
494 + }
495 + #else
496 + #pragma unroll
497 + for (int i = 0; i < kElemsPerLane; i++) {
498 + q[i] = nvfp4Quotient(x[i], sf_scale, sf_value, sf_rcp);
499 + }
500 + #endif
501 + }
502 + uint16_t packed;
503 + #if MINIMAX_M3_NVFP4_NATIVE_CVT
504 + // Low nibble = even element; cvt packs its first source in the high nibble.
505 + asm("{\n"
506 + " .reg .b8 lo, hi;\n"
507 + " cvt.rn.satfinite.e2m1x2.f32 lo, %2, %1;\n"
508 + " cvt.rn.satfinite.e2m1x2.f32 hi, %4, %3;\n"
509 + " mov.b16 %0, {lo, hi};\n"
510 + "}"
511 + : "=h"(packed)
512 + : "f"(q[0]), "f"(q[1]), "f"(q[2]), "f"(q[3]));
513 + #else
514 + uint32_t codes = 0;
515 + #pragma unroll
516 + for (int i = 0; i < kElemsPerLane; i++) codes |= e2m1Code(q[i]) << (4 * i);
517 + packed = static_cast<uint16_t>(codes);
518 + #endif
519 + packed = clearNegZeroNibbles(packed);
520 +
521 + uint8_t* const slot = page + static_cast<int64_t>(2 * head + kv) *
522 + block_size *
523 + (kNvfp4DataBytes + kNvfp4ScaleBytes);
524 + uint8_t* const data =
525 + slot + token * kNvfp4DataBytes + laneId * (kElemsPerLane / 2);
526 + *reinterpret_cast<uint16_t*>(data) = packed;
527 +
528 + if (laneId % kNvfp4LanesPerGroup == 0) {
529 + int const group = laneId / kNvfp4LanesPerGroup;
530 + int offset;
531 + if (kv == 0) {
532 + offset = token * kNvfp4ScaleBytes + group;
533 + } else {
534 + constexpr int kGroupsPerQuad = kNvfp4ScaleBytes / 4;
535 + int const swizzled_t = (token / 4) * 4 + group / kGroupsPerQuad;
536 + int const swizzled_s = (group % kGroupsPerQuad) * 4 + token % 4;
537 + offset = swizzled_t * kNvfp4ScaleBytes + swizzled_s;
538 + }
539 + slot[block_size * kNvfp4DataBytes + offset] = sf_byte;
540 + }
541 +}
542 +#endif // USE_ROCM
543 +
544 // ────────────────────────────────────────────────────────────────────────────
545 // Kernel
546 // ────────────────────────────────────────────────────────────────────────────
@@ -307,9 +563,8 @@ __device__ __forceinline__ void storeScaledQElemsFp8(
563 // branch and writes index_q/index_k outputs. Skip-index-topk reuse layers keep
564 // kHasIndex=true but set kProcessIndex=false.
565 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
310 - typename out_idx_t, bool kHasIndex, bool kInsertKV,
311 - bool kProcessIndex,
312 - bool kFp8Idx>
566 + bool kNvfp4, typename out_idx_t, bool kHasIndex, bool kInsertKV,
567 + bool kProcessIndex, bool kFp8Idx>
568 __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
569 scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse)
570 scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
@@ -332,7 +587,10 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
587 // The content (last) dim is always innermost-contiguous (stride 1), so the
588 // NHD/HND layout choice is captured by the head/token strides.
589 int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
335 - int64_t const kv_s_dim) {
590 + int64_t const kv_s_dim,
591 + // NVFP4 main cache only: per-layer K/V dequantization (global) scales.
592 + float const* __restrict__ kv_k_scale,
593 + float const* __restrict__ kv_v_scale) {
594 #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
595 // _typeConvert<BFloat16> is unavailable on pre-Ampere; the M3 kernel only
596 // runs with bf16/fp16 inputs in practice. Discard the bf16 body there.
@@ -352,8 +610,11 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
610 int const idx_slots = kProcessIndex ? niq + 1 : 0;
611 int const slots_per_token = nq + nkv + v_slots + idx_slots;
612
355 - int const tokenIdx = globalWarpIdx / slots_per_token;
356 - int const slot = globalWarpIdx % slots_per_token;
613 + // Unsigned: both are non-negative, and the signed division's sign handling
614 + // would cost every warp.
615 + unsigned const warp_u = globalWarpIdx, slots_u = slots_per_token;
616 + int const tokenIdx = warp_u / slots_u;
617 + int const slot = warp_u % slots_u;
618 if (tokenIdx >= num_tokens) return;
619
620 // Slot boundaries.
@@ -420,10 +681,13 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
681 // flat TMA descriptor can address them as [tokens*heads, head_dim]; this
682 // folds the de-interleaving into the store the kernel already does, instead
683 // of a separate q.contiguous() copy. Everything else stays in place.
684 + // Given q_fp8_out but no q_out, q is emitted only in fp8 (null store_ptr).
685 scalar_t* store_ptr = row_ptr;
686 if (isQ && q_out != nullptr) {
687 store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
688 slot * kHeadDim;
689 + } else if (isQ && q_fp8_out != nullptr) {
690 + store_ptr = nullptr;
691 } else if (isIQ && index_q_out != nullptr) {
692 // bf16 index_q_out: gather here. fp8: written by the explicit fp8 store.
693 if constexpr (!kFp8Idx) {
@@ -433,6 +697,13 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
697 }
698 }
699
700 + // NVFP4: the per-layer global scale is not produced by the predecessor,
701 + // so load it before the PDL wait.
702 + float nvfp4_scale = 0.0f;
703 + if constexpr (kNvfp4) {
704 + if (isK || isV) nvfp4_scale = isK ? *kv_k_scale : *kv_v_scale;
705 + }
706 +
707 // PDL: wait for the predecessor kernel (the qkv-projection GEMM that
708 // produces ``qkv``) to finish before touching any global memory. No-op
709 // when PDL is not enabled on the launch. The CUDA runtime wrapper emits
@@ -451,16 +722,17 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
722 normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
723 cos_ptr, /*apply_norm=*/norm_w != nullptr);
724 if constexpr (kFp8Idx) {
454 - // index_q is e4m3 bytes; Q/K (and in-place index_k) stay scalar_t.
725 + // index_q is E4M3(scalar_t(q)) bytes; Q/K (and in-place index_k) stay
726 + // scalar_t.
727 if (isIQ && index_q_out != nullptr) {
456 - storeElemsFp8(index_q_out +
457 - static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
458 - (slot - iq_begin) * kHeadDim + dim_base,
459 - elems);
460 - } else {
728 + storeIndexQElemsFp8<scalar_t>(
729 + index_q_out + static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
730 + (slot - iq_begin) * kHeadDim + dim_base,
731 + elems);
732 + } else if (store_ptr != nullptr) {
733 storeElems<scalar_t>(store_ptr + dim_base, elems);
734 }
463 - } else {
735 + } else if (store_ptr != nullptr) {
736 storeElems<scalar_t>(store_ptr + dim_base, elems);
737 }
738 if (isQ && q_fp8_out != nullptr) {
@@ -487,6 +759,16 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
759 } else {
760 storeElems<scalar_t>(index_cache + sm * kHeadDim + dim_base, elems);
761 }
762 + } else if constexpr (kNvfp4) {
763 +#ifndef USE_ROCM
764 + // kv_cache is [num_blocks, 2*nkv, block_size, 72] bytes (HND), slot
765 + // 2*head + side; kv_s_block is its page stride.
766 + storeNvfp4CacheElems<scalar_t>(
767 + reinterpret_cast<uint8_t*>(kv_cache) +
768 + (sm / block_size) * kv_s_block,
769 + isK ? 0 : 1, head, static_cast<int>(sm % block_size), block_size,
770 + laneId, elems, nvfp4_scale);
771 +#endif
772 } else if (isK || isV) {
773 // kv_cache logical shape [num_blocks, nkv, block_size, 2*head_dim].
774 // Paging is logical (block = sm/block_size, token = sm%block_size);
@@ -515,7 +797,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
797 // ────────────────────────────────────────────────────────────────────────────
798 // Launch wrapper
799 // ────────────────────────────────────────────────────────────────────────────
518 -template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
800 +template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
801 + bool kNvfp4 = false>
802 void launchFusedMiniMaxM3(
803 scalar_t* qkv, scalar_t* q_out, uint8_t* q_fp8_out, void* index_q_out,
804 scalar_t const* q_norm_w, scalar_t const* k_norm_w,
@@ -527,7 +810,8 @@ void launchFusedMiniMaxM3(
810 int const nq, int const nkv, int const niq, int const block_size,
811 int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
812 int64_t const kv_s_dim, bool const has_index, bool const insert_kv,
530 - bool const process_index, bool const fp8_idx, cudaStream_t stream) {
813 + bool const process_index, bool const fp8_idx, float const* kv_k_scale,
814 + float const* kv_v_scale, cudaStream_t stream) {
815 // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the
816 // void* pointers per instantiation in the LAUNCH macro.
817 // Slot count must match the kernel's compile-time gating.
@@ -559,35 +843,44 @@ void launchFusedMiniMaxM3(
843 config.attrs = attrs;
844 config.numAttrs = (sm_version >= 90) ? 1 : 0;
845
562 - #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \
563 - cudaLaunchKernelEx( \
564 - &config, \
565 - fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
566 - HAS_INDEX, INSERT, \
567 - PROCESS_INDEX, FP8>, \
568 - qkv, q_out, q_fp8_out, reinterpret_cast<OUT_T*>(index_q_out), \
569 - q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
570 - slot_mapping, index_slot_mapping, kv_cache, \
571 - reinterpret_cast<OUT_T*>(index_cache), eps, q_fp8_inv_scale, \
572 - rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
573 - kv_s_head, kv_s_token, kv_s_dim)
846 + #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \
847 + cudaLaunchKernelEx( \
848 + &config, \
849 + fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, \
850 + kNvfp4, OUT_T, HAS_INDEX, \
851 + INSERT, PROCESS_INDEX, FP8>, \
852 + qkv, q_out, q_fp8_out, reinterpret_cast<OUT_T*>(index_q_out), \
853 + q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
854 + slot_mapping, index_slot_mapping, kv_cache, \
855 + reinterpret_cast<OUT_T*>(index_cache), eps, q_fp8_inv_scale, \
856 + rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
857 + kv_s_head, kv_s_token, kv_s_dim, kv_k_scale, kv_v_scale)
858 #else
859 // ROCm: standard kernel launch syntax (no PDL/stream serialization).
860 // clang-format off
861 #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \
862 fusedMiniMaxM3QNormRopeKVInsertKernel< \
579 - scalar_t, cache_t, kv_dt, OUT_T, HAS_INDEX, INSERT, PROCESS_INDEX, \
580 - FP8><<<grid, kBlockSize, 0, stream>>>( \
863 + scalar_t, cache_t, kv_dt, kNvfp4, OUT_T, HAS_INDEX, INSERT, \
864 + PROCESS_INDEX, FP8><<<grid, kBlockSize, 0, stream>>>( \
865 qkv, q_out, q_fp8_out, reinterpret_cast<OUT_T*>(index_q_out), \
866 q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
867 slot_mapping, index_slot_mapping, kv_cache, \
868 reinterpret_cast<OUT_T*>(index_cache), eps, q_fp8_inv_scale, \
869 rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
586 - kv_s_head, kv_s_token, kv_s_dim)
870 + kv_s_head, kv_s_token, kv_s_dim, kv_k_scale, kv_v_scale)
871 // clang-format on
872 #endif
873
590 - if (has_index) {
874 + if constexpr (kNvfp4) {
875 + // NVFP4 main cache: sparse serving only (the host checks insert_kv).
876 + if (!process_index) {
877 + LAUNCH(true, true, false, false, scalar_t);
878 + } else if (fp8_idx) {
879 + LAUNCH(true, true, true, true, uint8_t);
880 + } else {
881 + LAUNCH(true, true, true, false, scalar_t);
882 + }
883 + } else if (has_index) {
884 if (!process_index) {
885 if (insert_kv) {
886 LAUNCH(true, true, false, false, scalar_t);
@@ -621,8 +914,9 @@ void launchFusedMiniMaxM3(
914 } // namespace vllm
915
916 // clang-format off
624 -#define CALL_FUSED_MINIMAX_M3(_RAW_T, CACHE_T, KV_DTYPE) \
625 - vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>( \
917 +#define CALL_FUSED_MINIMAX_M3_IMPL(CACHE_T, KV_DTYPE, NVFP4) \
918 + vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE, \
919 + NVFP4>( \
920 reinterpret_cast<st*>(qkv.data_ptr()), \
921 q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr()) : nullptr, \
922 q_fp8_out.has_value() \
@@ -654,7 +948,10 @@ void launchFusedMiniMaxM3(
948 static_cast<float>(eps), 1.0f / static_cast<float>(q_fp8_scale), \
949 static_cast<int>(rotary_dim), num_tokens, nq, nkv, niq, \
950 static_cast<int>(block_size), kv_s_block, kv_s_head, kv_s_token, \
657 - kv_s_dim, has_index, insert_kv, process_index, fp8_idx, stream)
951 + kv_s_dim, has_index, insert_kv, process_index, fp8_idx, kv_k_scale_ptr, \
952 + kv_v_scale_ptr, stream)
953 +#define CALL_FUSED_MINIMAX_M3(_RAW_T, CACHE_T, KV_DTYPE) \
954 + CALL_FUSED_MINIMAX_M3_IMPL(CACHE_T, KV_DTYPE, false)
955 // clang-format on
956
957 // ────────────────────────────────────────────────────────────────────────────
@@ -681,7 +978,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
978 const std::string& kv_cache_dtype, bool skip_index_branch,
979 std::optional<torch::stable::Tensor>
980 q_fp8_out, // [N, nq*128] contiguous E4M3
684 - double q_fp8_scale) {
981 + double q_fp8_scale,
982 + std::optional<torch::stable::Tensor> kv_k_scale, // [] f32, NVFP4 only
983 + std::optional<torch::stable::Tensor> kv_v_scale) {
984 STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
985 "qkv must be contiguous CUDA");
986 STD_TORCH_CHECK(
@@ -721,8 +1020,10 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
1020 bool const has_index = niq > 0;
1021 bool const insert_kv = kv_cache.has_value();
1022 bool const process_index = has_index && !skip_index_branch;
1023 + bool const nvfp4_kv = kv_cache_dtype == "nvfp4";
1024 vllm::Fp8KVCacheDataType const kv_dt =
725 - vllm::get_fp8_kv_cache_data_type(kv_cache_dtype);
1025 + nvfp4_kv ? vllm::Fp8KVCacheDataType::kAuto
1026 + : vllm::get_fp8_kv_cache_data_type(kv_cache_dtype);
1027 int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
1028 int const expected_row =
1029 (nq + 2 * nkv + (has_index ? niq + 1 : 0)) * kHeadDim;
@@ -769,8 +1070,13 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
1070 index_slot_mapping->numel() == slot_mapping->numel()),
1071 "index_slot_mapping must be int64 CUDA with slot_mapping length");
1072 }
772 - // Main attention KV cache: auto matches qkv, fp8 uses uint8 storage.
773 - if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
1073 + // Main attention KV cache: auto matches qkv, fp8 and nvfp4 use uint8
1074 + // storage.
1075 + if (nvfp4_kv) {
1076 + STD_TORCH_CHECK(
1077 + kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte,
1078 + "nvfp4 kv_cache must use uint8 storage");
1079 + } else if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
1080 STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
1081 "auto kv_cache dtype must match qkv");
1082 } else {
@@ -790,6 +1096,27 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
1096 STD_TORCH_CHECK(kv_cache->dim() == 4 && kv_cache->stride(3) == 1,
1097 "kv_cache must be [nb,nkv,bs,2*head_dim] with contiguous "
1098 "content dim (stride(3)==1)");
1099 + if (nvfp4_kv) {
1100 + // Packed HND pages [nb, 2*nkv, bs, 72]: slot 2*head + side is one
1101 + // contiguous [bs, 72] block, which the store addresses directly.
1102 + int64_t const full_dim = kHeadDim / 2 + kHeadDim / 16;
1103 + STD_TORCH_CHECK(
1104 + kv_cache->size(1) == 2 * nkv && kv_cache->size(2) == block_size &&
1105 + kv_cache->size(3) == full_dim &&
1106 + kv_cache->stride(2) == full_dim &&
1107 + kv_cache->stride(1) == block_size * full_dim &&
1108 + block_size % 4 == 0,
1109 + "nvfp4 kv_cache must be HND [nb, 2*num_kv_heads, block_size, 72] "
1110 + "with contiguous pages and block_size % 4 == 0");
1111 + for (auto const* scale : {&kv_k_scale, &kv_v_scale}) {
1112 + STD_TORCH_CHECK(
1113 + scale->has_value() && (*scale)->is_cuda() &&
1114 + (*scale)->scalar_type() ==
1115 + torch::headeronly::ScalarType::Float &&
1116 + (*scale)->numel() >= 1,
1117 + "nvfp4 kv_cache requires float32 CUDA kv_k_scale and kv_v_scale");
1118 + }
1119 + }
1120 kv_s_block = kv_cache->stride(0);
1121 kv_s_head = kv_cache->stride(1);
1122 kv_s_token = kv_cache->stride(2);
@@ -854,10 +1181,31 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
1181 "fp8 index path: index_q_out must be fp8 e4m3");
1182 }
1183
1184 + STD_TORCH_CHECK(!nvfp4_kv || insert_kv,
1185 + "nvfp4 kv_cache_dtype requires kv_cache (insert mode)");
1186 + float const* kv_k_scale_ptr =
1187 + nvfp4_kv ? kv_k_scale->const_data_ptr<float>() : nullptr;
1188 + float const* kv_v_scale_ptr =
1189 + nvfp4_kv ? kv_v_scale->const_data_ptr<float>() : nullptr;
1190 +
1191 const torch::stable::accelerator::DeviceGuard device_guard(
1192 qkv.get_device_index());
1193 auto stream = get_current_cuda_stream(qkv.get_device_index());
1194
1195 + if (nvfp4_kv) {
1196 +#ifndef USE_ROCM
1197 + VLLM_STABLE_DISPATCH_HALF_TYPES(
1198 + qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert_nvfp4", [&] {
1199 + using st = scalar_t;
1200 + CALL_FUSED_MINIMAX_M3_IMPL(uint8_t, vllm::Fp8KVCacheDataType::kAuto,
1201 + true);
1202 + });
1203 +#else
1204 + STD_TORCH_CHECK(false, "nvfp4 kv_cache is not supported on ROCm");
1205 +#endif
1206 + return;
1207 + }
1208 +
1209 VLLM_STABLE_DISPATCH_HALF_TYPES(
1210 qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert", [&] {
1211 using st = scalar_t;
@@ -867,3 +1215,4 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
1215 }
1216
1217 #undef CALL_FUSED_MINIMAX_M3
1218 +#undef CALL_FUSED_MINIMAX_M3_IMPL
csrc/libtorch_stable/ops.h
+3 -1
@@ -381,7 +381,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
381 std::optional<torch::stable::Tensor> q_out,
382 std::optional<torch::stable::Tensor> index_q_out,
383 const std::string& kv_cache_dtype, bool skip_index_branch,
384 - std::optional<torch::stable::Tensor> q_fp8_out, double q_fp8_scale);
384 + std::optional<torch::stable::Tensor> q_fp8_out, double q_fp8_scale,
385 + std::optional<torch::stable::Tensor> kv_k_scale,
386 + std::optional<torch::stable::Tensor> kv_v_scale);
387
388 #ifdef VLLM_ENABLE_FUSED_KDA_DECODE
389 void fused_kda_decode(
csrc/libtorch_stable/torch_bindings.cpp
+2 -1
@@ -513,7 +513,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
513 "Tensor!? kv_cache, Tensor!? index_cache, "
514 "int block_size, Tensor!? q_out, Tensor!? index_q_out, "
515 "str kv_cache_dtype, bool skip_index_branch=False, "
516 - "Tensor!? q_fp8_out=None, float q_fp8_scale=1.0) -> ()");
516 + "Tensor!? q_fp8_out=None, float q_fp8_scale=1.0, "
517 + "Tensor? kv_k_scale=None, Tensor? kv_v_scale=None) -> ()");
518
519 #ifdef VLLM_ENABLE_FUSED_KDA_DECODE
520 ops.def(
tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
+441 -1
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Correctness tests for MiniMax M3 CUTLASS sparse decode."""
3 +"""Correctness tests for MiniMax M3 MSA sparse attention (CUTLASS decode and
4 +NVFP4 KV cache)."""
5
6 import math
7 from types import SimpleNamespace
@@ -8,6 +9,7 @@ from types import SimpleNamespace
9 import pytest
10 import torch
11
12 +from tests.kernels.quantization.nvfp4_utils import dequant_nvfp4_kv_cache
13 from vllm import _custom_ops as ops
14 from vllm.config import AttentionConfig, CUDAGraphMode
15 from vllm.forward_context import ForwardContext, override_forward_context
@@ -18,6 +20,7 @@ from vllm.models.minimax_m3.common.sparse_attention import (
20 MiniMaxM3SparseBackend,
21 MiniMaxM3SparseMetadata,
22 MiniMaxM3SparseMetadataBuilder,
23 + MiniMaxM3SparsePrefillMetadata,
24 MiniMaxM3SparseTritonImpl,
25 select_main_backend_and_impl_cls,
26 )
@@ -28,6 +31,7 @@ from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention
31 from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import (
32 MSACutlassDecodePlanCache,
33 msa_cutlass_sparse_decode,
34 + nvfp4_kv_cache_views,
35 prepare_decode_metadata,
36 should_prepare_decode_metadata,
37 )
@@ -36,6 +40,7 @@ from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
40 MiniMaxM3SparseMSADecodeMetadata,
41 MiniMaxM3SparseMSAImpl,
42 MiniMaxM3SparseMSAMetadataBuilder,
43 + MiniMaxM3SparseMSANvfp4Backend,
44 )
45 from vllm.platforms import current_platform
46 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -164,6 +169,53 @@ def test_msa_cutlass_decode_static_dispatch_requires_sm100(
169 )
170
171
172 +@pytest.mark.parametrize("decode_backend", ["cutlass", "triton"])
173 +def test_msa_cutlass_decode_static_dispatch_nvfp4_has_no_fallback(
174 + decode_backend: str,
175 +) -> None:
176 + """NVFP4 has no Triton decode, so it plans CUTLASS for any batch size."""
177 + assert should_prepare_decode_metadata(
178 + 1,
179 + DEFAULT_QUERY_LEN,
180 + decode_backend=decode_backend, # type: ignore[arg-type]
181 + num_q_heads=16,
182 + num_kv_heads=1,
183 + kv_cache_dtype="nvfp4",
184 + page_size=BLOCK_SIZE,
185 + topk_blocks=TOPK,
186 + )
187 +
188 +
189 +def test_nvfp4_selects_msa_nvfp4_backend_and_requires_sm100(
190 + monkeypatch: pytest.MonkeyPatch,
191 +) -> None:
192 + backend_cls, impl_cls = select_main_backend_and_impl_cls(
193 + topk_blocks=TOPK,
194 + kv_cache_dtype="nvfp4",
195 + num_kv_heads=1,
196 + )
197 + assert backend_cls is MiniMaxM3SparseMSANvfp4Backend
198 + assert impl_cls is MiniMaxM3SparseMSAImpl
199 + with pytest.raises(ValueError, match="nvfp4_4over6"):
200 + select_main_backend_and_impl_cls(
201 + topk_blocks=TOPK,
202 + kv_cache_dtype="nvfp4_4over6",
203 + num_kv_heads=1,
204 + )
205 +
206 + monkeypatch.setattr(
207 + current_platform,
208 + "is_device_capability_family",
209 + lambda _: False,
210 + )
211 + with pytest.raises(ValueError, match="NVFP4"):
212 + select_main_backend_and_impl_cls(
213 + topk_blocks=TOPK,
214 + kv_cache_dtype="nvfp4",
215 + num_kv_heads=1,
216 + )
217 +
218 +
219 @pytest.mark.parametrize(
220 ("backend", "expected"),
221 [
@@ -639,3 +691,391 @@ def test_msa_cutlass_decode_matches_triton_with_interleaved_cache(
691 graph.replay()
692 current_platform.synchronize()
693 torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02)
694 +
695 +
696 +def _make_paged_layout(
697 + seq_lens_list: list[int],
698 +) -> tuple[torch.Tensor, int]:
699 + """Block table over a shuffled physical page pool."""
700 + pages_per_request = [math.ceil(seq_len / BLOCK_SIZE) for seq_len in seq_lens_list]
701 + num_pages = sum(pages_per_request)
702 + block_table = torch.zeros(
703 + len(seq_lens_list), max(pages_per_request), dtype=torch.int32, device="cuda"
704 + )
705 + physical_pages = torch.randperm(num_pages, dtype=torch.int32, device="cuda")
706 + offset = 0
707 + for request, request_pages in enumerate(pages_per_request):
708 + block_table[request, :request_pages] = physical_pages[
709 + offset : offset + request_pages
710 + ]
711 + offset += request_pages
712 + return block_table, num_pages
713 +
714 +
715 +def _make_nvfp4_kv_cache(
716 + num_pages: int,
717 + num_kv_heads: int,
718 + kv_cache_dtype: str,
719 + k_scale: torch.Tensor,
720 + v_scale: torch.Tensor,
721 +) -> tuple[torch.Tensor, torch.Tensor]:
722 + """Write random K/V through the NVFP4 ``reshape_and_cache_flash`` writer,
723 + whose pages hold every head's data then every head's scales per side, and
724 + repack them into head slots.
725 +
726 + Returns the packed HND cache ``[pages, 2 * Hkv, 128, 72]`` and its BF16
727 + dequantization ``[pages, Hkv, 128, 2 * 128]`` for the Triton reference.
728 + """
729 + num_slots = num_pages * BLOCK_SIZE
730 + key = torch.randn(
731 + num_slots, num_kv_heads, HEAD_DIM, dtype=torch.bfloat16, device="cuda"
732 + )
733 + value = torch.randn_like(key)
734 + written = torch.zeros(
735 + num_pages,
736 + 2 * num_kv_heads,
737 + BLOCK_SIZE,
738 + HEAD_DIM // 2 + HEAD_DIM // 16,
739 + dtype=torch.uint8,
740 + device="cuda",
741 + )
742 + k_cache, v_cache = written.transpose(1, 2).split(num_kv_heads, dim=-2)
743 + ops.reshape_and_cache_flash(
744 + key,
745 + value,
746 + k_cache,
747 + v_cache,
748 + torch.arange(num_slots, dtype=torch.int64, device="cuda"),
749 + kv_cache_dtype,
750 + k_scale,
751 + v_scale,
752 + )
753 +
754 + kv_cache = torch.empty_like(written)
755 + head_slot_views = iter(nvfp4_kv_cache_views(kv_cache))
756 + dequantized = []
757 + for side, global_scale, swizzled in (
758 + (written[:, :num_kv_heads], k_scale, False),
759 + (written[:, num_kv_heads:], v_scale, True),
760 + ):
761 + data = side.as_strided(
762 + (num_pages, num_kv_heads, BLOCK_SIZE, HEAD_DIM // 2),
763 + (side.stride(0), BLOCK_SIZE * HEAD_DIM // 2, HEAD_DIM // 2, 1),
764 + )
765 + scale = side.as_strided(
766 + (num_pages, num_kv_heads, BLOCK_SIZE, HEAD_DIM // 16),
767 + (side.stride(0), BLOCK_SIZE * HEAD_DIM // 16, HEAD_DIM // 16, 1),
768 + side.storage_offset() + num_kv_heads * BLOCK_SIZE * HEAD_DIM // 2,
769 + )
770 + next(head_slot_views).copy_(data)
771 + next(head_slot_views).copy_(scale)
772 + dequantized.append(
773 + dequant_nvfp4_kv_cache(
774 + data,
775 + scale,
776 + global_scale.item(),
777 + HEAD_DIM,
778 + BLOCK_SIZE,
779 + swizzled_scales=swizzled,
780 + )
781 + )
782 + return kv_cache, torch.cat(dequantized, dim=-1).to(torch.bfloat16)
783 +
784 +
785 +def test_nvfp4_kv_cache_views_address_head_slots() -> None:
786 + """Slot 2h (K) and 2h + 1 (V) hold head h's data then scales, so each head
787 + is one contiguous run of the page: a TP rank's page is a byte slice."""
788 + num_kv_heads, slot_bytes = 4, BLOCK_SIZE * (HEAD_DIM // 2 + HEAD_DIM // 16)
789 + kv_cache = torch.randint(
790 + 0,
791 + 256,
792 + (3, 2 * num_kv_heads, BLOCK_SIZE, HEAD_DIM // 2 + HEAD_DIM // 16),
793 + dtype=torch.uint8,
794 + )
795 + k_data, k_sf, v_data, v_sf = nvfp4_kv_cache_views(kv_cache)
796 + runs = kv_cache.view(3, num_kv_heads, 2 * slot_bytes)
797 + for head in range(num_kv_heads):
798 + expected = torch.cat(
799 + [v[:, head].flatten(1) for v in (k_data, k_sf, v_data, v_sf)], dim=1
800 + )
801 + assert torch.equal(runs[:, head], expected)
802 +
803 +
804 +@pytest.mark.parametrize(
805 + ("num_q_heads", "num_kv_heads", "num_request_pairs", "query_len", "capture"),
806 + [
807 + pytest.param(64, 4, 1, 1, True, id="tp1-small-batch-query-len-1"),
808 + pytest.param(64, 4, 8, 8, False, id="tp1-query-len-8"),
809 + pytest.param(16, 1, 16, 4, True, id="tp4-query-len-4"),
810 + ],
811 +)
812 +def test_msa_cutlass_decode_nvfp4_matches_triton_on_dequantized_cache(
813 + num_q_heads: int,
814 + num_kv_heads: int,
815 + num_request_pairs: int,
816 + query_len: int,
817 + capture: bool,
818 +) -> None:
819 + """CUTLASS decode reads vLLM's packed NVFP4 pages, including swizzled V
820 + scales and device global scales, as their exact dequantized values."""
821 + torch.manual_seed(0)
822 + seq_lens_list = [257, 513] * num_request_pairs
823 + seq_lens_cpu = torch.tensor(seq_lens_list, dtype=torch.int32)
824 + seq_lens = seq_lens_cpu.cuda()
825 + block_table, num_pages = _make_paged_layout(seq_lens_list)
826 + k_scale = torch.tensor(0.5, dtype=torch.float32, device="cuda")
827 + v_scale = torch.tensor(2.0, dtype=torch.float32, device="cuda")
828 + kv_cache, kv_reference = _make_nvfp4_kv_cache(
829 + num_pages, num_kv_heads, "nvfp4", k_scale, v_scale
830 + )
831 +
832 + num_query_tokens = len(seq_lens_list) * query_len
833 + query = torch.randn(
834 + num_query_tokens, num_q_heads, HEAD_DIM, dtype=torch.bfloat16, device="cuda"
835 + )
836 + query_fp8 = query.to(torch.float8_e4m3fn)
837 + topk = _make_topk(seq_lens_list, num_kv_heads, query_len)
838 + expected = torch.empty_like(query)
839 + minimax_m3_sparse_attn_decode(
840 + query_fp8.to(torch.bfloat16),
841 + kv_reference,
842 + topk.transpose(0, 1),
843 + block_table,
844 + seq_lens,
845 + num_kv_heads,
846 + SM_SCALE,
847 + expected,
848 + query_len,
849 + )
850 +
851 + assert should_prepare_decode_metadata(
852 + len(seq_lens_list),
853 + query_len,
854 + decode_backend="triton",
855 + num_q_heads=num_q_heads,
856 + num_kv_heads=num_kv_heads,
857 + kv_cache_dtype="nvfp4",
858 + page_size=BLOCK_SIZE,
859 + topk_blocks=TOPK,
860 + )
861 + metadata = prepare_decode_metadata(
862 + block_table,
863 + seq_lens,
864 + seq_lens_cpu,
865 + query_len,
866 + num_q_heads=num_q_heads,
867 + num_kv_heads=num_kv_heads,
868 + page_size=BLOCK_SIZE,
869 + topk_blocks=TOPK,
870 + )
871 + actual = torch.zeros_like(query)
872 +
873 + def run() -> None:
874 + msa_cutlass_sparse_decode(
875 + query_fp8,
876 + kv_cache,
877 + topk,
878 + actual,
879 + metadata,
880 + scale=SM_SCALE,
881 + q_scale_float=1.0,
882 + k_scale_float=1.0,
883 + v_scale_float=1.0,
884 + k_scale=k_scale,
885 + v_scale=v_scale,
886 + )
887 +
888 + if capture:
889 + run()
890 + graph = torch.cuda.CUDAGraph()
891 + with torch.cuda.graph(graph):
892 + run()
893 + actual.zero_()
894 + graph.replay()
895 + else:
896 + run()
897 + current_platform.synchronize()
898 + torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.05)
899 +
900 +
901 +@pytest.mark.parametrize(
902 + ("num_q_heads", "num_kv_heads"),
903 + [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")],
904 +)
905 +def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache(
906 + num_q_heads: int,
907 + num_kv_heads: int,
908 +) -> None:
909 + """MSA prefill reads NVFP4 pages as their exact dequantized values,
910 + including chunked prefill over cached context, and takes q from its fp8
911 + copy (the only one the fused insert writes on this path)."""
912 + torch.manual_seed(0)
913 + q_lens = [300, 128, 77]
914 + seq_lens_list = [300, 640, 1100]
915 + block_table, num_pages = _make_paged_layout(seq_lens_list)
916 + k_scale = torch.tensor(0.5, dtype=torch.float32, device="cuda")
917 + v_scale = torch.tensor(2.0, dtype=torch.float32, device="cuda")
918 + kv_cache, kv_reference = _make_nvfp4_kv_cache(
919 + num_pages, num_kv_heads, "nvfp4", k_scale, v_scale
920 + )
921 +
922 + total_q = sum(q_lens)
923 + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda")
924 + context_lens = seq_lens - torch.tensor(q_lens, dtype=torch.int32, device="cuda")
925 + cu_seqlens_q = torch.zeros(len(q_lens) + 1, dtype=torch.int32, device="cuda")
926 + cu_seqlens_q[1:] = torch.tensor(q_lens, device="cuda").cumsum(0)
927 + cu_seqlens_k = torch.zeros_like(cu_seqlens_q)
928 + cu_seqlens_k[1:] = seq_lens.cumsum(0)
929 + prefill = MiniMaxM3SparsePrefillMetadata(
930 + cu_seqlens_q=cu_seqlens_q,
931 + cu_seqlens_k=cu_seqlens_k,
932 + seq_lens=seq_lens,
933 + context_lens=context_lens,
934 + block_table=block_table,
935 + max_query_len=max(q_lens),
936 + max_seq_len=max(seq_lens_list),
937 + total_kv_blocks=sum(math.ceil(n / BLOCK_SIZE) for n in seq_lens_list),
938 + )
939 + metadata = MiniMaxM3SparseMetadata(
940 + seq_lens=seq_lens,
941 + max_seq_len=max(seq_lens_list),
942 + slot_mapping=torch.empty(total_q, dtype=torch.int64, device="cuda"),
943 + num_actual_tokens=total_q,
944 + num_decodes=0,
945 + num_decode_tokens=0,
946 + num_prefills=len(q_lens),
947 + num_prefill_tokens=total_q,
948 + prefill=prefill,
949 + )
950 +
951 + # Every causally visible page (<= TOPK per query), ending with the local one.
952 + topk = torch.full(
953 + (total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device="cuda"
954 + )
955 + token = 0
956 + for q_len, seq_len in zip(q_lens, seq_lens_list):
957 + for local_q in range(q_len):
958 + visible = (seq_len - q_len + local_q) // BLOCK_SIZE + 1
959 + assert visible <= TOPK
960 + topk[token, :, :visible] = torch.arange(visible, device="cuda")
961 + token += 1
962 +
963 + layer_name = "model.layers.0.self_attn.attn"
964 + q_scale = 0.5
965 + layer = SimpleNamespace(
966 + layer_name=layer_name,
967 + topk_indices_buffer=topk,
968 + _q_scale_float=q_scale,
969 + _k_scale=k_scale,
970 + _v_scale=v_scale,
971 + )
972 + query_fp8 = (
973 + torch.randn(
974 + total_q, num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda"
975 + )
976 + / q_scale
977 + ).to(torch.float8_e4m3fn)
978 + query = query_fp8.to(torch.bfloat16) * q_scale
979 + common = dict(
980 + num_heads=num_q_heads,
981 + head_size=HEAD_DIM,
982 + scale=SM_SCALE,
983 + num_kv_heads=num_kv_heads,
984 + topk_blocks=TOPK,
985 + sparse_block_size=BLOCK_SIZE,
986 + )
987 + forward_context = ForwardContext(
988 + no_compile_layers={},
989 + attn_metadata={layer_name: metadata},
990 + slot_mapping={},
991 + )
992 + with override_forward_context(forward_context):
993 + expected = MiniMaxM3SparseTritonImpl(**common).forward(
994 + layer, query, kv_reference, torch.empty_like(query)
995 + )
996 + actual = MiniMaxM3SparseMSAImpl(**common, kv_cache_dtype="nvfp4").forward(
997 + layer,
998 + torch.full_like(query, float("nan")),
999 + kv_cache,
1000 + torch.zeros_like(query),
1001 + query_fp8=query_fp8,
1002 + )
1003 + current_platform.synchronize()
1004 + torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.05)
1005 +
1006 +
1007 +def test_msa_triton_decode_fallback_reads_dequantized_query_fp8() -> None:
1008 + """With CUTLASS decode on but not planned for the batch, the Triton decode
1009 + fallback attends with the dequantized fp8 q, the only q written."""
1010 + torch.manual_seed(0)
1011 + num_q_heads, num_kv_heads, q_scale = 16, 1, 0.5
1012 + seq_lens_list = [257, 513, 129]
1013 + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda")
1014 + block_table, num_pages = _make_paged_layout(seq_lens_list)
1015 + kv_cache = (
1016 + torch.randn(
1017 + num_pages,
1018 + num_kv_heads,
1019 + BLOCK_SIZE,
1020 + 2 * HEAD_DIM,
1021 + dtype=torch.bfloat16,
1022 + device="cuda",
1023 + )
1024 + * 0.25
1025 + ).to(torch.float8_e4m3fn)
1026 + num_tokens = len(seq_lens_list)
1027 + metadata = MiniMaxM3SparseMetadata(
1028 + seq_lens=seq_lens,
1029 + max_seq_len=max(seq_lens_list),
1030 + slot_mapping=torch.empty(num_tokens, dtype=torch.int64, device="cuda"),
1031 + num_actual_tokens=num_tokens,
1032 + num_decodes=num_tokens,
1033 + num_decode_tokens=num_tokens,
1034 + num_prefills=0,
1035 + num_prefill_tokens=0,
1036 + decode=MiniMaxM3SparseMSADecodeMetadata(
1037 + seq_lens=seq_lens,
1038 + block_table=block_table,
1039 + decode_query_len=1,
1040 + msa_cutlass=None,
1041 + ),
1042 + )
1043 + layer_name = "model.layers.0.self_attn.attn"
1044 + layer = SimpleNamespace(
1045 + layer_name=layer_name,
1046 + topk_indices_buffer=_make_topk(seq_lens_list, num_kv_heads, 1),
1047 + _q_scale_float=q_scale,
1048 + )
1049 + query_fp8 = torch.randn(
1050 + num_tokens, num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda"
1051 + ).to(torch.float8_e4m3fn)
1052 + query = query_fp8.to(torch.bfloat16) * q_scale
1053 + common = dict(
1054 + num_heads=num_q_heads,
1055 + head_size=HEAD_DIM,
1056 + scale=SM_SCALE,
1057 + num_kv_heads=num_kv_heads,
1058 + kv_cache_dtype="fp8",
1059 + topk_blocks=TOPK,
1060 + sparse_block_size=BLOCK_SIZE,
1061 + )
1062 + impl = MiniMaxM3SparseMSAImpl(**common, msa_decode_backend="cutlass")
1063 + assert impl.use_cutlass_decode
1064 + forward_context = ForwardContext(
1065 + no_compile_layers={},
1066 + attn_metadata={layer_name: metadata},
1067 + slot_mapping={},
1068 + )
1069 + with override_forward_context(forward_context):
1070 + expected = MiniMaxM3SparseTritonImpl(**common).forward(
1071 + layer, query, kv_cache, torch.empty_like(query)
1072 + )
1073 + actual = impl.forward(
1074 + layer,
1075 + torch.full_like(query, float("nan")),
1076 + kv_cache,
1077 + torch.zeros_like(query),
1078 + query_fp8=query_fp8,
1079 + )
1080 + current_platform.synchronize()
1081 + 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
@@ -320,6 +320,38 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
320 index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0
321 )
322
323 + # Without q_out, q is emitted only in fp8: the same bytes, while qkv keeps
324 + # its raw q and every other output matches the run above.
325 + outputs = [qkv_orig.clone(), torch.empty_like(q_fp8), torch.empty_like(index_q)]
326 + caches = [torch.zeros_like(kv_cache), torch.zeros_like(index_cache)]
327 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
328 + outputs[0],
329 + q_w,
330 + k_w,
331 + cos_sin,
332 + positions,
333 + num_heads,
334 + num_kv_heads,
335 + ROTARY_DIM,
336 + eps,
337 + iq_w,
338 + ik_w,
339 + num_idx_heads,
340 + slot_mapping,
341 + index_slot_mapping,
342 + *caches,
343 + block_size,
344 + None,
345 + outputs[2],
346 + kv_cache_dtype,
347 + q_fp8_out=outputs[1],
348 + q_fp8_scale=0.5,
349 + )
350 + for actual, expected in zip(
351 + outputs + caches, [qkv, q_fp8, index_q, kv_cache, index_cache]
352 + ):
353 + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
354 +
355
356 @pytest.mark.parametrize("num_tokens", [1, 64, 513])
357 @pytest.mark.parametrize("block_size", [16, 64])
@@ -441,8 +473,8 @@ def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype):
473 # ── Test 4: fp8 (e4m3) index outputs ─────────────────────────────────────────
474 # The fp8 score path stores index_q and the index-K cache as e4m3 while q/k/v +
475 # q_out stay bf16. Asserts: (1) q/k/v/q_out are bit-identical to the bf16 run
444 -# (the index dtype must not perturb the main branch), and (2) the e4m3 index
445 -# outputs dequantize close to the bf16 reference.
476 +# (the index dtype must not perturb the main branch), and (2) index_q is the
477 +# e4m3 cast of the bf16 run's index_q, and the index-K cache is within fp8 ulp.
478
479
480 @pytest.mark.skipif(
@@ -528,6 +560,196 @@ def test_sparse_full_fp8_index(num_tokens, block_size):
560 torch.testing.assert_close(qkv_fp, qkv_bf, rtol=0, atol=0)
561 torch.testing.assert_close(kvc_fp, kvc_bf, rtol=0, atol=0)
562
531 - # (2) Dequantized e4m3 index outputs match the bf16 reference within fp8 ulp.
532 - torch.testing.assert_close(iq_fp.float(), iq_bf.float(), rtol=0.13, atol=0.05)
563 + # (2) index_q is e4m3(bf16(q)); the index-K cache is within fp8 ulp.
564 + assert torch.equal(
565 + iq_fp.view(torch.uint8), iq_bf.to(torch.float8_e4m3fn).view(torch.uint8)
566 + )
567 torch.testing.assert_close(idxc_fp.float(), idxc_bf.float(), rtol=0.13, atol=0.05)
568 +
569 +
570 +# ── Test 5: sparse mode with an NVFP4 main KV cache ──────────────────────────
571 +# The NVFP4 inserts quantize x * sf_scale, where sf_scale = 1 / the global
572 +# scale: per 16-value block,
573 +# scale = min(448, max(E4M3(max(amax, 1e-12) / 6), 1/512)) and
574 +# code = sign | RNE E2M1 of min(|x / scale|, 6), with zero stored as +0.
575 +
576 +E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)
577 +
578 +
579 +def _div(a, b):
580 + # Division by a Python scalar would become a reciprocal multiply.
581 + return a / torch.as_tensor(b, dtype=a.dtype, device=a.device)
582 +
583 +
584 +def nvfp4_ref(x, sf_scale):
585 + """Returns the E2M1 codes [..., 128] and E4M3 scale bytes [..., 8] of x."""
586 + blocks = x.float().unflatten(-1, (-1, 16))
587 + amax = blocks.abs().amax(-1, keepdim=True).clamp_min(1e-12)
588 + scale = (sf_scale * _div(amax, 6.0)).clamp(max=448.0)
589 + scale = scale.to(torch.float8_e4m3fn).float().clamp(min=2.0**-9)
590 + q = _div(blocks * sf_scale, scale)
591 + # Alternating > / >= rounds exact midpoints to the even code.
592 + mag = sum(
593 + (q.abs() >= m) if i % 2 else (q.abs() > m) for i, m in enumerate(E2M1_MIDPOINTS)
594 + ).to(torch.uint8)
595 + codes = torch.where(mag > 0, mag | (torch.signbit(q).to(torch.uint8) << 3), 0)
596 + scale_bytes = scale.squeeze(-1).to(torch.float8_e4m3fn).view(torch.uint8)
597 + return codes.flatten(-2), scale_bytes
598 +
599 +
600 +def nvfp4_pages(sides, slot_mapping, num_blocks, block_size):
601 + """Expected NVFP4 pages [nb, len(sides) * nkv, bs, 72] for the K (then V)
602 + side given as (x [N, nkv, 128], sf_scale). Slot h * len(sides) + side is
603 + head h's packed E2M1 data [bs, 64] (low nibble = even element), then its
604 + E4M3 scales [bs, 8]: linear for K, 4x4 token-quad swizzled for V."""
605 + nkv, ns, device = sides[0][0].shape[1], HEAD_DIM // 16, slot_mapping.device
606 + row_bytes, data_bytes = HEAD_DIM // 2 + ns, block_size * HEAD_DIM // 2
607 + pages = torch.zeros(
608 + num_blocks,
609 + nkv,
610 + len(sides),
611 + block_size * row_bytes,
612 + dtype=torch.uint8,
613 + device=device,
614 + )
615 + valid = slot_mapping >= 0
616 + blk, tok = slot_mapping[valid] // block_size, slot_mapping[valid] % block_size
617 + g, t = torch.arange(ns, device=device), tok[:, None]
618 + scale_offsets = (t * ns + g, ((t // 4) * 4 + g // 2) * ns + (g % 2) * 4 + t % 4)
619 + for side, (x, sf_scale) in enumerate(sides):
620 + codes, scales = nvfp4_ref(x[valid], sf_scale)
621 + data = pages[:, :, side, :data_bytes].view(num_blocks, nkv, block_size, -1)
622 + data[blk, :, tok] = codes[..., 0::2] | (codes[..., 1::2] << 4)
623 + scale_region = pages[:, :, side, data_bytes:]
624 + for h in range(nkv):
625 + scale_region[blk[:, None], h, scale_offsets[side]] = scales[:, h]
626 + return pages.view(num_blocks, -1, block_size, row_bytes)
627 +
628 +
629 +def nvfp4_edge_values(num_rows, global_scale):
630 + """[num_rows, 128] bf16 whose 16-value blocks, divided by global_scale, hit
631 + the quantizer's corner cases: exact E2M1 midpoints under non-power-of-2 scales
632 + (blocks 0-1), amax under the 1/512 scale floor (2), all -0 (3), and
633 + saturation with |x| > 4096 up to bf16 max (4) or with inf (5)."""
634 + g = torch.Generator().manual_seed(0)
635 + x = torch.randn(num_rows, 8, 16, generator=g)
636 + sign = torch.where(torch.rand(x.shape, generator=g) < 0.5, -1.0, 1.0)
637 + steps = torch.tensor(E2M1_MIDPOINTS + (0.5, 1.0, 3.0))
638 + mant = 1 + torch.randint(1, 8, (num_rows, 2, 1), generator=g) / 8
639 + scale = mant * 2.0 ** torch.randint(-6, 3, (num_rows, 2, 1), generator=g)
640 + x[:, :2] = steps[torch.randint(len(steps), (num_rows, 2, 16), generator=g)] * scale
641 + x[:, :2, :1] = 6 * scale # amax = 6 * scale, so the block scale is `scale`
642 + x[:, 2] = torch.rand(num_rows, 16, generator=g) * 2.0**-8
643 + x[:, 4] *= 5000
644 + x = x * sign * global_scale
645 + x[:, 3] = -0.0
646 + x[:, 4, 0] = sign[:, 4, 0] * torch.finfo(torch.bfloat16).max
647 + x[:, 5, 0] = sign[:, 5, 0] * float("inf")
648 + return x.flatten(1).to(torch.bfloat16)
649 +
650 +
651 +@pytest.mark.skipif(
652 + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
653 + reason="NVFP4 KV caches need SM100",
654 +)
655 +@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
656 +@pytest.mark.parametrize("num_kv_heads", [1, 4])
657 +@pytest.mark.parametrize("block_size", [16, 128])
658 +@pytest.mark.parametrize("kv_scales", [(0.5, 2.0), (1.0, 1.0)])
659 +def test_sparse_full_nvfp4(num_tokens, num_kv_heads, block_size, kv_scales):
660 + """The fused NVFP4 insert writes the reference quantizer's bytes (packed
661 + E2M1, linear K / swizzled V E4M3 block scales) for the kernel's own
662 + post-RoPE k and for a raw v that hits the quantizer's corner cases, and
663 + skips padded slots. Unit global scales take the kernel's separate fast
664 + path."""
665 + kv_cache_dtype = "nvfp4"
666 + torch.manual_seed(5)
667 + device, dtype, eps = "cuda", torch.bfloat16, 1e-6
668 + base, max_pos = 5_000_000.0, 4096
669 + num_heads, num_idx_heads = 4 * num_kv_heads, 1
670 +
671 + weights = [
672 + torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1 for _ in range(4)
673 + ]
674 + cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
675 + positions = torch.randint(
676 + 0, max_pos, (num_tokens,), dtype=torch.int64, device=device
677 + )
678 + qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
679 + iqsz = num_idx_heads * HEAD_DIM
680 + qkv = torch.randn(
681 + num_tokens, qsz + 2 * kvsz + iqsz + HEAD_DIM, dtype=dtype, device=device
682 + )
683 + k_scale, v_scale = (
684 + torch.tensor(s, dtype=torch.float32, device=device) for s in kv_scales
685 + )
686 + qkv[:, qsz + kvsz : qsz + 2 * kvsz] = nvfp4_edge_values(
687 + num_tokens * num_kv_heads, v_scale.item()
688 + ).view(num_tokens, kvsz)
689 + qkv_orig = qkv.clone()
690 +
691 + num_blocks = (num_tokens + block_size - 1) // block_size + 1
692 + kv_cache = torch.zeros(
693 + num_blocks,
694 + 2 * num_kv_heads,
695 + block_size,
696 + HEAD_DIM // 2 + HEAD_DIM // 16,
697 + dtype=torch.uint8,
698 + device=device,
699 + )
700 + index_cache = torch.zeros(
701 + num_blocks, block_size, HEAD_DIM, dtype=dtype, device=device
702 + )
703 + slot_mapping = torch.randperm(
704 + num_blocks * block_size, dtype=torch.int64, device=device
705 + )[:num_tokens]
706 + slot_mapping[::5] = -1 # padded / unscheduled tokens are not written
707 + index_slot_mapping = torch.randperm(
708 + num_blocks * block_size, dtype=torch.int64, device=device
709 + )[:num_tokens]
710 +
711 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
712 + qkv,
713 + *weights[:2],
714 + cos_sin,
715 + positions,
716 + num_heads,
717 + num_kv_heads,
718 + ROTARY_DIM,
719 + eps,
720 + *weights[2:],
721 + num_idx_heads,
722 + slot_mapping,
723 + index_slot_mapping,
724 + kv_cache,
725 + index_cache,
726 + block_size,
727 + torch.empty(num_tokens, qsz, dtype=dtype, device=device),
728 + torch.empty(num_tokens, iqsz, dtype=dtype, device=device),
729 + kv_cache_dtype,
730 + kv_k_scale=k_scale,
731 + kv_v_scale=v_scale,
732 + )
733 +
734 + _, k_out, v_out, _, index_k = qkv.split([qsz, kvsz, kvsz, iqsz, HEAD_DIM], dim=-1)
735 + torch.testing.assert_close(
736 + v_out, qkv_orig.split([qsz, kvsz, kvsz, iqsz, HEAD_DIM], dim=-1)[2]
737 + )
738 + expected_kv_cache = nvfp4_pages(
739 + [
740 + (k_out.reshape(num_tokens, num_kv_heads, HEAD_DIM), 1 / k_scale),
741 + (v_out.reshape(num_tokens, num_kv_heads, HEAD_DIM), 1 / v_scale),
742 + ],
743 + slot_mapping,
744 + num_blocks,
745 + block_size,
746 + )
747 + assert torch.equal(kv_cache, expected_kv_cache), (
748 + f"{int((kv_cache != expected_kv_cache).sum())} NVFP4 cache bytes differ"
749 + )
750 +
751 + expected_index_cache = torch.zeros_like(index_cache).view(-1, HEAD_DIM)
752 + expected_index_cache[index_slot_mapping] = index_k
753 + torch.testing.assert_close(
754 + index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0
755 + )
vllm/_custom_ops.py
+12 -1
@@ -2587,6 +2587,8 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
2587 skip_index_branch: bool = False,
2588 q_fp8_out: torch.Tensor | None = None,
2589 q_fp8_scale: float = 1.0,
2590 + kv_k_scale: torch.Tensor | None = None,
2591 + kv_v_scale: torch.Tensor | None = None,
2592 ) -> None:
2593 """Fused MiniMax-M3 attention pre-processing (in-place).
2594
@@ -2610,7 +2612,14 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
2612 attention's flat TMA descriptor.
2613
2614 If ``q_fp8_out`` is given, the same normalized q is also written in FP8
2613 - E4M3 using ``q_fp8_scale`` as its dequantization scale.
2615 + E4M3 using ``q_fp8_scale`` as its dequantization scale. Without ``q_out``,
2616 + q is then written only in FP8 and the q slice of ``qkv`` is left as is.
2617 +
2618 + ``kv_cache_dtype="nvfp4"`` quantizes k/v into the packed HND
2619 + ``[num_blocks, 2 * num_kv_heads, block_size, 72]`` uint8 cache using the
2620 + device dequantization scales ``kv_k_scale``/``kv_v_scale``. Slot
2621 + ``2 * head + side`` (K = 0, V = 1) holds that head's E2M1 data followed by
2622 + its E4M3 block scales, so every head is one contiguous run of the page.
2623
2624 When ``skip_index_branch`` is true, sparse rows still keep their packed
2625 ``[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(
2650 skip_index_branch,
2651 q_fp8_out,
2652 q_fp8_scale,
2653 + kv_k_scale,
2654 + kv_v_scale,
2655 )
2656
2657
vllm/models/minimax_m3/common/sparse_attention.py
+19 -3
@@ -56,6 +56,7 @@ from vllm.v1.attention.backends.utils import (
56 from vllm.v1.kv_cache_interface import (
57 AttentionSpec,
58 KVCacheLayout,
59 + get_kv_quant_mode,
60 is_quantized_kv_cache,
61 )
62
@@ -531,10 +532,10 @@ def select_main_backend_and_impl_cls(
532 """Pick the main attention backend and implementation.
533
534 Blackwell (SM100) uses the MSA attend for supported top-k block counts
534 - when the KV cache is BF16 or FP8 E4M3; MI355 uses AITER sparse PA
535 + when the KV cache is BF16, FP8 E4M3 or NVFP4; MI355 uses AITER sparse PA
536 with shuffle KV cache layout; Other platforms and FP8 E5M2 fall
536 - back to Triton. The MSA modules are imported lazily to avoid import errors
537 - on unsupported platforms.
537 + back to Triton. NVFP4 is MSA-only. The MSA modules are imported lazily to
538 + avoid import errors on unsupported platforms.
539 """
540 use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(
541 num_kv_heads, emits_sparse_block_table=emits_sparse_block_table
@@ -545,6 +546,18 @@ def select_main_backend_and_impl_cls(
546 and topk_blocks in (4, 8, 16, 32)
547 and kv_cache_dtype != "fp8_e5m2"
548 )
549 + use_nvfp4 = get_kv_quant_mode(kv_cache_dtype).is_nvfp4
550 + if use_nvfp4 and kv_cache_dtype != "nvfp4":
551 + raise ValueError(
552 + "MiniMax M3 sparse attention supports only the 'nvfp4' NVFP4 KV "
553 + f"cache dtype, got {kv_cache_dtype}"
554 + )
555 + if use_nvfp4 and (use_aiter_sparse_pa or not use_msa):
556 + raise ValueError(
557 + "MiniMax M3 sparse attention supports an NVFP4 KV cache only with "
558 + f"the SM100 MSA backend (kv_cache_dtype={kv_cache_dtype}, "
559 + f"topk_blocks={topk_blocks})"
560 + )
561 selected = (
562 "AITER_SPARSE_PA" if use_aiter_sparse_pa else ("MSA" if use_msa else "Triton")
563 )
@@ -565,8 +578,11 @@ def select_main_backend_and_impl_cls(
578 from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
579 MiniMaxM3SparseMSABackend,
580 MiniMaxM3SparseMSAImpl,
581 + MiniMaxM3SparseMSANvfp4Backend,
582 )
583
584 + if use_nvfp4:
585 + return MiniMaxM3SparseMSANvfp4Backend, MiniMaxM3SparseMSAImpl
586 return MiniMaxM3SparseMSABackend, MiniMaxM3SparseMSAImpl
587 return MiniMaxM3SparseBackend, MiniMaxM3SparseTritonImpl
588
vllm/models/minimax_m3/nvidia/model.py
+7 -1
@@ -500,6 +500,8 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
500 # of wrapping the generic Attention module. Keep the same runtime scale
501 # attributes so FP8 KV reads can honor vLLM's per-layer descale contract.
502 set_default_quant_scales(self, register_buffer=True)
503 + # The fused insert quantizes k/v into NVFP4 pages with these scales.
504 + self.use_nvfp4_kv = get_kv_quant_mode(self.kv_cache_dtype).is_nvfp4
505 # Indexer side-cache dtype, mirroring --kv-cache-dtype for the main
506 # cache (--attention-config '{"indexer_kv_dtype": ...}').
507 self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
@@ -611,6 +613,8 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
613 main_slot_mapping = fwd_slot_mapping[self.layer_name]
614 index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
615 q = qkv.new_empty((num_tokens, self.q_size))
616 + # With query_fp8, q is emitted only in fp8; the attend dequantizes it
617 + # into ``q`` for the tokens that run on bf16-query kernels.
618 query_fp8 = self._allocate_query_fp8(qkv)
619 # index_q matches the index-K cache dtype (e4m3 for the fp8 score path);
620 # the fused kernel emits fp8 directly when this buffer is e4m3.
@@ -636,11 +640,13 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
640 self.kv_cache,
641 self.indexer.index_cache.kv_cache,
642 self.kv_cache.size(2), # paged-cache block size
639 - q,
643 + q if query_fp8 is None else None,
644 index_q,
645 self.kv_cache_dtype,
646 q_fp8_out=query_fp8,
647 q_fp8_scale=self._q_scale_float,
648 + kv_k_scale=self._k_scale if self.use_nvfp4_kv else None,
649 + kv_v_scale=self._v_scale if self.use_nvfp4_kv else None,
650 )
651
652 output = torch.empty_like(q)
vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py
+52 -7
@@ -12,6 +12,7 @@ import torch
12 from vllm.config.attention import MiniMaxM3MSADecodeBackend
13 from vllm.platforms import current_platform
14 from vllm.triton_utils import tl, triton
15 +from vllm.v1.kv_cache_interface import get_kv_quant_mode
16
17 _MAX_NUM_Q_HEADS = 64
18 _MAX_NUM_KV_HEADS = 4
@@ -23,9 +24,36 @@ _TOPK = 16
24 _MAX_QUERY_HEAD_ROWS = 65536
25 _MAX_DECODE_QUERY_LEN = 32
26 # Kernel benchmarks put the CUTLASS crossover at 16 requests for TP1 and TP4.
27 +# NVFP4 has no Triton fallback, so it takes CUTLASS at every batch size.
28 _MIN_CUTLASS_BATCH_SIZE = 16
29
30
31 +def is_nvfp4_kv_cache(kv_cache_dtype: str) -> bool:
32 + return get_kv_quant_mode(kv_cache_dtype).is_nvfp4
33 +
34 +
35 +def nvfp4_kv_cache_views(
36 + kv_cache: torch.Tensor,
37 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
38 + """Return the K data, K scale, V data and V scale views of an NVFP4 cache.
39 +
40 + ``kv_cache`` is uint8 ``[pages, 2 * Hkv, page_size, 72]``; slot
41 + ``2 * head + side`` (K = 0, V = 1) is that head's E2M1 data
42 + ``[page_size, 64]`` followed by its E4M3 block scales ``[page_size, 8]``,
43 + the one page layout ``fmha_sm100`` reads.
44 + """
45 + from vllm.third_party.fmha_sm100.nvfp4_kv import nvfp4_head_slot_views
46 +
47 + return nvfp4_head_slot_views(*nvfp4_kv_cache_slots(kv_cache))
48 +
49 +
50 +def nvfp4_kv_cache_slots(
51 + kv_cache: torch.Tensor,
52 +) -> tuple[torch.Tensor, torch.Tensor]:
53 + """Return the per-head K and V slots of an NVFP4 cache."""
54 + return kv_cache[:, 0::2], kv_cache[:, 1::2]
55 +
56 +
57 @dataclass
58 class MSACutlassDecodeMetadata:
59 plan: Any
@@ -91,6 +119,9 @@ class MSACutlassDecodePlanCache:
119 use_fp8_kvcache=True,
120 split_prefill_decode=False,
121 device=device,
122 + # The Q8KV4 route keeps plan-time lengths and its own page offsets,
123 + # which this cached plan does not refresh.
124 + decode_backend="kv_mode3",
125 )
126
127 plan_info = plan[3]
@@ -214,11 +245,12 @@ def supports_cutlass_sparse_decode(
245 topk_blocks: int,
246 ) -> bool:
247 """Return whether static model geometry supports CUTLASS sparse decode."""
248 + nvfp4 = is_nvfp4_kv_cache(kv_cache_dtype)
249 return (
218 - decode_backend == "cutlass"
250 + (decode_backend == "cutlass" or nvfp4)
251 and current_platform.is_cuda()
252 and current_platform.is_device_capability_family(100)
221 - and kv_cache_dtype in ("fp8", "fp8_e4m3")
253 + and (kv_cache_dtype in ("fp8", "fp8_e4m3") or nvfp4)
254 and _supported_head_geometry(num_q_heads, num_kv_heads)
255 and page_size == _PAGE_SIZE
256 and topk_blocks == _TOPK
@@ -248,7 +280,7 @@ def should_prepare_decode_metadata(
280 topk_blocks=topk_blocks,
281 )
282 and 1 <= decode_query_len <= _MAX_DECODE_QUERY_LEN
251 - and batch_size >= _MIN_CUTLASS_BATCH_SIZE
283 + and (batch_size >= _MIN_CUTLASS_BATCH_SIZE or is_nvfp4_kv_cache(kv_cache_dtype))
284 and total_q * num_q_heads <= _MAX_QUERY_HEAD_ROWS
285 )
286
@@ -292,9 +324,22 @@ def msa_cutlass_sparse_decode(
324 q_scale_float: float,
325 k_scale_float: float,
326 v_scale_float: float,
327 + k_scale: torch.Tensor | None = None,
328 + v_scale: torch.Tensor | None = None,
329 ) -> None:
296 - """Run CUTLASS sparse decode with metadata prepared by the MSA builder."""
297 - key, value = kv_cache.split(_HEAD_DIM, dim=-1)
330 + """Run CUTLASS sparse decode with metadata prepared by the MSA builder.
331 +
332 + NVFP4 caches are uint8 head-slot pages (see ``nvfp4_kv_cache_views``) and
333 + take the device global scales ``k_scale``/``v_scale``.
334 + """
335 + if kv_cache.dtype == torch.uint8:
336 + assert k_scale is not None and v_scale is not None
337 + key, value = nvfp4_kv_cache_slots(kv_cache)
338 + k_scale_arg: float | torch.Tensor = k_scale
339 + v_scale_arg: float | torch.Tensor = v_scale
340 + else:
341 + key, value = kv_cache.split(_HEAD_DIM, dim=-1)
342 + k_scale_arg, v_scale_arg = k_scale_float, v_scale_float
343
344 from vllm.third_party.fmha_sm100.api import fmha_sm100
345
@@ -310,7 +355,7 @@ def msa_cutlass_sparse_decode(
355 output_o=True,
356 sm_scale=scale,
357 q_scale=q_scale_float,
313 - k_scale=k_scale_float,
314 - v_scale=v_scale_float,
358 + k_scale=k_scale_arg,
359 + v_scale=v_scale_arg,
360 o_scale=1.0,
361 )
vllm/models/minimax_m3/nvidia/sparse_attention_msa.py
+108 -11
@@ -4,15 +4,18 @@
4
5 Prefill attends with ``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``).
6 Decode uses Triton split-K by default, with an opt-in CUTLASS ``fmha_sm100``
7 -path for regular decode and speculative verification.
7 +path for regular decode and speculative verification. NVFP4 KV caches read the
8 +vLLM packed pages directly in both prefill and CUTLASS decode.
9 """
10
10 -from dataclasses import dataclass
11 +from dataclasses import dataclass, replace
12 +from typing import ClassVar
13
14 import torch
15
16 from vllm.config import VllmConfig
17 from vllm.config.attention import MiniMaxM3MSADecodeBackend
18 +from vllm.config.cache import CacheDType
19 from vllm.forward_context import get_forward_context
20 from vllm.logger import init_logger
21 from vllm.models.minimax_m3.common.ops.sparse_attn import (
@@ -29,20 +32,33 @@ from vllm.models.minimax_m3.common.sparse_attention import (
32 from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import (
33 MSACutlassDecodeMetadata,
34 MSACutlassDecodePlanCache,
35 + is_nvfp4_kv_cache,
36 msa_cutlass_sparse_decode,
37 + nvfp4_kv_cache_views,
38 prepare_decode_metadata,
39 should_prepare_decode_metadata,
40 supports_cutlass_sparse_decode,
41 )
42 +from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim
43 from vllm.v1.attention.backend import (
44 AttentionLayer,
45 CommonAttentionMetadata,
46 )
41 -from vllm.v1.kv_cache_interface import AttentionSpec
47 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
48
49 logger = init_logger(__name__)
50
51
52 +def _dequantize_query(
53 + query_fp8: torch.Tensor, q_scale: float, out: torch.Tensor
54 +) -> torch.Tensor:
55 + """Write ``query_fp8 * q_scale`` into ``out`` (E4M3 is exact in BF16)."""
56 + out.copy_(query_fp8)
57 + if q_scale != 1.0:
58 + out.mul_(q_scale)
59 + return out
60 +
61 +
62 class MiniMaxM3SparseMSABackend(MiniMaxM3SparseBackend):
63 """MiniMax M3 backend with NVIDIA MSA-specific decode metadata."""
64
@@ -51,6 +67,30 @@ class MiniMaxM3SparseMSABackend(MiniMaxM3SparseBackend):
67 return MiniMaxM3SparseMSAMetadataBuilder
68
69
70 +class MiniMaxM3SparseMSANvfp4Backend(MiniMaxM3SparseMSABackend):
71 + """MSA backend over HND NVFP4 KV cache pages in per-head K/V slots."""
72 +
73 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["nvfp4"]
74 +
75 + @classmethod
76 + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
77 + """Store each head's K and V as two slots of packed data + block scales
78 + (see ``nvfp4_kv_cache_views``)."""
79 + if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_nvfp4:
80 + return spec
81 + return replace(
82 + spec,
83 + num_head_slots=2 * spec.num_kv_heads,
84 + state_content_bytes=nvfp4_kv_cache_full_dim(spec.head_size)
85 + * get_dtype_size(spec.dtype),
86 + )
87 +
88 + @classmethod
89 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None:
90 + # Every head's K/V slots are one contiguous run of the page.
91 + return (KVCacheLayout.LBHNC,)
92 +
93 +
94 class MiniMaxM3SparseCutlassBackend(MiniMaxM3SparseMSABackend):
95 """Attention-backend alias selecting CUTLASS MSA sparse decode."""
96
@@ -166,6 +206,7 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
206 topk_blocks=topk_blocks,
207 sparse_block_size=sparse_block_size,
208 )
209 + self.use_nvfp4_kv = is_nvfp4_kv_cache(kv_cache_dtype)
210 self.use_cutlass_decode = supports_cutlass_sparse_decode(
211 decode_backend=msa_decode_backend,
212 num_q_heads=self.num_heads,
@@ -174,6 +215,12 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
215 page_size=self.block_size,
216 topk_blocks=self.topk_blocks,
217 )
218 + if self.use_nvfp4_kv and not self.use_cutlass_decode:
219 + raise ValueError(
220 + "MiniMax M3 NVFP4 KV cache requires CUTLASS MSA sparse decode "
221 + f"(num_heads={self.num_heads}, num_kv_heads={self.num_kv_heads}, "
222 + f"block_size={self.block_size}, topk_blocks={self.topk_blocks})"
223 + )
224 logger.info_once(
225 "MiniMax M3 MSA sparse decode selected %s",
226 "CUTLASS" if self.use_cutlass_decode else "Triton",
@@ -218,9 +265,17 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
265 hd = self.head_size
266 q = query[:num_tokens].view(-1, self.num_heads, hd)
267 out = output[:num_tokens].view(-1, self.num_heads, hd)
221 - kv_cache = (
222 - kv_cache.view(self.kv_cache_fp8_dtype) if self.use_fp8_kv else kv_cache
268 + # Given query_fp8, it is the only q the fused insert wrote: CUTLASS
269 + # decode reads it directly, and the bf16-query kernels read it
270 + # dequantized into ``query``.
271 + q_fp8 = (
272 + None
273 + if query_fp8 is None
274 + else query_fp8[:num_tokens].view(-1, self.num_heads, hd)
275 )
276 + q_scale_float = getattr(layer, "_q_scale_float", 1.0)
277 + if self.use_fp8_kv and not self.use_nvfp4_kv:
278 + kv_cache = kv_cache.view(self.kv_cache_fp8_dtype)
279 k_scale = getattr(layer, "_k_scale", None) if self.use_fp8_kv else None
280 v_scale = getattr(layer, "_v_scale", None) if self.use_fp8_kv else None
281
@@ -234,21 +289,31 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
289 else None
290 )
291 if self.use_cutlass_decode and msa_metadata is not None:
237 - assert query_fp8 is not None
292 + assert q_fp8 is not None
293 msa_cutlass_sparse_decode(
239 - query_fp8[:nd].view(-1, self.num_heads, hd),
294 + q_fp8[:nd],
295 kv_cache,
296 topk[:nd],
297 out[:nd],
298 msa_metadata,
299 scale=self.scale,
245 - q_scale_float=getattr(layer, "_q_scale_float", 1.0),
300 + q_scale_float=q_scale_float,
301 k_scale_float=getattr(layer, "_k_scale_float", 1.0),
302 v_scale_float=getattr(layer, "_v_scale_float", 1.0),
303 + k_scale=k_scale if self.use_nvfp4_kv else None,
304 + v_scale=v_scale if self.use_nvfp4_kv else None,
305 + )
306 + elif self.use_nvfp4_kv:
307 + raise RuntimeError(
308 + "MiniMax M3 NVFP4 KV cache has no Triton decode fallback; "
309 + f"CUTLASS decode was not planned for {main_md.num_decodes} "
310 + f"requests with query length {d.decode_query_len}"
311 )
312 else:
313 minimax_m3_sparse_attn_decode(
251 - q[:nd],
314 + q[:nd]
315 + if q_fp8 is None
316 + else _dequantize_query(q_fp8[:nd], q_scale_float, q[:nd]),
317 kv_cache,
318 topk[:nd].transpose(0, 1),
319 d.block_table,
@@ -266,6 +331,7 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
331 from vllm.third_party.fmha_sm100.sparse import (
332 build_k2q_csr,
333 sparse_atten_func,
334 + sparse_atten_nvfp4_kv_func,
335 )
336
337 p = main_md.prefill
@@ -273,8 +339,11 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
339 # [H, prefill, MK] transposed view; build_k2q_csr consumes the
340 # strided view directly (topK stays innermost-contiguous).
341 prefill_topk = topk[nd:num_tokens].transpose(0, 1)
276 - qp = q[nd:]
277 - k_cache, v_cache = kv_cache.split(self.head_size, dim=-1)
342 + qp = (
343 + q[nd:]
344 + if q_fp8 is None
345 + else _dequantize_query(q_fp8[nd:], q_scale_float, q[nd:])
346 + )
347 k2q_row_ptr, k2q_q_indices, schedule = build_k2q_csr(
348 prefill_topk,
349 p.cu_seqlens_q,
@@ -287,6 +356,34 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
356 qhead_per_kv=qp.shape[1] // self.num_kv_heads,
357 return_schedule=True,
358 )
359 + if self.use_nvfp4_kv:
360 + k_data, k_sf, v_data, v_sf = nvfp4_kv_cache_views(kv_cache)
361 + sparse_atten_nvfp4_kv_func(
362 + qp,
363 + k_data,
364 + v_data,
365 + k_sf,
366 + v_sf,
367 + k_scale,
368 + v_scale,
369 + k2q_row_ptr,
370 + k2q_q_indices,
371 + self.topk_blocks,
372 + cu_seqlens_q=p.cu_seqlens_q,
373 + cu_seqlens_k=p.cu_seqlens_k,
374 + max_seqlen_q=p.max_query_len,
375 + max_seqlen_k=p.max_seq_len,
376 + blk_kv=SPARSE_BLOCK_SIZE,
377 + causal=True,
378 + softmax_scale=self.scale,
379 + page_table=p.block_table,
380 + seqused_k=p.seq_lens,
381 + schedule=schedule,
382 + out=out[nd:],
383 + kv_layout="vllm",
384 + )
385 + return output
386 + k_cache, v_cache = kv_cache.split(self.head_size, dim=-1)
387 sparse_atten_func(
388 qp,
389 k_cache,