A high-throughput and memory-efficient inference and serving engine for LLMs
[M3] Enable FP8 sparse GQA (#45744)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Thien Tran committed
Jun 17, 2026 at 12:38 UTC
4c626633159887b0f2c962058c17c78f1434556d
10 files changed
+179
-178
cmake/external_projects/fmha_sm100.cmake
+8
-10
@@ -32,19 +32,17 @@ message(STATUS "fmha_sm100 is available at ${fmha_sm100_SOURCE_DIR}")
32
33
add_custom_target(fmha_sm100)
34
35
+set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100")
36
+
37
install(FILES
36
- "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100/__init__.py"
37
- "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100/sparse.py"
38
+ "${FMHA_SM100_PY_ROOT}/__init__.py"
39
+ "${FMHA_SM100_PY_ROOT}/sparse.py"
40
DESTINATION vllm/third_party/fmha_sm100
41
COMPONENT fmha_sm100)
42
41
-install(DIRECTORY "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100/cute/"
43
+install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/"
44
DESTINATION vllm/third_party/fmha_sm100/cute
45
COMPONENT fmha_sm100
44
- FILES_MATCHING
45
- REGEX "/__pycache__(/.*)?$" EXCLUDE
46
- REGEX ".*\\.pyc$" EXCLUDE
47
- PATTERN "example.py" EXCLUDE
48
- PATTERN "test_*.py" EXCLUDE
49
- PATTERN "*.py"
50
- PATTERN "build_k2q_csr.cu")
46
+ PATTERN "__pycache__" EXCLUDE
47
+ PATTERN "*.pyc" EXCLUDE
48
+ PATTERN ".git*" EXCLUDE)
csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
+80
-40
@@ -58,8 +58,15 @@
58
59
#include "../cuda_compat.h"
60
#include "../type_convert.cuh"
61
+#include "../attention/dtype_fp8.cuh"
62
#include "dispatch_utils.h"
63
64
+#ifdef USE_ROCM
65
+ #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
66
+#else
67
+ #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
68
+#endif
69
+
70
#ifndef FINAL_MASK
71
#ifdef USE_ROCM
72
#define FINAL_MASK 0xffffffffffffffffULL
@@ -186,6 +193,21 @@ __device__ __forceinline__ void storeElems(
193
*reinterpret_cast<uint2*>(dst) = v;
194
}
195
196
+template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
197
+__device__ __forceinline__ void storeCacheElems(
198
+ cache_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
199
+ if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
200
+ // kAuto means unquantized KV cache here: cache_t == scalar_t, so store the
201
+ // model dtype directly. FP8 cache dtypes use the conversion path below.
202
+ storeElems<scalar_t>(reinterpret_cast<scalar_t*>(dst), elems);
203
+ } else {
204
+#pragma unroll
205
+ for (int i = 0; i < kElemsPerLane; i++) {
206
+ dst[i] = fp8::scaled_convert<cache_t, float, kv_dt>(elems[i], 1.0f);
207
+ }
208
+ }
209
+}
210
+
211
// ────────────────────────────────────────────────────────────────────────────
212
// Kernel
213
// ────────────────────────────────────────────────────────────────────────────
@@ -202,7 +224,8 @@ __device__ __forceinline__ void storeElems(
224
// V : nkv only if kInsertKV (V-cache insert; no warps in dense)
225
// IQ: niq only if kIsSparse (norm+RoPE)
226
// IK: 1 only if kIsSparse (norm+RoPE; +index-cache insert)
205
-template <typename scalar_t, bool kIsSparse, bool kInsertKV>
227
+template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
228
+ bool kIsSparse, bool kInsertKV>
229
__global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
230
scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse)
231
scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
@@ -215,7 +238,7 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
238
int64_t const* __restrict__ positions, // [N] i64
239
int64_t const* __restrict__ slot_mapping, // main K/V slots or nullptr
240
int64_t const* __restrict__ index_slot_mapping, // index K slots/nullptr
218
- scalar_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr
241
+ cache_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr
242
scalar_t* __restrict__ index_cache, // [nb*bs, 128] or nullptr
243
float const eps, int const rotary_dim, int const num_tokens, int const nq,
244
int const nkv, int const niq, int const block_size,
@@ -355,7 +378,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
378
int const kv = isK ? 0 : 1;
379
int64_t const off =
380
b * kv_s_block + kv * kv_s_kv + t * kv_s_token + head * kv_s_head;
358
- storeElems<scalar_t>(kv_cache + off + dim_base, elems);
381
+ storeCacheElems<scalar_t, cache_t, kv_dt>(kv_cache + off + dim_base,
382
+ elems);
383
}
384
}
385
}
@@ -373,13 +397,13 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
397
// ────────────────────────────────────────────────────────────────────────────
398
// Launch wrapper
399
// ────────────────────────────────────────────────────────────────────────────
376
-template <typename scalar_t>
400
+template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
401
void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
402
scalar_t const* q_norm_w, scalar_t const* k_norm_w,
403
scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
404
scalar_t const* cos_sin_cache,
405
int64_t const* positions, int64_t const* slot_mapping,
382
- int64_t const* index_slot_mapping, scalar_t* kv_cache,
406
+ int64_t const* index_slot_mapping, cache_t* kv_cache,
407
scalar_t* index_cache, float const eps,
408
int const rotary_dim, int const num_tokens,
409
int const nq, int const nkv, int const niq,
@@ -419,7 +443,8 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
443
#define LAUNCH(IS_SPARSE, INSERT) \
444
cudaLaunchKernelEx( \
445
&config, \
422
- fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT>, \
446
+ fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, \
447
+ IS_SPARSE, INSERT>, \
448
qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, \
449
cos_sin_cache, positions, slot_mapping, index_slot_mapping, kv_cache, \
450
index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, \
@@ -428,7 +453,8 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
453
// ROCm: standard kernel launch syntax (no PDL/stream serialization).
454
// clang-format off
455
#define LAUNCH(IS_SPARSE, INSERT) \
431
- fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT> \
456
+ fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, \
457
+ IS_SPARSE, INSERT> \
458
<<<grid, kBlockSize, 0, stream>>>( \
459
qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, \
460
ik_norm_w, cos_sin_cache, positions, slot_mapping, \
@@ -455,6 +481,33 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
481
} // namespace minimax_m3_fused_ops
482
} // namespace vllm
483
484
+#define CALL_FUSED_MINIMAX_M3(_RAW_T, CACHE_T, KV_DTYPE) \
485
+ vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>( \
486
+ reinterpret_cast<st*>(qkv.data_ptr()), \
487
+ q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr()) : nullptr, \
488
+ index_q_out.has_value() ? reinterpret_cast<st*>(index_q_out->data_ptr()) \
489
+ : nullptr, \
490
+ reinterpret_cast<st const*>(q_norm_weight.data_ptr()), \
491
+ reinterpret_cast<st const*>(k_norm_weight.data_ptr()), \
492
+ has_index ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr()) \
493
+ : nullptr, \
494
+ has_index ? reinterpret_cast<st const*>(index_k_norm_weight->data_ptr()) \
495
+ : nullptr, \
496
+ reinterpret_cast<st const*>(cos_sin_cache.data_ptr()), \
497
+ reinterpret_cast<int64_t const*>(positions.data_ptr()), \
498
+ insert_kv ? reinterpret_cast<int64_t const*>(slot_mapping->data_ptr()) \
499
+ : nullptr, \
500
+ insert_kv ? reinterpret_cast<int64_t const*>( \
501
+ effective_index_slot_mapping->data_ptr()) \
502
+ : nullptr, \
503
+ insert_kv ? reinterpret_cast<CACHE_T*>(kv_cache->data_ptr()) : nullptr, \
504
+ (insert_kv && has_index) \
505
+ ? reinterpret_cast<st*>(index_cache->data_ptr()) \
506
+ : nullptr, \
507
+ static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq, \
508
+ nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv, kv_s_token, \
509
+ kv_s_head, has_index, insert_kv, stream)
510
+
511
// ────────────────────────────────────────────────────────────────────────────
512
// Torch op wrapper
513
// ────────────────────────────────────────────────────────────────────────────
@@ -475,9 +528,14 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
528
int64_t block_size,
529
std::optional<torch::stable::Tensor> q_out, // [N, nq*128] contiguous
530
std::optional<torch::stable::Tensor>
478
- index_q_out) { // [N, niq*128] contiguous
531
+ index_q_out, // [N, niq*128] contiguous
532
+ const std::string& kv_cache_dtype) {
533
STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
534
"qkv must be contiguous CUDA");
535
+ STD_TORCH_CHECK(
536
+ qkv.scalar_type() == torch::headeronly::ScalarType::Half ||
537
+ qkv.scalar_type() == torch::headeronly::ScalarType::BFloat16,
538
+ "qkv must be float16 or bfloat16");
539
STD_TORCH_CHECK(
540
positions.is_cuda() &&
541
positions.scalar_type() == torch::headeronly::ScalarType::Long,
@@ -510,6 +568,8 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
568
// (1 head)]) right after [q|k|v] in the same row; the dense layer does not.
569
bool const has_index = niq > 0;
570
bool const insert_kv = kv_cache.has_value();
571
+ vllm::Fp8KVCacheDataType const kv_dt =
572
+ vllm::get_fp8_kv_cache_data_type(kv_cache_dtype);
573
int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
574
int const expected_row =
575
(nq + 2 * nkv + (has_index ? niq + 1 : 0)) * kHeadDim;
@@ -552,8 +612,14 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
612
torch::headeronly::ScalarType::Long &&
613
index_slot_mapping->numel() == slot_mapping->numel()),
614
"index_slot_mapping must be int64 CUDA with slot_mapping length");
555
- STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
556
- "kv_cache dtype must match qkv (bf16 cache only)");
615
+ if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
616
+ STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
617
+ "auto kv_cache dtype must match qkv");
618
+ } else {
619
+ STD_TORCH_CHECK(
620
+ kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte,
621
+ "fp8 kv_cache must use uint8 storage");
622
+ }
623
STD_TORCH_CHECK(index_cache.has_value() &&
624
index_cache->scalar_type() == qkv.scalar_type(),
625
"insert mode requires matching index_cache");
@@ -601,35 +667,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
667
VLLM_STABLE_DISPATCH_HALF_TYPES(
668
qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert", [&] {
669
using st = scalar_t;
604
- vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st>(
605
- reinterpret_cast<st*>(qkv.data_ptr()),
606
- q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr())
607
- : nullptr,
608
- index_q_out.has_value()
609
- ? reinterpret_cast<st*>(index_q_out->data_ptr())
610
- : nullptr,
611
- reinterpret_cast<st const*>(q_norm_weight.data_ptr()),
612
- reinterpret_cast<st const*>(k_norm_weight.data_ptr()),
613
- has_index
614
- ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr())
615
- : nullptr,
616
- has_index
617
- ? reinterpret_cast<st const*>(index_k_norm_weight->data_ptr())
618
- : nullptr,
619
- reinterpret_cast<st const*>(cos_sin_cache.data_ptr()),
620
- reinterpret_cast<int64_t const*>(positions.data_ptr()),
621
- insert_kv
622
- ? reinterpret_cast<int64_t const*>(slot_mapping->data_ptr())
623
- : nullptr,
624
- insert_kv ? reinterpret_cast<int64_t const*>(
625
- effective_index_slot_mapping->data_ptr())
626
- : nullptr,
627
- insert_kv ? reinterpret_cast<st*>(kv_cache->data_ptr()) : nullptr,
628
- (insert_kv && has_index)
629
- ? reinterpret_cast<st*>(index_cache->data_ptr())
630
- : nullptr,
631
- static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens,
632
- nq, nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv,
633
- kv_s_token, kv_s_head, has_index, insert_kv, stream);
670
+ DISPATCH_BY_KV_CACHE_DTYPE(qkv.scalar_type(), kv_cache_dtype,
671
+ CALL_FUSED_MINIMAX_M3);
672
});
673
}
674
+
675
+#undef CALL_FUSED_MINIMAX_M3
csrc/libtorch_stable/ops.h
+5
-1
@@ -3,6 +3,9 @@
3
#include <torch/csrc/stable/library.h>
4
#include <torch/csrc/stable/tensor.h>
5
6
+#include <optional>
7
+#include <string>
8
+
9
void per_token_group_quant_fp8(const torch::stable::Tensor& input,
10
torch::stable::Tensor& output_q,
11
torch::stable::Tensor& output_s,
@@ -297,7 +300,8 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
300
std::optional<torch::stable::Tensor> kv_cache,
301
std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
302
std::optional<torch::stable::Tensor> q_out,
300
- std::optional<torch::stable::Tensor> index_q_out);
303
+ std::optional<torch::stable::Tensor> index_q_out,
304
+ const std::string& kv_cache_dtype);
305
306
// Sampler kernels (shared CUDA/ROCm)
307
void apply_repetition_penalties_(
csrc/libtorch_stable/torch_bindings.cpp
+2
-1
@@ -471,7 +471,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
471
"int num_index_heads, "
472
"Tensor? slot_mapping, Tensor? index_slot_mapping, "
473
"Tensor!? kv_cache, Tensor!? index_cache, "
474
- "int block_size, Tensor!? q_out, Tensor!? index_q_out) -> ()");
474
+ "int block_size, Tensor!? q_out, Tensor!? index_q_out, "
475
+ "str kv_cache_dtype) -> ()");
476
477
// Apply repetition penalties to logits in-place.
478
ops.def(
setup.py
+1
-1
@@ -1168,7 +1168,7 @@ package_data = {
1168
"third_party/deep_gemm/include/**/*.h",
1169
"third_party/deep_gemm/include/**/*.hpp",
1170
# fmha_sm100 sparse CuTe-DSL helper kernels (vendored via cmake)
1171
- "third_party/fmha_sm100/cute/src/sm100/build_k2q_csr/build_k2q_csr.cu",
1171
+ "third_party/fmha_sm100/cute/**/*.cu",
1172
]
1173
}
1174
tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
+49
-13
@@ -106,7 +106,16 @@ def test_dense_norm_rope(num_tokens, num_heads, num_kv_heads):
106
qkv_orig = qkv.clone()
107
108
ops.fused_minimax_m3_qknorm_rope_kv_insert(
109
- qkv, q_w, k_w, cos_sin, positions, num_heads, num_kv_heads, ROTARY_DIM, eps
109
+ qkv,
110
+ q_w,
111
+ k_w,
112
+ cos_sin,
113
+ positions,
114
+ num_heads,
115
+ num_kv_heads,
116
+ ROTARY_DIM,
117
+ eps,
118
+ kv_cache_dtype="auto",
119
)
120
q_out, k_out, v_out = qkv.split([qsz, kvsz, kvsz], dim=-1)
121
@@ -133,7 +142,8 @@ def test_dense_norm_rope(num_tokens, num_heads, num_kv_heads):
142
143
@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
144
@pytest.mark.parametrize("block_size", [16, 64])
136
-def test_sparse_full(num_tokens, block_size):
145
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
146
+def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
147
torch.manual_seed(1)
148
device, dtype, eps = "cuda", torch.bfloat16, 1e-6
149
base, max_pos = 5_000_000.0, 4096
@@ -158,8 +168,15 @@ def test_sparse_full(num_tokens, block_size):
168
splits = [qsz, kvsz, kvsz, iqsz, iksz]
169
170
num_blocks = (num_tokens + block_size - 1) // block_size + 1
171
+ kv_cache_storage_dtype = torch.uint8 if kv_cache_dtype == "fp8" else dtype
172
kv_cache = torch.zeros(
162
- num_blocks, 2, block_size, num_kv_heads, HEAD_DIM, dtype=dtype, device=device
173
+ num_blocks,
174
+ 2,
175
+ block_size,
176
+ num_kv_heads,
177
+ HEAD_DIM,
178
+ dtype=kv_cache_storage_dtype,
179
+ device=device,
180
)
181
index_cache = torch.zeros(
182
num_blocks, block_size, HEAD_DIM, dtype=dtype, device=device
@@ -195,11 +212,12 @@ def test_sparse_full(num_tokens, block_size):
212
block_size,
213
q_out,
214
index_q,
215
+ kv_cache_dtype,
216
)
217
218
# ── norm+rope parity. q/index_q land in their gather buffers; k/index_k are
219
# rewritten in place inside qkv. ──
202
- _, k_out, _, _, index_k = qkv.split(splits, dim=-1)
220
+ _, k_out, v_out, _, index_k = qkv.split(splits, dim=-1)
221
q_in, k_in, v_in, iq_orig, ik_orig = qkv_orig.split(splits, dim=-1)
222
q_ref = norm_rope_ref(
223
q_in.view(num_tokens, num_heads, HEAD_DIM), q_w, positions, cos_sin, eps
@@ -230,15 +248,33 @@ def test_sparse_full(num_tokens, block_size):
248
# ── 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].
233
- idx_flat = index_cache.view(num_blocks * block_size, HEAD_DIM)
251
k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
252
v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM) # v is raw (no norm/rope)
236
- for t in range(num_tokens):
237
- s = slot_mapping[t].item()
238
- b, pos = s // block_size, s % block_size
239
- torch.testing.assert_close(
240
- kv_cache[b, 0, pos], k_ref_h[t], rtol=1e-2, atol=1e-2
253
+ if kv_cache_dtype == "fp8":
254
+ expected_kv_cache = torch.zeros_like(kv_cache)
255
+ scale = torch.ones((), device=device)
256
+ ops.reshape_and_cache_flash(
257
+ k_out.view(num_tokens, num_kv_heads, HEAD_DIM),
258
+ v_out.view(num_tokens, num_kv_heads, HEAD_DIM),
259
+ expected_kv_cache[:, 0],
260
+ expected_kv_cache[:, 1],
261
+ slot_mapping,
262
+ kv_cache_dtype,
263
+ scale,
264
+ scale,
265
)
242
- torch.testing.assert_close(kv_cache[b, 1, pos], v_ref_h[t], rtol=0, atol=0)
243
- index_s = index_slot_mapping[t].item()
244
- torch.testing.assert_close(idx_flat[index_s], ik_ref[t], rtol=1e-2, atol=1e-2)
266
+ torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
267
+ else:
268
+ for t in range(num_tokens):
269
+ s = slot_mapping[t].item()
270
+ b, pos = s // block_size, s % block_size
271
+ torch.testing.assert_close(
272
+ kv_cache[b, 0, pos], k_ref_h[t], rtol=1e-2, atol=1e-2
273
+ )
274
+ torch.testing.assert_close(kv_cache[b, 1, pos], v_ref_h[t], rtol=0, atol=0)
275
+
276
+ expected_index_cache = torch.zeros_like(index_cache).view(-1, HEAD_DIM)
277
+ expected_index_cache[index_slot_mapping] = index_k
278
+ torch.testing.assert_close(
279
+ index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0
280
+ )
vllm/_custom_ops.py
+5
-2
@@ -2634,6 +2634,7 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
2634
block_size: int = 0,
2635
q_out: torch.Tensor | None = None,
2636
index_q_out: torch.Tensor | None = None,
2637
+ kv_cache_dtype: str = "auto",
2638
) -> None:
2639
"""Fused MiniMax-M3 attention pre-processing (in-place).
2640
@@ -2645,8 +2646,9 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
2646
index_k]`` — the index branch is read straight out of ``qkv``.
2647
2648
When ``kv_cache`` is given (sparse serving), also scatter-inserts the
2648
- normed/roped k & v into the paged bf16 KV cache by ``slot_mapping`` and the
2649
- index key into ``index_cache`` by ``index_slot_mapping``. If
2649
+ normed/roped k & v into the paged KV cache by ``slot_mapping`` and the
2650
+ index key into ``index_cache`` by ``index_slot_mapping``. ``kv_cache_dtype``
2651
+ selects the cache storage/conversion path. If
2652
``index_slot_mapping`` is omitted, ``slot_mapping`` is used for both caches.
2653
2654
If ``q_out`` / ``index_q_out`` (contiguous ``[N, nq*128]`` / ``[N,
@@ -2675,6 +2677,7 @@ def fused_minimax_m3_qknorm_rope_kv_insert(
2677
block_size,
2678
q_out,
2679
index_q_out,
2680
+ kv_cache_dtype,
2681
)
2682
2683
vllm/models/minimax_m3/amd/model.py
+7
-63
@@ -426,6 +426,7 @@ class MiniMaxM3Attention(nn.Module):
426
self.num_kv_heads,
427
self.rotary_emb.rotary_dim,
428
self.q_norm.variance_epsilon,
429
+ kv_cache_dtype="auto",
430
)
431
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
432
attn_output = self.attn(q, k, v)
@@ -533,11 +534,6 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
534
self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
535
self.kv_cache_dtype, vllm_config.model_config
536
)
536
- # fp8 main-K/V cache: the fused qknorm+rope+kv-insert op is bf16-cache-only
537
- # (asserts kv_cache dtype == qkv), so on the fp8 path we run it in
538
- # norm+rope-only mode and write the cache via the fp8-capable
539
- # reshape_and_cache_flash in _insert_kv. (index cache stays bf16.)
540
- self._fp8_kv = "fp8" in self.kv_cache_dtype
537
538
self.attn_backend = MiniMaxM3SparseBackend
539
# Indexer and main attention are separate impls. On ROCm the SM100 gate
@@ -592,37 +588,6 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
588
kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
589
)
590
595
- def _insert_kv(
596
- self,
597
- key: torch.Tensor,
598
- value: torch.Tensor,
599
- index_key: torch.Tensor,
600
- main_slot_mapping: torch.Tensor,
601
- index_slot_mapping: torch.Tensor,
602
- ) -> None:
603
- """Write main K/V (fp8-quantizing) and index-K into their paged caches.
604
-
605
- Used only on the fp8-KV path: the fused #20 op is bf16-cache-only, so it
606
- runs in norm+rope-only mode and the (already normed/roped) k/v/index_k are
607
- written here via ``reshape_and_cache_flash`` (which honors kv_cache_dtype,
608
- unit scale -- matching the fp8 read path added in #33). Mirrors the
609
- pre-#20 unfused insert. The index cache stays bf16 (no quant).
610
- """
611
- key_cache, value_cache = self.kv_cache.unbind(1)
612
- scale = torch.ones((), device=key.device)
613
- ops.reshape_and_cache_flash(
614
- key.view(-1, self.num_kv_heads, self.head_dim),
615
- value.view(-1, self.num_kv_heads, self.head_dim),
616
- key_cache,
617
- value_cache,
618
- main_slot_mapping,
619
- self.kv_cache_dtype,
620
- scale,
621
- scale,
622
- )
623
- idx_cache = self.indexer.index_cache.kv_cache.view(-1, self.idx_head_dim)
624
- idx_cache[index_slot_mapping] = index_key.to(idx_cache.dtype)
625
-
591
def forward(
592
self,
593
positions: torch.Tensor,
@@ -636,11 +601,9 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
601
# of the single fused ``qkv`` tensor. Once the paged caches are bound the
602
# kernel also inserts k/v and the index key into them (each with its own
603
# slot_mapping); the memory-profiling run (caches unbound, no slot_mapping)
639
- # short-circuits to zeros below. Replaces the
640
- # q_norm/k_norm/rotary_emb/index_*_norm/index_rotary_emb/_insert_kv chain.
641
- # (#20 fused_minimax_m3_qknorm_rope_kv_insert; HIP/CDNA path. The main and
642
- # index slot mappings are read from the forward context's slot_mapping
643
- # dict, matching the breakable-cudagraph path -- see nvidia/model.py.)
604
+ # short-circuits to zeros below. The main and index slot mappings are read
605
+ # from the forward context's slot_mapping dict, matching the
606
+ # breakable-cudagraph path -- see nvidia/model.py.
607
cos_sin_cache = self.rotary_emb.cos_sin_cache
608
rotary_dim = self.rotary_emb.rotary_dim
609
eps = self.q_norm.variance_epsilon
@@ -658,12 +621,6 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
621
index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
622
q = qkv.new_empty((num_tokens, self.q_size))
623
index_q = qkv.new_empty((num_tokens, self.index_q_size))
661
- # On the fp8-KV path the fused op cannot write the (fp8) cache, so pass
662
- # kv_cache/index_cache = None -> insert_kv=False (norm+rope only): it still
663
- # de-interleaves q/index_q and rewrites the normed/roped k & index_k in
664
- # place in qkv, leaving v raw (correct -- v is never normed/roped). We then
665
- # write the cache via _insert_kv below.
666
- insert_via_fused = not self._fp8_kv
624
ops.fused_minimax_m3_qknorm_rope_kv_insert(
625
qkv,
626
self.q_norm.weight,
@@ -679,26 +636,13 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
636
self.num_idx_heads,
637
main_slot_mapping,
638
index_slot_mapping,
682
- self.kv_cache if insert_via_fused else None,
683
- self.indexer.index_cache.kv_cache if insert_via_fused else None,
639
+ self.kv_cache,
640
+ self.indexer.index_cache.kv_cache,
641
self.kv_cache.size(2), # paged-cache block size
642
q,
643
index_q,
644
+ self.kv_cache_dtype,
645
)
688
- if not insert_via_fused:
689
- # Extract the normed/roped k, raw v, normed/roped index_k from qkv
690
- # ([q | k | v | index_q | index_k], all head_dim=128) and fp8-insert.
691
- kv = self.num_kv_heads * self.head_dim
692
- # These are strided views into qkv (row stride = full qkv width), but
693
- # their last dim is contiguous, so `_insert_kv`'s `.view(-1, nkv,
694
- # head_dim)` works on them and `reshape_and_cache_flash` honors the
695
- # input stride -- no `.contiguous()` needed (verified bit-identical;
696
- # avoids a [N, kv] copy per step on the fp8-KV path).
697
- k = qkv[:, self.q_size : self.q_size + kv]
698
- v = qkv[:, self.q_size + kv : self.q_size + 2 * kv]
699
- ik0 = self.q_size + 2 * kv + self.index_q_size
700
- index_k = qkv[:, ik0 : ik0 + self.num_idx_heads * self.idx_head_dim]
701
- self._insert_kv(k, v, index_k, main_slot_mapping, index_slot_mapping)
646
647
output = torch.empty_like(q)
648
attn_output = self._run_attention(q, index_q, output)
vllm/models/minimax_m3/common/sparse_attention.py
+16
-4
@@ -22,6 +22,7 @@ import torch
22
from vllm.config import VllmConfig
23
from vllm.config.cache import CacheDType
24
from vllm.forward_context import get_forward_context
25
+from vllm.logger import init_logger
26
from vllm.models.minimax_m3.common.ops.sparse_attn import (
27
SPARSE_BLOCK_SIZE,
28
minimax_m3_sparse_attn,
@@ -44,6 +45,8 @@ from vllm.v1.attention.backends.utils import (
45
)
46
from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
47
48
+logger = init_logger(__name__)
49
+
50
51
class MiniMaxM3SparseBackend(AttentionBackend):
52
"""Block-sparse GQA backend for MiniMax M3 sparse attention layers."""
@@ -385,16 +388,25 @@ def select_main_impl_cls(
388
) -> type[MiniMaxM3SparseImpl]:
389
"""Pick the main attend impl off the main KV-cache dtype.
390
388
- bf16 on Blackwell (SM100) uses the MSA attend; fp8 or non-Blackwell falls
391
+ Blackwell (SM100) uses the MSA attend for supported top-k block counts
392
+ when the KV cache is BF16 or FP8 E4M3; non-Blackwell and FP8 E5M2 fall
393
back to Triton. The MSA module is imported lazily so AMD/non-SM100 never
394
import fmha_sm100.
395
"""
392
- if (
396
+ use_msa = (
397
current_platform.is_cuda()
398
and current_platform.is_device_capability_family(100)
399
and topk_blocks in (4, 8, 16, 32)
396
- and not is_quantized_kv_cache(kv_cache_dtype)
397
- ):
400
+ and kv_cache_dtype != "fp8_e5m2"
401
+ )
402
+ selected = "MSA" if use_msa else "Triton"
403
+ logger.info_once(
404
+ "MiniMax M3 sparse attention selected %s (kv_cache_dtype=%s, topk_blocks=%s)",
405
+ selected,
406
+ kv_cache_dtype,
407
+ topk_blocks,
408
+ )
409
+ if use_msa:
410
from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
411
MiniMaxM3SparseMSAImpl,
412
)
vllm/models/minimax_m3/nvidia/model.py
+6
-43
@@ -66,10 +66,7 @@ from vllm.model_executor.models.utils import (
66
maybe_prefix,
67
)
68
from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model
69
-from vllm.models.minimax_m3.common.indexer import (
70
- MiniMaxM3Indexer,
71
- MiniMaxM3IndexerMetadata,
72
-)
69
+from vllm.models.minimax_m3.common.indexer import MiniMaxM3Indexer
70
from vllm.models.minimax_m3.common.mm_preprocess import (
71
MiniMaxM3VLDummyInputsBuilder,
72
MiniMaxM3VLMultiModalProcessor,
@@ -78,7 +75,6 @@ from vllm.models.minimax_m3.common.mm_preprocess import (
75
from vllm.models.minimax_m3.common.sparse_attention import (
76
MiniMaxM3SparseBackend,
77
MiniMaxM3SparseImpl,
81
- MiniMaxM3SparseMetadata,
78
select_main_impl_cls,
79
)
80
from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
@@ -375,6 +371,7 @@ class MiniMaxM3Attention(nn.Module):
371
self.num_kv_heads,
372
self.rotary_emb.rotary_dim,
373
self.q_norm.variance_epsilon,
374
+ kv_cache_dtype="auto",
375
)
376
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
377
attn_output = self.attn(q, k, v)
@@ -545,40 +542,6 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
542
kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
543
)
544
548
- def _insert_kv(
549
- self, key: torch.Tensor, value: torch.Tensor, index_key: torch.Tensor
550
- ) -> None:
551
- """Write main K/V and index-K into their paged caches.
552
-
553
- No-op during the profiling run, where caches are not yet bound and
554
- ``attn_metadata`` is None.
555
- """
556
- attn_metadata = get_forward_context().attn_metadata
557
- if not isinstance(attn_metadata, dict):
558
- return
559
- main_meta = attn_metadata[self.layer_name]
560
- index_meta = attn_metadata[self.indexer.index_cache.prefix]
561
- assert isinstance(main_meta, MiniMaxM3SparseMetadata)
562
- assert isinstance(index_meta, MiniMaxM3IndexerMetadata)
563
-
564
- # Identity scale: unused for the bf16 cache, required arg of the op.
565
- key_cache, value_cache = self.kv_cache.unbind(1)
566
- scale = torch.ones((), device=key.device)
567
- ops.reshape_and_cache_flash(
568
- key.view(-1, self.num_kv_heads, self.head_dim),
569
- value.view(-1, self.num_kv_heads, self.head_dim),
570
- key_cache,
571
- value_cache,
572
- main_meta.slot_mapping,
573
- self.kv_cache_dtype,
574
- scale,
575
- scale,
576
- )
577
-
578
- # Index-key cache: single vector per token, scatter by slot.
579
- idx_cache = self.indexer.index_cache.kv_cache.view(-1, self.idx_head_dim)
580
- idx_cache[index_meta.slot_mapping] = index_key.to(idx_cache.dtype)
581
-
545
def forward(
546
self,
547
positions: torch.Tensor,
@@ -592,10 +555,9 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
555
# of the single fused ``qkv`` tensor (the "5 results"). Once the paged
556
# caches are bound the kernel also inserts k/v and the index key into
557
# them; the initial memory-profiling run (caches unbound, no slot_mapping)
595
- # short-circuits to zeros below. Replaces the
596
- # q_norm/k_norm/rotary_emb/index_*_norm/index_rotary_emb/_insert_kv
597
- # sequence. k/v and index_k are rewritten in place inside qkv (and
598
- # scatter-inserted into the caches); q and index_q are de-interleaved
558
+ # short-circuits to zeros below. k/v and index_k are rewritten in place
559
+ # inside qkv (and scatter-inserted into the caches); q and index_q are
560
+ # de-interleaved
561
# straight into the dedicated contiguous ``q``/``index_q`` buffers below.
562
563
cos_sin_cache = self.rotary_emb.cos_sin_cache
@@ -635,6 +597,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
597
self.kv_cache.size(2), # paged-cache block size
598
q,
599
index_q,
600
+ self.kv_cache_dtype,
601
)
602
603
output = torch.empty_like(q)