From bb4844dba3c620490f190207f00a597ab7d08507 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Tue, 16 Jun 2026 06:30:27 +0000 Subject: [PATCH] [Model] Add fmha_sm100 MSA indexer backend + fp8 index cache for MiniMax M3 Add an SM100/Blackwell lightning-indexer impl that computes the per-128-block QK max-scores with fmha_sm100's score-only (OnlyScore) path and selects the top-k blocks with the existing Triton minimax_m3_index_topk kernel, mirroring how the main MSA attention pairs the SM100 attend with Triton. Decode and prefill requests are split manually (decode-first batch) and each side gets its own _fmha_sm100_plan / _fmha_sm100 call. Auto-selected on SM100 when topk_blocks in (4, 8, 16, 32) for both bf16 and fp8 index caches; falls back to the Triton indexer otherwise. The builder declares AttentionCGSupport.NEVER (eager; the attention is broken out of the graph by _run_attention). Extend the fused qknorm+rope+kv-insert kernel to optionally emit fp8 (e4m3) for the index-K cache and index-Q via a direct cast with no scale tensors (RMSNorm outputs are O(1) and scalar scales do not change top-k ordering). Only the index outputs go fp8; q/k/v and q_out stay bf16 and bit-identical to the existing path. MiniMaxM3IndexerCache now accepts fp8 caches and the model allocates index_q in the cache dtype. Tests: test_fmha_sm100_indexer_matches_reference (bf16/fp8 x prefill/decode) and test_msa_indexer_impl_matches_triton (full impl parity vs the Triton indexer through the real metadata builders); fp8 fused-kernel parity is covered in test_fused_minimax_m3_qknorm_rope_kv_insert. AI assistance (Claude Code) was used for this change. Signed-off-by: Yongye Zhu Co-Authored-By: Claude Opus 4.8 Signed-off-by: Yongye Zhu --- ...minimax_m3_qknorm_rope_kv_insert_kernel.cu | 203 ++++++++++---- tests/kernels/attention/test_minimax_m3.py | 245 +++++++++++++++++ ..._fused_minimax_m3_qknorm_rope_kv_insert.py | 96 +++++++ vllm/models/minimax_m3/common/indexer.py | 37 ++- vllm/models/minimax_m3/nvidia/indexer_msa.py | 248 ++++++++++++++++++ vllm/models/minimax_m3/nvidia/model.py | 7 +- 6 files changed, 776 insertions(+), 60 deletions(-) create mode 100644 vllm/models/minimax_m3/nvidia/indexer_msa.py diff --git a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu index 06c8048cd90..4261b702cd7 100644 --- a/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu @@ -67,6 +67,13 @@ #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" #endif +// Direct float -> E4M3 FP8 conversion for the indexer Q / index-K outputs. +#ifndef USE_ROCM + #include +#else + #include +#endif + #ifndef FINAL_MASK #ifdef USE_ROCM #define FINAL_MASK 0xffffffffffffffffULL @@ -75,6 +82,19 @@ #endif #endif +#ifdef USE_ROCM +// ROCm-compatible direct float -> E4M3 FP8 conversion (mirrors the DeepSeek V4 +// fused kernel). +__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) { + #if defined(HIP_FP8_TYPE_OCP) + __hip_fp8_e4m3 fp8_val(val); + #else + __hip_fp8_e4m3_fnuz fp8_val(val); + #endif + return reinterpret_cast(fp8_val); +} +#endif + namespace vllm { namespace minimax_m3_fused_ops { @@ -193,6 +213,8 @@ __device__ __forceinline__ void storeElems( *reinterpret_cast(dst) = v; } +// Main K/V cache store. kAuto = unquantized (cache_t == scalar_t); fp8 cache +// dtypes use the scaled-convert path with identity scale. template __device__ __forceinline__ void storeCacheElems( cache_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) { @@ -208,6 +230,32 @@ __device__ __forceinline__ void storeCacheElems( } } +// Store 4 fp32 registers -> 4 contiguous E4M3 FP8 bytes (direct cast, +// saturating to ±448). Used for the fp8 indexer-Q / index-K outputs; no scale +// (RMSNorm outputs are O(1) and the score path only needs relative block +// ordering). +__device__ __forceinline__ void storeElemsFp8( + uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) { + constexpr float kFp8Max = 448.0f; +#ifndef USE_ROCM + __nv_fp8x2_storage_t out2[kElemsPerLane / 2]; + #pragma unroll + for (int i = 0; i < kElemsPerLane / 2; i++) { + float2 vv = make_float2(elems[2 * i], elems[2 * i + 1]); + vv.x = fminf(fmaxf(vv.x, -kFp8Max), kFp8Max); + vv.y = fminf(fmaxf(vv.y, -kFp8Max), kFp8Max); + out2[i] = __nv_cvt_float2_to_fp8x2(vv, __NV_SATFINITE, __NV_E4M3); + } + *reinterpret_cast(dst) = *reinterpret_cast(out2); +#else + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + float vv = fminf(fmaxf(elems[i], -kFp8Max), kFp8Max); + dst[i] = rocm_cvt_float_to_fp8_e4m3(vv); + } +#endif +} + // ──────────────────────────────────────────────────────────────────────────── // Kernel // ──────────────────────────────────────────────────────────────────────────── @@ -224,12 +272,14 @@ __device__ __forceinline__ void storeCacheElems( // V : nkv only if kInsertKV (V-cache insert; no warps in dense) // IQ: niq only if kIsSparse (norm+RoPE) // IK: 1 only if kIsSparse (norm+RoPE; +index-cache insert) +// cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx: +// indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte). template + typename out_idx_t, bool kIsSparse, bool kInsertKV, bool kFp8Idx> __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse) - scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr - scalar_t* __restrict__ index_q_out, // [N, niq*128] contiguous, or nullptr + scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr + out_idx_t* __restrict__ index_q_out, // [N, niq*128]; scalar_t or e4m3 byte scalar_t const* __restrict__ q_norm_w, scalar_t const* __restrict__ k_norm_w, scalar_t const* __restrict__ iq_norm_w, @@ -238,8 +288,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( int64_t const* __restrict__ positions, // [N] i64 int64_t const* __restrict__ slot_mapping, // main K/V slots or nullptr int64_t const* __restrict__ index_slot_mapping, // index K slots/nullptr - cache_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr - scalar_t* __restrict__ index_cache, // [nb*bs, 128] or nullptr + cache_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr + out_idx_t* __restrict__ index_cache, // [nb*bs, 128]; scalar_t or e4m3 byte float const eps, int const rotary_dim, int const num_tokens, int const nq, int const nkv, int const niq, int const block_size, // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128]. @@ -334,9 +384,12 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( store_ptr = q_out + static_cast(tokenIdx) * nq * kHeadDim + slot * kHeadDim; } else if (isIQ && index_q_out != nullptr) { - store_ptr = index_q_out + - static_cast(tokenIdx) * niq * kHeadDim + - (slot - iq_begin) * kHeadDim; + // bf16 index_q_out: gather here. fp8: written by the explicit fp8 store. + if constexpr (!kFp8Idx) { + store_ptr = index_q_out + + static_cast(tokenIdx) * niq * kHeadDim + + (slot - iq_begin) * kHeadDim; + } } // PDL: wait for the predecessor kernel (the qkv-projection GEMM that @@ -356,7 +409,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim; normAndRope(elems, laneId, eps, norm_w, do_rope, rotary_dim, cos_ptr, /*apply_norm=*/norm_w != nullptr); - storeElems(store_ptr + dim_base, elems); + if constexpr (kFp8Idx) { + // index_q is e4m3 bytes; Q/K (and in-place index_k) stay scalar_t. + if (isIQ && index_q_out != nullptr) { + storeElemsFp8(index_q_out + + static_cast(tokenIdx) * niq * kHeadDim + + (slot - iq_begin) * kHeadDim + dim_base, + elems); + } else { + storeElems(store_ptr + dim_base, elems); + } + } else { + storeElems(store_ptr + dim_base, elems); + } } // ── Cache inserts (sparse serving only). ─────────────────────────────── @@ -367,8 +432,11 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( : (isIK ? index_slot_mapping[tokenIdx] : -1); if (sm >= 0) { // skip padded / unscheduled tokens if (isIK) { - scalar_t* dst = index_cache + sm * kHeadDim + dim_base; - storeElems(dst, elems); + if constexpr (kFp8Idx) { + storeElemsFp8(index_cache + sm * kHeadDim + dim_base, elems); + } else { + storeElems(index_cache + sm * kHeadDim + dim_base, elems); + } } else if (isK || isV) { // kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim]. // Paging is logical (block = sm/block_size, token = sm%block_size); @@ -398,19 +466,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( // Launch wrapper // ──────────────────────────────────────────────────────────────────────────── template -void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out, - scalar_t const* q_norm_w, scalar_t const* k_norm_w, - scalar_t const* iq_norm_w, scalar_t const* ik_norm_w, - scalar_t const* cos_sin_cache, - int64_t const* positions, int64_t const* slot_mapping, - int64_t const* index_slot_mapping, cache_t* kv_cache, - scalar_t* index_cache, float const eps, - int const rotary_dim, int const num_tokens, - int const nq, int const nkv, int const niq, - int const block_size, int64_t const kv_s_block, - int64_t const kv_s_kv, int64_t const kv_s_token, - int64_t const kv_s_head, bool const has_index, - bool const insert_kv, cudaStream_t stream) { +void launchFusedMiniMaxM3( + scalar_t* qkv, scalar_t* q_out, void* index_q_out, scalar_t const* q_norm_w, + scalar_t const* k_norm_w, scalar_t const* iq_norm_w, + scalar_t const* ik_norm_w, scalar_t const* cos_sin_cache, + int64_t const* positions, int64_t const* slot_mapping, + int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache, + float const eps, int const rotary_dim, int const num_tokens, int const nq, + int const nkv, int const niq, int const block_size, + int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token, + int64_t const kv_s_head, bool const has_index, bool const insert_kv, + bool const fp8_idx, cudaStream_t stream) { + // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the + // void* pointers per instantiation in the LAUNCH macro. // Slot count must match the kernel's compile-time gating. int const v_slots = insert_kv ? nkv : 0; 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, config.attrs = attrs; config.numAttrs = (sm_version >= 90) ? 1 : 0; - #define LAUNCH(IS_SPARSE, INSERT) \ - cudaLaunchKernelEx( \ - &config, \ - fusedMiniMaxM3QNormRopeKVInsertKernel, \ - qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, \ - cos_sin_cache, positions, slot_mapping, index_slot_mapping, kv_cache, \ - index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, \ - kv_s_block, kv_s_kv, kv_s_token, kv_s_head) + #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T) \ + cudaLaunchKernelEx( \ + &config, \ + fusedMiniMaxM3QNormRopeKVInsertKernel, \ + qkv, q_out, reinterpret_cast(index_q_out), q_norm_w, k_norm_w, \ + iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping, \ + index_slot_mapping, kv_cache, reinterpret_cast(index_cache), \ + eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \ + kv_s_kv, kv_s_token, kv_s_head) #else // ROCm: standard kernel launch syntax (no PDL/stream serialization). // clang-format off - #define LAUNCH(IS_SPARSE, INSERT) \ - fusedMiniMaxM3QNormRopeKVInsertKernel \ + #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T) \ + fusedMiniMaxM3QNormRopeKVInsertKernel \ <<>>( \ - qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, \ - ik_norm_w, cos_sin_cache, positions, slot_mapping, \ - index_slot_mapping, kv_cache, index_cache, eps, rotary_dim, \ + qkv, q_out, reinterpret_cast(index_q_out), q_norm_w, \ + k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \ + slot_mapping, index_slot_mapping, kv_cache, \ + reinterpret_cast(index_cache), eps, rotary_dim, \ num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv, \ kv_s_token, kv_s_head) // clang-format on @@ -466,14 +536,22 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out, if (has_index) { if (insert_kv) { - LAUNCH(true, true); // sparse serving + if (fp8_idx) { + LAUNCH(true, true, true, uint8_t); // sparse serving, fp8 index outputs + } else { + LAUNCH(true, true, false, scalar_t); // sparse serving, bf16 + } } else { - LAUNCH(true, false); // sparse profiling + if (fp8_idx) { + LAUNCH(true, false, true, uint8_t); // sparse profiling, fp8 index_q + } else { + LAUNCH(true, false, false, scalar_t); // sparse profiling, bf16 + } } } else { // Dense layer: never has an index branch and never inserts here (the // generic Attention layer owns the KV insert). - LAUNCH(false, false); + LAUNCH(false, false, false, scalar_t); } #undef LAUNCH } @@ -485,8 +563,9 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out, vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3( \ reinterpret_cast(qkv.data_ptr()), \ q_out.has_value() ? reinterpret_cast(q_out->data_ptr()) : nullptr, \ - index_q_out.has_value() ? reinterpret_cast(index_q_out->data_ptr()) \ - : nullptr, \ + index_q_out.has_value() \ + ? reinterpret_cast(index_q_out->data_ptr()) \ + : nullptr, \ reinterpret_cast(q_norm_weight.data_ptr()), \ reinterpret_cast(k_norm_weight.data_ptr()), \ has_index ? reinterpret_cast(index_q_norm_weight->data_ptr()) \ @@ -502,11 +581,11 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out, : nullptr, \ insert_kv ? reinterpret_cast(kv_cache->data_ptr()) : nullptr, \ (insert_kv && has_index) \ - ? reinterpret_cast(index_cache->data_ptr()) \ + ? reinterpret_cast(index_cache->data_ptr()) \ : nullptr, \ static_cast(eps), static_cast(rotary_dim), num_tokens, nq, \ nkv, niq, static_cast(block_size), kv_s_block, kv_s_kv, kv_s_token, \ - kv_s_head, has_index, insert_kv, stream) + kv_s_head, has_index, insert_kv, fp8_idx, stream) // ──────────────────────────────────────────────────────────────────────────── // Torch op wrapper @@ -612,6 +691,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert( torch::headeronly::ScalarType::Long && index_slot_mapping->numel() == slot_mapping->numel()), "index_slot_mapping must be int64 CUDA with slot_mapping length"); + // Main attention KV cache: auto matches qkv, fp8 uses uint8 storage. if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) { STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(), "auto kv_cache dtype must match qkv"); @@ -620,9 +700,13 @@ void fused_minimax_m3_qknorm_rope_kv_insert( kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte, "fp8 kv_cache must use uint8 storage"); } - STD_TORCH_CHECK(index_cache.has_value() && - index_cache->scalar_type() == qkv.scalar_type(), - "insert mode requires matching index_cache"); + // Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3. + STD_TORCH_CHECK( + index_cache.has_value() && + (index_cache->scalar_type() == qkv.scalar_type() || + index_cache->scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn), + "insert mode requires index_cache matching qkv dtype or fp8 e4m3"); STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1, "kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous " "head_dim (stride(4)==1)"); @@ -652,14 +736,31 @@ void fused_minimax_m3_qknorm_rope_kv_insert( "index_q_out requires the index branch (num_index_heads > 0)"); STD_TORCH_CHECK( index_q_out->is_cuda() && index_q_out->is_contiguous() && - index_q_out->scalar_type() == qkv.scalar_type(), - "index_q_out must be a contiguous CUDA tensor matching qkv dtype"); + (index_q_out->scalar_type() == qkv.scalar_type() || + index_q_out->scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn), + "index_q_out must be contiguous CUDA, qkv dtype or fp8 e4m3"); STD_TORCH_CHECK(index_q_out->numel() == static_cast(num_tokens) * niq * kHeadDim, "index_q_out must have num_tokens * num_index_heads * 128 " "elements"); } + // fp8 index path: the index-K cache and index-Q outputs are e4m3 bytes while + // q/k/v + q_out stay qkv dtype. Both index outputs must agree. + auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn; + bool const fp8_idx = + (index_cache.has_value() && index_cache->scalar_type() == kFp8) || + (index_q_out.has_value() && index_q_out->scalar_type() == kFp8); + if (fp8_idx) { + STD_TORCH_CHECK( + !index_cache.has_value() || index_cache->scalar_type() == kFp8, + "fp8 index path: index_cache must be fp8 e4m3"); + STD_TORCH_CHECK( + !index_q_out.has_value() || index_q_out->scalar_type() == kFp8, + "fp8 index path: index_q_out must be fp8 e4m3"); + } + const torch::stable::accelerator::DeviceGuard device_guard( qkv.get_device_index()); auto stream = get_current_cuda_stream(qkv.get_device_index()); diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 1246f1721c2..07be65649b4 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -244,6 +244,251 @@ def test_prefill_index_topk_correctness(): _assert_topk_indices_equal_unordered(actual, expected) +# MSA indexer (SM100): fmha_sm100 OnlyScore for the per-block scores, then the +# Triton minimax_m3_index_topk for selection (no sparse_topk_select). Uses a +# deterministic construction (idx_q == 1, distinct e4m3-exact per-block values) +# so scores are strictly monotonic in the block id -> exact top-k agreement. +def _fmha_indexer_topk( + idx_q: torch.Tensor, # [total_q, H, 128] bf16/e4m3 + index_cache: torch.Tensor, # [num_pages, 128, 128] bf16/e4m3 + block_table: torch.Tensor, + q_lens: torch.Tensor, + seq_lens: torch.Tensor, + prefix_lens: torch.Tensor, + sm_scale: float, + topk: int, +) -> torch.Tensor: + """Replicate MiniMaxM3IndexerMSAImpl's score path (single decode/prefill side).""" + from vllm.third_party.fmha_sm100.api import _fmha_sm100, _fmha_sm100_plan + + num_idx_heads, head_dim = idx_q.shape[1], idx_q.shape[2] + nvp = [(s + 127) // 128 for s in seq_lens.tolist()] + kv_indices = torch.cat([block_table[r, : nvp[r]] for r in range(len(nvp))]).to( + torch.int32 + ) + + qo = q_lens.cpu().to(torch.int32) + kv = seq_lens.cpu().to(torch.int32) + plan = _fmha_sm100_plan( + qo, + kv, + num_idx_heads, + num_kv_heads=1, + qo_offset=kv - qo, + page_size=128, + output_maxscore=True, + causal=True, + num_kv_splits=1, + ) + k_pages = index_cache.view(index_cache.shape[0], 1, 128, head_dim) + _, max_score = _fmha_sm100( + idx_q, + k_pages, + k_pages, + plan, + kv_indices=kv_indices, + output_o=False, + output_maxscore=True, + sm_scale=sm_scale, + ) + + batch = q_lens.numel() + cu = torch.zeros(batch + 1, dtype=torch.int32, device=idx_q.device) + cu[1:] = q_lens.to(torch.int32).cumsum(0) + # max_score [H, k_tiles, total_q] -> transpose to [H, total_q, k_tiles]. + return minimax_m3_index_topk( + max_score.transpose(1, 2), + cu, + prefix_lens.to(torch.int32), + int(q_lens.max()), + topk, + 0, # init_blocks + 0, # local_blocks + ) + + +# e4m3-exact, strictly-increasing per-block values: with idx_q == 1 (also exact) +# the per-block scores are exact and distinct in BOTH bf16 and e4m3, so the fp8 +# score path selects the same top-k as the reference (no quantization ties). +_E4M3_EXACT_VALUES = [ + *range(1, 17), # 1..16 (step 1) + *range(18, 33, 2), # 18..32 (step 2) + *range(36, 65, 4), # 36..64 (step 4) + *range(72, 129, 8), # 72..128 (step 8) +] + + +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="fmha_sm100 indexer requires SM100 (Blackwell).", +) +@pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize( + ("q_lens", "prefix_lens"), + [ + ((4, 3), (2048, 2560)), # prefill: every token sees >= 16 causal blocks + ((1, 1, 1), (2048, 3000, 4096)), # decode: one query token per request + ], +) +def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype): + torch.manual_seed(0) + num_idx_heads, head_dim = 4, HEAD_DIM + device = "cuda" + + q_lens_t = torch.tensor(q_lens, device=device, dtype=torch.int32) + prefix_lens_t = torch.tensor(prefix_lens, device=device, dtype=torch.int32) + seq_lens = prefix_lens_t + q_lens_t + batch = len(q_lens) + max_blocks = (int(seq_lens.max()) + BLOCK_SIZE - 1) // BLOCK_SIZE + assert max_blocks <= len(_E4M3_EXACT_VALUES) + num_pages = batch * max_blocks + block_table = torch.randperm(num_pages, device=device, dtype=torch.int32).reshape( + batch, max_blocks + ) + + idx_q = torch.ones( + int(q_lens_t.sum()), num_idx_heads, head_dim, device=device, dtype=index_dtype + ) + index_cache = torch.empty( + num_pages, BLOCK_SIZE, head_dim, device=device, dtype=index_dtype + ) + for r in range(batch): + for b in range(max_blocks): + index_cache[block_table[r, b]] = float(_E4M3_EXACT_VALUES[b]) + + sm_scale = head_dim**-0.5 + actual = _fmha_indexer_topk( + idx_q, + index_cache, + block_table, + q_lens_t, + seq_lens, + prefix_lens_t, + sm_scale, + TOPK, + ) + expected = _reference_index_topk( + idx_q, + index_cache, + block_table, + q_lens_t, + seq_lens, + prefix_lens_t, + TOPK, + init_blocks=0, + local_blocks=0, + sm_scale=sm_scale, + ) + _assert_topk_indices_equal_unordered(actual, expected) + + +# Full impl-level parity: drive both MiniMaxM3IndexerMSAImpl (fmha_sm100 score + +# Triton top-k) and MiniMaxM3IndexerTritonImpl through their real metadata +# builders on the SAME CommonAttentionMetadata + index cache, and assert the +# selected blocks agree. This exercises all the metadata the impl/kernels consume +# (decode/prefill split, cu_seqlens_q rebasing, prefix_lens, kv_indices gather, +# decode_pages split) -- a metadata bug on either side shifts the causal window +# or the block->page mapping and breaks the comparison. +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="fmha_sm100 indexer requires SM100 (Blackwell).", +) +@pytest.mark.parametrize("topk", [8, 16]) +def test_msa_indexer_impl_matches_triton(topk, monkeypatch): + import vllm.models.minimax_m3.common.indexer as indexer_mod + from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, + ) + from vllm.config import set_current_vllm_config + from vllm.forward_context import set_forward_context + from vllm.models.minimax_m3.common.indexer import ( + MiniMaxM3IndexerTritonImpl, + MiniMaxM3IndexerTritonMetadataBuilder, + ) + from vllm.models.minimax_m3.nvidia.indexer_msa import ( + MiniMaxM3IndexerMSAImpl, + MiniMaxM3IndexerMSAMetadataBuilder, + ) + + torch.manual_seed(0) + device = torch.device("cuda") + num_idx_heads, head_dim = 4, HEAD_DIM + # TP=1: avoid requiring an initialized distributed group in a unit test. + monkeypatch.setattr(indexer_mod, "get_tensor_model_parallel_world_size", lambda: 1) + + vllm_config = create_vllm_config( + block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=8192 + ) + vllm_config.model_config.hf_config.sparse_attention_config = { + "sparse_num_index_heads": num_idx_heads + } + + # Decode-first mixed batch: 2 decode reqs (q_len 1) then 2 prefill reqs. Long + # prefixes so every token sees > TOPK causal blocks (non-trivial selection). + batch = BatchSpec(seq_lens=[2305, 2561, 2624, 2720], query_lens=[1, 1, 64, 96]) + common = create_common_attn_metadata( + batch, BLOCK_SIZE, device, arange_block_indices=True + ) + num_tokens = batch.compute_num_tokens() + + # Deterministic index cache: distinct, monotonic per-logical-block values so + # the top-k is unambiguous (both kernels pick the same blocks, no fp ties). + block_table = common.block_table_tensor + num_pages = int(block_table.max().item()) + 1 + index_cache = torch.zeros( + num_pages, BLOCK_SIZE, head_dim, device=device, dtype=DTYPE + ) + for r, seq_len in enumerate(batch.seq_lens): + for b in range((seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE): + index_cache[block_table[r, b]] = float(b + 1) + index_q = torch.ones( + num_tokens, num_idx_heads * head_dim, device=device, dtype=DTYPE + ) + + spec = MLAAttentionSpec( + block_size=BLOCK_SIZE, num_kv_heads=1, head_size=head_dim, dtype=DTYPE + ) + impl_kwargs = dict( + num_kv_heads=num_idx_heads, + scale=head_dim**-0.5, + topk_blocks=topk, + sparse_block_size=BLOCK_SIZE, + num_index_heads=num_idx_heads, + index_head_dim=head_dim, + init_blocks=0, + local_blocks=0, + ) + + with set_current_vllm_config(vllm_config): + msa_impl = MiniMaxM3IndexerMSAImpl(prefix="idx_msa", **impl_kwargs) + triton_impl = MiniMaxM3IndexerTritonImpl(prefix="idx_triton", **impl_kwargs) + msa_builder = MiniMaxM3IndexerMSAMetadataBuilder( + spec, [msa_impl.index_cache.prefix], vllm_config, device + ) + triton_builder = MiniMaxM3IndexerTritonMetadataBuilder( + spec, [triton_impl.index_cache.prefix], vllm_config, device + ) + + # Both impls score against the same index keys. + msa_impl.index_cache.kv_cache = index_cache + triton_impl.index_cache.kv_cache = index_cache + + attn_metadata = { + msa_impl.index_cache.prefix: msa_builder.build(0, common), + triton_impl.index_cache.prefix: triton_builder.build(0, common), + } + with set_forward_context(attn_metadata, vllm_config): + msa_decode, msa_prefill = msa_impl(index_q) + tri_decode, tri_prefill = triton_impl(index_q) + + assert msa_decode is not None and tri_decode is not None + assert msa_prefill is not None and tri_prefill is not None + _assert_topk_indices_equal_unordered(msa_decode, tri_decode) + _assert_topk_indices_equal_unordered(msa_prefill, tri_prefill) + + @pytest.mark.parametrize( ("decode_query_len", "max_decode_query_len"), [ diff --git a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py index 96729614f82..3f79d4f5db0 100644 --- a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py @@ -278,3 +278,99 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype): torch.testing.assert_close( index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0 ) + + +# ── Test 3: fp8 (e4m3) index outputs ───────────────────────────────────────── +# The fp8 score path stores index_q and the index-K cache as e4m3 while q/k/v + +# q_out stay bf16. Asserts: (1) q/k/v/q_out are bit-identical to the bf16 run +# (the index dtype must not perturb the main branch), and (2) the e4m3 index +# outputs dequantize close to the bf16 reference. + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9), + reason="e4m3 conversion requires CUDA SM89+.", +) +@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513]) +@pytest.mark.parametrize("block_size", [16, 64]) +def test_sparse_full_fp8_index(num_tokens, block_size): + torch.manual_seed(1) + device, dtype, eps = "cuda", torch.bfloat16, 1e-6 + base, max_pos = 5_000_000.0, 4096 + num_heads, num_kv_heads, num_idx_heads = 16, 4, 4 + + q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1 + k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1 + iq_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1 + ik_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1 + cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device) + positions = torch.randint( + 0, max_pos, (num_tokens,), dtype=torch.int64, device=device + ) + + qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM + iqsz, iksz = num_idx_heads * HEAD_DIM, HEAD_DIM + qkv0 = torch.randn( + num_tokens, qsz + 2 * kvsz + iqsz + iksz, dtype=dtype, device=device + ) + + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.randperm( + num_blocks * block_size, dtype=torch.int64, device=device + )[:num_tokens] + index_slot_mapping = torch.roll(slot_mapping, shifts=1) + + def run(index_dtype): + qkv = qkv0.clone() + kv_cache = torch.zeros( + num_blocks, + 2, + block_size, + num_kv_heads, + HEAD_DIM, + dtype=dtype, + device=device, + ) + index_cache = torch.zeros( + num_blocks, block_size, HEAD_DIM, dtype=index_dtype, device=device + ) + q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device) + index_q = torch.empty(num_tokens, iqsz, dtype=index_dtype, device=device) + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, + q_w, + k_w, + cos_sin, + positions, + num_heads, + num_kv_heads, + ROTARY_DIM, + eps, + iq_w, + ik_w, + num_idx_heads, + slot_mapping, + index_slot_mapping, + kv_cache, + index_cache, + block_size, + q_out, + index_q, + ) + return qkv, kv_cache, index_cache, q_out, index_q + + qkv_bf, kvc_bf, idxc_bf, qo_bf, iq_bf = run(torch.bfloat16) + qkv_fp, kvc_fp, idxc_fp, qo_fp, iq_fp = run(torch.float8_e4m3fn) + + assert iq_fp.dtype == torch.float8_e4m3fn + assert idxc_fp.dtype == torch.float8_e4m3fn + + # (1) The main branch (q/k/v in qkv, q_out, kv cache) must be bit-identical: + # the index output dtype must not perturb anything else. + torch.testing.assert_close(qo_fp, qo_bf, rtol=0, atol=0) + torch.testing.assert_close(qkv_fp, qkv_bf, rtol=0, atol=0) + torch.testing.assert_close(kvc_fp, kvc_bf, rtol=0, atol=0) + + # (2) Dequantized e4m3 index outputs match the bf16 reference within fp8 ulp. + torch.testing.assert_close(iq_fp.float(), iq_bf.float(), rtol=0.13, atol=0.05) + torch.testing.assert_close(idxc_fp.float(), idxc_bf.float(), rtol=0.13, atol=0.05) diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index 4da52805604..c989a1f4b11 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -31,6 +31,7 @@ from vllm.models.minimax_m3.common.ops.index_topk import ( minimax_m3_index_score, minimax_m3_index_topk, ) +from vllm.platforms import current_platform from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -120,16 +121,20 @@ class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase): backend_cls: type[AttentionBackend] = MiniMaxM3IndexerBackend, ) -> None: super().__init__() - if indexer_kv_dtype != "bf16": + if indexer_kv_dtype in ("fp8", "fp8_e4m3"): + cache_dtype = torch.float8_e4m3fn + elif indexer_kv_dtype == "bf16": + cache_dtype = torch.bfloat16 + else: raise NotImplementedError( - f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported yet " - "for the MiniMax M3 indexer cache (only 'bf16')." + f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the " + "MiniMax M3 indexer cache (only 'bf16' or 'fp8'/'fp8_e4m3')." ) self.kv_cache = torch.tensor([]) self.head_dim = head_dim self.indexer_kv_dtype = indexer_kv_dtype - # Storage dtype for the side cache (bf16 today; quantized layouts later). - self.dtype = torch.bfloat16 + # Side-cache storage dtype: bf16, or e4m3 for the fp8 score path. + self.dtype = cache_dtype self.prefix = prefix self.cache_config = cache_config # Impl-chosen backend -> each impl gets its own builder (get_attn_backend). @@ -438,18 +443,33 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl): def select_indexer_impl_cls( *, + topk_blocks: int, indexer_kv_dtype: IndexerKVDType = "bf16", ) -> type[MiniMaxM3IndexerImpl]: - """Pick the indexer impl off the index-cache dtype. + """Pick the indexer impl off the platform, top-k count, and cache dtype. - The SM100 MSA indexer score path is disabled for now; use the local Triton - indexer. If re-enabled, add a NVIDIA-specific ``MiniMaxM3IndexerImpl`` here. + On Blackwell (SM100) with ``topk_blocks`` in ``(4, 8, 16, 32)`` (matching the + main MSA attend), the fmha_sm100 score path + Triton top-k is used for both + bf16 and fp8 index caches. Everything else falls back to the Triton indexer + (bf16 only). """ if indexer_kv_dtype in ("mxfp4", "nvfp4"): raise NotImplementedError( f"indexer_kv_dtype={indexer_kv_dtype!r} needs the (not-yet-added) " "CuteDSL indexer impl." ) + if ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and topk_blocks in (4, 8, 16, 32) + and indexer_kv_dtype in ("bf16", "fp8", "fp8_e4m3") + ): + # Lazy import so AMD / non-SM100 never import fmha_sm100. + from vllm.models.minimax_m3.nvidia.indexer_msa import ( + MiniMaxM3IndexerMSAImpl, + ) + + return MiniMaxM3IndexerMSAImpl if indexer_kv_dtype != "bf16": raise NotImplementedError( f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the " @@ -483,6 +503,7 @@ class MiniMaxM3Indexer(nn.Module): ) -> None: super().__init__() impl_cls = select_indexer_impl_cls( + topk_blocks=topk_blocks, indexer_kv_dtype=indexer_kv_dtype, ) self.impl = impl_cls( diff --git a/vllm/models/minimax_m3/nvidia/indexer_msa.py b/vllm/models/minimax_m3/nvidia/indexer_msa.py new file mode 100644 index 00000000000..54d49137e8c --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/indexer_msa.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MSA (SM100/Blackwell) indexer impl for MiniMax M3. + +The lightning indexer's per-128-block QK max-score is computed with +``fmha_sm100``'s score-only (``OnlyScore``) path; the top-k block selection +reuses the existing Triton ``minimax_m3_index_topk`` kernel (it handles the +per-token causal window + forced init/local blocks for any ``topk``). This +mirrors how ``MiniMaxM3SparseMSAImpl`` pairs the SM100 attend with Triton. + +Decode and prefill requests are split manually (the batch is decode-first) and +each side gets its own ``_fmha_sm100_plan`` / ``_fmha_sm100`` call -- the public +``fmha_sm100`` wrapper's mixed-batch split + max-score merge is bypassed. + +The plan is built eagerly at metadata-build time (it ``.tolist()``s the segment +lengths and allocates fresh workspaces), so this builder is not +cudagraph-replay-safe -- it declares ``AttentionCGSupport.NEVER`` and the +attention runs eager (broken out of the graph by ``_run_attention``). +``fmha_sm100`` imports are function-local so this module is import-safe on +AMD / non-SM100. +""" + +from dataclasses import dataclass +from typing import ClassVar + +import torch + +from vllm.forward_context import get_forward_context +from vllm.models.minimax_m3.common.indexer import ( + MiniMaxM3IndexerBackend, + MiniMaxM3IndexerImpl, + MiniMaxM3IndexerMetadata, + MiniMaxM3IndexerMetadataBuilder, +) +from vllm.models.minimax_m3.common.ops.index_topk import minimax_m3_index_topk +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + CommonAttentionMetadata, +) +from vllm.v1.attention.backends.utils import split_decodes_and_prefills + +# Page size == sparse block size == index-K block; fmha tile id == M3 block id. +PAGE_SIZE = 128 + + +class MiniMaxM3IndexerMSABackend(MiniMaxM3IndexerBackend): + """Indexer side-cache backend selecting the MSA builder.""" + + @staticmethod + def get_builder_cls() -> type["MiniMaxM3IndexerMSAMetadataBuilder"]: + return MiniMaxM3IndexerMSAMetadataBuilder + + +@dataclass +class MiniMaxM3IndexerMSASubMetadata: + """Per-side (decode or prefill) fmha score plan + Triton top-k inputs.""" + + plan: dict # _fmha_sm100_plan PlanInfo + cu_seqlens_q: torch.Tensor # [n + 1] int32, rebased to 0 + prefix_lens: torch.Tensor # [n] int32, context tokens before this side + max_query_len: int + + +@dataclass +class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata): + """Indexer metadata with separate decode/prefill fmha score plans.""" + + # Whole-batch flattened physical page table [sum_pages] int32, request-major + # (decode reqs first); split at ``decode_pages`` for the two _fmha_sm100 runs. + kv_indices: torch.Tensor | None = None + decode_pages: int = 0 + decode_metadata: MiniMaxM3IndexerMSASubMetadata | None = None + prefill_metadata: MiniMaxM3IndexerMSASubMetadata | None = None + + +class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): + """Builds separate decode/prefill fmha_sm100 score plans (eager only).""" + + # Plans/workspaces are allocated fresh each build() and the run reads + # per-request metadata off the plan, so this is not cudagraph-replay-safe. + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> MiniMaxM3IndexerMSAMetadata: + from vllm.third_party.fmha_sm100.api import _fmha_sm100_plan + + num_reqs = common_attn_metadata.num_reqs + num_tokens = common_attn_metadata.num_actual_tokens + seq_lens = common_attn_metadata.seq_lens + block_table = common_attn_metadata.block_table_tensor + query_start_loc = common_attn_metadata.query_start_loc + + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( + split_decodes_and_prefills( + common_attn_metadata, + decode_threshold=self.reorder_batch_threshold, + require_uniform=True, + ) + ) + assert num_decodes + num_prefills == num_reqs + assert num_decode_tokens + num_prefill_tokens == num_tokens + + # Context (prefix) lengths into the stable buffer; sliced per side below. + context_lens = self.context_len_buffer[:num_reqs] + context_lens.copy_( + common_attn_metadata.compute_num_computed_tokens(), non_blocking=True + ) + + # Exact per-request lengths (host): the plan .tolist()s these and uses the + # exact KV lengths for run-time causal masking. + qsl_cpu = common_attn_metadata.query_start_loc_cpu[: num_reqs + 1] + qo_lens_cpu = (qsl_cpu[1:] - qsl_cpu[:-1]).to(torch.int32) + kv_lens_cpu = seq_lens[:num_reqs].cpu().to(torch.int32) + nvp = (kv_lens_cpu + PAGE_SIZE - 1) // PAGE_SIZE + + # Whole-batch request-major flat page table; decode pages come first. + max_blocks = block_table.shape[1] + cols = torch.arange(max_blocks, device=block_table.device) + valid = cols[None, :] < nvp.to(block_table.device)[:, None] + kv_indices = block_table[valid].to(torch.int32) + decode_pages = int(nvp[:num_decodes].sum()) + + def build_side_metadata( + lo: int, hi: int + ) -> MiniMaxM3IndexerMSASubMetadata | None: + if hi <= lo: + return None + side_qo = qo_lens_cpu[lo:hi] + side_kv = kv_lens_cpu[lo:hi] + plan = _fmha_sm100_plan( + side_qo, + side_kv, + self.num_index_heads, + num_kv_heads=1, + qo_offset=side_kv - side_qo, # bottom-right causal + page_size=PAGE_SIZE, + output_maxscore=True, + causal=True, + num_kv_splits=1, + ) + cu = query_start_loc[lo : hi + 1] - query_start_loc[lo] + return MiniMaxM3IndexerMSASubMetadata( + plan=plan, + cu_seqlens_q=cu.to(torch.int32), + prefix_lens=context_lens[lo:hi], + max_query_len=int(side_qo.max()), + ) + + # Decode requests precede prefill requests in the reordered batch. + decode_metadata = build_side_metadata(0, num_decodes) + prefill_metadata = build_side_metadata(num_decodes, num_reqs) + + return MiniMaxM3IndexerMSAMetadata( + seq_lens=seq_lens, + max_seq_len=common_attn_metadata.max_seq_len, + slot_mapping=common_attn_metadata.slot_mapping, + num_actual_tokens=num_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + kv_indices=kv_indices, + decode_pages=decode_pages, + decode_metadata=decode_metadata, + prefill_metadata=prefill_metadata, + ) + + +class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl): + """fmha_sm100 OnlyScore for the per-block scores; Triton top-k selection.""" + + indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerMSABackend + + def forward( + self, + index_query: torch.Tensor, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + from vllm.third_party.fmha_sm100.api import _fmha_sm100 + + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + return None, None # profiling run; caches unbound + md = attn_metadata[self.index_cache.prefix] + assert isinstance(md, MiniMaxM3IndexerMSAMetadata) + + num_tokens = md.num_actual_tokens + nd = md.num_decode_tokens + index_q = index_query[:num_tokens].view( + -1, self.num_index_heads, self.index_head_dim + ) + # Index-K cache (num_blocks, 128, D) -> paged MQA (num_blocks, 1, 128, D). + kv = self.index_cache.kv_cache + k_pages = kv.view(kv.shape[0], 1, PAGE_SIZE, self.index_head_dim) + kv_indices = md.kv_indices + + def score_topk( + meta: MiniMaxM3IndexerMSASubMetadata, + query: torch.Tensor, + page_table: torch.Tensor | None, + ) -> torch.Tensor: + # OnlyScore -> max_score [num_index_heads, max_k_tiles, num_tokens]. + _, max_score = _fmha_sm100( + query, + k_pages, + k_pages, # V placeholder; not read in OnlyScore + meta.plan, + kv_indices=page_table, + output_o=False, + output_maxscore=True, + sm_scale=self.scale, + ) + # Triton top-k wants [num_index_heads, num_tokens, max_block]; the + # transpose is a strided view (the kernel reads via strides). One + # 128-token KV tile == one M3 sparse block. + return minimax_m3_index_topk( + max_score.transpose(1, 2), + meta.cu_seqlens_q, + meta.prefix_lens, + meta.max_query_len, + self.topk_blocks, + self.init_blocks, + self.local_blocks, + ) + + def run_decode() -> torch.Tensor | None: + decode_metadata = md.decode_metadata + if decode_metadata is None: + return None + decode_pages = ( + kv_indices[: md.decode_pages] if kv_indices is not None else None + ) + return score_topk(decode_metadata, index_q[:nd], decode_pages) + + def run_prefill() -> torch.Tensor | None: + prefill_metadata = md.prefill_metadata + if prefill_metadata is None: + return None + prefill_pages = ( + kv_indices[md.decode_pages :] if kv_indices is not None else None + ) + return score_topk(prefill_metadata, index_q[nd:], prefill_pages) + + return run_decode(), run_prefill() diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index aaced78ed7c..a2b58acd632 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -576,7 +576,12 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase): main_slot_mapping = fwd_slot_mapping[self.layer_name] index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix] q = qkv.new_empty((num_tokens, self.q_size)) - index_q = qkv.new_empty((num_tokens, self.index_q_size)) + # index_q matches the index-K cache dtype (e4m3 for the fp8 score path); + # the fused kernel emits fp8 directly when this buffer is e4m3. + index_q = qkv.new_empty( + (num_tokens, self.index_q_size), + dtype=self.indexer.index_cache.dtype, + ) ops.fused_minimax_m3_qknorm_rope_kv_insert( qkv, self.q_norm.weight,