@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ee5a89f4d7

[ROCm][MiniMax-M3] Add AITER sparse paged attention (#47287)

Signed-off-by: Tan Pin Siang <tanpinsiang@gmail.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: Hongxia Yang <hongxia.yang@amd.com> Co-authored-by: Jun Kang Chow <junkangchow@gmail.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>

Tan Pin Siang committed Jul 13, 2026 at 10:27 UTC ee5a89f4d7b818af7d92770dbf6aa1b41b503b4a
15 files changed +1553 -127
csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
+126 -75
@@ -39,12 +39,10 @@
39 * The IQ/IK warps address the index_q/index_k sub-blocks *inside* qkv at the
40 * fixed physical offsets (nq+2*nkv)*128 and (nq+2*nkv+niq)*128.
41 *
42 - * Dense vs sparse is a compile-time choice via the ``kIsSparse``/``kInsertKV``
43 - * template bools (3 instantiations: dense <false,false>, sparse-profiling
44 - * <true,false>, sparse-serving <true,true>), so the index slots, the V slots
45 - * and the cache inserts fold away entirely on paths that don't use them. The
46 - * dense layer passes no caches/index: norm+RoPE happens in place and the
47 - * generic ``Attention`` layer owns the cache write.
42 + * Dense vs sparse row layout and index-branch processing are separate template
43 + * choices. Skip-index-topk reuse layers still have sparse rows and insert main
44 + * K/V cache entries, but compile away index_q/index_k work and index-cache
45 + * writes.
46 *
47 * Q/K and (sparse) index_q/index_k are all rewritten in place inside the fused
48 * ``qkv`` tensor. Caches (bf16) are scatter-written by slot.
@@ -223,10 +221,25 @@ __device__ __forceinline__ void storeCacheElems(
221 // model dtype directly. FP8 cache dtypes use the conversion path below.
222 storeElems<scalar_t>(reinterpret_cast<scalar_t*>(dst), elems);
223 } else {
226 -#pragma unroll
224 +#ifdef USE_ROCM
225 + // Match ROCm's model-dtype materialization before FP8 cache conversion.
226 + using Converter = vllm::_typeConvert<scalar_t>;
227 + using rounded_t = typename Converter::hip_type;
228 + rounded_t rounded[kElemsPerLane];
229 + #pragma unroll
230 + for (int i = 0; i < kElemsPerLane; i++) {
231 + rounded[i] = Converter::convert(elems[i]);
232 + }
233 + #pragma unroll
234 + for (int i = 0; i < kElemsPerLane; i++) {
235 + dst[i] = fp8::scaled_convert<cache_t, rounded_t, kv_dt>(rounded[i], 1.0f);
236 + }
237 +#else
238 + #pragma unroll
239 for (int i = 0; i < kElemsPerLane; i++) {
240 dst[i] = fp8::scaled_convert<cache_t, float, kv_dt>(elems[i], 1.0f);
241 }
242 +#endif
243 }
244 }
245
@@ -262,20 +275,25 @@ __device__ __forceinline__ void storeElemsFp8(
275 // Grid: 1D, ceil(num_tokens * slots_per_token / warps_per_block).
276 // Each warp = one (token, slot).
277 //
265 -// `kIsSparse` and `kInsertKV` are compile-time template bools, so all the
266 -// branch decisions that distinguish the dense layer from the sparse layer
267 -// (index slots, KV/index inserts, V slots) fold away per instantiation.
268 -// Three instantiations are built: dense <false,false>, sparse-profiling
269 -// <true,false> and sparse-serving <true,true>. Slots per token:
278 +// `kHasIndex`, `kProcessIndex`, and `kInsertKV` are compile-time template
279 +// bools, so branch decisions that distinguish the dense layer from the sparse
280 +// layer (index slots, KV/index inserts, V slots) fold away per instantiation.
281 +// Slots per token:
282 // Q : nq (always — norm+RoPE)
283 // K : nkv (always — norm+RoPE; +K-cache insert)
284 // 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)
285 +// IQ: niq only if kProcessIndex (norm+RoPE)
286 +// IK: 1 only if kProcessIndex (norm+RoPE; +index-cache insert)
287 // cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx:
288 // indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte).
289 +// kHasIndex means the qkv row is laid out as sparse [q|k|v|index_q|index_k].
290 +// kProcessIndex controls whether this launch actually norms/ropes the index
291 +// branch and writes index_q/index_k outputs. Skip-index-topk reuse layers keep
292 +// kHasIndex=true but set kProcessIndex=false.
293 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
278 - typename out_idx_t, bool kIsSparse, bool kInsertKV, bool kFp8Idx>
294 + typename out_idx_t, bool kHasIndex, bool kInsertKV,
295 + bool kProcessIndex,
296 + bool kFp8Idx>
297 __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
298 scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse)
299 scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
@@ -308,9 +326,12 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
326 int const laneId = threadIdx.x % 32;
327 int const globalWarpIdx = blockIdx.x * warpsPerBlock + (threadIdx.x / 32);
328
329 + static_assert(!kProcessIndex || kHasIndex,
330 + "index processing requires sparse row layout");
331 +
332 // Slot layout (compile-time gated: dense has neither V nor index slots).
333 int const v_slots = kInsertKV ? nkv : 0;
313 - int const idx_slots = kIsSparse ? niq + 1 : 0;
334 + int const idx_slots = kProcessIndex ? niq + 1 : 0;
335 int const slots_per_token = nq + nkv + v_slots + idx_slots;
336
337 int const tokenIdx = globalWarpIdx / slots_per_token;
@@ -321,14 +342,14 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
342 int const k_begin = nq;
343 int const v_begin = nq + nkv; // valid only when kInsertKV
344 int const iq_begin = nq + nkv + v_slots; // index block start
324 - int const ik_slot = iq_begin + niq; // valid only when kIsSparse
345 + int const ik_slot = iq_begin + niq; // valid only when kProcessIndex
346
347 bool const isQ = slot < k_begin;
348 bool const isK = slot >= k_begin && slot < v_begin;
349 bool isV = false;
350 if constexpr (kInsertKV) isV = slot >= v_begin && slot < v_begin + nkv;
351 bool isIQ = false, isIK = false;
331 - if constexpr (kIsSparse) {
352 + if constexpr (kProcessIndex) {
353 isIQ = slot >= iq_begin && slot < ik_slot;
354 isIK = slot == ik_slot;
355 }
@@ -336,7 +357,7 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
357 int const dim_base = laneId * kElemsPerLane;
358 // Physical row width of qkv: the dense layer packs [q|k|v]; the sparse
359 // layer additionally packs [index_q (niq heads) | index_k (1 head)].
339 - int const qkv_row = (nq + 2 * nkv + (kIsSparse ? (niq + 1) : 0)) * kHeadDim;
360 + int const qkv_row = (nq + 2 * nkv + (kHasIndex ? (niq + 1) : 0)) * kHeadDim;
361
362 // ── Resolve source pointer + per-branch parameters. ────────────────────
363 scalar_t* row_ptr = nullptr; // in-place output location
@@ -367,10 +388,13 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
388 row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
389 (nq + 2 * nkv + ih) * kHeadDim;
390 norm_w = iq_norm_w;
370 - } else { // isIK -- single shared index key at (nq+2*nkv+niq)*128.
391 + } else if (isIK) {
392 + // Single shared index key at (nq+2*nkv+niq)*128.
393 row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
394 (nq + 2 * nkv + niq) * kHeadDim;
395 norm_w = ik_norm_w;
396 + } else {
397 + return;
398 }
399
400 // Store destination. Q and index_q are gathered into dedicated contiguous
@@ -426,9 +450,12 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
450 // ── Cache inserts (sparse serving only). ───────────────────────────────
451 if constexpr (kInsertKV) {
452 // Guard (not early-return) so every thread reaches the PDL trigger below.
429 - int64_t const sm = (isK || isV)
430 - ? slot_mapping[tokenIdx]
431 - : (isIK ? index_slot_mapping[tokenIdx] : -1);
453 + int64_t sm = -1;
454 + if (isK || isV) {
455 + sm = slot_mapping[tokenIdx];
456 + } else if constexpr (kProcessIndex) {
457 + if (isIK) sm = index_slot_mapping[tokenIdx];
458 + }
459 if (sm >= 0) { // skip padded / unscheduled tokens
460 if (isIK) {
461 if constexpr (kFp8Idx) {
@@ -475,12 +502,12 @@ void launchFusedMiniMaxM3(
502 int const nkv, int const niq, int const block_size,
503 int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token,
504 int64_t const kv_s_dim, bool const has_index, bool const insert_kv,
478 - bool const fp8_idx, cudaStream_t stream) {
505 + bool const process_index, bool const fp8_idx, cudaStream_t stream) {
506 // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the
507 // void* pointers per instantiation in the LAUNCH macro.
508 // Slot count must match the kernel's compile-time gating.
509 int const v_slots = insert_kv ? nkv : 0;
483 - int const idx_slots = has_index ? niq + 1 : 0;
510 + int const idx_slots = process_index ? niq + 1 : 0;
511 int const slots_per_token = nq + nkv + v_slots + idx_slots;
512
513 constexpr int kBlockSize = 256;
@@ -507,11 +534,12 @@ void launchFusedMiniMaxM3(
534 config.attrs = attrs;
535 config.numAttrs = (sm_version >= 90) ? 1 : 0;
536
510 - #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T) \
537 + #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \
538 cudaLaunchKernelEx( \
539 &config, \
540 fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
514 - IS_SPARSE, INSERT, FP8>, \
541 + HAS_INDEX, INSERT, \
542 + PROCESS_INDEX, FP8>, \
543 qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, k_norm_w, \
544 iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping, \
545 index_slot_mapping, kv_cache, reinterpret_cast<OUT_T*>(index_cache), \
@@ -520,37 +548,45 @@ void launchFusedMiniMaxM3(
548 #else
549 // ROCm: standard kernel launch syntax (no PDL/stream serialization).
550 // clang-format off
523 - #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T) \
524 - fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
525 - IS_SPARSE, INSERT, FP8> \
526 - <<<grid, kBlockSize, 0, stream>>>( \
527 - qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, \
528 - k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
529 - slot_mapping, index_slot_mapping, kv_cache, \
530 - reinterpret_cast<OUT_T*>(index_cache), eps, rotary_dim, \
531 - num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_head, \
532 - kv_s_token, kv_s_dim)
551 + #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \
552 + fusedMiniMaxM3QNormRopeKVInsertKernel< \
553 + scalar_t, cache_t, kv_dt, OUT_T, HAS_INDEX, INSERT, PROCESS_INDEX, \
554 + FP8><<<grid, kBlockSize, 0, stream>>>( \
555 + qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, \
556 + k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
557 + slot_mapping, index_slot_mapping, kv_cache, \
558 + reinterpret_cast<OUT_T*>(index_cache), eps, rotary_dim, num_tokens, \
559 + nq, nkv, niq, block_size, kv_s_block, kv_s_head, kv_s_token, \
560 + kv_s_dim)
561 // clang-format on
562 #endif
563
564 if (has_index) {
537 - if (insert_kv) {
565 + if (!process_index) {
566 + if (insert_kv) {
567 + LAUNCH(true, true, false, false, scalar_t);
568 + } else {
569 + LAUNCH(true, false, false, false, scalar_t);
570 + }
571 + } else if (insert_kv) {
572 if (fp8_idx) {
539 - LAUNCH(true, true, true, uint8_t); // sparse serving, fp8 index outputs
573 + LAUNCH(true, true, true, true,
574 + uint8_t); // sparse serving, fp8 index outputs
575 } else {
541 - LAUNCH(true, true, false, scalar_t); // sparse serving, bf16
576 + LAUNCH(true, true, true, false, scalar_t); // sparse serving, bf16
577 }
578 } else {
579 if (fp8_idx) {
545 - LAUNCH(true, false, true, uint8_t); // sparse profiling, fp8 index_q
580 + LAUNCH(true, false, true, true,
581 + uint8_t); // sparse profiling, fp8 index_q
582 } else {
547 - LAUNCH(true, false, false, scalar_t); // sparse profiling, bf16
583 + LAUNCH(true, false, true, false, scalar_t); // sparse profiling, bf16
584 }
585 }
586 } else {
587 // Dense layer: never has an index branch and never inserts here (the
588 // generic Attention layer owns the KV insert).
553 - LAUNCH(false, false, false, scalar_t);
589 + LAUNCH(false, false, false, false, scalar_t);
590 }
591 #undef LAUNCH
592 }
@@ -558,6 +594,7 @@ void launchFusedMiniMaxM3(
594 } // namespace minimax_m3_fused_ops
595 } // namespace vllm
596
597 +// clang-format off
598 #define CALL_FUSED_MINIMAX_M3(_RAW_T, CACHE_T, KV_DTYPE) \
599 vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>( \
600 reinterpret_cast<st*>(qkv.data_ptr()), \
@@ -567,24 +604,29 @@ void launchFusedMiniMaxM3(
604 : nullptr, \
605 reinterpret_cast<st const*>(q_norm_weight.data_ptr()), \
606 reinterpret_cast<st const*>(k_norm_weight.data_ptr()), \
570 - has_index ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr()) \
571 - : nullptr, \
572 - has_index ? reinterpret_cast<st const*>(index_k_norm_weight->data_ptr()) \
573 - : nullptr, \
607 + process_index \
608 + ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr()) \
609 + : nullptr, \
610 + process_index \
611 + ? reinterpret_cast<st const*>(index_k_norm_weight->data_ptr()) \
612 + : nullptr, \
613 reinterpret_cast<st const*>(cos_sin_cache.data_ptr()), \
614 reinterpret_cast<int64_t const*>(positions.data_ptr()), \
615 insert_kv ? reinterpret_cast<int64_t const*>(slot_mapping->data_ptr()) \
616 : nullptr, \
578 - insert_kv ? reinterpret_cast<int64_t const*>( \
579 - effective_index_slot_mapping->data_ptr()) \
580 - : nullptr, \
617 + (insert_kv && process_index) \
618 + ? reinterpret_cast<int64_t const*>( \
619 + effective_index_slot_mapping->data_ptr()) \
620 + : nullptr, \
621 insert_kv ? reinterpret_cast<CACHE_T*>(kv_cache->data_ptr()) : nullptr, \
582 - (insert_kv && has_index) \
622 + (insert_kv && process_index) \
623 ? reinterpret_cast<void*>(index_cache->data_ptr()) \
624 : nullptr, \
625 static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq, \
626 nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_head, \
587 - kv_s_token, kv_s_dim, has_index, insert_kv, fp8_idx, stream)
627 + kv_s_token, kv_s_dim, has_index, insert_kv, process_index, fp8_idx, \
628 + stream)
629 +// clang-format on
630
631 // ────────────────────────────────────────────────────────────────────────────
632 // Torch op wrapper
@@ -607,7 +649,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
649 std::optional<torch::stable::Tensor> q_out, // [N, nq*128] contiguous
650 std::optional<torch::stable::Tensor>
651 index_q_out, // [N, niq*128] contiguous
610 - const std::string& kv_cache_dtype) {
652 + const std::string& kv_cache_dtype, bool skip_index_branch) {
653 STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
654 "qkv must be contiguous CUDA");
655 STD_TORCH_CHECK(
@@ -646,6 +688,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
688 // (1 head)]) right after [q|k|v] in the same row; the dense layer does not.
689 bool const has_index = niq > 0;
690 bool const insert_kv = kv_cache.has_value();
691 + bool const process_index = has_index && !skip_index_branch;
692 vllm::Fp8KVCacheDataType const kv_dt =
693 vllm::get_fp8_kv_cache_data_type(kv_cache_dtype);
694 int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
@@ -661,7 +704,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
704 STD_TORCH_CHECK(
705 !insert_kv || has_index,
706 "insert mode (kv_cache) requires the index branch (sparse layer)");
664 - if (has_index) {
707 + STD_TORCH_CHECK(has_index || !skip_index_branch,
708 + "skip_index_branch requires sparse qkv rows");
709 + if (process_index) {
710 STD_TORCH_CHECK(
711 index_q_norm_weight.has_value() && index_k_norm_weight.has_value(),
712 "index branch requires both index norm weights");
@@ -683,13 +728,15 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
728 slot_mapping.has_value() && slot_mapping->is_cuda() &&
729 slot_mapping->scalar_type() == torch::headeronly::ScalarType::Long,
730 "insert mode requires int64 CUDA slot_mapping");
686 - STD_TORCH_CHECK(
687 - !index_slot_mapping.has_value() ||
688 - (index_slot_mapping->is_cuda() &&
689 - index_slot_mapping->scalar_type() ==
690 - torch::headeronly::ScalarType::Long &&
691 - index_slot_mapping->numel() == slot_mapping->numel()),
692 - "index_slot_mapping must be int64 CUDA with slot_mapping length");
731 + if (process_index) {
732 + STD_TORCH_CHECK(
733 + !index_slot_mapping.has_value() ||
734 + (index_slot_mapping->is_cuda() &&
735 + index_slot_mapping->scalar_type() ==
736 + torch::headeronly::ScalarType::Long &&
737 + index_slot_mapping->numel() == slot_mapping->numel()),
738 + "index_slot_mapping must be int64 CUDA with slot_mapping length");
739 + }
740 // Main attention KV cache: auto matches qkv, fp8 uses uint8 storage.
741 if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
742 STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
@@ -700,12 +747,14 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
747 "fp8 kv_cache must use uint8 storage");
748 }
749 // Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3.
703 - STD_TORCH_CHECK(
704 - index_cache.has_value() &&
705 - (index_cache->scalar_type() == qkv.scalar_type() ||
706 - index_cache->scalar_type() ==
707 - torch::headeronly::ScalarType::Float8_e4m3fn),
708 - "insert mode requires index_cache matching qkv dtype or fp8 e4m3");
750 + if (process_index) {
751 + STD_TORCH_CHECK(
752 + index_cache.has_value() &&
753 + (index_cache->scalar_type() == qkv.scalar_type() ||
754 + index_cache->scalar_type() ==
755 + torch::headeronly::ScalarType::Float8_e4m3fn),
756 + "insert mode requires index_cache matching qkv dtype or fp8 e4m3");
757 + }
758 STD_TORCH_CHECK(kv_cache->dim() == 4 && kv_cache->stride(3) == 1,
759 "kv_cache must be [nb,nkv,bs,2*head_dim] with contiguous "
760 "content dim (stride(3)==1)");
@@ -713,9 +762,11 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
762 kv_s_head = kv_cache->stride(1);
763 kv_s_token = kv_cache->stride(2);
764 kv_s_dim = kv_cache->stride(3);
716 - effective_index_slot_mapping = index_slot_mapping.has_value()
717 - ? &index_slot_mapping.value()
718 - : &slot_mapping.value();
765 + if (process_index) {
766 + effective_index_slot_mapping = index_slot_mapping.has_value()
767 + ? &index_slot_mapping.value()
768 + : &slot_mapping.value();
769 + }
770 }
771 // Optional contiguous gather targets: when given, the normed/roped q (and
772 // index_q) are written here instead of in place, so callers avoid a separate
@@ -730,9 +781,8 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
781 "q_out must have num_tokens * num_heads * 128 elements");
782 }
783 if (index_q_out.has_value()) {
733 - STD_TORCH_CHECK(
734 - has_index,
735 - "index_q_out requires the index branch (num_index_heads > 0)");
784 + STD_TORCH_CHECK(process_index,
785 + "index_q_out requires index branch processing");
786 STD_TORCH_CHECK(
787 index_q_out->is_cuda() && index_q_out->is_contiguous() &&
788 (index_q_out->scalar_type() == qkv.scalar_type() ||
@@ -749,8 +799,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
799 // q/k/v + q_out stay qkv dtype. Both index outputs must agree.
800 auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn;
801 bool const fp8_idx =
752 - (index_cache.has_value() && index_cache->scalar_type() == kFp8) ||
753 - (index_q_out.has_value() && index_q_out->scalar_type() == kFp8);
802 + process_index &&
803 + ((index_cache.has_value() && index_cache->scalar_type() == kFp8) ||
804 + (index_q_out.has_value() && index_q_out->scalar_type() == kFp8));
805 if (fp8_idx) {
806 STD_TORCH_CHECK(
807 !index_cache.has_value() || index_cache->scalar_type() == kFp8,
csrc/libtorch_stable/ops.h
+1 -1
@@ -313,7 +313,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
313 std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
314 std::optional<torch::stable::Tensor> q_out,
315 std::optional<torch::stable::Tensor> index_q_out,
316 - const std::string& kv_cache_dtype);
316 + const std::string& kv_cache_dtype, bool skip_index_branch);
317
318 // Sampler kernels (shared CUDA/ROCm)
319 void apply_repetition_penalties_(
csrc/libtorch_stable/torch_bindings.cpp
+1 -1
@@ -466,7 +466,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
466 "Tensor? slot_mapping, Tensor? index_slot_mapping, "
467 "Tensor!? kv_cache, Tensor!? index_cache, "
468 "int block_size, Tensor!? q_out, Tensor!? index_q_out, "
469 - "str kv_cache_dtype) -> ()");
469 + "str kv_cache_dtype, bool skip_index_branch=False) -> ()");
470
471 // Apply repetition penalties to logits in-place.
472 ops.def(
tests/kernels/attention/test_minimax_m3.py
+41
@@ -820,6 +820,47 @@ def test_main_backend_layout_contract():
820 )
821
822
823 +def test_aiter_sparse_pa_layout_contract(monkeypatch):
824 + """The shuffle-only AITER path retains separately contiguous K/V storage."""
825 + import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
826 +
827 + monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
828 + monkeypatch.setattr(
829 + sparse_attn_mod.rocm_aiter_ops,
830 + "is_shuffle_kv_cache_enabled",
831 + lambda: True,
832 + )
833 +
834 + nb, bs, h, d = 7, BLOCK_SIZE, 1, HEAD_DIM
835 + logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
836 + order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
837 + assert logical == (nb, 2, bs, h, d)
838 + assert order == (1, 0, 2, 3, 4)
839 +
840 + physical_shape = tuple(logical[i] for i in order)
841 + inv_order = [order.index(i) for i in range(len(order))]
842 + raw = torch.empty(physical_shape, device="cuda", dtype=DTYPE)
843 + logical_view = raw.permute(*inv_order)
844 + key_cache, value_cache = logical_view.unbind(1)
845 + assert key_cache.is_contiguous()
846 + assert value_cache.is_contiguous()
847 +
848 +
849 +def test_aiter_sparse_pa_rejects_multiple_kv_heads(monkeypatch):
850 + """Do not pair AITER's separated cache layout with the Triton fallback."""
851 + import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
852 +
853 + monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
854 + monkeypatch.setattr(
855 + sparse_attn_mod.rocm_aiter_ops,
856 + "is_shuffle_kv_cache_enabled",
857 + lambda: True,
858 + )
859 +
860 + with pytest.raises(ValueError, match="num_kv_heads == 1"):
861 + MiniMaxM3SparseBackend.get_kv_cache_shape(7, BLOCK_SIZE, 2, HEAD_DIM)
862 +
863 +
864 def test_main_backend_unknown_layout_raises(monkeypatch):
865 """An unrecognized layout (injected past env-var validation) is rejected."""
866 import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
+130 -13
@@ -171,10 +171,9 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
171 kv_cache_storage_dtype = torch.uint8 if kv_cache_dtype == "fp8" else dtype
172 kv_cache = torch.zeros(
173 num_blocks,
174 - 2,
175 - block_size,
174 num_kv_heads,
177 - HEAD_DIM,
175 + block_size,
176 + 2 * HEAD_DIM,
177 dtype=kv_cache_storage_dtype,
178 device=device,
179 )
@@ -246,18 +245,21 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
245 torch.testing.assert_close(index_k, ik_ref, rtol=1e-2, atol=1e-2)
246
247 # ── Cache inserts. ──
249 - # Main cache layout is [num_blocks, 2, block_size, num_kv_heads, head_dim]
250 - # (the K/V axis sits *before* block_size); index cache is [nb, bs, head_dim].
248 + # Main cache layout is [num_blocks, num_kv_heads, block_size, 2*head_dim];
249 + # index cache is [num_blocks, block_size, head_dim].
250 k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
251 v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM) # v is raw (no norm/rope)
252 if kv_cache_dtype == "fp8":
253 expected_kv_cache = torch.zeros_like(kv_cache)
254 + expected_k_cache, expected_v_cache = expected_kv_cache.transpose(1, 2).split(
255 + HEAD_DIM, dim=-1
256 + )
257 scale = torch.ones((), device=device)
258 ops.reshape_and_cache_flash(
259 k_out.view(num_tokens, num_kv_heads, HEAD_DIM),
260 v_out.view(num_tokens, num_kv_heads, HEAD_DIM),
259 - expected_kv_cache[:, 0],
260 - expected_kv_cache[:, 1],
261 + expected_k_cache,
262 + expected_v_cache,
263 slot_mapping,
264 kv_cache_dtype,
265 scale,
@@ -269,9 +271,11 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
271 s = slot_mapping[t].item()
272 b, pos = s // block_size, s % block_size
273 torch.testing.assert_close(
272 - kv_cache[b, 0, pos], k_ref_h[t], rtol=1e-2, atol=1e-2
274 + kv_cache[b, :, pos, :HEAD_DIM], k_ref_h[t], rtol=1e-2, atol=1e-2
275 + )
276 + torch.testing.assert_close(
277 + kv_cache[b, :, pos, HEAD_DIM:], v_ref_h[t], rtol=0, atol=0
278 )
274 - torch.testing.assert_close(kv_cache[b, 1, pos], v_ref_h[t], rtol=0, atol=0)
279
280 expected_index_cache = torch.zeros_like(index_cache).view(-1, HEAD_DIM)
281 expected_index_cache[index_slot_mapping] = index_k
@@ -280,7 +284,121 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
284 )
285
286
283 -# ── Test 3: fp8 (e4m3) index outputs ─────────────────────────────────────────
287 +@pytest.mark.parametrize("num_tokens", [1, 64, 513])
288 +@pytest.mark.parametrize("block_size", [16, 64])
289 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
290 +def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype):
291 + torch.manual_seed(2)
292 + device, dtype, eps = "cuda", torch.bfloat16, 1e-6
293 + base, max_pos = 5_000_000.0, 4096
294 + num_heads, num_kv_heads, num_idx_heads = 16, 4, 4
295 +
296 + q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
297 + k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
298 + cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
299 + positions = torch.randint(
300 + 0, max_pos, (num_tokens,), dtype=torch.int64, device=device
301 + )
302 +
303 + qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
304 + iqsz, iksz = num_idx_heads * HEAD_DIM, HEAD_DIM
305 + qkv = torch.randn(
306 + num_tokens, qsz + 2 * kvsz + iqsz + iksz, dtype=dtype, device=device
307 + )
308 + qkv_orig = qkv.clone()
309 + splits = [qsz, kvsz, kvsz, iqsz, iksz]
310 +
311 + num_blocks = (num_tokens + block_size - 1) // block_size + 1
312 + kv_cache_storage_dtype = torch.uint8 if kv_cache_dtype == "fp8" else dtype
313 + kv_cache = torch.zeros(
314 + num_blocks,
315 + num_kv_heads,
316 + block_size,
317 + 2 * HEAD_DIM,
318 + dtype=kv_cache_storage_dtype,
319 + device=device,
320 + )
321 + index_cache = torch.randn(
322 + num_blocks, block_size, HEAD_DIM, dtype=dtype, device=device
323 + )
324 + index_cache_orig = index_cache.clone()
325 + slot_mapping = torch.randperm(
326 + num_blocks * block_size, dtype=torch.int64, device=device
327 + )[:num_tokens]
328 + q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device)
329 +
330 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
331 + qkv,
332 + q_w,
333 + k_w,
334 + cos_sin,
335 + positions,
336 + num_heads,
337 + num_kv_heads,
338 + ROTARY_DIM,
339 + eps,
340 + num_index_heads=num_idx_heads,
341 + slot_mapping=slot_mapping,
342 + kv_cache=kv_cache,
343 + index_cache=index_cache,
344 + block_size=block_size,
345 + q_out=q_out,
346 + kv_cache_dtype=kv_cache_dtype,
347 + skip_index_branch=True,
348 + )
349 +
350 + _, k_out, v_out, index_q_out, index_k_out = qkv.split(splits, dim=-1)
351 + q_in, k_in, v_in, index_q_in, index_k_in = qkv_orig.split(splits, dim=-1)
352 + q_ref = norm_rope_ref(
353 + q_in.view(num_tokens, num_heads, HEAD_DIM), q_w, positions, cos_sin, eps
354 + ).view(num_tokens, qsz)
355 + k_ref = norm_rope_ref(
356 + k_in.view(num_tokens, num_kv_heads, HEAD_DIM),
357 + k_w,
358 + positions,
359 + cos_sin,
360 + eps,
361 + ).view(num_tokens, kvsz)
362 +
363 + torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2)
364 + torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2)
365 + torch.testing.assert_close(v_out, v_in, rtol=0, atol=0)
366 + torch.testing.assert_close(index_q_out, index_q_in, rtol=0, atol=0)
367 + torch.testing.assert_close(index_k_out, index_k_in, rtol=0, atol=0)
368 + torch.testing.assert_close(index_cache, index_cache_orig, rtol=0, atol=0)
369 +
370 + if kv_cache_dtype == "fp8":
371 + expected_kv_cache = torch.zeros_like(kv_cache)
372 + expected_k_cache, expected_v_cache = expected_kv_cache.transpose(1, 2).split(
373 + HEAD_DIM, dim=-1
374 + )
375 + scale = torch.ones((), device=device)
376 + ops.reshape_and_cache_flash(
377 + k_out.view(num_tokens, num_kv_heads, HEAD_DIM),
378 + v_out.view(num_tokens, num_kv_heads, HEAD_DIM),
379 + expected_k_cache,
380 + expected_v_cache,
381 + slot_mapping,
382 + kv_cache_dtype,
383 + scale,
384 + scale,
385 + )
386 + torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
387 + else:
388 + k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
389 + v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM)
390 + for t in range(num_tokens):
391 + s = slot_mapping[t].item()
392 + b, pos = s // block_size, s % block_size
393 + torch.testing.assert_close(
394 + kv_cache[b, :, pos, :HEAD_DIM], k_ref_h[t], rtol=1e-2, atol=1e-2
395 + )
396 + torch.testing.assert_close(
397 + kv_cache[b, :, pos, HEAD_DIM:], v_ref_h[t], rtol=0, atol=0
398 + )
399 +
400 +
401 +# ── Test 4: fp8 (e4m3) index outputs ─────────────────────────────────────────
402 # The fp8 score path stores index_q and the index-K cache as e4m3 while q/k/v +
403 # q_out stay bf16. Asserts: (1) q/k/v/q_out are bit-identical to the bf16 run
404 # (the index dtype must not perturb the main branch), and (2) the e4m3 index
@@ -324,10 +442,9 @@ def test_sparse_full_fp8_index(num_tokens, block_size):
442 qkv = qkv0.clone()
443 kv_cache = torch.zeros(
444 num_blocks,
327 - 2,
328 - block_size,
445 num_kv_heads,
330 - HEAD_DIM,
446 + block_size,
447 + 2 * HEAD_DIM,
448 dtype=dtype,
449 device=device,
450 )
tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py new
+189
@@ -0,0 +1,189 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from vllm.platforms import current_platform
8 +
9 +if not torch.cuda.is_available() or not current_platform.supports_fp8():
10 + pytest.skip(
11 + "MiniMax-M3 FP8 sparse attention scale tests require an FP8 GPU.",
12 + allow_module_level=True,
13 + )
14 +
15 +if current_platform.is_rocm():
16 + from vllm.models.minimax_m3.amd.ops.sparse_attn import (
17 + SPARSE_BLOCK_SIZE,
18 + minimax_m3_sparse_attn,
19 + minimax_m3_sparse_attn_decode,
20 + )
21 +else:
22 + from vllm.models.minimax_m3.common.ops.sparse_attn import (
23 + SPARSE_BLOCK_SIZE,
24 + minimax_m3_sparse_attn,
25 + minimax_m3_sparse_attn_decode,
26 + )
27 +
28 +
29 +DEVICE = "cuda"
30 +DTYPE = torch.bfloat16
31 +HEAD_DIM = 128
32 +NUM_KV_HEADS = 1
33 +NUM_HEADS = 2
34 +K_SCALE = 0.25
35 +V_SCALE = 0.5
36 +
37 +
38 +def _scale_tensors(mode: str, num_blocks: int):
39 + if mode == "scalar":
40 + k_scale = torch.tensor(K_SCALE, dtype=torch.float32, device=DEVICE)
41 + v_scale = torch.tensor(V_SCALE, dtype=torch.float32, device=DEVICE)
42 + else:
43 + shape = (NUM_KV_HEADS, num_blocks * SPARSE_BLOCK_SIZE)
44 + k_scale = torch.full(shape, K_SCALE, dtype=torch.float32, device=DEVICE)
45 + v_scale = torch.full(shape, V_SCALE, dtype=torch.float32, device=DEVICE)
46 + return k_scale, v_scale
47 +
48 +
49 +def _make_kv_cache(num_blocks: int, seed: int):
50 + torch.manual_seed(seed)
51 + fp8_dtype = current_platform.fp8_dtype()
52 + kv_ref = torch.randn(
53 + num_blocks,
54 + NUM_KV_HEADS,
55 + SPARSE_BLOCK_SIZE,
56 + 2 * HEAD_DIM,
57 + dtype=DTYPE,
58 + device=DEVICE,
59 + )
60 + kv_fp8 = torch.empty_like(kv_ref, dtype=fp8_dtype)
61 + kv_fp8[..., :HEAD_DIM] = (kv_ref[..., :HEAD_DIM].float() / K_SCALE).to(fp8_dtype)
62 + kv_fp8[..., HEAD_DIM:] = (kv_ref[..., HEAD_DIM:].float() / V_SCALE).to(fp8_dtype)
63 +
64 + kv_dequant = torch.empty_like(kv_ref)
65 + kv_dequant[..., :HEAD_DIM] = (kv_fp8[..., :HEAD_DIM].float() * K_SCALE).to(DTYPE)
66 + kv_dequant[..., HEAD_DIM:] = (kv_fp8[..., HEAD_DIM:].float() * V_SCALE).to(DTYPE)
67 + return kv_fp8, kv_dequant
68 +
69 +
70 +@pytest.mark.parametrize("scale_mode", ["scalar", "per_token_head"])
71 +@torch.inference_mode()
72 +def test_minimax_m3_sparse_prefill_fp8_kv_scales(scale_mode: str):
73 + total_q = 17
74 + num_blocks = 1
75 + torch.manual_seed(0)
76 + q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, dtype=DTYPE, device=DEVICE) * 0.1
77 + kv_fp8, kv_dequant = _make_kv_cache(num_blocks, seed=1)
78 + k_scale, v_scale = _scale_tensors(scale_mode, num_blocks)
79 +
80 + topk = torch.zeros(NUM_KV_HEADS, total_q, 1, dtype=torch.int32, device=DEVICE)
81 + block_table = torch.zeros(1, 1, dtype=torch.int32, device=DEVICE)
82 + cu_seqlens = torch.tensor([0, total_q], dtype=torch.int32, device=DEVICE)
83 + seq_lens = torch.tensor([total_q], dtype=torch.int32, device=DEVICE)
84 + prefix_lens = torch.zeros(1, dtype=torch.int32, device=DEVICE)
85 + got = torch.empty_like(q)
86 + ref = torch.empty_like(q)
87 + unscaled = torch.empty_like(q)
88 +
89 + minimax_m3_sparse_attn(
90 + q,
91 + kv_fp8,
92 + topk,
93 + block_table,
94 + cu_seqlens,
95 + seq_lens,
96 + prefix_lens,
97 + total_q,
98 + NUM_KV_HEADS,
99 + HEAD_DIM**-0.5,
100 + got,
101 + k_scale=k_scale,
102 + v_scale=v_scale,
103 + )
104 + minimax_m3_sparse_attn(
105 + q,
106 + kv_dequant,
107 + topk,
108 + block_table,
109 + cu_seqlens,
110 + seq_lens,
111 + prefix_lens,
112 + total_q,
113 + NUM_KV_HEADS,
114 + HEAD_DIM**-0.5,
115 + ref,
116 + )
117 + minimax_m3_sparse_attn(
118 + q,
119 + kv_fp8,
120 + topk,
121 + block_table,
122 + cu_seqlens,
123 + seq_lens,
124 + prefix_lens,
125 + total_q,
126 + NUM_KV_HEADS,
127 + HEAD_DIM**-0.5,
128 + unscaled,
129 + )
130 +
131 + torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2)
132 + assert not torch.allclose(unscaled, ref, rtol=1e-1, atol=1e-1)
133 +
134 +
135 +@pytest.mark.parametrize("scale_mode", ["scalar", "per_token_head"])
136 +@torch.inference_mode()
137 +def test_minimax_m3_sparse_decode_fp8_kv_scales(scale_mode: str):
138 + total_q = 2
139 + num_blocks = 1
140 + torch.manual_seed(2)
141 + q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, dtype=DTYPE, device=DEVICE) * 0.1
142 + kv_fp8, kv_dequant = _make_kv_cache(num_blocks, seed=3)
143 + k_scale, v_scale = _scale_tensors(scale_mode, num_blocks)
144 +
145 + topk = torch.zeros(NUM_KV_HEADS, total_q, 1, dtype=torch.int32, device=DEVICE)
146 + block_table = torch.zeros(total_q, 1, dtype=torch.int32, device=DEVICE)
147 + seq_lens = torch.tensor([64, 128], dtype=torch.int32, device=DEVICE)
148 + got = torch.empty_like(q)
149 + ref = torch.empty_like(q)
150 + unscaled = torch.empty_like(q)
151 +
152 + minimax_m3_sparse_attn_decode(
153 + q,
154 + kv_fp8,
155 + topk,
156 + block_table,
157 + seq_lens,
158 + NUM_KV_HEADS,
159 + HEAD_DIM**-0.5,
160 + got,
161 + decode_query_len=1,
162 + k_scale=k_scale,
163 + v_scale=v_scale,
164 + )
165 + minimax_m3_sparse_attn_decode(
166 + q,
167 + kv_dequant,
168 + topk,
169 + block_table,
170 + seq_lens,
171 + NUM_KV_HEADS,
172 + HEAD_DIM**-0.5,
173 + ref,
174 + decode_query_len=1,
175 + )
176 + minimax_m3_sparse_attn_decode(
177 + q,
178 + kv_fp8,
179 + topk,
180 + block_table,
181 + seq_lens,
182 + NUM_KV_HEADS,
183 + HEAD_DIM**-0.5,
184 + unscaled,
185 + decode_query_len=1,
186 + )
187 +
188 + torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2)
189 + assert not torch.allclose(unscaled, ref, rtol=1e-1, atol=1e-1)
vllm/_custom_ops.py
+7
@@ -2516,6 +2516,7 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
2516 q_out: torch.Tensor | None = None,
2517 index_q_out: torch.Tensor | None = None,
2518 kv_cache_dtype: str = "auto",
2519 + skip_index_branch: bool = False,
2520 ) -> None:
2521 """Fused MiniMax-M3 attention pre-processing (in-place).
2522
@@ -2537,6 +2538,11 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
2538 instead of in place — folding the de-interleave into this kernel's store so
2539 callers skip a separate ``.contiguous()`` copy before the SM100 sparse
2540 attention's flat TMA descriptor.
2541 +
2542 + When ``skip_index_branch`` is true, sparse rows still keep their packed
2543 + ``[index_q | index_k]`` tail, but the kernel only processes the main q/k/v
2544 + branches and main KV cache. This is used by MiniMax-M3 index-topk reuse
2545 + layers that consume top-k block ids selected by an earlier sparse layer.
2546 """
2547 torch.ops._C.fused_minimax_m3_qknorm_rope_kv_insert(
2548 qkv,
@@ -2559,6 +2565,7 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
2565 q_out,
2566 index_q_out,
2567 kv_cache_dtype,
2568 + skip_index_branch,
2569 )
2570
2571
vllm/models/minimax_m3/amd/model.py
+214 -25
@@ -35,6 +35,7 @@ from vllm.config import (
35 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36 from vllm.forward_context import get_forward_context
37 from vllm.model_executor.layers.attention import Attention
38 +from vllm.model_executor.layers.attention.attention import set_default_quant_scales
39 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
40 from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
41 fused_allreduce_gemma_rms_norm,
@@ -93,6 +94,7 @@ from vllm.models.minimax_m3.common.mm_preprocess import (
94 from vllm.models.minimax_m3.common.sparse_attention import (
95 MiniMaxM3SparseBackend,
96 MiniMaxM3SparseImpl,
97 + minimax_m3_use_aiter_sparse_pa,
98 select_main_impl_cls,
99 )
100 from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
@@ -103,6 +105,7 @@ from vllm.v1.kv_cache_interface import (
105 FullAttentionSpec,
106 KVCacheSpec,
107 get_kv_quant_mode,
108 + is_quantized_kv_cache,
109 )
110
111
@@ -636,6 +639,11 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
639 self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
640 self.kv_cache_dtype, vllm_config.model_config
641 )
642 + # MiniMax-M3 sparse attention owns its KV-cache insert/read path instead
643 + # of wrapping the generic Attention module. Keep the same runtime scale
644 + # attributes so FP8 KV reads can honor vLLM's per-layer descale contract.
645 + self.calculate_kv_scales = False
646 + set_default_quant_scales(self, register_buffer=True)
647
648 # Shared top-k buffer: the indexer writes the selected blocks into it and
649 # the attend impl reads them back (no Python value crosses the break).
@@ -647,6 +655,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
655 self.impl: MiniMaxM3SparseImpl = select_main_impl_cls( # type: ignore[assignment]
656 topk_blocks=sparse_cfg["sparse_topk_blocks"],
657 kv_cache_dtype=self.kv_cache_dtype,
658 + num_kv_heads=self.num_kv_heads,
659 )(
660 self.num_heads,
661 self.head_dim,
@@ -656,6 +665,10 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
665 topk_blocks=sparse_cfg["sparse_topk_blocks"],
666 sparse_block_size=sparse_cfg["sparse_block_size"],
667 )
668 + self.use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(self.num_kv_heads)
669 + self.kv_cache_k = torch.tensor([])
670 + self.kv_cache_v = torch.tensor([])
671 + self._aiter_sparse_pa_cache_data_ptr = 0
672 # Self-contained nn.Module: owns its side cache, selects its impl in init
673 # (Triton on ROCm, where the SM100 gate is always False).
674 self.indexer = MiniMaxM3Indexer(
@@ -694,6 +707,91 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
707 kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
708 )
709
710 + def _ensure_aiter_sparse_pa_kv_cache(self) -> None:
711 + if self.kv_cache.numel() == 0:
712 + return
713 + if self._aiter_sparse_pa_cache_data_ptr == self.kv_cache.data_ptr():
714 + return
715 +
716 + kv_cache = self.kv_cache
717 + if is_quantized_kv_cache(self.kv_cache_dtype):
718 + kv_cache = kv_cache.view(self.impl.kv_cache_fp8_dtype)
719 + key_cache, value_cache = kv_cache.unbind(1)
720 + if not key_cache.is_contiguous() or not value_cache.is_contiguous():
721 + raise RuntimeError(
722 + "MiniMax-M3 AITER sparse PA requires K/V-separated KV cache "
723 + "storage. Set VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1 before "
724 + "initializing the engine."
725 + )
726 +
727 + x = 16 // key_cache.element_size()
728 + if self.head_dim % x != 0:
729 + raise RuntimeError(
730 + "MiniMax-M3 AITER sparse PA requires head_dim divisible by "
731 + f"16 / dtype_size, got head_dim={self.head_dim}, x={x}"
732 + )
733 + num_blocks = key_cache.shape[0]
734 + num_phys16 = num_blocks * 8
735 + self.kv_cache_k = key_cache.view(
736 + num_phys16,
737 + self.num_kv_heads,
738 + self.head_dim // x,
739 + 16,
740 + x,
741 + )
742 + self.kv_cache_v = value_cache.view(
743 + num_phys16,
744 + self.num_kv_heads,
745 + 16 // x,
746 + self.head_dim,
747 + x,
748 + )
749 + self._aiter_sparse_pa_cache_data_ptr = self.kv_cache.data_ptr()
750 +
751 + def get_aiter_sparse_pa_kv_cache(self) -> tuple[torch.Tensor, torch.Tensor]:
752 + self._ensure_aiter_sparse_pa_kv_cache()
753 + return self.kv_cache_k, self.kv_cache_v
754 +
755 + def _insert_aiter_sparse_pa_kv(
756 + self,
757 + k: torch.Tensor,
758 + v: torch.Tensor,
759 + index_k: torch.Tensor | None,
760 + slot_mapping: torch.Tensor,
761 + index_slot_mapping: torch.Tensor | None,
762 + ) -> None:
763 + if self.kv_cache.numel() == 0:
764 + return
765 + from aiter import reshape_and_cache
766 +
767 + from vllm.models.minimax_m3.amd.ops.sparse_pa import (
768 + minimax_m3_insert_index_cache,
769 + )
770 +
771 + key_cache, value_cache = self.get_aiter_sparse_pa_kv_cache()
772 + kv_cache_dtype = (
773 + self.kv_cache_dtype
774 + if is_quantized_kv_cache(self.kv_cache_dtype)
775 + else "auto"
776 + )
777 + reshape_and_cache(
778 + k.contiguous(),
779 + v.contiguous(),
780 + key_cache,
781 + value_cache,
782 + slot_mapping,
783 + kv_cache_dtype=kv_cache_dtype,
784 + k_scale=getattr(self, "_k_scale", None),
785 + v_scale=getattr(self, "_v_scale", None),
786 + asm_layout=True,
787 + )
788 + if index_k is None or index_slot_mapping is None:
789 + return
790 + index_cache = self.indexer.index_cache.kv_cache
791 + if index_cache.numel() == 0:
792 + return
793 + minimax_m3_insert_index_cache(index_k, index_cache, index_slot_mapping)
794 +
795 def forward(
796 self,
797 positions: torch.Tensor,
@@ -724,31 +822,121 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
822 return qkv.new_zeros((num_tokens, self.hidden_size))
823
824 main_slot_mapping = fwd_slot_mapping[self.layer_name]
727 - index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
825 q = qkv.new_empty((num_tokens, self.q_size))
729 - index_q = qkv.new_empty((num_tokens, self.index_q_size))
730 - ops.fused_minimax_m3_qknorm_rope_kv_insert(
731 - qkv,
732 - self.q_norm.weight,
733 - self.k_norm.weight,
734 - cos_sin_cache,
735 - positions,
736 - self.num_heads,
737 - self.num_kv_heads,
738 - rotary_dim,
739 - eps,
740 - self.index_q_norm.weight,
741 - self.index_k_norm.weight,
742 - self.num_idx_heads,
743 - main_slot_mapping,
744 - index_slot_mapping,
745 - self.kv_cache,
746 - self.indexer.index_cache.kv_cache,
747 - self.kv_cache.size(2), # paged-cache block size
748 - q,
749 - index_q,
750 - self.kv_cache_dtype,
751 - )
826 + if self.skip_index_topk:
827 + index_q = None
828 + if self.use_aiter_sparse_pa:
829 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
830 + qkv,
831 + self.q_norm.weight,
832 + self.k_norm.weight,
833 + cos_sin_cache,
834 + positions,
835 + self.num_heads,
836 + self.num_kv_heads,
837 + rotary_dim,
838 + eps,
839 + num_index_heads=self.num_idx_heads,
840 + q_out=q,
841 + skip_index_branch=True,
842 + )
843 + k_start = self.q_size
844 + v_start = k_start + self.kv_size
845 + k = qkv[:, k_start:v_start].view(
846 + num_tokens, self.num_kv_heads, self.head_dim
847 + )
848 + v = qkv[:, v_start : v_start + self.kv_size].view(
849 + num_tokens, self.num_kv_heads, self.head_dim
850 + )
851 + self._insert_aiter_sparse_pa_kv(
852 + k,
853 + v,
854 + None,
855 + main_slot_mapping,
856 + None,
857 + )
858 + else:
859 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
860 + qkv,
861 + self.q_norm.weight,
862 + self.k_norm.weight,
863 + cos_sin_cache,
864 + positions,
865 + self.num_heads,
866 + self.num_kv_heads,
867 + rotary_dim,
868 + eps,
869 + num_index_heads=self.num_idx_heads,
870 + slot_mapping=main_slot_mapping,
871 + kv_cache=self.kv_cache,
872 + block_size=self.kv_cache.size(2), # paged-cache block size
873 + q_out=q,
874 + kv_cache_dtype=self.kv_cache_dtype,
875 + skip_index_branch=True,
876 + )
877 + else:
878 + index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
879 + index_q = qkv.new_empty((num_tokens, self.index_q_size))
880 + if self.use_aiter_sparse_pa:
881 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
882 + qkv,
883 + self.q_norm.weight,
884 + self.k_norm.weight,
885 + cos_sin_cache,
886 + positions,
887 + self.num_heads,
888 + self.num_kv_heads,
889 + rotary_dim,
890 + eps,
891 + self.index_q_norm.weight,
892 + self.index_k_norm.weight,
893 + self.num_idx_heads,
894 + q_out=q,
895 + index_q_out=index_q,
896 + kv_cache_dtype=self.kv_cache_dtype,
897 + )
898 + k_start = self.q_size
899 + v_start = k_start + self.kv_size
900 + index_k_start = v_start + self.kv_size + self.index_q_size
901 + k = qkv[:, k_start:v_start].view(
902 + num_tokens, self.num_kv_heads, self.head_dim
903 + )
904 + v = qkv[:, v_start : v_start + self.kv_size].view(
905 + num_tokens, self.num_kv_heads, self.head_dim
906 + )
907 + index_k = qkv[
908 + :, index_k_start : index_k_start + self.idx_head_dim
909 + ].view(num_tokens, self.idx_head_dim)
910 + self._insert_aiter_sparse_pa_kv(
911 + k,
912 + v,
913 + index_k,
914 + main_slot_mapping,
915 + index_slot_mapping,
916 + )
917 + else:
918 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
919 + qkv,
920 + self.q_norm.weight,
921 + self.k_norm.weight,
922 + cos_sin_cache,
923 + positions,
924 + self.num_heads,
925 + self.num_kv_heads,
926 + rotary_dim,
927 + eps,
928 + self.index_q_norm.weight,
929 + self.index_k_norm.weight,
930 + self.num_idx_heads,
931 + main_slot_mapping,
932 + index_slot_mapping,
933 + self.kv_cache,
934 + self.indexer.index_cache.kv_cache,
935 + self.kv_cache.size(2), # paged-cache block size
936 + q,
937 + index_q,
938 + self.kv_cache_dtype,
939 + )
940
941 output = torch.empty_like(q)
942 attn_output = self._run_attention(q, index_q, output)
@@ -759,7 +947,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
947 def _run_attention(
948 self,
949 query: torch.Tensor,
762 - index_query: torch.Tensor,
950 + index_query: torch.Tensor | None,
951 output: torch.Tensor,
952 ) -> torch.Tensor:
953 # Single eager break around both: their split-K kernels read per-request
@@ -768,6 +956,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
956 # When skip_index_topk is set (ATOM index_topk_freq), reuse the selection
957 # the preceding compute layer wrote into the shared buffer this forward.
958 if not self.skip_index_topk:
959 + assert index_query is not None
960 self.indexer(index_query)
961 return self.impl.forward(self, query, self.kv_cache, output)
962
vllm/models/minimax_m3/amd/ops/sparse_attn.py
+65 -8
@@ -12,7 +12,9 @@ import torch
12
13 from vllm.models.minimax_m3.common.ops.sparse_attn import (
14 _FP8_DTYPES,
15 + _KV_SCALE_NONE,
16 SPARSE_BLOCK_SIZE,
17 + _kv_scale_args,
18 minimax_m3_sparse_attn_decode,
19 )
20 from vllm.platforms.rocm import on_gfx950, on_mi3xx
@@ -70,7 +72,9 @@ def _sparse_attn_prefill_kwargs() -> dict:
72 @triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"])
73 def _gqa_sparse_fwd_kernel(
74 q_ptr, # [total_q, num_heads, head_dim]
73 - kv_cache_ptr, # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim]
75 + kv_cache_ptr, # main cache: [num_blocks, num_kv_heads, 128, 2*head_dim]
76 + k_scale_ptr,
77 + v_scale_ptr,
78 t_ptr, # topk_idx: [num_kv_heads, total_q, topk]
79 o_ptr, # [total_q, num_heads, head_dim]
80 block_table_ptr, # [num_reqs, max_blocks]
@@ -88,10 +92,13 @@ def _gqa_sparse_fwd_kernel(
92 stride_qh,
93 stride_qd,
94 stride_kv_blk,
91 - stride_kv_kv,
92 - stride_kv_pos,
95 stride_kv_h,
96 + stride_kv_pos,
97 stride_kv_d,
98 + stride_ks_h,
99 + stride_ks_t,
100 + stride_vs_h,
101 + stride_vs_t,
102 stride_th,
103 stride_tn,
104 stride_tk,
@@ -106,6 +113,7 @@ def _gqa_sparse_fwd_kernel(
113 BLOCK_SIZE_T: tl.constexpr,
114 BLOCK_SIZE_QH: tl.constexpr,
115 USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load
116 + KV_SCALE_MODE: tl.constexpr, # 0: none, 1: scalar, 2: [kv_head, token]
117 SUB_K: tl.constexpr, # CDNA only: KV sub-tile width (see _IS_MI3XX)
118 ):
119 sm_scale_log2e = sm_scale * 1.4426950409
@@ -161,7 +169,6 @@ def _gqa_sparse_fwd_kernel(
169 pos_mask_sub = pos_sub < seq_len
170 k_sub = tl.load(
171 kv_base
164 - + 0 * stride_kv_kv
172 + off_sub[None, :] * stride_kv_pos
173 + off_d[:, None] * stride_kv_d,
174 mask=d_mask[:, None] & pos_mask_sub[None, :],
@@ -169,6 +176,17 @@ def _gqa_sparse_fwd_kernel(
176 )
177 if USE_FP8:
178 k_sub = k_sub.to(q.dtype)
179 + if KV_SCALE_MODE == 1:
180 + k_sub = (k_sub * tl.load(k_scale_ptr)).to(q.dtype)
181 + elif KV_SCALE_MODE == 2:
182 + k_scale = tl.load(
183 + k_scale_ptr
184 + + pid_kh * stride_ks_h
185 + + (page * BLOCK_SIZE_K + off_sub) * stride_ks_t,
186 + mask=pos_mask_sub,
187 + other=1.0,
188 + )
189 + k_sub = (k_sub * k_scale[None, :]).to(q.dtype)
190 off_q_sub = (
191 tl.arange(0, BLOCK_SIZE_Q)[:, None]
192 + pid_q_j * BLOCK_SIZE_Q
@@ -187,14 +205,24 @@ def _gqa_sparse_fwd_kernel(
205 acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None]
206 v_sub = tl.load(
207 kv_base
190 - + 1 * stride_kv_kv
208 + off_sub[:, None] * stride_kv_pos
192 - + off_d[None, :] * stride_kv_d,
209 + + (head_dim + off_d[None, :]) * stride_kv_d,
210 mask=pos_mask_sub[:, None] & d_mask[None, :],
211 other=0.0,
212 )
213 if USE_FP8:
214 v_sub = v_sub.to(q.dtype)
215 + if KV_SCALE_MODE == 1:
216 + v_sub = (v_sub * tl.load(v_scale_ptr)).to(q.dtype)
217 + elif KV_SCALE_MODE == 2:
218 + v_scale = tl.load(
219 + v_scale_ptr
220 + + pid_kh * stride_vs_h
221 + + (page * BLOCK_SIZE_K + off_sub) * stride_vs_t,
222 + mask=pos_mask_sub,
223 + other=1.0,
224 + )
225 + v_sub = (v_sub * v_scale[:, None]).to(q.dtype)
226 acc_o += tl.dot(p_sub.to(v_sub.dtype), v_sub)
227 m_i = m_ij
228 lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
@@ -214,7 +242,7 @@ def _gqa_sparse_fwd_kernel(
242 @torch.no_grad()
243 def minimax_m3_sparse_attn(
244 q: torch.Tensor, # [total_q, num_heads, head_dim]
217 - kv_cache: torch.Tensor, # [num_blocks, 2, 128, num_kv_heads, head_dim]
245 + kv_cache: torch.Tensor, # [num_blocks, num_kv_heads, 128, 2*head_dim]
246 topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk]
247 block_table: torch.Tensor, # [batch, max_blocks]
248 cu_seqlens_q: torch.Tensor, # [batch+1] int32
@@ -224,6 +252,8 @@ def minimax_m3_sparse_attn(
252 num_kv_heads: int,
253 sm_scale: float,
254 output: torch.Tensor, # [total_q, num_heads, head_dim]
255 + k_scale: torch.Tensor | None = None,
256 + v_scale: torch.Tensor | None = None,
257 ) -> None:
258 """GQA block-sparse attention over the selected blocks. block_size_q == 1."""
259 total_q, num_heads, head_dim = q.shape
@@ -231,10 +261,33 @@ def minimax_m3_sparse_attn(
261 topk = topk_idx.shape[-1]
262 gqa_group_size = num_heads // num_kv_heads
263 use_fp8 = kv_cache.dtype in _FP8_DTYPES
264 + (
265 + k_scale_arg,
266 + v_scale_arg,
267 + stride_ks_h,
268 + stride_ks_t,
269 + stride_vs_h,
270 + stride_vs_t,
271 + kv_scale_mode,
272 + ) = (
273 + _kv_scale_args(output, num_kv_heads, k_scale, v_scale)
274 + if use_fp8
275 + else (
276 + output,
277 + output,
278 + 0,
279 + 0,
280 + 0,
281 + 0,
282 + _KV_SCALE_NONE,
283 + )
284 + )
285 grid = (max_query_len, num_kv_heads, batch)
286 _gqa_sparse_fwd_kernel[grid](
287 q,
288 kv_cache,
289 + k_scale_arg,
290 + v_scale_arg,
291 topk_idx,
292 output,
293 block_table,
@@ -255,7 +308,10 @@ def minimax_m3_sparse_attn(
308 kv_cache.stride(1),
309 kv_cache.stride(2),
310 kv_cache.stride(3),
258 - kv_cache.stride(4),
311 + stride_ks_h,
312 + stride_ks_t,
313 + stride_vs_h,
314 + stride_vs_t,
315 topk_idx.stride(0),
316 topk_idx.stride(1),
317 topk_idx.stride(2),
@@ -266,6 +322,7 @@ def minimax_m3_sparse_attn(
322 BLOCK_SIZE_Q=1,
323 BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
324 USE_FP8=use_fp8,
325 + KV_SCALE_MODE=kv_scale_mode,
326 SUB_K=_SPARSE_ATTN_SUB_K,
327 **_sparse_attn_prefill_kwargs(),
328 )
vllm/models/minimax_m3/amd/ops/sparse_pa.py new
+466
@@ -0,0 +1,466 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""AITER page-16 sparse paged-attention helpers for MiniMax-M3 on ROCm."""
4 +
5 +import torch
6 +
7 +try:
8 + from vllm.triton_utils import tl, triton
9 +except ModuleNotFoundError:
10 + import triton
11 + import triton.language as tl
12 +
13 +from vllm.models.minimax_m3.common.ops.sparse_attn import SPARSE_BLOCK_SIZE
14 +
15 +ASM_PAGE_SIZE = 16
16 +PAGES_PER_SPARSE_BLOCK = SPARSE_BLOCK_SIZE // ASM_PAGE_SIZE
17 +
18 +_FP8_DTYPES = {
19 + dtype
20 + for dtype in (
21 + getattr(torch, "float8_e4m3fn", None),
22 + getattr(torch, "float8_e4m3fnuz", None),
23 + getattr(torch, "float8_e5m2", None),
24 + getattr(torch, "float8_e5m2fnuz", None),
25 + )
26 + if dtype is not None
27 +}
28 +
29 +
30 +def _is_fp8_kv_cache_tensor(kv_cache: torch.Tensor) -> bool:
31 + return kv_cache.dtype in _FP8_DTYPES
32 +
33 +
34 +@triton.jit
35 +def _build_sparse_block_table_kernel(
36 + topk_ptr, # [1, batch, topk] int32, selected logical 128-block ids
37 + block_table_ptr, # [batch, max_blocks] int32, logical 128-page table
38 + seq_lens_ptr, # [batch] int32
39 + sparse_bt_ptr, # [batch, topk * 8] int32, physical 16-page table
40 + sparse_ctx_ptr, # [batch] int32
41 + max_topk,
42 + stride_topk_n,
43 + stride_topk_k,
44 + stride_bt_b,
45 + stride_sbt_b,
46 + SPARSE_BLOCK_SIZE_C: tl.constexpr,
47 + PAGES_PER_BLOCK: tl.constexpr,
48 + BLOCK_SIZE_T: tl.constexpr,
49 +):
50 + pid_b = tl.program_id(0)
51 + seq_len = tl.load(seq_lens_ptr + pid_b)
52 + last_blk = (seq_len - 1) // SPARSE_BLOCK_SIZE_C
53 +
54 + topk_row = topk_ptr + pid_b * stride_topk_n
55 + bt_row = block_table_ptr + pid_b * stride_bt_b
56 + sbt_row = sparse_bt_ptr + pid_b * stride_sbt_b
57 +
58 + off_t = tl.arange(0, BLOCK_SIZE_T)
59 + blk = tl.load(topk_row + off_t * stride_topk_k, mask=off_t < max_topk, other=-1)
60 + valid = blk >= 0
61 + is_tail = valid & (blk == last_blk)
62 + is_full = valid & (blk != last_blk)
63 +
64 + n_full = tl.sum(is_full.to(tl.int32), axis=0)
65 + n_valid = tl.sum(valid.to(tl.int32), axis=0)
66 + earlier_full = tl.cumsum(is_full.to(tl.int32), axis=0) - is_full.to(tl.int32)
67 + slot = tl.where(is_full, earlier_full, n_full)
68 +
69 + logical_page = tl.load(bt_row + blk, mask=valid, other=0).to(tl.int32)
70 + base_phys = logical_page * PAGES_PER_BLOCK
71 + dst_base = slot * PAGES_PER_BLOCK
72 +
73 + for j in tl.static_range(PAGES_PER_BLOCK):
74 + tl.store(sbt_row + dst_base + j, base_phys + j, mask=valid)
75 +
76 + n_used = n_valid * PAGES_PER_BLOCK
77 + off_w = tl.arange(0, BLOCK_SIZE_T * PAGES_PER_BLOCK)
78 + row_width = max_topk * PAGES_PER_BLOCK
79 + tl.store(
80 + sbt_row + off_w,
81 + tl.zeros_like(off_w),
82 + mask=(off_w >= n_used) & (off_w < row_width),
83 + )
84 +
85 + tail_tokens = seq_len - last_blk * SPARSE_BLOCK_SIZE_C
86 + has_tail = tl.sum(is_tail.to(tl.int32), axis=0) > 0
87 + ctx = n_full * SPARSE_BLOCK_SIZE_C + tl.where(has_tail, tail_tokens, 0)
88 + ctx = tl.where(has_tail, ctx, tl.minimum(n_valid * SPARSE_BLOCK_SIZE_C, seq_len))
89 + tl.store(sparse_ctx_ptr + pid_b, ctx)
90 +
91 +
92 +@torch.no_grad()
93 +def minimax_m3_build_sparse_block_table(
94 + topk_idx: torch.Tensor, # [1, batch, topk]
95 + block_table: torch.Tensor, # [batch, max_blocks]
96 + seq_lens: torch.Tensor, # [batch]
97 +) -> tuple[torch.Tensor, torch.Tensor]:
98 + """Compact selected logical sparse blocks into physical page-16 tables."""
99 + assert topk_idx.shape[0] == 1, "AITER sparse PA requires num_kv_heads == 1"
100 + batch = topk_idx.shape[1]
101 + topk = topk_idx.shape[-1]
102 + width = topk * PAGES_PER_SPARSE_BLOCK
103 + sparse_bt = torch.empty((batch, width), dtype=torch.int32, device=topk_idx.device)
104 + sparse_ctx = torch.empty((batch,), dtype=torch.int32, device=topk_idx.device)
105 + _build_sparse_block_table_kernel[(batch,)](
106 + topk_idx,
107 + block_table,
108 + seq_lens,
109 + sparse_bt,
110 + sparse_ctx,
111 + topk,
112 + topk_idx.stride(1),
113 + topk_idx.stride(2),
114 + block_table.stride(0),
115 + sparse_bt.stride(0),
116 + SPARSE_BLOCK_SIZE_C=SPARSE_BLOCK_SIZE,
117 + PAGES_PER_BLOCK=PAGES_PER_SPARSE_BLOCK,
118 + BLOCK_SIZE_T=triton.next_power_of_2(topk),
119 + )
120 + return sparse_bt, sparse_ctx
121 +
122 +
123 +@triton.jit
124 +def _build_sparse_block_table_prefill_kernel(
125 + topk_ptr, # [1, total_q, topk]
126 + block_table_ptr, # [batch, max_blocks]
127 + req_id_ptr, # [total_q]
128 + abs_pos_ptr, # [total_q]
129 + sparse_bt_ptr, # [total_q, topk * 8]
130 + sparse_ctx_ptr, # [total_q]
131 + max_topk,
132 + stride_topk_n,
133 + stride_topk_k,
134 + stride_bt_b,
135 + stride_sbt_n,
136 + SPARSE_BLOCK_SIZE_C: tl.constexpr,
137 + PAGES_PER_BLOCK: tl.constexpr,
138 + BLOCK_SIZE_T: tl.constexpr,
139 +):
140 + pid_n = tl.program_id(0)
141 + req_id = tl.load(req_id_ptr + pid_n)
142 + abs_pos = tl.load(abs_pos_ptr + pid_n)
143 + causal_len = abs_pos + 1
144 + self_blk = abs_pos // SPARSE_BLOCK_SIZE_C
145 +
146 + topk_row = topk_ptr + pid_n * stride_topk_n
147 + bt_row = block_table_ptr + req_id * stride_bt_b
148 + sbt_row = sparse_bt_ptr + pid_n * stride_sbt_n
149 +
150 + off_t = tl.arange(0, BLOCK_SIZE_T)
151 + blk = tl.load(topk_row + off_t * stride_topk_k, mask=off_t < max_topk, other=-1)
152 + valid = (blk >= 0) & (blk <= self_blk)
153 + is_tail = valid & (blk == self_blk)
154 + is_full = valid & (blk < self_blk)
155 +
156 + n_full = tl.sum(is_full.to(tl.int32), axis=0)
157 + n_valid = tl.sum(valid.to(tl.int32), axis=0)
158 + earlier_full = tl.cumsum(is_full.to(tl.int32), axis=0) - is_full.to(tl.int32)
159 + slot = tl.where(is_full, earlier_full, n_full)
160 +
161 + logical_page = tl.load(bt_row + blk, mask=valid, other=0).to(tl.int32)
162 + base_phys = logical_page * PAGES_PER_BLOCK
163 + dst_base = slot * PAGES_PER_BLOCK
164 +
165 + for j in tl.static_range(PAGES_PER_BLOCK):
166 + tl.store(sbt_row + dst_base + j, base_phys + j, mask=valid)
167 +
168 + n_used = n_valid * PAGES_PER_BLOCK
169 + off_w = tl.arange(0, BLOCK_SIZE_T * PAGES_PER_BLOCK)
170 + row_width = max_topk * PAGES_PER_BLOCK
171 + tl.store(
172 + sbt_row + off_w,
173 + tl.zeros_like(off_w),
174 + mask=(off_w >= n_used) & (off_w < row_width),
175 + )
176 +
177 + tail_tokens = causal_len - self_blk * SPARSE_BLOCK_SIZE_C
178 + has_tail = tl.sum(is_tail.to(tl.int32), axis=0) > 0
179 + ctx = n_full * SPARSE_BLOCK_SIZE_C + tl.where(has_tail, tail_tokens, 0)
180 + ctx = tl.where(has_tail, ctx, tl.minimum(n_valid * SPARSE_BLOCK_SIZE_C, causal_len))
181 + tl.store(sparse_ctx_ptr + pid_n, ctx)
182 +
183 +
184 +@torch.no_grad()
185 +def minimax_m3_build_sparse_block_table_prefill(
186 + topk_idx: torch.Tensor, # [1, total_q, topk]
187 + block_table: torch.Tensor, # [batch, max_blocks]
188 + query_req_id: torch.Tensor, # [total_q]
189 + query_abs_pos: torch.Tensor, # [total_q]
190 +) -> tuple[torch.Tensor, torch.Tensor]:
191 + """Build one page-16 sparse block table row per prefill query token."""
192 + assert topk_idx.shape[0] == 1, "AITER sparse PA requires num_kv_heads == 1"
193 + total_q = topk_idx.shape[1]
194 + topk = topk_idx.shape[-1]
195 + width = topk * PAGES_PER_SPARSE_BLOCK
196 + sparse_bt = torch.empty((total_q, width), dtype=torch.int32, device=topk_idx.device)
197 + sparse_ctx = torch.empty((total_q,), dtype=torch.int32, device=topk_idx.device)
198 + _build_sparse_block_table_prefill_kernel[(total_q,)](
199 + topk_idx,
200 + block_table,
201 + query_req_id,
202 + query_abs_pos,
203 + sparse_bt,
204 + sparse_ctx,
205 + topk,
206 + topk_idx.stride(1),
207 + topk_idx.stride(2),
208 + block_table.stride(0),
209 + sparse_bt.stride(0),
210 + SPARSE_BLOCK_SIZE_C=SPARSE_BLOCK_SIZE,
211 + PAGES_PER_BLOCK=PAGES_PER_SPARSE_BLOCK,
212 + BLOCK_SIZE_T=triton.next_power_of_2(topk),
213 + )
214 + return sparse_bt, sparse_ctx
215 +
216 +
217 +@triton.jit
218 +def _insert_index_cache_kernel(
219 + index_k_ptr, # [num_tokens, head_dim]
220 + index_cache_ptr, # [num_blocks, block_size, head_dim]
221 + slot_mapping_ptr, # [num_tokens]
222 + stride_k_t,
223 + stride_k_d,
224 + stride_c_b,
225 + stride_c_t,
226 + stride_c_d,
227 + stride_slot_t,
228 + CACHE_BLOCK_SIZE: tl.constexpr,
229 + HEAD_DIM: tl.constexpr,
230 + BLOCK_D: tl.constexpr,
231 +):
232 + pid_t = tl.program_id(0)
233 + offs_d = tl.arange(0, BLOCK_D)
234 + slot = tl.load(slot_mapping_ptr + pid_t * stride_slot_t)
235 + block_id = slot // CACHE_BLOCK_SIZE
236 + block_off = slot - block_id * CACHE_BLOCK_SIZE
237 +
238 + src = index_k_ptr + pid_t * stride_k_t + offs_d * stride_k_d
239 + dst = (
240 + index_cache_ptr
241 + + block_id * stride_c_b
242 + + block_off * stride_c_t
243 + + offs_d * stride_c_d
244 + )
245 + mask = (slot >= 0) & (offs_d < HEAD_DIM)
246 + value = tl.load(src, mask=offs_d < HEAD_DIM, other=0.0)
247 + tl.store(dst, value, mask=mask)
248 +
249 +
250 +@torch.no_grad()
251 +def minimax_m3_insert_index_cache(
252 + index_k: torch.Tensor,
253 + index_cache: torch.Tensor,
254 + index_slot_mapping: torch.Tensor,
255 +) -> None:
256 + """Scatter index keys into MiniMax-M3's key-only side cache."""
257 + if index_k.numel() == 0 or index_cache.numel() == 0:
258 + return
259 + if index_k.dim() != 2 or index_cache.dim() != 3:
260 + raise ValueError("MiniMax-M3 index cache insert expects [N,D] and [B,T,D]")
261 + if index_k.shape[1] != index_cache.shape[2]:
262 + raise ValueError("MiniMax-M3 index key dim must match index cache head dim")
263 + if index_slot_mapping.dim() != 1 or index_slot_mapping.shape[0] != index_k.shape[0]:
264 + raise ValueError("MiniMax-M3 index slot mapping must be a length-N vector")
265 + if index_cache.stride(2) != 1:
266 + raise ValueError("MiniMax-M3 index cache requires contiguous head dimension")
267 +
268 + head_dim = index_k.shape[1]
269 + _insert_index_cache_kernel[(index_k.shape[0],)](
270 + index_k,
271 + index_cache,
272 + index_slot_mapping,
273 + index_k.stride(0),
274 + index_k.stride(1),
275 + index_cache.stride(0),
276 + index_cache.stride(1),
277 + index_cache.stride(2),
278 + index_slot_mapping.stride(0),
279 + CACHE_BLOCK_SIZE=index_cache.shape[1],
280 + HEAD_DIM=head_dim,
281 + BLOCK_D=triton.next_power_of_2(head_dim),
282 + num_warps=4,
283 + )
284 +
285 +
286 +def _gluon_scale_arg(
287 + scale: torch.Tensor | None,
288 + *,
289 + num_phys_pages: int,
290 + num_kv_heads: int,
291 +) -> torch.Tensor | None:
292 + if scale is None:
293 + return None
294 + if scale.numel() == 1:
295 + return scale
296 + if scale.dim() != 2 or scale.shape[0] != num_kv_heads:
297 + raise ValueError(
298 + "MiniMax-M3 AITER sparse PA supports scalar KV scales or "
299 + "[num_kv_heads, max_kv_tokens] scales"
300 + )
301 + max_tokens = num_phys_pages * ASM_PAGE_SIZE
302 + if scale.shape[1] < max_tokens:
303 + raise ValueError(
304 + "MiniMax-M3 AITER sparse PA KV scale tensor is smaller than the "
305 + f"cache token capacity ({scale.shape[1]} < {max_tokens})"
306 + )
307 + scale = scale[:, :max_tokens]
308 + return (
309 + scale.transpose(0, 1)
310 + .contiguous()
311 + .view(num_phys_pages, ASM_PAGE_SIZE, num_kv_heads)
312 + .permute(0, 2, 1)
313 + .contiguous()
314 + .unsqueeze(-1)
315 + )
316 +
317 +
318 +@torch.no_grad()
319 +def minimax_m3_sparse_attn_decode_aiter(
320 + q: torch.Tensor, # [batch, num_heads, head_dim]
321 + k_cache: torch.Tensor, # [phys16, num_kv_heads, head_dim // x, 16, x]
322 + v_cache: torch.Tensor, # [phys16, num_kv_heads, 16 // x, head_dim, x]
323 + topk_idx: torch.Tensor, # [1, batch, topk]
324 + block_table: torch.Tensor, # [batch, max_blocks]
325 + seq_lens: torch.Tensor, # [batch]
326 + num_kv_heads: int,
327 + sm_scale: float,
328 + output: torch.Tensor, # [batch, num_heads, head_dim]
329 + k_scale: torch.Tensor | None = None,
330 + v_scale: torch.Tensor | None = None,
331 +) -> None:
332 + sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table(
333 + topk_idx, block_table, seq_lens
334 + )
335 + _run_gluon_decode(
336 + q,
337 + k_cache,
338 + v_cache,
339 + sparse_bt,
340 + sparse_ctx,
341 + num_kv_heads,
342 + sm_scale,
343 + output,
344 + k_scale,
345 + v_scale,
346 + )
347 +
348 +
349 +@torch.no_grad()
350 +def minimax_m3_sparse_attn_prefill_aiter(
351 + q: torch.Tensor, # [total_q, num_heads, head_dim]
352 + k_cache: torch.Tensor,
353 + v_cache: torch.Tensor,
354 + topk_idx: torch.Tensor, # [1, total_q, topk]
355 + block_table: torch.Tensor,
356 + cu_seqlens_q: torch.Tensor,
357 + prefix_lens: torch.Tensor,
358 + num_kv_heads: int,
359 + sm_scale: float,
360 + output: torch.Tensor,
361 + k_scale: torch.Tensor | None = None,
362 + v_scale: torch.Tensor | None = None,
363 +) -> None:
364 + total_q = q.shape[0]
365 + pos = torch.arange(total_q, dtype=torch.int32, device=q.device)
366 + query_req_id = torch.searchsorted(cu_seqlens_q[1:].contiguous(), pos, right=True)
367 + query_req_id = query_req_id.to(torch.int32)
368 + query_abs_pos = (prefix_lens[query_req_id] + (pos - cu_seqlens_q[query_req_id])).to(
369 + torch.int32
370 + )
371 + sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_prefill(
372 + topk_idx, block_table, query_req_id, query_abs_pos
373 + )
374 + _run_gluon_decode(
375 + q,
376 + k_cache,
377 + v_cache,
378 + sparse_bt,
379 + sparse_ctx,
380 + num_kv_heads,
381 + sm_scale,
382 + output,
383 + k_scale,
384 + v_scale,
385 + )
386 +
387 +
388 +def _run_gluon_decode(
389 + q: torch.Tensor,
390 + k_cache: torch.Tensor,
391 + v_cache: torch.Tensor,
392 + sparse_bt: torch.Tensor,
393 + sparse_ctx: torch.Tensor,
394 + num_kv_heads: int,
395 + sm_scale: float,
396 + output: torch.Tensor,
397 + k_scale: torch.Tensor | None,
398 + v_scale: torch.Tensor | None,
399 +) -> None:
400 + from aiter import dtypes as aiter_dtypes
401 + from aiter.ops.triton.gluon.pa_decode_gluon import (
402 + get_recommended_splits,
403 + pa_decode_gluon,
404 + )
405 +
406 + if not q.is_contiguous():
407 + q = q.contiguous()
408 + if not output.is_contiguous():
409 + raise ValueError("MiniMax-M3 AITER sparse PA output must be contiguous")
410 +
411 + total_q, num_q_heads, head_size = q.shape
412 + if head_size != 128:
413 + raise ValueError("MiniMax-M3 AITER sparse PA requires head_dim == 128")
414 + group_size = num_q_heads // num_kv_heads
415 +
416 + nphys16, hkv = k_cache.shape[0], k_cache.shape[1]
417 + k_cache_view = k_cache.view(nphys16 * hkv, 1, *k_cache.shape[2:])
418 + v_cache_view = v_cache.view(nphys16 * hkv, 1, *v_cache.shape[2:])
419 + q_view = q.view(total_q * num_kv_heads, group_size, head_size)
420 + out_view = output.view(total_q * num_kv_heads, group_size, head_size)
421 +
422 + num_seqs = total_q * num_kv_heads
423 + max_context_partition_num = get_recommended_splits(num_seqs, 1)
424 + context_partition_size = 256
425 + intermediate_shape = (num_seqs, 1, max_context_partition_num, group_size)
426 + exp_sums = torch.empty(intermediate_shape, dtype=torch.float32, device=q.device)
427 + max_logits = torch.empty_like(exp_sums)
428 + temporary_output = torch.empty(
429 + *intermediate_shape, head_size, dtype=q.dtype, device=q.device
430 + )
431 +
432 + is_fp8 = _is_fp8_kv_cache_tensor(k_cache)
433 + compute_type = aiter_dtypes.fp8 if is_fp8 else q.dtype
434 + if is_fp8:
435 + k_scale_arg = _gluon_scale_arg(
436 + k_scale, num_phys_pages=nphys16, num_kv_heads=hkv
437 + )
438 + v_scale_arg = _gluon_scale_arg(
439 + v_scale, num_phys_pages=nphys16, num_kv_heads=hkv
440 + )
441 + else:
442 + k_scale_arg = v_scale_arg = None
443 +
444 + pa_decode_gluon(
445 + output=out_view,
446 + query=q_view,
447 + key_cache=k_cache_view,
448 + value_cache=v_cache_view,
449 + context_lengths=sparse_ctx,
450 + block_tables=sparse_bt,
451 + softmax_scale=sm_scale,
452 + query_length=1,
453 + max_context_partition_num=max_context_partition_num,
454 + context_partition_size=context_partition_size,
455 + compute_type=compute_type,
456 + query_scale=None,
457 + key_scale=k_scale_arg,
458 + value_scale=v_scale_arg,
459 + exp_sums=exp_sums,
460 + max_logits=max_logits,
461 + temporary_output=temporary_output,
462 + alibi_slopes=None,
463 + sinks=None,
464 + sliding_window=-1,
465 + ps=True,
466 + )
vllm/models/minimax_m3/amd/sparse_attention_msa.py new
+94
@@ -0,0 +1,94 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MSA AITER block-sparse attend for MiniMax M3."""
4 +
5 +import torch
6 +
7 +from vllm.forward_context import get_forward_context
8 +from vllm.models.minimax_m3.common.sparse_attention import (
9 + MiniMaxM3SparseImpl,
10 + MiniMaxM3SparseMetadata,
11 +)
12 +from vllm.v1.attention.backend import (
13 + AttentionLayer,
14 +)
15 +
16 +
17 +class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
18 + """ROCm AITER page-16 SHUFFLE sparse paged attention."""
19 +
20 + def forward(
21 + self,
22 + layer: AttentionLayer,
23 + query: torch.Tensor,
24 + kv_cache: torch.Tensor,
25 + output: torch.Tensor,
26 + ) -> torch.Tensor:
27 + from vllm.models.minimax_m3.amd.ops.sparse_pa import (
28 + minimax_m3_sparse_attn_decode_aiter,
29 + minimax_m3_sparse_attn_prefill_aiter,
30 + )
31 +
32 + attn_metadata = get_forward_context().attn_metadata
33 + if not isinstance(attn_metadata, dict):
34 + return output
35 + main_md = attn_metadata[layer.layer_name] # type: ignore[attr-defined]
36 + assert isinstance(main_md, MiniMaxM3SparseMetadata)
37 +
38 + nd = main_md.num_decode_tokens
39 + num_tokens = main_md.num_actual_tokens
40 + topk = layer.topk_indices_buffer # type: ignore[attr-defined]
41 + assert topk is not None
42 + if self.num_kv_heads != 1:
43 + raise NotImplementedError(
44 + "MiniMax-M3 AITER sparse PA currently requires per-rank "
45 + f"num_kv_heads == 1, got {self.num_kv_heads}"
46 + )
47 +
48 + hd = self.head_size
49 + q = query[:num_tokens].view(-1, self.num_heads, hd)
50 + out = output[:num_tokens].view(-1, self.num_heads, hd)
51 + k_cache, v_cache = layer.get_aiter_sparse_pa_kv_cache() # type: ignore[attr-defined]
52 + k_scale = getattr(layer, "_k_scale", None) if self.use_fp8_kv else None
53 + v_scale = getattr(layer, "_v_scale", None) if self.use_fp8_kv else None
54 +
55 + if main_md.num_decodes > 0:
56 + d = main_md.decode
57 + assert d is not None
58 + if d.decode_query_len != 1:
59 + raise NotImplementedError(
60 + "MiniMax-M3 AITER sparse PA does not support speculative "
61 + f"decode_query_len={d.decode_query_len}"
62 + )
63 + minimax_m3_sparse_attn_decode_aiter(
64 + q[:nd],
65 + k_cache,
66 + v_cache,
67 + topk[:, :nd, :],
68 + d.block_table,
69 + d.seq_lens,
70 + self.num_kv_heads,
71 + self.scale,
72 + out[:nd],
73 + k_scale=k_scale,
74 + v_scale=v_scale,
75 + )
76 +
77 + if main_md.num_prefills > 0:
78 + p = main_md.prefill
79 + assert p is not None
80 + minimax_m3_sparse_attn_prefill_aiter(
81 + q[nd:],
82 + k_cache,
83 + v_cache,
84 + topk[:, nd:num_tokens, :],
85 + p.block_table,
86 + p.cu_seqlens_q,
87 + p.context_lens,
88 + self.num_kv_heads,
89 + self.scale,
90 + out[nd:],
91 + k_scale=k_scale,
92 + v_scale=v_scale,
93 + )
94 + return output
vllm/models/minimax_m3/common/ops/sparse_attn.py
+160
@@ -52,6 +52,8 @@ _FP8_DTYPES = (
52 def _gqa_sparse_fwd_kernel(
53 q_ptr, # [total_q, num_heads, head_dim]
54 kv_cache_ptr, # main cache: [num_blocks, num_kv_heads, 128, 2*head_dim]
55 + k_scale_ptr,
56 + v_scale_ptr,
57 t_ptr, # topk_idx: [num_kv_heads, total_q, topk]
58 o_ptr, # [total_q, num_heads, head_dim]
59 block_table_ptr, # [num_reqs, max_blocks]
@@ -72,6 +74,10 @@ def _gqa_sparse_fwd_kernel(
74 stride_kv_h,
75 stride_kv_pos,
76 stride_kv_d,
77 + stride_ks_h,
78 + stride_ks_t,
79 + stride_vs_h,
80 + stride_vs_t,
81 stride_th,
82 stride_tn,
83 stride_tk,
@@ -85,6 +91,7 @@ def _gqa_sparse_fwd_kernel(
91 BLOCK_SIZE_H: tl.constexpr,
92 BLOCK_SIZE_QH: tl.constexpr,
93 USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load
94 + KV_SCALE_MODE: tl.constexpr, # 0: none, 1: scalar, 2: [kv_head, token]
95 ):
96 sm_scale_log2e = sm_scale * 1.4426950409
97 pid_q = tl.program_id(0)
@@ -148,6 +155,17 @@ def _gqa_sparse_fwd_kernel(
155 )
156 if USE_FP8:
157 k = k.to(q.dtype)
158 + if KV_SCALE_MODE == 1:
159 + k = (k * tl.load(k_scale_ptr)).to(q.dtype)
160 + elif KV_SCALE_MODE == 2:
161 + k_scale = tl.load(
162 + k_scale_ptr
163 + + pid_kh * stride_ks_h
164 + + (page * BLOCK_SIZE_K + off_n) * stride_ks_t,
165 + mask=pos_mask,
166 + other=1.0,
167 + )
168 + k = (k * k_scale[None, :]).to(q.dtype)
169 qk = tl.zeros((BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_K), dtype=tl.float32)
170 # causal: q_abs_pos - k_off >= block_start (c)
171 qk += tl.where(off_q[:, None, :] >= c, 0, float("-inf"))
@@ -169,6 +187,17 @@ def _gqa_sparse_fwd_kernel(
187 )
188 if USE_FP8:
189 v = v.to(q.dtype)
190 + if KV_SCALE_MODE == 1:
191 + v = (v * tl.load(v_scale_ptr)).to(q.dtype)
192 + elif KV_SCALE_MODE == 2:
193 + v_scale = tl.load(
194 + v_scale_ptr
195 + + pid_kh * stride_vs_h
196 + + (page * BLOCK_SIZE_K + off_n) * stride_vs_t,
197 + mask=pos_mask,
198 + other=1.0,
199 + )
200 + v = (v * v_scale[:, None]).to(q.dtype)
201 acc_o += tl.dot(p.to(v.dtype), v)
202 m_i = m_ij
203 lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
@@ -205,6 +234,8 @@ def _gqa_sparse_fwd_kernel(
234 def _gqa_sparse_decode_kernel(
235 q_ptr, # [total_q, num_heads, head_dim]
236 kv_cache_ptr, # main cache: [num_blocks, num_kv_heads, 128, 2*head_dim]
237 + k_scale_ptr,
238 + v_scale_ptr,
239 t_ptr, # topk_idx: [num_kv_heads, total_q, topk]
240 o_ptr, # partial out: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim]
241 lse_ptr, # partial lse (log2): [NUM_TOPK_CHUNKS, total_q, num_heads]
@@ -223,6 +254,10 @@ def _gqa_sparse_decode_kernel(
254 stride_kv_h,
255 stride_kv_pos,
256 stride_kv_d,
257 + stride_ks_h,
258 + stride_ks_t,
259 + stride_vs_h,
260 + stride_vs_t,
261 stride_th,
262 stride_tn,
263 stride_tk,
@@ -239,6 +274,7 @@ def _gqa_sparse_decode_kernel(
274 BLOCK_SIZE_H: tl.constexpr,
275 BLOCK_SIZE_D: tl.constexpr,
276 USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load
277 + KV_SCALE_MODE: tl.constexpr, # 0: none, 1: scalar, 2: [kv_head, token]
278 USE_PDL: tl.constexpr,
279 ):
280 sm_scale_log2e = sm_scale * 1.4426950409
@@ -305,6 +341,17 @@ def _gqa_sparse_decode_kernel(
341 )
342 if USE_FP8:
343 k = k.to(q.dtype)
344 + if KV_SCALE_MODE == 1:
345 + k = (k * tl.load(k_scale_ptr)).to(q.dtype)
346 + elif KV_SCALE_MODE == 2:
347 + k_scale = tl.load(
348 + k_scale_ptr
349 + + pid_kh * stride_ks_h
350 + + (page * BLOCK_SIZE_K + off_n) * stride_ks_t,
351 + mask=pos_mask,
352 + other=1.0,
353 + )
354 + k = (k * k_scale[None, :]).to(q.dtype)
355 qk = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_K), dtype=tl.float32)
356 qk += tl.where(pos_mask[None, :], 0, float("-inf"))
357 qk += tl.dot(q, k) * sm_scale_log2e
@@ -323,6 +370,17 @@ def _gqa_sparse_decode_kernel(
370 )
371 if USE_FP8:
372 v = v.to(q.dtype)
373 + if KV_SCALE_MODE == 1:
374 + v = (v * tl.load(v_scale_ptr)).to(q.dtype)
375 + elif KV_SCALE_MODE == 2:
376 + v_scale = tl.load(
377 + v_scale_ptr
378 + + pid_kh * stride_vs_h
379 + + (page * BLOCK_SIZE_K + off_n) * stride_vs_t,
380 + mask=pos_mask,
381 + other=1.0,
382 + )
383 + v = (v * v_scale[:, None]).to(q.dtype)
384 acc_o += tl.dot(p.to(v.dtype), v)
385 m_i = m_ij
386 lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
@@ -410,6 +468,48 @@ def _merge_topk_attn_out_kernel(
468 # ---------------------------------------------------------------------------
469 # Python wrappers
470 # ---------------------------------------------------------------------------
471 +_KV_SCALE_NONE = 0
472 +_KV_SCALE_SCALAR = 1
473 +_KV_SCALE_PER_TOKEN_HEAD = 2
474 +
475 +
476 +def _kv_scale_args(
477 + output: torch.Tensor,
478 + num_kv_heads: int,
479 + k_scale: torch.Tensor | None,
480 + v_scale: torch.Tensor | None,
481 +) -> tuple[torch.Tensor, torch.Tensor, int, int, int, int, int]:
482 + if k_scale is None and v_scale is None:
483 + return output, output, 0, 0, 0, 0, _KV_SCALE_NONE
484 + if k_scale is None or v_scale is None:
485 + raise ValueError("k_scale and v_scale must be both provided or both None")
486 + if k_scale.device != output.device or v_scale.device != output.device:
487 + raise ValueError("k_scale and v_scale must be on the same device as output")
488 + if k_scale.numel() == 1 and v_scale.numel() == 1:
489 + return k_scale, v_scale, 0, 0, 0, 0, _KV_SCALE_SCALAR
490 + if k_scale.dim() == 2 and v_scale.dim() == 2:
491 + if k_scale.shape[0] != num_kv_heads or v_scale.shape[0] != num_kv_heads:
492 + raise ValueError(
493 + "per-token/head KV scales must have shape "
494 + f"[{num_kv_heads}, max_kv_tokens]"
495 + )
496 + if k_scale.shape != v_scale.shape:
497 + raise ValueError("k_scale and v_scale must have matching shapes")
498 + return (
499 + k_scale,
500 + v_scale,
501 + k_scale.stride(0),
502 + k_scale.stride(1),
503 + v_scale.stride(0),
504 + v_scale.stride(1),
505 + _KV_SCALE_PER_TOKEN_HEAD,
506 + )
507 + raise ValueError(
508 + "MiniMax-M3 sparse attention supports scalar KV scales or "
509 + "[num_kv_heads, max_kv_tokens] per-token/head scales"
510 + )
511 +
512 +
513 @torch.no_grad()
514 def minimax_m3_sparse_attn(
515 q: torch.Tensor, # [total_q, num_heads, head_dim]
@@ -423,6 +523,8 @@ def minimax_m3_sparse_attn(
523 num_kv_heads: int,
524 sm_scale: float,
525 output: torch.Tensor, # [total_q, num_heads, head_dim]
526 + k_scale: torch.Tensor | None = None,
527 + v_scale: torch.Tensor | None = None,
528 ) -> None:
529 """GQA block-sparse attention over the selected blocks. block_size_q == 1."""
530 total_q, num_heads, head_dim = q.shape
@@ -430,10 +532,33 @@ def minimax_m3_sparse_attn(
532 topk = topk_idx.shape[-1]
533 gqa_group_size = num_heads // num_kv_heads
534 use_fp8 = kv_cache.dtype in _FP8_DTYPES
535 + (
536 + k_scale_arg,
537 + v_scale_arg,
538 + stride_ks_h,
539 + stride_ks_t,
540 + stride_vs_h,
541 + stride_vs_t,
542 + kv_scale_mode,
543 + ) = (
544 + _kv_scale_args(output, num_kv_heads, k_scale, v_scale)
545 + if use_fp8
546 + else (
547 + output,
548 + output,
549 + 0,
550 + 0,
551 + 0,
552 + 0,
553 + _KV_SCALE_NONE,
554 + )
555 + )
556 grid = (max_query_len, num_kv_heads, batch)
557 _gqa_sparse_fwd_kernel[grid](
558 q,
559 kv_cache,
560 + k_scale_arg,
561 + v_scale_arg,
562 topk_idx,
563 output,
564 block_table,
@@ -454,6 +579,10 @@ def minimax_m3_sparse_attn(
579 kv_cache.stride(1),
580 kv_cache.stride(2),
581 kv_cache.stride(3),
582 + stride_ks_h,
583 + stride_ks_t,
584 + stride_vs_h,
585 + stride_vs_t,
586 topk_idx.stride(0),
587 topk_idx.stride(1),
588 topk_idx.stride(2),
@@ -464,6 +593,7 @@ def minimax_m3_sparse_attn(
593 BLOCK_SIZE_Q=1,
594 BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
595 USE_FP8=use_fp8,
596 + KV_SCALE_MODE=kv_scale_mode,
597 )
598
599
@@ -478,6 +608,8 @@ def minimax_m3_sparse_attn_decode(
608 sm_scale: float,
609 output: torch.Tensor, # [total_q, num_heads, head_dim]
610 decode_query_len: int,
611 + k_scale: torch.Tensor | None = None,
612 + v_scale: torch.Tensor | None = None,
613 ) -> None:
614 """GQA block-sparse attention for decode (split-K over the top-k blocks)."""
615 total_q, num_heads, head_dim = q.shape
@@ -485,6 +617,27 @@ def minimax_m3_sparse_attn_decode(
617 max_topk = topk_idx.shape[-1]
618 gqa_group_size = num_heads // num_kv_heads
619 use_fp8 = kv_cache.dtype in _FP8_DTYPES
620 + (
621 + k_scale_arg,
622 + v_scale_arg,
623 + stride_ks_h,
624 + stride_ks_t,
625 + stride_vs_h,
626 + stride_vs_t,
627 + kv_scale_mode,
628 + ) = (
629 + _kv_scale_args(output, num_kv_heads, k_scale, v_scale)
630 + if use_fp8
631 + else (
632 + output,
633 + output,
634 + 0,
635 + 0,
636 + 0,
637 + 0,
638 + _KV_SCALE_NONE,
639 + )
640 + )
641 use_pdl = current_platform.is_arch_support_pdl()
642 # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA
643 # SM9+); this ROCm Triton rejects it even when False ("Keyword argument
@@ -505,6 +658,8 @@ def minimax_m3_sparse_attn_decode(
658 _gqa_sparse_decode_kernel[grid](
659 q,
660 kv_cache,
661 + k_scale_arg,
662 + v_scale_arg,
663 topk_idx,
664 o_partial,
665 lse_partial,
@@ -523,6 +678,10 @@ def minimax_m3_sparse_attn_decode(
678 kv_cache.stride(1),
679 kv_cache.stride(2),
680 kv_cache.stride(3),
681 + stride_ks_h,
682 + stride_ks_t,
683 + stride_vs_h,
684 + stride_vs_t,
685 topk_idx.stride(0),
686 topk_idx.stride(1),
687 topk_idx.stride(2),
@@ -537,6 +696,7 @@ def minimax_m3_sparse_attn_decode(
696 BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
697 NUM_TOPK_CHUNKS=num_topk_chunks,
698 USE_FP8=use_fp8,
699 + KV_SCALE_MODE=kv_scale_mode,
700 USE_PDL=use_pdl,
701 **pdl_launch,
702 )
vllm/models/minimax_m3/common/sparse_attention.py
+48 -4
@@ -20,6 +20,7 @@ from typing import ClassVar
20
21 import torch
22
23 +from vllm._aiter_ops import rocm_aiter_ops
24 from vllm.config import VllmConfig
25 from vllm.config.cache import CacheDType
26 from vllm.forward_context import get_forward_context
@@ -58,6 +59,21 @@ from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
59 logger = init_logger(__name__)
60
61
62 +def _minimax_m3_aiter_sparse_pa_requested() -> bool:
63 + return rocm_aiter_ops.is_enabled() and rocm_aiter_ops.is_shuffle_kv_cache_enabled()
64 +
65 +
66 +def minimax_m3_use_aiter_sparse_pa(num_kv_heads: int) -> bool:
67 + """Whether to use the ROCm AITER page-16 sparse PA prototype."""
68 + requested = _minimax_m3_aiter_sparse_pa_requested()
69 + if requested and num_kv_heads != 1:
70 + raise ValueError(
71 + "MiniMax M3 AITER sparse paged attention requires "
72 + f"num_kv_heads == 1 per tensor-parallel rank, got {num_kv_heads}."
73 + )
74 + return requested
75 +
76 +
77 class MiniMaxM3SparseBackend(AttentionBackend):
78 """Block-sparse GQA backend for MiniMax M3 sparse attention layers."""
79
@@ -104,6 +120,12 @@ class MiniMaxM3SparseBackend(AttentionBackend):
120 head_size: int,
121 cache_dtype_str: str = "auto",
122 ) -> tuple[int, ...]:
123 + if minimax_m3_use_aiter_sparse_pa(num_kv_heads):
124 + # AITER's assembly paged-attention kernels require independently
125 + # contiguous K and V storage. Keep that specialized layout behind
126 + # the shuffle flag while every other implementation uses the
127 + # packed-content contract introduced by #44455.
128 + return (num_blocks, 2, block_size, num_kv_heads, head_size)
129 # K and V are packed into the content dim: logical (B, H, N, 2*hs).
130 return (num_blocks, num_kv_heads, block_size, 2 * head_size)
131
@@ -116,6 +138,11 @@ class MiniMaxM3SparseBackend(AttentionBackend):
138 # layout we want.
139 if include_num_layers_dimension:
140 raise NotImplementedError # no cross-layer KV blocks in M3
141 + if _minimax_m3_aiter_sparse_pa_requested():
142 + # The AITER page-16 sparse PA path reinterprets the K and V slices
143 + # as separate SHUFFLE caches. Keep K/V physically separated so
144 + # kv_cache[:, 0] and kv_cache[:, 1] are contiguous byte ranges.
145 + return (1, 0, 2, 3, 4)
146 cache_layout = get_kv_cache_layout()
147 if cache_layout == "NHD":
148 # (num_blocks, block_size, num_kv_heads, 2*head_size)
@@ -365,6 +392,8 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
392 kv_cache = (
393 kv_cache.view(self.kv_cache_fp8_dtype) if self.use_fp8_kv else kv_cache
394 )
395 + k_scale = getattr(layer, "_k_scale", None) if self.use_fp8_kv else None
396 + v_scale = getattr(layer, "_v_scale", None) if self.use_fp8_kv else None
397
398 # Decode [:nd]: split-K over the selected blocks (request-major chunks).
399 if main_md.num_decodes > 0:
@@ -380,6 +409,8 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
409 self.scale,
410 out[:nd],
411 d.decode_query_len,
412 + k_scale=k_scale,
413 + v_scale=v_scale,
414 )
415
416 # Prefill [nd:]: cu_seqlens_q already rebased to 0.
@@ -398,6 +429,8 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
429 self.num_kv_heads,
430 self.scale,
431 out[nd:],
432 + k_scale=k_scale,
433 + v_scale=v_scale,
434 )
435 return output
436
@@ -406,27 +439,38 @@ def select_main_impl_cls(
439 *,
440 topk_blocks: int,
441 kv_cache_dtype: str,
442 + num_kv_heads: int,
443 ) -> type[MiniMaxM3SparseImpl]:
444 """Pick the main attend impl off the main KV-cache dtype.
445
446 Blackwell (SM100) uses the MSA attend for supported top-k block counts
413 - when the KV cache is BF16 or FP8 E4M3; non-Blackwell and FP8 E5M2 fall
414 - back to Triton. The MSA module is imported lazily so AMD/non-SM100 never
415 - import fmha_sm100.
447 + when the KV cache is BF16 or FP8 E4M3; MI355 uses AITER sparse PA
448 + with shuffle KV cache layout; Other platforms and FP8 E5M2 fall
449 + back to Triton. The MSA modules are imported lazily avoid import errors
450 + on unsupported platforms.
451 """
452 + use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(num_kv_heads)
453 use_msa = (
454 current_platform.is_cuda()
455 and current_platform.is_device_capability_family(100)
456 and topk_blocks in (4, 8, 16, 32)
457 and kv_cache_dtype != "fp8_e5m2"
458 )
423 - selected = "MSA" if use_msa else "Triton"
459 + selected = (
460 + "AITER_SPARSE_PA" if use_aiter_sparse_pa else ("MSA" if use_msa else "Triton")
461 + )
462 logger.info_once(
463 "MiniMax M3 sparse attention selected %s (kv_cache_dtype=%s, topk_blocks=%s)",
464 selected,
465 kv_cache_dtype,
466 topk_blocks,
467 )
468 + if use_aiter_sparse_pa:
469 + from vllm.models.minimax_m3.amd.sparse_attention_msa import (
470 + MiniMaxM3SparseAiterPAImpl,
471 + )
472 +
473 + return MiniMaxM3SparseAiterPAImpl
474 if use_msa:
475 from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
476 MiniMaxM3SparseMSAImpl,
vllm/models/minimax_m3/nvidia/model.py
+7
@@ -25,6 +25,7 @@ from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
25 from vllm.forward_context import get_forward_context
26 from vllm.model_executor.layers.activation import SiluAndMulWithClamp
27 from vllm.model_executor.layers.attention import Attention
28 +from vllm.model_executor.layers.attention.attention import set_default_quant_scales
29 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
30 from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
31 fused_allreduce_gemma_rms_norm,
@@ -491,6 +492,11 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
492 self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
493 self.kv_cache_dtype, vllm_config.model_config
494 )
495 + # MiniMax-M3 sparse attention owns its KV-cache insert/read path instead
496 + # of wrapping the generic Attention module. Keep the same runtime scale
497 + # attributes so FP8 KV reads can honor vLLM's per-layer descale contract.
498 + self.calculate_kv_scales = False
499 + set_default_quant_scales(self, register_buffer=True)
500 # Indexer side-cache dtype, mirroring --kv-cache-dtype for the main
501 # cache (--attention-config '{"indexer_kv_dtype": ...}').
502 self.indexer_kv_dtype = vllm_config.attention_config.indexer_kv_dtype
@@ -506,6 +512,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
512 self.impl: MiniMaxM3SparseImpl = select_main_impl_cls( # type: ignore[assignment]
513 topk_blocks=sparse_cfg["sparse_topk_blocks"],
514 kv_cache_dtype=self.kv_cache_dtype,
515 + num_kv_heads=self.num_kv_heads,
516 )(
517 self.num_heads,
518 self.head_dim,
vllm/models/minimax_m3/nvidia/sparse_attention_msa.py
+4
@@ -49,6 +49,8 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
49 kv_cache = (
50 kv_cache.view(self.kv_cache_fp8_dtype) if self.use_fp8_kv else kv_cache
51 )
52 + k_scale = getattr(layer, "_k_scale", None) if self.use_fp8_kv else None
53 + v_scale = getattr(layer, "_v_scale", None) if self.use_fp8_kv else None
54
55 # Decode [:nd]: Triton split-K placeholder (no MSA decode yet).
56 if main_md.num_decodes > 0:
@@ -64,6 +66,8 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
66 self.scale,
67 out[:nd],
68 d.decode_query_len,
69 + k_scale=k_scale,
70 + v_scale=v_scale,
71 )
72
73 # Prefill [nd:]: MSA sparse FMHA over the selected blocks.