@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6691f087a6

[Minimax-M3] BF16/FP8 Indexer using MSA (#45892)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>

Yongye Zhu committed Jun 23, 2026 at 10:28 UTC 6691f087a65bc161192ced91360bf11313828258
13 files changed +1048 -104
cmake/external_projects/fmha_sm100.cmake
+26 -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 544eee5e09ae2dfa774d5b06739013f9b7402c57
20 + GIT_TAG fee783153f3efe57e3e933c5cb7e267a7cebcfb5
21 GIT_PROGRESS TRUE
22 CONFIGURE_COMMAND ""
23 BUILD_COMMAND ""
@@ -36,13 +36,38 @@ set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100")
36
37 install(FILES
38 "${FMHA_SM100_PY_ROOT}/__init__.py"
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}/sparse.py"
43 + "${FMHA_SM100_PY_ROOT}/sparse_fmha_adapter.py"
44 DESTINATION vllm/third_party/fmha_sm100
45 COMPONENT fmha_sm100)
46
47 +install(DIRECTORY "${FMHA_SM100_PY_ROOT}/csrc/"
48 + DESTINATION vllm/third_party/fmha_sm100/csrc
49 + COMPONENT fmha_sm100
50 + PATTERN "__pycache__" EXCLUDE
51 + PATTERN "*.pyc" EXCLUDE
52 + PATTERN ".git*" EXCLUDE)
53 +
54 install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/"
55 DESTINATION vllm/third_party/fmha_sm100/cute
56 COMPONENT fmha_sm100
57 PATTERN "__pycache__" EXCLUDE
58 PATTERN "*.pyc" EXCLUDE
59 PATTERN ".git*" EXCLUDE)
60 +
61 +install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/include/"
62 + DESTINATION vllm/third_party/fmha_sm100/cutlass/include
63 + COMPONENT fmha_sm100
64 + PATTERN "__pycache__" EXCLUDE
65 + PATTERN "*.pyc" EXCLUDE
66 + PATTERN ".git*" EXCLUDE)
67 +
68 +install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/tools/util/include/"
69 + DESTINATION vllm/third_party/fmha_sm100/cutlass/tools/util/include
70 + COMPONENT fmha_sm100
71 + PATTERN "__pycache__" EXCLUDE
72 + PATTERN "*.pyc" EXCLUDE
73 + PATTERN ".git*" EXCLUDE)
csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
+152 -51
@@ -67,6 +67,13 @@
67 #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
68 #endif
69
70 +// Direct float -> E4M3 FP8 conversion for the indexer Q / index-K outputs.
71 +#ifndef USE_ROCM
72 + #include <cuda_fp8.h>
73 +#else
74 + #include <hip/hip_fp8.h>
75 +#endif
76 +
77 #ifndef FINAL_MASK
78 #ifdef USE_ROCM
79 #define FINAL_MASK 0xffffffffffffffffULL
@@ -75,6 +82,19 @@
82 #endif
83 #endif
84
85 +#ifdef USE_ROCM
86 +// ROCm-compatible direct float -> E4M3 FP8 conversion (mirrors the DeepSeek V4
87 +// fused kernel).
88 +__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
89 + #if defined(HIP_FP8_TYPE_OCP)
90 + __hip_fp8_e4m3 fp8_val(val);
91 + #else
92 + __hip_fp8_e4m3_fnuz fp8_val(val);
93 + #endif
94 + return reinterpret_cast<uint8_t&>(fp8_val);
95 +}
96 +#endif
97 +
98 namespace vllm {
99 namespace minimax_m3_fused_ops {
100
@@ -193,6 +213,8 @@ __device__ __forceinline__ void storeElems(
213 *reinterpret_cast<uint2*>(dst) = v;
214 }
215
216 +// Main K/V cache store. kAuto = unquantized (cache_t == scalar_t); fp8 cache
217 +// dtypes use the scaled-convert path with identity scale.
218 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
219 __device__ __forceinline__ void storeCacheElems(
220 cache_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
@@ -208,6 +230,32 @@ __device__ __forceinline__ void storeCacheElems(
230 }
231 }
232
233 +// Store 4 fp32 registers -> 4 contiguous E4M3 FP8 bytes (direct cast,
234 +// saturating to ±448). Used for the fp8 indexer-Q / index-K outputs; no scale
235 +// (RMSNorm outputs are O(1) and the score path only needs relative block
236 +// ordering).
237 +__device__ __forceinline__ void storeElemsFp8(
238 + uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
239 + constexpr float kFp8Max = 448.0f;
240 +#ifndef USE_ROCM
241 + __nv_fp8x2_storage_t out2[kElemsPerLane / 2];
242 + #pragma unroll
243 + for (int i = 0; i < kElemsPerLane / 2; i++) {
244 + float2 vv = make_float2(elems[2 * i], elems[2 * i + 1]);
245 + vv.x = fminf(fmaxf(vv.x, -kFp8Max), kFp8Max);
246 + vv.y = fminf(fmaxf(vv.y, -kFp8Max), kFp8Max);
247 + out2[i] = __nv_cvt_float2_to_fp8x2(vv, __NV_SATFINITE, __NV_E4M3);
248 + }
249 + *reinterpret_cast<uint32_t*>(dst) = *reinterpret_cast<uint32_t const*>(out2);
250 +#else
251 + #pragma unroll
252 + for (int i = 0; i < kElemsPerLane; i++) {
253 + float vv = fminf(fmaxf(elems[i], -kFp8Max), kFp8Max);
254 + dst[i] = rocm_cvt_float_to_fp8_e4m3(vv);
255 + }
256 +#endif
257 +}
258 +
259 // ────────────────────────────────────────────────────────────────────────────
260 // Kernel
261 // ────────────────────────────────────────────────────────────────────────────
@@ -224,12 +272,14 @@ __device__ __forceinline__ void storeCacheElems(
272 // V : nkv only if kInsertKV (V-cache insert; no warps in dense)
273 // IQ: niq only if kIsSparse (norm+RoPE)
274 // IK: 1 only if kIsSparse (norm+RoPE; +index-cache insert)
275 +// cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx:
276 +// indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte).
277 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
228 - bool kIsSparse, bool kInsertKV>
278 + typename out_idx_t, bool kIsSparse, bool kInsertKV, bool kFp8Idx>
279 __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
280 scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse)
231 - scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
232 - scalar_t* __restrict__ index_q_out, // [N, niq*128] contiguous, or nullptr
281 + scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
282 + out_idx_t* __restrict__ index_q_out, // [N, niq*128]; scalar_t or e4m3 byte
283 scalar_t const* __restrict__ q_norm_w,
284 scalar_t const* __restrict__ k_norm_w,
285 scalar_t const* __restrict__ iq_norm_w,
@@ -238,8 +288,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
288 int64_t const* __restrict__ positions, // [N] i64
289 int64_t const* __restrict__ slot_mapping, // main K/V slots or nullptr
290 int64_t const* __restrict__ index_slot_mapping, // index K slots/nullptr
241 - cache_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr
242 - scalar_t* __restrict__ index_cache, // [nb*bs, 128] or nullptr
291 + cache_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr
292 + out_idx_t* __restrict__ index_cache, // [nb*bs, 128]; scalar_t or e4m3 byte
293 float const eps, int const rotary_dim, int const num_tokens, int const nq,
294 int const nkv, int const niq, int const block_size,
295 // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
@@ -334,9 +384,12 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
384 store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
385 slot * kHeadDim;
386 } else if (isIQ && index_q_out != nullptr) {
337 - store_ptr = index_q_out +
338 - static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
339 - (slot - iq_begin) * kHeadDim;
387 + // bf16 index_q_out: gather here. fp8: written by the explicit fp8 store.
388 + if constexpr (!kFp8Idx) {
389 + store_ptr = index_q_out +
390 + static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
391 + (slot - iq_begin) * kHeadDim;
392 + }
393 }
394
395 // PDL: wait for the predecessor kernel (the qkv-projection GEMM that
@@ -356,7 +409,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
409 scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim;
410 normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
411 cos_ptr, /*apply_norm=*/norm_w != nullptr);
359 - storeElems<scalar_t>(store_ptr + dim_base, elems);
412 + if constexpr (kFp8Idx) {
413 + // index_q is e4m3 bytes; Q/K (and in-place index_k) stay scalar_t.
414 + if (isIQ && index_q_out != nullptr) {
415 + storeElemsFp8(index_q_out +
416 + static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
417 + (slot - iq_begin) * kHeadDim + dim_base,
418 + elems);
419 + } else {
420 + storeElems<scalar_t>(store_ptr + dim_base, elems);
421 + }
422 + } else {
423 + storeElems<scalar_t>(store_ptr + dim_base, elems);
424 + }
425 }
426
427 // ── Cache inserts (sparse serving only). ───────────────────────────────
@@ -367,8 +432,11 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
432 : (isIK ? index_slot_mapping[tokenIdx] : -1);
433 if (sm >= 0) { // skip padded / unscheduled tokens
434 if (isIK) {
370 - scalar_t* dst = index_cache + sm * kHeadDim + dim_base;
371 - storeElems<scalar_t>(dst, elems);
435 + if constexpr (kFp8Idx) {
436 + storeElemsFp8(index_cache + sm * kHeadDim + dim_base, elems);
437 + } else {
438 + storeElems<scalar_t>(index_cache + sm * kHeadDim + dim_base, elems);
439 + }
440 } else if (isK || isV) {
441 // kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
442 // Paging is logical (block = sm/block_size, token = sm%block_size);
@@ -398,19 +466,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
466 // Launch wrapper
467 // ────────────────────────────────────────────────────────────────────────────
468 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
401 -void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
402 - scalar_t const* q_norm_w, scalar_t const* k_norm_w,
403 - scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
404 - scalar_t const* cos_sin_cache,
405 - int64_t const* positions, int64_t const* slot_mapping,
406 - int64_t const* index_slot_mapping, cache_t* kv_cache,
407 - scalar_t* index_cache, float const eps,
408 - int const rotary_dim, int const num_tokens,
409 - int const nq, int const nkv, int const niq,
410 - int const block_size, int64_t const kv_s_block,
411 - int64_t const kv_s_kv, int64_t const kv_s_token,
412 - int64_t const kv_s_head, bool const has_index,
413 - bool const insert_kv, cudaStream_t stream) {
469 +void launchFusedMiniMaxM3(
470 + scalar_t* qkv, scalar_t* q_out, void* index_q_out, scalar_t const* q_norm_w,
471 + scalar_t const* k_norm_w, scalar_t const* iq_norm_w,
472 + scalar_t const* ik_norm_w, scalar_t const* cos_sin_cache,
473 + int64_t const* positions, int64_t const* slot_mapping,
474 + int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache,
475 + float const eps, int const rotary_dim, int const num_tokens, int const nq,
476 + int const nkv, int const niq, int const block_size,
477 + int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
478 + int64_t const kv_s_head, bool const has_index, bool const insert_kv,
479 + bool const fp8_idx, cudaStream_t stream) {
480 + // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the
481 + // void* pointers per instantiation in the LAUNCH macro.
482 // Slot count must match the kernel's compile-time gating.
483 int const v_slots = insert_kv ? nkv : 0;
484 int const idx_slots = has_index ? niq + 1 : 0;
@@ -440,25 +508,27 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
508 config.attrs = attrs;
509 config.numAttrs = (sm_version >= 90) ? 1 : 0;
510
443 - #define LAUNCH(IS_SPARSE, INSERT) \
444 - cudaLaunchKernelEx( \
445 - &config, \
446 - fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, \
447 - IS_SPARSE, INSERT>, \
448 - qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, \
449 - cos_sin_cache, positions, slot_mapping, index_slot_mapping, kv_cache, \
450 - index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, \
451 - kv_s_block, kv_s_kv, kv_s_token, kv_s_head)
511 + #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T) \
512 + cudaLaunchKernelEx( \
513 + &config, \
514 + fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
515 + IS_SPARSE, INSERT, FP8>, \
516 + qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, k_norm_w, \
517 + iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping, \
518 + index_slot_mapping, kv_cache, reinterpret_cast<OUT_T*>(index_cache), \
519 + eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \
520 + kv_s_kv, kv_s_token, kv_s_head)
521 #else
522 // ROCm: standard kernel launch syntax (no PDL/stream serialization).
523 // clang-format off
455 - #define LAUNCH(IS_SPARSE, INSERT) \
456 - fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, \
457 - IS_SPARSE, INSERT> \
524 + #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T) \
525 + fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
526 + IS_SPARSE, INSERT, FP8> \
527 <<<grid, kBlockSize, 0, stream>>>( \
459 - qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, \
460 - ik_norm_w, cos_sin_cache, positions, slot_mapping, \
461 - index_slot_mapping, kv_cache, index_cache, eps, rotary_dim, \
528 + qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, \
529 + k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
530 + slot_mapping, index_slot_mapping, kv_cache, \
531 + reinterpret_cast<OUT_T*>(index_cache), eps, rotary_dim, \
532 num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv, \
533 kv_s_token, kv_s_head)
534 // clang-format on
@@ -466,14 +536,22 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
536
537 if (has_index) {
538 if (insert_kv) {
469 - LAUNCH(true, true); // sparse serving
539 + if (fp8_idx) {
540 + LAUNCH(true, true, true, uint8_t); // sparse serving, fp8 index outputs
541 + } else {
542 + LAUNCH(true, true, false, scalar_t); // sparse serving, bf16
543 + }
544 } else {
471 - LAUNCH(true, false); // sparse profiling
545 + if (fp8_idx) {
546 + LAUNCH(true, false, true, uint8_t); // sparse profiling, fp8 index_q
547 + } else {
548 + LAUNCH(true, false, false, scalar_t); // sparse profiling, bf16
549 + }
550 }
551 } else {
552 // Dense layer: never has an index branch and never inserts here (the
553 // generic Attention layer owns the KV insert).
476 - LAUNCH(false, false);
554 + LAUNCH(false, false, false, scalar_t);
555 }
556 #undef LAUNCH
557 }
@@ -485,8 +563,9 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
563 vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>( \
564 reinterpret_cast<st*>(qkv.data_ptr()), \
565 q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr()) : nullptr, \
488 - index_q_out.has_value() ? reinterpret_cast<st*>(index_q_out->data_ptr()) \
489 - : nullptr, \
566 + index_q_out.has_value() \
567 + ? reinterpret_cast<void*>(index_q_out->data_ptr()) \
568 + : nullptr, \
569 reinterpret_cast<st const*>(q_norm_weight.data_ptr()), \
570 reinterpret_cast<st const*>(k_norm_weight.data_ptr()), \
571 has_index ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr()) \
@@ -502,11 +581,11 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
581 : nullptr, \
582 insert_kv ? reinterpret_cast<CACHE_T*>(kv_cache->data_ptr()) : nullptr, \
583 (insert_kv && has_index) \
505 - ? reinterpret_cast<st*>(index_cache->data_ptr()) \
584 + ? reinterpret_cast<void*>(index_cache->data_ptr()) \
585 : nullptr, \
586 static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq, \
587 nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv, kv_s_token, \
509 - kv_s_head, has_index, insert_kv, stream)
588 + kv_s_head, has_index, insert_kv, fp8_idx, stream)
589
590 // ────────────────────────────────────────────────────────────────────────────
591 // Torch op wrapper
@@ -612,6 +691,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
691 torch::headeronly::ScalarType::Long &&
692 index_slot_mapping->numel() == slot_mapping->numel()),
693 "index_slot_mapping must be int64 CUDA with slot_mapping length");
694 + // Main attention KV cache: auto matches qkv, fp8 uses uint8 storage.
695 if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
696 STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
697 "auto kv_cache dtype must match qkv");
@@ -620,9 +700,13 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
700 kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte,
701 "fp8 kv_cache must use uint8 storage");
702 }
623 - STD_TORCH_CHECK(index_cache.has_value() &&
624 - index_cache->scalar_type() == qkv.scalar_type(),
625 - "insert mode requires matching index_cache");
703 + // Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3.
704 + STD_TORCH_CHECK(
705 + index_cache.has_value() &&
706 + (index_cache->scalar_type() == qkv.scalar_type() ||
707 + index_cache->scalar_type() ==
708 + torch::headeronly::ScalarType::Float8_e4m3fn),
709 + "insert mode requires index_cache matching qkv dtype or fp8 e4m3");
710 STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
711 "kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
712 "head_dim (stride(4)==1)");
@@ -652,14 +736,31 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
736 "index_q_out requires the index branch (num_index_heads > 0)");
737 STD_TORCH_CHECK(
738 index_q_out->is_cuda() && index_q_out->is_contiguous() &&
655 - index_q_out->scalar_type() == qkv.scalar_type(),
656 - "index_q_out must be a contiguous CUDA tensor matching qkv dtype");
739 + (index_q_out->scalar_type() == qkv.scalar_type() ||
740 + index_q_out->scalar_type() ==
741 + torch::headeronly::ScalarType::Float8_e4m3fn),
742 + "index_q_out must be contiguous CUDA, qkv dtype or fp8 e4m3");
743 STD_TORCH_CHECK(index_q_out->numel() ==
744 static_cast<int64_t>(num_tokens) * niq * kHeadDim,
745 "index_q_out must have num_tokens * num_index_heads * 128 "
746 "elements");
747 }
748
749 + // fp8 index path: the index-K cache and index-Q outputs are e4m3 bytes while
750 + // q/k/v + q_out stay qkv dtype. Both index outputs must agree.
751 + auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn;
752 + bool const fp8_idx =
753 + (index_cache.has_value() && index_cache->scalar_type() == kFp8) ||
754 + (index_q_out.has_value() && index_q_out->scalar_type() == kFp8);
755 + if (fp8_idx) {
756 + STD_TORCH_CHECK(
757 + !index_cache.has_value() || index_cache->scalar_type() == kFp8,
758 + "fp8 index path: index_cache must be fp8 e4m3");
759 + STD_TORCH_CHECK(
760 + !index_q_out.has_value() || index_q_out->scalar_type() == kFp8,
761 + "fp8 index path: index_q_out must be fp8 e4m3");
762 + }
763 +
764 const torch::stable::accelerator::DeviceGuard device_guard(
765 qkv.get_device_index());
766 auto stream = get_current_cuda_stream(qkv.get_device_index());
setup.py
+8
@@ -1171,7 +1171,15 @@ package_data = {
1171 "third_party/deep_gemm/include/**/*.h",
1172 "third_party/deep_gemm/include/**/*.hpp",
1173 # fmha_sm100 sparse CuTe-DSL helper kernels (vendored via cmake)
1174 + "third_party/fmha_sm100/csrc/**/*.cu",
1175 + "third_party/fmha_sm100/csrc/**/*.h",
1176 + "third_party/fmha_sm100/csrc/**/*.jinja",
1177 + "third_party/fmha_sm100/csrc/**/*.cu.jinja",
1178 "third_party/fmha_sm100/cute/**/*.cu",
1179 + "third_party/fmha_sm100/cutlass/include/**/*.h",
1180 + "third_party/fmha_sm100/cutlass/include/**/*.hpp",
1181 + "third_party/fmha_sm100/cutlass/tools/util/include/**/*.h",
1182 + "third_party/fmha_sm100/cutlass/tools/util/include/**/*.hpp",
1183 ]
1184 }
1185
tests/kernels/attention/test_minimax_m3.py
+325 -1
@@ -134,6 +134,7 @@ def _reference_index_topk(
134 topk: int,
135 init_blocks: int,
136 local_blocks: int,
137 + sm_scale: float = 1.0,
138 ) -> torch.Tensor:
139 total_q, num_idx_heads, _ = idx_q.shape
140 out = torch.full(
@@ -149,7 +150,7 @@ def _reference_index_topk(
150 num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
151 pages = block_table[req_id, :num_blocks]
152 k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1)
152 - score = torch.einsum("qhd,kd->hqk", q.float(), k.float())
153 + score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale
154
155 q_pos = prefix_len + torch.arange(q_len, device=idx_q.device)
156 k_pos = torch.arange(k.shape[0], device=idx_q.device)
@@ -244,6 +245,270 @@ def test_prefill_index_topk_correctness():
245 _assert_topk_indices_equal_unordered(actual, expected)
246
247
248 +# MSA indexer (SM100): fmha_sm100 OnlyScore for the per-block scores, then the
249 +# Triton minimax_m3_index_topk for selection (no sparse_topk_select). Uses a
250 +# deterministic construction (idx_q == 1, distinct e4m3-exact per-block values)
251 +# so scores are strictly monotonic in the block id -> exact top-k agreement.
252 +def _fmha_indexer_topk(
253 + idx_q: torch.Tensor, # [total_q, H, 128] bf16/e4m3
254 + index_cache: torch.Tensor, # [num_pages, 128, 128] bf16/e4m3
255 + block_table: torch.Tensor,
256 + q_lens: torch.Tensor,
257 + seq_lens: torch.Tensor,
258 + prefix_lens: torch.Tensor,
259 + sm_scale: float,
260 + topk: int,
261 +) -> torch.Tensor:
262 + """Replicate MiniMaxM3IndexerMSAImpl's score path (single decode/prefill side)."""
263 + from vllm.third_party.fmha_sm100.api import _fmha_sm100, _fmha_sm100_plan
264 +
265 + num_idx_heads, head_dim = idx_q.shape[1], idx_q.shape[2]
266 + nvp = [(s + 127) // 128 for s in seq_lens.tolist()]
267 + kv_indices = torch.cat([block_table[r, : nvp[r]] for r in range(len(nvp))]).to(
268 + torch.int32
269 + )
270 +
271 + qo = q_lens.cpu().to(torch.int32)
272 + kv = seq_lens.cpu().to(torch.int32)
273 + plan = _fmha_sm100_plan(
274 + qo,
275 + kv,
276 + num_idx_heads,
277 + num_kv_heads=1,
278 + qo_offset=kv - qo,
279 + page_size=128,
280 + output_maxscore=True,
281 + causal=True,
282 + num_kv_splits=1,
283 + )
284 + k_pages = index_cache.view(index_cache.shape[0], 1, 128, head_dim)
285 + _, max_score = _fmha_sm100(
286 + idx_q,
287 + k_pages,
288 + k_pages,
289 + plan,
290 + kv_indices=kv_indices,
291 + output_o=False,
292 + output_maxscore=True,
293 + sm_scale=sm_scale,
294 + )
295 +
296 + batch = q_lens.numel()
297 + cu = torch.zeros(batch + 1, dtype=torch.int32, device=idx_q.device)
298 + cu[1:] = q_lens.to(torch.int32).cumsum(0)
299 + # max_score [H, k_tiles, total_q] -> transpose to [H, total_q, k_tiles].
300 + return minimax_m3_index_topk(
301 + max_score.transpose(1, 2),
302 + cu,
303 + prefix_lens.to(torch.int32),
304 + int(q_lens.max()),
305 + topk,
306 + 0, # init_blocks
307 + 0, # local_blocks
308 + )
309 +
310 +
311 +# e4m3-exact, strictly-increasing per-block values: with idx_q == 1 (also exact)
312 +# the per-block scores are exact and distinct in BOTH bf16 and e4m3, so the fp8
313 +# score path selects the same top-k as the reference (no quantization ties).
314 +_E4M3_EXACT_VALUES = [
315 + *range(1, 17), # 1..16 (step 1)
316 + *range(18, 33, 2), # 18..32 (step 2)
317 + *range(36, 65, 4), # 36..64 (step 4)
318 + *range(72, 129, 8), # 72..128 (step 8)
319 +]
320 +
321 +
322 +@pytest.mark.skipif(
323 + not current_platform.is_device_capability_family(100),
324 + reason="fmha_sm100 indexer requires SM100 (Blackwell).",
325 +)
326 +@pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn])
327 +@pytest.mark.parametrize(
328 + ("q_lens", "prefix_lens"),
329 + [
330 + ((4, 3), (2048, 2560)), # prefill: every token sees >= 16 causal blocks
331 + ((1, 1, 1), (2048, 3000, 4096)), # decode: one query token per request
332 + ],
333 +)
334 +def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype):
335 + torch.manual_seed(0)
336 + num_idx_heads, head_dim = 4, HEAD_DIM
337 + device = "cuda"
338 +
339 + q_lens_t = torch.tensor(q_lens, device=device, dtype=torch.int32)
340 + prefix_lens_t = torch.tensor(prefix_lens, device=device, dtype=torch.int32)
341 + seq_lens = prefix_lens_t + q_lens_t
342 + batch = len(q_lens)
343 + max_blocks = (int(seq_lens.max()) + BLOCK_SIZE - 1) // BLOCK_SIZE
344 + assert max_blocks <= len(_E4M3_EXACT_VALUES)
345 + num_pages = batch * max_blocks
346 + block_table = torch.randperm(num_pages, device=device, dtype=torch.int32).reshape(
347 + batch, max_blocks
348 + )
349 +
350 + idx_q = torch.ones(
351 + int(q_lens_t.sum()), num_idx_heads, head_dim, device=device, dtype=index_dtype
352 + )
353 + index_cache = torch.empty(
354 + num_pages, BLOCK_SIZE, head_dim, device=device, dtype=index_dtype
355 + )
356 + for r in range(batch):
357 + for b in range(max_blocks):
358 + index_cache[block_table[r, b]] = float(_E4M3_EXACT_VALUES[b])
359 +
360 + sm_scale = head_dim**-0.5
361 + actual = _fmha_indexer_topk(
362 + idx_q,
363 + index_cache,
364 + block_table,
365 + q_lens_t,
366 + seq_lens,
367 + prefix_lens_t,
368 + sm_scale,
369 + TOPK,
370 + )
371 + expected = _reference_index_topk(
372 + idx_q,
373 + index_cache,
374 + block_table,
375 + q_lens_t,
376 + seq_lens,
377 + prefix_lens_t,
378 + TOPK,
379 + init_blocks=0,
380 + local_blocks=0,
381 + sm_scale=sm_scale,
382 + )
383 + _assert_topk_indices_equal_unordered(actual, expected)
384 +
385 +
386 +# Full impl-level parity: drive both MiniMaxM3IndexerMSAImpl (fmha_sm100 score +
387 +# Triton top-k) and MiniMaxM3IndexerTritonImpl through their real metadata
388 +# builders on the SAME CommonAttentionMetadata + index cache, and assert the
389 +# selected blocks agree. This exercises all the metadata the impl/kernels consume
390 +# (decode/prefill split, cu_seqlens_q rebasing, prefix_lens, kv_indices gather,
391 +# decode_pages split) -- a metadata bug on either side shifts the causal window
392 +# or the block->page mapping and breaks the comparison.
393 +@pytest.mark.skipif(
394 + not current_platform.is_device_capability_family(100),
395 + reason="fmha_sm100 indexer requires SM100 (Blackwell).",
396 +)
397 +@pytest.mark.parametrize("topk", [8, 16])
398 +def test_msa_indexer_impl_matches_triton(topk, monkeypatch):
399 + import vllm.models.minimax_m3.common.indexer as indexer_mod
400 + from tests.v1.attention.utils import (
401 + BatchSpec,
402 + create_common_attn_metadata,
403 + create_vllm_config,
404 + )
405 + from vllm.config import set_current_vllm_config
406 + from vllm.forward_context import set_forward_context
407 + from vllm.models.minimax_m3.common.indexer import (
408 + MiniMaxM3IndexerTritonImpl,
409 + MiniMaxM3IndexerTritonMetadataBuilder,
410 + )
411 + from vllm.models.minimax_m3.nvidia.indexer_msa import (
412 + MiniMaxM3IndexerMSAImpl,
413 + MiniMaxM3IndexerMSAMetadataBuilder,
414 + )
415 +
416 + torch.manual_seed(0)
417 + device = torch.device("cuda")
418 + num_idx_heads, head_dim = 4, HEAD_DIM
419 + # TP=1: avoid requiring an initialized distributed group in a unit test.
420 + monkeypatch.setattr(indexer_mod, "get_tensor_model_parallel_world_size", lambda: 1)
421 +
422 + vllm_config = create_vllm_config(
423 + block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=8192
424 + )
425 + vllm_config.model_config.hf_config.sparse_attention_config = {
426 + "sparse_num_index_heads": num_idx_heads
427 + }
428 +
429 + # Decode-first mixed batch: 2 decode reqs (q_len 1) then 2 prefill reqs. Long
430 + # prefixes so every token sees > TOPK causal blocks (non-trivial selection).
431 + batch = BatchSpec(seq_lens=[2305, 2561, 2624, 2720], query_lens=[1, 1, 64, 96])
432 + common = create_common_attn_metadata(
433 + batch, BLOCK_SIZE, device, arange_block_indices=True
434 + )
435 + num_tokens = batch.compute_num_tokens()
436 +
437 + # Deterministic index cache: distinct, monotonic per-logical-block values so
438 + # the top-k is unambiguous (both kernels pick the same blocks, no fp ties).
439 + block_table = common.block_table_tensor
440 + num_pages = int(block_table.max().item()) + 1
441 + index_cache = torch.zeros(
442 + num_pages, BLOCK_SIZE, head_dim, device=device, dtype=DTYPE
443 + )
444 + for r, seq_len in enumerate(batch.seq_lens):
445 + for b in range((seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE):
446 + index_cache[block_table[r, b]] = float(b + 1)
447 + index_q = torch.ones(
448 + num_tokens, num_idx_heads * head_dim, device=device, dtype=DTYPE
449 + )
450 +
451 + spec = MLAAttentionSpec(
452 + block_size=BLOCK_SIZE, num_kv_heads=1, head_size=head_dim, dtype=DTYPE
453 + )
454 + impl_kwargs = dict(
455 + num_kv_heads=num_idx_heads,
456 + scale=head_dim**-0.5,
457 + topk_blocks=topk,
458 + sparse_block_size=BLOCK_SIZE,
459 + num_index_heads=num_idx_heads,
460 + index_head_dim=head_dim,
461 + init_blocks=0,
462 + local_blocks=0,
463 + )
464 +
465 + with set_current_vllm_config(vllm_config):
466 + msa_impl = MiniMaxM3IndexerMSAImpl(prefix="idx_msa", **impl_kwargs)
467 + triton_impl = MiniMaxM3IndexerTritonImpl(prefix="idx_triton", **impl_kwargs)
468 + msa_builder = MiniMaxM3IndexerMSAMetadataBuilder(
469 + spec, [msa_impl.index_cache.prefix], vllm_config, device
470 + )
471 + triton_builder = MiniMaxM3IndexerTritonMetadataBuilder(
472 + spec, [triton_impl.index_cache.prefix], vllm_config, device
473 + )
474 +
475 + # Both impls score against the same index keys.
476 + msa_impl.index_cache.kv_cache = index_cache
477 + triton_impl.index_cache.kv_cache = index_cache
478 +
479 + # Exercise the shared persistent top-k buffer for BOTH impls: each must write
480 + # decode ([:, :nd]) and prefill ([:, nd:]) into its buffer and return views.
481 + # Separate buffers so the two forwards don't clobber each other.
482 + nd = sum(q for q in batch.query_lens if q <= 1)
483 + msa_impl.topk_indices_buffer = torch.full(
484 + (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
485 + )
486 + triton_impl.topk_indices_buffer = torch.full(
487 + (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
488 + )
489 +
490 + attn_metadata = {
491 + msa_impl.index_cache.prefix: msa_builder.build(0, common),
492 + triton_impl.index_cache.prefix: triton_builder.build(0, common),
493 + }
494 + with set_forward_context(attn_metadata, vllm_config):
495 + msa_decode, msa_prefill = msa_impl(index_q)
496 + tri_decode, tri_prefill = triton_impl(index_q)
497 +
498 + assert msa_decode is not None and tri_decode is not None
499 + assert msa_prefill is not None and tri_prefill is not None
500 + _assert_topk_indices_equal_unordered(msa_decode, tri_decode)
501 + _assert_topk_indices_equal_unordered(msa_prefill, tri_prefill)
502 + # decode/prefill outputs are views into each impl's persistent buffer.
503 + for impl, dec, pre in (
504 + (msa_impl, msa_decode, msa_prefill),
505 + (triton_impl, tri_decode, tri_prefill),
506 + ):
507 + buf = impl.topk_indices_buffer
508 + assert dec.data_ptr() == buf[:, :nd, :].data_ptr()
509 + assert pre.data_ptr() == buf[:, nd:, :].data_ptr()
510 +
511 +
512 @pytest.mark.parametrize(
513 ("decode_query_len", "max_decode_query_len"),
514 [
@@ -317,6 +582,65 @@ def test_decode_index_topk_correctness(
582 _assert_topk_indices_equal_unordered(actual, expected)
583
584
585 +@pytest.mark.skipif(
586 + not current_platform.is_device_capability_family(100),
587 + reason="fp8 e4m3 indexer cache is the SM100 (MSA) path.",
588 +)
589 +@pytest.mark.parametrize("num_idx_heads", [1, 4])
590 +def test_decode_index_topk_fp8(num_idx_heads: int):
591 + """The fp8 (e4m3) indexer cache feeds the Triton decode kernel on the MSA
592 + path. The kernel must score in fp32 (no scaling) so its top-k matches a
593 + reference computed from the dequantized fp8 values."""
594 + torch.manual_seed(0)
595 + topk, init_blocks, local_blocks, head_dim = 8, 0, 1, 128
596 + decode_query_len = 1
597 + active_seq_lens = torch.tensor((129, 1025, 4097), device="cuda", dtype=torch.int32)
598 + q_lens = torch.full_like(active_seq_lens, decode_query_len)
599 + prefix_lens = active_seq_lens - decode_query_len
600 + batch = active_seq_lens.numel()
601 + max_seq_len = int(active_seq_lens.max())
602 + max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
603 + num_pages = batch * max_blocks
604 + block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(
605 + batch, max_blocks
606 + )
607 + idx_q = torch.randn(
608 + batch * decode_query_len, num_idx_heads, head_dim, device="cuda"
609 + ).to(torch.float8_e4m3fn)
610 + index_kv_cache = torch.randn(num_pages, BLOCK_SIZE, head_dim, device="cuda").to(
611 + torch.float8_e4m3fn
612 + )
613 +
614 + actual = minimax_m3_index_decode(
615 + idx_q,
616 + index_kv_cache,
617 + block_table,
618 + active_seq_lens,
619 + max_seq_len=max_seq_len,
620 + topk=topk,
621 + init_blocks=init_blocks,
622 + local_blocks=local_blocks,
623 + num_kv_heads=num_idx_heads,
624 + sm_scale=head_dim**-0.5,
625 + decode_query_len=decode_query_len,
626 + )
627 + # Reference from the DEQUANTIZED fp8 values (the kernel computes the fp8 QK
628 + # in fp32, so it must match an fp32 matmul of the same e4m3 values).
629 + expected = _reference_index_topk(
630 + idx_q.float(),
631 + index_kv_cache.float(),
632 + block_table,
633 + q_lens,
634 + active_seq_lens,
635 + prefix_lens,
636 + topk,
637 + init_blocks,
638 + local_blocks,
639 + head_dim**-0.5,
640 + )
641 + _assert_topk_indices_equal_unordered(actual, expected)
642 +
643 +
644 # Sparse attention kernels.
645 def _reference_sparse_attn(
646 q: torch.Tensor,
tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
+96
@@ -278,3 +278,99 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
278 torch.testing.assert_close(
279 index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0
280 )
281 +
282 +
283 +# ── Test 3: fp8 (e4m3) index outputs ─────────────────────────────────────────
284 +# The fp8 score path stores index_q and the index-K cache as e4m3 while q/k/v +
285 +# q_out stay bf16. Asserts: (1) q/k/v/q_out are bit-identical to the bf16 run
286 +# (the index dtype must not perturb the main branch), and (2) the e4m3 index
287 +# outputs dequantize close to the bf16 reference.
288 +
289 +
290 +@pytest.mark.skipif(
291 + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9),
292 + reason="e4m3 conversion requires CUDA SM89+.",
293 +)
294 +@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
295 +@pytest.mark.parametrize("block_size", [16, 64])
296 +def test_sparse_full_fp8_index(num_tokens, block_size):
297 + torch.manual_seed(1)
298 + device, dtype, eps = "cuda", torch.bfloat16, 1e-6
299 + base, max_pos = 5_000_000.0, 4096
300 + num_heads, num_kv_heads, num_idx_heads = 16, 4, 4
301 +
302 + q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
303 + k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
304 + iq_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
305 + ik_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
306 + cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
307 + positions = torch.randint(
308 + 0, max_pos, (num_tokens,), dtype=torch.int64, device=device
309 + )
310 +
311 + qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
312 + iqsz, iksz = num_idx_heads * HEAD_DIM, HEAD_DIM
313 + qkv0 = torch.randn(
314 + num_tokens, qsz + 2 * kvsz + iqsz + iksz, dtype=dtype, device=device
315 + )
316 +
317 + num_blocks = (num_tokens + block_size - 1) // block_size + 1
318 + slot_mapping = torch.randperm(
319 + num_blocks * block_size, dtype=torch.int64, device=device
320 + )[:num_tokens]
321 + index_slot_mapping = torch.roll(slot_mapping, shifts=1)
322 +
323 + def run(index_dtype):
324 + qkv = qkv0.clone()
325 + kv_cache = torch.zeros(
326 + num_blocks,
327 + 2,
328 + block_size,
329 + num_kv_heads,
330 + HEAD_DIM,
331 + dtype=dtype,
332 + device=device,
333 + )
334 + index_cache = torch.zeros(
335 + num_blocks, block_size, HEAD_DIM, dtype=index_dtype, device=device
336 + )
337 + q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device)
338 + index_q = torch.empty(num_tokens, iqsz, dtype=index_dtype, device=device)
339 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
340 + qkv,
341 + q_w,
342 + k_w,
343 + cos_sin,
344 + positions,
345 + num_heads,
346 + num_kv_heads,
347 + ROTARY_DIM,
348 + eps,
349 + iq_w,
350 + ik_w,
351 + num_idx_heads,
352 + slot_mapping,
353 + index_slot_mapping,
354 + kv_cache,
355 + index_cache,
356 + block_size,
357 + q_out,
358 + index_q,
359 + )
360 + return qkv, kv_cache, index_cache, q_out, index_q
361 +
362 + qkv_bf, kvc_bf, idxc_bf, qo_bf, iq_bf = run(torch.bfloat16)
363 + qkv_fp, kvc_fp, idxc_fp, qo_fp, iq_fp = run(torch.float8_e4m3fn)
364 +
365 + assert iq_fp.dtype == torch.float8_e4m3fn
366 + assert idxc_fp.dtype == torch.float8_e4m3fn
367 +
368 + # (1) The main branch (q/k/v in qkv, q_out, kv cache) must be bit-identical:
369 + # the index output dtype must not perturb anything else.
370 + torch.testing.assert_close(qo_fp, qo_bf, rtol=0, atol=0)
371 + torch.testing.assert_close(qkv_fp, qkv_bf, rtol=0, atol=0)
372 + torch.testing.assert_close(kvc_fp, kvc_bf, rtol=0, atol=0)
373 +
374 + # (2) Dequantized e4m3 index outputs match the bf16 reference within fp8 ulp.
375 + torch.testing.assert_close(iq_fp.float(), iq_bf.float(), rtol=0.13, atol=0.05)
376 + torch.testing.assert_close(idxc_fp.float(), idxc_bf.float(), rtol=0.13, atol=0.05)
vllm/envs.py
+3 -4
@@ -1580,10 +1580,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
1580 os.getenv("VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS", "1")
1581 ),
1582 # Enforce function parameter schemas in structural-tag based tool calling.
1583 - "VLLM_ENFORCE_STRICT_TOOL_CALLING": lambda: os.getenv(
1584 - "VLLM_ENFORCE_STRICT_TOOL_CALLING", "True"
1585 - ).lower()
1586 - in ("true", "1"),
1583 + "VLLM_ENFORCE_STRICT_TOOL_CALLING": lambda: (
1584 + os.getenv("VLLM_ENFORCE_STRICT_TOOL_CALLING", "True").lower() in ("true", "1")
1585 + ),
1586 # Control the max chunk bytes (in MB) for the rpc message queue.
1587 # Object larger than this threshold will be broadcast to worker
1588 # processes via zmq.
vllm/models/minimax_m3/amd/model.py
+28 -3
@@ -457,6 +457,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
457 quant_config: QuantizationConfig | None = None,
458 prefix: str = "",
459 cache_config: CacheConfig | None = None,
460 + topk_indices_buffer: torch.Tensor | None = None,
461 ) -> None:
462 super().__init__()
463 self.hidden_size = config.hidden_size
@@ -535,6 +536,9 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
536 self.kv_cache_dtype, vllm_config.model_config
537 )
538
539 + # Shared top-k buffer: the indexer writes the selected blocks into it and
540 + # the attend impl reads them back (no Python value crosses the break).
541 + self.topk_indices_buffer = topk_indices_buffer
542 self.attn_backend = MiniMaxM3SparseBackend
543 # Indexer and main attention are separate impls. On ROCm the SM100 gate
544 # is always False, so both pick Triton and the index cache stays bf16.
@@ -565,6 +569,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
569 local_blocks=sparse_cfg.get("sparse_local_block", 0),
570 score_type=sparse_cfg.get("sparse_score_type", "max"),
571 cache_config=cache_config,
572 + topk_indices_buffer=topk_indices_buffer,
573 )
574
575 # Register the main K/V cache so the KV-cache manager allocates it.
@@ -657,9 +662,10 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
662 output: torch.Tensor,
663 ) -> torch.Tensor:
664 # Single eager break around both: their split-K kernels read per-request
660 - # metadata and can't be captured into a cudagraph.
661 - topk_idx = self.indexer(index_query)
662 - return self.impl.forward(self, query, self.kv_cache, topk_idx, output)
665 + # metadata and can't be captured into a cudagraph. The indexer writes its
666 + # top-k into the shared ``topk_indices_buffer``; the attend reads it back.
667 + self.indexer(index_query)
668 + return self.impl.forward(self, query, self.kv_cache, output)
669
670
671 class MiniMaxM3DecoderLayer(nn.Module):
@@ -671,6 +677,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
677 quant_config: QuantizationConfig | None = None,
678 force_sparse_attn: bool = False,
679 force_moe: bool = False,
680 + topk_indices_buffer: torch.Tensor | None = None,
681 ) -> None:
682 super().__init__()
683 self.hidden_size = config.hidden_size
@@ -690,6 +697,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
697 quant_config=quant_config,
698 prefix=f"{prefix}.self_attn",
699 cache_config=cache_config,
700 + topk_indices_buffer=topk_indices_buffer,
701 )
702 else:
703 self.self_attn = MiniMaxM3Attention(
@@ -771,6 +779,22 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
779 prefix=f"{prefix}.embed_tokens",
780 )
781
782 + # Reserved top-k indices buffer shared by all sparse-attention indexer
783 + # layers (mirrors DeepseekV4); the indexer writes its per-head decode/
784 + # prefill block selection into it, the attend reads it back.
785 + sparse_cfg = getattr(config, "sparse_attention_config", None)
786 + if sparse_cfg is not None:
787 + tp_size = get_tensor_model_parallel_world_size()
788 + num_index_heads = max(1, sparse_cfg["sparse_num_index_heads"] // tp_size)
789 + self.topk_indices_buffer = torch.empty(
790 + num_index_heads,
791 + vllm_config.scheduler_config.max_num_batched_tokens,
792 + sparse_cfg["sparse_topk_blocks"],
793 + dtype=torch.int32,
794 + )
795 + else:
796 + self.topk_indices_buffer = None
797 +
798 self.start_layer, self.end_layer, self.layers = make_layers(
799 config.num_hidden_layers,
800 lambda prefix: MiniMaxM3DecoderLayer(
@@ -778,6 +802,7 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
802 prefix,
803 cache_config=cache_config,
804 quant_config=quant_config,
805 + topk_indices_buffer=self.topk_indices_buffer,
806 ),
807 prefix=f"{prefix}.layers",
808 )
vllm/models/minimax_m3/common/indexer.py
+60 -8
@@ -25,12 +25,14 @@ from vllm.config.attention import IndexerKVDType
25 from vllm.config.cache import CacheDType
26 from vllm.distributed import get_tensor_model_parallel_world_size
27 from vllm.forward_context import get_forward_context
28 +from vllm.logger import init_logger
29 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
30 from vllm.models.minimax_m3.common.ops.index_topk import (
31 minimax_m3_index_decode,
32 minimax_m3_index_score,
33 minimax_m3_index_topk,
34 )
35 +from vllm.platforms import current_platform
36 from vllm.v1.attention.backend import (
37 AttentionBackend,
38 AttentionCGSupport,
@@ -46,6 +48,8 @@ from vllm.v1.kv_cache_interface import (
48 MLAAttentionSpec,
49 )
50
51 +logger = init_logger(__name__)
52 +
53
54 class MiniMaxM3IndexerBackend(AttentionBackend):
55 """Indexer side-cache backend (key-only)."""
@@ -120,16 +124,20 @@ class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase):
124 backend_cls: type[AttentionBackend] = MiniMaxM3IndexerBackend,
125 ) -> None:
126 super().__init__()
123 - if indexer_kv_dtype != "bf16":
127 + if indexer_kv_dtype in ("fp8", "fp8_e4m3"):
128 + cache_dtype = torch.float8_e4m3fn
129 + elif indexer_kv_dtype == "bf16":
130 + cache_dtype = torch.bfloat16
131 + else:
132 raise NotImplementedError(
125 - f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported yet "
126 - "for the MiniMax M3 indexer cache (only 'bf16')."
133 + f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the "
134 + "MiniMax M3 indexer cache (only 'bf16' or 'fp8'/'fp8_e4m3')."
135 )
136 self.kv_cache = torch.tensor([])
137 self.head_dim = head_dim
138 self.indexer_kv_dtype = indexer_kv_dtype
131 - # Storage dtype for the side cache (bf16 today; quantized layouts later).
132 - self.dtype = torch.bfloat16
139 + # Side-cache storage dtype: bf16, or e4m3 for the fp8 score path.
140 + self.dtype = cache_dtype
141 self.prefix = prefix
142 self.cache_config = cache_config
143 # Impl-chosen backend -> each impl gets its own builder (get_attn_backend).
@@ -344,6 +352,7 @@ class MiniMaxM3IndexerImpl(nn.Module):
352 score_type: str = "max",
353 cache_config: CacheConfig | None = None,
354 indexer_kv_dtype: IndexerKVDType = "bf16",
355 + topk_indices_buffer: torch.Tensor | None = None,
356 ) -> None:
357 super().__init__()
358 self.num_kv_heads = num_kv_heads
@@ -356,6 +365,9 @@ class MiniMaxM3IndexerImpl(nn.Module):
365 self.num_index_heads = num_index_heads
366 self.index_head_dim = index_head_dim
367 self.indexer_kv_dtype = indexer_kv_dtype
368 + # Shared, stable-address top-k output buffer (set by the model for the
369 + # cudagraph-safe MSA impl); None -> impl allocates fresh (eager).
370 + self.topk_indices_buffer = topk_indices_buffer
371 # Owns the side cache (registers itself in the static forward context).
372 self.index_cache = MiniMaxM3IndexerCache(
373 head_dim=index_head_dim,
@@ -392,6 +404,10 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
404 )
405 kv = self.index_cache.kv_cache
406
407 + # Both sides write into the single shared persistent topk_indices_buffer
408 + # (decode at [:, :nd], prefill at [:, nd:]) and return views into it; the
409 + # kernels' out= writes out[:, :total_q]. None -> allocate fresh.
410 + buf = self.topk_indices_buffer
411 decode_topk: torch.Tensor | None = None
412 prefill_topk: torch.Tensor | None = None
413 if index_md.num_decodes > 0:
@@ -409,6 +425,7 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
425 self.num_kv_heads,
426 d.decode_query_len,
427 d.max_decode_query_len,
428 + out=buf,
429 )
430 if index_md.num_prefills > 0:
431 p = index_md.prefill
@@ -432,29 +449,61 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
449 self.topk_blocks,
450 self.init_blocks,
451 self.local_blocks,
452 + out=buf[:, nd:, :] if buf is not None else None,
453 )
454 return decode_topk, prefill_topk
455
456
457 def select_indexer_impl_cls(
458 *,
459 + topk_blocks: int,
460 indexer_kv_dtype: IndexerKVDType = "bf16",
461 ) -> type[MiniMaxM3IndexerImpl]:
443 - """Pick the indexer impl off the index-cache dtype.
462 + """Pick the indexer impl off the platform, top-k count, and cache dtype.
463
445 - The SM100 MSA indexer score path is disabled for now; use the local Triton
446 - indexer. If re-enabled, add a NVIDIA-specific ``MiniMaxM3IndexerImpl`` here.
464 + On Blackwell (SM100) with ``topk_blocks`` in ``(4, 8, 16, 32)`` (matching the
465 + main MSA attend), the fmha_sm100 score path + Triton top-k is used for both
466 + bf16 and fp8 index caches. Everything else falls back to the Triton indexer
467 + (bf16 only).
468 """
469 if indexer_kv_dtype in ("mxfp4", "nvfp4"):
470 raise NotImplementedError(
471 f"indexer_kv_dtype={indexer_kv_dtype!r} needs the (not-yet-added) "
472 "CuteDSL indexer impl."
473 )
474 + is_sm100 = (
475 + current_platform.is_cuda() and current_platform.is_device_capability_family(100)
476 + )
477 + use_msa = (
478 + is_sm100
479 + and topk_blocks in (4, 8, 16, 32)
480 + and indexer_kv_dtype in ("bf16", "fp8", "fp8_e4m3")
481 + )
482 + if use_msa:
483 + # Lazy import so AMD / non-SM100 never import fmha_sm100.
484 + from vllm.models.minimax_m3.nvidia.indexer_msa import (
485 + MiniMaxM3IndexerMSAImpl,
486 + )
487 +
488 + logger.info_once(
489 + "MiniMax M3 indexer: selected MSA (fmha_sm100 score + Triton top-k) "
490 + "[topk_blocks=%d, indexer_kv_dtype=%s]",
491 + topk_blocks,
492 + indexer_kv_dtype,
493 + )
494 + return MiniMaxM3IndexerMSAImpl
495 if indexer_kv_dtype != "bf16":
496 raise NotImplementedError(
497 f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the "
498 "Triton indexer impl."
499 )
500 + logger.info_once(
501 + "MiniMax M3 indexer: selected Triton (no fmha_sm100) "
502 + "[topk_blocks=%d, indexer_kv_dtype=%s, sm100=%s]",
503 + topk_blocks,
504 + indexer_kv_dtype,
505 + is_sm100,
506 + )
507 return MiniMaxM3IndexerTritonImpl
508
509
@@ -480,9 +529,11 @@ class MiniMaxM3Indexer(nn.Module):
529 score_type: str = "max",
530 cache_config: CacheConfig | None = None,
531 indexer_kv_dtype: IndexerKVDType = "bf16",
532 + topk_indices_buffer: torch.Tensor | None = None,
533 ) -> None:
534 super().__init__()
535 impl_cls = select_indexer_impl_cls(
536 + topk_blocks=topk_blocks,
537 indexer_kv_dtype=indexer_kv_dtype,
538 )
539 self.impl = impl_cls(
@@ -498,6 +549,7 @@ class MiniMaxM3Indexer(nn.Module):
549 score_type=score_type,
550 cache_config=cache_config,
551 indexer_kv_dtype=indexer_kv_dtype,
552 + topk_indices_buffer=topk_indices_buffer,
553 )
554
555 @property
vllm/models/minimax_m3/common/ops/index_topk.py
+30 -12
@@ -373,7 +373,10 @@ def _decode_index_score_kernel(
373 + off_k[:, None] * stride_ik_pos
374 + off_d * stride_ik_d,
375 ) # [N,D]
376 - kq = tl.dot(k, q) # [N,HQ]
376 + # fp32 accumulation is required for the fp8 (e4m3) index cache: q/k are
377 + # loaded in their stored dtype (bf16 or e4m3) and the MMA accumulates in
378 + # fp32 so the per-block max score is exact for the fp8 indexer too.
379 + kq = tl.dot(k, q, out_dtype=tl.float32) # [N,HQ]
380 kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf"))
381 score = tl.max(kq, axis=0) # [HQ]
382 is_visible_block = blk < num_blocks_q
@@ -709,16 +712,25 @@ def minimax_m3_index_topk(
712 topk: int,
713 init_blocks: int,
714 local_blocks: int,
715 + out: torch.Tensor | None = None,
716 ) -> torch.Tensor:
713 - """Select index top-k from a precomputed score tensor."""
717 + """Select index top-k from a precomputed score tensor.
718 +
719 + When ``out`` is provided (a ``[num_idx_heads, >=total_q, topk]`` buffer), the
720 + result is written into ``out[:, :total_q, :]`` instead of a fresh tensor --
721 + used to keep the top-k output at a stable address for cudagraph capture.
722 + """
723 num_idx_heads = score.shape[0]
724 batch = cu_seqlens_q.shape[0] - 1
725 total_q = score.shape[1]
717 - topk_idx = torch.empty(
718 - (num_idx_heads, total_q, topk),
719 - dtype=torch.int32,
720 - device=score.device,
721 - )
726 + if out is not None:
727 + topk_idx = out[:, :total_q, :]
728 + else:
729 + topk_idx = torch.empty(
730 + (num_idx_heads, total_q, topk),
731 + dtype=torch.int32,
732 + device=score.device,
733 + )
734 # block_size_q == 1 -> query blocks coincide with query tokens.
735 grid_topk = (max_query_len, batch, num_idx_heads)
736 _topk_index_kernel[grid_topk](
@@ -757,10 +769,13 @@ def minimax_m3_index_decode(
769 num_kv_heads: int,
770 decode_query_len: int,
771 max_decode_query_len: int,
772 + out: torch.Tensor | None = None,
773 ) -> torch.Tensor:
774 """Decode index block-score + top-k, both split-K (cudagraph-safe).
775
776 Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad).
777 + When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into
778 + ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating.
779 """
780 total_q, num_idx_heads, head_dim = idx_q.shape
781 assert num_idx_heads == num_kv_heads, (
@@ -834,11 +849,14 @@ def minimax_m3_index_decode(
849 **score_kwargs,
850 )
851
837 - topk_idx = torch.empty(
838 - (num_idx_heads, total_q, topk),
839 - dtype=torch.int32,
840 - device=idx_q.device,
841 - )
852 + if out is not None:
853 + topk_idx = out[:, :total_q, :]
854 + else:
855 + topk_idx = torch.empty(
856 + (num_idx_heads, total_q, topk),
857 + dtype=torch.int32,
858 + device=idx_q.device,
859 + )
860 # Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts
861 # pow2(num_topk_chunks * pow2(topk)) candidates.
862 TOPK_TARGET_GRID = 64
vllm/models/minimax_m3/common/sparse_attention.py
+22 -15
@@ -2,10 +2,11 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Main block-sparse GQA attention for MiniMax M3 sparse layers.
4
5 -The lightning indexer (``indexer.py``) selects the top-k KV blocks; this module
6 -holds the main attention that attends only to those blocks: the paged K/V cache
7 -backend, its metadata + builder, and the impl that consumes the indexer's
8 -``topk_idx``. The Triton attend kernel lives here; the SM100 (MSA)
5 +The lightning indexer (``indexer.py``) selects the top-k KV blocks (written into
6 +the shared ``layer.topk_indices_buffer``); this module holds the main attention
7 +that attends only to those blocks: the paged K/V cache backend, its metadata +
8 +builder, and the impl that reads the indexer's top-k from that buffer. The Triton
9 +attend kernel lives here; the SM100 (MSA)
10 ``build_k2q_csr`` + ``sparse_atten_func`` attend lives in
11 ``nvidia/sparse_attention_msa.py``.
12
@@ -272,9 +273,10 @@ class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
273 """Abstract base for block-sparse GQA over the indexer-selected blocks.
274
275 Inherits ``AttentionImplBase`` for a custom forward signature (the layer
275 - pre-inserts K/V and runs the indexer, so forward takes the queries +
276 - ``topk_idx``). The Triton and MSA subclasses each own a full ``forward`` --
277 - no shared forward code.
276 + pre-inserts K/V and runs the indexer, which writes the selected blocks into
277 + the shared ``layer.topk_indices_buffer``; the attend reads them back from
278 + there). The Triton and MSA subclasses each own a full ``forward`` -- no
279 + shared forward code.
280 """
281
282 def __init__(
@@ -311,10 +313,14 @@ class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
313 layer: AttentionLayer,
314 query: torch.Tensor,
315 kv_cache: torch.Tensor,
314 - topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
316 output: torch.Tensor,
317 ) -> torch.Tensor:
317 - """Attend the queries to the indexer-selected blocks. Per kernel."""
318 + """Attend the queries to the indexer-selected blocks. Per kernel.
319 +
320 + The indexer has already written the top-k block ids into
321 + ``layer.topk_indices_buffer`` (decode at ``[:, :nd]``, prefill at
322 + ``[:, nd:num_tokens]``); the attend reads them from there.
323 + """
324 raise NotImplementedError
325
326
@@ -326,7 +332,6 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
332 layer: AttentionLayer,
333 query: torch.Tensor,
334 kv_cache: torch.Tensor,
329 - topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
335 output: torch.Tensor,
336 ) -> torch.Tensor:
337 attn_metadata = get_forward_context().attn_metadata
@@ -334,10 +339,12 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
339 return output # profiling run; caches unbound
340 main_md = attn_metadata[layer.layer_name] # type: ignore[attr-defined]
341 assert isinstance(main_md, MiniMaxM3SparseMetadata)
337 - decode_topk, prefill_topk = topk_idx
342
343 nd = main_md.num_decode_tokens
344 num_tokens = main_md.num_actual_tokens
345 + # Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:].
346 + topk = layer.topk_indices_buffer # type: ignore[attr-defined]
347 + assert topk is not None
348 hd = self.head_size
349 q = query[:num_tokens].view(-1, self.num_heads, hd)
350 out = output[:num_tokens].view(-1, self.num_heads, hd)
@@ -348,11 +355,11 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
355 # Decode [:nd]: split-K over the selected blocks (request-major chunks).
356 if main_md.num_decodes > 0:
357 d = main_md.decode
351 - assert d is not None and decode_topk is not None
358 + assert d is not None
359 minimax_m3_sparse_attn_decode(
360 q[:nd],
361 kv_cache,
355 - decode_topk,
362 + topk[:, :nd, :],
363 d.block_table,
364 d.seq_lens,
365 self.num_kv_heads,
@@ -364,11 +371,11 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
371 # Prefill [nd:]: cu_seqlens_q already rebased to 0.
372 if main_md.num_prefills > 0:
373 p = main_md.prefill
367 - assert p is not None and prefill_topk is not None
374 + assert p is not None
375 minimax_m3_sparse_attn(
376 q[nd:],
377 kv_cache,
371 - prefill_topk,
378 + topk[:, nd:num_tokens, :],
379 p.block_table,
380 p.cu_seqlens_q,
381 p.seq_lens,
vllm/models/minimax_m3/nvidia/indexer_msa.py new
+251
@@ -0,0 +1,251 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MSA (SM100/Blackwell) indexer impl for MiniMax M3.
4 +
5 +Prefill scores with ``fmha_sm100``'s score-only (``OnlyScore``) path then selects
6 +top-k blocks with the Triton ``minimax_m3_index_topk`` kernel -- fmha is much
7 +faster than Triton for the wide prefill score (benchmarked ~3-5x).
8 +
9 +Decode uses the Triton fused ``minimax_m3_index_decode`` (the same kernel the
10 +Triton indexer impl uses): for q_len==1 it is a purpose-built vector x matrix
11 +score (no wasted tensor-core tiles) with a 256-way split-K and a fused split-K
12 +top-k, which beats fmha's OnlyScore (wasted MMA on a single query, 64-split cap)
13 +by ~1.1-3.7x. It is cudagraph-safe by construction (shape-constant split grids)
14 +and writes the shared ``topk_indices_buffer`` via ``out=``.
15 +
16 +``fmha_sm100`` imports are function-local so this module is import-safe on
17 +AMD / non-SM100.
18 +"""
19 +
20 +from dataclasses import dataclass
21 +from typing import ClassVar
22 +
23 +import torch
24 +
25 +from vllm.forward_context import get_forward_context
26 +from vllm.models.minimax_m3.common.indexer import (
27 + MiniMaxM3IndexerBackend,
28 + MiniMaxM3IndexerDecodeMetadata,
29 + MiniMaxM3IndexerImpl,
30 + MiniMaxM3IndexerMetadata,
31 + MiniMaxM3IndexerMetadataBuilder,
32 +)
33 +from vllm.models.minimax_m3.common.ops.index_topk import (
34 + minimax_m3_index_decode,
35 + minimax_m3_index_topk,
36 +)
37 +from vllm.v1.attention.backend import (
38 + AttentionBackend,
39 + AttentionCGSupport,
40 + CommonAttentionMetadata,
41 +)
42 +from vllm.v1.attention.backends.utils import split_decodes_and_prefills
43 +
44 +# Page size == sparse block size == index-K block; fmha tile id == M3 block id.
45 +PAGE_SIZE = 128
46 +
47 +
48 +class MiniMaxM3IndexerMSABackend(MiniMaxM3IndexerBackend):
49 + """Indexer side-cache backend selecting the MSA builder."""
50 +
51 + @staticmethod
52 + def get_builder_cls() -> type["MiniMaxM3IndexerMSAMetadataBuilder"]:
53 + return MiniMaxM3IndexerMSAMetadataBuilder
54 +
55 +
56 +@dataclass
57 +class MiniMaxM3IndexerMSAPrefillMetadata:
58 + """fmha score plan + Triton top-k inputs for the prefill side (eager)."""
59 +
60 + plan: dict # fmha_sm100 PlanInfo
61 + cu_seqlens_q: torch.Tensor # [num_prefills + 1] int32, rebased to 0
62 + prefix_lens: torch.Tensor # [num_prefills] int32, context tokens
63 + max_query_len: int
64 + page_table: torch.Tensor # flat physical page indices for the prefill side
65 +
66 +
67 +@dataclass
68 +class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata):
69 + """Decode reuses the inherited base ``decode`` field (the Triton decode
70 + metadata); ``prefill_msa`` carries the fmha score plan for the prefill side
71 + (the base ``prefill`` field is unused on this path)."""
72 +
73 + prefill_msa: MiniMaxM3IndexerMSAPrefillMetadata | None = None
74 +
75 +
76 +class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
77 + """Decode metadata is the cudagraph-safe Triton decode metadata; the prefill
78 + fmha plan is built eagerly (prefill batches are not captured)."""
79 +
80 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
81 +
82 + def build(
83 + self,
84 + common_prefix_len: int,
85 + common_attn_metadata: CommonAttentionMetadata,
86 + fast_build: bool = False,
87 + ) -> MiniMaxM3IndexerMSAMetadata:
88 + num_reqs = common_attn_metadata.num_reqs
89 + num_tokens = common_attn_metadata.num_actual_tokens
90 + seq_lens = common_attn_metadata.seq_lens
91 + block_table = common_attn_metadata.block_table_tensor
92 + query_start_loc = common_attn_metadata.query_start_loc
93 +
94 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
95 + split_decodes_and_prefills(
96 + common_attn_metadata,
97 + decode_threshold=self.reorder_batch_threshold,
98 + require_uniform=True,
99 + )
100 + )
101 + assert num_decodes + num_prefills == num_reqs
102 + assert num_decode_tokens + num_prefill_tokens == num_tokens
103 +
104 + # Context (prefix) lengths into the stable cudagraph buffer.
105 + context_lens = self.context_len_buffer[:num_reqs]
106 + context_lens.copy_(
107 + common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
108 + )
109 +
110 + decode: MiniMaxM3IndexerDecodeMetadata | None = None
111 + if num_decodes > 0:
112 + qsl_cpu = common_attn_metadata.query_start_loc_cpu
113 + query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes]
114 + decode_query_len = int(query_lens_cpu[0].item())
115 + assert decode_query_len > 0
116 + assert torch.all(
117 + (query_lens_cpu == decode_query_len) | (query_lens_cpu == 0)
118 + )
119 + decode = MiniMaxM3IndexerDecodeMetadata(
120 + seq_lens=seq_lens[:num_decodes],
121 + block_table=block_table[:num_decodes],
122 + max_seq_len=common_attn_metadata.max_seq_len,
123 + decode_query_len=decode_query_len,
124 + max_decode_query_len=self.max_decode_query_len,
125 + )
126 +
127 + prefill: MiniMaxM3IndexerMSAPrefillMetadata | None = None
128 + if num_prefills > 0:
129 + # Prefill is eager (not captured); the host lengths it needs (and the
130 + # _fmha_sm100_plan .tolist() inside) make the D->H sync acceptable.
131 + from vllm.third_party.fmha_sm100.api import _fmha_sm100_plan
132 +
133 + lo, hi = num_decodes, num_reqs
134 + qsl_cpu = common_attn_metadata.query_start_loc_cpu[: num_reqs + 1]
135 + qo_lens_cpu = (qsl_cpu[1:] - qsl_cpu[:-1]).to(torch.int32)
136 + kv_lens_cpu = seq_lens[:num_reqs].cpu().to(torch.int32)
137 + nvp = (kv_lens_cpu + PAGE_SIZE - 1) // PAGE_SIZE
138 + side_qo = qo_lens_cpu[lo:hi]
139 + side_kv = kv_lens_cpu[lo:hi]
140 + plan = _fmha_sm100_plan(
141 + side_qo,
142 + side_kv,
143 + self.num_index_heads,
144 + num_kv_heads=1,
145 + qo_offset=side_kv - side_qo, # bottom-right causal
146 + page_size=PAGE_SIZE,
147 + output_maxscore=True,
148 + causal=True,
149 + num_kv_splits=1,
150 + )
151 + cols = torch.arange(block_table.shape[1], device=block_table.device)
152 + valid = cols[None, :] < nvp[lo:hi].to(block_table.device)[:, None]
153 + prefill = MiniMaxM3IndexerMSAPrefillMetadata(
154 + plan=plan,
155 + cu_seqlens_q=(query_start_loc[lo : hi + 1] - query_start_loc[lo]).to(
156 + torch.int32
157 + ),
158 + prefix_lens=context_lens[lo:hi],
159 + max_query_len=int(side_qo.max()),
160 + page_table=block_table[lo:hi][valid].to(torch.int32),
161 + )
162 +
163 + return MiniMaxM3IndexerMSAMetadata(
164 + seq_lens=seq_lens,
165 + max_seq_len=common_attn_metadata.max_seq_len,
166 + slot_mapping=common_attn_metadata.slot_mapping,
167 + num_actual_tokens=num_tokens,
168 + num_decodes=num_decodes,
169 + num_decode_tokens=num_decode_tokens,
170 + num_prefills=num_prefills,
171 + num_prefill_tokens=num_prefill_tokens,
172 + decode=decode,
173 + prefill_msa=prefill,
174 + )
175 +
176 +
177 +class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl):
178 + """Decode: Triton fused score+top-k. Prefill: fmha_sm100 OnlyScore + top-k."""
179 +
180 + indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerMSABackend
181 +
182 + def forward(
183 + self,
184 + index_query: torch.Tensor,
185 + ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
186 + attn_metadata = get_forward_context().attn_metadata
187 + if not isinstance(attn_metadata, dict):
188 + return None, None # profiling run; caches unbound
189 + md = attn_metadata[self.index_cache.prefix]
190 + assert isinstance(md, MiniMaxM3IndexerMSAMetadata)
191 +
192 + num_tokens = md.num_actual_tokens
193 + nd = md.num_decode_tokens
194 + index_q = index_query[:num_tokens].view(
195 + -1, self.num_index_heads, self.index_head_dim
196 + )
197 + kv = self.index_cache.kv_cache
198 + # Both sides write into the single shared persistent topk_indices_buffer:
199 + # decode at [:, :nd], prefill at [:, nd:] (each kernel writes [:, :total_q]).
200 + buf = self.topk_indices_buffer
201 +
202 + decode_topk: torch.Tensor | None = None
203 + if md.decode is not None:
204 + d = md.decode
205 + decode_topk = minimax_m3_index_decode(
206 + index_q[:nd],
207 + kv,
208 + d.block_table,
209 + d.seq_lens,
210 + d.max_seq_len,
211 + self.topk_blocks,
212 + self.init_blocks,
213 + self.local_blocks,
214 + self.num_kv_heads,
215 + d.decode_query_len,
216 + d.max_decode_query_len,
217 + out=buf,
218 + )
219 +
220 + prefill_topk: torch.Tensor | None = None
221 + if md.prefill_msa is not None:
222 + from vllm.third_party.fmha_sm100.api import _fmha_sm100
223 +
224 + p = md.prefill_msa
225 + # Index-K cache (num_blocks, 128, D) -> paged MQA (num_blocks,1,128,D).
226 + k_pages = kv.view(kv.shape[0], 1, PAGE_SIZE, self.index_head_dim)
227 + _, max_score = _fmha_sm100(
228 + index_q[nd:],
229 + k_pages,
230 + k_pages, # V placeholder; not read in OnlyScore
231 + p.plan,
232 + kv_indices=p.page_table,
233 + output_o=False,
234 + output_maxscore=True,
235 + sm_scale=self.scale,
236 + )
237 + # Triton top-k wants [num_index_heads, num_tokens, max_block]; the
238 + # transpose is a strided view (the kernel reads via strides).
239 + out = buf[:, nd:, :] if buf is not None else None
240 + prefill_topk = minimax_m3_index_topk(
241 + max_score.transpose(1, 2),
242 + p.cu_seqlens_q,
243 + p.prefix_lens,
244 + p.max_query_len,
245 + self.topk_blocks,
246 + self.init_blocks,
247 + self.local_blocks,
248 + out=out,
249 + )
250 +
251 + return decode_topk, prefill_topk
vllm/models/minimax_m3/nvidia/model.py
+40 -4
@@ -402,6 +402,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
402 quant_config: QuantizationConfig | None = None,
403 prefix: str = "",
404 cache_config: CacheConfig | None = None,
405 + topk_indices_buffer: torch.Tensor | None = None,
406 ) -> None:
407 super().__init__()
408 self.hidden_size = config.hidden_size
@@ -489,6 +490,10 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
490 # cache (--attention-config '{"indexer_kv_dtype": ...}').
491 self.indexer_kv_dtype = vllm_config.attention_config.indexer_kv_dtype
492
493 + # Shared top-k buffer: the indexer writes the selected blocks into it and
494 + # the attend impl reads them back (so nothing crosses the eager break as a
495 + # Python value, which would freeze at capture).
496 + self.topk_indices_buffer = topk_indices_buffer
497 self.attn_backend = MiniMaxM3SparseBackend
498 # Indexer (top-k selection) and main attention are separate impls, each
499 # picking Triton vs MSA off its cache dtype. impl is AttentionImplBase
@@ -519,6 +524,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
524 score_type=sparse_cfg.get("sparse_score_type", "max"),
525 cache_config=cache_config,
526 indexer_kv_dtype=self.indexer_kv_dtype,
527 + topk_indices_buffer=topk_indices_buffer,
528 )
529
530 # Register the main K/V cache so the KV-cache manager allocates it.
@@ -576,7 +582,12 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
582 main_slot_mapping = fwd_slot_mapping[self.layer_name]
583 index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
584 q = qkv.new_empty((num_tokens, self.q_size))
579 - index_q = qkv.new_empty((num_tokens, self.index_q_size))
585 + # index_q matches the index-K cache dtype (e4m3 for the fp8 score path);
586 + # the fused kernel emits fp8 directly when this buffer is e4m3.
587 + index_q = qkv.new_empty(
588 + (num_tokens, self.index_q_size),
589 + dtype=self.indexer.index_cache.dtype,
590 + )
591 ops.fused_minimax_m3_qknorm_rope_kv_insert(
592 qkv,
593 self.q_norm.weight,
@@ -613,9 +624,10 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
624 output: torch.Tensor,
625 ) -> torch.Tensor:
626 # Single eager break around both: their split-K kernels read per-request
616 - # metadata and can't be captured into a cudagraph.
617 - topk_idx = self.indexer(index_query)
618 - return self.impl.forward(self, query, self.kv_cache, topk_idx, output)
627 + # metadata and can't be captured into a cudagraph. The indexer writes its
628 + # top-k into the shared ``topk_indices_buffer``; the attend reads it back.
629 + self.indexer(index_query)
630 + return self.impl.forward(self, query, self.kv_cache, output)
631
632
633 class MiniMaxM3DecoderLayer(nn.Module):
@@ -627,6 +639,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
639 force_sparse_attn: bool = False,
640 force_moe: bool = False,
641 is_mtp_block: bool = False,
642 + topk_indices_buffer: torch.Tensor | None = None,
643 ) -> None:
644 super().__init__()
645 if is_mtp_block:
@@ -662,6 +675,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
675 quant_config=quant_config,
676 prefix=f"{prefix}.self_attn",
677 cache_config=cache_config,
678 + topk_indices_buffer=topk_indices_buffer,
679 )
680 else:
681 self.self_attn = MiniMaxM3Attention(
@@ -747,11 +761,33 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
761 prefix=f"{prefix}.embed_tokens",
762 )
763
764 + # Reserved top-k indices buffer shared by all sparse-attention indexer
765 + # layers (mirrors DeepseekV4); kept at a stable address so the indexer's
766 + # top-k output survives cudagraph capture/replay. Shape matches the
767 + # per-head index top-k output [num_index_heads, total_q, topk].
768 + sparse_cfg = getattr(config, "sparse_attention_config", None)
769 + if sparse_cfg is not None:
770 + tp_size = get_tensor_model_parallel_world_size()
771 + num_index_heads = max(1, sparse_cfg["sparse_num_index_heads"] // tp_size)
772 + # Pad tokens to a multiple of 4 so the buffer head stride stays
773 + # int4-aligned for build_k2q_csr's vectorised int4 loads.
774 + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
775 + padded_num_tokens = (max_num_batched_tokens + 3) // 4 * 4
776 + self.topk_indices_buffer = torch.empty(
777 + num_index_heads,
778 + padded_num_tokens,
779 + sparse_cfg["sparse_topk_blocks"],
780 + dtype=torch.int32,
781 + )
782 + else:
783 + self.topk_indices_buffer = None
784 +
785 self.start_layer, self.end_layer, self.layers = make_layers(
786 config.num_hidden_layers,
787 lambda prefix: MiniMaxM3DecoderLayer(
788 vllm_config=vllm_config,
789 prefix=prefix,
790 + topk_indices_buffer=self.topk_indices_buffer,
791 ),
792 prefix=f"{prefix}.layers",
793 )
vllm/models/minimax_m3/nvidia/sparse_attention_msa.py
+7 -5
@@ -29,7 +29,6 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
29 layer: AttentionLayer,
30 query: torch.Tensor,
31 kv_cache: torch.Tensor,
32 - topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
32 output: torch.Tensor,
33 ) -> torch.Tensor:
34 attn_metadata = get_forward_context().attn_metadata
@@ -37,10 +36,12 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
36 return output # profiling run; caches unbound
37 main_md = attn_metadata[layer.layer_name] # type: ignore[attr-defined]
38 assert isinstance(main_md, MiniMaxM3SparseMetadata)
40 - decode_topk, prefill_topk = topk_idx
39
40 nd = main_md.num_decode_tokens
41 num_tokens = main_md.num_actual_tokens
42 + # Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:].
43 + topk = layer.topk_indices_buffer # type: ignore[attr-defined]
44 + assert topk is not None
45 hd = self.head_size
46 q = query[:num_tokens].view(-1, self.num_heads, hd)
47 out = output[:num_tokens].view(-1, self.num_heads, hd)
@@ -51,11 +52,11 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
52 # Decode [:nd]: Triton split-K placeholder (no MSA decode yet).
53 if main_md.num_decodes > 0:
54 d = main_md.decode
54 - assert d is not None and decode_topk is not None
55 + assert d is not None
56 minimax_m3_sparse_attn_decode(
57 q[:nd],
58 kv_cache,
58 - decode_topk,
59 + topk[:, :nd, :],
60 d.block_table,
61 d.seq_lens,
62 self.num_kv_heads,
@@ -72,7 +73,8 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
73 )
74
75 p = main_md.prefill
75 - assert p is not None and prefill_topk is not None
76 + assert p is not None
77 + prefill_topk = topk[:, nd:num_tokens, :]
78 qp = q[nd:]
79 k_cache = kv_cache[:, 0].transpose(1, 2)
80 v_cache = kv_cache[:, 1].transpose(1, 2)