@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0a1c5034f5
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Model] Add MiniMax M3 support (#45381)

Signed-off-by: youkaichao <youkaichao@gmail.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Signed-off-by: functionstackx <47992694+functionstackx@users.noreply.github.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: functionstackx <47992694+functionstackx@users.noreply.github.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>

youkaichao committed Jun 16, 2026 at 01:01 UTC 0a1c5034f5e4fe736db672010cda33d9d850f87e
108 files changed +14746 -323
.gitignore
+3
@@ -15,6 +15,9 @@ vllm/third_party/flashmla/flash_mla_interface.py
15 # DeepGEMM vendored package built from source
16 vllm/third_party/deep_gemm/
17
18 +# fmha_sm100 vendored package built from source
19 +vllm/third_party/fmha_sm100/
20 +
21 # triton jit
22 .triton
23
CMakeLists.txt
+2
@@ -440,6 +440,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
440 "csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
441 "csrc/libtorch_stable/pos_encoding_kernels.cu"
442 "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
443 + "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu"
444 "csrc/libtorch_stable/layernorm_kernels.cu"
445 "csrc/libtorch_stable/layernorm_quant_kernels.cu"
446 "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
@@ -1398,6 +1399,7 @@ endif()
1399 # For CUDA we also build and ship some external projects.
1400 if (VLLM_GPU_LANG STREQUAL "CUDA")
1401 include(cmake/external_projects/deepgemm.cmake)
1402 + include(cmake/external_projects/fmha_sm100.cmake)
1403 include(cmake/external_projects/flashmla.cmake)
1404 include(cmake/external_projects/qutlass.cmake)
1405
cmake/external_projects/fmha_sm100.cmake new
+50
@@ -0,0 +1,50 @@
1 +include(FetchContent)
2 +
3 +# If FMHA_SM100_SRC_DIR is set, fmha_sm100 is installed from that directory
4 +# instead of downloading. This is useful for local MSA development.
5 +if(DEFINED ENV{FMHA_SM100_SRC_DIR})
6 + set(FMHA_SM100_SRC_DIR $ENV{FMHA_SM100_SRC_DIR})
7 +endif()
8 +
9 +if(FMHA_SM100_SRC_DIR)
10 + FetchContent_Declare(
11 + fmha_sm100
12 + SOURCE_DIR ${FMHA_SM100_SRC_DIR}
13 + CONFIGURE_COMMAND ""
14 + BUILD_COMMAND ""
15 + )
16 +else()
17 + FetchContent_Declare(
18 + fmha_sm100
19 + GIT_REPOSITORY https://github.com/vllm-project/MSA.git
20 + GIT_TAG 544eee5e09ae2dfa774d5b06739013f9b7402c57
21 + GIT_PROGRESS TRUE
22 + CONFIGURE_COMMAND ""
23 + BUILD_COMMAND ""
24 + )
25 +endif()
26 +
27 +FetchContent_GetProperties(fmha_sm100)
28 +if(NOT fmha_sm100_POPULATED)
29 + FetchContent_Populate(fmha_sm100)
30 +endif()
31 +message(STATUS "fmha_sm100 is available at ${fmha_sm100_SOURCE_DIR}")
32 +
33 +add_custom_target(fmha_sm100)
34 +
35 +install(FILES
36 + "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100/__init__.py"
37 + "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100/sparse.py"
38 + DESTINATION vllm/third_party/fmha_sm100
39 + COMPONENT fmha_sm100)
40 +
41 +install(DIRECTORY "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100/cute/"
42 + DESTINATION vllm/third_party/fmha_sm100/cute
43 + 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")
csrc/libtorch_stable/activation_kernels.cu
+77 -41
@@ -10,11 +10,20 @@
10
11 namespace vllm {
12
13 -template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
13 +// `alpha` and `beta` are applied to opposite operands:
14 +// - alpha lives INSIDE the activation (the activated half): the gated
15 +// activation computes act_half * sigmoid(alpha * act_half).
16 +// - beta is added to the OTHER (non-activated) half before the multiply.
17 +// So the result is always ACT(act_half, alpha) * (other_half + beta).
18 +// Which half is which depends on `act_first` (see below). Defaults
19 +// alpha=1.0, beta=0.0 reproduce the plain SwiGLU/GeGLU behavior.
20 +template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float),
21 bool act_first, bool HAS_CLAMP>
22 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
23 const scalar_t& y,
17 - const float limit) {
24 + const float limit,
25 + const float alpha,
26 + const float beta) {
27 if constexpr (act_first) {
28 scalar_t gate = x;
29 scalar_t up = y;
@@ -22,7 +31,9 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
31 gate = (scalar_t)fminf((float)gate, limit);
32 up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
33 }
25 - return ACT_FN(gate) * up;
34 + // act_first: gate is the activated half -> alpha applies to gate;
35 + // beta is added to up (the non-activated half).
36 + return (scalar_t)(ACT_FN(gate, alpha) * ((float)up + beta));
37 } else {
38 scalar_t gate = x;
39 scalar_t up = y;
@@ -30,55 +41,68 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
41 gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
42 up = (scalar_t)fminf((float)up, limit);
43 }
33 - return gate * ACT_FN(up);
44 + // !act_first: up is the activated half -> alpha applies to up;
45 + // beta is added to gate (the non-activated half).
46 + return (scalar_t)(((float)gate + beta) * ACT_FN(up, alpha));
47 }
48 }
49
37 -template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
50 +template <typename packed_t,
51 + packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
52 bool act_first, bool HAS_CLAMP>
53 __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
54 const packed_t& y,
41 - const float limit) {
55 + const float limit,
56 + const float alpha,
57 + const float beta) {
58 if constexpr (act_first) {
59 packed_t gate = x;
60 packed_t up = y;
61 + float2 u = cast_to_float2(up);
62 if constexpr (HAS_CLAMP) {
63 float2 g = cast_to_float2(gate);
47 - float2 u = cast_to_float2(up);
64 g.x = fminf(g.x, limit);
65 g.y = fminf(g.y, limit);
66 u.x = fmaxf(fminf(u.x, limit), -limit);
67 u.y = fmaxf(fminf(u.y, limit), -limit);
68 gate = cast_to_packed<packed_t>(g);
53 - up = cast_to_packed<packed_t>(u);
69 }
55 - return packed_mul(PACKED_ACT_FN(gate), up);
70 + // act_first: gate is the activated half -> alpha applies to gate;
71 + // beta is added to up (the non-activated half).
72 + float2 activated = cast_to_float2(PACKED_ACT_FN(gate, alpha));
73 + activated.x *= u.x + beta;
74 + activated.y *= u.y + beta;
75 + return cast_to_packed<packed_t>(activated);
76 } else {
77 packed_t gate = x;
78 packed_t up = y;
79 + float2 g = cast_to_float2(gate);
80 if constexpr (HAS_CLAMP) {
60 - float2 g = cast_to_float2(gate);
81 float2 u = cast_to_float2(up);
82 g.x = fmaxf(fminf(g.x, limit), -limit);
83 g.y = fmaxf(fminf(g.y, limit), -limit);
84 u.x = fminf(u.x, limit);
85 u.y = fminf(u.y, limit);
66 - gate = cast_to_packed<packed_t>(g);
86 up = cast_to_packed<packed_t>(u);
87 }
69 - return packed_mul(gate, PACKED_ACT_FN(up));
88 + // !act_first: up is the activated half -> alpha applies to up;
89 + // beta is added to gate (the non-activated half).
90 + float2 activated = cast_to_float2(PACKED_ACT_FN(up, alpha));
91 + activated.x *= g.x + beta;
92 + activated.y *= g.y + beta;
93 + return cast_to_packed<packed_t>(activated);
94 }
95 }
96
97 // Activation and gating kernel template.
98 template <typename scalar_t, typename packed_t,
75 - scalar_t (*ACT_FN)(const scalar_t&),
76 - packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
77 - bool use_vec, bool HAS_CLAMP, bool use_256b = false>
99 + scalar_t (*ACT_FN)(const scalar_t&, const float),
100 + packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
101 + bool act_first, bool use_vec, bool HAS_CLAMP, bool use_256b = false>
102 __global__ void act_and_mul_kernel(
103 scalar_t* __restrict__ out, // [..., d]
104 const scalar_t* __restrict__ input, // [..., 2, d]
81 - const int d, const float limit) {
105 + const int d, const float limit, const float alpha, const float beta) {
106 const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
107 const scalar_t* y_ptr = x_ptr + d;
108 scalar_t* out_ptr = out + blockIdx.x * d;
@@ -105,7 +129,7 @@ __global__ void act_and_mul_kernel(
129 for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
130 x.elts[j] =
131 packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
108 - x.elts[j], y.elts[j], limit);
132 + x.elts[j], y.elts[j], limit, alpha, beta);
133 }
134 if constexpr (use_256b) {
135 st256(x, &out_vec[i]);
@@ -118,29 +142,34 @@ __global__ void act_and_mul_kernel(
142 for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
143 const scalar_t x = VLLM_LDG(&x_ptr[idx]);
144 const scalar_t y = VLLM_LDG(&y_ptr[idx]);
121 - out_ptr[idx] =
122 - compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
145 + out_ptr[idx] = compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(
146 + x, y, limit, alpha, beta);
147 }
148 }
149 }
150
151 +// Gated activations take an `alpha` argument that scales the sigmoid input
152 +// (`x * sigmoid(alpha * x)`). alpha defaults to 1.0 at all call sites, which
153 +// is exactly SiLU; only the clamp path (silu_and_mul_with_clamp) passes a
154 +// non-default alpha. Activations that do not use alpha simply ignore it.
155 template <typename T>
128 -__device__ __forceinline__ T silu_kernel(const T& x) {
129 - // x * sigmoid(x)
130 - return (T)(((float)x) / (1.0f + expf((float)-x)));
156 +__device__ __forceinline__ T silu_kernel(const T& x, const float alpha) {
157 + // x * sigmoid(alpha * x)
158 + return (T)(((float)x) / (1.0f + expf((float)-x * alpha)));
159 }
160
161 template <typename packed_t>
134 -__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val) {
135 - // x * sigmoid(x)
162 +__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val,
163 + const float alpha) {
164 + // x * sigmoid(alpha * x)
165 float2 fval = cast_to_float2(val);
137 - fval.x = fval.x / (1.0f + expf(-fval.x));
138 - fval.y = fval.y / (1.0f + expf(-fval.y));
166 + fval.x = fval.x / (1.0f + expf(-fval.x * alpha));
167 + fval.y = fval.y / (1.0f + expf(-fval.y * alpha));
168 return cast_to_packed<packed_t>(fval);
169 }
170
171 template <typename T>
143 -__device__ __forceinline__ T gelu_kernel(const T& x) {
172 +__device__ __forceinline__ T gelu_kernel(const T& x, const float /*alpha*/) {
173 // Equivalent to PyTorch GELU with 'none' approximation.
174 // Refer to:
175 // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
@@ -150,7 +179,8 @@ __device__ __forceinline__ T gelu_kernel(const T& x) {
179 }
180
181 template <typename packed_t>
153 -__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
182 +__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val,
183 + const float /*alpha*/) {
184 // Equivalent to PyTorch GELU with 'none' approximation.
185 // Refer to:
186 // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
@@ -162,7 +192,8 @@ __device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
192 }
193
194 template <typename T>
165 -__device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
195 +__device__ __forceinline__ T gelu_tanh_kernel(const T& x,
196 + const float /*alpha*/) {
197 // Equivalent to PyTorch GELU with 'tanh' approximation.
198 // Refer to:
199 // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
@@ -176,7 +207,7 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
207
208 template <typename packed_t>
209 __device__ __forceinline__ packed_t
179 -packed_gelu_tanh_kernel(const packed_t& val) {
210 +packed_gelu_tanh_kernel(const packed_t& val, const float /*alpha*/) {
211 // Equivalent to PyTorch GELU with 'tanh' approximation.
212 // Refer to:
213 // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
@@ -202,7 +233,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
233 // clamped (max only) and up input is clamped (both sides) before the
234 // activation function is applied.
235 #define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \
205 - HAS_CLAMP, LIMIT) \
236 + HAS_CLAMP, LIMIT, ALPHA, BETA) \
237 auto dtype = input.scalar_type(); \
238 int d = input.size(-1) / 2; \
239 int64_t num_tokens = input.numel() / input.size(-1); \
@@ -230,7 +261,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
261 PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
262 ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
263 out.mutable_data_ptr<scalar_t>(), \
233 - input.const_data_ptr<scalar_t>(), d, LIMIT); \
264 + input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA); \
265 }); \
266 } else { \
267 VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
@@ -240,7 +271,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
271 PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
272 ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
273 out.mutable_data_ptr<scalar_t>(), \
243 - input.const_data_ptr<scalar_t>(), d, LIMIT); \
274 + input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA); \
275 }); \
276 } \
277 } else { \
@@ -252,7 +283,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
283 PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
284 ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
285 out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(), \
255 - d, LIMIT); \
286 + d, LIMIT, ALPHA, BETA); \
287 }); \
288 }
289
@@ -260,14 +291,18 @@ void silu_and_mul(torch::stable::Tensor& out, // [..., d]
291 torch::stable::Tensor& input) // [..., 2 * d]
292 {
293 LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
263 - true, false, 0.0f);
294 + true, false, 0.0f, 1.0f, 0.0f);
295 }
296
297 void silu_and_mul_clamp(torch::stable::Tensor& out, // [..., d]
298 torch::stable::Tensor& input, // [..., 2 * d]
268 - double limit) {
299 + double limit, double alpha, double beta) {
300 + // out = (gate.clamp(max=limit) * sigmoid(alpha * gate.clamp(max=limit)))
301 + // * (up.clamp(+-limit) + beta)
302 + // alpha=1.0, beta=0.0 reduce this to silu(gate) * up.
303 LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
270 - true, true, (float)limit);
304 + true, true, (float)limit, (float)alpha,
305 + (float)beta);
306 }
307
308 void mul_and_silu(torch::stable::Tensor& out, // [..., d]
@@ -276,21 +311,22 @@ void mul_and_silu(torch::stable::Tensor& out, // [..., d]
311 // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
312 // applies the silu to the latter half of the input.
313 LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
279 - false, false, 0.0f);
314 + false, false, 0.0f, 1.0f, 0.0f);
315 }
316
317 void gelu_and_mul(torch::stable::Tensor& out, // [..., d]
318 torch::stable::Tensor& input) // [..., 2 * d]
319 {
320 LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
286 - true, false, 0.0f);
321 + true, false, 0.0f, 1.0f, 0.0f);
322 }
323
324 void gelu_tanh_and_mul(torch::stable::Tensor& out, // [..., d]
325 torch::stable::Tensor& input) // [..., 2 * d]
326 {
292 - LAUNCH_ACTIVATION_GATE_KERNEL(
293 - vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
327 + LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
328 + vllm::packed_gelu_tanh_kernel, true, false,
329 + 0.0f, 1.0f, 0.0f);
330 }
331
332 namespace vllm {
csrc/libtorch_stable/fp32_router_gemm.cu
+44 -41
@@ -175,49 +175,52 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
175 }
176
177 // ---------------------------------------------------------------------------
178 -// Explicit instantiations: M=1..32, E=256, H=3072, for both input types
178 +// Explicit instantiations: M=1..32, for both input types, for the supported
179 +// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5] and (128, 6144) [MiniMax-M3].
180 // ---------------------------------------------------------------------------
181
181 -#define INSTANTIATE(T, M) \
182 - template void invokeFp32RouterGemm<T, M, 256, 3072>( \
183 - float*, T const*, float const*, cudaStream_t);
184 -
185 -#define INSTANTIATE_ALL(T) \
186 - INSTANTIATE(T, 1) \
187 - INSTANTIATE(T, 2) \
188 - INSTANTIATE(T, 3) \
189 - INSTANTIATE(T, 4) \
190 - INSTANTIATE(T, 5) \
191 - INSTANTIATE(T, 6) \
192 - INSTANTIATE(T, 7) \
193 - INSTANTIATE(T, 8) \
194 - INSTANTIATE(T, 9) \
195 - INSTANTIATE(T, 10) \
196 - INSTANTIATE(T, 11) \
197 - INSTANTIATE(T, 12) \
198 - INSTANTIATE(T, 13) \
199 - INSTANTIATE(T, 14) \
200 - INSTANTIATE(T, 15) \
201 - INSTANTIATE(T, 16) \
202 - INSTANTIATE(T, 17) \
203 - INSTANTIATE(T, 18) \
204 - INSTANTIATE(T, 19) \
205 - INSTANTIATE(T, 20) \
206 - INSTANTIATE(T, 21) \
207 - INSTANTIATE(T, 22) \
208 - INSTANTIATE(T, 23) \
209 - INSTANTIATE(T, 24) \
210 - INSTANTIATE(T, 25) \
211 - INSTANTIATE(T, 26) \
212 - INSTANTIATE(T, 27) \
213 - INSTANTIATE(T, 28) \
214 - INSTANTIATE(T, 29) \
215 - INSTANTIATE(T, 30) \
216 - INSTANTIATE(T, 31) \
217 - INSTANTIATE(T, 32)
218 -
219 -INSTANTIATE_ALL(float)
220 -INSTANTIATE_ALL(__nv_bfloat16)
182 +#define INSTANTIATE(T, M, E, H) \
183 + template void invokeFp32RouterGemm<T, M, E, H>(float*, T const*, \
184 + float const*, cudaStream_t);
185 +
186 +#define INSTANTIATE_ALL(T, E, H) \
187 + INSTANTIATE(T, 1, E, H) \
188 + INSTANTIATE(T, 2, E, H) \
189 + INSTANTIATE(T, 3, E, H) \
190 + INSTANTIATE(T, 4, E, H) \
191 + INSTANTIATE(T, 5, E, H) \
192 + INSTANTIATE(T, 6, E, H) \
193 + INSTANTIATE(T, 7, E, H) \
194 + INSTANTIATE(T, 8, E, H) \
195 + INSTANTIATE(T, 9, E, H) \
196 + INSTANTIATE(T, 10, E, H) \
197 + INSTANTIATE(T, 11, E, H) \
198 + INSTANTIATE(T, 12, E, H) \
199 + INSTANTIATE(T, 13, E, H) \
200 + INSTANTIATE(T, 14, E, H) \
201 + INSTANTIATE(T, 15, E, H) \
202 + INSTANTIATE(T, 16, E, H) \
203 + INSTANTIATE(T, 17, E, H) \
204 + INSTANTIATE(T, 18, E, H) \
205 + INSTANTIATE(T, 19, E, H) \
206 + INSTANTIATE(T, 20, E, H) \
207 + INSTANTIATE(T, 21, E, H) \
208 + INSTANTIATE(T, 22, E, H) \
209 + INSTANTIATE(T, 23, E, H) \
210 + INSTANTIATE(T, 24, E, H) \
211 + INSTANTIATE(T, 25, E, H) \
212 + INSTANTIATE(T, 26, E, H) \
213 + INSTANTIATE(T, 27, E, H) \
214 + INSTANTIATE(T, 28, E, H) \
215 + INSTANTIATE(T, 29, E, H) \
216 + INSTANTIATE(T, 30, E, H) \
217 + INSTANTIATE(T, 31, E, H) \
218 + INSTANTIATE(T, 32, E, H)
219 +
220 +INSTANTIATE_ALL(float, 256, 3072)
221 +INSTANTIATE_ALL(__nv_bfloat16, 256, 3072)
222 +INSTANTIATE_ALL(float, 128, 6144)
223 +INSTANTIATE_ALL(__nv_bfloat16, 128, 6144)
224
225 #undef INSTANTIATE_ALL
226 #undef INSTANTIATE
csrc/libtorch_stable/fp32_router_gemm_entry.cu
+42 -18
@@ -22,36 +22,42 @@ inline int getSMVersion() {
22
23 } // namespace
24
25 -static constexpr int FP32_NUM_EXPERTS = 256;
26 -static constexpr int FP32_HIDDEN_DIM = 3072;
25 static constexpr int FP32_MAX_TOKENS = 32;
26
27 +// Supported (hidden_dim, num_experts) pairs (must match the instantiations in
28 +// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3.
29 +static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) {
30 + return (hidden_dim == 3072 && num_experts == 256) ||
31 + (hidden_dim == 6144 && num_experts == 128);
32 +}
33 +
34 // Forward declarations — 4 template params must match fp32_router_gemm.cu
35 template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
36 void invokeFp32RouterGemm(float* output, InputT const* mat_a,
37 float const* mat_b, cudaStream_t stream);
38
34 -// LoopUnroller templated on InputT
35 -template <typename InputT, int kBegin, int kEnd>
39 +// LoopUnroller templated on InputT, kNumExperts and kHiddenDim
40 +template <typename InputT, int kNumExperts, int kHiddenDim, int kBegin,
41 + int kEnd>
42 struct Fp32LoopUnroller {
43 static void unroll(int num_tokens, float* output, InputT const* mat_a,
44 float const* mat_b, cudaStream_t stream) {
45 if (num_tokens == kBegin) {
40 - invokeFp32RouterGemm<InputT, kBegin, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
46 + invokeFp32RouterGemm<InputT, kBegin, kNumExperts, kHiddenDim>(
47 output, mat_a, mat_b, stream);
48 } else {
43 - Fp32LoopUnroller<InputT, kBegin + 1, kEnd>::unroll(num_tokens, output,
44 - mat_a, mat_b, stream);
49 + Fp32LoopUnroller<InputT, kNumExperts, kHiddenDim, kBegin + 1,
50 + kEnd>::unroll(num_tokens, output, mat_a, mat_b, stream);
51 }
52 }
53 };
54
49 -template <typename InputT, int kEnd>
50 -struct Fp32LoopUnroller<InputT, kEnd, kEnd> {
55 +template <typename InputT, int kNumExperts, int kHiddenDim, int kEnd>
56 +struct Fp32LoopUnroller<InputT, kNumExperts, kHiddenDim, kEnd, kEnd> {
57 static void unroll(int num_tokens, float* output, InputT const* mat_a,
58 float const* mat_b, cudaStream_t stream) {
59 if (num_tokens == kEnd) {
54 - invokeFp32RouterGemm<InputT, kEnd, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
60 + invokeFp32RouterGemm<InputT, kEnd, kNumExperts, kHiddenDim>(
61 output, mat_a, mat_b, stream);
62 } else {
63 throw std::invalid_argument(
@@ -60,6 +66,23 @@ struct Fp32LoopUnroller<InputT, kEnd, kEnd> {
66 }
67 };
68
69 +// Dispatch over the supported (num_experts, hidden_dim) pairs.
70 +template <typename InputT>
71 +void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens,
72 + float* output, InputT const* mat_a,
73 + float const* mat_b, cudaStream_t stream) {
74 + if (num_experts == 256 && hidden_dim == 3072) {
75 + Fp32LoopUnroller<InputT, 256, 3072, 1, FP32_MAX_TOKENS>::unroll(
76 + num_tokens, output, mat_a, mat_b, stream);
77 + } else if (num_experts == 128 && hidden_dim == 6144) {
78 + Fp32LoopUnroller<InputT, 128, 6144, 1, FP32_MAX_TOKENS>::unroll(
79 + num_tokens, output, mat_a, mat_b, stream);
80 + } else {
81 + throw std::invalid_argument(
82 + "fp32_router_gemm: unsupported (hidden_dim, num_experts) pair");
83 + }
84 +}
85 +
86 void fp32_router_gemm(
87 torch::stable::Tensor& output, // [num_tokens, num_experts]
88 torch::stable::Tensor const& mat_a, // [num_tokens, hidden_dim]
@@ -85,10 +108,10 @@ void fp32_router_gemm(
108 STD_TORCH_CHECK(
109 mat_a.size(1) == mat_b.size(1),
110 "fp32_router_gemm: mat_a and mat_b must have the same hidden_dim");
88 - STD_TORCH_CHECK(hidden_dim == FP32_HIDDEN_DIM,
89 - "fp32_router_gemm: expected hidden_dim=3072");
90 - STD_TORCH_CHECK(num_experts == FP32_NUM_EXPERTS,
91 - "fp32_router_gemm: expected num_experts=256");
111 + STD_TORCH_CHECK(
112 + fp32_router_gemm_supported(hidden_dim, num_experts),
113 + "fp32_router_gemm: supported (hidden_dim, num_experts) pairs are "
114 + "(3072, 256) and (6144, 128)");
115 STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
116 "fp32_router_gemm: num_tokens must be in [0, 32]");
117 STD_TORCH_CHECK(
@@ -113,12 +136,13 @@ void fp32_router_gemm(
136 if (mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
137 auto const* mat_a_ptr =
138 reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
116 - Fp32LoopUnroller<__nv_bfloat16, 1, FP32_MAX_TOKENS>::unroll(
117 - num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
139 + dispatchFp32RouterGemm<__nv_bfloat16>(num_experts, hidden_dim, num_tokens,
140 + out_ptr, mat_a_ptr, mat_b_ptr,
141 + stream);
142 } else {
143 auto const* mat_a_ptr = reinterpret_cast<float const*>(mat_a.data_ptr());
120 - Fp32LoopUnroller<float, 1, FP32_MAX_TOKENS>::unroll(
121 - num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
144 + dispatchFp32RouterGemm<float>(num_experts, hidden_dim, num_tokens, out_ptr,
145 + mat_a_ptr, mat_b_ptr, stream);
146 }
147 }
148
csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu new
+635
@@ -0,0 +1,635 @@
1 +/*
2 + * SPDX-License-Identifier: Apache-2.0
3 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 + *
5 + * Horizontally-fused MiniMax-M3 attention pre-processing kernel.
6 + *
7 + * Replaces the per-token Python sequence in
8 + * ``MiniMaxM3SparseAttention.forward`` / ``MiniMaxM3Attention.forward``:
9 + *
10 + * q = q_norm(q); k = k_norm(k); q, k = rotary_emb(pos, q, k)
11 + * index_q = index_q_norm(index_q); index_k = index_k_norm(index_k)
12 + * index_q, index_k = rotary_emb(pos, index_q, index_k)
13 + * _insert_kv(k, v, index_k)
14 + *
15 + * All branches share head_dim=128 and the *same* partial-NeoX RoPE table
16 + * (``rotary_dim`` rotated, the trailing dims pass through). The four norms
17 + * are Gemma-style RMSNorm (``x * rsqrt(mean(x^2)+eps) * (1 + weight)``) with
18 + * independent weights.
19 + *
20 + * Everything lives in a single fused ``qkv`` tensor. The sparse layer's
21 + * fused projection (MinimaxM3QKVParallelLinearWithIndexer) emits, per token::
22 + *
23 + * [ q | k | v | index_q | index_k ] (the "5 results")
24 + *
25 + * while the dense layer emits just ``[ q | k | v ]``. The kernel reads the
26 + * index branch straight out of that packed row -- no separate index tensors.
27 + *
28 + * One kernel, one grid; each warp owns one (token, head-slot) pair. Slot
29 + * enumeration per token:
30 + * [0, nq) Q heads -> norm(q_w) + RoPE, write
31 + * qkv [nq, nq+nkv) K heads -> norm(k_w) + RoPE, write
32 + * qkv
33 + * (+ insert into key cache)
34 + * [nq+nkv, nq+2*nkv) V heads -> insert into value cache
35 + * IQ heads (niq) -> norm(iq_w) + RoPE, write iq
36 + * IK (1) -> norm(ik_w) + RoPE
37 + * (+ insert into index cache)
38 + *
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.
48 + *
49 + * Q/K and (sparse) index_q/index_k are all rewritten in place inside the fused
50 + * ``qkv`` tensor. Caches (bf16) are scatter-written by slot.
51 + */
52 +
53 +#include <cmath>
54 +#include <cuda_runtime.h>
55 +#include <type_traits>
56 +
57 +#include "torch_utils.h"
58 +
59 +#include "../cuda_compat.h"
60 +#include "../type_convert.cuh"
61 +#include "dispatch_utils.h"
62 +
63 +#ifndef FINAL_MASK
64 + #ifdef USE_ROCM
65 + #define FINAL_MASK 0xffffffffffffffffULL
66 + #else
67 + #define FINAL_MASK 0xffffffffu
68 + #endif
69 +#endif
70 +
71 +namespace vllm {
72 +namespace minimax_m3_fused_ops {
73 +
74 +namespace {
75 +inline int getSMVersion() {
76 + auto* props = get_device_prop();
77 + return props->major * 10 + props->minor;
78 +}
79 +} // namespace
80 +
81 +// ────────────────────────────────────────────────────────────────────────────
82 +// Constants (hard-coded for MiniMax-M3-preview).
83 +// ────────────────────────────────────────────────────────────────────────────
84 +constexpr int kHeadDim = 128;
85 +constexpr int kNumLanes = 32;
86 +constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 4
87 +
88 +// ────────────────────────────────────────────────────────────────────────────
89 +// Helpers
90 +// ────────────────────────────────────────────────────────────────────────────
91 +__device__ __forceinline__ float warpReduceSum(float val) {
92 +#pragma unroll
93 + for (int mask = 16; mask > 0; mask >>= 1) {
94 + val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
95 + }
96 + return val;
97 +}
98 +
99 +// Gemma RMSNorm over the full head (no-op when ``weight == nullptr``), rounded
100 +// back to scalar_t like the materialized unfused norm output, followed by
101 +// partial NeoX RoPE on the leading ``rotary_dim`` dims. Each lane owns
102 +// ``kElemsPerLane`` contiguous dims [laneId*4, laneId*4+4).
103 +template <typename scalar_t>
104 +__device__ __forceinline__ void normAndRope(
105 + float (&elems)[kElemsPerLane], int const laneId, float const eps,
106 + scalar_t const* __restrict__ weight, // [kHeadDim] or nullptr (no norm)
107 + bool const do_rope, int const rotary_dim,
108 + scalar_t const* __restrict__ cos_ptr, // cos_sin_cache + pos*rotary_dim
109 + bool const apply_norm) {
110 + // ── Gemma RMSNorm: x * rsqrt(mean(x^2)+eps) * (1 + w) ──────────────────
111 + if (apply_norm) {
112 + float sumsq = 0.0f;
113 +#pragma unroll
114 + for (int i = 0; i < kElemsPerLane; i++) sumsq += elems[i] * elems[i];
115 + sumsq = warpReduceSum(sumsq);
116 + float const rms_rcp = rsqrtf(sumsq / static_cast<float>(kHeadDim) + eps);
117 +#pragma unroll
118 + for (int i = 0; i < kElemsPerLane; i++) {
119 + int const dim = laneId * kElemsPerLane + i;
120 + float const w = 1.0f + static_cast<float>(weight[dim]);
121 + elems[i] = elems[i] * rms_rcp * w;
122 + }
123 + }
124 +
125 + // ── Partial NeoX RoPE on dims [0, rotary_dim) ──────────────────────────
126 + // half = rotary_dim/2. Pair (i, i+half) for i in [0, half). Lane L owns
127 + // dims [4L, 4L+4); since half is a multiple of 4, a lane lies wholly in the
128 + // first half (own=x[i]) or second half (own=x[i+half]); its partner lives
129 + // ``half/4`` lanes away (XOR with that distance).
130 + if (do_rope) {
131 + int const half = rotary_dim / 2;
132 + int const dim0 = laneId * kElemsPerLane;
133 + bool const in_rope = dim0 < rotary_dim;
134 + int const lane_xor = half / kElemsPerLane; // partner-lane distance
135 +
136 + float partner[kElemsPerLane];
137 +#pragma unroll
138 + for (int i = 0; i < kElemsPerLane; i++) {
139 + partner[i] = __shfl_xor_sync(FINAL_MASK, elems[i], lane_xor, 32);
140 + }
141 + if (in_rope) {
142 + bool const first_half = dim0 < half;
143 + int const i_base = first_half ? dim0 : (dim0 - half); // cos/sin index
144 + scalar_t const* sin_ptr = cos_ptr + half;
145 +#pragma unroll
146 + for (int i = 0; i < kElemsPerLane; i++) {
147 + float const c = static_cast<float>(cos_ptr[i_base + i]);
148 + float const s = static_cast<float>(sin_ptr[i_base + i]);
149 + if (first_half) {
150 + elems[i] = elems[i] * c - partner[i] * s;
151 + } else {
152 + elems[i] = elems[i] * c + partner[i] * s;
153 + }
154 + }
155 + }
156 + }
157 +}
158 +
159 +// Load 4 contiguous bf16 -> 4 fp32 registers.
160 +template <typename scalar_t>
161 +__device__ __forceinline__ void loadElems(scalar_t const* __restrict__ src,
162 + float (&elems)[kElemsPerLane]) {
163 + using Converter = vllm::_typeConvert<scalar_t>;
164 + uint2 v = *reinterpret_cast<uint2 const*>(src);
165 + auto const* p =
166 + reinterpret_cast<typename Converter::packed_hip_type const*>(&v);
167 +#pragma unroll
168 + for (int i = 0; i < kElemsPerLane / 2; i++) {
169 + float2 f2 = Converter::convert(p[i]);
170 + elems[2 * i] = f2.x;
171 + elems[2 * i + 1] = f2.y;
172 + }
173 +}
174 +
175 +// Store 4 fp32 registers -> 4 contiguous bf16.
176 +template <typename scalar_t>
177 +__device__ __forceinline__ void storeElems(
178 + scalar_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
179 + using Converter = vllm::_typeConvert<scalar_t>;
180 + uint2 v;
181 + auto* p = reinterpret_cast<typename Converter::packed_hip_type*>(&v);
182 +#pragma unroll
183 + for (int i = 0; i < kElemsPerLane / 2; i++) {
184 + p[i] = Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1]));
185 + }
186 + *reinterpret_cast<uint2*>(dst) = v;
187 +}
188 +
189 +// ────────────────────────────────────────────────────────────────────────────
190 +// Kernel
191 +// ────────────────────────────────────────────────────────────────────────────
192 +// Grid: 1D, ceil(num_tokens * slots_per_token / warps_per_block).
193 +// Each warp = one (token, slot).
194 +//
195 +// `kIsSparse` and `kInsertKV` are compile-time template bools, so all the
196 +// branch decisions that distinguish the dense layer from the sparse layer
197 +// (index slots, KV/index inserts, V slots) fold away per instantiation.
198 +// Three instantiations are built: dense <false,false>, sparse-profiling
199 +// <true,false> and sparse-serving <true,true>. Slots per token:
200 +// Q : nq (always — norm+RoPE)
201 +// K : nkv (always — norm+RoPE; +K-cache insert)
202 +// V : nkv only if kInsertKV (V-cache insert; no warps in dense)
203 +// IQ: niq only if kIsSparse (norm+RoPE)
204 +// IK: 1 only if kIsSparse (norm+RoPE; +index-cache insert)
205 +template <typename scalar_t, bool kIsSparse, bool kInsertKV>
206 +__global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
207 + scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse)
208 + scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
209 + scalar_t* __restrict__ index_q_out, // [N, niq*128] contiguous, or nullptr
210 + scalar_t const* __restrict__ q_norm_w,
211 + scalar_t const* __restrict__ k_norm_w,
212 + scalar_t const* __restrict__ iq_norm_w,
213 + scalar_t const* __restrict__ ik_norm_w,
214 + scalar_t const* __restrict__ cos_sin_cache, // [max_pos, rotary_dim]
215 + int64_t const* __restrict__ positions, // [N] i64
216 + int64_t const* __restrict__ slot_mapping, // main K/V slots or nullptr
217 + int64_t const* __restrict__ index_slot_mapping, // index K slots/nullptr
218 + scalar_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr
219 + scalar_t* __restrict__ index_cache, // [nb*bs, 128] or nullptr
220 + float const eps, int const rotary_dim, int const num_tokens, int const nq,
221 + int const nkv, int const niq, int const block_size,
222 + // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
223 + // The head_dim (last) dim is always innermost-contiguous (stride 1), so the
224 + // NHD/HND layout choice is fully captured by these four strides: NHD keeps
225 + // s_token < s_head, HND swaps them. dim_base addresses head_dim directly.
226 + int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
227 + int64_t const kv_s_head) {
228 +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
229 + // _typeConvert<BFloat16> is unavailable on pre-Ampere; the M3 kernel only
230 + // runs with bf16/fp16 inputs in practice. Discard the bf16 body there.
231 + if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
232 + return;
233 + } else {
234 +#endif
235 + int const warpsPerBlock = blockDim.x / 32;
236 + int const laneId = threadIdx.x % 32;
237 + int const globalWarpIdx = blockIdx.x * warpsPerBlock + (threadIdx.x / 32);
238 +
239 + // Slot layout (compile-time gated: dense has neither V nor index slots).
240 + int const v_slots = kInsertKV ? nkv : 0;
241 + int const idx_slots = kIsSparse ? niq + 1 : 0;
242 + int const slots_per_token = nq + nkv + v_slots + idx_slots;
243 +
244 + int const tokenIdx = globalWarpIdx / slots_per_token;
245 + int const slot = globalWarpIdx % slots_per_token;
246 + if (tokenIdx >= num_tokens) return;
247 +
248 + // Slot boundaries.
249 + int const k_begin = nq;
250 + int const v_begin = nq + nkv; // valid only when kInsertKV
251 + int const iq_begin = nq + nkv + v_slots; // index block start
252 + int const ik_slot = iq_begin + niq; // valid only when kIsSparse
253 +
254 + bool const isQ = slot < k_begin;
255 + bool const isK = slot >= k_begin && slot < v_begin;
256 + bool isV = false;
257 + if constexpr (kInsertKV) isV = slot >= v_begin && slot < v_begin + nkv;
258 + bool isIQ = false, isIK = false;
259 + if constexpr (kIsSparse) {
260 + isIQ = slot >= iq_begin && slot < ik_slot;
261 + isIK = slot == ik_slot;
262 + }
263 +
264 + int const dim_base = laneId * kElemsPerLane;
265 + // Physical row width of qkv: the dense layer packs [q|k|v]; the sparse
266 + // layer additionally packs [index_q (niq heads) | index_k (1 head)].
267 + int const qkv_row = (nq + 2 * nkv + (kIsSparse ? (niq + 1) : 0)) * kHeadDim;
268 +
269 + // ── Resolve source pointer + per-branch parameters. ────────────────────
270 + scalar_t* row_ptr = nullptr; // in-place output location
271 + scalar_t const* norm_w = nullptr; // nullptr -> skip norm (V)
272 + bool do_rope = true;
273 + int head = 0; // kv head index for inserts
274 +
275 + if (isQ) {
276 + row_ptr =
277 + qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
278 + norm_w = q_norm_w;
279 + } else if (isK) {
280 + head = slot - k_begin;
281 + row_ptr =
282 + qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
283 + norm_w = k_norm_w;
284 + } else if (isV) {
285 + // qkv V section starts at slot index (nq + nkv): slot * kHeadDim is the
286 + // correct in-tensor offset.
287 + head = slot - v_begin;
288 + row_ptr =
289 + qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
290 + norm_w = nullptr; // V: no norm, no rope
291 + do_rope = false;
292 + } else if (isIQ) {
293 + // index_q sub-block lives at physical offset (nq+2*nkv)*128 in qkv.
294 + int const ih = slot - iq_begin;
295 + row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
296 + (nq + 2 * nkv + ih) * kHeadDim;
297 + norm_w = iq_norm_w;
298 + } else { // isIK -- single shared index key at (nq+2*nkv+niq)*128.
299 + row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
300 + (nq + 2 * nkv + niq) * kHeadDim;
301 + norm_w = ik_norm_w;
302 + }
303 +
304 + // Store destination. Q and index_q are gathered into dedicated contiguous
305 + // output buffers (when provided) so the downstream SM100 sparse kernel's
306 + // flat TMA descriptor can address them as [tokens*heads, head_dim]; this
307 + // folds the de-interleaving into the store the kernel already does, instead
308 + // of a separate q.contiguous() copy. Everything else stays in place.
309 + scalar_t* store_ptr = row_ptr;
310 + if (isQ && q_out != nullptr) {
311 + store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
312 + slot * kHeadDim;
313 + } else if (isIQ && index_q_out != nullptr) {
314 + store_ptr = index_q_out +
315 + static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
316 + (slot - iq_begin) * kHeadDim;
317 + }
318 +
319 + // PDL: wait for the predecessor kernel (the qkv-projection GEMM that
320 + // produces ``qkv``) to finish before touching any global memory. No-op
321 + // when PDL is not enabled on the launch. The CUDA runtime wrapper emits
322 + // the griddepcontrol.wait PTX with the required memory clobber internally.
323 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
324 + cudaGridDependencySynchronize();
325 +#endif
326 +
327 + // ── Load -> norm+rope (fp32) -> store back in place. ───────────────────
328 + float elems[kElemsPerLane];
329 + loadElems<scalar_t>(row_ptr + dim_base, elems);
330 +
331 + if (!isV) {
332 + int64_t const pos = positions[tokenIdx];
333 + scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim;
334 + normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
335 + cos_ptr, /*apply_norm=*/norm_w != nullptr);
336 + storeElems<scalar_t>(store_ptr + dim_base, elems);
337 + }
338 +
339 + // ── Cache inserts (sparse serving only). ───────────────────────────────
340 + if constexpr (kInsertKV) {
341 + // Guard (not early-return) so every thread reaches the PDL trigger below.
342 + int64_t const sm = (isK || isV)
343 + ? slot_mapping[tokenIdx]
344 + : (isIK ? index_slot_mapping[tokenIdx] : -1);
345 + if (sm >= 0) { // skip padded / unscheduled tokens
346 + if (isIK) {
347 + scalar_t* dst = index_cache + sm * kHeadDim + dim_base;
348 + storeElems<scalar_t>(dst, elems);
349 + } else if (isK || isV) {
350 + // kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
351 + // Paging is logical (block = sm/block_size, token = sm%block_size);
352 + // the physical NHD/HND layout is honoured via the passed strides.
353 + int64_t const b = sm / block_size;
354 + int64_t const t = sm % block_size;
355 + int const kv = isK ? 0 : 1;
356 + int64_t const off =
357 + 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);
359 + }
360 + }
361 + }
362 +
363 + // PDL: signal that this kernel is done so a dependent successor may launch
364 + // early. No-op when PDL is not enabled on the launch.
365 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
366 + cudaTriggerProgrammaticLaunchCompletion();
367 +#endif
368 +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
369 + }
370 +#endif
371 +}
372 +
373 +// ────────────────────────────────────────────────────────────────────────────
374 +// Launch wrapper
375 +// ────────────────────────────────────────────────────────────────────────────
376 +template <typename scalar_t>
377 +void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
378 + scalar_t const* q_norm_w, scalar_t const* k_norm_w,
379 + scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
380 + scalar_t const* cos_sin_cache,
381 + int64_t const* positions, int64_t const* slot_mapping,
382 + int64_t const* index_slot_mapping, scalar_t* kv_cache,
383 + scalar_t* index_cache, float const eps,
384 + int const rotary_dim, int const num_tokens,
385 + int const nq, int const nkv, int const niq,
386 + int const block_size, int64_t const kv_s_block,
387 + int64_t const kv_s_kv, int64_t const kv_s_token,
388 + int64_t const kv_s_head, bool const has_index,
389 + bool const insert_kv, cudaStream_t stream) {
390 + // Slot count must match the kernel's compile-time gating.
391 + int const v_slots = insert_kv ? nkv : 0;
392 + int const idx_slots = has_index ? niq + 1 : 0;
393 + int const slots_per_token = nq + nkv + v_slots + idx_slots;
394 +
395 + constexpr int kBlockSize = 256;
396 + constexpr int kWarpsPerBlock = kBlockSize / 32;
397 + int64_t const total_warps =
398 + static_cast<int64_t>(num_tokens) * slots_per_token;
399 + int const grid =
400 + static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
401 + if (grid == 0) return;
402 +
403 +#ifndef USE_ROCM
404 + // PDL: enable programmatic stream serialization whenever the hardware
405 + // supports it (SM90+). On pre-Hopper GPUs the attribute is unavailable, so
406 + // leave numAttrs = 0 and launch as a regular kernel via cudaLaunchKernelEx.
407 + static int const sm_version = getSMVersion();
408 + cudaLaunchConfig_t config;
409 + config.gridDim = dim3(grid);
410 + config.blockDim = dim3(kBlockSize);
411 + config.dynamicSmemBytes = 0;
412 + config.stream = stream;
413 + cudaLaunchAttribute attrs[1];
414 + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
415 + attrs[0].val.programmaticStreamSerializationAllowed = 1;
416 + config.attrs = attrs;
417 + config.numAttrs = (sm_version >= 90) ? 1 : 0;
418 +
419 + #define LAUNCH(IS_SPARSE, INSERT) \
420 + cudaLaunchKernelEx( \
421 + &config, \
422 + fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT>, \
423 + qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, \
424 + cos_sin_cache, positions, slot_mapping, index_slot_mapping, kv_cache, \
425 + index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, \
426 + kv_s_block, kv_s_kv, kv_s_token, kv_s_head)
427 +#else
428 + // ROCm: standard kernel launch syntax (no PDL/stream serialization).
429 + // clang-format off
430 + #define LAUNCH(IS_SPARSE, INSERT) \
431 + fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT> \
432 + <<<grid, kBlockSize, 0, stream>>>( \
433 + qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, \
434 + ik_norm_w, cos_sin_cache, positions, slot_mapping, \
435 + index_slot_mapping, kv_cache, index_cache, eps, rotary_dim, \
436 + num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv, \
437 + kv_s_token, kv_s_head)
438 + // clang-format on
439 +#endif
440 +
441 + if (has_index) {
442 + if (insert_kv) {
443 + LAUNCH(true, true); // sparse serving
444 + } else {
445 + LAUNCH(true, false); // sparse profiling
446 + }
447 + } else {
448 + // Dense layer: never has an index branch and never inserts here (the
449 + // generic Attention layer owns the KV insert).
450 + LAUNCH(false, false);
451 + }
452 +#undef LAUNCH
453 +}
454 +
455 +} // namespace minimax_m3_fused_ops
456 +} // namespace vllm
457 +
458 +// ────────────────────────────────────────────────────────────────────────────
459 +// Torch op wrapper
460 +// ────────────────────────────────────────────────────────────────────────────
461 +void fused_minimax_m3_qknorm_rope_kv_insert(
462 + torch::stable::Tensor& qkv, // [N, qkv_row] (packs index if sparse)
463 + torch::stable::Tensor const& q_norm_weight, // [128]
464 + torch::stable::Tensor const& k_norm_weight, // [128]
465 + torch::stable::Tensor const& cos_sin_cache, // [max_pos, rotary_dim]
466 + torch::stable::Tensor const& positions, // [N] i64
467 + int64_t num_heads, int64_t num_kv_heads, int64_t rotary_dim, double eps,
468 + std::optional<torch::stable::Tensor> index_q_norm_weight, // [128]
469 + std::optional<torch::stable::Tensor> index_k_norm_weight, // [128]
470 + int64_t num_index_heads, // niq; 0 => dense
471 + std::optional<torch::stable::Tensor> slot_mapping, // [N] i64
472 + std::optional<torch::stable::Tensor> index_slot_mapping, // [N] i64
473 + std::optional<torch::stable::Tensor> kv_cache, // [nb,2,bs,nkv,128]
474 + std::optional<torch::stable::Tensor> index_cache, // [nb,bs,128]
475 + int64_t block_size,
476 + std::optional<torch::stable::Tensor> q_out, // [N, nq*128] contiguous
477 + std::optional<torch::stable::Tensor>
478 + index_q_out) { // [N, niq*128] contiguous
479 + STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
480 + "qkv must be contiguous CUDA");
481 + STD_TORCH_CHECK(
482 + positions.is_cuda() &&
483 + positions.scalar_type() == torch::headeronly::ScalarType::Long,
484 + "positions must be int64 CUDA");
485 + STD_TORCH_CHECK(cos_sin_cache.is_cuda() && cos_sin_cache.is_contiguous(),
486 + "cos_sin_cache must be contiguous CUDA");
487 + STD_TORCH_CHECK(cos_sin_cache.scalar_type() == qkv.scalar_type(),
488 + "cos_sin_cache dtype must match qkv");
489 + STD_TORCH_CHECK(
490 + cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == rotary_dim,
491 + "cos_sin_cache shape [max_pos, rotary_dim]");
492 +
493 + STD_TORCH_CHECK(q_norm_weight.scalar_type() == qkv.scalar_type() &&
494 + k_norm_weight.scalar_type() == qkv.scalar_type(),
495 + "q/k norm weight dtype must match qkv");
496 + STD_TORCH_CHECK(
497 + q_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim &&
498 + k_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim,
499 + "q/k norm weight must have 128 elements");
500 + STD_TORCH_CHECK(rotary_dim > 0 && rotary_dim % 8 == 0 &&
501 + rotary_dim <= vllm::minimax_m3_fused_ops::kHeadDim,
502 + "rotary_dim must be a positive multiple of 8 and <= 128");
503 +
504 + int const num_tokens = static_cast<int>(qkv.size(0));
505 + int const nq = static_cast<int>(num_heads);
506 + int const nkv = static_cast<int>(num_kv_heads);
507 + int const niq = static_cast<int>(num_index_heads);
508 +
509 + // The sparse layer packs the index branch ([index_q (niq heads) | index_k
510 + // (1 head)]) right after [q|k|v] in the same row; the dense layer does not.
511 + bool const has_index = niq > 0;
512 + bool const insert_kv = kv_cache.has_value();
513 + int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
514 + int const expected_row =
515 + (nq + 2 * nkv + (has_index ? niq + 1 : 0)) * kHeadDim;
516 + STD_TORCH_CHECK(qkv.size(1) == expected_row,
517 + "qkv last dim must be (num_heads + 2*num_kv_heads"
518 + " + num_index_heads + 1) * 128 for sparse, "
519 + "(num_heads + 2*num_kv_heads) * 128 for dense");
520 +
521 + // Only the sparse layer inserts here (dense lets the generic Attention layer
522 + // own the KV write); there is no dense+insert kernel instantiation.
523 + STD_TORCH_CHECK(
524 + !insert_kv || has_index,
525 + "insert mode (kv_cache) requires the index branch (sparse layer)");
526 + if (has_index) {
527 + STD_TORCH_CHECK(
528 + index_q_norm_weight.has_value() && index_k_norm_weight.has_value(),
529 + "index branch requires both index norm weights");
530 + STD_TORCH_CHECK(index_q_norm_weight->scalar_type() == qkv.scalar_type() &&
531 + index_k_norm_weight->scalar_type() == qkv.scalar_type(),
532 + "index norm weights dtype must match qkv");
533 + STD_TORCH_CHECK(index_q_norm_weight->numel() == kHeadDim &&
534 + index_k_norm_weight->numel() == kHeadDim,
535 + "index norm weights must have 128 elements");
536 + }
537 + // kv_cache strides (logical shape [nb, 2, bs, nkv, head_dim]). Read straight
538 + // off the tensor so the kernel honours whatever physical layout the attention
539 + // backend allocated (NHD: stride order (0,1,2,3,4); HND: (0,1,3,2,4)). No new
540 + // op argument is needed -- the strides ride along with the tensor itself.
541 + int64_t kv_s_block = 0, kv_s_kv = 0, kv_s_token = 0, kv_s_head = 0;
542 + torch::stable::Tensor const* effective_index_slot_mapping = nullptr;
543 + if (insert_kv) {
544 + STD_TORCH_CHECK(
545 + slot_mapping.has_value() && slot_mapping->is_cuda() &&
546 + slot_mapping->scalar_type() == torch::headeronly::ScalarType::Long,
547 + "insert mode requires int64 CUDA slot_mapping");
548 + STD_TORCH_CHECK(
549 + !index_slot_mapping.has_value() ||
550 + (index_slot_mapping->is_cuda() &&
551 + index_slot_mapping->scalar_type() ==
552 + torch::headeronly::ScalarType::Long &&
553 + index_slot_mapping->numel() == slot_mapping->numel()),
554 + "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)");
557 + STD_TORCH_CHECK(index_cache.has_value() &&
558 + index_cache->scalar_type() == qkv.scalar_type(),
559 + "insert mode requires matching index_cache");
560 + STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
561 + "kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
562 + "head_dim (stride(4)==1)");
563 + kv_s_block = kv_cache->stride(0);
564 + kv_s_kv = kv_cache->stride(1);
565 + kv_s_token = kv_cache->stride(2);
566 + kv_s_head = kv_cache->stride(3);
567 + effective_index_slot_mapping = index_slot_mapping.has_value()
568 + ? &index_slot_mapping.value()
569 + : &slot_mapping.value();
570 + }
571 + // Optional contiguous gather targets: when given, the normed/roped q (and
572 + // index_q) are written here instead of in place, so callers avoid a separate
573 + // .contiguous() copy. index_q_out only makes sense on the sparse path.
574 + if (q_out.has_value()) {
575 + STD_TORCH_CHECK(
576 + q_out->is_cuda() && q_out->is_contiguous() &&
577 + q_out->scalar_type() == qkv.scalar_type(),
578 + "q_out must be a contiguous CUDA tensor matching qkv dtype");
579 + STD_TORCH_CHECK(
580 + q_out->numel() == static_cast<int64_t>(num_tokens) * nq * kHeadDim,
581 + "q_out must have num_tokens * num_heads * 128 elements");
582 + }
583 + if (index_q_out.has_value()) {
584 + STD_TORCH_CHECK(
585 + has_index,
586 + "index_q_out requires the index branch (num_index_heads > 0)");
587 + STD_TORCH_CHECK(
588 + index_q_out->is_cuda() && index_q_out->is_contiguous() &&
589 + index_q_out->scalar_type() == qkv.scalar_type(),
590 + "index_q_out must be a contiguous CUDA tensor matching qkv dtype");
591 + STD_TORCH_CHECK(index_q_out->numel() ==
592 + static_cast<int64_t>(num_tokens) * niq * kHeadDim,
593 + "index_q_out must have num_tokens * num_index_heads * 128 "
594 + "elements");
595 + }
596 +
597 + const torch::stable::accelerator::DeviceGuard device_guard(
598 + qkv.get_device_index());
599 + auto stream = get_current_cuda_stream(qkv.get_device_index());
600 +
601 + VLLM_STABLE_DISPATCH_HALF_TYPES(
602 + qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert", [&] {
603 + 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);
634 + });
635 +}
csrc/libtorch_stable/ops.h
+20 -1
@@ -281,6 +281,24 @@ minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
281 int64_t const nranks, double const eps);
282 #endif
283
284 +// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV /
285 +// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs
286 +// the index branch and scatters k/v/index_k into their paged caches.
287 +void fused_minimax_m3_qknorm_rope_kv_insert(
288 + torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight,
289 + torch::stable::Tensor const& k_norm_weight,
290 + torch::stable::Tensor const& cos_sin_cache,
291 + torch::stable::Tensor const& positions, int64_t num_heads,
292 + int64_t num_kv_heads, int64_t rotary_dim, double eps,
293 + std::optional<torch::stable::Tensor> index_q_norm_weight,
294 + std::optional<torch::stable::Tensor> index_k_norm_weight,
295 + int64_t num_index_heads, std::optional<torch::stable::Tensor> slot_mapping,
296 + std::optional<torch::stable::Tensor> index_slot_mapping,
297 + std::optional<torch::stable::Tensor> kv_cache,
298 + std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
299 + std::optional<torch::stable::Tensor> q_out,
300 + std::optional<torch::stable::Tensor> index_q_out);
301 +
302 // Sampler kernels (shared CUDA/ROCm)
303 void apply_repetition_penalties_(
304 torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -346,7 +364,8 @@ void free_shared_buffer(int64_t buffer);
364 // Activation kernels (shared CUDA/ROCm)
365 void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
366 void silu_and_mul_clamp(torch::stable::Tensor& out,
349 - torch::stable::Tensor& input, double limit);
367 + torch::stable::Tensor& input, double limit,
368 + double alpha = 1.0, double beta = 0.0);
369 void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
370 void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
371 void gelu_tanh_and_mul(torch::stable::Tensor& out,
csrc/libtorch_stable/torch_bindings.cpp
+18 -2
@@ -461,6 +461,18 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
461 "float eps) -> (Tensor, Tensor)");
462 #endif
463
464 + // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert.
465 + ops.def(
466 + "fused_minimax_m3_qknorm_rope_kv_insert("
467 + "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, "
468 + "Tensor cos_sin_cache, Tensor positions, int num_heads, "
469 + "int num_kv_heads, int rotary_dim, float eps, "
470 + "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, "
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) -> ()");
475 +
476 // Apply repetition penalties to logits in-place.
477 ops.def(
478 "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -488,9 +500,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
500 ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
501
502 // SwiGLU activation with input clamping.
503 + // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to
504 + // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up.
505 ops.def(
492 - "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
493 - "-> ()");
506 + "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
507 + "float alpha=1.0, float beta=0.0) -> ()");
508
509 // Activation function used in GeGLU with `none` approximation.
510 ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
@@ -679,6 +693,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
693 ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms));
694 ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
695 #endif
696 + ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
697 + TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
698
699 // Sampler kernels (shared CUDA/ROCm)
700 ops.impl("apply_repetition_penalties_",
csrc/ops.h
+2 -1
@@ -50,7 +50,8 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
50
51 void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
52
53 -void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
53 +void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
54 + double alpha = 1.0, double beta = 0.0);
55
56 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
57 torch::Tensor& scale);
docs/design/attention_backends.md
+14 -2
@@ -170,8 +170,8 @@ Priority is **1 = highest** (tried first).
170 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
171 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
172 | `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
173 -| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
174 -| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
173 +| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
174 +| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
175 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
176 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
177 | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
@@ -188,6 +188,18 @@ Priority is **1 = highest** (tried first).
188 >
189 > **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise.
190
191 +## MiniMax M3 Sparse Attention Backends
192 +
193 +Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")
194 +layers. It is wired in directly by the model and is not part of the
195 +automatic priority lists above. A lightning indexer scores KV blocks, the
196 +top-k blocks (plus fixed init/local blocks) are selected, and attention
197 +attends only to those blocks; index keys live in a separate side cache.
198 +
199 +| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
200 +| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
201 +| `MINIMAX_M3_SPARSE` | bf16, fp16 | `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 128 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
202 +
203 ## MLA (Multi-head Latent Attention) Backends
204
205 MLA uses separate backends for prefill and decode phases.
pyproject.toml
+2
@@ -162,6 +162,8 @@ dout = "dout"
162 Pn = "Pn"
163 arange = "arange"
164 thw = "thw"
165 +# temporal position ids (parallels hpos/wpos in vision RoPE)
166 +tpos = "tpos"
167 subtile = "subtile"
168 HSA = "HSA"
169 setp = "setp"
requirements/common.txt
+1
@@ -29,6 +29,7 @@ xgrammar >= 0.2.1, < 1.0.0; platform_machine == "x86_64" or platform_machine ==
29 typing_extensions >= 4.10
30 filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
31 partial-json-parser # used for parsing partial JSON outputs
32 +jsonschema >= 4.23.0 # required for MiniMax M3 tool schema validation
33 pyzmq >= 25.0.0
34 msgspec
35 mistral_common[image] >= 1.11.3
requirements/test/cuda.txt
+1
@@ -360,6 +360,7 @@ jsonpointer==3.0.0
360 # via jsonschema
361 jsonschema==4.23.0
362 # via
363 + # -c requirements/common.txt
364 # hypothesis-jsonschema
365 # mistral-common
366 # ray
requirements/test/rocm.txt
+2
@@ -439,6 +439,8 @@ jsonpointer==3.1.0
439 # via jsonschema
440 jsonschema==4.26.0
441 # via
442 + # -c requirements/common.txt
443 + # -r requirements/test/../common.txt
444 # hypothesis-jsonschema
445 # mcp
446 # mistral-common
requirements/test/xpu.txt
+1
@@ -229,6 +229,7 @@ jsonlines==4.0.0
229 # via lm-eval
230 jsonschema==4.26.0
231 # via
232 + # -c requirements/common.txt
233 # hypothesis-jsonschema
234 # mistral-common
235 # schemathesis
rust/src/chat/src/lib.rs
+2 -2
@@ -277,7 +277,7 @@ mod tests {
277 )
278 .unwrap_err();
279
280 - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, phi4_mini_json, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
280 + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
281 }
282
283 #[test]
@@ -288,6 +288,6 @@ mod tests {
288 )
289 .unwrap_err();
290
291 - expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
291 + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
292 }
293 }
rust/src/chat/src/parser/reasoning/mod.rs
+7 -3
@@ -5,9 +5,9 @@ use std::sync::LazyLock;
5 pub use vllm_reasoning_parser::{
6 CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
7 DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
8 - KimiReasoningParser, MiniMaxM2ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser,
9 - ReasoningDelta, ReasoningError, ReasoningParser, SeedOssReasoningParser, Step3ReasoningParser,
10 - Step3p5ReasoningParser,
8 + KimiReasoningParser, MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser,
9 + NemotronV3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta, ReasoningError,
10 + ReasoningParser, SeedOssReasoningParser, Step3ReasoningParser, Step3p5ReasoningParser,
11 };
12 use vllm_tokenizer::DynTokenizer;
13
@@ -24,6 +24,7 @@ pub mod names {
24 pub const KIMI: &str = "kimi";
25 pub const KIMI_K2: &str = "kimi_k2";
26 pub const MINIMAX_M2: &str = "minimax_m2";
27 + pub const MINIMAX_M3: &str = "minimax_m3";
28 pub const NEMOTRON_V3: &str = "nemotron_v3";
29 pub const QWEN3: &str = "qwen3";
30 pub const SEED_OSS: &str = "seed_oss";
@@ -62,6 +63,7 @@ impl ReasoningParserFactory {
63 .register_parser::<KimiReasoningParser>(names::KIMI)
64 .register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
65 .register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
66 + .register_parser::<MiniMaxM3ReasoningParser>(names::MINIMAX_M3)
67 .register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)
68 .register_parser::<Qwen3ReasoningParser>(names::QWEN3)
69 .register_parser::<SeedOssReasoningParser>(names::SEED_OSS)
@@ -90,6 +92,8 @@ impl ReasoningParserFactory {
92 .register_pattern("step3", names::STEP3)
93 .register_pattern("seed-oss", names::SEED_OSS)
94 .register_pattern("seedoss", names::SEED_OSS)
95 + .register_pattern("minimax-m3", names::MINIMAX_M3)
96 + .register_pattern("mm-m3", names::MINIMAX_M3)
97 .register_pattern("minimax", names::MINIMAX_M2)
98 .register_pattern("mm-m2", names::MINIMAX_M2)
99 .register_pattern("cohere", names::COHERE_CMD)
rust/src/chat/src/parser/reasoning/tests.rs
+15
@@ -34,10 +34,12 @@ fn factory_contains_and_lists_registered_parsers() {
34 assert!(factory.contains(names::DEEPSEEK_V4));
35 assert!(factory.contains(names::SEED_OSS));
36 assert!(factory.contains(names::STEP3P5));
37 + assert!(factory.contains(names::MINIMAX_M3));
38 assert!(factory.list().contains(&names::QWEN3.to_string()));
39 assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string()));
40 assert!(factory.list().contains(&names::SEED_OSS.to_string()));
41 assert!(factory.list().contains(&names::STEP3P5.to_string()));
42 + assert!(factory.list().contains(&names::MINIMAX_M3.to_string()));
43 }
44
45 #[test]
@@ -88,6 +90,19 @@ fn factory_routes_seed_oss_models() {
90 );
91 }
92
93 +#[test]
94 +fn factory_resolves_minimax_m3_before_generic_minimax() {
95 + let factory = ReasoningParserFactory::new();
96 + assert_eq!(
97 + factory.resolve_name_for_model("MiniMaxAI/Minimax-M3-preview"),
98 + Some(names::MINIMAX_M3)
99 + );
100 + assert_eq!(
101 + factory.resolve_name_for_model("mm-m3"),
102 + Some(names::MINIMAX_M3)
103 + );
104 +}
105 +
106 #[test]
107 fn factory_rejects_unknown_parser_names() {
108 let tokenizer = Arc::new(FakeTokenizer);
rust/src/chat/src/parser/tool/mod.rs
+7 -2
@@ -6,8 +6,9 @@ pub use vllm_tool_parser::{
6 DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
7 Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser,
8 HyV3ToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser,
9 - MinimaxM2ToolParser, MistralToolParser, Phi4MiniJsonToolParser, Qwen3CoderToolParser,
10 - Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError, ToolParserOutput,
9 + MinimaxM2ToolParser, MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser,
10 + Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError,
11 + ToolParserOutput,
12 };
13
14 use crate::parser::ParserFactory;
@@ -32,6 +33,7 @@ pub mod names {
33 pub const LLAMA3_JSON: &str = "llama3_json";
34 pub const LLAMA4_JSON: &str = "llama4_json";
35 pub const MINIMAX_M2: &str = "minimax_m2";
36 + pub const MINIMAX_M3: &str = "minimax_m3";
37 pub const MISTRAL: &str = "mistral";
38 pub const PHI4_MINI_JSON: &str = "phi4_mini_json";
39 pub const QWEN3_CODER: &str = "qwen3_coder";
@@ -73,6 +75,7 @@ impl ToolParserFactory {
75 .register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
76 .register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
77 .register_parser::<MinimaxM2ToolParser>(names::MINIMAX_M2)
78 + .register_parser::<MinimaxM3ToolParser>(names::MINIMAX_M3)
79 .register_parser::<MistralToolParser>(names::MISTRAL)
80 .register_parser::<Phi4MiniJsonToolParser>(names::PHI4_MINI_JSON)
81 .register_parser::<Qwen3XmlToolParser>(names::QWEN3_XML)
@@ -111,6 +114,8 @@ impl ToolParserFactory {
114 .register_pattern("gemma-4", names::GEMMA4)
115 .register_pattern("granite-4", names::GRANITE4)
116 .register_pattern("kimi-k2", names::KIMI_K2)
117 + .register_pattern("minimax-m3", names::MINIMAX_M3)
118 + .register_pattern("mm-m3", names::MINIMAX_M3)
119 .register_pattern("minimax", names::MINIMAX_M2)
120 .register_pattern("mm-m2", names::MINIMAX_M2);
121
rust/src/chat/src/parser/tool/tests.rs
+8
@@ -157,6 +157,14 @@ fn factory_new_resolves_default_patterns() {
157 factory.resolve_name_for_model("tencent/Hy3-preview"),
158 Some(names::HY_V3)
159 );
160 + assert_eq!(
161 + factory.resolve_name_for_model("MiniMax/MiniMax-M3-Text"),
162 + Some(names::MINIMAX_M3)
163 + );
164 + assert_eq!(
165 + factory.resolve_name_for_model("org/mm-m3-base"),
166 + Some(names::MINIMAX_M3)
167 + );
168 assert_eq!(
169 factory.resolve_name_for_model("MiniMax/MiniMax-M2-01"),
170 Some(names::MINIMAX_M2)
rust/src/reasoning-parser/src/lib.rs
+2
@@ -19,6 +19,7 @@ mod deepseek_r1;
19 mod delimited;
20 mod gemma4;
21 mod kimi;
22 +mod minimax_m3;
23 mod qwen3;
24 mod seed_oss;
25 mod step3p5;
@@ -31,6 +32,7 @@ pub use self::deepseek_r1::DeepSeekR1ReasoningParser;
32 pub(crate) use self::delimited::DelimitedReasoningParser;
33 pub use self::gemma4::Gemma4ReasoningParser;
34 pub use self::kimi::KimiReasoningParser;
35 +pub use self::minimax_m3::MiniMaxM3ReasoningParser;
36 pub use self::qwen3::Qwen3ReasoningParser;
37 pub use self::seed_oss::SeedOssReasoningParser;
38 pub use self::step3p5::Step3p5ReasoningParser;
rust/src/reasoning-parser/src/minimax_m3.rs new
+98
@@ -0,0 +1,98 @@
1 +use vllm_tokenizer::DynTokenizer;
2 +
3 +use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
4 +
5 +const M3_THINK_START: &str = "<mm:think>";
6 +const M3_THINK_END: &str = "</mm:think>";
7 +
8 +/// Reasoning parser for MiniMax M3 style outputs.
9 +///
10 +/// MiniMax M3 uses `<mm:think>...</mm:think>` delimiters. Its chat template may
11 +/// prefill either delimiter depending on the requested thinking mode, so the
12 +/// shared delimited parser derives the starting state from the rendered prompt.
13 +pub struct MiniMaxM3ReasoningParser {
14 + inner: DelimitedReasoningParser,
15 + /// True until the first response text is classified. Only this position may
16 + /// drop a stray `</mm:think>` emitted at the start of a response.
17 + at_response_start: bool,
18 + /// Holds an initial suffix like `</mm` while it may still complete into the
19 + /// leading closer on a later chunk.
20 + leading_end_buffer: String,
21 +}
22 +
23 +impl MiniMaxM3ReasoningParser {
24 + /// Create a MiniMax M3 parser backed by the shared delimited state machine.
25 + pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
26 + Ok(Self {
27 + inner: DelimitedReasoningParser::new(tokenizer, M3_THINK_START, M3_THINK_END, false)?,
28 + at_response_start: true,
29 + leading_end_buffer: String::new(),
30 + })
31 + }
32 +
33 + /// Drop a response-leading `</mm:think>` while preserving later unmatched
34 + /// closers as ordinary content.
35 + fn push_inner(&mut self, delta: &str) -> ReasoningDelta {
36 + if self.at_response_start && !self.inner.in_reasoning() {
37 + self.leading_end_buffer.push_str(delta);
38 + let buffered = std::mem::take(&mut self.leading_end_buffer);
39 +
40 + if buffered.is_empty() {
41 + return ReasoningDelta::default();
42 + }
43 + if let Some(rest) = buffered.strip_prefix(M3_THINK_END) {
44 + self.at_response_start = false;
45 + return self.inner.push(rest);
46 + }
47 + if M3_THINK_END.starts_with(buffered.as_str()) {
48 + self.leading_end_buffer = buffered;
49 + return ReasoningDelta::default();
50 + }
51 +
52 + self.at_response_start = false;
53 + return self.inner.push(&buffered);
54 + }
55 +
56 + self.inner.push(delta)
57 + }
58 +}
59 +
60 +fn append_delta(target: &mut ReasoningDelta, delta: ReasoningDelta) {
61 + if let Some(reasoning) = delta.reasoning {
62 + target.push_reasoning(&reasoning);
63 + }
64 + if let Some(content) = delta.content {
65 + target.push_content(&content);
66 + }
67 +}
68 +
69 +impl ReasoningParser for MiniMaxM3ReasoningParser {
70 + fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
71 + where
72 + Self: Sized + 'static,
73 + {
74 + Ok(Box::new(Self::new(tokenizer)?))
75 + }
76 +
77 + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
78 + self.inner.initialize(prompt_token_ids);
79 + self.at_response_start = true;
80 + self.leading_end_buffer.clear();
81 + Ok(())
82 + }
83 +
84 + fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
85 + Ok(self.push_inner(delta))
86 + }
87 +
88 + fn finish(&mut self) -> Result<ReasoningDelta> {
89 + let mut delta = ReasoningDelta::default();
90 + if !self.leading_end_buffer.is_empty() {
91 + let pending = std::mem::take(&mut self.leading_end_buffer);
92 + self.at_response_start = false;
93 + append_delta(&mut delta, self.inner.push(&pending));
94 + }
95 + append_delta(&mut delta, self.inner.finish());
96 + Ok(delta)
97 + }
98 +}
rust/src/reasoning-parser/src/tests.rs
+67 -1
@@ -3,7 +3,8 @@ use std::sync::Arc;
3 use vllm_tokenizer::Tokenizer;
4
5 use super::{
6 - DeepSeekR1ReasoningParser, DelimitedReasoningParser, Qwen3ReasoningParser, ReasoningParser,
6 + DeepSeekR1ReasoningParser, DelimitedReasoningParser, MiniMaxM3ReasoningParser,
7 + Qwen3ReasoningParser, ReasoningParser,
8 };
9
10 pub(crate) struct FakeTokenizer;
@@ -32,6 +33,8 @@ impl Tokenizer for FakeTokenizer {
33 "<|END_THINKING|>" => Some(4),
34 "◁think▷" => Some(5),
35 "◁/think▷" => Some(6),
36 + "<mm:think>" => Some(8),
37 + "</mm:think>" => Some(9),
38 "<seed:think>" => Some(10),
39 "</seed:think>" => Some(11),
40 _ => None,
@@ -161,3 +164,66 @@ fn deepseek_r1_stops_scanning_at_last_special_token() {
164 assert_eq!(delta.reasoning.as_deref(), Some("reason"));
165 assert_eq!(delta.content.as_deref(), Some("answer"));
166 }
167 +
168 +#[test]
169 +fn minimax_m3_handles_explicit_think_delimiters() {
170 + let tokenizer = Arc::new(FakeTokenizer);
171 + let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
172 +
173 + let delta = parser.push("<mm:think>reason</mm:think>answer").unwrap();
174 + assert_eq!(delta.reasoning.as_deref(), Some("reason"));
175 + assert_eq!(delta.content.as_deref(), Some("answer"));
176 +}
177 +
178 +#[test]
179 +fn minimax_m3_drops_leading_end_marker() {
180 + let tokenizer = Arc::new(FakeTokenizer);
181 + let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
182 +
183 + let delta = parser.push("</mm:think>answer").unwrap();
184 + assert_eq!(delta.reasoning, None);
185 + assert_eq!(delta.content.as_deref(), Some("answer"));
186 +}
187 +
188 +#[test]
189 +fn minimax_m3_preserves_non_leading_end_marker() {
190 + let tokenizer = Arc::new(FakeTokenizer);
191 + let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
192 +
193 + let delta = parser.push("XXX</mm:think>YYY").unwrap();
194 + assert_eq!(delta.reasoning, None);
195 + assert_eq!(delta.content.as_deref(), Some("XXX</mm:think>YYY"));
196 +}
197 +
198 +#[test]
199 +fn minimax_m3_drops_split_leading_end_marker() {
200 + let tokenizer = Arc::new(FakeTokenizer);
201 + let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
202 +
203 + assert!(parser.push("</mm").unwrap().is_empty());
204 + let delta = parser.push(":think>answer").unwrap();
205 + assert_eq!(delta.reasoning, None);
206 + assert_eq!(delta.content.as_deref(), Some("answer"));
207 +}
208 +
209 +#[test]
210 +fn minimax_m3_uses_prompt_prefilled_start_marker() {
211 + let tokenizer = Arc::new(FakeTokenizer);
212 + let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
213 + parser.initialize(&[8]).unwrap();
214 +
215 + let delta = parser.push("reason</mm:think>answer").unwrap();
216 + assert_eq!(delta.reasoning.as_deref(), Some("reason"));
217 + assert_eq!(delta.content.as_deref(), Some("answer"));
218 +}
219 +
220 +#[test]
221 +fn minimax_m3_uses_prompt_prefilled_end_marker() {
222 + let tokenizer = Arc::new(FakeTokenizer);
223 + let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
224 + parser.initialize(&[9]).unwrap();
225 +
226 + let delta = parser.push("answer").unwrap();
227 + assert_eq!(delta.reasoning, None);
228 + assert_eq!(delta.content.as_deref(), Some("answer"));
229 +}
rust/src/tool-parser/python/src/lib.rs
+2
@@ -38,6 +38,8 @@ macro_rules! tool_parser_factory {
38
39 // Export a tool parser to Python by registering it here.
40 tool_parser_factory! {
41 + MinimaxM3ToolParser,
42 +
43 // Below are the parsers just for testing purposes on Python side.
44 DeepSeekV4ToolParser,
45 KimiK2ToolParser,
rust/src/tool-parser/src/lib.rs
+2
@@ -10,6 +10,7 @@ mod hy_v3;
10 mod json;
11 mod kimi_k2;
12 mod minimax_m2;
13 +mod minimax_m3;
14 mod parameters;
15 mod qwen_coder;
16 #[cfg(any(test, feature = "test-util"))]
@@ -30,6 +31,7 @@ pub use json::{
31 };
32 pub use kimi_k2::KimiK2ToolParser;
33 pub use minimax_m2::MinimaxM2ToolParser;
34 +pub use minimax_m3::MinimaxM3ToolParser;
35 pub use qwen_coder::Qwen3CoderToolParser;
36 use serde::{Deserialize, Serialize};
37 use serde_json::Value;
rust/src/tool-parser/src/minimax_m3.rs new
+885
@@ -0,0 +1,885 @@
1 +use winnow::ascii::{multispace0 as ws0, multispace1 as ws1};
2 +use winnow::combinator::{alt, delimited, seq};
3 +use winnow::error::{ContextError, ErrMode};
4 +use winnow::prelude::*;
5 +use winnow::stream::Partial;
6 +use winnow::token::{literal, rest, take_until};
7 +
8 +use super::parameters::{ParamElement, ParamInput, ToolSchemas};
9 +use super::utils::{parse_buffered_event, safe_text_len};
10 +use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
11 +use crate::Tool;
12 +
13 +const NAMESPACE: &str = "]<]minimax[>[";
14 +const TOOL_CALL_START: &str = "]<]minimax[>[<tool_call>";
15 +const TOOL_CALL_END: &str = "]<]minimax[>[</tool_call>";
16 +const INVOKE_START: &str = "]<]minimax[>[<invoke";
17 +const INVOKE_END: &str = "]<]minimax[>[</invoke>";
18 +const ELEMENT_START: &str = "]<]minimax[>[<";
19 +const ELEMENT_END_START: &str = "]<]minimax[>[</";
20 +const MIXED_TEXT_FIELD: &str = "$text";
21 +
22 +type MinimaxM3Input<'i> = Partial<&'i str>;
23 +
24 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25 +enum MinimaxM3Mode {
26 + Text,
27 + ToolBlock,
28 + Done,
29 +}
30 +
31 +#[derive(Debug, Clone, PartialEq, Eq)]
32 +enum MinimaxM3Event {
33 + Text {
34 + len: usize,
35 + },
36 + ToolBlockStart,
37 + Invoke {
38 + name: String,
39 + params: Vec<(String, ParamInput)>,
40 + },
41 + ToolBlockEnd,
42 + IgnoredRest,
43 +}
44 +
45 +/// Tool parser for MiniMax M3 namespace-delimited XML-style tool calls.
46 +///
47 +/// Example tool call content with recursive parameters:
48 +///
49 +/// ```text
50 +/// ]<]minimax[>[<tool_call>
51 +/// ]<]minimax[>[<invoke name="create_order">
52 +/// ]<]minimax[>[<user_id>42]<]minimax[>[</user_id>
53 +/// ]<]minimax[>[<shipping>
54 +/// ]<]minimax[>[<city>Singapore]<]minimax[>[</city>
55 +/// ]<]minimax[>[<zip>018956]<]minimax[>[</zip>
56 +/// ]<]minimax[>[</shipping>
57 +/// ]<]minimax[>[<items>
58 +/// ]<]minimax[>[<item>
59 +/// ]<]minimax[>[<sku>book-001]<]minimax[>[</sku>
60 +/// ]<]minimax[>[<qty>2]<]minimax[>[</qty>
61 +/// ]<]minimax[>[</item>
62 +/// ]<]minimax[>[</items>
63 +/// ]<]minimax[>[</invoke>
64 +/// ]<]minimax[>[</tool_call>
65 +/// ```
66 +///
67 +/// With a schema where `shipping` is an object and `items` is an array of
68 +/// objects, recursive parameter conversion produces:
69 +///
70 +/// ```json
71 +/// {
72 +/// "user_id": 42,
73 +/// "shipping": {
74 +/// "city": "Singapore",
75 +/// "zip": 18956
76 +/// },
77 +/// "items": [
78 +/// {
79 +/// "sku": "book-001",
80 +/// "qty": 2
81 +/// }
82 +/// ]
83 +/// }
84 +/// ```
85 +///
86 +/// MiniMax M3 emits the namespace marker `]<]minimax[>[` before each structural
87 +/// tag. Arguments are emitted only after a full `<invoke>` block is parsed.
88 +pub struct MinimaxM3ToolParser {
89 + buffer: String,
90 + mode: MinimaxM3Mode,
91 + emitted_tool_count: usize,
92 + tool_parameters: ToolSchemas,
93 +}
94 +
95 +impl MinimaxM3ToolParser {
96 + /// Create a MiniMax M3 tool parser.
97 + pub fn new(tools: &[Tool]) -> Self {
98 + Self {
99 + buffer: String::new(),
100 + mode: MinimaxM3Mode::Text,
101 + emitted_tool_count: 0,
102 + tool_parameters: ToolSchemas::from_tools(tools),
103 + }
104 + }
105 +
106 + /// Apply one parsed MiniMax M3 event to parser state and output.
107 + fn apply_event(&mut self, event: MinimaxM3Event, output: &mut ToolParserOutput) -> Result<()> {
108 + match event {
109 + MinimaxM3Event::Text { len: consumed_len } => {
110 + output.normal_text.push_str(&self.buffer[..consumed_len]);
111 + }
112 + MinimaxM3Event::ToolBlockStart => self.mode = MinimaxM3Mode::ToolBlock,
113 + MinimaxM3Event::Invoke { name, params } => {
114 + let arguments = self.tool_parameters.convert_params_with_schema(&name, params);
115 + let arguments = serde_json::to_string(&arguments)
116 + .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
117 +
118 + output.calls.push(ToolCallDelta {
119 + tool_index: self.emitted_tool_count,
120 + name: Some(name),
121 + arguments,
122 + });
123 + self.emitted_tool_count += 1;
124 + }
125 + MinimaxM3Event::ToolBlockEnd => self.mode = MinimaxM3Mode::Done,
126 + MinimaxM3Event::IgnoredRest => {}
127 + }
128 + Ok(())
129 + }
130 +}
131 +
132 +impl ToolParser for MinimaxM3ToolParser {
133 + fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
134 + where
135 + Self: Sized + 'static,
136 + {
137 + Ok(Box::new(Self::new(tools)))
138 + }
139 +
140 + fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
141 + self.buffer.push_str(chunk);
142 +
143 + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
144 + parse_next_minimax_m3_event(input, self.mode)
145 + })? {
146 + self.apply_event(event, output)?;
147 + self.buffer.drain(..consumed_len);
148 + }
149 +
150 + Ok(())
151 + }
152 +
153 + fn finish(&mut self) -> Result<ToolParserOutput> {
154 + let mut output = ToolParserOutput::default();
155 + match self.mode {
156 + MinimaxM3Mode::Text => {
157 + output.normal_text.push_str(&self.buffer);
158 + }
159 + MinimaxM3Mode::ToolBlock => {
160 + if !self.buffer.trim_start().is_empty() {
161 + return Err(parsing_failed!("incomplete MiniMax M3 tool call"));
162 + }
163 + }
164 + MinimaxM3Mode::Done => {}
165 + }
166 + let _ = self.reset();
167 + Ok(output)
168 + }
169 +
170 + fn reset(&mut self) -> String {
171 + self.mode = MinimaxM3Mode::Text;
172 + self.emitted_tool_count = 0;
173 + std::mem::take(&mut self.buffer)
174 + }
175 +}
176 +
177 +/// Parse a MiniMax M3 event for the current parser mode.
178 +fn parse_next_minimax_m3_event(
179 + input: &mut MinimaxM3Input<'_>,
180 + mode: MinimaxM3Mode,
181 +) -> ModalResult<MinimaxM3Event> {
182 + match mode {
183 + MinimaxM3Mode::Text => parse_text_event(input),
184 + MinimaxM3Mode::ToolBlock => parse_tool_block_event(input),
185 + MinimaxM3Mode::Done => ignored_rest_event(input),
186 + }
187 +}
188 +
189 +/// Parse a text-mode MiniMax M3 event.
190 +fn parse_text_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
191 + alt((tool_block_start_event, safe_text_event)).parse_next(input)
192 +}
193 +
194 +/// Parse a MiniMax M3 tool-block start marker.
195 +fn tool_block_start_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
196 + literal(TOOL_CALL_START).value(MinimaxM3Event::ToolBlockStart).parse_next(input)
197 +}
198 +
199 +/// Parse a safe text run before the next MiniMax M3 marker.
200 +fn safe_text_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
201 + safe_text_len(input, TOOL_CALL_START).map(|len| MinimaxM3Event::Text { len })
202 +}
203 +
204 +/// Parse one event inside a MiniMax M3 tool block.
205 +fn parse_tool_block_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
206 + alt((tool_block_end_event, invoke_event)).parse_next(input)
207 +}
208 +
209 +/// Parse a MiniMax M3 tool-block end marker.
210 +fn tool_block_end_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
211 + (ws0, literal(TOOL_CALL_END))
212 + .value(MinimaxM3Event::ToolBlockEnd)
213 + .parse_next(input)
214 +}
215 +
216 +/// Parse a complete MiniMax M3 invoke block.
217 +fn invoke_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
218 + let (name, body) = seq!(
219 + _: ws0,
220 + _: literal(INVOKE_START),
221 + _: (ws1, literal("name=")),
222 + partial_attr_value,
223 + _: literal(">"),
224 + take_until(0.., INVOKE_END),
225 + _: literal(INVOKE_END),
226 + )
227 + .parse_next(input)?;
228 + let params = parse_invoke_params(body)?;
229 +
230 + Ok(MinimaxM3Event::Invoke {
231 + name: name.trim().to_string(),
232 + params,
233 + })
234 +}
235 +
236 +/// Parse all parameter elements inside a complete MiniMax M3 invoke body.
237 +fn parse_invoke_params(invoke_body: &str) -> ModalResult<Vec<(String, ParamInput)>> {
238 + let mut input = invoke_body;
239 + let mut elements = Vec::new();
240 +
241 + loop {
242 + let _ = ws0.parse_next(&mut input)?;
243 + if input.is_empty() {
244 + break;
245 + }
246 + if input.starts_with(ELEMENT_START) {
247 + elements.push(parameter_element(&mut input)?);
248 + continue;
249 + }
250 + if input.starts_with(NAMESPACE) {
251 + return malformed();
252 + }
253 + // Be tolerant: ordinary text at an invokeparameter boundary ends this invoke.
254 + // Keep parsed parameters and drop the remaining invoke body.
255 + break;
256 + }
257 +
258 + Ok(elements.into_iter().map(|element| (element.name, element.value)).collect())
259 +}
260 +
261 +/// Parse a MiniMax M3 parameter element.
262 +fn parameter_element(input: &mut &str) -> ModalResult<ParamElement> {
263 + let name = open_element_tag(input)?.to_string();
264 + let value = element_body(input, &name)?;
265 + close_element_tag(input, &name)?;
266 + Ok(ParamElement { name, value })
267 +}
268 +
269 +/// Parse a MiniMax M3 opening element tag.
270 +fn open_element_tag<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
271 + let name = seq!(
272 + _: literal(ELEMENT_START),
273 + take_until(1.., ">"),
274 + _: literal(">"),
275 + )
276 + .parse_next(input)?;
277 +
278 + let name = name.0;
279 + if name.starts_with('/') || name.trim().is_empty() {
280 + return malformed();
281 + }
282 +
283 + Ok(name)
284 +}
285 +
286 +/// Parse a MiniMax M3 closing element tag.
287 +fn close_element_tag(input: &mut &str, name: &str) -> ModalResult<()> {
288 + literal(ELEMENT_END_START).void().parse_next(input)?;
289 + literal(name).void().parse_next(input)?;
290 + literal(">").void().parse_next(input)
291 +}
292 +
293 +/// Parse the body of one MiniMax M3 element.
294 +fn element_body(input: &mut &str, closing_name: &str) -> ModalResult<ParamInput> {
295 + let close_tag = format!("{ELEMENT_END_START}{closing_name}>");
296 + let mut text = String::new();
297 + let mut elements = Vec::new();
298 +
299 + loop {
300 + text.push_str(text_until_namespace(input)?);
301 +
302 + if input.starts_with(&close_tag) {
303 + // Close tag reached, end of element body.
304 + break;
305 + }
306 + if input.starts_with(ELEMENT_START) {
307 + // Child element start reached, parse child element recursively.
308 + elements.push(parameter_element(input)?);
309 + continue;
310 + }
311 + if input.starts_with(NAMESPACE) {
312 + // Unexpected namespace marker.
313 + return malformed();
314 + }
315 + }
316 +
317 + if elements.is_empty() {
318 + Ok(ParamInput::Text(text))
319 + } else {
320 + if !text.trim().is_empty() {
321 + push_mixed_text_element(&mut elements, text);
322 + }
323 + Ok(ParamInput::Elements(elements))
324 + }
325 +}
326 +
327 +/// Parse text until the next MiniMax M3 namespace marker.
328 +fn text_until_namespace<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
329 + take_until(0.., NAMESPACE).parse_next(input)
330 +}
331 +
332 +/// Preserve mixed text content under a reserved object field.
333 +///
334 +/// By default, the field name is `$text`, but if that collides with an existing
335 +/// child element name, prepend `$` until there is no collision.
336 +fn push_mixed_text_element(elements: &mut Vec<ParamElement>, text: String) {
337 + let mut name = MIXED_TEXT_FIELD.to_string();
338 + while elements.iter().any(|element| element.name == name) {
339 + name.insert(0, '$');
340 + }
341 + elements.push(ParamElement {
342 + name,
343 + value: ParamInput::Text(text),
344 + });
345 +}
346 +
347 +/// Parse a quoted or unquoted XML attribute value from partial streaming input.
348 +fn partial_attr_value<'i>(input: &mut MinimaxM3Input<'i>) -> ModalResult<&'i str> {
349 + alt((
350 + delimited(literal("\""), take_until(1.., "\""), literal("\"")),
351 + delimited(literal("'"), take_until(1.., "'"), literal("'")),
352 + take_until(1.., ">"),
353 + ))
354 + .parse_next(input)
355 +}
356 +
357 +/// Parse ignored rest after the MiniMax M3 tool block ends.
358 +fn ignored_rest_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
359 + rest.value(MinimaxM3Event::IgnoredRest).parse_next(input)
360 +}
361 +
362 +fn malformed<T>() -> ModalResult<T> {
363 + Err(ErrMode::Cut(ContextError::new()))
364 +}
365 +
366 +#[cfg(test)]
367 +mod tests {
368 + use expect_test::expect;
369 + use serde_json::{Value, json};
370 + use thiserror_ext::AsReport;
371 +
372 + use super::{
373 + ELEMENT_END_START, ELEMENT_START, INVOKE_END, INVOKE_START, MinimaxM3ToolParser,
374 + TOOL_CALL_END, TOOL_CALL_START, ToolParser,
375 + };
376 + use crate::test_utils::{collect_stream, split_by_chars, test_tools};
377 + use crate::{Tool, ToolParserTestExt as _};
378 +
379 + fn element(name: &str, body: &str) -> String {
380 + format!("{ELEMENT_START}{name}>{body}{ELEMENT_END_START}{name}>")
381 + }
382 +
383 + fn invoke(function_name: &str, body: &str) -> String {
384 + format!("{INVOKE_START} name=\"{function_name}\">{body}{INVOKE_END}")
385 + }
386 +
387 + fn build_tool_block(invokes: &[(&str, String)]) -> String {
388 + let invokes = invokes
389 + .iter()
390 + .map(|(function_name, body)| invoke(function_name, body))
391 + .collect::<Vec<_>>()
392 + .join("\n");
393 + format!("{TOOL_CALL_START}\n{invokes}\n{TOOL_CALL_END}")
394 + }
395 +
396 + fn m3_test_tools() -> Vec<Tool> {
397 + let mut tools = test_tools();
398 + tools.push(Tool {
399 + name: "create_order".to_string(),
400 + description: None,
401 + parameters: json!({
402 + "type": "object",
403 + "properties": {
404 + "user_id": { "type": "integer" },
405 + "urgent": { "type": "boolean" },
406 + "note": { "type": "string" },
407 + "shipping": {
408 + "type": "object",
409 + "properties": {
410 + "city": { "type": "string" },
411 + "zip": { "type": "integer" }
412 + }
413 + },
414 + "items": {
415 + "type": "array",
416 + "items": {
417 + "type": "object",
418 + "properties": {
419 + "sku": { "type": "string" },
420 + "qty": { "type": "integer" }
421 + }
422 + }
423 + },
424 + "metadata": {
425 + "type": "object",
426 + "additionalProperties": { "type": "integer" }
427 + },
428 + "duplicate_demo": {
429 + "type": "object",
430 + "properties": {
431 + "tag": { "type": "string" }
432 + }
433 + },
434 + "schema_mismatch_array": {
435 + "type": "array",
436 + "items": { "type": "integer" }
437 + }
438 + }
439 + }),
440 + strict: None,
441 + });
442 + tools
443 + }
444 +
445 + fn order_arguments() -> String {
446 + let shipping = element(
447 + "shipping",
448 + &format!(
449 + "{}{}",
450 + element("city", "Singapore"),
451 + element("zip", "018956")
452 + ),
453 + );
454 + let first_item = element(
455 + "item",
456 + &format!("{}{}", element("sku", "book-001"), element("qty", "2")),
457 + );
458 + let second_item = element(
459 + "item",
460 + &format!("{}{}", element("sku", "pen-007"), element("qty", "5")),
461 + );
462 + let items = element("items", &format!("{first_item}{second_item}"));
463 + let metadata = element(
464 + "metadata",
465 + &format!("{}{}", element("score", "42"), element("rank", "7")),
466 + );
467 + let duplicate_demo = element(
468 + "duplicate_demo",
469 + &format!("{}{}", element("tag", "a"), element("tag", "b")),
470 + );
471 + let schema_mismatch_array = element(
472 + "schema_mismatch_array",
473 + &format!("{}{}", element("x", "1"), element("x", "2")),
474 + );
475 +
476 + [
477 + element("user_id", "42"),
478 + element("urgent", "true"),
479 + element("note", "Please leave at front desk."),
480 + shipping,
481 + items,
482 + metadata,
483 + duplicate_demo,
484 + schema_mismatch_array,
485 + element(
486 + "unknown_struct",
487 + &format!("{}{}", element("a", "1"), element("a", "2")),
488 + ),
489 + ]
490 + .join("")
491 + }
492 +
493 + #[test]
494 + fn minimax_m3_parse_complete_without_tool_call_keeps_text() {
495 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
496 + let output = parser.parse_complete("Hello, world!").unwrap();
497 +
498 + assert_eq!(output.normal_text, "Hello, world!");
499 + assert!(output.calls.is_empty());
500 + }
501 +
502 + #[test]
503 + fn minimax_m3_parse_complete_extracts_single_tool_call() {
504 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
505 + let output = parser
506 + .parse_complete(&build_tool_block(&[(
507 + "get_weather",
508 + format!("{}{}", element("city", "Seattle"), element("days", "5")),
509 + )]))
510 + .unwrap();
511 +
512 + assert!(output.normal_text.is_empty());
513 + assert_eq!(output.calls.len(), 1);
514 + assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
515 + assert_eq!(
516 + serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
517 + json!({ "city": "Seattle", "days": 5 })
518 + );
519 + }
520 +
521 + #[test]
522 + fn minimax_m3_parse_complete_preserves_prefix_and_ignores_trailing_text() {
523 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
524 + let output = format!(
525 + "Let me check. {} This trailing text is ignored.",
526 + build_tool_block(&[("get_weather", element("city", "Seattle"))])
527 + );
528 + let output = parser.parse_complete(&output).unwrap();
529 +
530 + assert_eq!(output.normal_text, "Let me check. ");
531 + assert_eq!(output.calls.len(), 1);
532 + }
533 +
534 + #[test]
535 + fn minimax_m3_parse_complete_extracts_multiple_invokes() {
536 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
537 + let output = parser
538 + .parse_complete(&build_tool_block(&[
539 + ("get_weather", element("city", "Seattle")),
540 + ("get_weather", element("city", "NYC")),
541 + ]))
542 + .unwrap();
543 +
544 + assert_eq!(output.calls.len(), 2);
545 + assert_eq!(output.calls[0].tool_index, 0);
546 + assert_eq!(output.calls[1].tool_index, 1);
547 + assert_eq!(
548 + serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
549 + json!({ "city": "Seattle" })
550 + );
551 + assert_eq!(
552 + serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
553 + json!({ "city": "NYC" })
554 + );
555 + }
556 +
557 + #[test]
558 + fn minimax_m3_invoke_body_junk_drops_rest_of_invoke() {
559 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
560 + let output = parser
561 + .parse_complete(&build_tool_block(&[(
562 + "get_weather",
563 + [
564 + element("city", "Seattle"),
565 + "I need to use the city above.".to_string(),
566 + element("days", "5"),
567 + ]
568 + .join(""),
569 + )]))
570 + .unwrap();
571 +
572 + assert_eq!(
573 + serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
574 + json!({ "city": "Seattle" })
575 + );
576 + }
577 +
578 + #[test]
579 + fn minimax_m3_parse_complete_converts_schema_types() {
580 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
581 + let output = parser
582 + .parse_complete(&build_tool_block(&[(
583 + "convert",
584 + [
585 + element("whole", "5.0"),
586 + element("flag", "true"),
587 + element("payload", r#"{"nested":true}"#),
588 + element("items", "[1,2]"),
589 + element("empty", "42"),
590 + ]
591 + .join(""),
592 + )]))
593 + .unwrap();
594 +
595 + assert_eq!(
596 + serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
597 + json!({
598 + "whole": 5.0,
599 + "flag": true,
600 + "payload": { "nested": true },
601 + "items": [1, 2],
602 + "empty": "42",
603 + })
604 + );
605 + }
606 +
607 + #[test]
608 + fn minimax_m3_parse_complete_converts_nested_arguments() {
609 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
610 + let output = parser
611 + .parse_complete(&build_tool_block(&[("create_order", order_arguments())]))
612 + .unwrap();
613 +
614 + assert_eq!(
615 + serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
616 + json!({
617 + "user_id": 42,
618 + "urgent": true,
619 + "note": "Please leave at front desk.",
620 + "shipping": {
621 + "city": "Singapore",
622 + "zip": 18956
623 + },
624 + "items": [
625 + {
626 + "sku": "book-001",
627 + "qty": 2
628 + },
629 + {
630 + "sku": "pen-007",
631 + "qty": 5
632 + }
633 + ],
634 + "metadata": {
635 + "score": 42,
636 + "rank": 7
637 + },
638 + "duplicate_demo": {
639 + "tag": ["a", "b"]
640 + },
641 + "schema_mismatch_array": [1, 2],
642 + "unknown_struct": {
643 + "a": ["1", "2"]
644 + }
645 + })
646 + );
647 + }
648 +
649 + #[test]
650 + fn minimax_m3_parse_complete_handles_multiline_leaf_parameters() {
651 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
652 + let output = parser
653 + .parse_complete(&build_tool_block(&[(
654 + "calculate_area",
655 + [
656 + element("shape", "\nrectangle\n"),
657 + element("dimensions", r#"{"width":10,"height":20}"#),
658 + element("precision", "2"),
659 + ]
660 + .join(""),
661 + )]))
662 + .unwrap();
663 +
664 + assert_eq!(
665 + serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
666 + json!({
667 + "shape": "\nrectangle\n",
668 + "dimensions": { "width": 10, "height": 20 },
669 + "precision": 2,
670 + })
671 + );
672 + }
673 +
674 + #[test]
675 + fn minimax_m3_streaming_extracts_single_tool_call() {
676 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
677 + let output = collect_stream(
678 + &mut parser,
679 + &[
680 + TOOL_CALL_START,
681 + &invoke("get_weather", &element("city", "Seattle")),
682 + TOOL_CALL_END,
683 + ],
684 + );
685 +
686 + assert!(output.normal_text.is_empty());
687 + assert_eq!(output.calls.len(), 1);
688 + assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
689 + assert_eq!(
690 + serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
691 + json!({ "city": "Seattle" })
692 + );
693 + }
694 +
695 + #[test]
696 + fn minimax_m3_streaming_preserves_prefix_text() {
697 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
698 + let output = collect_stream(
699 + &mut parser,
700 + &[
701 + "Let me check. ",
702 + TOOL_CALL_START,
703 + &invoke("get_weather", &element("city", "Seattle")),
704 + TOOL_CALL_END,
705 + ],
706 + );
707 +
708 + assert_eq!(output.normal_text, "Let me check. ");
709 + assert_eq!(output.calls.len(), 1);
710 + }
711 +
712 + #[test]
713 + fn minimax_m3_streaming_without_tool_call_emits_text_incrementally() {
714 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
715 + let output = collect_stream(&mut parser, &["Hello, ", "world!"]);
716 +
717 + assert_eq!(output.normal_text, "Hello, world!");
718 + assert!(output.calls.is_empty());
719 + }
720 +
721 + #[test]
722 + fn minimax_m3_streaming_handles_marker_split_across_chunks() {
723 + let text = build_tool_block(&[("get_weather", element("city", "Seattle"))]);
724 + let chunks = split_by_chars(&text, 3);
725 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
726 + let output = collect_stream(&mut parser, &chunks);
727 +
728 + assert_eq!(output.calls.len(), 1);
729 + assert!(output.normal_text.is_empty());
730 + }
731 +
732 + #[test]
733 + fn minimax_m3_streaming_extracts_multiple_invokes_in_order() {
734 + let text = build_tool_block(&[
735 + ("get_weather", element("city", "Seattle")),
736 + ("get_weather", element("city", "NYC")),
737 + ]);
738 + let chunks = split_by_chars(&text, 7);
739 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
740 + let output = collect_stream(&mut parser, &chunks);
741 +
742 + assert_eq!(output.calls.len(), 2);
743 + assert_eq!(output.calls[0].tool_index, 0);
744 + assert_eq!(output.calls[1].tool_index, 1);
745 + }
746 +
747 + #[test]
748 + fn minimax_m3_streaming_does_not_emit_incomplete_tool_call() {
749 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
750 + let output = parser
751 + .parse_chunk(&format!(
752 + "{TOOL_CALL_START}{INVOKE_START} name=\"get_weather\">"
753 + ))
754 + .unwrap();
755 +
756 + assert!(output.normal_text.is_empty());
757 + assert!(output.calls.is_empty());
758 + }
759 +
760 + #[test]
761 + fn minimax_m3_streaming_ignores_text_after_tool_block() {
762 + let text = format!(
763 + "{} ignored",
764 + build_tool_block(&[("get_weather", element("city", "Seattle"))])
765 + );
766 + let chunks = split_by_chars(&text, 5);
767 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
768 + let output = collect_stream(&mut parser, &chunks);
769 +
770 + assert!(output.normal_text.is_empty());
771 + assert_eq!(output.calls.len(), 1);
772 + }
773 +
774 + #[test]
775 + fn minimax_m3_finish_fails_incomplete_tool_call() {
776 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
777 + parser
778 + .parse_chunk(&format!(
779 + "{TOOL_CALL_START}{INVOKE_START} name=\"get_weather\">"
780 + ))
781 + .unwrap();
782 +
783 + assert!(parser.finish().is_err());
784 + }
785 +
786 + #[test]
787 + fn minimax_m3_finish_recovers_after_bare_tool_block_start() {
788 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
789 + parser.parse_chunk(TOOL_CALL_START).unwrap();
790 +
791 + let output = parser.finish().unwrap();
792 + assert!(output.normal_text.is_empty());
793 + assert!(output.calls.is_empty());
794 + }
795 +
796 + #[test]
797 + fn minimax_m3_finish_recovers_completed_invoke_with_whitespace_tail() {
798 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
799 + let output = parser
800 + .parse_complete(&format!(
801 + "{}\n{}\n \n",
802 + TOOL_CALL_START,
803 + invoke("get_weather", &element("city", "Seattle"))
804 + ))
805 + .unwrap();
806 +
807 + assert_eq!(output.calls.len(), 1);
808 + assert_eq!(
809 + serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
810 + json!({ "city": "Seattle" })
811 + );
812 + }
813 +
814 + #[test]
815 + fn minimax_m3_finish_fails_partial_outer_end_marker() {
816 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
817 + parser
818 + .parse_chunk(&format!(
819 + "{}\n{}\n{}",
820 + TOOL_CALL_START,
821 + invoke("get_weather", &element("city", "Seattle")),
822 + &TOOL_CALL_END[..3]
823 + ))
824 + .unwrap();
825 +
826 + assert!(parser.finish().is_err());
827 + }
828 +
829 + #[test]
830 + fn minimax_m3_malformed_tool_call_fails_fast() {
831 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
832 + let error = parser
833 + .parse_chunk(&format!(
834 + "{TOOL_CALL_START}{ELEMENT_START}bad>{TOOL_CALL_END}"
835 + ))
836 + .unwrap_err();
837 +
838 + expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
839 + }
840 +
841 + #[test]
842 + fn minimax_m3_mixed_content_is_preserved_as_text_field() {
843 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
844 + let body = element(
845 + "payload",
846 + &format!("text before {} text after", element("child", "value")),
847 + );
848 + let output = parser.parse_complete(&build_tool_block(&[("convert", body)])).unwrap();
849 +
850 + assert_eq!(
851 + serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
852 + json!({
853 + "payload": {
854 + "child": "value",
855 + "$text": "text before text after"
856 + }
857 + })
858 + );
859 + }
860 +
861 + #[test]
862 + fn minimax_m3_mixed_text_field_avoids_child_name_collision() {
863 + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
864 + let body = element(
865 + "payload",
866 + &format!(
867 + "text{}{}",
868 + element("$text", "child text"),
869 + element("child", "value")
870 + ),
871 + );
872 + let output = parser.parse_complete(&build_tool_block(&[("convert", body)])).unwrap();
873 +
874 + assert_eq!(
875 + serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
876 + json!({
877 + "payload": {
878 + "$text": "child text",
879 + "$$text": "text",
880 + "child": "value"
881 + }
882 + })
883 + );
884 + }
885 +}
setup.py
+19
@@ -432,6 +432,19 @@ class cmake_build_ext(build_ext):
432 dirs_exist_ok=True,
433 )
434
435 + # copy vendored fmha_sm100 package from build_lib to source tree
436 + # for editable installs
437 + fmha_sm100_build = os.path.join(
438 + self.build_lib, "vllm", "third_party", "fmha_sm100"
439 + )
440 + if os.path.exists(fmha_sm100_build):
441 + print(f"Copying {fmha_sm100_build} to vllm/third_party/fmha_sm100")
442 + shutil.copytree(
443 + fmha_sm100_build,
444 + "vllm/third_party/fmha_sm100",
445 + dirs_exist_ok=True,
446 + )
447 +
448
449 class precompiled_build_ext(build_ext):
450 """Disables extension building when using precompiled binaries."""
@@ -787,6 +800,7 @@ class precompiled_wheel_utils:
800 )
801 # DeepGEMM: extract all files (.py, .so, .cuh, .h, .hpp, etc.)
802 deep_gemm_regex = re.compile(r"vllm/third_party/deep_gemm/.*")
803 + fmha_sm100_regex = re.compile(r"vllm/third_party/fmha_sm100/.*")
804 file_members = []
805 for member in wheel.filelist:
806 if member.filename in exact_members:
@@ -812,6 +826,7 @@ class precompiled_wheel_utils:
826 or triton_kernels_regex.match(member.filename)
827 or flashmla_regex.match(member.filename)
828 or deep_gemm_regex.match(member.filename)
829 + or fmha_sm100_regex.match(member.filename)
830 ):
831 file_members.append(member)
832
@@ -1120,6 +1135,8 @@ if _is_cuda():
1135 # DeepGEMM requires CUDA 12.3+ (SM90/SM100)
1136 # Optional since it won't build on unsupported architectures
1137 ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True))
1138 + # fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party.
1139 + ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True))
1140
1141 if _is_cpu():
1142 import platform
@@ -1150,6 +1167,8 @@ package_data = {
1167 "third_party/deep_gemm/include/**/*.cuh",
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",
1172 ]
1173 }
1174
tests/kernels/attention/test_minimax_m3.py new
+854
@@ -0,0 +1,854 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Correctness tests for MiniMax M3 sparse prefill attention kernels."""
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm import _custom_ops as ops
9 +from vllm.models.minimax_m3.common.indexer import (
10 + MiniMaxM3IndexerBackend,
11 +)
12 +from vllm.models.minimax_m3.common.ops.index_topk import (
13 + minimax_m3_index_decode,
14 + minimax_m3_index_score,
15 + minimax_m3_index_topk,
16 +)
17 +from vllm.models.minimax_m3.common.ops.sparse_attn import (
18 + minimax_m3_sparse_attn,
19 + minimax_m3_sparse_attn_decode,
20 +)
21 +from vllm.models.minimax_m3.common.sparse_attention import (
22 + MiniMaxM3SparseBackend,
23 +)
24 +from vllm.platforms import current_platform
25 +from vllm.v1.attention.backends.utils import set_kv_cache_layout
26 +from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
27 +from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache
28 +from vllm.v1.worker.utils import AttentionGroup
29 +
30 +if not (current_platform.is_cuda() or current_platform.is_rocm()):
31 + pytest.skip(
32 + "MiniMax M3 attention kernels require CUDA or ROCm.",
33 + allow_module_level=True,
34 + )
35 +
36 +
37 +@pytest.fixture
38 +def kv_layout(request):
39 + """Set the global KV cache layout for one test and restore it after."""
40 + set_kv_cache_layout(request.param)
41 + try:
42 + yield request.param
43 + finally:
44 + set_kv_cache_layout(None)
45 +
46 +
47 +def _stride_order_for(backend: type[MiniMaxM3SparseBackend], ndim: int) -> tuple:
48 + """Mirror the allocator's stride-order resolution (identity fallback)."""
49 + try:
50 + stride_order = backend.get_kv_cache_stride_order()
51 + assert len(stride_order) == ndim
52 + except (AttributeError, NotImplementedError):
53 + stride_order = tuple(range(ndim))
54 + return stride_order
55 +
56 +
57 +def _allocate_main_kv_via_contract(
58 + num_pages: int, device: torch.device | str = "cuda"
59 +) -> torch.Tensor:
60 + """Build the main KV cache exactly as the production allocator does for the
61 + currently active layout: allocate the physical (permuted) tensor, then
62 + expose the inverse-permuted logical-NHD view the backend sees."""
63 + logical_shape = MiniMaxM3SparseBackend.get_kv_cache_shape(
64 + num_pages, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
65 + )
66 + stride_order = _stride_order_for(MiniMaxM3SparseBackend, len(logical_shape))
67 + physical_shape = tuple(logical_shape[i] for i in stride_order)
68 + inv_order = [stride_order.index(i) for i in range(len(stride_order))]
69 + raw = torch.randn(physical_shape, device=device, dtype=DTYPE)
70 + return raw.permute(*inv_order)
71 +
72 +
73 +NUM_Q_HEADS = 32
74 +NUM_KV_HEADS = 2
75 +HEAD_DIM = 128
76 +BLOCK_SIZE = 128
77 +DTYPE = torch.bfloat16
78 +SM_SCALE = HEAD_DIM**-0.5
79 +TOPK = 16
80 +
81 +
82 +# Index top-k kernels.
83 +def _reference_index_topk(
84 + idx_q: torch.Tensor,
85 + index_kv_cache: torch.Tensor,
86 + block_table: torch.Tensor,
87 + q_lens: torch.Tensor,
88 + seq_lens: torch.Tensor,
89 + prefix_lens: torch.Tensor,
90 + topk: int,
91 + init_blocks: int,
92 + local_blocks: int,
93 + sm_scale: float,
94 +) -> torch.Tensor:
95 + total_q, num_idx_heads, _ = idx_q.shape
96 + out = torch.full(
97 + (num_idx_heads, total_q, topk), -1, device=idx_q.device, dtype=torch.int32
98 + )
99 +
100 + q_start = 0
101 + for req_id, (q_len, seq_len, prefix_len) in enumerate(
102 + zip(q_lens.tolist(), seq_lens.tolist(), prefix_lens.tolist())
103 + ):
104 + q_end = q_start + q_len
105 + q = idx_q[q_start:q_end]
106 + num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
107 + pages = block_table[req_id, :num_blocks]
108 + k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1)
109 + score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale
110 +
111 + q_pos = prefix_len + torch.arange(q_len, device=idx_q.device)
112 + k_pos = torch.arange(k.shape[0], device=idx_q.device)
113 + score.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf"))
114 + score = score.reshape(num_idx_heads, q_len, num_blocks, BLOCK_SIZE)
115 + score_tensor = score.max(dim=3).values
116 +
117 + valid_blocks = (q_pos + BLOCK_SIZE) // BLOCK_SIZE
118 + for local_q, num_valid_blocks in enumerate(valid_blocks.tolist()):
119 + end = min(init_blocks, num_valid_blocks)
120 + score_tensor[:, local_q, :end] = 1e30
121 + start = max(0, num_valid_blocks - local_blocks)
122 + score_tensor[:, local_q, start:num_valid_blocks] = 1e29
123 +
124 + k = min(topk, num_valid_blocks)
125 + topk_idx = score_tensor[:, local_q].topk(k, dim=1).indices
126 + out[:, q_start + local_q, :k] = topk_idx
127 + q_start = q_end
128 +
129 + return out
130 +
131 +
132 +def _assert_topk_indices_equal_unordered(
133 + actual: torch.Tensor,
134 + expected: torch.Tensor,
135 +) -> None:
136 + """Compare selected sparse blocks without requiring a deterministic order."""
137 + assert actual.shape == expected.shape
138 + actual_flat = actual.cpu().reshape(-1, actual.shape[-1]).tolist()
139 + expected_flat = expected.cpu().reshape(-1, expected.shape[-1]).tolist()
140 + for actual_row, expected_row in zip(actual_flat, expected_flat):
141 + assert set(actual_row) == set(expected_row)
142 +
143 +
144 +def test_prefill_index_topk_correctness():
145 + topk = 6
146 + init_blocks = 0
147 + local_blocks = 1
148 + num_idx_heads = 2
149 + head_dim = 16
150 + q_lens = torch.tensor((4, 3), device="cuda", dtype=torch.int32)
151 + prefix_lens = torch.tensor((0, 1024), device="cuda", dtype=torch.int32)
152 + seq_lens = prefix_lens + q_lens
153 + batch = q_lens.numel()
154 + max_seq_len = seq_lens.max().item()
155 + max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
156 + num_pages = batch * max_blocks
157 +
158 + cu_seqlens = torch.zeros(batch + 1, device="cuda", dtype=torch.int32)
159 + cu_seqlens[1:] = q_lens.cumsum(0)
160 + block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(
161 + batch, max_blocks
162 + )
163 + idx_q = torch.ones(q_lens.sum().item(), num_idx_heads, head_dim, device="cuda")
164 + index_kv_cache = torch.empty(num_pages, BLOCK_SIZE, head_dim, device="cuda")
165 + for req_id in range(batch):
166 + for block_id in range(max_blocks):
167 + page = block_table[req_id, block_id]
168 + index_kv_cache[page].fill_(block_id + 1)
169 +
170 + score = minimax_m3_index_score(
171 + idx_q,
172 + index_kv_cache,
173 + block_table,
174 + cu_seqlens,
175 + seq_lens,
176 + prefix_lens,
177 + max_query_len=q_lens.max().item(),
178 + max_seq_len=max_seq_len,
179 + num_kv_heads=num_idx_heads,
180 + sm_scale=head_dim**-0.5,
181 + )
182 + actual = minimax_m3_index_topk(
183 + score,
184 + cu_seqlens,
185 + prefix_lens,
186 + max_query_len=q_lens.max().item(),
187 + topk=topk,
188 + init_blocks=init_blocks,
189 + local_blocks=local_blocks,
190 + )
191 + expected = _reference_index_topk(
192 + idx_q,
193 + index_kv_cache,
194 + block_table,
195 + q_lens,
196 + seq_lens,
197 + prefix_lens,
198 + topk,
199 + init_blocks,
200 + local_blocks,
201 + head_dim**-0.5,
202 + )
203 + _assert_topk_indices_equal_unordered(actual, expected)
204 +
205 +
206 +@pytest.mark.parametrize("decode_query_len", [1, 4])
207 +@pytest.mark.parametrize("num_padded_reqs", [0, 2])
208 +def test_decode_index_topk_correctness(
209 + decode_query_len: int,
210 + num_padded_reqs: int,
211 +):
212 + topk = 6
213 + init_blocks = 0
214 + local_blocks = 1
215 + num_idx_heads = 2
216 + head_dim = 16
217 + active_seq_lens = torch.tensor((7, 129, 1025), device="cuda", dtype=torch.int32)
218 + q_lens = torch.full_like(active_seq_lens, decode_query_len)
219 + prefix_lens = active_seq_lens - decode_query_len
220 + active_batch = active_seq_lens.numel()
221 + batch = active_batch + num_padded_reqs
222 + seq_lens = torch.cat(
223 + [
224 + active_seq_lens,
225 + torch.zeros(num_padded_reqs, device="cuda", dtype=torch.int32),
226 + ]
227 + )
228 + max_seq_len = active_seq_lens.max().item()
229 + max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
230 + num_pages = active_batch * max_blocks
231 +
232 + active_block_table = torch.randperm(
233 + num_pages, device="cuda", dtype=torch.int32
234 + ).reshape(active_batch, max_blocks)
235 + block_table = torch.zeros(batch, max_blocks, device="cuda", dtype=torch.int32)
236 + block_table[:active_batch] = active_block_table
237 + idx_q = torch.randn(
238 + batch * decode_query_len, num_idx_heads, head_dim, device="cuda"
239 + )
240 + index_kv_cache = torch.randn(num_pages, BLOCK_SIZE, head_dim, device="cuda")
241 +
242 + actual = minimax_m3_index_decode(
243 + idx_q,
244 + index_kv_cache,
245 + block_table,
246 + seq_lens,
247 + max_seq_len=max_seq_len,
248 + topk=topk,
249 + init_blocks=init_blocks,
250 + local_blocks=local_blocks,
251 + num_kv_heads=num_idx_heads,
252 + sm_scale=head_dim**-0.5,
253 + decode_query_len=decode_query_len,
254 + )
255 + expected = torch.full_like(actual, -1)
256 + active_tokens = active_batch * decode_query_len
257 + expected[:, :active_tokens] = _reference_index_topk(
258 + idx_q[:active_tokens],
259 + index_kv_cache,
260 + block_table[:active_batch],
261 + q_lens,
262 + active_seq_lens,
263 + prefix_lens,
264 + topk,
265 + init_blocks,
266 + local_blocks,
267 + head_dim**-0.5,
268 + )
269 + _assert_topk_indices_equal_unordered(actual, expected)
270 +
271 +
272 +# Sparse attention kernels.
273 +def _reference_sparse_attn(
274 + q: torch.Tensor,
275 + kv_cache: torch.Tensor,
276 + topk_idx: torch.Tensor,
277 + block_table: torch.Tensor,
278 + q_lens: torch.Tensor,
279 + seq_lens: torch.Tensor,
280 + prefix_lens: torch.Tensor,
281 +) -> torch.Tensor:
282 + out = torch.empty_like(q, dtype=torch.float32)
283 + gqa_group_size = NUM_Q_HEADS // NUM_KV_HEADS
284 + q_start = 0
285 + for req_id, (q_len, seq_len, prefix_len) in enumerate(
286 + zip(q_lens.tolist(), seq_lens.tolist(), prefix_lens.tolist())
287 + ):
288 + q_end = q_start + q_len
289 + q_req = q[q_start:q_end]
290 + positions = torch.arange(seq_len, device="cuda")
291 + pages = block_table[req_id, positions // BLOCK_SIZE]
292 + rows = positions % BLOCK_SIZE
293 + k_req = kv_cache[pages, 0, rows]
294 + v_req = kv_cache[pages, 1, rows].float()
295 +
296 + q_pos = prefix_len + torch.arange(q_len, device="cuda")
297 + key_blocks = positions // BLOCK_SIZE
298 + causal_mask = positions.unsqueeze(0) <= q_pos.unsqueeze(1)
299 +
300 + for kv_head in range(NUM_KV_HEADS):
301 + selected = topk_idx[kv_head, q_start:q_end]
302 + selected_mask = (key_blocks[None, :, None] == selected[:, None, :]).any(-1)
303 + mask = causal_mask & selected_mask
304 + head_start = kv_head * gqa_group_size
305 + head_end = head_start + gqa_group_size
306 +
307 + q_heads = q_req[:, head_start:head_end].transpose(0, 1)
308 + k_head = k_req[:, kv_head].T.expand(gqa_group_size, -1, -1)
309 + scores = torch.bmm(q_heads, k_head, out_dtype=torch.float32)
310 + scores = scores.transpose(0, 1) * SM_SCALE
311 + probs = torch.softmax(
312 + scores.masked_fill(~mask[:, None, :], -float("inf")), -1
313 + )
314 + out[q_start:q_end, head_start:head_end] = torch.einsum(
315 + "qhk,kd->qhd", probs, v_req[:, kv_head]
316 + )
317 + q_start += q_len
318 + return out.to(q.dtype)
319 +
320 +
321 +@pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
322 +@pytest.mark.parametrize(
323 + ("q_lens", "kv_lens"),
324 + [
325 + ((129, 257), (129, 257)),
326 + ((65, 129, 257), (129, 257, 385)),
327 + ],
328 +)
329 +def test_prefill_sparse_attention_correctness(
330 + kv_layout: str,
331 + q_lens: tuple[int, ...],
332 + kv_lens: tuple[int, ...],
333 +):
334 + assert len(q_lens) == len(kv_lens)
335 + assert all(kv_len >= q_len for q_len, kv_len in zip(q_lens, kv_lens))
336 +
337 + # Build paged-KV metadata, including a non-identity page order.
338 + batch = len(q_lens)
339 + pages_per_req = [(kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE for kv_len in kv_lens]
340 + max_blocks = max(pages_per_req)
341 + num_pages = sum(pages_per_req)
342 + physical_pages = torch.randperm(num_pages, device="cuda", dtype=torch.int32)
343 + block_table = torch.zeros(batch, max_blocks, device="cuda", dtype=torch.int32)
344 + base_page = 0
345 + for req_id, num_req_pages in enumerate(pages_per_req):
346 + block_table[req_id, :num_req_pages] = physical_pages[
347 + base_page : base_page + num_req_pages
348 + ]
349 + base_page += num_req_pages
350 +
351 + q_lens_t = torch.tensor(q_lens, device="cuda", dtype=torch.int32)
352 + seq_lens = torch.tensor(kv_lens, device="cuda", dtype=torch.int32)
353 + prefix_lens = seq_lens - q_lens_t
354 + cu_seqlens = torch.zeros(batch + 1, device="cuda", dtype=torch.int32)
355 + cu_seqlens[1:] = q_lens_t.cumsum(0)
356 + total_q = sum(q_lens)
357 + max_seqlen_q = max(q_lens)
358 +
359 + q_shape = (total_q, NUM_Q_HEADS, HEAD_DIM)
360 + q = torch.randn(q_shape, device="cuda", dtype=DTYPE)
361 + # Allocate the main KV cache through the backend layout contract so the
362 + # physical storage matches the active layout (contiguous NHD or strided
363 + # HND), while the kernels and reference see the logical-NHD view.
364 + kv_cache = _allocate_main_kv_via_contract(num_pages)
365 +
366 + # Build sparse block indices with the same contract as the real M3 indexer:
367 + # one forced local block, then score-selected older causal blocks.
368 + topk_shape = (NUM_KV_HEADS, total_q, TOPK)
369 + topk_idx = torch.full(topk_shape, -1, device="cuda", dtype=torch.int32)
370 + q_start = 0
371 + for q_len, prefix_len in zip(q_lens_t.tolist(), prefix_lens.tolist()):
372 + for local_q in range(q_len):
373 + current_block = (prefix_len + local_q) // BLOCK_SIZE
374 + older_blocks = torch.randperm(
375 + current_block, device="cuda", dtype=torch.int32
376 + )
377 + selected = torch.cat(
378 + [
379 + torch.tensor([current_block], device="cuda", dtype=torch.int32),
380 + older_blocks[: TOPK - 1],
381 + ]
382 + )
383 + topk_idx[:, q_start + local_q, : selected.numel()] = selected
384 + q_start += q_len
385 +
386 + actual = torch.empty_like(q)
387 + minimax_m3_sparse_attn(
388 + q,
389 + kv_cache,
390 + topk_idx,
391 + block_table,
392 + cu_seqlens,
393 + seq_lens,
394 + prefix_lens,
395 + max_seqlen_q,
396 + NUM_KV_HEADS,
397 + SM_SCALE,
398 + actual,
399 + )
400 +
401 + expected = _reference_sparse_attn(
402 + q,
403 + kv_cache,
404 + topk_idx,
405 + block_table,
406 + q_lens_t,
407 + seq_lens,
408 + prefix_lens,
409 + )
410 + torch.accelerator.synchronize()
411 +
412 + error = (actual.float() - expected.float()).abs()
413 + assert error.mean().item() < 2.5e-4
414 + assert error.max().item() < 1.7e-2
415 +
416 +
417 +def test_main_backend_layout_contract():
418 + """The main sparse backend exposes the logical-NHD shape and the
419 + flash_attn-style stride order for each layout."""
420 + nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
421 + logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
422 + assert logical == (nb, 2, bs, h, d)
423 + # The old HND-ordered shape is no longer the logical shape.
424 + assert logical != (nb, 2, h, bs, d)
425 +
426 + try:
427 + set_kv_cache_layout("HND")
428 + assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 3, 2, 4)
429 + set_kv_cache_layout("NHD")
430 + assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 2, 3, 4)
431 + finally:
432 + set_kv_cache_layout(None)
433 +
434 + for layout in ("NHD", "HND"):
435 + try:
436 + set_kv_cache_layout(layout)
437 + order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
438 + finally:
439 + set_kv_cache_layout(None)
440 + # Valid permutation: no duplicates, covers every axis.
441 + assert set(order) == set(range(len(order)))
442 +
443 + # M3 has no cross-layer KV blocks.
444 + with pytest.raises(NotImplementedError):
445 + MiniMaxM3SparseBackend.get_kv_cache_stride_order(
446 + include_num_layers_dimension=True
447 + )
448 +
449 +
450 +def test_main_backend_unknown_layout_raises(monkeypatch):
451 + """An unrecognized layout (injected past env-var validation) is rejected."""
452 + import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
453 +
454 + monkeypatch.setattr(sparse_attn_mod, "get_kv_cache_layout", lambda: "BOGUS")
455 + with pytest.raises(ValueError, match="Unknown cache layout format"):
456 + MiniMaxM3SparseBackend.get_kv_cache_stride_order()
457 +
458 +
459 +def test_indexer_backend_stride_order_is_identity():
460 + """The 3-dim indexer cache must not inherit the parent's 5-element stride
461 + order; it overrides to the 3-element identity so the allocator keeps the
462 + contiguous layout."""
463 + assert MiniMaxM3IndexerBackend.get_kv_cache_stride_order() == (0, 1, 2)
464 +
465 + # Cross-layer (per-layer-stacked) KV blocks are not supported.
466 + with pytest.raises(NotImplementedError):
467 + MiniMaxM3IndexerBackend.get_kv_cache_stride_order(
468 + include_num_layers_dimension=True
469 + )
470 +
471 + # The stride order matches the 3-dim indexer shape rank.
472 + indexer_shape = MiniMaxM3IndexerBackend.get_kv_cache_shape(
473 + 5, BLOCK_SIZE, 1, HEAD_DIM
474 + )
475 + assert len(indexer_shape) == 3
476 + assert _stride_order_for(MiniMaxM3IndexerBackend, len(indexer_shape)) == (0, 1, 2)
477 +
478 +
479 +def test_hnd_allocation_is_byte_identical_to_transpose():
480 + """Under HND the backend-visible logical view is byte-identical to the
481 + pre-change allocate-HND-then-transpose(2, 3) workaround."""
482 + nb, bs, h, d = 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
483 + logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
484 + try:
485 + set_kv_cache_layout("HND")
486 + stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
487 + finally:
488 + set_kv_cache_layout(None)
489 +
490 + physical_shape = tuple(logical[i] for i in stride_order)
491 + # The physical (permuted) shape equals the old hardcoded HND shape.
492 + assert physical_shape == (nb, 2, h, bs, d)
493 +
494 + inv_order = [stride_order.index(i) for i in range(len(stride_order))]
495 + raw = torch.empty(physical_shape, device="cuda", dtype=DTYPE)
496 + view = raw.permute(*inv_order)
497 + expected = raw.view((nb, 2, h, bs, d)).transpose(2, 3)
498 +
499 + assert view.shape == expected.shape
500 + assert view.stride() == expected.stride()
501 + assert view.storage_offset() == expected.storage_offset()
502 +
503 + # Negative: the identity (wrong) stride order under HND does not reproduce
504 + # the transpose view.
505 + wrong_view = raw.view(logical)
506 + assert wrong_view.stride() != expected.stride()
507 +
508 +
509 +def test_main_cache_is_block_first_and_unpadded():
510 + """The allocator's contiguous-view branch (not the padded-strided branch)
511 + is used for the main GQA cache: its spec is unpadded and the physical
512 + layout keeps num_blocks as the first dimension under both layouts."""
513 + from vllm.v1.kv_cache_interface import FullAttentionSpec
514 +
515 + spec = FullAttentionSpec(
516 + block_size=BLOCK_SIZE,
517 + num_kv_heads=NUM_KV_HEADS,
518 + head_size=HEAD_DIM,
519 + head_size_v=HEAD_DIM,
520 + dtype=DTYPE,
521 + )
522 + # Unpadded -> allocator uses kv_tensor.view(...) rather than as_strided().
523 + assert spec.page_size_padded is None
524 +
525 + logical = MiniMaxM3SparseBackend.get_kv_cache_shape(
526 + 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
527 + )
528 + for layout in ("NHD", "HND"):
529 + try:
530 + set_kv_cache_layout(layout)
531 + order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
532 + finally:
533 + set_kv_cache_layout(None)
534 + inv_order = [order.index(i) for i in range(len(order))]
535 + # Physical first dim is num_blocks (block-first); required by the
536 + # padded-strided branch's block-first assumption if it were ever taken.
537 + assert inv_order[0] == 0
538 + assert logical[order[0]] == logical[0]
539 +
540 +
541 +def _build_decode_inputs(
542 + seq_lens_list: tuple[int, ...],
543 + decode_query_len: int = 1,
544 + num_padded_reqs: int = 0,
545 +):
546 + """Shared decode setup: uniform query tokens per request, a non-identity
547 + block table, and topk indices selecting the current block plus older causal
548 + blocks for each query token."""
549 + active_batch = len(seq_lens_list)
550 + batch = active_batch + num_padded_reqs
551 + pages_per_req = [(s + BLOCK_SIZE - 1) // BLOCK_SIZE for s in seq_lens_list]
552 + max_blocks = max(pages_per_req)
553 + num_pages = sum(pages_per_req)
554 + physical_pages = torch.randperm(num_pages, device="cuda", dtype=torch.int32)
555 + block_table = torch.zeros(batch, max_blocks, device="cuda", dtype=torch.int32)
556 + base_page = 0
557 + for req_id, num_req_pages in enumerate(pages_per_req):
558 + block_table[req_id, :num_req_pages] = physical_pages[
559 + base_page : base_page + num_req_pages
560 + ]
561 + base_page += num_req_pages
562 +
563 + seq_lens = torch.tensor(
564 + (*seq_lens_list, *([0] * num_padded_reqs)),
565 + device="cuda",
566 + dtype=torch.int32,
567 + )
568 + q = torch.randn(
569 + batch * decode_query_len, NUM_Q_HEADS, HEAD_DIM, device="cuda", dtype=DTYPE
570 + )
571 +
572 + topk_idx = torch.full(
573 + (NUM_KV_HEADS, batch * decode_query_len, TOPK),
574 + -1,
575 + device="cuda",
576 + dtype=torch.int32,
577 + )
578 + token_id = 0
579 + for req_id, seq_len in enumerate(seq_lens_list):
580 + for local_q in range(decode_query_len):
581 + query_pos = seq_len - decode_query_len + local_q
582 + current_block = query_pos // BLOCK_SIZE
583 + older_blocks = torch.randperm(
584 + current_block, device="cuda", dtype=torch.int32
585 + )
586 + selected = torch.cat(
587 + [
588 + torch.tensor([current_block], device="cuda", dtype=torch.int32),
589 + older_blocks[: TOPK - 1],
590 + ]
591 + )
592 + topk_idx[:, token_id, : selected.numel()] = selected
593 + token_id += 1
594 +
595 + return q, block_table, seq_lens, topk_idx, num_pages
596 +
597 +
598 +@pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
599 +@pytest.mark.parametrize(
600 + "seq_lens_list",
601 + [(130, 257), (129, 200, 384)],
602 +)
603 +@pytest.mark.parametrize("decode_query_len", [1, 4])
604 +@pytest.mark.parametrize("num_padded_reqs", [0, 2])
605 +def test_decode_sparse_attention_correctness(
606 + kv_layout: str,
607 + seq_lens_list: tuple[int, ...],
608 + decode_query_len: int,
609 + num_padded_reqs: int,
610 +):
611 + """Decode (split-K) parity under both layouts: this is the only coverage of
612 + the decode-site cache feed, and the strided HND case fails if the kernel
613 + ignores the cache strides."""
614 + torch.manual_seed(0)
615 + q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(
616 + seq_lens_list, decode_query_len, num_padded_reqs
617 + )
618 + kv_cache = _allocate_main_kv_via_contract(num_pages)
619 +
620 + actual = torch.empty_like(q)
621 + minimax_m3_sparse_attn_decode(
622 + q,
623 + kv_cache,
624 + topk_idx,
625 + block_table,
626 + seq_lens,
627 + NUM_KV_HEADS,
628 + SM_SCALE,
629 + actual,
630 + decode_query_len,
631 + )
632 +
633 + # Reuse the prefill reference: decode is a uniform query chunk ending at
634 + # seq_len - 1 for each request.
635 + active_batch = len(seq_lens_list)
636 + active_tokens = active_batch * decode_query_len
637 + q_lens_t = torch.full(
638 + (len(seq_lens_list),), decode_query_len, device="cuda", dtype=torch.int32
639 + )
640 + active_seq_lens = seq_lens[:active_batch]
641 + prefix_lens = active_seq_lens - q_lens_t
642 + expected = _reference_sparse_attn(
643 + q[:active_tokens],
644 + kv_cache,
645 + topk_idx[:, :active_tokens],
646 + block_table[:active_batch],
647 + q_lens_t,
648 + active_seq_lens,
649 + prefix_lens,
650 + )
651 + torch.accelerator.synchronize()
652 +
653 + error = (actual[:active_tokens].float() - expected.float()).abs()
654 + assert error.mean().item() < 2.5e-4
655 + assert error.max().item() < 1.7e-2
656 +
657 +
658 +def test_decode_wrong_layout_breaks_parity():
659 + """Negative (AC-3/AC-5): consuming the physical HND buffer as if it were
660 + already contiguous-NHD (i.e. skipping the allocator's inverse permute)
661 + reorders the K/V content, so the decode output no longer matches the
662 + reference computed on the correct logical view. The mislabeled tensor keeps
663 + the same shape as the correct view, so the kernel stays in bounds."""
664 + torch.manual_seed(0)
665 + seq_lens_list = (130, 257)
666 + q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(seq_lens_list)
667 +
668 + # Physical HND storage [blocks, 2, heads, block, dim].
669 + phys = torch.randn(
670 + (num_pages, 2, NUM_KV_HEADS, BLOCK_SIZE, HEAD_DIM), device="cuda", dtype=DTYPE
671 + )
672 + # Correct logical-NHD view (strided) vs. the same bytes mislabeled as a
673 + # contiguous-NHD cache — same shape, different content mapping.
674 + correct = phys.permute(0, 1, 3, 2, 4)
675 + wrong = phys.reshape(num_pages, 2, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM)
676 +
677 + q_lens_t = torch.ones(len(seq_lens_list), device="cuda", dtype=torch.int32)
678 + prefix_lens = seq_lens - q_lens_t
679 + expected = _reference_sparse_attn(
680 + q, correct, topk_idx, block_table, q_lens_t, seq_lens, prefix_lens
681 + )
682 +
683 + actual = torch.empty_like(q)
684 + minimax_m3_sparse_attn_decode(
685 + q, wrong, topk_idx, block_table, seq_lens, NUM_KV_HEADS, SM_SCALE, actual, 1
686 + )
687 + torch.accelerator.synchronize()
688 + assert (actual.float() - expected.float()).abs().max().item() > 1.7e-2
689 +
690 +
691 +def _make_attn_group(backend, spec):
692 + return AttentionGroup(
693 + backend=backend,
694 + layer_names=["main"],
695 + kv_cache_spec=spec,
696 + kv_cache_group_id=0,
697 + )
698 +
699 +
700 +def test_main_cache_byte_identical_through_production_allocator():
701 + """AC-2: drive the real allocator (`_reshape_kv_cache`) for the M3 main
702 + `FullAttentionSpec` under HND and assert the backend-visible view has the
703 + same shape, stride, and storage offset as the pre-change
704 + allocate-HND-then-transpose path; the indexer `MLAAttentionSpec` allocates
705 + through the same path to its 3-dim shape."""
706 + nb = 4
707 + spec = FullAttentionSpec(
708 + block_size=BLOCK_SIZE,
709 + num_kv_heads=NUM_KV_HEADS,
710 + head_size=HEAD_DIM,
711 + head_size_v=HEAD_DIM,
712 + dtype=DTYPE,
713 + )
714 + raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
715 + group = _make_attn_group(MiniMaxM3SparseBackend, spec)
716 + try:
717 + set_kv_cache_layout("HND")
718 + kv_caches = _reshape_kv_cache([group], {"main": raw}, "auto", [BLOCK_SIZE], {})
719 + finally:
720 + set_kv_cache_layout(None)
721 + view = kv_caches["main"]
722 +
723 + oracle = raw.view(DTYPE).view((nb, 2, NUM_KV_HEADS, BLOCK_SIZE, HEAD_DIM))
724 + oracle = oracle.transpose(2, 3)
725 + assert tuple(view.shape) == tuple(oracle.shape)
726 + assert view.stride() == oracle.stride()
727 + assert view.storage_offset() == oracle.storage_offset()
728 +
729 + # Indexer cache allocates through the same path under both layouts.
730 + ispec = MLAAttentionSpec(
731 + block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
732 + )
733 + for layout in ("NHD", "HND"):
734 + iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
735 + igroup = AttentionGroup(
736 + backend=MiniMaxM3IndexerBackend,
737 + layer_names=["idx"],
738 + kv_cache_spec=ispec,
739 + kv_cache_group_id=0,
740 + )
741 + try:
742 + set_kv_cache_layout(layout)
743 + iout = _reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {})
744 + finally:
745 + set_kv_cache_layout(None)
746 + assert tuple(iout["idx"].shape) == (nb, BLOCK_SIZE, HEAD_DIM)
747 +
748 +
749 +def test_indexer_inherited_stride_order_trips_allocator_assert():
750 + """AC-4 negative: without the indexer override, the inherited 5-element
751 + stride order trips the allocator's `len(stride_order) == len(shape)` assert
752 + for the 3-dim indexer shape; the `AssertionError` is NOT swallowed by the
753 + allocator's `(AttributeError, NotImplementedError)` fallback."""
754 +
755 + class _BrokenIndexerBackend(MiniMaxM3IndexerBackend):
756 + # Simulate inheriting the parent's 5-element stride order.
757 + get_kv_cache_stride_order = staticmethod(
758 + MiniMaxM3SparseBackend.get_kv_cache_stride_order
759 + )
760 +
761 + nb = 4
762 + ispec = MLAAttentionSpec(
763 + block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
764 + )
765 + iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
766 + igroup = AttentionGroup(
767 + backend=_BrokenIndexerBackend,
768 + layer_names=["idx"],
769 + kv_cache_spec=ispec,
770 + kv_cache_group_id=0,
771 + )
772 + try:
773 + set_kv_cache_layout("HND")
774 + with pytest.raises(AssertionError):
775 + _reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {})
776 + finally:
777 + set_kv_cache_layout(None)
778 +
779 +
780 +def test_padded_main_cache_is_flagged():
781 + """AC-2.1 negative: the M3 main cache relies on the allocator's
782 + contiguous-view branch (`page_size_padded is None`). A spec that sets
783 + `page_size_padded` is explicitly flagged rather than silently wrong-strided."""
784 +
785 + def _require_unpadded_block_first(spec, stride_order):
786 + inv_order = [stride_order.index(i) for i in range(len(stride_order))]
787 + assert spec.page_size_padded is None, (
788 + "main GQA cache must be unpadded to use the contiguous-view "
789 + "allocator branch"
790 + )
791 + assert inv_order[0] == 0, "main GQA cache must remain block-first"
792 +
793 + try:
794 + set_kv_cache_layout("HND")
795 + stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
796 + finally:
797 + set_kv_cache_layout(None)
798 +
799 + good = FullAttentionSpec(
800 + block_size=BLOCK_SIZE,
801 + num_kv_heads=NUM_KV_HEADS,
802 + head_size=HEAD_DIM,
803 + head_size_v=HEAD_DIM,
804 + dtype=DTYPE,
805 + )
806 + _require_unpadded_block_first(good, stride_order) # passes
807 +
808 + padded = FullAttentionSpec(
809 + block_size=BLOCK_SIZE,
810 + num_kv_heads=NUM_KV_HEADS,
811 + head_size=HEAD_DIM,
812 + head_size_v=HEAD_DIM,
813 + dtype=DTYPE,
814 + page_size_padded=good.page_size_bytes + 128,
815 + )
816 + with pytest.raises(AssertionError):
817 + _require_unpadded_block_first(padded, stride_order)
818 +
819 +
820 +@pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
821 +def test_reshape_and_cache_flash_write_persists(kv_layout: str):
822 + """AC-5 write path: the `reshape_and_cache_flash` write site now consumes
823 + `self.kv_cache.unbind(1)` directly. Writing through those views must persist
824 + into the bound storage (read back through an independent logical view) under
825 + both layouts — a `.contiguous()` copy of the unbind slice would leave the
826 + bound storage unchanged."""
827 + torch.manual_seed(0)
828 + num_pages = 4
829 + kv_cache = _allocate_main_kv_via_contract(num_pages)
830 + with torch.no_grad():
831 + kv_cache.zero_()
832 +
833 + # Exactly the production write-site code under test.
834 + key_cache, value_cache = kv_cache.unbind(1)
835 +
836 + num_tokens = 12
837 + slot_mapping = torch.randperm(num_pages * BLOCK_SIZE, device="cuda")[
838 + :num_tokens
839 + ].to(torch.int64)
840 + key = torch.randn(num_tokens, NUM_KV_HEADS, HEAD_DIM, device="cuda", dtype=DTYPE)
841 + value = torch.randn(num_tokens, NUM_KV_HEADS, HEAD_DIM, device="cuda", dtype=DTYPE)
842 + scale = torch.ones((), device="cuda")
843 + ops.reshape_and_cache_flash(
844 + key, value, key_cache, value_cache, slot_mapping, "auto", scale, scale
845 + )
846 + torch.accelerator.synchronize()
847 +
848 + # Read back through the independent logical view; proves the writes landed
849 + # in the engine-bound storage, not a detached copy.
850 + for t in range(num_tokens):
851 + slot = int(slot_mapping[t].item())
852 + blk, intra = divmod(slot, BLOCK_SIZE)
853 + torch.testing.assert_close(kv_cache[blk, 0, intra], key[t])
854 + torch.testing.assert_close(kv_cache[blk, 1, intra], value[t])
tests/kernels/core/test_fused_allreduce_gemma_rms_norm.py new
+109
@@ -0,0 +1,109 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the manual AllReduce + GemmaRMSNorm fusion used by MiniMax M3.
4 +
5 +``fused_allreduce_gemma_rms_norm`` must match the unfused model path, i.e.
6 +``GemmaRMSNorm(all_reduce(partial), residual)``, both on the flashinfer fast
7 +path (TP>1 with flashinfer + NVSwitch) and on the eager fallback (TP==1, or when
8 +flashinfer is unavailable / the GPU has no NVSwitch).
9 +"""
10 +
11 +import pytest
12 +import torch
13 +from torch.multiprocessing import spawn
14 +
15 +from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
16 +from vllm.distributed import cleanup_dist_env_and_memory
17 +from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
18 +from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
19 + fused_allreduce_gemma_rms_norm,
20 +)
21 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm
22 +from vllm.platforms import current_platform
23 +from vllm.utils.network_utils import get_open_port
24 +from vllm.utils.torch_utils import set_random_seed
25 +
26 +
27 +@ensure_current_vllm_config()
28 +def _worker_fused_ar_norm(
29 + local_rank,
30 + world_size,
31 + port,
32 + num_tokens,
33 + hidden_size,
34 + dtype,
35 + seed,
36 + eps,
37 +):
38 + """Per-rank worker: compare the fused helper vs all_reduce + GemmaRMSNorm."""
39 + device = torch.device(f"cuda:{local_rank}")
40 + torch.accelerator.set_device_index(device)
41 + init_test_distributed_environment(
42 + world_size, 1, local_rank, port, local_rank=local_rank
43 + )
44 +
45 + # Norm weights are identical across ranks (replicated GemmaRMSNorm).
46 + set_random_seed(seed)
47 + norm = GemmaRMSNorm(hidden_size, eps=eps).cuda().to(dtype)
48 + with torch.no_grad():
49 + norm.weight.normal_(mean=0.0, std=0.1)
50 +
51 + # Residual is shared across ranks; the partial o_proj output differs per rank
52 + # (each rank holds a partial sum that all_reduce combines).
53 + torch.manual_seed(seed + 7)
54 + residual = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
55 + torch.manual_seed(seed + 1000 + local_rank)
56 + partial = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
57 +
58 + # Reference: the unfused model path.
59 + reduced = tensor_model_parallel_all_reduce(partial.clone())
60 + ref_out, ref_res = norm(reduced, residual.clone())
61 +
62 + # Fused helper (flashinfer fast path when available, else fallback).
63 + out, res = fused_allreduce_gemma_rms_norm(partial.clone(), residual.clone(), norm)
64 + torch.accelerator.synchronize()
65 +
66 + torch.testing.assert_close(out, ref_out, atol=2e-2, rtol=2e-2)
67 + torch.testing.assert_close(res, ref_res, atol=2e-2, rtol=2e-2)
68 +
69 + cleanup_dist_env_and_memory()
70 +
71 +
72 +@pytest.mark.skipif(
73 + not current_platform.is_cuda(),
74 + reason="CUDA required",
75 +)
76 +# world_size=1 exercises the TP==1 identity branch on a single GPU; >1 exercises
77 +# the all_reduce + GemmaRMSNorm equivalence (flashinfer kernel or fallback).
78 +@pytest.mark.parametrize("world_size", [1, 2, 4])
79 +@pytest.mark.parametrize("num_tokens", [1, 128, 333])
80 +@pytest.mark.parametrize("hidden_size", [2048, 4096])
81 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
82 +@pytest.mark.parametrize("eps", [1e-6])
83 +@pytest.mark.parametrize("seed", [42])
84 +def test_fused_allreduce_gemma_rms_norm(
85 + world_size,
86 + num_tokens,
87 + hidden_size,
88 + dtype,
89 + eps,
90 + seed,
91 +):
92 + num_gpus = current_platform.device_count()
93 + if num_gpus < world_size:
94 + pytest.skip(f"Need >= {world_size} GPUs, have {num_gpus}")
95 + port = str(get_open_port())
96 + spawn(
97 + _worker_fused_ar_norm,
98 + args=(
99 + world_size,
100 + port,
101 + num_tokens,
102 + hidden_size,
103 + dtype,
104 + seed,
105 + eps,
106 + ),
107 + nprocs=world_size,
108 + join=True,
109 + )
tests/kernels/test_fp32_router_gemm.py
+21 -15
@@ -1,6 +1,9 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Tests for fp32_router_gemm kernel: activation×weight→fp32, H=3072, E=256.
3 +"""Tests for fp32_router_gemm kernel: activation×weight→fp32.
4 +
5 +Supported (hidden_size, num_experts) pairs:
6 + (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
7
8 Correctness baseline: torch.matmul in float64.
9 """
@@ -10,8 +13,8 @@ import torch
13
14 from vllm._custom_ops import fp32_router_gemm
15
13 -NUM_EXPERTS = 256
14 -HIDDEN_DIM = 3072
16 +# (hidden_size, num_experts)
17 +SHAPES = [(3072, 256), (6144, 128)]
18 # Absolute tolerance for fp32 kernel vs float64 reference
19 ATOL_FP32 = 2e-4
20 ATOL_BF16 = 2e-2 # bf16 activation has lower precision
@@ -30,49 +33,52 @@ def _ref(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
33 return torch.nn.functional.linear(mat_a.float(), mat_b.float())
34
35
36 +@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
37 @pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
34 -def test_fp32_activation(num_tokens: int):
38 +def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int):
39 """fp32 activation → fp32 output should match reference closely."""
40 _requires_sm90()
41 torch.manual_seed(42)
42 device = torch.device("cuda")
39 - mat_a = torch.randn(num_tokens, HIDDEN_DIM, dtype=torch.float32, device=device)
40 - mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device)
43 + mat_a = torch.randn(num_tokens, hidden_dim, dtype=torch.float32, device=device)
44 + mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
45
46 out = fp32_router_gemm(mat_a, mat_b)
47 ref = _ref(mat_a, mat_b)
48
45 - assert out.shape == (num_tokens, NUM_EXPERTS)
49 + assert out.shape == (num_tokens, num_experts)
50 assert out.dtype == torch.float32
51 torch.testing.assert_close(out, ref, atol=ATOL_FP32, rtol=0)
52
53
54 +@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
55 @pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
51 -def test_bf16_activation(num_tokens: int):
56 +def test_bf16_activation(num_tokens: int, hidden_dim: int, num_experts: int):
57 """bf16 activation → fp32 output should match reference within bf16 error."""
58 _requires_sm90()
59 torch.manual_seed(42)
60 device = torch.device("cuda")
61 mat_a_bf16 = torch.randn(
57 - num_tokens, HIDDEN_DIM, dtype=torch.bfloat16, device=device
62 + num_tokens, hidden_dim, dtype=torch.bfloat16, device=device
63 )
59 - mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device)
64 + mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
65
66 out = fp32_router_gemm(mat_a_bf16, mat_b)
67 ref = _ref(mat_a_bf16, mat_b).to(device)
68
64 - assert out.shape == (num_tokens, NUM_EXPERTS)
69 + assert out.shape == (num_tokens, num_experts)
70 assert out.dtype == torch.float32
71 torch.testing.assert_close(out, ref, atol=ATOL_BF16, rtol=0)
72
73
69 -def test_output_shape_and_dtype():
74 +@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES)
75 +def test_output_shape_and_dtype(hidden_dim: int, num_experts: int):
76 """Basic shape and dtype checks."""
77 _requires_sm90()
78 device = torch.device("cuda")
73 - mat_a = torch.randn(4, HIDDEN_DIM, dtype=torch.float32, device=device)
74 - mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device)
79 + mat_a = torch.randn(4, hidden_dim, dtype=torch.float32, device=device)
80 + mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device)
81 out = fp32_router_gemm(mat_a, mat_b)
76 - assert out.shape == (4, NUM_EXPERTS)
82 + assert out.shape == (4, num_experts)
83 assert out.dtype == torch.float32
84 assert out.device.type == "cuda"
tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py new
+244
@@ -0,0 +1,244 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit test for the horizontally-fused MiniMax-M3 attention pre-processing
4 +kernel:
5 +
6 + fused_minimax_m3_qknorm_rope_kv_insert
7 + - q / k / index_q / index_k: Gemma RMSNorm + partial NeoX RoPE (in place)
8 + - sparse (insert) mode: scatter k/v into the paged bf16 KV cache and the
9 + index key into the index cache by its own slot mapping.
10 +
11 +Reference: PyTorch Gemma RMSNorm with the same dtype materialization boundary
12 +as the unfused path, followed by vLLM CUDA rotary_embedding-style NeoX RoPE.
13 +"""
14 +
15 +import pytest
16 +import torch
17 +
18 +import vllm._custom_ops as ops
19 +
20 +HEAD_DIM = 128
21 +ROTARY_DIM = 64
22 +
23 +
24 +def _op_available() -> bool:
25 + return hasattr(torch.ops._C, "fused_minimax_m3_qknorm_rope_kv_insert")
26 +
27 +
28 +pytestmark = pytest.mark.skipif(
29 + not torch.cuda.is_available() or not _op_available(),
30 + reason="CUDA not available or fused MiniMax-M3 op not built in",
31 +)
32 +
33 +
34 +def make_cos_sin_cache(max_pos, rotary_dim, base, dtype, device):
35 + inv_freq = 1.0 / (
36 + base
37 + ** (
38 + torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device)
39 + / rotary_dim
40 + )
41 + )
42 + t = torch.arange(max_pos, dtype=torch.float32, device=device)
43 + freqs = torch.einsum("i,j->ij", t, inv_freq) # [max_pos, rotary_dim/2]
44 + cache = torch.cat((freqs.cos(), freqs.sin()), dim=-1) # [max_pos, rotary_dim]
45 + return cache.to(dtype)
46 +
47 +
48 +def gemma_rmsnorm(x, weight, eps):
49 + """x: [..., 128]; weight: [128]. Returns original dtype."""
50 + xf = x.float()
51 + var = xf.pow(2).mean(dim=-1, keepdim=True)
52 + out = xf * torch.rsqrt(var + eps)
53 + out = out * (1.0 + weight.float())
54 + return out.to(x.dtype)
55 +
56 +
57 +def apply_rope_neox_partial(x, positions, cos_sin_cache, rotary_dim):
58 + """NeoX-style RoPE on the leading rotary_dim dims; rest pass through.
59 +
60 + x: [num_tokens, num_heads, head_dim]
61 + cos_sin_cache: [max_pos, rotary_dim] (cos||sin), read as float (matches the
62 + kernel, which loads the bf16 cache and converts to fp32).
63 + """
64 + half = rotary_dim // 2
65 + cs = cos_sin_cache[positions].float() # [num_tokens, rotary_dim]
66 + cos = cs[..., :half].unsqueeze(1) # [nt, 1, half]
67 + sin = cs[..., half:].unsqueeze(1)
68 +
69 + rot = x[..., :rotary_dim].float()
70 + x1 = rot[..., :half]
71 + x2 = rot[..., half:]
72 + o1 = x1 * cos - x2 * sin
73 + o2 = x2 * cos + x1 * sin
74 + out = x.clone()
75 + out[..., :half] = o1
76 + out[..., half:rotary_dim] = o2
77 + return out.to(x.dtype)
78 +
79 +
80 +def norm_rope_ref(x, weight, positions, cos_sin_cache, eps):
81 + """[nt, nheads, 128] -> Gemma norm + neox partial rope."""
82 + normed = gemma_rmsnorm(x, weight, eps)
83 + roped = apply_rope_neox_partial(normed, positions, cos_sin_cache, ROTARY_DIM)
84 + return roped
85 +
86 +
87 +# ── Test 1: dense mode (norm+rope only, no index, no insert) ─────────────────
88 +
89 +
90 +@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
91 +@pytest.mark.parametrize("num_heads,num_kv_heads", [(8, 2), (16, 4), (64, 4)])
92 +def test_dense_norm_rope(num_tokens, num_heads, num_kv_heads):
93 + torch.manual_seed(0)
94 + device, dtype, eps = "cuda", torch.bfloat16, 1e-6
95 + base, max_pos = 5_000_000.0, 4096
96 +
97 + q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
98 + k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
99 + cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
100 + positions = torch.randint(
101 + 0, max_pos, (num_tokens,), dtype=torch.int64, device=device
102 + )
103 +
104 + qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
105 + qkv = torch.randn(num_tokens, qsz + 2 * kvsz, dtype=dtype, device=device)
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
110 + )
111 + q_out, k_out, v_out = qkv.split([qsz, kvsz, kvsz], dim=-1)
112 +
113 + q_in, k_in, v_in = qkv_orig.split([qsz, kvsz, kvsz], dim=-1)
114 + q_ref = norm_rope_ref(
115 + q_in.view(num_tokens, num_heads, HEAD_DIM), q_w, positions, cos_sin, eps
116 + ).view(num_tokens, qsz)
117 + k_ref = norm_rope_ref(
118 + k_in.view(num_tokens, num_kv_heads, HEAD_DIM),
119 + k_w,
120 + positions,
121 + cos_sin,
122 + eps,
123 + ).view(num_tokens, kvsz)
124 +
125 + torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2)
126 + torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2)
127 + # V is untouched.
128 + torch.testing.assert_close(v_out, v_in, rtol=0, atol=0)
129 +
130 +
131 +# ── Test 2: sparse mode (full: index branch + cache inserts) ─────────────────
132 +
133 +
134 +@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
135 +@pytest.mark.parametrize("block_size", [16, 64])
136 +def test_sparse_full(num_tokens, block_size):
137 + torch.manual_seed(1)
138 + device, dtype, eps = "cuda", torch.bfloat16, 1e-6
139 + base, max_pos = 5_000_000.0, 4096
140 + num_heads, num_kv_heads, num_idx_heads = 16, 4, 4
141 +
142 + q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
143 + k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
144 + iq_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
145 + ik_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
146 + cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
147 + positions = torch.randint(
148 + 0, max_pos, (num_tokens,), dtype=torch.int64, device=device
149 + )
150 +
151 + qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
152 + iqsz, iksz = num_idx_heads * HEAD_DIM, HEAD_DIM
153 + # Single fused tensor packing [q | k | v | index_q | index_k].
154 + qkv = torch.randn(
155 + num_tokens, qsz + 2 * kvsz + iqsz + iksz, dtype=dtype, device=device
156 + )
157 + qkv_orig = qkv.clone()
158 + splits = [qsz, kvsz, kvsz, iqsz, iksz]
159 +
160 + num_blocks = (num_tokens + block_size - 1) // block_size + 1
161 + kv_cache = torch.zeros(
162 + num_blocks, 2, block_size, num_kv_heads, HEAD_DIM, dtype=dtype, device=device
163 + )
164 + index_cache = torch.zeros(
165 + num_blocks, block_size, HEAD_DIM, dtype=dtype, device=device
166 + )
167 + slot_mapping = torch.randperm(
168 + num_blocks * block_size, dtype=torch.int64, device=device
169 + )[:num_tokens]
170 + index_slot_mapping = torch.roll(slot_mapping, shifts=1)
171 +
172 + # Contiguous gather targets: the kernel writes the normed/roped q and
173 + # index_q here (de-interleaved from the packed qkv); k/v/index_k stay in
174 + # place inside qkv and are scatter-inserted into the caches.
175 + q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device)
176 + index_q = torch.empty(num_tokens, iqsz, dtype=dtype, device=device)
177 +
178 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
179 + qkv,
180 + q_w,
181 + k_w,
182 + cos_sin,
183 + positions,
184 + num_heads,
185 + num_kv_heads,
186 + ROTARY_DIM,
187 + eps,
188 + iq_w,
189 + ik_w,
190 + num_idx_heads,
191 + slot_mapping,
192 + index_slot_mapping,
193 + kv_cache,
194 + index_cache,
195 + block_size,
196 + q_out,
197 + index_q,
198 + )
199 +
200 + # ── norm+rope parity. q/index_q land in their gather buffers; k/index_k are
201 + # rewritten in place inside qkv. ──
202 + _, k_out, _, _, index_k = qkv.split(splits, dim=-1)
203 + q_in, k_in, v_in, iq_orig, ik_orig = qkv_orig.split(splits, dim=-1)
204 + q_ref = norm_rope_ref(
205 + q_in.view(num_tokens, num_heads, HEAD_DIM), q_w, positions, cos_sin, eps
206 + ).view(num_tokens, qsz)
207 + k_ref = norm_rope_ref(
208 + k_in.view(num_tokens, num_kv_heads, HEAD_DIM),
209 + k_w,
210 + positions,
211 + cos_sin,
212 + eps,
213 + ).view(num_tokens, kvsz)
214 + iq_ref = norm_rope_ref(
215 + iq_orig.view(num_tokens, num_idx_heads, HEAD_DIM),
216 + iq_w,
217 + positions,
218 + cos_sin,
219 + eps,
220 + ).view(num_tokens, num_idx_heads * HEAD_DIM)
221 + ik_ref = norm_rope_ref(
222 + ik_orig.view(num_tokens, 1, HEAD_DIM), ik_w, positions, cos_sin, eps
223 + ).view(num_tokens, HEAD_DIM)
224 +
225 + torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2)
226 + torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2)
227 + torch.testing.assert_close(index_q, iq_ref, rtol=1e-2, atol=1e-2)
228 + torch.testing.assert_close(index_k, ik_ref, rtol=1e-2, atol=1e-2)
229 +
230 + # ── Cache inserts. ──
231 + # Main cache layout is [num_blocks, 2, block_size, num_kv_heads, head_dim]
232 + # (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)
234 + k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
235 + 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
241 + )
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)
tests/kernels/test_minimax_m3_amd_ops.py new
+337
@@ -0,0 +1,337 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Reference-vs-optimized unit tests for the MiniMax-M3 AMD/ROCm fused kernels.
4 +
5 +Each optimized kernel added for the ROCm port has a slow PyTorch reference; the
6 +tests assert the two agree within tolerance:
7 +
8 + * Gemma RMSNorm (plain + fused-add-residual) -> fp32 PyTorch normalize
9 + * SwiGLU-OAI (split layout) -> fp32 PyTorch elementwise
10 + * Fused MXFP8 activation quant (Triton) -> _mxfp8_e4m3_quantize_torch
11 + * Native MXFP8 linear (dot_scaled) -> dequant-to-bf16 @ matmul
12 + * Native MXFP8 MoE (dot_scaled grouped GEMM) -> dequant-to-bf16 MoE math
13 +
14 +The native MXFP8 GEMMs also guard the ``dot_scaled`` rhs-scale orientation: the
15 +scale is loaded ``[N, K//32]`` and passed WITHOUT transpose; a stray ``.T``
16 +makes the shape ``[K//32, N]`` and Triton raises before producing output, so any
17 +regression there fails these tests loudly.
18 +
19 +Hardware scope: the whole module is ROCm-only (these are the AMD path; NVIDIA
20 +uses the FlashInfer kernels). The norm/activation/quant kernels run on any ROCm
21 +arch; the native MXFP8 ``dot_scaled`` linear/MoE tests are additionally gated to
22 +CDNA4 gfx95x (``@requires_gfx950``) since gfx942 uses the BF16 emulation path.
23 +
24 +Run: pytest tests/kernels/test_minimax_m3_amd_ops.py -v
25 +"""
26 +
27 +import pytest
28 +import torch
29 +
30 +from vllm.platforms import current_platform
31 +
32 +if not current_platform.is_rocm():
33 + pytest.skip("MiniMax-M3 AMD fused ops require ROCm.", allow_module_level=True)
34 +if not torch.cuda.is_available():
35 + pytest.skip("Requires a GPU.", allow_module_level=True)
36 +
37 +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( # noqa: E402
38 + _mxfp8_e4m3_quantize_torch,
39 + _mxfp8_e4m3_quantize_triton,
40 + dequant_mxfp8_to_bf16,
41 +)
42 +from vllm.models.minimax_m3.amd.ops import ( # noqa: E402
43 + gemma_fused_add_rmsnorm,
44 + gemma_rmsnorm,
45 + swiglu_oai_split,
46 +)
47 +from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import _num_warps # noqa: E402
48 +
49 +DEVICE = "cuda"
50 +EPS = 1e-6
51 +
52 +
53 +def _gcn_arch() -> str:
54 + try:
55 + return torch.cuda.get_device_properties(0).gcnArchName
56 + except Exception: # pragma: no cover - no device / non-AMD
57 + return ""
58 +
59 +
60 +# The pure-Triton norm/activation/quant kernels run on any ROCm arch (CDNA3
61 +# gfx942 and CDNA4 gfx950). The native MXFP8 ``dot_scaled`` GEMMs (linear + MoE)
62 +# use CDNA4 hardware microscaling and are gated to gfx95x in the source
63 +# (``RocmDotScaledMxfp8LinearKernel.is_supported``; the MoE oracle routes gfx942
64 +# to the BF16 emulation path instead) — so those tests are gfx950-only.
65 +requires_gfx950 = pytest.mark.skipif(
66 + "gfx95" not in _gcn_arch(),
67 + reason="native MXFP8 dot_scaled is a CDNA4 (gfx95x) feature; "
68 + "gfx942 uses the BF16 emulation path instead.",
69 +)
70 +
71 +
72 +def _relerr(a: torch.Tensor, b: torch.Tensor) -> float:
73 + a = a.float()
74 + b = b.float()
75 + return ((a - b).norm() / (b.norm() + 1e-8)).item()
76 +
77 +
78 +# --------------------------------------------------------------------------- #
79 +# Gemma RMSNorm
80 +# --------------------------------------------------------------------------- #
81 +def _ref_gemma_rmsnorm(x, w, eps, residual=None):
82 + orig_dtype = x.dtype
83 + xf = x.float()
84 + res_out = None
85 + if residual is not None:
86 + xf = xf + residual.float()
87 + res_out = xf.to(orig_dtype)
88 + xf = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps)
89 + xf = xf * (1.0 + w.float())
90 + out = xf.to(orig_dtype)
91 + return out if residual is None else (out, res_out)
92 +
93 +
94 +@pytest.mark.parametrize("shape", [(1, 4096), (37, 6144), (128, 2048)])
95 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
96 +@pytest.mark.parametrize("seed", [0, 1234])
97 +@torch.inference_mode()
98 +def test_gemma_rmsnorm(shape, dtype, seed):
99 + torch.manual_seed(seed)
100 + x = torch.randn(*shape, device=DEVICE, dtype=dtype)
101 + w = torch.randn(shape[-1], device=DEVICE, dtype=dtype) * 0.1
102 + got = gemma_rmsnorm(x, w, EPS)
103 + ref = _ref_gemma_rmsnorm(x, w, EPS)
104 + assert got.shape == x.shape
105 + assert _relerr(got, ref) < 5e-3
106 +
107 +
108 +@pytest.mark.parametrize("shape", [(1, 6144), (64, 4096)])
109 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
110 +@torch.inference_mode()
111 +def test_gemma_fused_add_rmsnorm(shape, dtype):
112 + torch.manual_seed(0)
113 + x = torch.randn(*shape, device=DEVICE, dtype=dtype)
114 + res = torch.randn(*shape, device=DEVICE, dtype=dtype)
115 + w = torch.randn(shape[-1], device=DEVICE, dtype=dtype) * 0.1
116 + got_out, got_res = gemma_fused_add_rmsnorm(x, res, w, EPS)
117 + ref_out, ref_res = _ref_gemma_rmsnorm(x, w, EPS, residual=res)
118 + assert _relerr(got_out, ref_out) < 5e-3
119 + # residual_out is the pre-norm sum (x + res): bit-for-bit identical cast.
120 + assert torch.equal(got_res, ref_res)
121 +
122 +
123 +@torch.inference_mode()
124 +def test_gemma_rmsnorm_per_head_strided():
125 + """q_norm/k_norm normalize a non-contiguous ``qkv.split`` slice over head_dim."""
126 + torch.manual_seed(0)
127 + T, H, D, kv = 7, 48, 128, 8
128 + total = (H + 2 * kv) * D
129 + qkv = torch.randn(T, total, device=DEVICE, dtype=torch.bfloat16)
130 + q = qkv[..., : H * D] # non-contiguous view (row stride == total)
131 + q_by_head = q.view(T, H, D)
132 + assert not q_by_head.is_contiguous()
133 + w = torch.randn(D, device=DEVICE, dtype=torch.bfloat16) * 0.1
134 + got = gemma_rmsnorm(q_by_head, w, EPS)
135 + ref = _ref_gemma_rmsnorm(q_by_head, w, EPS)
136 + assert got.shape == q_by_head.shape
137 + assert _relerr(got, ref) < 5e-3
138 +
139 +
140 +def test_num_warps_monotonic():
141 + assert _num_warps(128) <= _num_warps(2048) <= _num_warps(8192)
142 +
143 +
144 +# --------------------------------------------------------------------------- #
145 +# SwiGLU-OAI (split layout)
146 +# --------------------------------------------------------------------------- #
147 +def _ref_swiglu(gate_up, alpha, beta, limit):
148 + d = gate_up.shape[-1] // 2
149 + gate = gate_up[..., :d].float()
150 + up = gate_up[..., d:].float()
151 + if limit is not None:
152 + gate = gate.clamp(max=limit)
153 + up = up.clamp(min=-limit, max=limit)
154 + return (gate * torch.sigmoid(alpha * gate) * (up + beta)).to(gate_up.dtype)
155 +
156 +
157 +@pytest.mark.parametrize("m,inter", [(1, 768), (64, 1536), (128, 1024)])
158 +@pytest.mark.parametrize("limit", [7.0, None])
159 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
160 +@torch.inference_mode()
161 +def test_swiglu_oai_split(m, inter, limit, dtype):
162 + torch.manual_seed(0)
163 + gate_up = torch.randn(m, 2 * inter, device=DEVICE, dtype=dtype)
164 + got = swiglu_oai_split(gate_up, alpha=1.702, beta=1.0, limit=limit)
165 + ref = _ref_swiglu(gate_up, 1.702, 1.0, limit)
166 + assert got.shape == (m, inter)
167 + assert _relerr(got, ref) < 5e-3
168 +
169 +
170 +# --------------------------------------------------------------------------- #
171 +# Fused MXFP8 activation quant (Triton vs torch reference)
172 +# --------------------------------------------------------------------------- #
173 +@pytest.mark.parametrize("shape", [(64, 4096), (1, 6144), (333, 2048)])
174 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
175 +@torch.inference_mode()
176 +def test_mxfp8_quant_triton_matches_torch(shape, dtype):
177 + torch.manual_seed(0)
178 + x = torch.randn(*shape, device=DEVICE, dtype=dtype)
179 + xq_t, s_t = _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout=False)
180 + xq_k, s_k = _mxfp8_e4m3_quantize_triton(x)
181 + assert s_k.shape == s_t.shape == (shape[0], shape[1] // 32)
182 + # E8M0 block exponents share the floor(log2(amax))+127 algorithm; allow at
183 + # most a 1-step difference at exact powers of two.
184 + assert (s_k.int() - s_t.int()).abs().max().item() <= 1
185 + # Dequantized values agree to fp8 granularity.
186 + deq_t = dequant_mxfp8_to_bf16(xq_t, s_t)
187 + deq_k = dequant_mxfp8_to_bf16(xq_k, s_k)
188 + assert _relerr(deq_k, deq_t) < 1e-2
189 +
190 +
191 +# --------------------------------------------------------------------------- #
192 +# Native MXFP8 linear (dot_scaled) vs dequant-to-bf16 matmul
193 +# --------------------------------------------------------------------------- #
194 +@requires_gfx950
195 +@pytest.mark.parametrize("m,n,k", [(64, 256, 128), (37, 512, 256), (1, 6144, 4096)])
196 +@torch.inference_mode()
197 +def test_mxfp8_native_linear(m, n, k):
198 + from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
199 + _mxfp8_dot_scaled_linear,
200 + )
201 +
202 + torch.manual_seed(0)
203 + w_bf16 = torch.randn(n, k, device=DEVICE, dtype=torch.bfloat16) * 0.1
204 + w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
205 + x = torch.randn(m, k, device=DEVICE, dtype=torch.bfloat16) * 0.5
206 +
207 + got = _mxfp8_dot_scaled_linear(x, w_fp8, w_scale)
208 + # Reference: consume the SAME quantized weights (isolates activation-quant
209 + # noise) -> dequant to bf16, plain matmul.
210 + w_deq = dequant_mxfp8_to_bf16(w_fp8, w_scale)
211 + ref = torch.nn.functional.linear(x, w_deq).to(x.dtype)
212 + assert got.shape == (m, n)
213 + # Only the activation is re-quantized inside the kernel -> small MX noise.
214 + assert _relerr(got, ref) < 5e-2
215 +
216 +
217 +# --------------------------------------------------------------------------- #
218 +# Native MXFP8 MoE (dot_scaled grouped GEMM) vs dequant-to-bf16 MoE math
219 +# --------------------------------------------------------------------------- #
220 +def _ref_moe(x, w13, w2, topk_weights, topk_ids, alpha, beta, limit):
221 + T, H = x.shape
222 + inter = w2.shape[-1]
223 + top_k = topk_ids.shape[1]
224 + out = torch.zeros(T, H, device=x.device, dtype=torch.float32)
225 + for t in range(T):
226 + for j in range(top_k):
227 + e = int(topk_ids[t, j].item())
228 + g1 = x[t].float() @ w13[e].float().T # [2I]
229 + gate = g1[:inter]
230 + up = g1[inter:]
231 + if limit is not None:
232 + gate = gate.clamp(max=limit)
233 + up = up.clamp(min=-limit, max=limit)
234 + act = gate * torch.sigmoid(alpha * gate) * (up + beta)
235 + g2 = act @ w2[e].float().T # [H]
236 + out[t] += topk_weights[t, j].float() * g2
237 + return out.to(x.dtype)
238 +
239 +
240 +@requires_gfx950
241 +@pytest.mark.parametrize(
242 + "T,H,inter,E,top_k", [(8, 256, 512, 8, 2), (1, 512, 256, 16, 4)]
243 +)
244 +@torch.inference_mode()
245 +def test_mxfp8_native_moe(T, H, inter, E, top_k):
246 + from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
247 + fused_moe_mxfp8_native,
248 + )
249 +
250 + torch.manual_seed(0)
251 + alpha, beta, limit = 1.702, 1.0, 7.0
252 + w13_bf16 = torch.randn(E, 2 * inter, H, device=DEVICE, dtype=torch.bfloat16) * 0.1
253 + w2_bf16 = torch.randn(E, H, inter, device=DEVICE, dtype=torch.bfloat16) * 0.1
254 + w13_fp8, w13_scale = _mxfp8_e4m3_quantize_torch(
255 + w13_bf16, is_sf_swizzled_layout=False
256 + )
257 + w2_fp8, w2_scale = _mxfp8_e4m3_quantize_torch(w2_bf16, is_sf_swizzled_layout=False)
258 +
259 + x = torch.randn(T, H, device=DEVICE, dtype=torch.bfloat16) * 0.5
260 + logits = torch.randn(T, E, device=DEVICE, dtype=torch.float32)
261 + topk_weights, topk_ids = logits.softmax(dim=-1).topk(top_k, dim=-1)
262 + topk_weights = topk_weights.to(torch.float32)
263 + topk_ids = topk_ids.to(torch.int32)
264 +
265 + got = fused_moe_mxfp8_native(
266 + x,
267 + w13_fp8,
268 + w13_scale,
269 + w2_fp8,
270 + w2_scale,
271 + topk_weights,
272 + topk_ids,
273 + alpha=alpha,
274 + beta=beta,
275 + limit=limit,
276 + global_num_experts=E,
277 + expert_map=None,
278 + )
279 + # Reference consumes the dequantized weights (same bits the kernel reads).
280 + w13_deq = dequant_mxfp8_to_bf16(w13_fp8, w13_scale)
281 + w2_deq = dequant_mxfp8_to_bf16(w2_fp8, w2_scale)
282 + ref = _ref_moe(x, w13_deq, w2_deq, topk_weights, topk_ids, alpha, beta, limit)
283 + assert got.shape == (T, H)
284 + assert _relerr(got, ref) < 5e-2
285 +
286 +
287 +# --------------------------------------------------------------------------- #
288 +# MXFP8 linear emulation: BF16-at-load (default) vs per-step dequant + switch
289 +# --------------------------------------------------------------------------- #
290 +@pytest.mark.parametrize("shape", [(512, 2048), (1, 6144)])
291 +@pytest.mark.parametrize("act_dtype", [torch.bfloat16, torch.float16])
292 +@pytest.mark.parametrize("dequant_at_load", [True, False])
293 +@torch.inference_mode()
294 +def test_mxfp8_linear_emulation_bf16_at_load(
295 + shape, act_dtype, dequant_at_load, monkeypatch
296 +):
297 + """EmulationMxfp8LinearKernel load-time BF16 dequant (default) and the
298 + ``VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0`` per-step fallback must produce the
299 + same result; the dtype-match (BF16/FP16 activations) must also hold."""
300 + from vllm.model_executor.kernels.linear.mxfp8.emulation import (
301 + EmulationMxfp8LinearKernel,
302 + )
303 + from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
304 + Mxfp8LinearLayerConfig,
305 + )
306 +
307 + monkeypatch.setenv(
308 + "VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD", "1" if dequant_at_load else "0"
309 + )
310 + N, K = shape
311 + torch.manual_seed(0)
312 + w_bf16 = torch.randn(N, K, device=DEVICE, dtype=torch.bfloat16)
313 + w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
314 + assert w_scale.shape == (N, K // 32)
315 +
316 + # Reference: dequant once, plain linear in the activation dtype.
317 + w_ref = dequant_mxfp8_to_bf16(w_fp8, w_scale).to(act_dtype)
318 + x = torch.randn(7, K, device=DEVICE, dtype=act_dtype)
319 + out_ref = torch.nn.functional.linear(x, w_ref)
320 +
321 + layer = torch.nn.Module()
322 + layer.weight = torch.nn.Parameter(w_fp8.clone(), requires_grad=False)
323 + layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False)
324 +
325 + kernel = EmulationMxfp8LinearKernel(Mxfp8LinearLayerConfig())
326 + kernel.process_weights_after_loading(layer)
327 +
328 + if dequant_at_load:
329 + # weights converted to BF16 at load (>= 2-byte)
330 + assert layer.weight.element_size() >= 2
331 + else:
332 + # opt-out: weights stay 1-byte MXFP8, dequant happens per-step
333 + assert layer.weight.element_size() == 1
334 +
335 + out = kernel.apply_weights(layer, x)
336 + assert out.dtype == act_dtype # dtype-match preserved (no tl.dot/F.linear crash)
337 + assert _relerr(out.float(), out_ref.float()) < 2e-2
tests/models/multimodal/processing/test_minimax_m3.py new
+138
@@ -0,0 +1,138 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for MiniMax-M3 VL ``max_long_side_pixel`` resize support.
4 +
5 +These exercise the vendored processor directly (no checkpoint / GPU needed), so
6 +they validate the long-side resize spec and the resulting prompt-token counts
7 +deterministically.
8 +"""
9 +
10 +import pytest
11 +import torch
12 +
13 +from vllm.transformers_utils.processors.minimax_m3 import (
14 + IMAGE_MAX_TOTAL_PIXELS,
15 + MIN_SHORT_SIDE_PIXEL,
16 + VIDEO_MAX_TOTAL_PIXELS,
17 + MiniMaxM3VLImageProcessor,
18 + MiniMaxM3VLVideoProcessor,
19 + smart_resize,
20 +)
21 +
22 +# Long sides are multiples of patch_size*merge_size (28) so the rounding is
23 +# exact and the expected token counts are unambiguous.
24 +LONG_SIDES = [252, 504, 1008]
25 +MERGE2 = 2**2 # merge_size ** 2
26 +
27 +
28 +def _image_tokens(grid_thw) -> int:
29 + g = list(grid_thw)
30 + return int(g[0] * g[1] * g[2]) // MERGE2
31 +
32 +
33 +# --------------------------------------------------------------------------- #
34 +# smart_resize: the long-side spec (a) shrink / (b) enlarge / (c) hard cap
35 +# --------------------------------------------------------------------------- #
36 +def test_smart_resize_long_side_shrink():
37 + # (a) long side exceeds the cap -> shrink so the long side equals the cap.
38 + h, w = smart_resize(
39 + 2048, 1024, factor=28, max_long_side_pixel=1008, max_total_pixels=10**9
40 + )
41 + assert max(h, w) == 1008
42 + assert (h, w) == (1008, 504) # aspect ratio preserved
43 +
44 +
45 +def test_smart_resize_short_side_enlarge():
46 + # (b) long side within the cap but short side below the floor -> enlarge so
47 + # the short side reaches min_short_side_pixel.
48 + h, w = smart_resize(
49 + 200, 40, factor=28, max_long_side_pixel=1008, max_total_pixels=10**9
50 + )
51 + assert min(h, w) == MIN_SHORT_SIDE_PIXEL # 112
52 +
53 +
54 +def test_smart_resize_total_pixels_raises():
55 + # (c) still over the area cap after resizing -> raise instead of inferring.
56 + with pytest.raises(ValueError, match="max_total_pixels"):
57 + smart_resize(
58 + 5000,
59 + 5000,
60 + factor=28,
61 + max_long_side_pixel=4000,
62 + max_total_pixels=IMAGE_MAX_TOTAL_PIXELS,
63 + )
64 +
65 +
66 +def test_smart_resize_backward_compatible_area_bound():
67 + # Without max_long_side_pixel the original Qwen-style area bound is used.
68 + assert smart_resize(2048, 2048, factor=28, max_pixels=451584) == (672, 672)
69 +
70 +
71 +# --------------------------------------------------------------------------- #
72 +# Image processor: monotonic prompt-token counts for 252 < 504 < 1008
73 +# --------------------------------------------------------------------------- #
74 +def test_image_tokens_increase_with_max_long_side_pixel():
75 + proc = MiniMaxM3VLImageProcessor()
76 + counts = []
77 + for long_side in LONG_SIDES:
78 + patches = proc.get_number_of_image_patches(
79 + 2048, 2048, images_kwargs={"max_long_side_pixel": long_side}
80 + )
81 + counts.append(patches // MERGE2)
82 +
83 + assert counts == [81, 324, 1296]
84 + assert counts[0] < counts[1] < counts[2]
85 +
86 +
87 +def test_image_processor_defaults_match_spec():
88 + proc = MiniMaxM3VLImageProcessor()
89 + assert proc.max_long_side_pixel is None # opt-in
90 + assert proc.min_short_side_pixel == MIN_SHORT_SIDE_PIXEL
91 + assert proc.max_total_pixels == IMAGE_MAX_TOTAL_PIXELS
92 +
93 +
94 +def test_image_preprocess_pipeline_monotonic():
95 + proc = MiniMaxM3VLImageProcessor()
96 + image = torch.randint(0, 255, (3, 2048, 2048), dtype=torch.uint8)
97 + counts = []
98 + for long_side in LONG_SIDES:
99 + out = proc.preprocess(
100 + [image],
101 + do_resize=True,
102 + max_long_side_pixel=long_side,
103 + return_tensors="pt",
104 + )
105 + counts.append(_image_tokens(out["image_grid_thw"][0]))
106 + assert counts == [81, 324, 1296]
107 +
108 +
109 +# --------------------------------------------------------------------------- #
110 +# Video processor: same monotonic behavior + volumetric (w*h*frames) cap
111 +# --------------------------------------------------------------------------- #
112 +def test_video_tokens_increase_with_max_long_side_pixel():
113 + proc = MiniMaxM3VLVideoProcessor()
114 + assert proc.max_total_pixels == VIDEO_MAX_TOTAL_PIXELS
115 + video = torch.randint(0, 255, (4, 3, 2048, 2048), dtype=torch.uint8)
116 + counts = []
117 + for long_side in LONG_SIDES:
118 + out = proc.preprocess(
119 + videos=[video],
120 + do_resize=True,
121 + max_long_side_pixel=long_side,
122 + return_tensors="pt",
123 + )
124 + counts.append(_image_tokens(out["video_grid_thw"][0]))
125 + assert counts[0] < counts[1] < counts[2]
126 +
127 +
128 +def test_video_volumetric_cap_raises():
129 + proc = MiniMaxM3VLVideoProcessor()
130 + # 400 frames at a 1008-long-side square: 1008*1008*400 >> 301,056,000.
131 + video = torch.randint(0, 255, (400, 3, 2048, 2048), dtype=torch.uint8)
132 + with pytest.raises(ValueError, match="max_total_pixels"):
133 + proc.preprocess(
134 + videos=[video],
135 + do_resize=True,
136 + max_long_side_pixel=1008,
137 + return_tensors="pt",
138 + )
tests/models/registry.py
+15
@@ -420,6 +420,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
420 "MiniMaxAI/MiniMax-M2",
421 trust_remote_code=True,
422 ),
423 + "MiniMaxM3SparseForCausalLM": _HfExamplesInfo(
424 + "MiniMaxAI/MiniMax-M3",
425 + trust_remote_code=True,
426 + is_available_online=False,
427 + ),
428 "Ministral3ForCausalLM": _HfExamplesInfo("mistralai/Ministral-3-3B-Instruct-2512"),
429 "MistralForCausalLM": _HfExamplesInfo("mistralai/Mistral-7B-Instruct-v0.1"),
430 "MistralLarge3ForCausalLM": _HfExamplesInfo(
@@ -1099,6 +1104,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
1104 "MiniMaxAI/MiniMax-VL-01",
1105 trust_remote_code=True,
1106 ),
1107 + "MiniMaxM3SparseForConditionalGeneration": _HfExamplesInfo(
1108 + "MiniMaxAI/MiniMax-M3",
1109 + trust_remote_code=True,
1110 + is_available_online=False,
1111 + ),
1112 "Mistral3ForConditionalGeneration": _HfExamplesInfo(
1113 "mistralai/Mistral-Small-3.1-24B-Instruct-2503",
1114 extras={"fp8": "nm-testing/Mistral-Small-3.1-24B-Instruct-2503-FP8-dynamic"},
@@ -1601,6 +1611,11 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1611 speculative_model="XiaomiMiMo/MiMo-V2.5-Omni",
1612 is_available_online=False,
1613 ),
1614 + "MiniMaxM3MTP": _HfExamplesInfo(
1615 + "MiniMaxAI/MiniMax-M3",
1616 + trust_remote_code=True,
1617 + is_available_online=False,
1618 + ),
1619 "NemotronHMTPModel": _HfExamplesInfo(
1620 "nvidia/Nemotron-Super-Placeholder",
1621 speculative_model="nvidia/Nemotron-Super-Placeholder",
tests/reasoning/test_minimax_m3_reasoning_parser.py new
+320
@@ -0,0 +1,320 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import string
5 +from collections.abc import Sequence
6 +
7 +import pytest
8 +
9 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
10 +from vllm.reasoning import ReasoningParserManager
11 +from vllm.reasoning.minimax_m3_reasoning_parser import MiniMaxM3ReasoningParser
12 +
13 +pytestmark = pytest.mark.skip_global_cleanup
14 +
15 +
16 +class MiniMaxM3Tokenizer:
17 + """Small tokenizer with MiniMax M3 reasoning tags as special tokens."""
18 +
19 + special_tokens = ("<mm:think>", "</mm:think>")
20 +
21 + def __init__(self):
22 + self._token_to_id: dict[str, int] = {}
23 + self._id_to_token: dict[int, str] = {}
24 + for token in self.special_tokens:
25 + self._add_token(token)
26 + for char in string.printable:
27 + self._add_token(char)
28 +
29 + def _add_token(self, token: str) -> int:
30 + token_id = self._token_to_id.get(token)
31 + if token_id is None:
32 + token_id = len(self._token_to_id) + 1
33 + self._token_to_id[token] = token_id
34 + self._id_to_token[token_id] = token
35 + return token_id
36 +
37 + def get_vocab(self) -> dict[str, int]:
38 + return dict(self._token_to_id)
39 +
40 + def encode(
41 + self,
42 + text: str,
43 + truncation: bool | None = None,
44 + max_length: int | None = None,
45 + add_special_tokens: bool = True,
46 + ) -> list[int]:
47 + return [self._add_token(token) for token in self.tokenize(text)]
48 +
49 + def decode(
50 + self, ids: Sequence[int] | int, skip_special_tokens: bool = False
51 + ) -> str:
52 + if isinstance(ids, int):
53 + ids = [ids]
54 + return "".join(self._id_to_token[token_id] for token_id in ids)
55 +
56 + def tokenize(self, text: str) -> list[str]:
57 + tokens: list[str] = []
58 + pos = 0
59 + while pos < len(text):
60 + for special_token in self.special_tokens:
61 + if text.startswith(special_token, pos):
62 + tokens.append(special_token)
63 + pos += len(special_token)
64 + break
65 + else:
66 + tokens.append(text[pos])
67 + pos += 1
68 + return tokens
69 +
70 + def convert_ids_to_tokens(
71 + self,
72 + ids: Sequence[int],
73 + skip_special_tokens: bool = False,
74 + ) -> list[str]:
75 + return [self._id_to_token[token_id] for token_id in ids]
76 +
77 + def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]:
78 + if isinstance(tokens, str):
79 + return self._add_token(tokens)
80 + return [self._add_token(token) for token in tokens]
81 +
82 + def convert_tokens_to_string(self, tokens: list[str]) -> str:
83 + return "".join(tokens)
84 +
85 +
86 +def make_parser(
87 + chat_template_kwargs: dict[str, str] | None = None,
88 +) -> tuple[MiniMaxM3ReasoningParser, MiniMaxM3Tokenizer]:
89 + tokenizer = MiniMaxM3Tokenizer()
90 + return (
91 + MiniMaxM3ReasoningParser(tokenizer, chat_template_kwargs=chat_template_kwargs),
92 + tokenizer,
93 + )
94 +
95 +
96 +def run_streaming(
97 + parser: MiniMaxM3ReasoningParser,
98 + tokenizer: MiniMaxM3Tokenizer,
99 + chunks: list[str],
100 +) -> tuple[str | None, str | None, list[bool]]:
101 + previous_text = ""
102 + previous_token_ids: list[int] = []
103 + reasoning_parts: list[str] = []
104 + content_parts: list[str] = []
105 + reasoning_end_states: list[bool] = []
106 +
107 + for chunk in chunks:
108 + delta_token_ids = tokenizer.encode(chunk, add_special_tokens=False)
109 + current_text = previous_text + chunk
110 + current_token_ids = previous_token_ids + delta_token_ids
111 + delta = parser.extract_reasoning_streaming(
112 + previous_text=previous_text,
113 + current_text=current_text,
114 + delta_text=chunk,
115 + previous_token_ids=previous_token_ids,
116 + current_token_ids=current_token_ids,
117 + delta_token_ids=delta_token_ids,
118 + )
119 + reasoning_end_states.append(
120 + parser.is_reasoning_end_streaming(current_token_ids, delta_token_ids)
121 + )
122 +
123 + if delta is not None:
124 + if delta.reasoning is not None:
125 + reasoning_parts.append(delta.reasoning)
126 + if delta.content is not None:
127 + content_parts.append(delta.content)
128 +
129 + previous_text = current_text
130 + previous_token_ids = current_token_ids
131 +
132 + return (
133 + "".join(reasoning_parts) or None,
134 + "".join(content_parts) or None,
135 + reasoning_end_states,
136 + )
137 +
138 +
139 +def test_parser_registration():
140 + parser_cls = ReasoningParserManager.get_reasoning_parser("minimax_m3")
141 +
142 + assert parser_cls is MiniMaxM3ReasoningParser
143 +
144 +
145 +def test_nonstreaming_extracts_explicit_reasoning_block():
146 + parser, _ = make_parser()
147 + request = ChatCompletionRequest(messages=[], model="test-model")
148 +
149 + reasoning, content = parser.extract_reasoning(
150 + "<mm:think>plan</mm:think>answer", request
151 + )
152 +
153 + assert reasoning == "plan"
154 + assert content == "answer"
155 +
156 +
157 +def test_nonstreaming_without_start_tag_is_content():
158 + parser, _ = make_parser()
159 + request = ChatCompletionRequest(messages=[], model="test-model")
160 +
161 + reasoning, content = parser.extract_reasoning("plain answer", request)
162 +
163 + assert reasoning is None
164 + assert content == "plain answer"
165 +
166 +
167 +def test_nonstreaming_drops_leading_end_tag():
168 + parser, _ = make_parser()
169 + request = ChatCompletionRequest(messages=[], model="test-model")
170 +
171 + reasoning, content = parser.extract_reasoning("</mm:think>answer", request)
172 +
173 + assert reasoning is None
174 + assert content == "answer"
175 +
176 +
177 +def test_nonstreaming_non_leading_end_tag_is_content():
178 + parser, _ = make_parser()
179 + request = ChatCompletionRequest(messages=[], model="test-model")
180 +
181 + reasoning, content = parser.extract_reasoning("XXX</mm:think>YYY", request)
182 +
183 + assert reasoning is None
184 + assert content == "XXX</mm:think>YYY"
185 +
186 +
187 +def test_nonstreaming_enabled_mode_starts_in_reasoning():
188 + parser, _ = make_parser(chat_template_kwargs={"thinking_mode": "enabled"})
189 + request = ChatCompletionRequest(messages=[], model="test-model")
190 +
191 + reasoning, content = parser.extract_reasoning("plan</mm:think>answer", request)
192 +
193 + assert reasoning == "plan"
194 + assert content == "answer"
195 +
196 +
197 +def test_nonstreaming_open_reasoning_block():
198 + parser, _ = make_parser()
199 + request = ChatCompletionRequest(messages=[], model="test-model")
200 +
201 + reasoning, content = parser.extract_reasoning("<mm:think>still thinking", request)
202 +
203 + assert reasoning == "still thinking"
204 + assert content is None
205 +
206 +
207 +def test_streaming_reasoning_tags_are_not_returned():
208 + parser, tokenizer = make_parser()
209 +
210 + reasoning, content, end_states = run_streaming(
211 + parser,
212 + tokenizer,
213 + ["<mm:think>", "plan", "</mm:think>", "answer"],
214 + )
215 +
216 + assert reasoning == "plan"
217 + assert content == "answer"
218 + assert end_states == [False, False, True, True]
219 +
220 +
221 +def test_streaming_boundary_can_emit_reasoning_and_content():
222 + parser, tokenizer = make_parser()
223 +
224 + reasoning, content, end_states = run_streaming(
225 + parser,
226 + tokenizer,
227 + ["<mm:think>plan</mm:think>answer"],
228 + )
229 +
230 + assert reasoning == "plan"
231 + assert content == "answer"
232 + assert end_states == [True]
233 +
234 +
235 +def test_streaming_drops_leading_end_tag():
236 + parser, tokenizer = make_parser()
237 +
238 + reasoning, content, end_states = run_streaming(
239 + parser,
240 + tokenizer,
241 + ["</mm:think>", "answer"],
242 + )
243 +
244 + assert reasoning is None
245 + assert content == "answer"
246 + assert end_states == [True, True]
247 +
248 +
249 +def test_streaming_non_leading_end_tag_is_content():
250 + parser, tokenizer = make_parser()
251 +
252 + reasoning, content, end_states = run_streaming(
253 + parser,
254 + tokenizer,
255 + ["XXX</mm:think>YYY"],
256 + )
257 +
258 + assert reasoning is None
259 + assert content == "XXX</mm:think>YYY"
260 + assert end_states == [True]
261 +
262 +
263 +def test_streaming_enabled_mode_starts_in_reasoning():
264 + parser, tokenizer = make_parser(chat_template_kwargs={"thinking_mode": "enabled"})
265 +
266 + reasoning, content, end_states = run_streaming(
267 + parser,
268 + tokenizer,
269 + ["plan", "</mm:think>", "answer"],
270 + )
271 +
272 + assert reasoning == "plan"
273 + assert content == "answer"
274 + assert end_states == [False, True, True]
275 +
276 +
277 +def test_streaming_plain_content_ends_reasoning_phase():
278 + parser, tokenizer = make_parser()
279 +
280 + reasoning, content, end_states = run_streaming(
281 + parser,
282 + tokenizer,
283 + ["plain ", "answer"],
284 + )
285 +
286 + assert reasoning is None
287 + assert content == "plain answer"
288 + assert end_states == [True, True]
289 +
290 +
291 +def test_token_id_helpers():
292 + parser, tokenizer = make_parser()
293 + output_ids = tokenizer.encode(
294 + "<mm:think>abc</mm:think>def", add_special_tokens=False
295 + )
296 + open_reasoning_ids = tokenizer.encode("<mm:think>abc", add_special_tokens=False)
297 + content_ids = tokenizer.encode("plain", add_special_tokens=False)
298 +
299 + assert parser.is_reasoning_end(output_ids)
300 + assert not parser.is_reasoning_end(open_reasoning_ids)
301 + assert not parser.is_reasoning_end(content_ids)
302 + assert tokenizer.decode(parser.extract_content_ids(output_ids)) == "def"
303 + assert parser.extract_content_ids(open_reasoning_ids) == []
304 + assert parser.extract_content_ids(content_ids) == content_ids
305 + assert parser.count_reasoning_tokens(output_ids) == len(tokenizer.encode("abc"))
306 +
307 +
308 +def test_token_id_helpers_enabled_mode():
309 + parser, tokenizer = make_parser(chat_template_kwargs={"thinking_mode": "enabled"})
310 + output_ids = tokenizer.encode("abc</mm:think>def", add_special_tokens=False)
311 + open_reasoning_ids = tokenizer.encode("abc", add_special_tokens=False)
312 +
313 + assert parser.is_reasoning_end(output_ids)
314 + assert not parser.is_reasoning_end(open_reasoning_ids)
315 + assert tokenizer.decode(parser.extract_content_ids(output_ids)) == "def"
316 + assert parser.extract_content_ids(open_reasoning_ids) == []
317 + assert parser.count_reasoning_tokens(output_ids) == len(tokenizer.encode("abc"))
318 + assert parser.count_reasoning_tokens(open_reasoning_ids) == len(
319 + tokenizer.encode("abc")
320 + )
tests/tool_parsers/test_minimax_m3_tool_parser.py new
+261
@@ -0,0 +1,261 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import json
5 +from typing import Any
6 +
7 +import pytest
8 +
9 +from vllm.entrypoints.openai.chat_completion.protocol import (
10 + ChatCompletionToolsParam,
11 + FunctionDefinition,
12 +)
13 +from vllm.entrypoints.openai.engine.protocol import DeltaMessage
14 +from vllm.tool_parsers import ToolParserManager
15 +from vllm.tool_parsers.minimax_m3_tool_parser import MinimaxM3ToolParser
16 +
17 +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
18 +
19 +NS = "]<]minimax[>["
20 +EOS_ID = 99
21 +
22 +
23 +class FakeTokenizer:
24 + """Minimal fake tokenizer for unit tests."""
25 +
26 + def __init__(self):
27 + self.model_tokenizer = True
28 + self.vocab: dict[str, int] = {}
29 +
30 + def get_vocab(self) -> dict[str, int]:
31 + return self.vocab
32 +
33 +
34 +def sample_tools() -> list[ChatCompletionToolsParam]:
35 + return [
36 + ChatCompletionToolsParam(
37 + function=FunctionDefinition(
38 + name="create_order",
39 + parameters={
40 + "type": "object",
41 + "properties": {
42 + "user_id": {"type": "integer"},
43 + "urgent": {"type": "boolean"},
44 + "note": {"type": "string"},
45 + "shipping": {
46 + "type": "object",
47 + "properties": {
48 + "city": {"type": "string"},
49 + "zip": {"type": "integer"},
50 + },
51 + },
52 + "items": {
53 + "type": "array",
54 + "items": {
55 + "type": "object",
56 + "properties": {
57 + "sku": {"type": "string"},
58 + "qty": {"type": "integer"},
59 + },
60 + },
61 + },
62 + "metadata": {
63 + "type": "object",
64 + "additionalProperties": {"type": "string"},
65 + },
66 + "duplicate_demo": {"type": "object"},
67 + },
68 + },
69 + ),
70 + )
71 + ]
72 +
73 +
74 +@pytest.fixture
75 +def parser() -> MinimaxM3ToolParser:
76 + return MinimaxM3ToolParser(FakeTokenizer(), tools=sample_tools())
77 +
78 +
79 +def build_order_call() -> str:
80 + return (
81 + f"{NS}<tool_call>\n"
82 + f'{NS}<invoke name="create_order">'
83 + f"{NS}<user_id>42{NS}</user_id>"
84 + f"{NS}<urgent>true{NS}</urgent>"
85 + f"{NS}<note>Please leave at front desk.{NS}</note>"
86 + f"{NS}<shipping>"
87 + f"{NS}<city>Singapore{NS}</city>"
88 + f"{NS}<zip>018956{NS}</zip>"
89 + f"{NS}</shipping>"
90 + f"{NS}<items>"
91 + f"{NS}<item>{NS}<sku>book-001{NS}</sku>{NS}<qty>2{NS}</qty>{NS}</item>"
92 + f"{NS}<item>{NS}<sku>pen-007{NS}</sku>{NS}<qty>5{NS}</qty>{NS}</item>"
93 + f"{NS}</items>"
94 + f"{NS}<metadata>"
95 + f"{NS}<source>mobile{NS}</source>"
96 + f"{NS}<campaign>may-launch{NS}</campaign>"
97 + f"{NS}</metadata>"
98 + f"{NS}<duplicate_demo>"
99 + f"{NS}<tag>a{NS}</tag>"
100 + f"{NS}<tag>b{NS}</tag>"
101 + f"{NS}</duplicate_demo>"
102 + f"{NS}</invoke>\n"
103 + f"{NS}</tool_call>"
104 + )
105 +
106 +
107 +def build_order_invocation(user_id: int) -> str:
108 + return (
109 + f'{NS}<invoke name="create_order">'
110 + f"{NS}<user_id>{user_id}{NS}</user_id>"
111 + f"{NS}</invoke>"
112 + )
113 +
114 +
115 +def build_multiple_order_call() -> str:
116 + return (
117 + f"{NS}<tool_call>\n"
118 + f"{build_order_invocation(1)}\n"
119 + f"{build_order_invocation(2)}\n"
120 + f"{NS}</tool_call>"
121 + )
122 +
123 +
124 +def _feed(
125 + parser: MinimaxM3ToolParser, chunks: list[str | tuple[str, list[int]]]
126 +) -> list[DeltaMessage]:
127 + previous = ""
128 + results: list[DeltaMessage] = []
129 + for chunk in chunks:
130 + if isinstance(chunk, tuple):
131 + delta, delta_ids = chunk
132 + else:
133 + delta = chunk
134 + delta_ids = []
135 +
136 + current = previous + delta
137 + result = parser.extract_tool_calls_streaming(
138 + previous_text=previous,
139 + current_text=current,
140 + delta_text=delta,
141 + previous_token_ids=[],
142 + current_token_ids=[],
143 + delta_token_ids=delta_ids,
144 + request=None,
145 + )
146 + if result is not None:
147 + results.append(result)
148 + previous = current
149 + return results
150 +
151 +
152 +def _collect_content(results: list[DeltaMessage]) -> str:
153 + return "".join(result.content for result in results if result.content)
154 +
155 +
156 +def _collect_tool_calls(results: list[DeltaMessage]) -> dict[int, dict[str, Any]]:
157 + tool_calls: dict[int, dict[str, Any]] = {}
158 + for result in results:
159 + for tool_call in result.tool_calls or []:
160 + tool_calls.setdefault(
161 + tool_call.index,
162 + {"id": None, "name": "", "arguments": ""},
163 + )
164 + if tool_call.id:
165 + tool_calls[tool_call.index]["id"] = tool_call.id
166 + if tool_call.function:
167 + if tool_call.function.name:
168 + tool_calls[tool_call.index]["name"] += tool_call.function.name
169 + if tool_call.function.arguments:
170 + tool_calls[tool_call.index]["arguments"] += (
171 + tool_call.function.arguments
172 + )
173 + return tool_calls
174 +
175 +
176 +def test_minimax_m3_parser_registered():
177 + assert ToolParserManager.get_tool_parser("minimax_m3") is MinimaxM3ToolParser
178 +
179 +
180 +def test_non_streaming_nested_tool_call(parser):
181 + result = parser.extract_tool_calls(
182 + "I will create it.\n" + build_order_call(),
183 + request=None,
184 + )
185 +
186 + assert result.tools_called
187 + assert result.content == "I will create it.\n"
188 + assert len(result.tool_calls) == 1
189 + tool_call = result.tool_calls[0]
190 + assert tool_call.function.name == "create_order"
191 + assert json.loads(tool_call.function.arguments) == {
192 + "user_id": 42,
193 + "urgent": True,
194 + "note": "Please leave at front desk.",
195 + "shipping": {"city": "Singapore", "zip": 18956},
196 + "items": [
197 + {"sku": "book-001", "qty": 2},
198 + {"sku": "pen-007", "qty": 5},
199 + ],
200 + "metadata": {
201 + "source": "mobile",
202 + "campaign": "may-launch",
203 + },
204 + "duplicate_demo": {"tag": ["a", "b"]},
205 + }
206 +
207 +
208 +def test_non_streaming_without_tool_call_keeps_content(parser):
209 + result = parser.extract_tool_calls("plain response", request=None)
210 +
211 + assert not result.tools_called
212 + assert result.tool_calls == []
213 + assert result.content == "plain response"
214 +
215 +
216 +def test_non_streaming_multiple_tool_calls(parser):
217 + result = parser.extract_tool_calls(build_multiple_order_call(), request=None)
218 +
219 + assert result.tools_called
220 + assert result.content is None
221 + assert [tool_call.function.name for tool_call in result.tool_calls] == [
222 + "create_order",
223 + "create_order",
224 + ]
225 + assert [
226 + json.loads(tool_call.function.arguments)["user_id"]
227 + for tool_call in result.tool_calls
228 + ] == [1, 2]
229 +
230 +
231 +def test_streaming_without_tool_call_emits_text(parser):
232 + results = _feed(parser, ["plain ", "response"])
233 +
234 + assert _collect_content(results) == "plain response"
235 + assert _collect_tool_calls(results) == {}
236 +
237 +
238 +def test_streaming_nested_tool_call(parser):
239 + tool_call_text = build_order_call()
240 + results = _feed(
241 + parser,
242 + [
243 + "I will create it.\n",
244 + tool_call_text[:5],
245 + tool_call_text[5:17],
246 + tool_call_text[17:120],
247 + tool_call_text[120:],
248 + ("", [EOS_ID]),
249 + ],
250 + )
251 +
252 + assert _collect_content(results) == "I will create it.\n"
253 + tool_calls = _collect_tool_calls(results)
254 + assert len(tool_calls) == 1
255 + assert tool_calls[0]["name"] == "create_order"
256 + assert tool_calls[0]["id"] is not None
257 + assert json.loads(tool_calls[0]["arguments"]) == json.loads(
258 + parser.streamed_args_for_tool[0]
259 + )
260 + assert json.loads(parser.prev_tool_call_arr[0]["arguments"])["items"][1]["qty"] == 5
261 + assert results[-1].content is None
tools/pre_commit/generate_attention_backend_docs.py
+36 -4
@@ -30,6 +30,7 @@ REPO_ROOT = Path(__file__).parent.parent.parent
30 RELEVANT_PATTERNS = [
31 "vllm/v1/attention/backends/*.py",
32 "vllm/v1/attention/backends/**/*.py",
33 + "vllm/models/minimax_m3/common/sparse_attention.py",
34 "vllm/model_executor/layers/attention/mla_attention.py",
35 "vllm/platforms/cuda.py",
36 "tools/pre_commit/generate_attention_backend_docs.py",
@@ -1633,6 +1634,24 @@ def generate_mla_section(
1634 return "\n".join(lines)
1635
1636
1637 +def generate_minimax_section(backends: list[dict[str, Any]]) -> str:
1638 + """Generate the MiniMax M3 sparse attention section."""
1639 + lines = [
1640 + "## MiniMax M3 Sparse Attention Backends",
1641 + "",
1642 + 'Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")',
1643 + "layers. It is wired in directly by the model and is not part of the",
1644 + "automatic priority lists above. A lightning indexer scores KV blocks, the",
1645 + "top-k blocks (plus fixed init/local blocks) are selected, and attention",
1646 + "attends only to those blocks; index keys live in a separate side cache.",
1647 + "",
1648 + ]
1649 + columns = _build_columns(is_mla=False, has_versions=False)
1650 + lines.extend(_render_table(columns, backends))
1651 + lines.append("")
1652 + return "\n".join(lines)
1653 +
1654 +
1655 # ---------------------------------------------------------------------------
1656 # Top-level orchestration
1657 # ---------------------------------------------------------------------------
@@ -1669,15 +1688,24 @@ def generate_docs() -> str:
1688 if fi_features:
1689 all_backends = _expand_flashinfer_variants(all_backends, fi_features)
1690
1672 - # DeepSeek V4 (*_DSV4) decode backends get their own subsection rather than
1673 - # mixing into the main MLA / standard tables (the ROCm V4 backend isn't
1674 - # flagged is_mla by the AST heuristic, so filter purely on the name).
1691 + # DeepSeek V4 (*_DSV4) decode backends and MiniMax M3 sparse backends each
1692 + # get their own subsection rather than mixing into the main MLA / standard
1693 + # tables (the ROCm V4 backend isn't flagged is_mla by the AST heuristic, so
1694 + # filter purely on the name).
1695 def _is_v4(b: dict[str, Any]) -> bool:
1696 return b["name"].endswith("_DSV4")
1697
1698 + def _is_minimax(b: dict[str, Any]) -> bool:
1699 + return not b["is_mla"] and not _is_v4(b) and b["name"].startswith("MINIMAX")
1700 +
1701 v4_decode_backends = [b for b in all_backends if _is_v4(b)]
1702 + minimax_backends = [b for b in all_backends if _is_minimax(b)]
1703 mla_backends = [b for b in all_backends if b["is_mla"] and not _is_v4(b)]
1680 - non_mla_backends = [b for b in all_backends if not b["is_mla"] and not _is_v4(b)]
1704 + non_mla_backends = [
1705 + b
1706 + for b in all_backends
1707 + if not b["is_mla"] and not _is_v4(b) and not _is_minimax(b)
1708 + ]
1709
1710 # Generate documentation
1711 script_path = "tools/pre_commit/generate_attention_backend_docs.py"
@@ -1726,6 +1754,10 @@ def generate_docs() -> str:
1754 if footnotes:
1755 doc_lines.append("\n>\n".join(footnotes) + "\n")
1756
1757 + # Add MiniMax M3 sparse section (separate category after standard GQA)
1758 + if minimax_backends:
1759 + doc_lines.append(generate_minimax_section(minimax_backends))
1760 +
1761 # Add MLA section with prefill and decode backends
1762 doc_lines.append(
1763 generate_mla_section(mla_prefill_backends, mla_backends, v4_decode_backends)
vllm/_custom_ops.py
+64
@@ -2614,6 +2614,70 @@ def reshape_and_cache_flash(
2614 )
2615
2616
2617 +def fused_minimax_m3_qknorm_rope_kv_insert(
2618 + qkv: torch.Tensor,
2619 + q_norm_weight: torch.Tensor,
2620 + k_norm_weight: torch.Tensor,
2621 + cos_sin_cache: torch.Tensor,
2622 + positions: torch.Tensor,
2623 + num_heads: int,
2624 + num_kv_heads: int,
2625 + rotary_dim: int,
2626 + eps: float,
2627 + index_q_norm_weight: torch.Tensor | None = None,
2628 + index_k_norm_weight: torch.Tensor | None = None,
2629 + num_index_heads: int = 0,
2630 + slot_mapping: torch.Tensor | None = None,
2631 + index_slot_mapping: torch.Tensor | None = None,
2632 + kv_cache: torch.Tensor | None = None,
2633 + index_cache: torch.Tensor | None = None,
2634 + block_size: int = 0,
2635 + q_out: torch.Tensor | None = None,
2636 + index_q_out: torch.Tensor | None = None,
2637 +) -> None:
2638 + """Fused MiniMax-M3 attention pre-processing (in-place).
2639 +
2640 + Applies Gemma RMSNorm + partial NeoX RoPE to ``qkv`` in place. ``qkv`` is a
2641 + single fused tensor:
2642 +
2643 + - dense layer (``num_index_heads == 0``): ``[q | k | v]``;
2644 + - sparse layer (``num_index_heads > 0``): ``[q | k | v | index_q |
2645 + index_k]`` — the index branch is read straight out of ``qkv``.
2646 +
2647 + 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
2650 + ``index_slot_mapping`` is omitted, ``slot_mapping`` is used for both caches.
2651 +
2652 + If ``q_out`` / ``index_q_out`` (contiguous ``[N, nq*128]`` / ``[N,
2653 + niq*128]``) are given, the normed/roped q / index_q are written there
2654 + instead of in place — folding the de-interleave into this kernel's store so
2655 + callers skip a separate ``.contiguous()`` copy before the SM100 sparse
2656 + attention's flat TMA descriptor.
2657 + """
2658 + torch.ops._C.fused_minimax_m3_qknorm_rope_kv_insert(
2659 + qkv,
2660 + q_norm_weight,
2661 + k_norm_weight,
2662 + cos_sin_cache,
2663 + positions,
2664 + num_heads,
2665 + num_kv_heads,
2666 + rotary_dim,
2667 + eps,
2668 + index_q_norm_weight,
2669 + index_k_norm_weight,
2670 + num_index_heads,
2671 + slot_mapping,
2672 + index_slot_mapping,
2673 + kv_cache,
2674 + index_cache,
2675 + block_size,
2676 + q_out,
2677 + index_q_out,
2678 + )
2679 +
2680 +
2681 def concat_and_cache_mla(
2682 kv_c: torch.Tensor,
2683 k_pe: torch.Tensor,
vllm/config/attention.py
+6
@@ -9,6 +9,8 @@ from vllm.config.utils import config
9 from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
10 from vllm.v1.attention.backends.registry import AttentionBackendEnum
11
12 +IndexerKVDType = Literal["bf16", "fp8", "mxfp4", "nvfp4"]
13 +
14
15 @config
16 class AttentionConfig:
@@ -50,6 +52,10 @@ class AttentionConfig:
52 use_fp4_indexer_cache: bool = False
53 """If set, use fp4 indexer cache for dsv32 family model (not support yet)"""
54
55 + indexer_kv_dtype: IndexerKVDType = "bf16"
56 + """Data type for the sparse-attention indexer K cache. Quantized formats
57 + (fp8, mxfp4, nvfp4) require indexer kernel support in the backend."""
58 +
59 use_non_causal: bool = False
60 """Whether to use non-causal (bidirectional) attention."""
61
vllm/config/speculative.py
+31
@@ -45,6 +45,7 @@ MTPModelTypes = Literal[
45 "qwen3_next_mtp",
46 "qwen3_5_mtp",
47 "longcat_flash_mtp",
48 + "minimax_m3_mtp",
49 "mtp",
50 "pangu_ultra_moe_mtp",
51 "step3p5_mtp",
@@ -528,6 +529,36 @@ class SpeculativeConfig:
529 text_config.num_kv_shared_layers = 0
530 hf_config.update({"n_predict": 1, "architectures": ["Gemma4MTPModel"]})
531
532 + if (
533 + hf_config.model_type == "minimax_m3_vl"
534 + or initial_architecture == "MiniMaxM3SparseForConditionalGeneration"
535 + ):
536 + # MTP modules live on the language model of this VL checkpoint, so
537 + # promote text_config before rewriting it into an MTP config.
538 + quantization_config = getattr(hf_config, "quantization_config", None)
539 + hf_config = getattr(hf_config, "text_config", hf_config)
540 + if (
541 + quantization_config is not None
542 + and getattr(hf_config, "quantization_config", None) is None
543 + ):
544 + hf_config.update({"quantization_config": quantization_config})
545 + hf_config.model_type = "minimax_m3_mtp"
546 + n_predict = getattr(hf_config, "num_mtp_modules", 1)
547 + hf_config.update(
548 + {"n_predict": n_predict, "architectures": ["MiniMaxM3MTP"]}
549 + )
550 + elif (
551 + hf_config.model_type == "minimax_m3_mtp"
552 + or initial_architecture == "MiniMaxM3MTP"
553 + ):
554 + # Standalone MTP checkpoints already use a flat MTP config with no
555 + # VL wrapper / text_config to promote, so just normalize the
556 + # architecture and derive n_predict from num_mtp_modules.
557 + n_predict = getattr(hf_config, "num_mtp_modules", 1)
558 + hf_config.update(
559 + {"n_predict": n_predict, "architectures": ["MiniMaxM3MTP"]}
560 + )
561 +
562 return hf_config
563
564 def __post_init__(self):
vllm/config/vllm.py
+9 -3
@@ -1101,20 +1101,26 @@ class VllmConfig:
1101 )
1102 self.compilation_config.mode = CompilationMode.NONE
1103
1104 - # DeepSeek V4's model classes don't carry @support_torch_compile —
1104 + # For model classes don't carry @support_torch_compile —
1105 # the breakable cudagraph is the supported PIECEWISE path. Auto-enable
1106 # it unless the user has explicitly opted out via the env var.
1107 if (
1108 self.model_config is not None
1109 and "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ
1110 and any(
1111 - a in ("DeepseekV4ForCausalLM", "DeepSeekV4MTPModel")
1111 + a
1112 + in (
1113 + "DeepseekV4ForCausalLM",
1114 + "DeepSeekV4MTPModel",
1115 + "MiniMaxM3SparseForCausalLM",
1116 + "MiniMaxM3SparseForConditionalGeneration",
1117 + )
1118 for a in self.model_config.architectures
1119 )
1120 ):
1121 os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
1122 logger.info_once(
1117 - "Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1 for DeepSeek V4. "
1123 + "Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1. "
1124 "Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out."
1125 )
1126
vllm/envs.py
+10
@@ -113,6 +113,7 @@ if TYPE_CHECKING:
113 VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE: bool = True
114 VLLM_DISABLE_PYNCCL: bool = False
115 VLLM_USE_OINK_OPS: bool = False
116 + VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD: bool = True
117 VLLM_ROCM_USE_AITER: bool = False
118 VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False
119 VLLM_ROCM_USE_AITER_LINEAR: bool = True
@@ -1092,6 +1093,15 @@ environment_variables: dict[str, Callable[[], Any]] = {
1093 ),
1094 # Disable aiter ops unless specifically enabled.
1095 # Acts as a parent switch to enable the rest of the other operations.
1096 + # On hardware without a native MXFP8 kernel (e.g. ROCm gfx942 / MI300), the
1097 + # MXFP8 emulation path dequantizes weights MXFP8->BF16 once at load time and
1098 + # runs as a BF16 checkpoint (no per-step dequant). Set to 0 to fall back to
1099 + # per-step dequant: keeps the 1-byte MXFP8 weights (~half the weight memory)
1100 + # at the cost of dequantizing every forward step (much slower). Default on.
1101 + "VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD": lambda: (
1102 + os.getenv("VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD", "True").lower()
1103 + in ("true", "1")
1104 + ),
1105 "VLLM_ROCM_USE_AITER": lambda: (
1106 os.getenv("VLLM_ROCM_USE_AITER", "False").lower() in ("true", "1")
1107 ),
vllm/model_executor/kernels/linear/__init__.py
+6
@@ -93,6 +93,9 @@ from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
93 from vllm.model_executor.kernels.linear.mxfp8.marlin import (
94 MarlinMxfp8LinearKernel,
95 )
96 +from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
97 + RocmDotScaledMxfp8LinearKernel,
98 +)
99 from vllm.model_executor.kernels.linear.mxfp8.xpu import (
100 XPUMxFp8LinearKernel,
101 )
@@ -383,6 +386,9 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
386 EmulationMxfp8LinearKernel,
387 ],
388 PlatformEnum.ROCM: [
389 + # Native CDNA4 (gfx950) MX linear; is_supported() gates to gfx95x and
390 + # falls through to BF16 emulation (hipBLASLt) elsewhere / on regression.
391 + RocmDotScaledMxfp8LinearKernel,
392 EmulationMxfp8LinearKernel,
393 ],
394 PlatformEnum.XPU: [
vllm/model_executor/kernels/linear/mxfp8/emulation.py
+25 -1
@@ -33,6 +33,17 @@ class EmulationMxfp8LinearKernel(Mxfp8LinearKernel):
33
34 weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
35
36 + # Dequantize MXFP8 -> BF16 ONCE here, at load time, so apply_weights runs
37 + # a plain BF16 linear with no per-step dequant -- i.e. run as if from a
38 + # BF16 checkpoint. The 1-byte MXFP8 weight is replaced by BF16 (2x its
39 + # size, but linear weights are small vs the MoE experts); the tiny E8M0
40 + # scale is kept for the dtype/ndim asserts but is otherwise unused.
41 + # Opt out (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0) to keep the MXFP8
42 + # weight and dequant per-step in apply_weights instead.
43 + import vllm.envs as envs
44 +
45 + if envs.VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD:
46 + weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale)
47 layer.weight = Parameter(weight.contiguous(), requires_grad=False)
48 layer.weight_scale = Parameter(weight_scale, requires_grad=False)
49
@@ -42,6 +53,17 @@ class EmulationMxfp8LinearKernel(Mxfp8LinearKernel):
53 x: torch.Tensor,
54 bias: torch.Tensor | None = None,
55 ) -> torch.Tensor:
56 + weight = layer.weight
57 + # Load-time dequant path: weights are already BF16/FP16 (>= 2-byte), so
58 + # run a plain linear -- no per-step dequant. (MXFP8 weights are 1-byte.)
59 + if weight.element_size() >= 2:
60 + # F.linear requires x and weight share a dtype; .to() is a no-op when
61 + # they already match (e.g. both BF16).
62 + output = torch.nn.functional.linear(x, weight.to(x.dtype), bias)
63 + return output.to(x.dtype)
64 +
65 + # Fallback: weights still in MXFP8 -- dequant on the fly (other archs /
66 + # if a future caller skips the load-time conversion above).
67 weight_scale = layer.weight_scale
68 if weight_scale.dtype != MXFP8_SCALE_DTYPE:
69 raise ValueError(
@@ -55,6 +77,8 @@ class EmulationMxfp8LinearKernel(Mxfp8LinearKernel):
77 f"Ensure process_weights_after_loading was called."
78 )
79
58 - weight_bf16 = dequant_mxfp8_to_bf16(layer.weight, weight_scale)
80 + # Cast to x's dtype: dequant yields BF16, but F.linear needs both operands
81 + # to match (e.g. an FP16 model). No-op when x is already BF16.
82 + weight_bf16 = dequant_mxfp8_to_bf16(weight, weight_scale).to(x.dtype)
83 output = torch.nn.functional.linear(x, weight_bf16, bias)
84 return output.to(x.dtype)
vllm/model_executor/kernels/linear/mxfp8/flashinfer.py
-10
@@ -56,8 +56,6 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
56
57 input_shape = x.shape
58 input_2d = x.view(-1, K)
59 - M_orig = input_2d.shape[0]
60 -
59 min_dim = 128
60
61 assert min_dim <= K, (
@@ -72,11 +70,6 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
70 f"out_features is too small for mm_mxfp8."
71 )
72
75 - M_padded = ((M_orig + min_dim - 1) // min_dim) * min_dim
76 - if M_padded != M_orig:
77 - pad_rows = M_padded - M_orig
78 - input_2d = torch.nn.functional.pad(input_2d, (0, 0, 0, pad_rows))
79 -
73 input_mxfp8, input_scale = mxfp8_e4m3_quantize(
74 input_2d, is_sf_swizzled_layout=True
75 )
@@ -93,9 +86,6 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
86 backend="cutlass",
87 )
88
96 - if M_padded != M_orig:
97 - output = output[:M_orig, :]
98 -
89 if bias is not None:
90 output = output + bias
91
vllm/model_executor/kernels/linear/mxfp8/rocm_native.py new
+171
@@ -0,0 +1,171 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Native MXFP8 linear GEMM for AMD CDNA4 (gfx950) via Triton ``tl.dot_scaled``.
4 +
5 +Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16);
6 +activations are MXFP8-quantized per token. Uses the CDNA4 hardware microscaling
7 +matrix cores. Falls back (via the kernel selector) to the BF16
8 +``EmulationMxfp8LinearKernel`` on archs without native MX or for shapes with
9 +``K % 128 != 0``.
10 +"""
11 +
12 +import torch
13 +from torch.nn.parameter import Parameter
14 +
15 +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
16 + MXFP8_BLOCK_SIZE,
17 + MXFP8_SCALE_DTYPE,
18 + dequant_mxfp8_to_bf16,
19 + mxfp8_e4m3_quantize,
20 +)
21 +from vllm.platforms import current_platform
22 +from vllm.triton_utils import tl, triton
23 +
24 +from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
25 +
26 +
27 +@triton.jit
28 +def _mxfp8_linear_kernel(
29 + x_ptr,
30 + xs_ptr,
31 + w_ptr,
32 + ws_ptr,
33 + out_ptr,
34 + M,
35 + N,
36 + K,
37 + stride_xm,
38 + stride_xk,
39 + stride_xsm,
40 + stride_xsk,
41 + stride_wn,
42 + stride_wk,
43 + stride_wsn,
44 + stride_wsk,
45 + stride_om,
46 + stride_on,
47 + BLOCK_M: tl.constexpr,
48 + BLOCK_N: tl.constexpr,
49 + BLOCK_K: tl.constexpr,
50 +):
51 + pid_m = tl.program_id(0)
52 + pid_n = tl.program_id(1)
53 + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
54 + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
55 + offs_k = tl.arange(0, BLOCK_K)
56 + offs_sk = tl.arange(0, BLOCK_K // 32)
57 + m_mask = offs_m < M
58 + n_mask = offs_n < N
59 +
60 + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk
61 + xs_ptrs = xs_ptr + offs_m[:, None] * stride_xsm + offs_sk[None, :] * stride_xsk
62 + w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk
63 + ws_ptrs = ws_ptr + offs_n[:, None] * stride_wsn + offs_sk[None, :] * stride_wsk
64 +
65 + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
66 + for _ in range(0, tl.cdiv(K, BLOCK_K)):
67 + x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0)
68 + w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0)
69 + xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=0)
70 + ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0)
71 + acc += tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3")
72 + x_ptrs += BLOCK_K * stride_xk
73 + w_ptrs += BLOCK_K * stride_wk
74 + xs_ptrs += (BLOCK_K // 32) * stride_xsk
75 + ws_ptrs += (BLOCK_K // 32) * stride_wsk
76 +
77 + o_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
78 + tl.store(
79 + o_ptrs, acc.to(out_ptr.dtype.element_ty), mask=m_mask[:, None] & n_mask[None, :]
80 + )
81 +
82 +
83 +def _mxfp8_dot_scaled_linear(
84 + x: torch.Tensor, # [M, K] bf16/fp16
85 + w: torch.Tensor, # [N, K] fp8 e4m3
86 + w_scale: torch.Tensor, # [N, K//32] uint8 (E8M0)
87 +) -> torch.Tensor:
88 + M, K = x.shape
89 + N = w.shape[0]
90 + x_q, x_scale = mxfp8_e4m3_quantize(x)
91 + out = torch.empty((M, N), dtype=x.dtype, device=x.device)
92 + BLOCK_M, BLOCK_N, BLOCK_K = 64, 128, 128
93 + grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
94 + _mxfp8_linear_kernel[grid](
95 + x_q,
96 + x_scale,
97 + w,
98 + w_scale,
99 + out,
100 + M,
101 + N,
102 + K,
103 + x_q.stride(0),
104 + x_q.stride(1),
105 + x_scale.stride(0),
106 + x_scale.stride(1),
107 + w.stride(0),
108 + w.stride(1),
109 + w_scale.stride(0),
110 + w_scale.stride(1),
111 + out.stride(0),
112 + out.stride(1),
113 + BLOCK_M=BLOCK_M,
114 + BLOCK_N=BLOCK_N,
115 + BLOCK_K=BLOCK_K,
116 + num_warps=8,
117 + )
118 + return out
119 +
120 +
121 +class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
122 + """Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``."""
123 +
124 + @classmethod
125 + def is_supported(
126 + cls, compute_capability: int | None = None
127 + ) -> tuple[bool, str | None]:
128 + if not current_platform.is_rocm():
129 + return False, "not ROCm"
130 + # supports_mx() == gfx95x (CDNA4 native microscaling hardware). On other
131 + # archs dot_scaled would upcast to BF16, so the kernel selector falls
132 + # through to the BF16 emulation (hipBLASLt) path instead.
133 + if not current_platform.supports_mx():
134 + return False, "native MX requires CDNA4 (gfx95x)"
135 + return True, None
136 +
137 + @classmethod
138 + def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
139 + return True, None
140 +
141 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
142 + weight = layer.weight.data # [N, K] fp8
143 + N, K = weight.shape
144 + scale_k = K // MXFP8_BLOCK_SIZE
145 + weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
146 + layer.weight = Parameter(weight.contiguous(), requires_grad=False)
147 + layer.weight_scale = Parameter(weight_scale, requires_grad=False)
148 +
149 + def apply_weights(
150 + self,
151 + layer: torch.nn.Module,
152 + x: torch.Tensor,
153 + bias: torch.Tensor | None = None,
154 + ) -> torch.Tensor:
155 + if layer.weight_scale.dtype != MXFP8_SCALE_DTYPE:
156 + raise ValueError(
157 + f"Expected {MXFP8_SCALE_DTYPE} weight_scale, got "
158 + f"{layer.weight_scale.dtype}."
159 + )
160 + out_shape = (*x.shape[:-1], layer.weight.shape[0])
161 + x2d = x.reshape(-1, x.shape[-1])
162 + if x2d.shape[-1] % 128 == 0:
163 + out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale)
164 + else:
165 + # dot_scaled tiling needs K % 128 == 0; dequantize fallback otherwise.
166 + w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale)
167 + out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype)
168 + out = out.reshape(out_shape)
169 + if bias is not None:
170 + out = out + bias
171 + return out
vllm/model_executor/layers/activation.py
+22 -5
@@ -158,17 +158,28 @@ class SiluAndMulWithClamp(CustomOp):
158 Computes:
159 gate = clamp(x[..., :d], max=swiglu_limit)
160 up = clamp(x[..., d:], min=-swiglu_limit, max=swiglu_limit)
161 - out = silu(gate) * up
162 - where d = x.shape[-1] // 2.
161 + out = gate * sigmoid(alpha * gate) * (up + beta)
162 + where d = x.shape[-1] // 2. The defaults alpha=1.0, beta=0.0 reduce this to
163 + ``silu(gate) * up``; SwiGLU-OAI style models pass alpha (sigmoid scale) and
164 + beta=1.0 (up bias).
165
166 Shapes:
167 x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
168 return: (num_tokens, d) or (batch_size, seq_len, d)
169 """
170
169 - def __init__(self, swiglu_limit: float, *, compile_native: bool = True):
171 + def __init__(
172 + self,
173 + swiglu_limit: float,
174 + alpha: float = 1.0,
175 + beta: float = 0.0,
176 + *,
177 + compile_native: bool = True,
178 + ):
179 super().__init__(compile_native=compile_native)
180 self.swiglu_limit = float(swiglu_limit)
181 + self.alpha = float(alpha)
182 + self.beta = float(beta)
183 if current_platform.is_rocm() or current_platform.is_xpu():
184 self._forward_method = self.forward_native
185 elif current_platform.is_cuda_alike():
@@ -180,18 +191,24 @@ class SiluAndMulWithClamp(CustomOp):
191 d = x.shape[-1] // 2
192 gate = torch.clamp(x[..., :d], max=self.swiglu_limit)
193 up = torch.clamp(x[..., d:], min=-self.swiglu_limit, max=self.swiglu_limit)
183 - return F.silu(gate) * up
194 + return gate * torch.sigmoid(self.alpha * gate) * (up + self.beta)
195
196 def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
197 d = x.shape[-1] // 2
198 output_shape = x.shape[:-1] + (d,)
199 out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
189 - self.op(out, x, self.swiglu_limit)
200 + self.op(out, x, self.swiglu_limit, self.alpha, self.beta)
201 return out
202
203 def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
204 return self.forward_native(x)
205
206 + def extra_repr(self) -> str:
207 + return (
208 + f"swiglu_limit={self.swiglu_limit!r}, "
209 + f"alpha={self.alpha!r}, beta={self.beta!r}"
210 + )
211 +
212
213 # --8<-- [start:mul_and_silu]
214 @CustomOp.register("mul_and_silu")
vllm/model_executor/layers/attention/attention.py
+2
@@ -7,6 +7,7 @@ import torch
7 import torch.nn as nn
8
9 import vllm.envs as envs
10 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
11 from vllm.config import CacheConfig, get_current_vllm_config
12 from vllm.config.vllm import VllmConfig
13 from vllm.forward_context import ForwardContext, get_forward_context
@@ -730,6 +731,7 @@ direct_register_custom_op(
731 )
732
733
734 +@eager_break_during_capture
735 @maybe_transfer_kv_layer
736 def unified_attention_with_output(
737 query: torch.Tensor,
vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py new
+143
@@ -0,0 +1,143 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Manual fusion of tensor-parallel all-reduce with the following GemmaRMSNorm.
4 +
5 +Under tensor parallelism a ``RowParallelLinear`` (e.g. attention ``o_proj``)
6 +produces a per-rank partial sum that is all-reduced, and the result is then fed
7 +into a ``GemmaRMSNorm`` that adds the residual and normalizes. flashinfer ships a
8 +kernel that fuses all-reduce + residual-add + RMSNorm into a single launch; this
9 +helper drives it directly (no torch.compile pass) for models that run eager.
10 +
11 +Scope: attention output only, no quantization. When the flashinfer fast path is
12 +not applicable (TP==1, flashinfer/NVSwitch unavailable, unsupported dtype, or an
13 +oversize batch) it falls back to ``all_reduce`` + ``GemmaRMSNorm``, which is
14 +numerically identical to the unfused model path.
15 +"""
16 +
17 +import torch
18 +
19 +from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
20 +from vllm.distributed.parallel_state import (
21 + get_tensor_model_parallel_rank,
22 + get_tensor_model_parallel_world_size,
23 + get_tp_group,
24 +)
25 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm
26 +
27 +MiB = 1024 * 1024
28 +
29 +# flashinfer fused all-reduce + RMSNorm is wired as a registered custom op in
30 +# allreduce_rms_fusion; both that op and the workspace helpers only exist when
31 +# flashinfer.comm.allreduce_fusion is importable.
32 +try:
33 + from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
34 + flashinfer_trtllm_fused_allreduce_norm,
35 + )
36 + from vllm.distributed.device_communicators.flashinfer_all_reduce import (
37 + flashinfer_comm,
38 + get_fi_ar_workspace,
39 + )
40 +
41 + _AR_RESIDUAL_RMS_NORM = (
42 + flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm
43 + if flashinfer_comm is not None
44 + else None
45 + )
46 +except ImportError:
47 + flashinfer_trtllm_fused_allreduce_norm = None # type: ignore[assignment]
48 + get_fi_ar_workspace = None # type: ignore[assignment]
49 + _AR_RESIDUAL_RMS_NORM = None
50 +
51 +
52 +_FI_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16)
53 +
54 +
55 +def _max_token_num(tp_size: int, hidden_size: int, dtype: torch.dtype) -> int | None:
56 + """Workspace token budget for flashinfer fused all-reduce, or None if the
57 + current world size / device is unsupported. Mirrors ``FlashInferAllReduce``."""
58 + from vllm.config.compilation import PassConfig
59 +
60 + max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(tp_size)
61 + if not max_size_mb:
62 + return None
63 + element_size = torch.tensor([], dtype=dtype).element_size()
64 + return int(max_size_mb * MiB) // (hidden_size * element_size)
65 +
66 +
67 +def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool, int]:
68 + """Whether the flashinfer fused path applies; returns (ok, max_token_num)."""
69 + if (
70 + flashinfer_trtllm_fused_allreduce_norm is None
71 + or get_fi_ar_workspace is None
72 + or _AR_RESIDUAL_RMS_NORM is None
73 + ):
74 + return False, 0
75 + if (
76 + not hidden_states.is_cuda
77 + or hidden_states.dim() != 2
78 + or not hidden_states.is_contiguous()
79 + or hidden_states.dtype not in _FI_SUPPORTED_DTYPES
80 + ):
81 + return False, 0
82 +
83 + num_tokens, hidden_size = hidden_states.shape
84 + max_token_num = _max_token_num(tp_size, hidden_size, hidden_states.dtype)
85 + if max_token_num is None or num_tokens > max_token_num:
86 + return False, 0
87 +
88 + # Lazily create / fetch the (globally cached) workspace; returns None on
89 + # GPUs without NVSwitch, in which case we fall back gracefully.
90 + workspace = get_fi_ar_workspace(
91 + world_size=tp_size,
92 + rank=get_tensor_model_parallel_rank(),
93 + max_token_num=max_token_num,
94 + hidden_dim=hidden_size,
95 + dtype=hidden_states.dtype,
96 + group=get_tp_group().device_group,
97 + )
98 + if workspace is None:
99 + return False, 0
100 + return True, max_token_num
101 +
102 +
103 +def fused_allreduce_gemma_rms_norm(
104 + hidden_states: torch.Tensor,
105 + residual: torch.Tensor,
106 + norm: GemmaRMSNorm,
107 +) -> tuple[torch.Tensor, torch.Tensor]:
108 + """All-reduce ``hidden_states`` + add ``residual`` + GemmaRMSNorm, fused.
109 +
110 + ``hidden_states`` is the per-rank *partial* (un-reduced) output of a
111 + row-parallel linear; ``norm`` is the GemmaRMSNorm applied right after.
112 + Returns ``(normed_output, new_residual)``, equivalent to
113 + ``norm(all_reduce(hidden_states), residual)``.
114 + """
115 + tp_size = get_tensor_model_parallel_world_size()
116 + if tp_size == 1:
117 + # No all-reduce needed; identical to the unfused path.
118 + return norm(hidden_states, residual)
119 +
120 + ok, max_token_num = _can_use_flashinfer(hidden_states, tp_size)
121 + if ok:
122 + norm_out = torch.empty_like(hidden_states)
123 + # With norm_out provided, the kernel writes the new residual
124 + # (all_reduce(hidden_states) + residual) into the hidden_states buffer
125 + # and the normalized result into norm_out, leaving `residual` untouched.
126 + flashinfer_trtllm_fused_allreduce_norm(
127 + allreduce_in=hidden_states,
128 + residual=residual,
129 + rms_gamma=norm.weight,
130 + rms_eps=norm.variance_epsilon,
131 + world_size=tp_size,
132 + weight_bias=1.0, # GemmaRMSNorm-style
133 + launch_with_pdl=True,
134 + fp32_acc=True,
135 + max_token_num=max_token_num,
136 + pattern_code=_AR_RESIDUAL_RMS_NORM,
137 + norm_out=norm_out,
138 + )
139 + return norm_out, hidden_states
140 +
141 + # Fallback: explicit all-reduce + GemmaRMSNorm (matches the unfused model).
142 + reduced = tensor_model_parallel_all_reduce(hidden_states)
143 + return norm(reduced, residual)
vllm/model_executor/layers/fused_moe/activation.py
+25 -2
@@ -17,7 +17,12 @@ class MoEActivation(Enum):
17 GELU = "gelu"
18 GELU_TANH = "gelu_tanh"
19 RELU2 = "relu2"
20 + # SWIGLUOAI expects gate/up *interleaved* in w13 ([gate0, up0, gate1, ...]),
21 + # as in gpt-oss checkpoints. SWIGLUOAI_UNINTERLEAVE has identical math but
22 + # expects the *packed* layout ([all gates; all ups]), as produced by a
23 + # MergedColumnParallelLinear gate_up_proj (e.g. MiniMax-M3).
24 SWIGLUOAI = "swigluoai"
25 + SWIGLUOAI_UNINTERLEAVE = "swigluoai_uninterleave"
26 SWIGLUSTEP = "swiglustep"
27
28 # Non-gated activations (no mul with gate) expect input of shape [..., d]
@@ -73,6 +78,7 @@ _CUSTOM_OP_NAMES: dict[MoEActivation, str] = {
78 MoEActivation.GELU: "gelu_and_mul",
79 MoEActivation.GELU_TANH: "gelu_tanh_and_mul",
80 MoEActivation.SWIGLUOAI: "swigluoai_and_mul",
81 + MoEActivation.SWIGLUOAI_UNINTERLEAVE: "silu_and_mul_with_clamp",
82 MoEActivation.SWIGLUSTEP: "swiglustep_and_mul",
83 MoEActivation.RELU2: "relu2",
84 MoEActivation.SILU_NO_MUL: "silu_and_mul",
@@ -105,8 +111,17 @@ def apply_moe_activation(
111 activation: MoEActivation,
112 output: torch.Tensor,
113 input: torch.Tensor,
114 + *,
115 + clamp_limit: float | None = None,
116 + alpha: float = 1.0,
117 + beta: float = 0.0,
118 ) -> torch.Tensor:
109 - """Apply MoE activation function."""
119 + """Apply MoE activation function.
120 +
121 + ``clamp_limit``/``alpha``/``beta`` (from the quant config) drive the clamped
122 + SwiGLU kernels: ``SILU`` + ``clamp_limit`` and ``SWIGLUOAI_UNINTERLEAVE`` both
123 + map to ``silu_and_mul_with_clamp``. Other activations ignore them.
124 + """
125 assert input.dim() == 2, "Input must be 2D"
126 assert output.dim() == 2, "Output must be 2D"
127 if activation.is_gated:
@@ -122,13 +137,21 @@ def apply_moe_activation(
137
138 # Activations with gated multiplication (gate × activation(up))
139 if activation == MoEActivation.SILU:
125 - torch.ops._C.silu_and_mul(output, input)
140 + if clamp_limit is not None:
141 + # Fused silu(clamp(gate)) * clamp(up); equivalent to swiglu_limit_func.
142 + torch.ops._C.silu_and_mul_with_clamp(output, input, clamp_limit, 1.0, 0.0)
143 + else:
144 + torch.ops._C.silu_and_mul(output, input)
145 elif activation == MoEActivation.GELU:
146 torch.ops._C.gelu_and_mul(output, input)
147 elif activation == MoEActivation.GELU_TANH:
148 torch.ops._C.gelu_tanh_and_mul(output, input)
149 elif activation == MoEActivation.SWIGLUOAI:
150 torch.ops._C.swigluoai_and_mul(output, input)
151 + elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
152 + # SwiGLU-OAI on packed w13 (gate = first half, up = second half).
153 + assert clamp_limit is not None, "SWIGLUOAI_UNINTERLEAVE requires clamp_limit"
154 + torch.ops._C.silu_and_mul_with_clamp(output, input, clamp_limit, alpha, beta)
155 elif activation == MoEActivation.SWIGLUSTEP:
156 from vllm.model_executor.layers.activation import swiglustep_and_mul_triton
157
vllm/model_executor/layers/fused_moe/config.py
+17
@@ -900,6 +900,9 @@ def fp8_w8a16_moe_quant_config(
900 w1_bias: torch.Tensor | None = None,
901 w2_bias: torch.Tensor | None = None,
902 block_shape: list[int] | None = None,
903 + gemm1_alpha: float | None = None,
904 + gemm1_beta: float | None = None,
905 + gemm1_clamp_limit: float | None = None,
906 ) -> FusedMoEQuantConfig:
907 """
908 Construct a quant config for 16-bit float activations and fp8 weights.
@@ -925,6 +928,9 @@ def fp8_w8a16_moe_quant_config(
928 None,
929 w2_bias,
930 ),
931 + gemm1_alpha=gemm1_alpha,
932 + gemm1_beta=gemm1_beta,
933 + gemm1_clamp_limit=gemm1_clamp_limit,
934 )
935
936
@@ -979,15 +985,24 @@ def int4_w4afp8_moe_quant_config(
985 def biased_moe_quant_config(
986 w1_bias: torch.Tensor | None,
987 w2_bias: torch.Tensor | None,
988 + gemm1_alpha: float | None = None,
989 + gemm1_beta: float | None = None,
990 + gemm1_clamp_limit: float | None = None,
991 ) -> FusedMoEQuantConfig:
992 """
993 Construct a quant config for unquantized activations with biases.
994 +
995 + gemm1_alpha/gemm1_beta/gemm1_clamp_limit carry the SwiGLU gate params
996 + through to the fused activation kernel (e.g. swigluoai_uninterleave).
997 """
998 return FusedMoEQuantConfig(
999 _a1=FusedMoEQuantDesc(),
1000 _a2=FusedMoEQuantDesc(),
1001 _w1=FusedMoEQuantDesc(bias=w1_bias),
1002 _w2=FusedMoEQuantDesc(bias=w2_bias),
1003 + gemm1_alpha=gemm1_alpha,
1004 + gemm1_beta=gemm1_beta,
1005 + gemm1_clamp_limit=gemm1_clamp_limit,
1006 )
1007
1008
@@ -1268,6 +1283,8 @@ class FusedMoEConfig:
1283 # are filtered out by `FusedMoEExperts.is_supported_config` so the oracle
1284 # cannot silently select one and drop the clamp.
1285 swiglu_limit: float | None = None
1286 + swiglu_alpha: float | None = None
1287 + swiglu_beta: float | None = None
1288
1289 max_capture_size: int = 0
1290
vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
+86 -14
@@ -130,6 +130,9 @@ def _fwd_kernel_ep_scatter_2(
130 HIDDEN_SIZE_PAD: tl.constexpr,
131 SCALE_HIDDEN_SIZE: tl.constexpr,
132 SCALE_HIDDEN_SIZE_PAD: tl.constexpr,
133 + PACK_UE8M0: tl.constexpr,
134 + SCALE_PACKED_SIZE: tl.constexpr,
135 + SCALE_PACKED_SIZE_PAD: tl.constexpr,
136 ):
137 start_token_id = tl.program_id(0)
138 grid_num = tl.num_programs(0)
@@ -137,16 +140,47 @@ def _fwd_kernel_ep_scatter_2(
140 offset_in = tl.arange(0, HIDDEN_SIZE_PAD)
141 mask = offset_in < HIDDEN_SIZE
142
140 - offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
141 - mask_s = offset_in_s < SCALE_HIDDEN_SIZE
142 -
143 output_tensor_stride0 = output_tensor_stride0.to(tl.int64)
144
145 + if PACK_UE8M0:
146 + # One int32 per 4 consecutive 32-wide UE8M0 groups, stored MN-major.
147 + offs_pk = tl.arange(0, SCALE_PACKED_SIZE_PAD)
148 + mask_pk = offs_pk < SCALE_PACKED_SIZE
149 + else:
150 + offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
151 + mask_s = offset_in_s < SCALE_HIDDEN_SIZE
152 +
153 for token_id in range(start_token_id, total_token_num, grid_num):
154 to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask)
147 - to_copy_s = tl.load(
148 - recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s, mask=mask_s
149 - )
155 +
156 + if PACK_UE8M0:
157 + # Pack 4 UE8M0 bytes into one int32 (byte j = group 4*pk+j).
158 + base_s = recv_x_scale + token_id * recv_x_scale_stride0
159 + g0, g1 = offs_pk * 4, offs_pk * 4 + 1
160 + g2, g3 = offs_pk * 4 + 2, offs_pk * 4 + 3
161 + b0 = tl.load(
162 + base_s + g0 * recv_x_scale_stride1, mask=g0 < SCALE_HIDDEN_SIZE
163 + )
164 + b1 = tl.load(
165 + base_s + g1 * recv_x_scale_stride1, mask=g1 < SCALE_HIDDEN_SIZE
166 + )
167 + b2 = tl.load(
168 + base_s + g2 * recv_x_scale_stride1, mask=g2 < SCALE_HIDDEN_SIZE
169 + )
170 + b3 = tl.load(
171 + base_s + g3 * recv_x_scale_stride1, mask=g3 < SCALE_HIDDEN_SIZE
172 + )
173 + packed_s = (
174 + b0.to(tl.int32)
175 + | (b1.to(tl.int32) << 8)
176 + | (b2.to(tl.int32) << 16)
177 + | (b3.to(tl.int32) << 24)
178 + )
179 + else:
180 + to_copy_s = tl.load(
181 + recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s,
182 + mask=mask_s,
183 + )
184
185 for topk_index in tl.range(0, topk_num, 1, num_stages=4):
186 expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index)
@@ -164,11 +198,21 @@ def _fwd_kernel_ep_scatter_2(
198 output_tensor_ptr = (
199 output_tensor + dest_token_index_i64 * output_tensor_stride0
200 )
201 + tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
202 +
203 output_tensor_scale_ptr = (
204 output_tensor_scale + dest_token_index * output_tensor_scale_stride0
205 )
170 - tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
171 - tl.store(output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s)
206 + if PACK_UE8M0:
207 + tl.store(
208 + output_tensor_scale_ptr + offs_pk * output_tensor_scale_stride1,
209 + packed_s,
210 + mask=mask_pk,
211 + )
212 + else:
213 + tl.store(
214 + output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s
215 + )
216
217
218 @torch.no_grad()
@@ -183,9 +227,11 @@ def ep_scatter(
227 output_tensor_scale: torch.Tensor,
228 m_indices: torch.Tensor,
229 output_index: torch.Tensor,
230 + block_size: int = 128,
231 + pack_ue8m0: bool = False,
232 ):
233 BLOCK_E = 128 # token num of per expert is aligned to 128
188 - BLOCK_D = 128 # block size of quantization
234 + BLOCK_D = block_size # block size of activation-scale quantization
235 num_warps = 8
236 num_experts = num_recv_tokens_per_expert.shape[0]
237 hidden_size = recv_x.shape[1]
@@ -195,6 +241,10 @@ def ep_scatter(
241 assert m_indices.shape[0] % BLOCK_E == 0
242 assert expert_start_loc.shape[0] == num_experts
243
244 + # pack_ue8m0: scatter packs 4 UE8M0 bytes per int32; else copies scales as-is.
245 + scale_hidden_size = hidden_size // BLOCK_D
246 + scale_packed_size = (scale_hidden_size + 3) // 4 if pack_ue8m0 else 1
247 +
248 _fwd_kernel_ep_scatter_1[(grid,)](
249 num_recv_tokens_per_expert,
250 expert_start_loc,
@@ -234,8 +284,11 @@ def ep_scatter(
284 num_warps=num_warps,
285 HIDDEN_SIZE=hidden_size,
286 HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size),
237 - SCALE_HIDDEN_SIZE=hidden_size // BLOCK_D,
238 - SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size // BLOCK_D),
287 + SCALE_HIDDEN_SIZE=scale_hidden_size,
288 + SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size),
289 + PACK_UE8M0=pack_ue8m0,
290 + SCALE_PACKED_SIZE=scale_packed_size,
291 + SCALE_PACKED_SIZE_PAD=triton.next_power_of_2(scale_packed_size),
292 )
293 return
294
@@ -352,6 +405,7 @@ def deepgemm_moe_permute(
405 expert_map: torch.Tensor | None,
406 expert_tokens_meta: mk.ExpertTokensMetadata | None,
407 aq_out: torch.Tensor | None = None,
408 + block_size: int | None = None,
409 ):
410 assert aq.ndim == 2
411 assert topk_ids.dtype.is_signed, "The kernel uses -1 to represent invalid topk_ids"
@@ -359,6 +413,10 @@ def deepgemm_moe_permute(
413 device = aq.device
414
415 block_m, block_k = get_mk_alignment_for_contiguous_layout()
416 + # The activation-scale group size may differ from the M/K tile alignment
417 + # (e.g. MXFP8 uses a 32-element scale group while block_k stays 128).
418 + if block_size is not None:
419 + block_k = block_size
420
421 M_sum = compute_aligned_M(
422 M=topk_ids.size(0),
@@ -376,9 +434,21 @@ def deepgemm_moe_permute(
434 if aq_out is None:
435 aq_out = torch.empty((M_sum, H), device=device, dtype=aq.dtype)
436
379 - aq_scale_out = torch.empty(
380 - (M_sum, H // block_k), device=device, dtype=torch.float32
381 - )
437 + # uint8 UE8M0 (MXFP8) -> scatter packs into DeepGEMM's int32 MN-major
438 + # TMA-aligned layout; float32 (FP8/FP4) scattered row-major as-is.
439 + pack_ue8m0 = aq_scale.dtype == torch.uint8
440 + sf_k = H // block_k
441 + if pack_ue8m0:
442 + packed_sf_k = (sf_k + 3) // 4
443 + tma_aligned_mn = round_up(M_sum, 4)
444 + aq_scale_out = torch.empty_strided(
445 + (M_sum, packed_sf_k),
446 + (1, tma_aligned_mn),
447 + device=device,
448 + dtype=torch.int32,
449 + )
450 + else:
451 + aq_scale_out = torch.empty((M_sum, sf_k), device=device, dtype=torch.float32)
452
453 # DeepGEMM uses negative values in m_indices (here expert_ids) to mark
454 # completely invalid / padded blocks that should be skipped. We always
@@ -412,6 +482,8 @@ def deepgemm_moe_permute(
482 output_tensor_scale=aq_scale_out,
483 m_indices=expert_ids,
484 output_index=inv_perm,
485 + block_size=block_k,
486 + pack_ue8m0=pack_ue8m0,
487 )
488
489 return aq_out, aq_scale_out, expert_ids, inv_perm
vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py
+68 -14
@@ -33,7 +33,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
33 kFp8Dynamic128Sym,
34 kFp8Static128BlockSym,
35 kMxfp4Static,
36 + kMxfp8Dynamic,
37 + kMxfp8Static,
38 )
39 +from vllm.platforms import current_platform
40 from vllm.utils.deep_gemm import (
41 DeepGemmQuantScaleFMT,
42 get_mk_alignment_for_contiguous_layout,
@@ -123,12 +126,26 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
126
127 def __init__(self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig):
128 super().__init__(moe_config=moe_config, quant_config=quant_config)
126 - assert quant_config.block_shape == get_mk_alignment_for_contiguous_layout()
127 - assert quant_config.quant_dtype == torch.float8_e4m3fn
129 + # MXFP8: FP8 e4m3 values + UE8M0 1x32 block scales (Blackwell). Reuses
130 + # the same grouped GEMM (aliased to fp8_fp4) with recipe (1, 32).
131 + self.mxfp8 = quant_config.block_shape == [1, 32]
132 + if self.mxfp8:
133 + assert quant_config.quant_dtype == "mxfp8"
134 + else:
135 + assert quant_config.block_shape == get_mk_alignment_for_contiguous_layout()
136 + assert quant_config.quant_dtype == torch.float8_e4m3fn
137 assert not quant_config.per_act_token_quant
138 assert not quant_config.per_out_ch_quant
139
140 self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
141 + # Gated-activation params: silu == swigluoai with alpha=1, beta=0.
142 + # FP8 (silu) configs leave these None, reproducing plain silu.
143 + self.gemm1_alpha = (
144 + quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
145 + )
146 + self.gemm1_beta = (
147 + quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0
148 + )
149
150 @staticmethod
151 def activation_format() -> mk.FusedMoEActivationFormat:
@@ -147,14 +164,25 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
164 weight_key: QuantKey | None,
165 activation_key: QuantKey | None,
166 ) -> bool:
150 - SUPPORTED_W_A = [
151 - (kFp8Static128BlockSym, kFp8Dynamic128Sym),
152 - ]
153 - return (weight_key, activation_key) in SUPPORTED_W_A
167 + if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym):
168 + return True
169 + # MXFP8 1x32 uses the fp8_fp4 grouped GEMM with recipe (1, 32) — only
170 + # available on Blackwell (SM100).
171 + if (weight_key, activation_key) == (kMxfp8Static, kMxfp8Dynamic):
172 + return current_platform.is_device_capability_family(100)
173 + return False
174
175 @staticmethod
176 def _supports_activation(activation: MoEActivation) -> bool:
157 - return activation in [MoEActivation.SILU, MoEActivation.SWIGLUSTEP]
177 + # silu/swigluoai go through the fused alpha/beta kernel; swiglustep
178 + # uses the unfused activation path. The fused kernel reads packed w13
179 + # (gate = first half, up = second half), so it implements the
180 + # *uninterleaved* SwiGLU-OAI variant.
181 + return activation in [
182 + MoEActivation.SILU,
183 + MoEActivation.SWIGLUSTEP,
184 + MoEActivation.SWIGLUOAI_UNINTERLEAVE,
185 + ]
186
187 @staticmethod
188 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -179,7 +207,9 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
207 activation: MoEActivation,
208 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
209 assert self.block_shape is not None
182 - block_m = self.block_shape[0]
210 + # Use the contiguous-layout M alignment (matches apply()); block_shape[0]
211 + # is the quant block (1 for MXFP8) and would under-size the workspace.
212 + block_m = get_mk_alignment_for_contiguous_layout()[0]
213 M_sum = compute_aligned_M(
214 M, topk, local_num_experts, block_m, expert_tokens_meta
215 )
@@ -201,14 +231,24 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
231 M_sum, N = input.size()
232 activation_out_dim = self.adjust_N_for_activation(N, activation)
233
204 - # 1. DeepGemm UE8M0: fused SiLU+mul+clamp+quant+pack
234 + # silu and swigluoai are both expressible by the fused gated kernel via
235 + # (alpha, beta): silu uses alpha=1, beta=0; swigluoai uses config values.
236 + # The fused kernel reads packed w13, hence SWIGLUOAI_UNINTERLEAVE.
237 + fused_gated = activation in (
238 + MoEActivation.SILU,
239 + MoEActivation.SWIGLUOAI_UNINTERLEAVE,
240 + )
241 +
242 + # 1. DeepGemm UE8M0: fused gate+mul+clamp+quant+pack
243 if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
206 - if activation == MoEActivation.SILU:
244 + if fused_gated:
245 return fused_silu_mul_fp8_quant_packed(
246 input=input,
247 output_q=output,
248 group_size=block_k,
249 clamp_limit=self.gemm1_clamp_limit,
250 + alpha=self.gemm1_alpha,
251 + beta=self.gemm1_beta,
252 )
253 act_out = torch.empty(
254 (M_sum, activation_out_dim), dtype=input.dtype, device=input.device
@@ -221,14 +261,17 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
261 )
262 return a2q, a2q_scale
263
224 - # 2. Hopper / non‑E8M0: prefer the fused SiLU+mul+quant kernel
225 - if activation == MoEActivation.SILU:
264 + # 2. Hopper / non‑E8M0: prefer the fused gate+mul+quant kernel
265 + if fused_gated:
266 use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0
267 return silu_mul_per_token_group_quant_fp8_colmajor(
268 input=input,
269 output=output,
270 use_ue8m0=use_ue8m0,
271 clamp_limit=self.gemm1_clamp_limit,
272 + group_size=block_k,
273 + alpha=self.gemm1_alpha,
274 + beta=self.gemm1_beta,
275 )
276
277 # 3. fallback path for non-SiLU activations in non‑UE8M0 cases.
@@ -292,12 +335,23 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
335 expert_map=expert_map,
336 expert_tokens_meta=expert_tokens_meta,
337 aq_out=a1q_perm,
338 + # MXFP8 uses a 32-element activation-scale group (block_shape[1]);
339 + # FP8-block keeps the default (128) alignment.
340 + block_size=self.block_shape[1] if self.mxfp8 else None,
341 )
342 assert a1q.size(0) == M_sum
343
344 + # MXFP8 (1x32) drives the fp8_fp4-aliased grouped GEMM with recipe
345 + # (1, 32); the FP8 block path keeps the default (128) recipe.
346 + gemm_kwargs = (
347 + {"recipe_a": (1, self.block_shape[1]), "recipe_b": (1, self.block_shape[1])}
348 + if self.mxfp8
349 + else {}
350 + )
351 +
352 mm1_out = _resize_cache(workspace2, (M_sum, N))
353 m_grouped_fp8_gemm_nt_contiguous(
300 - (a1q, a1q_scale), (w1, self.w1_scale), mm1_out, expert_ids
354 + (a1q, a1q_scale), (w1, self.w1_scale), mm1_out, expert_ids, **gemm_kwargs
355 )
356
357 activation_out_dim = self.adjust_N_for_activation(N, activation)
@@ -310,7 +364,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
364
365 mm2_out = _resize_cache(workspace2, (M_sum, K))
366 m_grouped_fp8_gemm_nt_contiguous(
313 - (a2q, a2q_scale), (w2, self.w2_scale), mm2_out, expert_ids
367 + (a2q, a2q_scale), (w2, self.w2_scale), mm2_out, expert_ids, **gemm_kwargs
368 )
369
370 if apply_router_weight_on_input:
vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py
+5 -1
@@ -801,7 +801,11 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
801 return TopKWeightAndReduceDelegate()
802
803 def activation(
804 - self, activation: MoEActivation, output: torch.Tensor, input: torch.Tensor
804 + self,
805 + activation: MoEActivation,
806 + output: torch.Tensor,
807 + input: torch.Tensor,
808 + **kwargs,
809 ) -> None:
810 gemm1_clamp_limit = self.quant_config.gemm1_clamp_limit
811 if activation == MoEActivation.SILU and gemm1_clamp_limit is not None:
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+1
@@ -787,6 +787,7 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
787 activation: MoEActivation,
788 output: torch.Tensor,
789 input: torch.Tensor,
790 + **kwargs,
791 ) -> None:
792 quant_config = self.quant_config or FUSED_MOE_UNQUANTIZED_CONFIG
793 if activation == MoEActivation.SWIGLUOAI:
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+51 -23
@@ -28,10 +28,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
28 TopKWeightAndReduceDelegate,
29 TopKWeightAndReduceNoOP,
30 )
31 -from vllm.model_executor.layers.fused_moe.utils import (
32 - _resize_cache,
33 - swiglu_limit_func,
34 -)
31 +from vllm.model_executor.layers.fused_moe.utils import _resize_cache
32 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
33 get_marlin_input_dtype,
34 marlin_make_workspace_new,
@@ -74,9 +71,7 @@ def _fused_marlin_moe(
71 expert_ids: torch.Tensor,
72 num_tokens_post_padded: torch.Tensor,
73 activation: MoEActivation = MoEActivation.SILU,
77 - activation_func: Callable[
78 - [MoEActivation, torch.Tensor, torch.Tensor], None
79 - ] = apply_moe_activation,
74 + activation_func: Callable[..., None] = apply_moe_activation,
75 input_global_scale1: torch.Tensor | None = None,
76 input_global_scale2: torch.Tensor | None = None,
77 global_scale1: torch.Tensor | None = None,
@@ -94,6 +89,8 @@ def _fused_marlin_moe(
89 input_dtype: torch.dtype | None = None,
90 is_k_full: bool = True,
91 clamp_limit: float | None = None,
92 + gemm1_alpha: float = 1.0,
93 + gemm1_beta: float = 0.0,
94 ) -> torch.Tensor:
95 assert hidden_states.ndim == 2
96 M, K = hidden_states.size()
@@ -161,18 +158,16 @@ def _fused_marlin_moe(
158 use_fp32_reduce=True,
159 is_zp_float=False,
160 )
164 - if clamp_limit is not None and activation == MoEActivation.SILU:
165 - swiglu_limit_func(
166 - intermediate_cache2,
167 - intermediate_cache1.view(-1, w13_num_shards * N),
168 - clamp_limit,
169 - )
170 - else:
171 - activation_func(
172 - activation,
173 - intermediate_cache2,
174 - intermediate_cache1.view(-1, w13_num_shards * N),
175 - )
161 + # apply_moe_activation fuses the clamp/gate params: SILU + clamp_limit and
162 + # SWIGLUOAI_UNINTERLEAVE both map to the silu_and_mul_with_clamp kernel.
163 + activation_func(
164 + activation,
165 + intermediate_cache2,
166 + intermediate_cache1.view(-1, w13_num_shards * N),
167 + clamp_limit=clamp_limit,
168 + alpha=gemm1_alpha,
169 + beta=gemm1_beta,
170 + )
171
172 if output is None:
173 output = intermediate_cache3
@@ -238,9 +233,7 @@ def fused_marlin_moe(
233 apply_router_weight_on_input: bool = False,
234 global_num_experts: int = -1,
235 activation: MoEActivation = MoEActivation.SILU,
241 - activation_func: Callable[
242 - [MoEActivation, torch.Tensor, torch.Tensor], None
243 - ] = apply_moe_activation,
236 + activation_func: Callable[..., None] = apply_moe_activation,
237 moe_sum: Callable[[torch.Tensor, torch.Tensor], None] | None = None,
238 expert_map: torch.Tensor | None = None,
239 input_global_scale1: torch.Tensor | None = None,
@@ -260,6 +253,8 @@ def fused_marlin_moe(
253 output: torch.Tensor | None = None,
254 input_dtype: torch.dtype | None = None,
255 clamp_limit: float | None = None,
256 + gemm1_alpha: float = 1.0,
257 + gemm1_beta: float = 0.0,
258 ) -> torch.Tensor:
259 """
260 This function computes a Mixture of Experts (MoE) layer using two sets of
@@ -373,6 +368,8 @@ def fused_marlin_moe(
368 input_dtype=input_dtype,
369 is_k_full=is_k_full,
370 clamp_limit=clamp_limit,
371 + gemm1_alpha=gemm1_alpha,
372 + gemm1_beta=gemm1_beta,
373 ).view(-1, topk, K)
374
375 if output is None:
@@ -415,6 +412,8 @@ def batched_fused_marlin_moe(
412 output: torch.Tensor | None = None,
413 input_dtype: torch.dtype | None = None,
414 clamp_limit: float | None = None,
415 + gemm1_alpha: float = 1.0,
416 + gemm1_beta: float = 0.0,
417 ) -> torch.Tensor:
418 """
419 This function massages the inputs so the batched hidden_states can be
@@ -544,6 +543,8 @@ def batched_fused_marlin_moe(
543 input_dtype=input_dtype,
544 is_k_full=is_k_full,
545 clamp_limit=clamp_limit,
546 + gemm1_alpha=gemm1_alpha,
547 + gemm1_beta=gemm1_beta,
548 )
549
550 output = output.view(B, BATCH_TOKENS_MAX, K)
@@ -579,6 +580,15 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
580 self.is_k_full = is_k_full
581 self.input_dtype = get_marlin_input_dtype()
582 self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
583 + # Gated-activation params (used by SWIGLUOAI_UNINTERLEAVE on packed w13).
584 + # silu == swigluoai with alpha=1, beta=0; configs that don't set these
585 + # (plain silu) fall back to the silu identity.
586 + self.gemm1_alpha = (
587 + quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
588 + )
589 + self.gemm1_beta = (
590 + quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0
591 + )
592
593 super().__init__(
594 moe_config=moe_config,
@@ -627,6 +637,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
637 MoEActivation.GELU,
638 MoEActivation.GELU_TANH,
639 MoEActivation.SWIGLUOAI,
640 + MoEActivation.SWIGLUOAI_UNINTERLEAVE,
641 MoEActivation.SWIGLUSTEP,
642 MoEActivation.SILU_NO_MUL,
643 MoEActivation.GELU_NO_MUL,
@@ -787,6 +798,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
798 is_k_full=self.is_k_full,
799 input_dtype=self.input_dtype,
800 clamp_limit=self.gemm1_clamp_limit,
801 + gemm1_alpha=self.gemm1_alpha,
802 + gemm1_beta=self.gemm1_beta,
803 )
804 return
805
@@ -805,6 +818,10 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
818 act_enum: MoEActivation,
819 act_output: torch.Tensor,
820 act_input: torch.Tensor,
821 + *,
822 + clamp_limit: float | None = None,
823 + alpha: float = 1.0,
824 + beta: float = 0.0,
825 ) -> None:
826 # act_input = intermediate_cache1 (M*topk, 2N for gated)
827 # act_output = intermediate_cache2 (M*topk, N)
@@ -834,7 +851,14 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
851 "tlm": token_lora_mapping,
852 }
853 )
837 - self.activation(act_enum, act_output, act_input)
854 + self.activation(
855 + act_enum,
856 + act_output,
857 + act_input,
858 + clamp_limit=clamp_limit,
859 + alpha=alpha,
860 + beta=beta,
861 + )
862 lora_state["cache2"] = act_output
863
864 def moe_sum_with_lora(moe_out: torch.Tensor, out: torch.Tensor) -> None:
@@ -888,6 +912,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
912 is_k_full=self.is_k_full,
913 input_dtype=self.input_dtype,
914 clamp_limit=self.gemm1_clamp_limit,
915 + gemm1_alpha=self.gemm1_alpha,
916 + gemm1_beta=self.gemm1_beta,
917 )
918
919 def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
@@ -996,4 +1022,6 @@ class BatchedMarlinExperts(MarlinExpertsBase):
1022 input_dtype=self.input_dtype,
1023 is_k_full=self.is_k_full,
1024 clamp_limit=self.gemm1_clamp_limit,
1025 + gemm1_alpha=self.gemm1_alpha,
1026 + gemm1_beta=self.gemm1_beta,
1027 )
vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py new
+176
@@ -0,0 +1,176 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MXFP8 (1x32 block, E8M0 scale) MoE experts on Triton.
4 +
5 +``Mxfp8TritonExpertsBase`` stashes E8M0 weight scales for checkpoint layout.
6 +``Mxfp8EmulationTritonExperts`` dequantizes to BF16 and runs ``TritonExperts``
7 +for devices without a native MXFP8 MoE kernel (e.g. ROCm gfx942 / MI300).
8 +"""
9 +
10 +import torch
11 +
12 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13 +from vllm.logger import init_logger
14 +from vllm.model_executor.layers.fused_moe.config import (
15 + FusedMoEConfig,
16 + FusedMoEQuantConfig,
17 +)
18 +from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
19 +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
20 + dequant_mxfp8_to_bf16,
21 +)
22 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
23 + QuantKey,
24 + kMxfp8Dynamic,
25 + kMxfp8Static,
26 +)
27 +
28 +logger = init_logger(__name__)
29 +
30 +
31 +class Mxfp8TritonExpertsBase(TritonExperts):
32 + """Shared MXFP8 MoE setup: stash E8M0 scales, clear scales on ``quant_config``."""
33 +
34 + def __init__(
35 + self,
36 + moe_config: FusedMoEConfig,
37 + quant_config: FusedMoEQuantConfig,
38 + ):
39 + super().__init__(moe_config, quant_config)
40 + self.w1_scale_val = self.quant_config.w1_scale
41 + self.w2_scale_val = self.quant_config.w2_scale
42 + self.quant_config._w1.scale = None
43 + self.quant_config._w2.scale = None
44 +
45 + @staticmethod
46 + def _supports_quant_scheme(
47 + weight_key: QuantKey | None,
48 + activation_key: QuantKey | None,
49 + ) -> bool:
50 + return (weight_key, activation_key) == (kMxfp8Static, kMxfp8Dynamic)
51 +
52 + @staticmethod
53 + def _supports_activation(activation) -> bool:
54 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
55 +
56 + if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
57 + return True
58 + return TritonExperts._supports_activation(activation)
59 +
60 +
61 +class Mxfp8EmulationTritonExperts(Mxfp8TritonExpertsBase):
62 + """Dequantize MXFP8 weights to BF16 on the fly and run ``TritonExperts``."""
63 +
64 + def __init__(
65 + self,
66 + moe_config: FusedMoEConfig,
67 + quant_config: FusedMoEQuantConfig,
68 + ):
69 + super().__init__(moe_config, quant_config)
70 + logger.warning_once(
71 + "Using Mxfp8EmulationTritonExperts MoE backend. Weights are "
72 + "dequantized to BF16 on the fly; this is slower than a native "
73 + "MXFP8 MoE kernel and is intended for devices without one."
74 + )
75 +
76 + @property
77 + def quant_dtype(self) -> torch.dtype | str | None:
78 + # BF16 fallback: do not MXFP8-quantize activations in ``TritonExperts``.
79 + return None
80 +
81 + @property
82 + def block_shape(self) -> list[int] | None:
83 + return None
84 +
85 + @property
86 + def expects_unquantized_inputs(self) -> bool:
87 + return True
88 +
89 + @staticmethod
90 + def _supports_current_device() -> bool:
91 + return True
92 +
93 + def activation(
94 + self,
95 + activation,
96 + output: torch.Tensor,
97 + input: torch.Tensor,
98 + **kwargs,
99 + ):
100 + """Apply GEMM1 activation with quant-config alpha/beta/clamp."""
101 + from vllm.model_executor.layers.fused_moe.activation import (
102 + MoEActivation,
103 + apply_moe_activation,
104 + )
105 +
106 + if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
107 + limit = self.quant_config.gemm1_clamp_limit
108 + if limit is None:
109 + raise ValueError("SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit")
110 + alpha = self.quant_config.gemm1_alpha
111 + alpha = 1.702 if alpha is None else float(alpha)
112 + beta = self.quant_config.gemm1_beta
113 + beta = 1.0 if beta is None else float(beta)
114 + apply_moe_activation(
115 + activation,
116 + output,
117 + input,
118 + clamp_limit=float(limit),
119 + alpha=alpha,
120 + beta=beta,
121 + )
122 + return
123 + super().activation(activation, output, input)
124 +
125 + def apply(
126 + self,
127 + output: torch.Tensor,
128 + hidden_states: torch.Tensor,
129 + w1: torch.Tensor,
130 + w2: torch.Tensor,
131 + topk_weights: torch.Tensor,
132 + topk_ids: torch.Tensor,
133 + activation,
134 + global_num_experts: int,
135 + expert_map: torch.Tensor | None,
136 + a1q_scale: torch.Tensor | None,
137 + a2_scale: torch.Tensor | None,
138 + workspace13: torch.Tensor,
139 + workspace2: torch.Tensor,
140 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
141 + apply_router_weight_on_input: bool,
142 + ):
143 + # If the weights were already dequantized to BF16 at load time
144 + # (process_weights_after_loading on devices without a native MXFP8 MoE
145 + # kernel), use them directly -- no per-step dequant. MXFP8 weights are
146 + # 1-byte FP8 (element_size 1); BF16/FP16 are >= 2 bytes.
147 + if w1.element_size() >= 2:
148 + # tl.dot requires w and activations share a dtype; .to() is a no-op
149 + # when they already match (e.g. both BF16).
150 + w1_bf16 = w1.to(hidden_states.dtype)
151 + w2_bf16 = w2.to(hidden_states.dtype)
152 + else:
153 + w1_bf16 = dequant_mxfp8_to_bf16(w1, self.w1_scale_val).to(
154 + hidden_states.dtype
155 + )
156 + w2_bf16 = dequant_mxfp8_to_bf16(w2, self.w2_scale_val).to(
157 + hidden_states.dtype
158 + )
159 +
160 + super().apply(
161 + output=output,
162 + hidden_states=hidden_states,
163 + w1=w1_bf16,
164 + w2=w2_bf16,
165 + topk_weights=topk_weights,
166 + topk_ids=topk_ids,
167 + activation=activation,
168 + global_num_experts=global_num_experts,
169 + expert_map=expert_map,
170 + a1q_scale=None,
171 + a2_scale=None,
172 + workspace13=workspace13,
173 + workspace2=workspace2,
174 + expert_tokens_meta=expert_tokens_meta,
175 + apply_router_weight_on_input=apply_router_weight_on_input,
176 + )
vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py new
+326
@@ -0,0 +1,326 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Native MXFP8 (1x32 block, E8M0 scale) MoE for AMD CDNA4 (gfx950) via Triton
4 +``tl.dot_scaled`` (hardware microscaling matmul).
5 +
6 +The expert GEMMs consume the FP8 E4M3 weights and their E8M0 block scales
7 +directly (no dequant-to-BF16), and activations are MXFP8-quantized per token.
8 +On CDNA4 ``dot_scaled`` maps to the native MX matrix-core ops; on other archs
9 +Triton upcasts to BF16 (so this stays correct, just not faster) — but the
10 +oracle only selects this path on gfx950 and routes everything else to the
11 +BF16 ``Mxfp8EmulationTritonExperts`` fallback.
12 +
13 +Structure mirrors vLLM's ``fused_moe_kernel``: tokens are sorted by expert
14 +(``moe_align_block_size``); each program computes a ``[BLOCK_M, BLOCK_N]`` tile
15 +for one expert, accumulating over K with ``dot_scaled``. SwiGLU-OAI activation
16 +and the top-k weighted reduction run in PyTorch between/after the two GEMMs.
17 +"""
18 +
19 +import torch
20 +
21 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
22 +from vllm.logger import init_logger
23 +from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
24 + Mxfp8TritonExpertsBase,
25 +)
26 +from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
27 + moe_align_block_size,
28 +)
29 +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
30 + mxfp8_e4m3_quantize,
31 +)
32 +from vllm.platforms import current_platform
33 +from vllm.triton_utils import tl, triton
34 +
35 +logger = init_logger(__name__)
36 +
37 +
38 +@triton.jit
39 +def _mxfp8_grouped_gemm_kernel(
40 + a_ptr,
41 + a_scale_ptr,
42 + b_ptr,
43 + b_scale_ptr,
44 + c_ptr,
45 + topk_weights_ptr,
46 + sorted_token_ids_ptr,
47 + expert_ids_ptr,
48 + num_tokens_post_padded_ptr,
49 + N,
50 + K,
51 + num_valid_tokens,
52 + top_k,
53 + stride_am,
54 + stride_ak,
55 + stride_asm,
56 + stride_ask,
57 + stride_be,
58 + stride_bn,
59 + stride_bk,
60 + stride_bse,
61 + stride_bsn,
62 + stride_bsk,
63 + stride_cm,
64 + stride_cn,
65 + A_DIV: tl.constexpr,
66 + MUL_WEIGHT: tl.constexpr,
67 + BLOCK_M: tl.constexpr,
68 + BLOCK_N: tl.constexpr,
69 + BLOCK_K: tl.constexpr,
70 +):
71 + pid_m = tl.program_id(0)
72 + pid_n = tl.program_id(1)
73 + num_post = tl.load(num_tokens_post_padded_ptr)
74 + if pid_m * BLOCK_M >= num_post:
75 + return
76 +
77 + offs_tid = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
78 + offs_token = tl.load(sorted_token_ids_ptr + offs_tid).to(tl.int64)
79 + token_mask = offs_token < num_valid_tokens
80 + off_e = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
81 +
82 + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
83 + offs_k = tl.arange(0, BLOCK_K)
84 + offs_sk = tl.arange(0, BLOCK_K // 32)
85 + a_row = offs_token // A_DIV
86 +
87 + a_ptrs = a_ptr + a_row[:, None] * stride_am + offs_k[None, :] * stride_ak
88 + as_ptrs = a_scale_ptr + a_row[:, None] * stride_asm + offs_sk[None, :] * stride_ask
89 + b_ptrs = (
90 + b_ptr
91 + + off_e * stride_be
92 + + offs_n[:, None] * stride_bn
93 + + offs_k[None, :] * stride_bk
94 + )
95 + bs_ptrs = (
96 + b_scale_ptr
97 + + off_e * stride_bse
98 + + offs_n[:, None] * stride_bsn
99 + + offs_sk[None, :] * stride_bsk
100 + )
101 +
102 + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
103 + n_mask = offs_n < N
104 + for _ in range(0, tl.cdiv(K, BLOCK_K)):
105 + a = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0)
106 + b = tl.load(b_ptrs, mask=n_mask[:, None], other=0.0)
107 + asc = tl.load(as_ptrs, mask=token_mask[:, None], other=0)
108 + bsc = tl.load(bs_ptrs, mask=n_mask[:, None], other=0)
109 + acc += tl.dot_scaled(a, asc, "e4m3", b.T, bsc, "e4m3")
110 +
111 + a_ptrs += BLOCK_K * stride_ak
112 + b_ptrs += BLOCK_K * stride_bk
113 + as_ptrs += (BLOCK_K // 32) * stride_ask
114 + bs_ptrs += (BLOCK_K // 32) * stride_bsk
115 +
116 + if MUL_WEIGHT:
117 + w = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0)
118 + acc = acc * w[:, None]
119 +
120 + c_ptrs = c_ptr + offs_token[:, None] * stride_cm + offs_n[None, :] * stride_cn
121 + tl.store(
122 + c_ptrs,
123 + acc.to(c_ptr.dtype.element_ty),
124 + mask=token_mask[:, None] & n_mask[None, :],
125 + )
126 +
127 +
128 +def _grouped_gemm_mxfp8(
129 + a_q: torch.Tensor, # [M, K] fp8 e4m3
130 + a_scale: torch.Tensor, # [M, K//32] uint8 (E8M0)
131 + w: torch.Tensor, # [E, N, K] fp8 e4m3
132 + w_scale: torch.Tensor, # [E, N, K//32] uint8 (E8M0)
133 + sorted_token_ids: torch.Tensor,
134 + expert_ids: torch.Tensor,
135 + num_tokens_post_padded: torch.Tensor,
136 + num_valid_tokens: int,
137 + top_k: int,
138 + block_m: int,
139 + out_dtype: torch.dtype,
140 + a_div: int,
141 + mul_weight_by: torch.Tensor | None = None,
142 + expert_map: torch.Tensor | None = None,
143 +) -> torch.Tensor:
144 + M_routed = num_valid_tokens
145 + E, N, K = w.shape
146 + assert K % 128 == 0, f"MXFP8 native MoE requires K%128==0, got K={K}"
147 + # Under expert parallelism (expert_map set) tokens routed to non-local
148 + # experts are dropped from sorted_token_ids, so their output rows are never
149 + # written — zero them so the downstream reduction ignores their garbage.
150 + alloc = torch.zeros if expert_map is not None else torch.empty
151 + out = alloc((M_routed, N), dtype=out_dtype, device=a_q.device)
152 + BLOCK_N = 128
153 + BLOCK_K = 128
154 + grid = (triton.cdiv(sorted_token_ids.shape[0], block_m), triton.cdiv(N, BLOCK_N))
155 + _mxfp8_grouped_gemm_kernel[grid](
156 + a_q,
157 + a_scale,
158 + w,
159 + w_scale,
160 + out,
161 + mul_weight_by if mul_weight_by is not None else a_q,
162 + sorted_token_ids,
163 + expert_ids,
164 + num_tokens_post_padded,
165 + N,
166 + K,
167 + num_valid_tokens,
168 + top_k,
169 + a_q.stride(0),
170 + a_q.stride(1),
171 + a_scale.stride(0),
172 + a_scale.stride(1),
173 + w.stride(0),
174 + w.stride(1),
175 + w.stride(2),
176 + w_scale.stride(0),
177 + w_scale.stride(1),
178 + w_scale.stride(2),
179 + out.stride(0),
180 + out.stride(1),
181 + A_DIV=a_div,
182 + MUL_WEIGHT=mul_weight_by is not None,
183 + BLOCK_M=block_m,
184 + BLOCK_N=BLOCK_N,
185 + BLOCK_K=BLOCK_K,
186 + num_warps=8,
187 + )
188 + return out
189 +
190 +
191 +def fused_moe_mxfp8_native(
192 + hidden_states: torch.Tensor, # [T, H] bf16
193 + w13: torch.Tensor, # [E, 2I, H] fp8
194 + w13_scale: torch.Tensor, # [E, 2I, H//32] uint8
195 + w2: torch.Tensor, # [E, H, I] fp8
196 + w2_scale: torch.Tensor, # [E, H, I//32] uint8
197 + topk_weights: torch.Tensor, # [T, top_k]
198 + topk_ids: torch.Tensor, # [T, top_k] (global expert ids)
199 + *,
200 + alpha: float,
201 + beta: float,
202 + limit: float | None,
203 + global_num_experts: int,
204 + expert_map: torch.Tensor | None,
205 +) -> torch.Tensor:
206 + T, H = hidden_states.shape
207 + top_k = topk_ids.shape[1]
208 + M = T * top_k
209 +
210 + block_m = 64
211 + sorted_ids, expert_ids, num_post = moe_align_block_size(
212 + topk_ids,
213 + block_m,
214 + global_num_experts,
215 + expert_map,
216 + ignore_invalid_experts=expert_map is not None,
217 + )
218 +
219 + # GEMM1: x (mxfp8) @ w13^T -> [M, 2I]
220 + a_q, a_s = mxfp8_e4m3_quantize(hidden_states)
221 + g1 = _grouped_gemm_mxfp8(
222 + a_q,
223 + a_s,
224 + w13,
225 + w13_scale,
226 + sorted_ids,
227 + expert_ids,
228 + num_post,
229 + M,
230 + top_k,
231 + block_m,
232 + hidden_states.dtype,
233 + a_div=top_k,
234 + expert_map=expert_map,
235 + ) # [M, 2I]
236 +
237 + # SwiGLU-OAI (split layout: gate=g1[:, :I], up=g1[:, I:]) FUSED with the
238 + # GEMM2 MXFP8 activation-quant in one fp32 Triton pass — no bf16 ``act``
239 + # round-trip to HBM. Bit-exact vs the unfused swiglu+quant chain on measured
240 + # MoE shapes, and ~1.2-1.9x faster on that step in isolation. (Not the #22
241 + # ``silu_and_mul_with_clamp`` op: it rounds intermediates to bf16, rel ~3e-3.)
242 + # Lazy import: the amd.ops package pulls in the minimax_m3 platform dispatch,
243 + # only resolvable after the model module finishes loading.
244 + from vllm.models.minimax_m3.amd.ops import swiglu_oai_quantize_mxfp8
245 +
246 + # GEMM2: act (mxfp8) @ w2^T -> [M, H], weighted by topk_weights, then reduce.
247 + act_q, act_s = swiglu_oai_quantize_mxfp8(g1, alpha=alpha, beta=beta, limit=limit)
248 + g2 = _grouped_gemm_mxfp8(
249 + act_q,
250 + act_s,
251 + w2,
252 + w2_scale,
253 + sorted_ids,
254 + expert_ids,
255 + num_post,
256 + M,
257 + top_k,
258 + block_m,
259 + torch.float32,
260 + a_div=1,
261 + mul_weight_by=topk_weights.reshape(-1).to(torch.float32),
262 + expert_map=expert_map,
263 + ) # [M, H] == [T*top_k, H]
264 +
265 + return g2.view(T, top_k, H).sum(dim=1).to(hidden_states.dtype)
266 +
267 +
268 +class Mxfp8NativeTritonExperts(Mxfp8TritonExpertsBase):
269 + """Native MXFP8 MoE (CDNA4 ``dot_scaled``) on gfx950."""
270 +
271 + @property
272 + def quant_dtype(self) -> torch.dtype | str | None:
273 + return self.quant_config.quant_dtype
274 +
275 + @property
276 + def block_shape(self) -> list[int] | None:
277 + return self.quant_config.block_shape
278 +
279 + @property
280 + def expects_unquantized_inputs(self) -> bool:
281 + # Activations are MXFP8-quantized inside ``fused_moe_mxfp8_native``.
282 + return True
283 +
284 + @staticmethod
285 + def _supports_current_device() -> bool:
286 + return current_platform.is_rocm() and current_platform.supports_mx()
287 +
288 + def apply(
289 + self,
290 + output: torch.Tensor,
291 + hidden_states: torch.Tensor,
292 + w1: torch.Tensor,
293 + w2: torch.Tensor,
294 + topk_weights: torch.Tensor,
295 + topk_ids: torch.Tensor,
296 + activation,
297 + global_num_experts: int,
298 + expert_map: torch.Tensor | None,
299 + a1q_scale: torch.Tensor | None,
300 + a2_scale: torch.Tensor | None,
301 + workspace13: torch.Tensor,
302 + workspace2: torch.Tensor,
303 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
304 + apply_router_weight_on_input: bool,
305 + ):
306 + alpha = self.quant_config.gemm1_alpha
307 + alpha = 1.702 if alpha is None else float(alpha)
308 + beta = self.quant_config.gemm1_beta
309 + beta = 1.0 if beta is None else float(beta)
310 + limit = self.quant_config.gemm1_clamp_limit
311 + limit = None if limit is None else float(limit)
312 + out = fused_moe_mxfp8_native(
313 + hidden_states,
314 + w1,
315 + self.w1_scale_val,
316 + w2,
317 + self.w2_scale_val,
318 + topk_weights,
319 + topk_ids,
320 + alpha=alpha,
321 + beta=beta,
322 + limit=limit,
323 + global_num_experts=global_num_experts,
324 + expert_map=expert_map,
325 + )
326 + output.copy_(out)
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+32 -2
@@ -64,6 +64,15 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
64 self.quantization_emulation = False
65 super().__init__(moe_config, quant_config)
66
67 + self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
68 + # Gated-activation params: silu == swigluoai with alpha=1, beta=0.
69 + self.gemm1_alpha = (
70 + quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
71 + )
72 + self.gemm1_beta = (
73 + quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0
74 + )
75 +
76 @staticmethod
77 def activation_format() -> mk.FusedMoEActivationFormat:
78 return mk.FusedMoEActivationFormat.Standard
@@ -107,6 +116,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
116 MoEActivation.GELU,
117 MoEActivation.GELU_TANH,
118 MoEActivation.SWIGLUOAI,
119 + MoEActivation.SWIGLUOAI_UNINTERLEAVE,
120 MoEActivation.SWIGLUSTEP,
121 MoEActivation.SILU_NO_MUL,
122 MoEActivation.GELU_NO_MUL,
@@ -129,14 +139,34 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
139 return TopKWeightAndReduceNoOP()
140
141 def activation(
132 - self, activation: MoEActivation, output: torch.Tensor, input: torch.Tensor
142 + self,
143 + activation: MoEActivation,
144 + output: torch.Tensor,
145 + input: torch.Tensor,
146 + **kwargs,
147 ) -> None:
148 gemm1_clamp_limit = self.quant_config.gemm1_clamp_limit
149 if activation == MoEActivation.SILU and gemm1_clamp_limit is not None:
150 swiglu_limit_func(output, input, float(gemm1_clamp_limit))
151 return
152
139 - super().activation(activation, output, input)
153 + # SWIGLUOAI_UNINTERLEAVE routes to the silu_and_mul_with_clamp kernel and
154 + # needs the clamped-SwiGLU params (gemm1_clamp_limit/alpha/beta read from
155 + # the quant config in __init__) forwarded; without a clamp_limit it
156 + # asserts. Other activations ignore alpha/beta/clamp_limit.
157 + if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
158 + assert gemm1_clamp_limit is not None, (
159 + "SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit"
160 + )
161 +
162 + super().activation(
163 + activation,
164 + output,
165 + input,
166 + clamp_limit=gemm1_clamp_limit,
167 + alpha=self.gemm1_alpha,
168 + beta=self.gemm1_beta,
169 + )
170
171 def workspace_shapes(
172 self,
vllm/model_executor/layers/fused_moe/layer.py
+6
@@ -120,6 +120,8 @@ def FusedMoE(
120 scoring_func: str = "softmax",
121 routed_scaling_factor: float = 1.0,
122 swiglu_limit: float | None = None,
123 + swiglu_alpha: float | None = None,
124 + swiglu_beta: float | None = None,
125 e_score_correction_bias: torch.Tensor | None = None,
126 apply_router_weight_on_input: bool = False,
127 activation: str = "silu",
@@ -322,6 +324,8 @@ def FusedMoE(
324 device=vllm_config.device_config.device,
325 routing_method=router.routing_method_type, # Not ideal
326 swiglu_limit=swiglu_limit,
327 + swiglu_alpha=swiglu_alpha,
328 + swiglu_beta=swiglu_beta,
329 max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size,
330 )
331
@@ -353,6 +357,8 @@ def FusedMoE(
357 if not apply_routed_scale_to_output
358 else 1.0,
359 swiglu_limit=swiglu_limit,
360 + swiglu_alpha=swiglu_alpha,
361 + swiglu_beta=swiglu_beta,
362 # TODO get from router? needs to be truncated?
363 e_score_correction_bias=e_score_correction_bias,
364 apply_router_weight_on_input=apply_router_weight_on_input,
vllm/model_executor/layers/fused_moe/modular_kernel.py
+11 -2
@@ -880,9 +880,18 @@ class FusedMoEExpertsModular(FusedMoEExperts):
880 return N if not activation.is_gated else N // 2
881
882 def activation(
883 - self, activation: MoEActivation, output: torch.Tensor, input: torch.Tensor
883 + self,
884 + activation: MoEActivation,
885 + output: torch.Tensor,
886 + input: torch.Tensor,
887 + *,
888 + clamp_limit: float | None = None,
889 + alpha: float = 1.0,
890 + beta: float = 0.0,
891 ) -> None:
885 - apply_moe_activation(activation, output, input)
892 + apply_moe_activation(
893 + activation, output, input, clamp_limit=clamp_limit, alpha=alpha, beta=beta
894 + )
895
896 @abstractmethod
897 def finalize_weight_and_reduce_impl(self) -> TopKWeightAndReduce:
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+21 -4
@@ -52,6 +52,13 @@ class Fp8MoeBackend(Enum):
52 BATCHED_VLLM_CUTLASS = "BATCHED_VLLM_CUTLASS"
53 XPU = "XPU"
54 CPU = "CPU"
55 + # Dequantize-to-BF16 emulation for MXFP8 on devices without a native
56 + # MXFP8 MoE kernel (e.g. ROCm). Weights pass through unchanged here.
57 + EMULATION = "EMULATION"
58 + # MXFP8 MoE via a Triton ``dot_scaled`` kernel that lowers to CDNA4
59 + # (gfx950) native MX matrix-core ops. Weights stay in MXFP8 (no load-time
60 + # format conversion); the FP8 values + E8M0 scales are consumed directly.
61 + NATIVE_MXFP8 = "NATIVE_MXFP8"
62
63
64 def _get_priority_backends(
@@ -463,6 +470,10 @@ def convert_to_fp8_moe_kernel_format(
470 Fp8MoeBackend.VLLM_CUTLASS,
471 Fp8MoeBackend.BATCHED_VLLM_CUTLASS,
472 Fp8MoeBackend.XPU,
473 + # EMULATION dequantizes weights at runtime; NATIVE_MXFP8 consumes
474 + # the MXFP8 weights as-is — neither needs a load-time layout change.
475 + Fp8MoeBackend.EMULATION,
476 + Fp8MoeBackend.NATIVE_MXFP8,
477 ]:
478 raise ValueError(f"Unsupported FP8 MoE backend: {fp8_backend.value}")
479
@@ -481,6 +492,8 @@ def make_fp8_moe_quant_config(
492 per_act_token_quant: bool = False,
493 per_out_ch_quant: bool = False,
494 swiglu_limit: float | None = None,
495 + gemm1_alpha: float | None = None,
496 + gemm1_beta: float | None = None,
497 ) -> FusedMoEQuantConfig:
498 """
499 Create FusedMoEQuantConfig for the specified FP8 Backend.
@@ -503,6 +516,9 @@ def make_fp8_moe_quant_config(
516 w1_bias=w1_bias,
517 w2_bias=w2_bias,
518 block_shape=block_shape,
519 + gemm1_alpha=gemm1_alpha,
520 + gemm1_beta=gemm1_beta,
521 + gemm1_clamp_limit=swiglu_limit,
522 )
523
524 # Flashinfer CUTLASS per-tensor uses single dq scale
@@ -522,10 +538,9 @@ def make_fp8_moe_quant_config(
538 g2_alphas=(w2_scale * a2_scale).squeeze(),
539 gemm1_clamp_limit=swiglu_limit,
540 )
525 - # MXFP8 uses "mxfp8" quant_dtype so the prepare step dispatches to
526 - # _mxfp8_e4m3_quantize rather than standard FP8 block quantization.
527 - # Non-swizzled layout is required since the TRTLLM kernel expects
528 - # scales in (num_tokens, hidden_dim // 32) format.
541 + # MXFP8 (block [1, 32]) dispatches to the mxfp8 activation quant. Scales are
542 + # the non-swizzled (num_tokens, hidden_dim // 32) uint8 UE8M0 layout for all
543 + # backends; the DeepGEMM expert permute repacks them for the grouped GEMM.
544 if block_shape == [1, 32]:
545 return FusedMoEQuantConfig.make(
546 "mxfp8",
@@ -537,6 +552,8 @@ def make_fp8_moe_quant_config(
552 a2_scale=a2_scale,
553 block_shape=block_shape,
554 is_scale_swizzled=False,
555 + gemm1_alpha=gemm1_alpha,
556 + gemm1_beta=gemm1_beta,
557 gemm1_clamp_limit=swiglu_limit,
558 )
559
vllm/model_executor/layers/fused_moe/oracle/mxfp8.py
+49 -1
@@ -12,22 +12,43 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
12 kMxfp8Dynamic,
13 kMxfp8Static,
14 )
15 +from vllm.platforms import current_platform
16
17 logger = init_logger(__name__)
18
19 _SUPPORTED_BACKENDS = (
20 Fp8MoeBackend.FLASHINFER_TRTLLM,
21 + Fp8MoeBackend.DEEPGEMM,
22 Fp8MoeBackend.MARLIN,
23 Fp8MoeBackend.XPU,
24 )
25
26 _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = {
27 "flashinfer_trtllm": Fp8MoeBackend.FLASHINFER_TRTLLM,
28 + "deep_gemm": Fp8MoeBackend.DEEPGEMM,
29 "marlin": Fp8MoeBackend.MARLIN,
30 "xpu": Fp8MoeBackend.XPU,
31 }
32
33
34 +def _mxfp8_backend_to_kernel_cls(
35 + backend: Fp8MoeBackend,
36 +) -> list[type[mk.FusedMoEExperts]]:
37 + """Resolve the MXFP8 expert classes for a backend.
38 +
39 + DeepGEMM resolves directly to ``DeepGemmExperts`` (not the
40 + ``TritonOrDeepGemmExperts`` wrapper, whose Triton fallback cannot handle the
41 + MXFP8 1x32 scheme); all other backends defer to the FP8 resolver.
42 + """
43 + if backend == Fp8MoeBackend.DEEPGEMM:
44 + from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
45 + DeepGemmExperts,
46 + )
47 +
48 + return [DeepGemmExperts]
49 + return backend_to_kernel_cls(backend)
50 +
51 +
52 def _select_kernel_cls(
53 backend: Fp8MoeBackend,
54 config: FusedMoEConfig,
@@ -39,7 +60,7 @@ def _select_kernel_cls(
60 else mk.FusedMoEActivationFormat.Standard
61 )
62 last_reason: str | None = None
42 - for cls in backend_to_kernel_cls(backend):
63 + for cls in _mxfp8_backend_to_kernel_cls(backend):
64 supported, reason = cls.is_supported_config(
65 cls,
66 config,
@@ -55,6 +76,29 @@ def _select_kernel_cls(
76 )
77
78
79 +def _select_rocm_mxfp8_backend() -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:
80 + """ROCm fallback when vendor MXFP8 backends are unavailable."""
81 +
82 + if current_platform.supports_mx():
83 + from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
84 + Mxfp8NativeTritonExperts,
85 + )
86 +
87 + logger.info_once("Using native CDNA4 (gfx950) MXFP8 dot_scaled MoE backend.")
88 + return Fp8MoeBackend.NATIVE_MXFP8, Mxfp8NativeTritonExperts
89 +
90 + from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
91 + Mxfp8EmulationTritonExperts,
92 + )
93 +
94 + logger.info_once(
95 + "No native MXFP8 MoE backend available on this device; "
96 + "MXFP8 weights will be dequantized to BF16 once at load time and the "
97 + "MoE will run in BF16 (no per-step dequant)."
98 + )
99 + return Fp8MoeBackend.EMULATION, Mxfp8EmulationTritonExperts
100 +
101 +
102 def select_mxfp8_moe_backend(
103 config: FusedMoEConfig,
104 ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:
@@ -88,4 +132,8 @@ def select_mxfp8_moe_backend(
132 logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value)
133 return backend, experts_cls
134
135 + # simplify the logic for rocm, refactor later when more backends are supported
136 + if current_platform.is_rocm():
137 + return _select_rocm_mxfp8_backend()
138 +
139 raise ValueError("No MXFP8 MoE backends available.")
vllm/model_executor/layers/fused_moe/routed_experts.py
+4
@@ -72,6 +72,8 @@ class RoutedExperts(PluggableLayer):
72 scoring_func: str = "softmax",
73 routed_scaling_factor: float = 1.0,
74 swiglu_limit: float | None = None,
75 + swiglu_alpha: float | None = None,
76 + swiglu_beta: float | None = None,
77 e_score_correction_bias: torch.Tensor | None = None,
78 apply_router_weight_on_input: bool = False,
79 ):
@@ -103,6 +105,8 @@ class RoutedExperts(PluggableLayer):
105 self.scoring_func = scoring_func
106 self.routed_scaling_factor = routed_scaling_factor
107 self.swiglu_limit = swiglu_limit
108 + self.swiglu_alpha = swiglu_alpha
109 + self.swiglu_beta = swiglu_beta
110 self.e_score_correction_bias = e_score_correction_bias
111 self.apply_router_weight_on_input = apply_router_weight_on_input
112 # End random parameters
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+4 -5
@@ -29,9 +29,9 @@ class GateLinear(ReplicatedLinear):
29 DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
30 DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
31
32 - # Dimensions supported by the fp32 specialized kernel
33 - FP32_SUPPORTED_NUM_EXPERTS = [256]
34 - FP32_SUPPORTED_HIDDEN_SIZES = [3072]
32 + # (hidden_size, num_experts) pairs with an instantiated fp32 kernel:
33 + # (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3
34 + FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)}
35 FP32_MAX_TOKENS = 32
36
37 def __init__(
@@ -82,8 +82,7 @@ class GateLinear(ReplicatedLinear):
82 and self.weight.dtype == torch.float32
83 and current_platform.is_cuda()
84 and (is_hopper or is_blackwell)
85 - and output_size in self.FP32_SUPPORTED_NUM_EXPERTS
86 - and input_size in self.FP32_SUPPORTED_HIDDEN_SIZES
85 + and (input_size, output_size) in self.FP32_SUPPORTED_SHAPES
86 )
87
88 # cuBLAS bf16→fp32 eligibility
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+20 -8
@@ -11,7 +11,6 @@ import vllm.envs as envs
11 from vllm.logger import init_logger
12 from vllm.model_executor.custom_op import CustomOp
13 from vllm.model_executor.layers.fused_moe.config import (
14 - FUSED_MOE_UNQUANTIZED_CONFIG,
14 FusedMoEConfig,
15 FusedMoEQuantConfig,
16 biased_moe_quant_config,
@@ -184,11 +183,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
183
184 if not is_weight_update:
185 # Setup moe kernel only on the first call. For the unquantized
187 - # method, moe_quant_config is either the constant
188 - # FUSED_MOE_UNQUANTIZED_CONFIG or biased_moe_quant_config(...)
189 - # which references layer.w{13,2}_bias; since weight updates
190 - # mutate those bias tensors in place, the kernel does not need
191 - # to be re-built.
186 + # method, moe_quant_config carries no quantized scales -- only
187 + # optional w{13,2}_bias references and SwiGLU gate params. Since
188 + # weight updates mutate those bias tensors in place, the kernel
189 + # does not need to be re-built.
190 self.moe_quant_config = self.get_fused_moe_quant_config(layer)
191 assert self.moe_quant_config is not None
192 assert self.experts_cls is not None
@@ -272,13 +270,27 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
270 )
271
272 def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
273 + # SwiGLU/swigluoai gate params live on the layer; plumb them into the
274 + # quant config so the fused activation (e.g. swigluoai_uninterleave on
275 + # MiniMax-M3) receives gemm1_clamp_limit/alpha/beta.
276 + gemm1_alpha = getattr(layer, "swiglu_alpha", None)
277 + gemm1_beta = getattr(layer, "swiglu_beta", None)
278 + gemm1_clamp_limit = getattr(layer, "swiglu_limit", None)
279 +
280 if self.moe.has_bias:
281 return biased_moe_quant_config(
282 layer.w13_bias,
283 layer.w2_bias,
284 + gemm1_alpha=gemm1_alpha,
285 + gemm1_beta=gemm1_beta,
286 + gemm1_clamp_limit=gemm1_clamp_limit,
287 )
280 - else:
281 - return FUSED_MOE_UNQUANTIZED_CONFIG
288 +
289 + return FusedMoEQuantConfig.make(
290 + gemm1_alpha=gemm1_alpha,
291 + gemm1_beta=gemm1_beta,
292 + gemm1_clamp_limit=gemm1_clamp_limit,
293 + )
294
295 def apply(
296 self,
vllm/model_executor/layers/fused_moe/utils.py
+2
@@ -313,6 +313,8 @@ def moe_kernel_quantize_input(
313 "moe_kernel_quantize_input does not support quant_dtype='mxfp8' MOE "
314 "quantization emulation. Please open an issue."
315 )
316 + # Non-swizzled (M, K/32) uint8 UE8M0 scales; deepgemm_moe_permute packs
317 + # them for DeepGEMM, TRTLLM takes them as-is.
318 return _mxfp8_e4m3_quantize(
319 A,
320 A_scale,
vllm/model_executor/layers/linear.py
+186
@@ -1105,6 +1105,7 @@ class QKVParallelLinear(ColumnParallelLinear):
1105
1106 shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
1107 shard_size = self._get_shard_size_mapping(loaded_shard_id)
1108 + assert shard_offset is not None and shard_size is not None
1109
1110 if isinstance(param, BlockQuantScaleParameter):
1111 weight_block_size = getattr(self, "weight_block_size", None)
@@ -1302,6 +1303,191 @@ class QKVParallelLinear(ColumnParallelLinear):
1303 param_data.copy_(loaded_weight)
1304
1305
1306 +class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear):
1307 + """QKV projection fused with a lightning-indexer's index_q/index_k.
1308 +
1309 + NOTE: MiniMax-M3-specific. This is tailored to the M3 sparse-attention
1310 + layers (it assumes the indexer's head count equals the KV head count and
1311 + shares the main head_dim); it is not a general-purpose linear layer. It
1312 + lives here only to sit alongside QKVParallelLinear, whose sharding /
1313 + weight-loading machinery it reuses.
1314 +
1315 + A single column-parallel GEMM emits, per rank::
1316 +
1317 + [q | k | v | index_q | index_k]
1318 +
1319 + ``index_q`` must have the same head count as the KV heads
1320 + (``total_num_index_heads == total_num_kv_heads``) and ``index_head_size ==
1321 + head_size``, so it shards exactly like K/V -- including the KV-head
1322 + *replication* path when ``tp_size > total_num_kv_heads`` (this is what makes
1323 + a TP size greater than the KV-head count work). ``index_k`` is a single
1324 + shared head, replicated to every rank.
1325 + """
1326 +
1327 + def __init__(
1328 + self,
1329 + hidden_size: int,
1330 + head_size: int,
1331 + total_num_heads: int,
1332 + total_num_kv_heads: int,
1333 + total_num_index_heads: int,
1334 + index_head_size: int,
1335 + bias: bool = False,
1336 + quant_config: QuantizationConfig | None = None,
1337 + prefix: str = "",
1338 + ) -> None:
1339 + # index_q rides the KV-head sharding/replication path, so its head count
1340 + # must match the KV heads.
1341 + assert total_num_index_heads == total_num_kv_heads, (
1342 + "MinimaxM3QKVParallelLinearWithIndexer requires "
1343 + "total_num_index_heads == total_num_kv_heads"
1344 + )
1345 + self.hidden_size = hidden_size
1346 + self.head_size = head_size
1347 + self.v_head_size = head_size
1348 + self.total_num_heads = total_num_heads
1349 + self.total_num_kv_heads = total_num_kv_heads
1350 + self.total_num_index_heads = total_num_index_heads
1351 + self.index_head_size = index_head_size
1352 +
1353 + tp_size = get_tensor_model_parallel_world_size()
1354 + self.num_heads = divide(self.total_num_heads, tp_size)
1355 + if tp_size >= self.total_num_kv_heads:
1356 + self.num_kv_heads = 1
1357 + self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads)
1358 + else:
1359 + self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
1360 + self.num_kv_head_replicas = 1
1361 + # index_q shards identically to the KV heads.
1362 + self.num_index_heads = self.num_kv_heads
1363 +
1364 + # Global per-group sizes (replicated groups counted x tp_size, matching
1365 + # the QKVParallelLinear convention). index_k is a single replicated head.
1366 + q = self.num_heads * self.head_size
1367 + kv = self.num_kv_heads * self.head_size
1368 + iq = self.num_index_heads * self.index_head_size
1369 + ik = self.index_head_size
1370 + self.output_sizes = [
1371 + q * tp_size, # q
1372 + kv * tp_size, # k
1373 + kv * tp_size, # v
1374 + iq * tp_size, # index_q
1375 + ik * tp_size, # index_k (replicated)
1376 + ]
1377 +
1378 + # Skip QKVParallelLinear.__init__ (3-group layout); build the 5-group
1379 + # column-parallel weight directly.
1380 + ColumnParallelLinear.__init__(
1381 + self,
1382 + input_size=self.hidden_size,
1383 + output_size=sum(self.output_sizes),
1384 + bias=bias,
1385 + gather_output=False,
1386 + quant_config=quant_config,
1387 + prefix=prefix,
1388 + )
1389 +
1390 + def validate_shard_id(self, loaded_shard_id: str | None) -> None:
1391 + if loaded_shard_id is None:
1392 + return
1393 + if loaded_shard_id not in ("q", "k", "v", "index_q", "index_k"):
1394 + raise ValueError(
1395 + "Shard id for MinimaxM3QKVParallelLinearWithIndexer must be one of "
1396 + "'q', 'k', 'v', 'index_q', 'index_k'; got "
1397 + f"{loaded_shard_id}."
1398 + )
1399 +
1400 + def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int | None:
1401 + h = self.head_size
1402 + nq, nkv, nidx = self.num_heads, self.num_kv_heads, self.num_index_heads
1403 + return {
1404 + "q": 0,
1405 + "k": nq * h,
1406 + "v": (nq + nkv) * h,
1407 + "index_q": (nq + 2 * nkv) * h,
1408 + "index_k": (nq + 2 * nkv + nidx) * h,
1409 + }.get(loaded_shard_id)
1410 +
1411 + def _get_shard_size_mapping(self, loaded_shard_id: str) -> int | None:
1412 + h = self.head_size
1413 + return {
1414 + "q": self.num_heads * h,
1415 + "k": self.num_kv_heads * h,
1416 + "v": self.num_kv_heads * h,
1417 + "index_q": self.num_index_heads * h,
1418 + "index_k": self.index_head_size,
1419 + }.get(loaded_shard_id)
1420 +
1421 + def weight_loader_v2(
1422 + self,
1423 + param: BasevLLMParameter,
1424 + loaded_weight: torch.Tensor,
1425 + loaded_shard_id: str | None = None,
1426 + ) -> None:
1427 + self.validate_shard_id(loaded_shard_id)
1428 + # Index checkpoints are never pre-fused on disk; a shard id is always given.
1429 + assert loaded_shard_id in ("q", "k", "v", "index_q", "index_k")
1430 +
1431 + shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
1432 + shard_size = self._get_shard_size_mapping(loaded_shard_id)
1433 + assert shard_offset is not None and shard_size is not None
1434 + if isinstance(param, BlockQuantScaleParameter):
1435 + weight_block_size = getattr(self, "weight_block_size", None)
1436 + shard_size, shard_offset = adjust_block_scale_shard(
1437 + weight_block_size, shard_size, shard_offset
1438 + )
1439 +
1440 + # index_k is fully replicated: num_heads == tp_size makes
1441 + # load_qkv_weight pick shard_id_int == 0 on every rank. q/k/v/index_q ride
1442 + # the KV-head replication factor.
1443 + num_heads = (
1444 + self.tp_size if loaded_shard_id == "index_k" else self.num_kv_head_replicas
1445 + )
1446 + param.load_qkv_weight(
1447 + loaded_weight=loaded_weight,
1448 + num_heads=num_heads,
1449 + shard_id=loaded_shard_id,
1450 + shard_offset=shard_offset,
1451 + shard_size=shard_size,
1452 + tp_rank=self.tp_rank,
1453 + )
1454 +
1455 + def weight_loader(
1456 + self,
1457 + param: Parameter,
1458 + loaded_weight: torch.Tensor,
1459 + loaded_shard_id: str | None = None,
1460 + ) -> None:
1461 + # Unquantized (bf16) path. MXFP8 checkpoints use weight_loader_v2; this
1462 + # keeps an unquantized load correct too.
1463 + self.validate_shard_id(loaded_shard_id)
1464 + assert loaded_shard_id in ("q", "k", "v", "index_q", "index_k")
1465 + output_dim = getattr(param, "output_dim", None)
1466 + assert output_dim is not None
1467 +
1468 + shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
1469 + shard_size = self._get_shard_size_mapping(loaded_shard_id)
1470 + assert shard_offset is not None and shard_size is not None
1471 + if isinstance(param, BlockQuantScaleParameter):
1472 + weight_block_size = getattr(self, "weight_block_size", None)
1473 + shard_size, shard_offset = adjust_block_scale_shard(
1474 + weight_block_size, shard_size, shard_offset
1475 + )
1476 +
1477 + param_data = param.data.narrow(output_dim, shard_offset, shard_size)
1478 + if loaded_shard_id == "q":
1479 + shard_rank = self.tp_rank
1480 + elif loaded_shard_id == "index_k":
1481 + shard_rank = 0 # replicated to every rank
1482 + else:
1483 + shard_rank = self.tp_rank // self.num_kv_head_replicas
1484 + loaded_weight = loaded_weight.narrow(
1485 + output_dim, shard_rank * shard_size, shard_size
1486 + )
1487 + assert param_data.shape == loaded_weight.shape
1488 + param_data.copy_(loaded_weight)
1489 +
1490 +
1491 # --8<-- [start:row_parallel_linear]
1492 @PluggableLayer.register("row_parallel_linear")
1493 class RowParallelLinear(LinearBase):
vllm/model_executor/layers/quantization/__init__.py
+9 -8
@@ -163,17 +163,18 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
163 "deepseek_v4_fp8": DeepseekV4FP8Config,
164 "humming": HummingConfig,
165 "online": OnlineQuantizationConfig,
166 + # MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the
167 + # ModelOpt MXFP8 config (same format). The "mxfp8" online shorthand
168 + # below only applies to the `--quantization mxfp8` CLI path.
169 + "mxfp8": ModelOptMxFp8Config,
170 }
171
168 - # Register online shorthands as quantization methods so the user can
169 - # specify "LLM(..., quantization='fp8_per_tensor')" as shorthand for
170 - # creating a more complicated online quant config object.
172 + # Register online shorthands (e.g. "fp8_per_tensor") as quant methods.
173 + # setdefault so a shorthand that is also a checkpoint method (e.g. "mxfp8")
174 + # keeps its checkpoint config; the shorthand still works via the
175 + # `--quantization` CLI path in `resolve_quantization_config`.
176 for shorthand in _ONLINE_SHORTHANDS:
172 - assert shorthand not in method_to_config, (
173 - f"Online quant shorthand {shorthand!r} conflicts with an "
174 - f"existing quantization method"
175 - )
176 - method_to_config[shorthand] = OnlineQuantizationConfig
177 + method_to_config.setdefault(shorthand, OnlineQuantizationConfig)
178
179 # Update the `method_to_config` with customized quantization methods.
180 method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
vllm/model_executor/layers/quantization/modelopt.py
+77 -1
@@ -2,11 +2,12 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from fnmatch import fnmatch
5 -from typing import TYPE_CHECKING, Any
5 +from typing import TYPE_CHECKING, Any, cast
6
7 import torch
8 from torch.nn.parameter import Parameter
9
10 +import vllm.envs as envs
11 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12 from vllm.config import get_current_vllm_config
13 from vllm.logger import init_logger
@@ -27,6 +28,7 @@ from vllm.model_executor.layers.fused_moe import (
28 SharedExperts,
29 )
30 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
31 + Fp8MoeBackend,
32 convert_to_fp8_moe_kernel_format,
33 make_fp8_moe_kernel,
34 make_fp8_moe_quant_config,
@@ -1720,6 +1722,22 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase):
1722 return "modelopt_mxfp8"
1723 return None
1724
1725 + @classmethod
1726 + def from_config(cls, config: dict[str, Any]) -> "ModelOptMxFp8Config":
1727 + # MiniMax-style checkpoints tag `quant_method: "mxfp8"` + `ignored_layers`
1728 + # (same on-disk format as ModelOpt MXFP8); normalize to the ModelOpt
1729 + # schema and reuse the shared parser.
1730 + if "quantization" not in config and not config.get("quant_algo"):
1731 + config = {
1732 + "quant_method": "modelopt",
1733 + "quantization": {
1734 + "quant_algo": "MXFP8",
1735 + "kv_cache_quant_algo": config.get("kv_cache_quant_algo"),
1736 + "exclude_modules": config.get("ignored_layers", []) or [],
1737 + },
1738 + }
1739 + return cast("ModelOptMxFp8Config", super().from_config(config))
1740 +
1741 @classmethod
1742 def _from_config(
1743 cls,
@@ -1823,6 +1841,12 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase):
1841 layer.register_parameter("weight_scale", weight_scale)
1842
1843 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1844 + # Idempotent: the emulation kernel may dequant the weight to BF16 at load
1845 + # time (>=2-byte). If already converted, there is nothing left to do --
1846 + # avoid re-running the MXFP8-only validation/conversion below.
1847 + if layer.weight.element_size() >= 2:
1848 + return
1849 +
1850 # Validate weight tensor
1851 if layer.weight.ndim != 2:
1852 raise ValueError(
@@ -2065,6 +2089,44 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
2089 torch.stack(w2_scale_shuffled).contiguous(),
2090 )
2091
2092 + def _dequant_mxfp8_weights_to_bf16(self, layer: RoutedExperts) -> None:
2093 + """One-time MXFP8->BF16 weight dequant for the emulation path.
2094 +
2095 + On devices without a native MXFP8 MoE kernel (e.g. gfx942 / MI300),
2096 + ``Mxfp8EmulationTritonExperts`` otherwise dequantizes every expert
2097 + weight to BF16 on *every* forward step -- the dominant cost (conc1
2098 + ~1.3 tok/s). Doing the dequant once here and replacing the MXFP8
2099 + parameters with BF16 makes the MoE run exactly like a plain BF16
2100 + checkpoint (full precision, no per-step dequant); SwiGLU-OAI is still
2101 + applied by the experts' ``activation()`` override. The MXFP8 weights
2102 + are freed by ``replace_parameter`` (BF16 is 2x their size; the small
2103 + E8M0 scale tensors are left in place, unused).
2104 + """
2105 + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
2106 + dequant_mxfp8_to_bf16,
2107 + )
2108 +
2109 + target_dtype = getattr(layer, "orig_dtype", torch.bfloat16)
2110 + num_experts = layer.w13_weight.shape[0]
2111 +
2112 + # dequant_mxfp8_to_bf16 handles arbitrary leading dims (*x.shape[:-1]),
2113 + # so dequant the whole [E, N, K] weight in one vectorized call.
2114 + w13_bf16 = dequant_mxfp8_to_bf16(layer.w13_weight, layer.w13_weight_scale).to(
2115 + target_dtype
2116 + )
2117 + w2_bf16 = dequant_mxfp8_to_bf16(layer.w2_weight, layer.w2_weight_scale).to(
2118 + target_dtype
2119 + )
2120 +
2121 + replace_parameter(layer, "w13_weight", w13_bf16)
2122 + replace_parameter(layer, "w2_weight", w2_bf16)
2123 +
2124 + logger.info_once(
2125 + "MXFP8->BF16 load-time dequant complete (%d experts/layer); MoE "
2126 + "now runs in BF16 with no per-step dequant.",
2127 + num_experts,
2128 + )
2129 +
2130 def process_weights_after_loading(self, layer: RoutedExperts) -> None:
2131 # TODO(bnell): why is this required only for mxfp8?
2132 if getattr(layer, "_already_called_process_weights_after_loading", False):
@@ -2102,6 +2164,17 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
2164 routing_tables=layer._expert_routing_tables(),
2165 )
2166
2167 + # No native MXFP8 MoE kernel on this device (e.g. gfx942): the emulation
2168 + # experts would dequant MXFP8->BF16 every forward step. Convert the
2169 + # weights to BF16 once, here, so the MoE runs like a BF16 checkpoint.
2170 + # Opt out (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0) to keep the 1-byte
2171 + # MXFP8 weights and dequant per-step (~half the memory, much slower).
2172 + if (
2173 + self.mxfp8_backend == Fp8MoeBackend.EMULATION
2174 + and envs.VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD
2175 + ):
2176 + self._dequant_mxfp8_weights_to_bf16(layer)
2177 +
2178 def maybe_make_prepare_finalize(
2179 self,
2180 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
@@ -2131,6 +2204,9 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
2204 a1_scale=None,
2205 a2_scale=None,
2206 block_shape=self.weight_block_size,
2207 + swiglu_limit=getattr(layer, "swiglu_limit", None),
2208 + gemm1_alpha=getattr(layer, "swiglu_alpha", None),
2209 + gemm1_beta=getattr(layer, "swiglu_beta", None),
2210 )
2211
2212 def apply_monolithic(
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+138 -80
@@ -159,75 +159,104 @@ def _silu_mul_quant_fp8_packed_kernel(
159 output_q_stride_m,
160 output_scale_stride_k,
161 clamp_limit,
162 + alpha,
163 + beta,
164 N: tl.constexpr,
163 - NUM_GROUPS: tl.constexpr,
165 + GROUPS_PER_ROW: tl.constexpr,
166 + PACKS_PER_ROW: tl.constexpr,
167 fp8_min: tl.constexpr,
168 fp8_max: tl.constexpr,
169 GROUP_SIZE: tl.constexpr,
170 + PACKS_PER_CTA: tl.constexpr,
171 BLOCK_M: tl.constexpr,
172 HAS_CLAMP: tl.constexpr,
173 ):
170 - N_2: tl.constexpr = N // 2
171 -
172 - pid_pack = tl.program_id(0)
173 - pid_m = tl.program_id(1)
174 - m_offset = pid_m.to(tl.int64) * BLOCK_M
175 -
176 - if m_offset >= M:
177 - return
178 -
179 - offs_m = tl.arange(0, BLOCK_M)
180 - offs_n = tl.arange(0, GROUP_SIZE)
181 - row_mask = (m_offset + offs_m) < M
182 -
183 - base_row_offset = (m_offset + offs_m[:, None]) * input_stride_m
184 - base_out_offset = (m_offset + offs_m[:, None]) * output_q_stride_m
185 -
186 - packed_scale = tl.zeros((BLOCK_M,), dtype=tl.int32)
187 -
188 - for pack_idx in tl.static_range(4):
189 - group_id = pid_pack * 4 + pack_idx
190 -
191 - if group_id < NUM_GROUPS:
192 - n_offset = group_id * GROUP_SIZE
193 -
194 - act_ptrs = input_ptr + base_row_offset + n_offset + offs_n[None, :]
195 - act_in = tl.load(act_ptrs, mask=row_mask[:, None], other=0.0)
196 -
197 - mul_ptrs = act_ptrs + N_2
198 - mul_in = tl.load(mul_ptrs, mask=row_mask[:, None], other=0.0)
199 -
200 - act_f32 = act_in.to(tl.float32)
201 - mul_f32 = mul_in.to(tl.float32)
202 -
203 - if HAS_CLAMP:
204 - act_f32 = tl.minimum(act_f32, clamp_limit)
205 - mul_f32 = tl.clamp(mul_f32, -clamp_limit, clamp_limit)
206 -
207 - y = (act_f32 / (1.0 + tl.exp(-act_f32))) * mul_f32
208 - # Round through bf16 to match unfused precision path
209 - y = y.to(tl.bfloat16).to(tl.float32)
210 -
211 - absmax = tl.max(tl.abs(y), axis=1)
212 -
213 - scale_raw = tl.maximum(absmax / fp8_max, 1e-10)
214 - exponent = tl.ceil(tl.log2(scale_raw))
215 - scale = tl.math.exp2(exponent)
174 + GROUPS_PER_PACK: tl.constexpr = 4
175 + hidden_size: tl.constexpr = N // 2
176 +
177 + pack_tile = tl.program_id(0)
178 + row_start = tl.program_id(1).to(tl.int64) * BLOCK_M
179 + row_step = tl.num_programs(1).to(tl.int64) * BLOCK_M
180 +
181 + groups_per_cta: tl.constexpr = PACKS_PER_CTA * GROUPS_PER_PACK
182 + elems_per_cta: tl.constexpr = groups_per_cta * GROUP_SIZE
183 + col_start = pack_tile * elems_per_cta
184 + col_offsets = tl.arange(0, elems_per_cta)
185 + row_offsets = tl.arange(0, BLOCK_M)
186 + pack_offsets = tl.arange(0, PACKS_PER_CTA)
187 +
188 + col_mask = (col_start + col_offsets) < (GROUPS_PER_ROW * GROUP_SIZE)
189 +
190 + # persistent with grid_m-stride loop
191 + while row_start < M:
192 + rows = row_start + row_offsets
193 + row_mask = rows < M
194 + input_row_start = rows[:, None] * input_stride_m
195 + output_row_start = rows[:, None] * output_q_stride_m
196 +
197 + gate_flat = tl.load(
198 + input_ptr + input_row_start + col_start + col_offsets[None, :],
199 + mask=row_mask[:, None] & col_mask[None, :],
200 + other=0.0,
201 + )
202 + up_flat = tl.load(
203 + input_ptr
204 + + input_row_start
205 + + hidden_size
206 + + col_start
207 + + col_offsets[None, :],
208 + mask=row_mask[:, None] & col_mask[None, :],
209 + other=0.0,
210 + )
211
217 - y_q = tl.clamp(y / scale[:, None], fp8_min, fp8_max)
212 + gate = tl.reshape(gate_flat, (BLOCK_M, groups_per_cta, GROUP_SIZE)).to(
213 + tl.float32
214 + )
215 + up = tl.reshape(up_flat, (BLOCK_M, groups_per_cta, GROUP_SIZE)).to(tl.float32)
216 +
217 + if HAS_CLAMP:
218 + gate = tl.minimum(gate, clamp_limit)
219 + up = tl.clamp(up, -clamp_limit, clamp_limit)
220 +
221 + # Unified gated activation: silu == swigluoai with alpha=1, beta=0.
222 + # glu = gate * sigmoid(alpha * gate); y = (up + beta) * glu
223 + glu = gate / (1.0 + tl.exp(-gate * alpha))
224 + y = (up + beta) * glu
225 + # Round through bf16 to match unfused precision path
226 + y = y.to(tl.bfloat16).to(tl.float32)
227 +
228 + absmax = tl.max(tl.abs(y), axis=2)
229 + scale_raw = tl.maximum(absmax / fp8_max, 1e-10)
230 + exponent = tl.ceil(tl.log2(scale_raw))
231 + scale = tl.math.exp2(exponent)
232 +
233 + y_q = tl.clamp(y / scale[:, :, None], fp8_min, fp8_max)
234 +
235 + y_q_flat = tl.reshape(y_q, (BLOCK_M, elems_per_cta))
236 + tl.store(
237 + output_q_ptr + output_row_start + col_start + col_offsets[None, :],
238 + y_q_flat.to(output_q_ptr.dtype.element_ty),
239 + mask=row_mask[:, None] & col_mask[None, :],
240 + )
241
219 - out_q_ptrs = output_q_ptr + base_out_offset + n_offset + offs_n[None, :]
220 - tl.store(
221 - out_q_ptrs,
222 - y_q.to(output_q_ptr.dtype.element_ty),
223 - mask=row_mask[:, None],
224 - )
242 + scale_byte = tl.clamp(exponent + 127.0, 0.0, 255.0).to(tl.int32)
243 + scale_bytes = tl.reshape(scale_byte, (BLOCK_M, PACKS_PER_CTA, GROUPS_PER_PACK))
244 + shifts = tl.arange(0, GROUPS_PER_PACK) * 8
245 + packed_scale = tl.sum(scale_bytes << shifts[None, None, :], axis=2)
246
226 - exponent_biased = tl.clamp(exponent + 127.0, 0.0, 255.0).to(tl.int32)
227 - packed_scale = packed_scale | (exponent_biased << (pack_idx * 8))
247 + scale_pack = pack_tile * PACKS_PER_CTA + pack_offsets
248 + scale_ptrs = (
249 + output_scale_ptr
250 + + scale_pack[None, :] * output_scale_stride_k
251 + + rows[:, None]
252 + )
253 + tl.store(
254 + scale_ptrs,
255 + packed_scale,
256 + mask=row_mask[:, None] & (scale_pack[None, :] < PACKS_PER_ROW),
257 + )
258
229 - scale_ptrs = output_scale_ptr + pid_pack * output_scale_stride_k + m_offset + offs_m
230 - tl.store(scale_ptrs, packed_scale, mask=row_mask)
259 + row_start += row_step
260
261
262 def silu_mul_quant_fp8_packed_triton(
@@ -235,37 +264,48 @@ def silu_mul_quant_fp8_packed_triton(
264 group_size: int = 128,
265 output_q: torch.Tensor | None = None,
266 clamp_limit: float | None = None,
267 + alpha: float = 1.0,
268 + beta: float = 0.0,
269 ) -> tuple[torch.Tensor, torch.Tensor]:
270 assert input.dim() == 2
271 assert input.is_contiguous()
272
273 M, N = input.shape
243 - N_2 = N // 2
274 + hidden_size = N // 2
275
245 - assert N_2 % group_size == 0
276 + assert hidden_size % group_size == 0
277
278 fp8_dtype = torch.float8_e4m3fn
279 finfo = torch.finfo(fp8_dtype)
280 fp8_min, fp8_max = finfo.min, finfo.max
281
251 - num_groups_per_row = N_2 // group_size
252 - num_packed_groups = (num_groups_per_row + 3) // 4
253 - tma_aligned_M = ((M + 3) // 4) * 4
282 + groups_per_row = hidden_size // group_size
283 + groups_per_pack = 4 # pack 4 UE8M0 scales to a single INT32
284 + packs_per_row = triton.cdiv(groups_per_row, groups_per_pack)
285
286 if output_q is None:
256 - output_q = torch.empty((M, N_2), dtype=fp8_dtype, device=input.device)
287 + output_q = torch.empty((M, hidden_size), dtype=fp8_dtype, device=input.device)
288
289 + aligned_m = triton.cdiv(M, 4) * 4
290 output_scale_packed = torch.empty(
259 - (num_packed_groups, tma_aligned_M),
291 + (packs_per_row, aligned_m),
292 dtype=torch.int32,
293 device=input.device,
294 ).T[:M, :]
295
264 - BLOCK_M = 8
265 - grid = (num_packed_groups, (M + BLOCK_M - 1) // BLOCK_M)
266 -
267 - num_warps = max(4, group_size // 32)
296 + # Tuned for group_size=32 (MXFP8) and group_size=128 (DeepSeek-V4)
297 + num_warps = 4
298 num_stages = 2
299 + if group_size < 128:
300 + BM = 1
301 + packs_per_cta = 8
302 + else:
303 + BM = 1 if M < 512 else 4
304 + packs_per_cta = 2 if M < 512 else 1
305 +
306 + grid_n = triton.cdiv(packs_per_row, packs_per_cta)
307 + grid_m = min(triton.cdiv(M, BM), 4096)
308 + grid = (grid_n, grid_m)
309
310 has_clamp = clamp_limit is not None
311 _silu_mul_quant_fp8_packed_kernel[grid](
@@ -277,12 +317,16 @@ def silu_mul_quant_fp8_packed_triton(
317 output_q.stride(0),
318 output_scale_packed.stride(1),
319 clamp_limit if has_clamp else 0.0,
320 + alpha,
321 + beta,
322 N=N,
281 - NUM_GROUPS=num_groups_per_row,
323 + GROUPS_PER_ROW=groups_per_row,
324 + PACKS_PER_ROW=packs_per_row,
325 fp8_min=fp8_min,
326 fp8_max=fp8_max,
327 GROUP_SIZE=group_size,
285 - BLOCK_M=BLOCK_M,
328 + PACKS_PER_CTA=packs_per_cta,
329 + BLOCK_M=BM,
330 HAS_CLAMP=has_clamp,
331 num_warps=num_warps,
332 num_stages=num_stages,
@@ -303,6 +347,8 @@ def _silu_mul_per_token_group_quant_fp8_colmajor(
347 # Information for float8
348 eps,
349 clamp_limit,
350 + alpha,
351 + beta,
352 fp8_min: tl.constexpr,
353 fp8_max: tl.constexpr,
354 use_ue8m0: tl.constexpr,
@@ -348,10 +394,14 @@ def _silu_mul_per_token_group_quant_fp8_colmajor(
394 mul_in = tl.clamp(mul_in.to(tl.float32), -clamp_limit, clamp_limit).to(
395 y_ptr.dtype.element_ty
396 )
397 + # Unified gated activation: silu == swigluoai with alpha=1, beta=0.
398 + # glu = gate * sigmoid(alpha * gate); y = (up + beta) * glu
399 + # Keep glu/up at input precision (narrow before the mul) so the alpha=1,
400 + # beta=0 defaults match the C++ silu_and_mul path bit-for-bit.
401 act_in = act_in.to(tl.float32)
352 - one_f32 = tl.cast(1, tl.float32)
353 - silu_out = (act_in / (one_f32 + tl.exp(-act_in))).to(y_ptr.dtype.element_ty)
354 - y = (silu_out * mul_in).to(tl.float32)
402 + glu = (act_in / (1.0 + tl.exp(-act_in * alpha))).to(y_ptr.dtype.element_ty)
403 + up = (mul_in.to(tl.float32) + beta).to(y_ptr.dtype.element_ty)
404 + y = (glu * up).to(tl.float32)
405
406 # quant
407 _absmax = tl.maximum(tl.max(tl.abs(y), axis=1), eps)
@@ -379,11 +429,15 @@ def silu_mul_per_token_group_quant_fp8_colmajor(
429 use_ue8m0: bool | None = None,
430 eps: float = 1e-10,
431 clamp_limit: float | None = None,
432 + group_size: int = 128,
433 + alpha: float = 1.0,
434 + beta: float = 0.0,
435 ):
436 """
384 - silu+mul + block-fp8 quant with group size 128.
437 + Gated activation + block-fp8 quant. ``alpha``/``beta`` select the gate
438 + (silu: alpha=1, beta=0; swigluoai: alpha, beta from config).
439 """
386 - GROUP_SIZE = 128
440 + GROUP_SIZE = group_size
441 assert input.ndim == 2
442 if output is not None:
443 assert output.ndim == 2
@@ -431,6 +485,8 @@ def silu_mul_per_token_group_quant_fp8_colmajor(
485 output_scales.stride(-1),
486 eps,
487 clamp_limit if has_clamp else 0.0,
488 + alpha,
489 + beta,
490 fp8_min,
491 fp8_max,
492 use_ue8m0,
@@ -1015,9 +1071,10 @@ def deepgemm_post_process_fp8_weight_block(
1071 f"to be torch.float8_e4m3fn, got {wq.dtype} instead."
1072 )
1073
1018 - if ws.dtype == torch.float8_e8m0fnu:
1019 - # Scales already in E8M0 from checkpoint — upcast to fp32
1020 - # and skip requantization (weights already have power-of-two scales).
1074 + if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
1075 + # Scales already in E8M0 from checkpoint (float8_e8m0fnu, or raw E8M0
1076 + # bits as uint8 for MXFP8) — upcast to fp32 and skip requantization
1077 + # (weights already have power-of-two scales).
1078 ws = _upcast_e8m0_to_fp32(ws)
1079 else:
1080 assert ws.dtype == torch.float32, (
@@ -1057,7 +1114,8 @@ def deepgemm_post_process_fp8_weight_block(
1114 ws = ws.unsqueeze(0)
1115
1116 # From https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/utils/layout.hpp#L46
1060 - recipe = (1, 128, 128)
1117 + # (1, block_n, block_k): (1, 128, 128) for FP8 block, (1, 1, 32) for MXFP8.
1118 + recipe = (1, quant_block_shape[0], quant_block_shape[1])
1119
1120 # Ref : https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/apis/gemm.hpp
1121 # DeepGemm uses the `transform_sf_into_required_layout` function to
vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
+97
@@ -84,6 +84,92 @@ def _mxfp8_e4m3_quantize_torch(
84 return x_fp8, scales_uint8
85
86
87 +def _mxfp8_quant_triton_kernel():
88 + """Lazily-built Triton kernel: per-32-block E8M0 scale + FP8-E4M3 quant.
89 +
90 + Fuses what ``_mxfp8_e4m3_quantize_torch`` does in several elementwise passes
91 + into one launch. Each program handles ``[BLOCK_M, 32]`` (one MX block).
92 + """
93 + from vllm.triton_utils import tl, triton
94 +
95 + @triton.jit
96 + def _kernel(
97 + x_ptr,
98 + xq_ptr,
99 + s_ptr,
100 + M,
101 + K,
102 + sxm,
103 + sxk,
104 + sqm,
105 + sqk,
106 + ssm,
107 + ssk,
108 + BLOCK_M: tl.constexpr,
109 + ):
110 + pid_m = tl.program_id(0)
111 + pid_b = tl.program_id(1) # which 32-element block along K
112 + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
113 + offs_k = pid_b * 32 + tl.arange(0, 32)
114 + m_mask = offs_m < M
115 + x = tl.load(
116 + x_ptr + offs_m[:, None] * sxm + offs_k[None, :] * sxk,
117 + mask=m_mask[:, None],
118 + other=0.0,
119 + ).to(tl.float32)
120 + amax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-30) # [BLOCK_M]
121 + sb = tl.floor(tl.log2(amax)) + 127.0
122 + sb = tl.minimum(tl.maximum(sb, 0.0), 254.0)
123 + descale = tl.exp2(sb - 127.0)
124 + xq = (x / descale[:, None]).to(xq_ptr.dtype.element_ty)
125 + tl.store(
126 + xq_ptr + offs_m[:, None] * sqm + offs_k[None, :] * sqk,
127 + xq,
128 + mask=m_mask[:, None],
129 + )
130 + tl.store(s_ptr + offs_m * ssm + pid_b * ssk, sb.to(tl.uint8), mask=m_mask)
131 +
132 + return _kernel
133 +
134 +
135 +_MXFP8_QUANT_KERNEL = None
136 +
137 +
138 +def _mxfp8_e4m3_quantize_triton(
139 + x: torch.Tensor,
140 +) -> tuple[torch.Tensor, torch.Tensor]:
141 + """Fused 2D MXFP8 quant (non-swizzled, row-major [M, K//32] scales)."""
142 + from vllm.triton_utils import triton
143 +
144 + global _MXFP8_QUANT_KERNEL
145 + if _MXFP8_QUANT_KERNEL is None:
146 + _MXFP8_QUANT_KERNEL = _mxfp8_quant_triton_kernel()
147 +
148 + M, K = x.shape
149 + x = x.contiguous()
150 + xq = torch.empty((M, K), dtype=MXFP8_VALUE_DTYPE, device=x.device)
151 + scales = torch.empty(
152 + (M, K // MXFP8_BLOCK_SIZE), dtype=MXFP8_SCALE_DTYPE, device=x.device
153 + )
154 + BLOCK_M = 64
155 + grid = (triton.cdiv(M, BLOCK_M), K // MXFP8_BLOCK_SIZE)
156 + _MXFP8_QUANT_KERNEL[grid](
157 + x,
158 + xq,
159 + scales,
160 + M,
161 + K,
162 + x.stride(0),
163 + x.stride(1),
164 + xq.stride(0),
165 + xq.stride(1),
166 + scales.stride(0),
167 + scales.stride(1),
168 + BLOCK_M=BLOCK_M,
169 + )
170 + return xq, scales
171 +
172 +
173 def _mxfp8_e4m3_quantize_impl(
174 x: torch.Tensor,
175 is_sf_swizzled_layout: bool = False,
@@ -103,6 +189,17 @@ def _mxfp8_e4m3_quantize_impl(
189 x_scales = x_scales.view(x.size(0), -1)
190 return x_q, x_scales
191
192 + # ROCm: a single fused Triton kernel beats the multi-pass torch path for the
193 + # common 2D, non-swizzled activation-quant case (used by the native MX
194 + # linear/MoE). Falls back to torch otherwise (3D weights, swizzled layout).
195 + if (
196 + current_platform.is_rocm()
197 + and not is_sf_swizzled_layout
198 + and x.ndim == 2
199 + and x.shape[-1] % MXFP8_BLOCK_SIZE == 0
200 + ):
201 + return _mxfp8_e4m3_quantize_triton(x)
202 +
203 return _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout)
204
205
vllm/model_executor/models/registry.py
+9
@@ -164,6 +164,10 @@ _TEXT_GENERATION_MODELS = {
164 "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
165 "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
166 "MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"),
167 + "MiniMaxM3SparseForCausalLM": (
168 + "vllm.models.minimax_m3",
169 + "MiniMaxM3SparseForCausalLM",
170 + ),
171 "Ministral3ForCausalLM": ("mistral", "MistralForCausalLM"),
172 "MistralForCausalLM": ("mistral", "MistralForCausalLM"),
173 "MistralLarge3ForCausalLM": ("mistral_large_3", "MistralLarge3ForCausalLM"),
@@ -483,6 +487,10 @@ _MULTIMODAL_MODELS = {
487 "MantisForConditionalGeneration": ("llava", "MantisForConditionalGeneration"),
488 "MiDashengLMModel": ("midashenglm", "MiDashengLMModel"),
489 "MiMoV2OmniForCausalLM": ("mimo_v2_omni", "MiMoV2OmniForCausalLM"),
490 + "MiniMaxM3SparseForConditionalGeneration": (
491 + "vllm.models.minimax_m3",
492 + "MiniMaxM3SparseForConditionalGeneration",
493 + ),
494 "MiniMaxVL01ForConditionalGeneration": (
495 "minimax_vl_01",
496 "MiniMaxVL01ForConditionalGeneration",
@@ -620,6 +628,7 @@ _SPECULATIVE_DECODING_MODELS = {
628 "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
629 "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
630 "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"),
631 + "MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"),
632 "Gemma4MTPModel": ("gemma4_mtp", "Gemma4MTP"),
633 "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),
634 "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),
vllm/model_executor/warmup/kernel_warmup.py
+6
@@ -53,6 +53,10 @@ def _resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path:
53
54
55 def kernel_warmup(worker: "Worker"):
56 + from vllm.model_executor.warmup.minimax_m3_msa_warmup import (
57 + minimax_m3_msa_warmup,
58 + )
59 +
60 # Deep GEMM warmup
61 do_deep_gemm_warmup = (
62 envs.VLLM_USE_DEEP_GEMM
@@ -64,6 +68,8 @@ def kernel_warmup(worker: "Worker"):
68 max_tokens = worker.scheduler_config.max_num_batched_tokens
69 deep_gemm_warmup(model, max_tokens)
70
71 + minimax_m3_msa_warmup(worker)
72 +
73 enable_flashinfer_autotune = (
74 worker.vllm_config.kernel_config.enable_flashinfer_autotune
75 )
vllm/model_executor/warmup/minimax_m3_msa_warmup.py new
+43
@@ -0,0 +1,43 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import TYPE_CHECKING
5 +
6 +from vllm.logger import init_logger
7 +from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention
8 +from vllm.platforms import current_platform
9 +from vllm.tracing import instrument
10 +
11 +if TYPE_CHECKING:
12 + from vllm.v1.worker.gpu_worker import Worker
13 +
14 +logger = init_logger(__name__)
15 +
16 +
17 +@instrument(span_name="MiniMax M3 MSA warmup")
18 +def minimax_m3_msa_warmup(worker: "Worker") -> None:
19 + sparse_module = next(
20 + (
21 + module
22 + for module in worker.get_model().modules()
23 + if isinstance(module, MiniMaxM3SparseAttention)
24 + ),
25 + None,
26 + )
27 + if sparse_module is None:
28 + return
29 + if not (
30 + current_platform.is_cuda() and current_platform.is_device_capability_family(100)
31 + ):
32 + return
33 +
34 + logger.info("Warming up MiniMax M3 MSA kernels.")
35 +
36 + # Cover sparse prefill through the normal model path.
37 + worker.model_runner._dummy_run(
38 + num_tokens=16,
39 + skip_eplb=True,
40 + is_profile=True,
41 + force_attention=True,
42 + create_mixed_batch=True,
43 + )
vllm/models/minimax_m3/__init__.py new
+33
@@ -0,0 +1,33 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MiniMax M3 model — hardware-isolated entry point.
4 +
5 +The implementation lives under ``nvidia/`` and ``amd/``; this module picks the
6 +right one for the current platform and re-exports the public classes used by
7 +the model registry. (Mirrors ``vllm.models.deepseek_v4``.)
8 +"""
9 +
10 +from typing import TYPE_CHECKING
11 +
12 +from vllm.platforms import current_platform
13 +
14 +# The NVIDIA branch is the static default that type-checkers see; the ROCm
15 +# branch overrides it at runtime (kept type-compatible via type: ignore).
16 +if TYPE_CHECKING or not current_platform.is_rocm():
17 + from .nvidia.model import (
18 + MiniMaxM3SparseForCausalLM,
19 + MiniMaxM3SparseForConditionalGeneration,
20 + )
21 + from .nvidia.mtp import MiniMaxM3MTP
22 +else:
23 + from .amd.model import ( # type: ignore[assignment]
24 + MiniMaxM3SparseForCausalLM,
25 + MiniMaxM3SparseForConditionalGeneration,
26 + )
27 + from .amd.mtp import MiniMaxM3MTP # type: ignore[assignment]
28 +
29 +__all__ = [
30 + "MiniMaxM3MTP",
31 + "MiniMaxM3SparseForCausalLM",
32 + "MiniMaxM3SparseForConditionalGeneration",
33 +]
vllm/models/minimax_m3/amd/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/minimax_m3/amd/model.py new
+1216
@@ -0,0 +1,1216 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only MiniMax M3 (text backbone) model — AMD ROCm implementation.
4 +
5 +Self-contained per-platform impl (mirrors ``deepseek_v4/amd``). It is identical
6 +to ``../nvidia/model.py`` except for RMS normalization: FlashInfer's Gemma
7 +RMSNorm kernels are CUDA-only, so ``MiniMAXGemmaRMSNorm`` here uses a native
8 +(FlashInfer-free) implementation.
9 +
10 +The MiniMax-M3-preview config selects a single set of branches:
11 + * qk_norm_type == "per_head"
12 + * hidden_act == "swigluoai"
13 + * use_gemma_norm == True -> Gemma-style RMSNorm everywhere
14 + * attention_output_gate == False
15 + * scoring_func == "sigmoid" with a routing-bias correction term
16 + * sparse_attention_config present -> a subset of layers run the extra
17 + "index" attention branch.
18 +"""
19 +
20 +from collections.abc import Iterable
21 +
22 +import torch
23 +from torch import nn
24 +from transformers import PretrainedConfig
25 +
26 +from vllm import _custom_ops as ops
27 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
28 +from vllm.config import (
29 + CacheConfig,
30 + VllmConfig,
31 + get_current_vllm_config,
32 +)
33 +from vllm.distributed import get_tensor_model_parallel_world_size
34 +from vllm.forward_context import get_forward_context
35 +from vllm.model_executor.layers.attention import Attention
36 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
37 +from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
38 + fused_allreduce_gemma_rms_norm,
39 +)
40 +from vllm.model_executor.layers.fused_moe import (
41 + FusedMoE,
42 + GateLinear,
43 + fused_moe_make_expert_params_mapping,
44 +)
45 +from vllm.model_executor.layers.linear import (
46 + MergedColumnParallelLinear,
47 + MinimaxM3QKVParallelLinearWithIndexer,
48 + QKVParallelLinear,
49 + RowParallelLinear,
50 +)
51 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
52 +from vllm.model_executor.layers.quantization import QuantizationConfig
53 +from vllm.model_executor.layers.rotary_embedding import get_rope
54 +from vllm.model_executor.layers.vocab_parallel_embedding import (
55 + ParallelLMHead,
56 + VocabParallelEmbedding,
57 +)
58 +from vllm.model_executor.model_loader.weight_utils import (
59 + default_weight_loader,
60 + maybe_remap_kv_scale_name,
61 +)
62 +from vllm.model_executor.models.interfaces import (
63 + EagleModelMixin,
64 + MultiModalEmbeddings,
65 + SupportsEagle3,
66 + SupportsMultiModal,
67 +)
68 +from vllm.model_executor.models.utils import (
69 + AutoWeightsLoader,
70 + WeightsMapper,
71 + init_vllm_registered_model,
72 + is_pp_missing_parameter,
73 + make_layers,
74 + maybe_prefix,
75 +)
76 +from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model
77 +from vllm.models.minimax_m3.amd.ops import (
78 + gemma_fused_add_rmsnorm,
79 + gemma_rmsnorm,
80 + swiglu_oai_split,
81 +)
82 +from vllm.models.minimax_m3.common.indexer import MiniMaxM3Indexer
83 +from vllm.models.minimax_m3.common.mm_preprocess import (
84 + MiniMaxM3VLDummyInputsBuilder,
85 + MiniMaxM3VLMultiModalProcessor,
86 + MiniMaxM3VLProcessingInfo,
87 +)
88 +from vllm.models.minimax_m3.common.sparse_attention import (
89 + MiniMaxM3SparseBackend,
90 + MiniMaxM3SparseImpl,
91 + select_main_impl_cls,
92 +)
93 +from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
94 +from vllm.multimodal import MULTIMODAL_REGISTRY
95 +from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
96 +from vllm.v1.kv_cache_interface import (
97 + FullAttentionSpec,
98 + KVCacheSpec,
99 + get_kv_quant_mode,
100 +)
101 +
102 +
103 +def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]:
104 + """Layer ids whose attention runs the extra sparse "index" branch."""
105 + cfg = getattr(config, "sparse_attention_config", None)
106 + if not cfg:
107 + return set()
108 + freq = cfg.get("sparse_attention_freq")
109 + if freq is None:
110 + return set()
111 + return {i for i, f in enumerate(freq) if f != 0}
112 +
113 +
114 +def _is_moe_layer(config: PretrainedConfig, layer_id: int) -> bool:
115 + """Whether this layer's MLP is a sparse MoE block (vs a dense MLP)."""
116 + moe_layer_freq = getattr(config, "moe_layer_freq", None)
117 + if moe_layer_freq is None:
118 + return True
119 + return moe_layer_freq[layer_id] != 0
120 +
121 +
122 +def _build_rotary_emb(config: PretrainedConfig, head_dim: int):
123 + """Build the (partial NeoX) RoPE, honoring an optional ``rope_scaling`` config.
124 +
125 + Without scaling the cos/sin cache is sized to ``max_position_embeddings``
126 + (524288 native); a request whose positions exceed that reads the cache out of
127 + bounds and the worker hard-crashes (no Python traceback). When ``rope_scaling``
128 + is set (e.g. YaRN ``factor: 2`` to reach 1M), thread it into ``get_rope`` so the
129 + proper scaled embedding is built and its cache covers
130 + ``original_max_position_embeddings * factor`` positions. Default behavior
131 + (no scaling) is unchanged. Shared by the dense and sparse attention layers, and
132 + the index branch reuses the returned module.
133 +
134 + Note: for the VL checkpoint, set ``rope_scaling`` on the *text* config
135 + (``--hf-overrides '{"text_config":{"rope_scaling":{...}}}'``) -- that is the
136 + config the decoder reads here; a top-level override does not reach it.
137 + """
138 + rope_parameters = {
139 + "rope_theta": config.rope_theta,
140 + "partial_rotary_factor": config.partial_rotary_factor,
141 + }
142 + max_position = config.max_position_embeddings
143 + rope_scaling = getattr(config, "rope_scaling", None)
144 + if rope_scaling:
145 + rope_parameters.update(rope_scaling)
146 + # HF uses "rope_type" (older configs: "type"); get_rope reads "rope_type".
147 + if "rope_type" not in rope_parameters and "type" in rope_scaling:
148 + rope_parameters["rope_type"] = rope_scaling["type"]
149 + rope_parameters.setdefault(
150 + "original_max_position_embeddings", config.max_position_embeddings
151 + )
152 + factor = float(rope_scaling.get("factor", 1.0))
153 + # Cover the extended range (informational for get_rope's default branch;
154 + # the YaRN embedding sizes its own cache from original * factor).
155 + max_position = int(rope_parameters["original_max_position_embeddings"] * factor)
156 + return get_rope(
157 + head_dim,
158 + max_position=max_position,
159 + rope_parameters=rope_parameters,
160 + )
161 +
162 +
163 +class MiniMAXGemmaRMSNorm(nn.Module):
164 + """Gemma-style RMS normalization (native ROCm implementation).
165 +
166 + Normalizes in fp32 and scales by ``(1 + weight)`` — numerically equivalent
167 + to the FlashInfer ``gemma_rmsnorm`` / ``gemma_fused_add_rmsnorm`` kernels
168 + used in the NVIDIA path, which are unavailable on ROCm. When ``residual`` is
169 + given, the fused add + norm returns the updated ``(normed, residual)`` pair.
170 +
171 + The fp32 normalize + scale + (optional) residual-add run in a single fused
172 + Triton pass (``amd.ops.gemma_rmsnorm`` / ``gemma_fused_add_rmsnorm``) instead
173 + of a chain of elementwise PyTorch kernels.
174 + """
175 +
176 + def __init__(
177 + self,
178 + hidden_size: int,
179 + eps: float = 1e-6,
180 + ) -> None:
181 + super().__init__()
182 + self.weight = nn.Parameter(torch.zeros(hidden_size))
183 + self.variance_epsilon = eps
184 +
185 + def forward(
186 + self,
187 + x: torch.Tensor,
188 + residual: torch.Tensor | None = None,
189 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
190 + if residual is None:
191 + return gemma_rmsnorm(x, self.weight, self.variance_epsilon)
192 + return gemma_fused_add_rmsnorm(x, residual, self.weight, self.variance_epsilon)
193 +
194 +
195 +class MiniMaxM3MLP(nn.Module):
196 + """Dense SwiGLU-OAI MLP (used by the leading dense layers)."""
197 +
198 + def __init__(
199 + self,
200 + config: PretrainedConfig,
201 + intermediate_size: int,
202 + quant_config: QuantizationConfig | None = None,
203 + reduce_results: bool = True,
204 + prefix: str = "",
205 + ) -> None:
206 + super().__init__()
207 + self.gate_up_proj = MergedColumnParallelLinear(
208 + config.hidden_size,
209 + [intermediate_size] * 2,
210 + bias=False,
211 + quant_config=quant_config,
212 + prefix=f"{prefix}.gate_up_proj",
213 + )
214 + self.down_proj = RowParallelLinear(
215 + intermediate_size,
216 + config.hidden_size,
217 + bias=False,
218 + quant_config=quant_config,
219 + reduce_results=reduce_results,
220 + prefix=f"{prefix}.down_proj",
221 + )
222 + if config.hidden_act != "swigluoai":
223 + raise ValueError(
224 + f"Unsupported activation: {config.hidden_act}. "
225 + "Only swigluoai is supported."
226 + )
227 + # gate * sigmoid(alpha * gate) * (up + beta), with both halves clamped.
228 + # Kept as our fp32 Triton kernel (not the #22 SWIGLUOAI_UNINTERLEAVE op
229 + # ``silu_and_mul_with_clamp``): that op IS built on ROCm but rounds
230 + # intermediates to bf16 (rel ~3e-3 vs our fp32 ~1e-6), which costs gsm8k
231 + # accuracy since this activation feeds the MXFP8 quant + MoE.
232 + self.swiglu_alpha = config.swiglu_alpha
233 + self.swiglu_beta = config.swiglu_beta
234 + self.swiglu_limit = config.swiglu_limit
235 +
236 + def forward(self, x: torch.Tensor) -> torch.Tensor:
237 + gate_up, _ = self.gate_up_proj(x)
238 + x = swiglu_oai_split(
239 + gate_up,
240 + alpha=self.swiglu_alpha,
241 + beta=self.swiglu_beta,
242 + limit=self.swiglu_limit,
243 + )
244 + x, _ = self.down_proj(x)
245 + return x
246 +
247 +
248 +class MiniMaxM3MoE(nn.Module):
249 + """Sigmoid-routed MoE block with a routing-bias correction and a shared
250 + expert."""
251 +
252 + def __init__(
253 + self,
254 + config: PretrainedConfig,
255 + layer_id: int,
256 + quant_config: QuantizationConfig | None = None,
257 + prefix: str = "",
258 + ) -> None:
259 + super().__init__()
260 + self.tp_size = get_tensor_model_parallel_world_size()
261 + if self.tp_size > config.num_local_experts:
262 + raise ValueError(
263 + f"Tensor parallel size {self.tp_size} is greater than "
264 + f"the number of experts {config.num_local_experts}."
265 + )
266 +
267 + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
268 + self.n_shared_experts = getattr(config, "n_shared_experts", None)
269 +
270 + # Sigmoid routing uses a per-expert score-correction bias for selection.
271 + self.use_routing_bias = getattr(config, "use_routing_bias", False)
272 + if self.use_routing_bias:
273 + self.e_score_correction_bias = nn.Parameter(
274 + torch.empty(config.num_local_experts, dtype=torch.float32)
275 + )
276 + self.e_score_correction_bias.weight_loader = (
277 + MiniMaxM3MoE.ebias_weight_loader
278 + )
279 + else:
280 + self.e_score_correction_bias = None
281 +
282 + # Router weights are stored in fp32; GateLinear upcasts the bf16
283 + # activations and computes the gate in fp32 (fp32 router logits).
284 + self.gate = GateLinear(
285 + config.hidden_size,
286 + config.num_local_experts,
287 + bias=False,
288 + params_dtype=torch.float32,
289 + out_dtype=torch.float32,
290 + prefix=f"{prefix}.gate",
291 + )
292 +
293 + self.shared_experts: MiniMaxM3MLP | None = None
294 + if self.n_shared_experts:
295 + self.shared_experts = MiniMaxM3MLP(
296 + config=config,
297 + intermediate_size=config.intermediate_size * self.n_shared_experts,
298 + quant_config=quant_config,
299 + reduce_results=False,
300 + prefix=f"{prefix}.shared_experts",
301 + )
302 +
303 + self.experts = FusedMoE(
304 + num_experts=config.num_local_experts,
305 + top_k=config.num_experts_per_tok,
306 + hidden_size=config.hidden_size,
307 + intermediate_size=config.intermediate_size,
308 + scoring_func=config.scoring_func,
309 + e_score_correction_bias=self.e_score_correction_bias,
310 + renormalize=True,
311 + activation="swigluoai_uninterleave",
312 + swiglu_limit=config.swiglu_limit,
313 + swiglu_alpha=config.swiglu_alpha,
314 + swiglu_beta=config.swiglu_beta,
315 + routed_scaling_factor=self.routed_scaling_factor,
316 + apply_routed_scale_to_output=True,
317 + router_logits_dtype=self.gate.out_dtype,
318 + shared_experts=self.shared_experts,
319 + quant_config=quant_config,
320 + prefix=f"{prefix}.experts",
321 + )
322 +
323 + @staticmethod
324 + def ebias_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None:
325 + assert param.size() == loaded_weight.size()
326 + param.data.copy_(loaded_weight.to(torch.float32))
327 +
328 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
329 + num_tokens, hidden_dim = hidden_states.shape
330 + hidden_states = hidden_states.view(-1, hidden_dim)
331 +
332 + # router_logits: (num_tokens, n_experts); GateLinear casts to fp32.
333 + router_logits, _ = self.gate(hidden_states)
334 + final_hidden_states = self.experts(
335 + hidden_states=hidden_states, router_logits=router_logits
336 + )
337 +
338 + return final_hidden_states.view(num_tokens, hidden_dim)
339 +
340 +
341 +class MiniMaxM3Attention(nn.Module):
342 + """Dense attention with per-head QK norm and partial RoPE."""
343 +
344 + def __init__(
345 + self,
346 + config: PretrainedConfig,
347 + layer_id: int,
348 + quant_config: QuantizationConfig | None = None,
349 + prefix: str = "",
350 + cache_config: CacheConfig | None = None,
351 + ) -> None:
352 + super().__init__()
353 + self.hidden_size = config.hidden_size
354 + tp_size = get_tensor_model_parallel_world_size()
355 +
356 + self.total_num_heads = config.num_attention_heads
357 + assert self.total_num_heads % tp_size == 0
358 + self.num_heads = self.total_num_heads // tp_size
359 + self.total_num_kv_heads = config.num_key_value_heads
360 + if self.total_num_kv_heads >= tp_size:
361 + assert self.total_num_kv_heads % tp_size == 0
362 + else:
363 + assert tp_size % self.total_num_kv_heads == 0
364 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
365 + self.head_dim = config.head_dim
366 + self.q_size = self.num_heads * self.head_dim
367 + self.kv_size = self.num_kv_heads * self.head_dim
368 + self.scaling = self.head_dim**-0.5
369 +
370 + self.qkv_proj = QKVParallelLinear(
371 + self.hidden_size,
372 + self.head_dim,
373 + self.total_num_heads,
374 + self.total_num_kv_heads,
375 + bias=False,
376 + quant_config=quant_config,
377 + prefix=f"{prefix}.qkv_proj",
378 + )
379 + # reduce_results=False: the attention all-reduce is fused with the
380 + # following post_attention_layernorm (GemmaRMSNorm) in the decoder layer
381 + # via fused_allreduce_gemma_rms_norm.
382 + self.o_proj = RowParallelLinear(
383 + self.total_num_heads * self.head_dim,
384 + self.hidden_size,
385 + bias=False,
386 + reduce_results=False,
387 + quant_config=quant_config,
388 + prefix=f"{prefix}.o_proj",
389 + )
390 +
391 + # Per-head QK norm (qk_norm_type == "per_head", use_gemma_norm == True).
392 + self.q_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
393 + self.k_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
394 +
395 + # Partial RoPE: rotary_dim == head_dim * partial_rotary_factor. Honors
396 + # config.rope_scaling (e.g. YaRN) so long-context positions are covered.
397 + self.rotary_emb = _build_rotary_emb(config, self.head_dim)
398 +
399 + self.attn = Attention(
400 + self.num_heads,
401 + self.head_dim,
402 + self.scaling,
403 + num_kv_heads=self.num_kv_heads,
404 + cache_config=cache_config,
405 + quant_config=quant_config,
406 + prefix=f"{prefix}.attn",
407 + )
408 +
409 + def forward(
410 + self,
411 + positions: torch.Tensor,
412 + hidden_states: torch.Tensor,
413 + ) -> torch.Tensor:
414 + qkv, _ = self.qkv_proj(hidden_states)
415 + # Fused per-head Gemma QK-norm + partial NeoX RoPE on q/k, in place (dense
416 + # mode: no index branch, no KV-cache insert). Matches nvidia/model.py and
417 + # replaces the unfused split -> q_norm/k_norm -> rotary_emb chain; verified
418 + # bit-equivalent on ROCm (q/k rel ~2e-3 bf16 noise, v untouched).
419 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
420 + qkv,
421 + self.q_norm.weight,
422 + self.k_norm.weight,
423 + self.rotary_emb.cos_sin_cache,
424 + positions,
425 + self.num_heads,
426 + self.num_kv_heads,
427 + self.rotary_emb.rotary_dim,
428 + self.q_norm.variance_epsilon,
429 + )
430 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
431 + attn_output = self.attn(q, k, v)
432 + output, _ = self.o_proj(attn_output)
433 + return output
434 +
435 +
436 +class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
437 + """Block-sparse attention layer with the lightning-indexer branch.
438 +
439 + This is a merged attention layer: it owns the projections (qkv + index
440 + q/k), per-head QK norms and RoPE, *and* the attention-backend wiring that a
441 + generic ``Attention`` layer would normally provide — it binds the
442 + ``MiniMaxM3SparseBackend`` + main impl, registers the main paged K/V cache,
443 + and owns the lightning indexer (``MiniMaxM3Indexer``), which holds the
444 + index-key side cache.
445 +
446 + The index branch (index_{q,k}_proj + index_{q,k}_norm) feeds the sparse
447 + top-k block selection. M3 always disables the index value/output
448 + projections (``sparse_disable_index_value`` set for every sparse layer), so
449 + ``index_{v,o}_proj`` are never created.
450 + """
451 +
452 + def __init__(
453 + self,
454 + config: PretrainedConfig,
455 + layer_id: int,
456 + quant_config: QuantizationConfig | None = None,
457 + prefix: str = "",
458 + cache_config: CacheConfig | None = None,
459 + ) -> None:
460 + super().__init__()
461 + self.hidden_size = config.hidden_size
462 + tp_size = get_tensor_model_parallel_world_size()
463 +
464 + self.total_num_heads = config.num_attention_heads
465 + assert self.total_num_heads % tp_size == 0
466 + self.num_heads = self.total_num_heads // tp_size
467 + self.total_num_kv_heads = config.num_key_value_heads
468 + if self.total_num_kv_heads >= tp_size:
469 + assert self.total_num_kv_heads % tp_size == 0
470 + else:
471 + assert tp_size % self.total_num_kv_heads == 0
472 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
473 + self.head_dim = config.head_dim
474 + self.q_size = self.num_heads * self.head_dim
475 + self.kv_size = self.num_kv_heads * self.head_dim
476 + self.scaling = self.head_dim**-0.5
477 +
478 + # Sparse "index" branch dims. index_q has the same head count as the KV
479 + # heads (sparse_num_index_heads == num_key_value_heads), so it shards
480 + # identically -- including replication when tp_size > num_key_value_heads.
481 + sparse_cfg = config.sparse_attention_config
482 + self.total_idx_heads = sparse_cfg["sparse_num_index_heads"]
483 + self.num_idx_heads = self.num_kv_heads
484 + self.idx_head_dim = sparse_cfg["sparse_index_dim"]
485 + self.index_q_size = self.num_idx_heads * self.idx_head_dim
486 +
487 + # Single fused projection: q, k, v, index_q, index_k in one GEMM.
488 + self.qkv_proj = MinimaxM3QKVParallelLinearWithIndexer(
489 + self.hidden_size,
490 + self.head_dim,
491 + self.total_num_heads,
492 + self.total_num_kv_heads,
493 + self.total_idx_heads,
494 + self.idx_head_dim,
495 + bias=False,
496 + quant_config=quant_config,
497 + prefix=f"{prefix}.qkv_proj",
498 + )
499 + # reduce_results=False: the attention all-reduce is fused with the
500 + # following post_attention_layernorm (GemmaRMSNorm) in the decoder layer
501 + # via fused_allreduce_gemma_rms_norm.
502 + self.o_proj = RowParallelLinear(
503 + self.total_num_heads * self.head_dim,
504 + self.hidden_size,
505 + bias=False,
506 + reduce_results=False,
507 + quant_config=quant_config,
508 + prefix=f"{prefix}.o_proj",
509 + )
510 +
511 + # Per-head QK norm (qk_norm_type == "per_head", use_gemma_norm == True).
512 + self.q_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
513 + self.k_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
514 +
515 + # Partial RoPE: rotary_dim == head_dim * partial_rotary_factor. Honors
516 + # config.rope_scaling (e.g. YaRN) so long-context positions are covered.
517 + self.rotary_emb = _build_rotary_emb(config, self.head_dim)
518 +
519 + self.index_q_norm = MiniMAXGemmaRMSNorm(
520 + self.idx_head_dim, eps=config.rms_norm_eps
521 + )
522 + self.index_k_norm = MiniMAXGemmaRMSNorm(
523 + self.idx_head_dim, eps=config.rms_norm_eps
524 + )
525 + self.index_rotary_emb = self.rotary_emb
526 +
527 + # Attention-backend wiring.
528 + vllm_config = get_current_vllm_config()
529 + self.layer_name = f"{prefix}.attn"
530 + self.kv_cache_dtype = (
531 + cache_config.cache_dtype if cache_config is not None else "auto"
532 + )
533 + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
534 + self.kv_cache_dtype, vllm_config.model_config
535 + )
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
541 +
542 + self.attn_backend = MiniMaxM3SparseBackend
543 + # Indexer and main attention are separate impls. On ROCm the SM100 gate
544 + # is always False, so both pick Triton and the index cache stays bf16.
545 + # impl is AttentionImplBase (broader than AttentionLayerBase's annotation).
546 + self.impl: MiniMaxM3SparseImpl = select_main_impl_cls( # type: ignore[assignment]
547 + topk_blocks=sparse_cfg["sparse_topk_blocks"],
548 + kv_cache_dtype=self.kv_cache_dtype,
549 + )(
550 + self.num_heads,
551 + self.head_dim,
552 + self.scaling,
553 + self.num_kv_heads,
554 + kv_cache_dtype=self.kv_cache_dtype,
555 + topk_blocks=sparse_cfg["sparse_topk_blocks"],
556 + sparse_block_size=sparse_cfg["sparse_block_size"],
557 + )
558 + # Self-contained nn.Module: owns its side cache, selects its impl in init
559 + # (Triton on ROCm, where the SM100 gate is always False).
560 + self.indexer = MiniMaxM3Indexer(
561 + num_kv_heads=self.num_kv_heads,
562 + scale=self.scaling,
563 + topk_blocks=sparse_cfg["sparse_topk_blocks"],
564 + sparse_block_size=sparse_cfg["sparse_block_size"],
565 + num_index_heads=self.num_idx_heads,
566 + index_head_dim=self.idx_head_dim,
567 + prefix=self.layer_name,
568 + init_blocks=sparse_cfg.get("sparse_init_block", 0),
569 + local_blocks=sparse_cfg.get("sparse_local_block", 0),
570 + score_type=sparse_cfg.get("sparse_score_type", "max"),
571 + cache_config=cache_config,
572 + )
573 +
574 + # Register the main K/V cache so the KV-cache manager allocates it.
575 + compilation_config = vllm_config.compilation_config
576 + if self.layer_name in compilation_config.static_forward_context:
577 + raise ValueError(f"Duplicate layer name: {self.layer_name}")
578 + compilation_config.static_forward_context[self.layer_name] = self
579 + self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache
580 +
581 + def get_attn_backend(self) -> type[MiniMaxM3SparseBackend]:
582 + return self.attn_backend
583 +
584 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
585 + # Main GQA K/V cache. Block size may change after load, refresh it.
586 + return FullAttentionSpec(
587 + block_size=vllm_config.cache_config.block_size,
588 + num_kv_heads=self.num_kv_heads,
589 + head_size=self.head_dim,
590 + head_size_v=self.head_dim,
591 + dtype=self.kv_cache_torch_dtype,
592 + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
593 + )
594 +
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 +
626 + def forward(
627 + self,
628 + positions: torch.Tensor,
629 + hidden_states: torch.Tensor,
630 + ) -> torch.Tensor:
631 + # Single fused projection emitting [q | k | v | index_q | index_k].
632 + qkv, _ = self.qkv_proj(hidden_states)
633 +
634 + # Horizontally-fused per-head Gemma QK-norm + partial NeoX RoPE on the
635 + # main (q/k) and index (index_q/index_k) branches, all read straight out
636 + # of the single fused ``qkv`` tensor. Once the paged caches are bound the
637 + # kernel also inserts k/v and the index key into them (each with its own
638 + # 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.)
644 + cos_sin_cache = self.rotary_emb.cos_sin_cache
645 + rotary_dim = self.rotary_emb.rotary_dim
646 + eps = self.q_norm.variance_epsilon
647 + num_tokens = qkv.shape[0]
648 +
649 + fwd_slot_mapping = get_forward_context().slot_mapping
650 + if (
651 + not isinstance(fwd_slot_mapping, dict)
652 + or self.layer_name not in fwd_slot_mapping
653 + ):
654 + # Memory-profiling run: caches not yet bound, slot_mapping is empty.
655 + return qkv.new_zeros((num_tokens, self.hidden_size))
656 +
657 + main_slot_mapping = fwd_slot_mapping[self.layer_name]
658 + index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
659 + q = qkv.new_empty((num_tokens, self.q_size))
660 + 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
667 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
668 + qkv,
669 + self.q_norm.weight,
670 + self.k_norm.weight,
671 + cos_sin_cache,
672 + positions,
673 + self.num_heads,
674 + self.num_kv_heads,
675 + rotary_dim,
676 + eps,
677 + self.index_q_norm.weight,
678 + self.index_k_norm.weight,
679 + self.num_idx_heads,
680 + main_slot_mapping,
681 + 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,
684 + self.kv_cache.size(2), # paged-cache block size
685 + q,
686 + index_q,
687 + )
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)
702 +
703 + output = torch.empty_like(q)
704 + attn_output = self._run_attention(q, index_q, output)
705 + output, _ = self.o_proj(attn_output)
706 + return output
707 +
708 + @eager_break_during_capture
709 + def _run_attention(
710 + self,
711 + query: torch.Tensor,
712 + index_query: torch.Tensor,
713 + output: torch.Tensor,
714 + ) -> torch.Tensor:
715 + # Single eager break around both: their split-K kernels read per-request
716 + # metadata and can't be captured into a cudagraph.
717 + topk_idx = self.indexer(index_query)
718 + return self.impl.forward(self, query, self.kv_cache, topk_idx, output)
719 +
720 +
721 +class MiniMaxM3DecoderLayer(nn.Module):
722 + def __init__(
723 + self,
724 + config: PretrainedConfig,
725 + prefix: str,
726 + cache_config: CacheConfig | None = None,
727 + quant_config: QuantizationConfig | None = None,
728 + force_sparse_attn: bool = False,
729 + force_moe: bool = False,
730 + ) -> None:
731 + super().__init__()
732 + self.hidden_size = config.hidden_size
733 + # DecoderLayers are created with `make_layers` which passes the prefix
734 + # with the layer's index.
735 + layer_id = int(prefix.split(sep=".")[-1])
736 + self.layer_id = layer_id
737 +
738 + is_sparse_attention_layer = (
739 + force_sparse_attn or layer_id in _sparse_attention_layer_ids(config)
740 + )
741 +
742 + if is_sparse_attention_layer:
743 + self.self_attn = MiniMaxM3SparseAttention(
744 + config=config,
745 + layer_id=layer_id,
746 + quant_config=quant_config,
747 + prefix=f"{prefix}.self_attn",
748 + cache_config=cache_config,
749 + )
750 + else:
751 + self.self_attn = MiniMaxM3Attention(
752 + config=config,
753 + layer_id=layer_id,
754 + quant_config=quant_config,
755 + prefix=f"{prefix}.self_attn",
756 + cache_config=cache_config,
757 + )
758 +
759 + # Dense layers store the FFN under `mlp`; MoE layers under
760 + # `block_sparse_moe` -- matching the checkpoint's naming.
761 + self.is_moe_layer = force_moe or _is_moe_layer(config, layer_id)
762 + if self.is_moe_layer:
763 + self.block_sparse_moe = MiniMaxM3MoE(
764 + config=config,
765 + layer_id=layer_id,
766 + quant_config=quant_config,
767 + prefix=f"{prefix}.block_sparse_moe",
768 + )
769 + else:
770 + self.mlp = MiniMaxM3MLP(
771 + config=config,
772 + intermediate_size=config.dense_intermediate_size,
773 + quant_config=quant_config,
774 + prefix=f"{prefix}.mlp",
775 + )
776 +
777 + # config.use_gemma_norm is True for M3 -> Gemma-style RMSNorm.
778 + self.input_layernorm = MiniMAXGemmaRMSNorm(
779 + config.hidden_size, eps=config.rms_norm_eps
780 + )
781 + self.post_attention_layernorm = MiniMAXGemmaRMSNorm(
782 + config.hidden_size, eps=config.rms_norm_eps
783 + )
784 +
785 + def forward(
786 + self,
787 + positions: torch.Tensor,
788 + hidden_states: torch.Tensor,
789 + residual: torch.Tensor | None,
790 + ) -> tuple[torch.Tensor, torch.Tensor]:
791 + # Self Attention
792 + if residual is None:
793 + residual = hidden_states
794 + hidden_states = self.input_layernorm(hidden_states)
795 + else:
796 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
797 + hidden_states = self.self_attn(
798 + positions=positions,
799 + hidden_states=hidden_states,
800 + )
801 +
802 + hidden_states, residual = fused_allreduce_gemma_rms_norm(
803 + hidden_states, residual, self.post_attention_layernorm
804 + )
805 + ffn = self.block_sparse_moe if self.is_moe_layer else self.mlp
806 + hidden_states = ffn(hidden_states)
807 + return hidden_states, residual
808 +
809 +
810 +class MiniMaxM3Model(nn.Module, EagleModelMixin):
811 + fall_back_to_pt_during_load = False
812 +
813 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
814 + super().__init__()
815 +
816 + config = vllm_config.model_config.hf_text_config
817 + cache_config = vllm_config.cache_config
818 + quant_config = vllm_config.quant_config
819 + self.config = config
820 +
821 + self.vocab_size = config.vocab_size
822 +
823 + self.embed_tokens = VocabParallelEmbedding(
824 + config.vocab_size,
825 + config.hidden_size,
826 + quant_config=quant_config,
827 + prefix=f"{prefix}.embed_tokens",
828 + )
829 +
830 + self.start_layer, self.end_layer, self.layers = make_layers(
831 + config.num_hidden_layers,
832 + lambda prefix: MiniMaxM3DecoderLayer(
833 + config,
834 + prefix,
835 + cache_config=cache_config,
836 + quant_config=quant_config,
837 + ),
838 + prefix=f"{prefix}.layers",
839 + )
840 +
841 + self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
842 +
843 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
844 + return self.embed_tokens(input_ids)
845 +
846 + def forward(
847 + self,
848 + input_ids: torch.Tensor | None,
849 + positions: torch.Tensor,
850 + inputs_embeds: torch.Tensor | None = None,
851 + ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
852 + if inputs_embeds is not None:
853 + hidden_states = inputs_embeds
854 + else:
855 + hidden_states = self.embed_input_ids(input_ids)
856 + residual = None
857 +
858 + # EAGLE3 is not yet compatible with pipeline parallel
859 + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
860 + for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]):
861 + hidden_states, residual = layer(positions, hidden_states, residual)
862 + self._maybe_add_hidden_state(
863 + aux_hidden_states, idx + 1, hidden_states, residual
864 + )
865 +
866 + hidden_states, _ = self.norm(hidden_states, residual)
867 +
868 + if len(aux_hidden_states) > 0:
869 + return hidden_states, aux_hidden_states
870 + return hidden_states
871 +
872 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
873 + # Checkpoint experts use w1=gate, w2=down, w3=up.
874 + return fused_moe_make_expert_params_mapping(
875 + self,
876 + ckpt_gate_proj_name="w1",
877 + ckpt_down_proj_name="w2",
878 + ckpt_up_proj_name="w3",
879 + num_experts=self.config.num_local_experts,
880 + )
881 +
882 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
883 + # q/k/v_proj -> fused qkv_proj; gate_proj/up_proj -> fused gate_up_proj
884 + # (dense MLP and shared expert). On sparse layers the indexer
885 + # index_q/index_k_proj fold into the same fused qkv_proj
886 + # (MinimaxM3QKVParallelLinearWithIndexer); these entries simply never match on
887 + # dense layers, whose checkpoints have no index_*_proj weights. Leading
888 + # dots keep `q_proj`/`k_proj` from matching `index_q_proj`/`index_k_proj`
889 + # (preceded by `_`, not `.`).
890 + stacked_params_mapping: list[tuple[str, str, int | str]] = [
891 + # (param_name, shard_name, shard_id)
892 + (".qkv_proj", ".q_proj", "q"),
893 + (".qkv_proj", ".k_proj", "k"),
894 + (".qkv_proj", ".v_proj", "v"),
895 + (".qkv_proj", ".index_q_proj", "index_q"),
896 + (".qkv_proj", ".index_k_proj", "index_k"),
897 + (".gate_up_proj", ".gate_proj", 0),
898 + (".gate_up_proj", ".up_proj", 1),
899 + ]
900 +
901 + # (param_name, weight_name, expert_id, shard_id)
902 + expert_params_mapping = self.get_expert_mapping()
903 +
904 + params_dict = dict(self.named_parameters())
905 + loaded_params: set[str] = set()
906 + for name, loaded_weight in weights:
907 + # The MTP module is not modeled yet.
908 + if "mtp." in name:
909 + continue
910 +
911 + # The checkpoint stores block scales as ``weight_scale_inv``; the
912 + # ModelOpt MXFP8 layers expose them as ``weight_scale``.
913 + if "weight_scale_inv" in name:
914 + name = name.replace("weight_scale_inv", "weight_scale")
915 +
916 + for param_name, weight_name, shard_id in stacked_params_mapping:
917 + if weight_name not in name:
918 + continue
919 + # Routed experts (w1/w2/w3) are handled below; don't let the
920 + # stacked mapping rewrite them.
921 + if ("block_sparse_moe.experts." in name) and name not in params_dict:
922 + continue
923 + name = name.replace(weight_name, param_name)
924 + if name.endswith(".bias") and name not in params_dict:
925 + continue
926 + if is_pp_missing_parameter(name, self):
927 + continue
928 + if name not in params_dict:
929 + continue
930 + param = params_dict[name]
931 + weight_loader = param.weight_loader
932 + weight_loader(param, loaded_weight, shard_id)
933 + break
934 + else:
935 + for (
936 + param_name,
937 + weight_name,
938 + expert_id,
939 + expert_shard_id,
940 + ) in expert_params_mapping:
941 + if weight_name not in name:
942 + continue
943 + name = name.replace(weight_name, param_name)
944 + if is_pp_missing_parameter(name, self):
945 + continue
946 + if name not in params_dict:
947 + continue
948 + param = params_dict[name]
949 + weight_loader = param.weight_loader
950 + weight_loader(
951 + param,
952 + loaded_weight,
953 + name,
954 + shard_id=expert_shard_id,
955 + expert_id=expert_id,
956 + )
957 + break
958 + else:
959 + if name.endswith(".bias") and name not in params_dict:
960 + continue
961 + remapped = maybe_remap_kv_scale_name(name, params_dict)
962 + if remapped is None:
963 + continue
964 + name = remapped
965 + if is_pp_missing_parameter(name, self):
966 + continue
967 + # Modules not modeled yet (e.g. attention) are skipped until
968 + # they are ported.
969 + if name not in params_dict:
970 + continue
971 + param = params_dict[name]
972 + weight_loader = getattr(
973 + param, "weight_loader", default_weight_loader
974 + )
975 + weight_loader(param, loaded_weight)
976 + loaded_params.add(name)
977 + return loaded_params
978 +
979 +
980 +class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3):
981 + """MiniMax M3 (sparse/dense backbone) for causal language modeling."""
982 +
983 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
984 + super().__init__()
985 + config = vllm_config.model_config.hf_text_config
986 + quant_config = vllm_config.quant_config
987 + self.config = config
988 + self.quant_config = quant_config
989 + self.model = MiniMaxM3Model(
990 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
991 + )
992 + self.lm_head = ParallelLMHead(
993 + config.vocab_size,
994 + config.hidden_size,
995 + quant_config=quant_config,
996 + prefix=maybe_prefix(prefix, "lm_head"),
997 + )
998 + self.logits_processor = LogitsProcessor(config.vocab_size)
999 +
1000 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1001 + return self.model.embed_input_ids(input_ids)
1002 +
1003 + def forward(
1004 + self,
1005 + input_ids: torch.Tensor | None,
1006 + positions: torch.Tensor,
1007 + inputs_embeds: torch.Tensor | None = None,
1008 + **kwargs,
1009 + ) -> torch.Tensor:
1010 + return self.model(input_ids, positions, inputs_embeds)
1011 +
1012 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
1013 + return self.logits_processor(self.lm_head, hidden_states)
1014 +
1015 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1016 + return self.model.get_expert_mapping()
1017 +
1018 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1019 + loader = AutoWeightsLoader(self)
1020 + return loader.load_weights(weights)
1021 +
1022 +
1023 +# TODO(refactor): this VL wrapper is platform-agnostic and byte-identical to the
1024 +# NVIDIA copy — it only orchestrates the shared vision tower + the per-platform
1025 +# language model (resolved via ``init_vllm_registered_model``). Hoist it into
1026 +# ``common/`` to drop the amd/nvidia duplication once the split stabilizes.
1027 +@MULTIMODAL_REGISTRY.register_processor(
1028 + MiniMaxM3VLMultiModalProcessor,
1029 + info=MiniMaxM3VLProcessingInfo,
1030 + dummy_inputs=MiniMaxM3VLDummyInputsBuilder,
1031 +)
1032 +class MiniMaxM3SparseForConditionalGeneration(
1033 + nn.Module, SupportsMultiModal, SupportsEagle3
1034 +):
1035 + """Top-level (VL) entry point for MiniMax M3.
1036 +
1037 + Owns the shared MiniMax-M3 vision tower on ROCm and delegates text
1038 + generation to the AMD language-model path.
1039 + """
1040 +
1041 + # The vision tower runs replicated per rank under ``--mm-encoder-tp-mode
1042 + # data``; ``run_dp_sharded_mrope_vision_model`` shards the work across
1043 + # ranks (see ``_process_image_input`` / ``_process_video_input``).
1044 + supports_encoder_tp_data = True
1045 +
1046 + hf_to_vllm_mapper = WeightsMapper(
1047 + orig_to_new_prefix={
1048 + "multi_modal_projector.": "vision_tower.multi_modal_projector.",
1049 + "patch_merge_mlp.": "vision_tower.patch_merge_mlp.",
1050 + },
1051 + orig_to_new_substr={
1052 + ".mlp.fc1.": ".fc1.",
1053 + ".mlp.fc2.": ".fc2.",
1054 + },
1055 + )
1056 +
1057 + @classmethod
1058 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
1059 + if modality == "image":
1060 + return MiniMaxM3VLProcessingInfo.IMAGE_TOKEN
1061 + if modality == "video":
1062 + return MiniMaxM3VLProcessingInfo.VIDEO_TOKEN
1063 + raise ValueError(f"Unsupported modality: {modality!r}")
1064 +
1065 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1066 + super().__init__()
1067 + config = vllm_config.model_config.hf_config
1068 + self.config = config
1069 + self.quant_config = vllm_config.quant_config
1070 + self.multimodal_config = vllm_config.model_config.multimodal_config
1071 + assert self.multimodal_config is not None
1072 + self.use_data_parallel = self.multimodal_config.mm_encoder_tp_mode == "data"
1073 +
1074 + text_hidden_size = getattr(config.text_config, "hidden_size", None)
1075 + assert text_hidden_size is not None, "text_config.hidden_size is required"
1076 + projector_hidden_size = getattr(config, "projector_hidden_size", None)
1077 +
1078 + with self._mark_tower_model(vllm_config, {"image", "video"}):
1079 + vision_config = config.vision_config
1080 + self.vision_tower = MiniMaxVLVisionModel(
1081 + config=PretrainedConfig.from_dict(vision_config),
1082 + text_hidden_size=text_hidden_size,
1083 + projector_hidden_size=projector_hidden_size,
1084 + quant_config=self.quant_config,
1085 + prefix=maybe_prefix(prefix, "vision_tower"),
1086 + )
1087 +
1088 + self.language_model = init_vllm_registered_model(
1089 + vllm_config=vllm_config,
1090 + hf_config=config.text_config,
1091 + prefix=maybe_prefix(prefix, "language_model"),
1092 + architectures=["MiniMaxM3SparseForCausalLM"],
1093 + )
1094 +
1095 + def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None:
1096 + pixel_values = kwargs.pop("pixel_values", None)
1097 + image_grid_thw = kwargs.pop("image_grid_thw", None)
1098 + if pixel_values is None:
1099 + return None
1100 + return {"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}
1101 +
1102 + def _parse_and_validate_video_input(self, **kwargs: object) -> dict | None:
1103 + pixel_values_videos = kwargs.pop("pixel_values_videos", None)
1104 + video_grid_thw = kwargs.pop("video_grid_thw", None)
1105 + if pixel_values_videos is None:
1106 + return None
1107 + return {
1108 + "pixel_values_videos": pixel_values_videos,
1109 + "video_grid_thw": video_grid_thw,
1110 + }
1111 +
1112 + def _process_image_input(self, image_input: dict) -> tuple[torch.Tensor, ...]:
1113 + pixel_values: torch.Tensor = image_input["pixel_values"].type(
1114 + self.vision_tower.dtype
1115 + )
1116 + grid_thw: torch.Tensor = image_input["image_grid_thw"]
1117 + assert grid_thw.ndim == 2
1118 +
1119 + if self.use_data_parallel:
1120 + # Already returns a per-item tuple of embeddings.
1121 + return run_dp_sharded_mrope_vision_model(
1122 + self.vision_tower,
1123 + pixel_values,
1124 + grid_thw.tolist(),
1125 + rope_type="rope_3d",
1126 + )
1127 +
1128 + image_embeds = self.vision_tower(
1129 + pixel_values=pixel_values,
1130 + grid_thw=grid_thw.tolist(),
1131 + )
1132 +
1133 + # Split the concatenated output into one tensor per image item.
1134 + merge_size = self.vision_tower.spatial_merge_size
1135 + sizes = (grid_thw.prod(-1) // (merge_size * merge_size)).tolist()
1136 + return image_embeds.split(sizes)
1137 +
1138 + def _process_video_input(self, video_input: dict) -> tuple[torch.Tensor, ...]:
1139 + pixel_values: torch.Tensor = video_input["pixel_values_videos"].type(
1140 + self.vision_tower.dtype
1141 + )
1142 + grid_thw: torch.Tensor = video_input["video_grid_thw"]
1143 + assert grid_thw.ndim == 2
1144 +
1145 + if self.use_data_parallel:
1146 + # Already returns a per-item tuple of embeddings.
1147 + return run_dp_sharded_mrope_vision_model(
1148 + self.vision_tower,
1149 + pixel_values,
1150 + grid_thw.tolist(),
1151 + rope_type="rope_3d",
1152 + )
1153 +
1154 + video_embeds = self.vision_tower(
1155 + pixel_values=pixel_values,
1156 + grid_thw=grid_thw.tolist(),
1157 + )
1158 +
1159 + # Split the concatenated output into one tensor per video item.
1160 + merge_size = self.vision_tower.spatial_merge_size
1161 + sizes = (grid_thw.prod(-1) // (merge_size * merge_size)).tolist()
1162 + return video_embeds.split(sizes)
1163 +
1164 + def _parse_and_validate_multimodal_inputs(
1165 + self, **kwargs: object
1166 + ) -> dict[str, dict]:
1167 + mm_input_by_modality: dict[str, dict] = {}
1168 + for input_key in kwargs:
1169 + if input_key == "pixel_values" and "image" not in mm_input_by_modality:
1170 + image_input = self._parse_and_validate_image_input(**kwargs)
1171 + if image_input is not None:
1172 + mm_input_by_modality["image"] = image_input
1173 + if (
1174 + input_key == "pixel_values_videos"
1175 + and "video" not in mm_input_by_modality
1176 + ):
1177 + video_input = self._parse_and_validate_video_input(**kwargs)
1178 + if video_input is not None:
1179 + mm_input_by_modality["video"] = video_input
1180 + return mm_input_by_modality
1181 +
1182 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
1183 + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
1184 + if not mm_input_by_modality:
1185 + return []
1186 +
1187 + multimodal_embeddings: list[torch.Tensor] = []
1188 + for modality in mm_input_by_modality:
1189 + multimodal_input = mm_input_by_modality[modality]
1190 + if modality == "image":
1191 + image_embeddings = self._process_image_input(multimodal_input)
1192 + multimodal_embeddings.extend(image_embeddings)
1193 + if modality == "video":
1194 + video_embeddings = self._process_video_input(multimodal_input)
1195 + multimodal_embeddings.extend(video_embeddings)
1196 +
1197 + return tuple(multimodal_embeddings)
1198 +
1199 + def forward(
1200 + self,
1201 + input_ids: torch.Tensor | None,
1202 + positions: torch.Tensor,
1203 + inputs_embeds: torch.Tensor | None = None,
1204 + **kwargs,
1205 + ) -> torch.Tensor:
1206 + return self.language_model(input_ids, positions, inputs_embeds)
1207 +
1208 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
1209 + return self.language_model.compute_logits(hidden_states)
1210 +
1211 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1212 + return self.language_model.get_expert_mapping()
1213 +
1214 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1215 + loader = AutoWeightsLoader(self)
1216 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/models/minimax_m3/amd/mtp.py new
+330
@@ -0,0 +1,330 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MiniMax M3 MTP (multi-token prediction) draft model -- ROCm/AMD variant.
4 +
5 +Byte-identical to ``nvidia/mtp.py`` except this file lives under ``amd/`` so its
6 +``from .model import ...`` resolves to the self-contained AMD model (native Gemma
7 +RMSNorm, native MXFP8 MoE, Triton sparse attention). The MTP logic is
8 +platform-agnostic. (Mirrors ``vllm.models.deepseek_v4.amd.mtp``.)
9 +
10 +TODO(future, separate diff): since this is byte-identical to ``nvidia/mtp.py``,
11 +both copies could be consolidated into a single ``common/mtp.py`` that dispatches
12 +its model import (``..amd.model`` vs ``..nvidia.model``) via
13 +``current_platform.is_rocm()`` -- the same dispatch ``minimax_m3/__init__.py``
14 +uses. This was prototyped and VERIFIED working (``MiniMaxM3MTP`` resolves through
15 +``common.mtp`` to the AMD decoder layer / RMSNorm on ROCm), but it deletes the
16 +upstream ``nvidia/mtp.py`` and touches the NVIDIA load path, so it is deferred to
17 +a dedicated refactor diff to keep this AMD-enablement change NVIDIA-untouched.
18 +"""
19 +
20 +from collections.abc import Iterable
21 +
22 +import regex as re
23 +import torch
24 +import torch.nn as nn
25 +
26 +from vllm.config import VllmConfig
27 +from vllm.model_executor.layers.fused_moe import (
28 + fused_moe_make_expert_params_mapping,
29 +)
30 +from vllm.model_executor.layers.linear import (
31 + ReplicatedLinear,
32 +)
33 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
34 +from vllm.model_executor.layers.vocab_parallel_embedding import (
35 + ParallelLMHead,
36 + VocabParallelEmbedding,
37 +)
38 +from vllm.model_executor.model_loader.weight_utils import (
39 + default_weight_loader,
40 + maybe_remap_kv_scale_name,
41 +)
42 +from vllm.model_executor.models.utils import (
43 + maybe_prefix,
44 +)
45 +from vllm.sequence import IntermediateTensors
46 +
47 +from .model import (
48 + MiniMAXGemmaRMSNorm,
49 + MiniMaxM3DecoderLayer,
50 +)
51 +
52 +
53 +class MiniMaxM3MultiTokenPredictorLayer(nn.Module):
54 + def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
55 + super().__init__()
56 +
57 + assert vllm_config.speculative_config is not None
58 + config = vllm_config.speculative_config.draft_model_config.hf_config
59 + cache_config = vllm_config.cache_config
60 + quant_config = vllm_config.quant_config
61 +
62 + self.enorm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
63 + self.hnorm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
64 + self.eh_proj = ReplicatedLinear(
65 + config.hidden_size * 2,
66 + config.hidden_size,
67 + bias=False,
68 + quant_config=quant_config,
69 + prefix=f"{prefix}.eh_proj",
70 + )
71 + self.transformer_layer = MiniMaxM3DecoderLayer(
72 + config=config,
73 + prefix=prefix,
74 + cache_config=cache_config,
75 + quant_config=quant_config,
76 + force_sparse_attn=True,
77 + force_moe=True,
78 + )
79 + self.final_layernorm = MiniMAXGemmaRMSNorm(
80 + config.hidden_size, eps=config.rms_norm_eps
81 + )
82 +
83 + def forward(
84 + self,
85 + input_ids: torch.Tensor,
86 + positions: torch.Tensor,
87 + previous_hidden_states: torch.Tensor,
88 + inputs_embeds: torch.Tensor | None = None,
89 + spec_step_index: int = 0,
90 + ) -> torch.Tensor:
91 + assert inputs_embeds is not None
92 + # Mask out inputs at position 0, as not needed by MTP.
93 + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
94 +
95 + # Combine the normalized token embeddings with the normalized
96 + # previous hidden states.
97 + inputs_embeds = self.enorm(inputs_embeds)
98 + previous_hidden_states = self.hnorm(previous_hidden_states)
99 + hidden_states, _ = self.eh_proj(
100 + torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
101 + )
102 +
103 + # Apply transformer layer.
104 + hidden_states, residual = self.transformer_layer(
105 + positions=positions,
106 + hidden_states=hidden_states,
107 + residual=None,
108 + )
109 +
110 + hidden_states += residual
111 + return hidden_states
112 +
113 +
114 +class MiniMaxM3MultiTokenPredictor(nn.Module):
115 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
116 + super().__init__()
117 +
118 + assert vllm_config.speculative_config is not None
119 + # Use the draft (MTP) config, not the target model's. This is flat for a
120 + # standalone checkpoint, and the promoted text_config for a bundled one.
121 + config = vllm_config.speculative_config.draft_model_config.hf_config
122 + self.num_mtp_layers = config.num_mtp_modules
123 + self.layers = torch.nn.ModuleDict(
124 + {
125 + str(idx): MiniMaxM3MultiTokenPredictorLayer(
126 + vllm_config, f"{prefix}.layers.{idx}"
127 + )
128 + for idx in range(self.num_mtp_layers)
129 + }
130 + )
131 + self.embed_tokens = VocabParallelEmbedding(
132 + config.vocab_size,
133 + config.hidden_size,
134 + prefix=maybe_prefix(prefix, "embed_tokens"),
135 + )
136 +
137 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
138 + return self.embed_tokens(input_ids)
139 +
140 + def forward(
141 + self,
142 + input_ids: torch.Tensor,
143 + positions: torch.Tensor,
144 + previous_hidden_states: torch.Tensor,
145 + inputs_embeds: torch.Tensor | None = None,
146 + spec_step_idx: int = 0,
147 + ) -> torch.Tensor:
148 + if inputs_embeds is None:
149 + inputs_embeds = self.embed_tokens(input_ids)
150 + current_step_idx = spec_step_idx % self.num_mtp_layers
151 + return self.layers[str(current_step_idx)](
152 + input_ids,
153 + positions,
154 + previous_hidden_states,
155 + inputs_embeds,
156 + current_step_idx,
157 + )
158 +
159 +
160 +class MiniMaxM3MTP(nn.Module):
161 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
162 + super().__init__()
163 +
164 + assert vllm_config.speculative_config is not None
165 + self.config = vllm_config.speculative_config.draft_model_config.hf_config
166 + self.quant_config = vllm_config.quant_config
167 + self.model = MiniMaxM3MultiTokenPredictor(
168 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
169 + )
170 + self.lm_head = ParallelLMHead(
171 + self.config.vocab_size,
172 + self.config.hidden_size,
173 + quant_config=self.quant_config,
174 + prefix=maybe_prefix(prefix, "lm_head"),
175 + )
176 + self.logits_processor = LogitsProcessor(self.config.vocab_size)
177 +
178 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
179 + return self.model.embed_input_ids(input_ids)
180 +
181 + def forward(
182 + self,
183 + input_ids: torch.Tensor | None,
184 + positions: torch.Tensor,
185 + hidden_states: torch.Tensor,
186 + intermediate_tensors: IntermediateTensors | None = None,
187 + inputs_embeds: torch.Tensor | None = None,
188 + spec_step_idx: int = 0,
189 + ) -> torch.Tensor:
190 + return self.model(
191 + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
192 + )
193 +
194 + def compute_logits(
195 + self,
196 + hidden_states: torch.Tensor,
197 + spec_step_idx: int = 0,
198 + ) -> torch.Tensor | None:
199 + current_step_idx = spec_step_idx % self.model.num_mtp_layers
200 + mtp_layer = self.model.layers[str(current_step_idx)]
201 + return self.logits_processor(
202 + self.lm_head, mtp_layer.final_layernorm(hidden_states)
203 + )
204 +
205 + def _get_mtp_layer_idx_from_weight_name(self, name: str) -> int | None:
206 + """Return the MTP layer index in *.mtp.layers.{idx}.*, else None."""
207 + match = re.search(r"\.mtp\.layers\.(\d+)\.", name)
208 + return int(match.group(1)) if match else None
209 +
210 + def _map_checkpoint_name(self, name: str) -> str | None:
211 + """Map a full checkpoint key to this MTP module's parameter name.
212 +
213 + The MTP module only owns the *.mtp.layers.* weights plus the token
214 + embedding and LM head, which the checkpoint shares with the main model.
215 + Everything else belongs to other modules and is ignored here by returning
216 + None.
217 + """
218 + # In the bundled checkpoint, the MTP weights are prefixed with
219 + # "language_model". The standalone MTP checkpoint has no such prefix.
220 + # Strip it if present.
221 + name = name.removeprefix("language_model.")
222 +
223 + if name == "model.embed_tokens.weight":
224 + return "model.embed_tokens.weight"
225 + if name == "lm_head.weight":
226 + return "lm_head.weight"
227 + if "model.mtp.layers" in name:
228 + if "weight_scale_inv" in name:
229 + # The checkpoint stores block scales as "weight_scale_inv".
230 + # The ModelOpt MXFP8 layers expose them as "weight_scale".
231 + name = name.replace("weight_scale_inv", "weight_scale")
232 + # Strip "mtp" from prefix.
233 + return name.replace(".mtp.", ".")
234 + return None
235 +
236 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
237 + # Map q/k/v projections to qkv_proj, and gate/up projections to gate_up_proj.
238 + stacked_params_mapping: list[tuple[str, str, int | str]] = [
239 + # (param_name, shard_name, shard_id)
240 + (".qkv_proj", ".q_proj", "q"),
241 + (".qkv_proj", ".k_proj", "k"),
242 + (".qkv_proj", ".v_proj", "v"),
243 + (".qkv_proj", ".index_q_proj", "index_q"),
244 + (".qkv_proj", ".index_k_proj", "index_k"),
245 + (".gate_up_proj", ".gate_proj", 0),
246 + (".gate_up_proj", ".up_proj", 1),
247 + ]
248 +
249 + # Map expert weights w1/w2/w3 to gate/down/up.
250 + # (param_name, weight_name, expert_id, shard_id)
251 + expert_params_mapping = fused_moe_make_expert_params_mapping(
252 + self,
253 + ckpt_gate_proj_name="w1",
254 + ckpt_down_proj_name="w2",
255 + ckpt_up_proj_name="w3",
256 + num_experts=self.config.num_local_experts,
257 + )
258 +
259 + params_dict = dict(self.named_parameters())
260 + loaded_params: set[str] = set()
261 + loaded_mtp_layers: set[int] = set()
262 + for name, loaded_weight in weights:
263 + mtp_layer = self._get_mtp_layer_idx_from_weight_name(name)
264 + mapped_name = self._map_checkpoint_name(name)
265 + if mapped_name is None:
266 + # This weight does not belong to the MTP module, so skip it.
267 + continue
268 + name = mapped_name
269 +
270 + for param_name, weight_name, shard_id in stacked_params_mapping:
271 + if weight_name not in name:
272 + continue
273 +
274 + # Routed experts (w1/w2/w3) are handled below. Don't let the
275 + # stacked mapping rewrite them.
276 + if ("block_sparse_moe.experts." in name) and name not in params_dict:
277 + continue
278 + name = name.replace(weight_name, param_name)
279 + if name not in params_dict:
280 + continue
281 + param = params_dict[name]
282 + weight_loader = param.weight_loader
283 + weight_loader(param, loaded_weight, shard_id)
284 + break
285 + else:
286 + for (
287 + param_name,
288 + weight_name,
289 + expert_id,
290 + expert_shard_id,
291 + ) in expert_params_mapping:
292 + if weight_name not in name:
293 + continue
294 +
295 + name = name.replace(weight_name, param_name)
296 + if name not in params_dict:
297 + continue
298 + param = params_dict[name]
299 + weight_loader = param.weight_loader
300 + weight_loader(
301 + param,
302 + loaded_weight,
303 + name,
304 + shard_id=expert_shard_id,
305 + expert_id=expert_id,
306 + )
307 + break
308 + else:
309 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)
310 + if remapped_name is None or remapped_name not in params_dict:
311 + continue
312 + name = remapped_name
313 + param = params_dict[name]
314 + weight_loader = getattr(
315 + param, "weight_loader", default_weight_loader
316 + )
317 + weight_loader(param, loaded_weight)
318 +
319 + loaded_params.add(name)
320 + if mtp_layer is not None:
321 + loaded_mtp_layers.add(mtp_layer)
322 +
323 + # Validate that weights were loaded for each MTP layer.
324 + for layer_idx in range(self.model.num_mtp_layers):
325 + if layer_idx not in loaded_mtp_layers:
326 + raise ValueError(
327 + f"Failed to load MTP layer {layer_idx} weights from checkpoint."
328 + )
329 +
330 + return loaded_params
vllm/models/minimax_m3/amd/ops/__init__.py new
+24
@@ -0,0 +1,24 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""AMD/ROCm fused Triton ops for MiniMax-M3.
4 +
5 +These replace per-element PyTorch fallbacks (FlashInfer / fused HIP kernels are
6 +unavailable on ROCm) with single-pass Triton kernels to cut launch overhead and
7 +intermediate-tensor traffic during decode.
8 +"""
9 +
10 +from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import (
11 + gemma_fused_add_rmsnorm,
12 + gemma_rmsnorm,
13 +)
14 +from vllm.models.minimax_m3.amd.ops.swiglu_oai import (
15 + swiglu_oai_quantize_mxfp8,
16 + swiglu_oai_split,
17 +)
18 +
19 +__all__ = [
20 + "gemma_rmsnorm",
21 + "gemma_fused_add_rmsnorm",
22 + "swiglu_oai_split",
23 + "swiglu_oai_quantize_mxfp8",
24 +]
vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py new
+155
@@ -0,0 +1,155 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Fused Gemma-style RMSNorm for AMD ROCm via Triton.
4 +
5 +Gemma RMSNorm = normalize(x) * (1 + weight), computed in fp32. FlashInfer's
6 +``gemma_rmsnorm`` / ``gemma_fused_add_rmsnorm`` CUDA kernels are unavailable on
7 +ROCm, so the AMD path previously used a ~8-op PyTorch sequence (float cast, add,
8 +pow, mean, rsqrt, two muls, cast) — each a separate kernel launch materializing
9 +fp32 intermediates. These kernels collapse that into a single pass per row.
10 +
11 +Two entry points:
12 + * ``gemma_rmsnorm(x, w, eps)`` -> normalized tensor
13 + * ``gemma_fused_add_rmsnorm(x, res, w, eps)`` -> (normalized, x + res)
14 +
15 +Both normalize over the last dim and broadcast ``weight`` (shape [N]) over it,
16 +so they serve both the full-hidden norms (input/post-attn/final) and the
17 +per-head q_norm/k_norm (N == head_dim). Inputs may be non-contiguous views
18 +(e.g. ``qkv.split`` slices); strides are passed through and outputs are written
19 +contiguous.
20 +"""
21 +
22 +import torch
23 +
24 +from vllm.triton_utils import tl, triton
25 +
26 +
27 +@triton.jit
28 +def _gemma_rmsnorm_kernel(
29 + x_ptr,
30 + w_ptr,
31 + out_ptr,
32 + n_cols,
33 + stride_row,
34 + stride_col,
35 + eps,
36 + BLOCK_N: tl.constexpr,
37 +):
38 + row = tl.program_id(0)
39 + cols = tl.arange(0, BLOCK_N)
40 + mask = cols < n_cols
41 + x = tl.load(x_ptr + row * stride_row + cols * stride_col, mask=mask, other=0.0).to(
42 + tl.float32
43 + )
44 + var = tl.sum(x * x, axis=0) / n_cols
45 + rstd = 1.0 / tl.sqrt(var + eps)
46 + w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
47 + out = x * rstd * (1.0 + w)
48 + tl.store(
49 + out_ptr + row * n_cols + cols,
50 + out.to(out_ptr.dtype.element_ty),
51 + mask=mask,
52 + )
53 +
54 +
55 +@triton.jit
56 +def _gemma_fused_add_rmsnorm_kernel(
57 + x_ptr,
58 + res_ptr,
59 + w_ptr,
60 + out_ptr,
61 + res_out_ptr,
62 + n_cols,
63 + stride_xrow,
64 + stride_xcol,
65 + stride_rrow,
66 + stride_rcol,
67 + eps,
68 + BLOCK_N: tl.constexpr,
69 +):
70 + row = tl.program_id(0)
71 + cols = tl.arange(0, BLOCK_N)
72 + mask = cols < n_cols
73 + x = tl.load(
74 + x_ptr + row * stride_xrow + cols * stride_xcol, mask=mask, other=0.0
75 + ).to(tl.float32)
76 + r = tl.load(
77 + res_ptr + row * stride_rrow + cols * stride_rcol, mask=mask, other=0.0
78 + ).to(tl.float32)
79 + s = x + r
80 + # residual_out is the pre-norm sum (consumed by the next layer's add).
81 + tl.store(
82 + res_out_ptr + row * n_cols + cols,
83 + s.to(res_out_ptr.dtype.element_ty),
84 + mask=mask,
85 + )
86 + var = tl.sum(s * s, axis=0) / n_cols
87 + rstd = 1.0 / tl.sqrt(var + eps)
88 + w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
89 + out = s * rstd * (1.0 + w)
90 + tl.store(
91 + out_ptr + row * n_cols + cols,
92 + out.to(out_ptr.dtype.element_ty),
93 + mask=mask,
94 + )
95 +
96 +
97 +def _num_warps(block_n: int) -> int:
98 + if block_n >= 4096:
99 + return 16
100 + if block_n >= 1024:
101 + return 8
102 + return 4
103 +
104 +
105 +def gemma_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
106 + orig_shape = x.shape
107 + n = orig_shape[-1]
108 + x2 = x.reshape(-1, n)
109 + m = x2.shape[0]
110 + out = torch.empty((m, n), dtype=x.dtype, device=x.device)
111 + block_n = triton.next_power_of_2(n)
112 + _gemma_rmsnorm_kernel[(m,)](
113 + x2,
114 + weight,
115 + out,
116 + n,
117 + x2.stride(0),
118 + x2.stride(1),
119 + eps,
120 + BLOCK_N=block_n,
121 + num_warps=_num_warps(block_n),
122 + )
123 + return out.reshape(orig_shape)
124 +
125 +
126 +def gemma_fused_add_rmsnorm(
127 + x: torch.Tensor,
128 + residual: torch.Tensor,
129 + weight: torch.Tensor,
130 + eps: float,
131 +) -> tuple[torch.Tensor, torch.Tensor]:
132 + orig_shape = x.shape
133 + n = orig_shape[-1]
134 + x2 = x.reshape(-1, n)
135 + r2 = residual.reshape(-1, n)
136 + m = x2.shape[0]
137 + out = torch.empty((m, n), dtype=x.dtype, device=x.device)
138 + res_out = torch.empty((m, n), dtype=x.dtype, device=x.device)
139 + block_n = triton.next_power_of_2(n)
140 + _gemma_fused_add_rmsnorm_kernel[(m,)](
141 + x2,
142 + r2,
143 + weight,
144 + out,
145 + res_out,
146 + n,
147 + x2.stride(0),
148 + x2.stride(1),
149 + r2.stride(0),
150 + r2.stride(1),
151 + eps,
152 + BLOCK_N=block_n,
153 + num_warps=_num_warps(block_n),
154 + )
155 + return out.reshape(orig_shape), res_out.reshape(orig_shape)
vllm/models/minimax_m3/amd/ops/swiglu_oai.py new
+221
@@ -0,0 +1,221 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Fused SwiGLU-OAI activation (split layout) for AMD ROCm via Triton.
4 +
5 +SwiGLU-OAI on a ``[*, 2I]`` split-layout input (gate = first half, up = second
6 +half):
7 +
8 + gate = clamp(gate, max=limit)
9 + up = clamp(up, -limit, +limit)
10 + out = gate * sigmoid(alpha * gate) * (up + beta)
11 +
12 +On ROCm the dense MLP and the native MXFP8 MoE (between its two GEMMs) fell back
13 +to a chain of elementwise PyTorch ops with fp32 intermediates: vLLM's shared
14 +``SiluAndMulWithClamp`` blanket-routes ROCm to ``forward_native``, and the MoE
15 +applies the activation inline in PyTorch. This Triton kernel collapses that into
16 +a single pass producing the ``[*, I]`` output directly, and computes in fp32
17 +(rel ~1e-6 vs reference).
18 +
19 +Note: the vectorized ``torch.ops._C.silu_and_mul_with_clamp`` op IS built on
20 +ROCm and is ~1.2-2.2x faster in isolation, but the win is launch overhead that
21 +HIP graphs already eliminate — measured end-to-end throughput is identical
22 +(within noise), so we keep the fp32-accurate Triton kernel.
23 +"""
24 +
25 +import torch
26 +
27 +from vllm.triton_utils import tl, triton
28 +
29 +
30 +@triton.jit
31 +def _swiglu_oai_kernel(
32 + g_ptr,
33 + out_ptr,
34 + n_inter,
35 + stride_gm,
36 + stride_gn,
37 + stride_om,
38 + stride_on,
39 + alpha,
40 + beta,
41 + limit,
42 + HAS_LIMIT: tl.constexpr,
43 + BLOCK_I: tl.constexpr,
44 +):
45 + row = tl.program_id(0)
46 + pid_i = tl.program_id(1)
47 + cols = pid_i * BLOCK_I + tl.arange(0, BLOCK_I)
48 + mask = cols < n_inter
49 + gate = tl.load(g_ptr + row * stride_gm + cols * stride_gn, mask=mask, other=0.0).to(
50 + tl.float32
51 + )
52 + up = tl.load(
53 + g_ptr + row * stride_gm + (n_inter + cols) * stride_gn,
54 + mask=mask,
55 + other=0.0,
56 + ).to(tl.float32)
57 + if HAS_LIMIT:
58 + gate = tl.minimum(gate, limit)
59 + up = tl.minimum(tl.maximum(up, -limit), limit)
60 + out = gate * tl.sigmoid(alpha * gate) * (up + beta)
61 + tl.store(
62 + out_ptr + row * stride_om + cols * stride_on,
63 + out.to(out_ptr.dtype.element_ty),
64 + mask=mask,
65 + )
66 +
67 +
68 +@triton.jit
69 +def _swiglu_oai_quant_kernel(
70 + g_ptr,
71 + aq_ptr,
72 + as_ptr,
73 + M,
74 + n_inter,
75 + stride_gm,
76 + stride_gn,
77 + stride_qm,
78 + stride_qn,
79 + stride_sm,
80 + stride_sk,
81 + alpha,
82 + beta,
83 + limit,
84 + HAS_LIMIT: tl.constexpr,
85 + BLOCK_M: tl.constexpr,
86 +):
87 + """SwiGLU-OAI (split layout) fused with per-32-block MXFP8 (E4M3 + E8M0)
88 + quant. Each program handles ``[BLOCK_M, 32]`` of the ``[M, I]`` output (one
89 + MX block): it reads the matching gate/up columns from ``g1`` (``[M, 2I]``),
90 + computes the SwiGLU in fp32, then derives the block E8M0 scale and emits the
91 + FP8 values + scale in a single pass — no bf16 ``act`` round-trip to HBM.
92 + """
93 + pid_m = tl.program_id(0)
94 + pid_b = tl.program_id(1) # which 32-element block along I
95 + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
96 + offs_c = pid_b * 32 + tl.arange(0, 32)
97 + m_mask = offs_m < M
98 + gate = tl.load(
99 + g_ptr + offs_m[:, None] * stride_gm + offs_c[None, :] * stride_gn,
100 + mask=m_mask[:, None],
101 + other=0.0,
102 + ).to(tl.float32)
103 + up = tl.load(
104 + g_ptr + offs_m[:, None] * stride_gm + (n_inter + offs_c)[None, :] * stride_gn,
105 + mask=m_mask[:, None],
106 + other=0.0,
107 + ).to(tl.float32)
108 + if HAS_LIMIT:
109 + gate = tl.minimum(gate, limit)
110 + up = tl.minimum(tl.maximum(up, -limit), limit)
111 + act = gate * tl.sigmoid(alpha * gate) * (up + beta) # [BLOCK_M, 32] fp32
112 + amax = tl.maximum(tl.max(tl.abs(act), axis=1), 1e-30) # [BLOCK_M]
113 + sb = tl.minimum(tl.maximum(tl.floor(tl.log2(amax)) + 127.0, 0.0), 254.0)
114 + descale = tl.exp2(sb - 127.0)
115 + aq = (act / descale[:, None]).to(aq_ptr.dtype.element_ty)
116 + tl.store(
117 + aq_ptr + offs_m[:, None] * stride_qm + offs_c[None, :] * stride_qn,
118 + aq,
119 + mask=m_mask[:, None],
120 + )
121 + tl.store(
122 + as_ptr + offs_m * stride_sm + pid_b * stride_sk, sb.to(tl.uint8), mask=m_mask
123 + )
124 +
125 +
126 +def swiglu_oai_quantize_mxfp8(
127 + gate_up: torch.Tensor,
128 + alpha: float,
129 + beta: float,
130 + limit: float | None,
131 + block_m: int = 64,
132 +) -> tuple[torch.Tensor, torch.Tensor]:
133 + """SwiGLU-OAI on split-layout ``[M, 2I]`` fused with MXFP8 activation-quant.
134 +
135 + Returns ``(act_q [M, I] float8_e4m3fn, act_scale [M, I//32] uint8 E8M0)``,
136 + identical to ``mxfp8_e4m3_quantize(swiglu_oai_split(gate_up))`` but in a
137 + single Triton pass (no bf16 intermediate). Used between the two GEMMs of the
138 + native MXFP8 MoE. Numerically equivalent to the unfused chain (bit-exact on
139 + measured MoE shapes); marginally more accurate (fp32 act, no bf16 round-trip).
140 + """
141 + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
142 + MXFP8_BLOCK_SIZE,
143 + MXFP8_SCALE_DTYPE,
144 + MXFP8_VALUE_DTYPE,
145 + )
146 +
147 + two_i = gate_up.shape[-1]
148 + n_inter = two_i // 2
149 + assert n_inter % MXFP8_BLOCK_SIZE == 0, (
150 + f"fused swiglu+quant needs I % {MXFP8_BLOCK_SIZE} == 0, got I={n_inter}"
151 + )
152 + g1 = gate_up.reshape(-1, two_i).contiguous()
153 + M = g1.shape[0]
154 + aq = torch.empty((M, n_inter), dtype=MXFP8_VALUE_DTYPE, device=g1.device)
155 + asc = torch.empty(
156 + (M, n_inter // MXFP8_BLOCK_SIZE), dtype=MXFP8_SCALE_DTYPE, device=g1.device
157 + )
158 + grid = (triton.cdiv(M, block_m), n_inter // MXFP8_BLOCK_SIZE)
159 + _swiglu_oai_quant_kernel[grid](
160 + g1,
161 + aq,
162 + asc,
163 + M,
164 + n_inter,
165 + g1.stride(0),
166 + g1.stride(1),
167 + aq.stride(0),
168 + aq.stride(1),
169 + asc.stride(0),
170 + asc.stride(1),
171 + float(alpha),
172 + float(beta),
173 + 0.0 if limit is None else float(limit),
174 + HAS_LIMIT=limit is not None,
175 + BLOCK_M=block_m,
176 + num_warps=4,
177 + )
178 + return aq, asc
179 +
180 +
181 +def swiglu_oai_split(
182 + gate_up: torch.Tensor,
183 + alpha: float,
184 + beta: float,
185 + limit: float | None,
186 + out_dtype: torch.dtype | None = None,
187 +) -> torch.Tensor:
188 + """SwiGLU-OAI on a split-layout ``[*, 2I]`` tensor -> ``[*, I]``."""
189 + orig_shape = gate_up.shape
190 + two_i = orig_shape[-1]
191 + n_inter = two_i // 2
192 + x2 = gate_up.reshape(-1, two_i)
193 + m = x2.shape[0]
194 + dt = out_dtype if out_dtype is not None else gate_up.dtype
195 + out = torch.empty((m, n_inter), dtype=dt, device=gate_up.device)
196 + # Tile tuned on gfx950. The SwiGLU intermediate is sharded across tensor
197 + # parallel ranks (per-rank n_inter = I / tp: dense I=12288, MoE I=3072), and
198 + # a 512-wide tile (4 warps, ~2 elems/lane) only helps once the per-rank slice
199 + # is large enough to be bandwidth-bound — at TP=1 prefill that is ~1.25-1.35x
200 + # faster than 256. For small sharded slices (high TP) the kernel is launch-
201 + # bound (~12us) and a wide tile can slightly regress, so fall back to 256.
202 + # Decode is launch-bound at every TP. num_warps=8 underfills this tile, so it
203 + # is pinned to 4.
204 + block_i = 512 if n_inter >= 2048 else 256
205 + grid = (m, triton.cdiv(n_inter, block_i))
206 + _swiglu_oai_kernel[grid](
207 + x2,
208 + out,
209 + n_inter,
210 + x2.stride(0),
211 + x2.stride(1),
212 + out.stride(0),
213 + out.stride(1),
214 + float(alpha),
215 + float(beta),
216 + 0.0 if limit is None else float(limit),
217 + HAS_LIMIT=limit is not None,
218 + BLOCK_I=block_i,
219 + num_warps=4,
220 + )
221 + return out.reshape(*orig_shape[:-1], n_inter)
vllm/models/minimax_m3/common/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/minimax_m3/common/indexer.py new
+512
@@ -0,0 +1,512 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MiniMax M3 lightning indexer: side cache, metadata, and impl.
4 +
5 +The indexer scores KV blocks with the index heads and selects the top-k blocks
6 +(plus fixed init/local blocks) that the main block-sparse attention
7 +(``sparse_attention.py``) then attends to. It owns its own side cache
8 +(``MiniMaxM3IndexerCache``, one index-key vector per token), metadata, and
9 +metadata builder, mirroring how DeepSeek V4 keeps the indexer separate from the
10 +main attention.
11 +
12 +``MiniMaxM3Indexer`` is the ``nn.Module`` the attention layer holds (like
13 +``DeepseekV4Indexer``); it picks a kernel impl in ``__init__`` (via
14 +``select_indexer_impl_cls``) and delegates ``forward`` to it.
15 +"""
16 +
17 +from dataclasses import dataclass
18 +from typing import ClassVar
19 +
20 +import torch
21 +from torch import nn
22 +
23 +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
24 +from vllm.config.attention import IndexerKVDType
25 +from vllm.config.cache import CacheDType
26 +from vllm.distributed import get_tensor_model_parallel_world_size
27 +from vllm.forward_context import get_forward_context
28 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
29 +from vllm.models.minimax_m3.common.ops.index_topk import (
30 + minimax_m3_index_decode,
31 + minimax_m3_index_score,
32 + minimax_m3_index_topk,
33 +)
34 +from vllm.v1.attention.backend import (
35 + AttentionBackend,
36 + AttentionCGSupport,
37 + AttentionMetadata,
38 + AttentionMetadataBuilder,
39 + CommonAttentionMetadata,
40 + MultipleOf,
41 +)
42 +from vllm.v1.attention.backends.utils import split_decodes_and_prefills
43 +from vllm.v1.kv_cache_interface import (
44 + AttentionSpec,
45 + KVCacheSpec,
46 + MLAAttentionSpec,
47 +)
48 +
49 +
50 +class MiniMaxM3IndexerBackend(AttentionBackend):
51 + """Indexer side-cache backend (key-only)."""
52 +
53 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16, torch.float16]
54 + # bf16 today; mirrors the main backend to keep spec validation permissive.
55 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
56 + "bfloat16",
57 + "fp8",
58 + "fp8_e4m3",
59 + "fp8_e5m2",
60 + ]
61 +
62 + @staticmethod
63 + def get_name() -> str:
64 + return "MINIMAX_M3_SPARSE_INDEXER"
65 +
66 + @staticmethod
67 + def get_impl_cls() -> type["MiniMaxM3IndexerImpl"]:
68 + # Concrete impl chosen by select_indexer_impl_cls; base for introspection.
69 + return MiniMaxM3IndexerImpl
70 +
71 + @staticmethod
72 + def get_builder_cls() -> type["MiniMaxM3IndexerMetadataBuilder"]:
73 + return MiniMaxM3IndexerTritonMetadataBuilder
74 +
75 + @classmethod
76 + def get_supported_head_sizes(cls) -> list[int]:
77 + return [128]
78 +
79 + @staticmethod
80 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
81 + return [128]
82 +
83 + @classmethod
84 + def is_sparse(cls) -> bool:
85 + return True
86 +
87 + @staticmethod
88 + def get_kv_cache_shape(
89 + num_blocks: int,
90 + block_size: int,
91 + num_kv_heads: int,
92 + head_size: int,
93 + cache_dtype_str: str = "auto",
94 + ) -> tuple[int, ...]:
95 + return (num_blocks, block_size, head_size)
96 +
97 + @staticmethod
98 + def get_kv_cache_stride_order(
99 + include_num_layers_dimension: bool = False,
100 + ) -> tuple[int, ...]:
101 + if include_num_layers_dimension:
102 + # M3 does not use cross-layer (per-layer-stacked) KV blocks.
103 + raise NotImplementedError
104 + return (0, 1, 2)
105 +
106 +
107 +class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase):
108 + """Side KV cache for the indexer's per-token index keys (key-only).
109 +
110 + Registers itself in the static forward context so the KV-cache manager
111 + allocates it (like ``DeepseekV32IndexerCache``).
112 + """
113 +
114 + def __init__(
115 + self,
116 + head_dim: int,
117 + prefix: str,
118 + cache_config: CacheConfig | None = None,
119 + indexer_kv_dtype: IndexerKVDType = "bf16",
120 + backend_cls: type[AttentionBackend] = MiniMaxM3IndexerBackend,
121 + ) -> None:
122 + super().__init__()
123 + if indexer_kv_dtype != "bf16":
124 + raise NotImplementedError(
125 + f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported yet "
126 + "for the MiniMax M3 indexer cache (only 'bf16')."
127 + )
128 + self.kv_cache = torch.tensor([])
129 + self.head_dim = head_dim
130 + self.indexer_kv_dtype = indexer_kv_dtype
131 + # Storage dtype for the side cache (bf16 today; quantized layouts later).
132 + self.dtype = torch.bfloat16
133 + self.prefix = prefix
134 + self.cache_config = cache_config
135 + # Impl-chosen backend -> each impl gets its own builder (get_attn_backend).
136 + self.backend_cls = backend_cls
137 + compilation_config = get_current_vllm_config().compilation_config
138 + if prefix in compilation_config.static_forward_context:
139 + raise ValueError(f"Duplicate layer name: {prefix}")
140 + compilation_config.static_forward_context[prefix] = self
141 +
142 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
143 + # Key-only: MLAAttentionSpec budgets one vector/token (not 2x for K+V).
144 + return MLAAttentionSpec(
145 + block_size=vllm_config.cache_config.block_size,
146 + num_kv_heads=1,
147 + head_size=self.head_dim,
148 + dtype=self.dtype,
149 + )
150 +
151 + def forward(self) -> None: ...
152 +
153 + def get_attn_backend(self) -> type[AttentionBackend]:
154 + return self.backend_cls
155 +
156 +
157 +@dataclass
158 +class MiniMaxM3IndexerPrefillMetadata:
159 + """Per-prefill index-scoring state."""
160 +
161 + cu_seqlens_q: torch.Tensor # [num_prefills + 1] int32, rebased to 0
162 + seq_lens: torch.Tensor # [num_prefills] int32, total KV lengths
163 + context_lens: torch.Tensor # [num_prefills] int32 (cached/context tokens)
164 + block_table: torch.Tensor
165 + max_query_len: int
166 + max_seq_len: int
167 +
168 +
169 +@dataclass
170 +class MiniMaxM3IndexerDecodeMetadata:
171 + """Per-decode state (cudagraph-safe). ``decode_query_len`` is the uniform
172 + per-request query length (1, or 1 + num_speculative_tokens)."""
173 +
174 + seq_lens: torch.Tensor # [num_decodes] int32
175 + block_table: torch.Tensor
176 + max_seq_len: int
177 + decode_query_len: int
178 +
179 +
180 +@dataclass
181 +class MiniMaxM3IndexerMetadata(AttentionMetadata):
182 + """Indexer metadata, split into prefill and decode sub-metadata."""
183 +
184 + seq_lens: torch.Tensor
185 + max_seq_len: int
186 + slot_mapping: torch.Tensor
187 +
188 + num_actual_tokens: int # total query tokens (decode-first batch)
189 +
190 + # Split counts; identical to the main metadata's (same reorder threshold).
191 + num_decodes: int
192 + num_decode_tokens: int
193 + num_prefills: int
194 + num_prefill_tokens: int
195 +
196 + prefill: MiniMaxM3IndexerPrefillMetadata | None = None
197 + decode: MiniMaxM3IndexerDecodeMetadata | None = None
198 +
199 +
200 +class MiniMaxM3IndexerMetadataBuilder(
201 + AttentionMetadataBuilder[MiniMaxM3IndexerMetadata]
202 +):
203 + """Abstract base: shared setup only. The Triton and MSA builders are
204 + parallel subclasses that each own their full ``build`` (no shared code)."""
205 +
206 + # Full cudagraphs for uniform decode batches (incl. spec-decode verify).
207 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
208 + # Raised to 1 + num_speculative_tokens by _init_reorder_batch_threshold when
209 + # spec decode is on; matches the main builder so the splits agree.
210 + reorder_batch_threshold: int = 1
211 +
212 + def __init__(
213 + self,
214 + kv_cache_spec: AttentionSpec,
215 + layer_names: list[str],
216 + vllm_config: VllmConfig,
217 + device: torch.device,
218 + ) -> None:
219 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
220 + hf_config = vllm_config.model_config.hf_config
221 + text_config = getattr(hf_config, "text_config", hf_config)
222 + sparse_cfg = text_config.sparse_attention_config
223 + # Index-query head count from model config (cache spec has 1 vec/token).
224 + total_index_heads = sparse_cfg["sparse_num_index_heads"]
225 + tp_size = get_tensor_model_parallel_world_size()
226 + if total_index_heads >= tp_size:
227 + assert total_index_heads % tp_size == 0
228 + else:
229 + assert tp_size % total_index_heads == 0
230 + self.num_index_heads = max(1, total_index_heads // tp_size)
231 + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
232 +
233 + # Stable context-length buffer for decode cudagraph replays.
234 + self.context_len_buffer = torch.empty(
235 + vllm_config.scheduler_config.max_num_batched_tokens,
236 + dtype=torch.int32,
237 + device=device,
238 + )
239 +
240 +
241 +class MiniMaxM3IndexerTritonMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
242 + """Triton indexer metadata: no SM100 fmha_sm100 plan."""
243 +
244 + def build(
245 + self,
246 + common_prefix_len: int,
247 + common_attn_metadata: CommonAttentionMetadata,
248 + fast_build: bool = False,
249 + ) -> MiniMaxM3IndexerMetadata:
250 + num_reqs = common_attn_metadata.num_reqs
251 + num_tokens = common_attn_metadata.num_actual_tokens
252 + query_start_loc = common_attn_metadata.query_start_loc
253 + seq_lens = common_attn_metadata.seq_lens
254 + block_table = common_attn_metadata.block_table_tensor
255 +
256 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
257 + split_decodes_and_prefills(
258 + common_attn_metadata,
259 + decode_threshold=self.reorder_batch_threshold,
260 + require_uniform=True,
261 + )
262 + )
263 + assert num_decodes + num_prefills == num_reqs
264 + assert num_decode_tokens + num_prefill_tokens == num_tokens
265 +
266 + # Decode-first batch: context lengths into the stable cudagraph buffer.
267 + context_lens = self.context_len_buffer[:num_reqs]
268 + context_lens.copy_(
269 + common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
270 + )
271 +
272 + prefill_metadata: MiniMaxM3IndexerPrefillMetadata | None = None
273 + if num_prefills > 0:
274 + prefill_metadata = MiniMaxM3IndexerPrefillMetadata(
275 + cu_seqlens_q=(query_start_loc[num_decodes:] - num_decode_tokens).to(
276 + torch.int32
277 + ),
278 + seq_lens=seq_lens[num_decodes:],
279 + context_lens=context_lens[num_decodes:],
280 + block_table=block_table[num_decodes:],
281 + max_query_len=common_attn_metadata.max_query_len,
282 + max_seq_len=common_attn_metadata.max_seq_len,
283 + )
284 +
285 + decode_metadata: MiniMaxM3IndexerDecodeMetadata | None = None
286 + if num_decodes > 0:
287 + qsl_cpu = common_attn_metadata.query_start_loc_cpu
288 + query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes]
289 + decode_query_len = int(query_lens_cpu[0].item())
290 + assert decode_query_len > 0
291 + assert torch.all(
292 + (query_lens_cpu == decode_query_len) | (query_lens_cpu == 0)
293 + )
294 + assert num_decode_tokens == num_decodes * decode_query_len
295 + decode_metadata = MiniMaxM3IndexerDecodeMetadata(
296 + seq_lens=seq_lens[:num_decodes],
297 + block_table=block_table[:num_decodes],
298 + max_seq_len=common_attn_metadata.max_seq_len,
299 + decode_query_len=decode_query_len,
300 + )
301 +
302 + return MiniMaxM3IndexerMetadata(
303 + seq_lens=seq_lens,
304 + max_seq_len=common_attn_metadata.max_seq_len,
305 + slot_mapping=common_attn_metadata.slot_mapping,
306 + num_actual_tokens=num_tokens,
307 + num_decodes=num_decodes,
308 + num_decode_tokens=num_decode_tokens,
309 + num_prefills=num_prefills,
310 + num_prefill_tokens=num_prefill_tokens,
311 + prefill=prefill_metadata,
312 + decode=decode_metadata,
313 + )
314 +
315 +
316 +class MiniMaxM3IndexerImpl(nn.Module):
317 + """Abstract base for the indexer kernel impls.
318 +
319 + Each impl owns its side cache and reports its backend via
320 + ``indexer_backend_cls`` (so each gets its own builder). The Triton and MSA
321 + subclasses each own a full ``forward`` returning ``(decode_topk,
322 + prefill_topk)`` -- no shared forward code.
323 + """
324 +
325 + # Set by each impl so the side cache reports the matching backend + builder.
326 + indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerBackend
327 +
328 + def __init__(
329 + self,
330 + *,
331 + num_kv_heads: int,
332 + scale: float,
333 + topk_blocks: int,
334 + sparse_block_size: int,
335 + num_index_heads: int,
336 + index_head_dim: int,
337 + prefix: str,
338 + init_blocks: int = 0,
339 + local_blocks: int = 0,
340 + score_type: str = "max",
341 + cache_config: CacheConfig | None = None,
342 + indexer_kv_dtype: IndexerKVDType = "bf16",
343 + ) -> None:
344 + super().__init__()
345 + self.num_kv_heads = num_kv_heads
346 + self.scale = scale
347 + self.topk_blocks = topk_blocks
348 + self.block_size = sparse_block_size
349 + self.init_blocks = init_blocks
350 + self.local_blocks = local_blocks
351 + self.score_type = score_type
352 + self.num_index_heads = num_index_heads
353 + self.index_head_dim = index_head_dim
354 + self.indexer_kv_dtype = indexer_kv_dtype
355 + # Owns the side cache (registers itself in the static forward context).
356 + self.index_cache = MiniMaxM3IndexerCache(
357 + head_dim=index_head_dim,
358 + prefix=f"{prefix}.index_cache",
359 + cache_config=cache_config,
360 + indexer_kv_dtype=indexer_kv_dtype,
361 + backend_cls=type(self).indexer_backend_cls,
362 + )
363 +
364 + def forward(
365 + self,
366 + index_query: torch.Tensor,
367 + ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
368 + """Return ``(decode_topk, prefill_topk)``; implemented per kernel impl."""
369 + raise NotImplementedError
370 +
371 +
372 +class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
373 + """Triton indexer score + top-k for both prefill and decode."""
374 +
375 + def forward(
376 + self,
377 + index_query: torch.Tensor,
378 + ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
379 + attn_metadata = get_forward_context().attn_metadata
380 + if not isinstance(attn_metadata, dict):
381 + return None, None # profiling run; caches unbound
382 + index_md = attn_metadata[self.index_cache.prefix]
383 + assert isinstance(index_md, MiniMaxM3IndexerMetadata)
384 + num_tokens = index_md.num_actual_tokens
385 + nd = index_md.num_decode_tokens
386 + iq = index_query[:num_tokens].view(
387 + -1, self.num_index_heads, self.index_head_dim
388 + )
389 + kv = self.index_cache.kv_cache
390 +
391 + decode_topk: torch.Tensor | None = None
392 + prefill_topk: torch.Tensor | None = None
393 + if index_md.num_decodes > 0:
394 + d = index_md.decode
395 + assert d is not None
396 + decode_topk = minimax_m3_index_decode(
397 + iq[:nd],
398 + kv,
399 + d.block_table,
400 + d.seq_lens,
401 + d.max_seq_len,
402 + self.topk_blocks,
403 + self.init_blocks,
404 + self.local_blocks,
405 + self.num_kv_heads,
406 + self.scale,
407 + d.decode_query_len,
408 + )
409 + if index_md.num_prefills > 0:
410 + p = index_md.prefill
411 + assert p is not None
412 + score = minimax_m3_index_score(
413 + iq[nd:],
414 + kv,
415 + p.block_table,
416 + p.cu_seqlens_q,
417 + p.seq_lens,
418 + p.context_lens,
419 + p.max_query_len,
420 + p.max_seq_len,
421 + self.num_kv_heads,
422 + self.scale,
423 + )
424 + prefill_topk = minimax_m3_index_topk(
425 + score,
426 + p.cu_seqlens_q,
427 + p.context_lens,
428 + p.max_query_len,
429 + self.topk_blocks,
430 + self.init_blocks,
431 + self.local_blocks,
432 + )
433 + return decode_topk, prefill_topk
434 +
435 +
436 +def select_indexer_impl_cls(
437 + *,
438 + indexer_kv_dtype: IndexerKVDType = "bf16",
439 +) -> type[MiniMaxM3IndexerImpl]:
440 + """Pick the indexer impl off the index-cache dtype.
441 +
442 + The SM100 MSA indexer score path is disabled for now; use the local Triton
443 + indexer. If re-enabled, add a NVIDIA-specific ``MiniMaxM3IndexerImpl`` here.
444 + """
445 + if indexer_kv_dtype in ("mxfp4", "nvfp4"):
446 + raise NotImplementedError(
447 + f"indexer_kv_dtype={indexer_kv_dtype!r} needs the (not-yet-added) "
448 + "CuteDSL indexer impl."
449 + )
450 + if indexer_kv_dtype != "bf16":
451 + raise NotImplementedError(
452 + f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the "
453 + "Triton indexer impl."
454 + )
455 + return MiniMaxM3IndexerTritonImpl
456 +
457 +
458 +class MiniMaxM3Indexer(nn.Module):
459 + """Indexer module held by the attention layer (like ``DeepseekV4Indexer``).
460 +
461 + Picks the kernel impl in ``__init__`` (``select_indexer_impl_cls``) and
462 + delegates ``forward``; exposes the impl's side cache via ``index_cache``.
463 + """
464 +
465 + def __init__(
466 + self,
467 + *,
468 + num_kv_heads: int,
469 + scale: float,
470 + topk_blocks: int,
471 + sparse_block_size: int,
472 + num_index_heads: int,
473 + index_head_dim: int,
474 + prefix: str,
475 + init_blocks: int = 0,
476 + local_blocks: int = 0,
477 + score_type: str = "max",
478 + cache_config: CacheConfig | None = None,
479 + indexer_kv_dtype: IndexerKVDType = "bf16",
480 + ) -> None:
481 + super().__init__()
482 + impl_cls = select_indexer_impl_cls(
483 + indexer_kv_dtype=indexer_kv_dtype,
484 + )
485 + self.impl = impl_cls(
486 + num_kv_heads=num_kv_heads,
487 + scale=scale,
488 + topk_blocks=topk_blocks,
489 + sparse_block_size=sparse_block_size,
490 + num_index_heads=num_index_heads,
491 + index_head_dim=index_head_dim,
492 + prefix=prefix,
493 + init_blocks=init_blocks,
494 + local_blocks=local_blocks,
495 + score_type=score_type,
496 + cache_config=cache_config,
497 + indexer_kv_dtype=indexer_kv_dtype,
498 + )
499 +
500 + @property
501 + def index_cache(self) -> MiniMaxM3IndexerCache:
502 + return self.impl.index_cache
503 +
504 + @property
505 + def num_index_heads(self) -> int:
506 + return self.impl.num_index_heads
507 +
508 + def forward(
509 + self,
510 + index_query: torch.Tensor,
511 + ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
512 + return self.impl(index_query)
vllm/models/minimax_m3/common/mm_preprocess.py new
+514
@@ -0,0 +1,514 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import math
5 +from collections.abc import Mapping, Sequence
6 +from typing import cast
7 +
8 +import torch
9 +from transformers import BatchFeature
10 +from transformers.video_utils import VideoMetadata
11 +
12 +from vllm.config.multimodal import (
13 + BaseDummyOptions,
14 + ImageDummyOptions,
15 + VideoDummyOptions,
16 +)
17 +from vllm.inputs import MultiModalDataDict
18 +from vllm.multimodal.inputs import (
19 + MultiModalFieldConfig,
20 + MultiModalKwargsItems,
21 +)
22 +from vllm.multimodal.parse import (
23 + ImageSize,
24 + MultiModalDataItems,
25 + MultiModalDataParser,
26 +)
27 +from vllm.multimodal.processing import (
28 + BaseDummyInputsBuilder,
29 + BaseMultiModalProcessor,
30 + BaseProcessingInfo,
31 + PromptReplacement,
32 + PromptUpdate,
33 + PromptUpdateDetails,
34 +)
35 +from vllm.multimodal.video import (
36 + VIDEO_LOADER_REGISTRY,
37 + VideoBackend,
38 + VideoSourceMetadata,
39 + VideoTargetMetadata,
40 +)
41 +from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3Config
42 +from vllm.transformers_utils.processors.minimax_m3 import (
43 + MIN_SHORT_SIDE_PIXEL,
44 + MiniMaxM3VLImageProcessor,
45 + MiniMaxM3VLVideoProcessor,
46 + MiniMaxVLProcessor,
47 + smart_resize,
48 +)
49 +
50 +# Upper bound on the number of frames used to build the dummy video during
51 +# memory profiling. Sized to the worst-case video the processor accepts:
52 +# ``max_total_pixels // max_pixels_per_frame`` = 301,056,000 // 602,112 = 500
53 +# frames, each at the video processor's per-frame ``max_pixels`` (768 * 28 * 28
54 +# = 602,112). This reaches the true worst-case ~192,000 vision tokens, but only
55 +# because the dummy video is sized via ``get_video_size_with_most_features()``
56 +# (the video ``max_pixels`` bound), not the smaller image bound. Without a cap,
57 +# ``_get_max_video_frames(seq_len)`` with M3's large ``max_model_len`` yields
58 +# ~1400 frames, producing a multi-GB dummy tensor that overflows the
59 +# multimodal encoder cache.
60 +_MAX_FRAMES_PER_VIDEO = 500
61 +
62 +
63 +class MiniMaxM3VLProcessingInfo(BaseProcessingInfo):
64 + IMAGE_TOKEN = "]<]image[>["
65 + VIDEO_TOKEN = "]<]video[>["
66 + VISION_START_TOKEN = "]<]start of image[>["
67 + VISION_END_TOKEN = "]<]end of image[>["
68 +
69 + def get_hf_config(self) -> MiniMaxM3Config:
70 + return self.ctx.get_hf_config(MiniMaxM3Config)
71 +
72 + def get_hf_processor(self, **kwargs: object) -> MiniMaxVLProcessor:
73 + # The released checkpoint only ships the processor as remote code
74 + # (via ``auto_map``). Construct the vendored processor directly so the
75 + # model loads without ``--trust-remote-code``.
76 + return self.ctx.get_hf_processor(MiniMaxVLProcessor, **kwargs)
77 +
78 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
79 + return {"image": None, "video": None}
80 +
81 + def get_mm_max_tokens_per_item(
82 + self,
83 + seq_len: int,
84 + mm_counts: Mapping[str, int],
85 + ) -> Mapping[str, int]:
86 + return {
87 + "image": self.get_max_image_tokens(),
88 + "video": self.get_max_video_tokens(seq_len, mm_counts),
89 + }
90 +
91 + def get_image_processor(self, **kwargs: object) -> MiniMaxM3VLImageProcessor:
92 + return self.get_hf_processor(**kwargs).image_processor
93 +
94 + def get_video_processor(self, **kwargs: object) -> MiniMaxM3VLVideoProcessor:
95 + return self.get_hf_processor(**kwargs).video_processor
96 +
97 + def _get_vision_info(
98 + self,
99 + *,
100 + image_width: int,
101 + image_height: int,
102 + num_frames: int,
103 + image_processor,
104 + ) -> tuple[ImageSize, int]:
105 + """Compute resized image size and number of vision tokens.
106 +
107 + Mirrors the processor's Qwen-style ``smart_resize`` (area bound by
108 + ``max_pixels``) so token counts match the actual processor output.
109 + """
110 + patch_size: int = image_processor.patch_size
111 + merge_size: int = image_processor.merge_size
112 + temporal_patch_size: int = image_processor.temporal_patch_size
113 + factor = patch_size * merge_size
114 + max_pixels: int = image_processor.max_pixels
115 + # Long-side resize spec (opt-in). ``image_processor`` is the *video*
116 + # processor when counting video tokens, so read the bounds off it.
117 + max_long_side_pixel = getattr(image_processor, "max_long_side_pixel", None)
118 + min_short_side_pixel = getattr(
119 + image_processor, "min_short_side_pixel", MIN_SHORT_SIDE_PIXEL
120 + )
121 +
122 + new_h, new_w = smart_resize(
123 + image_height,
124 + image_width,
125 + factor=factor,
126 + max_pixels=max_pixels,
127 + max_long_side_pixel=max_long_side_pixel,
128 + min_short_side_pixel=min_short_side_pixel,
129 + # Token counting must not raise; the volumetric/area cap is enforced
130 + # in the processor's _preprocess on the real inputs.
131 + max_total_pixels=None,
132 + )
133 + grid_h = new_h // patch_size
134 + grid_w = new_w // patch_size
135 +
136 + # Pad frames to be divisible by temporal_patch_size
137 + padded_frames = num_frames + (-num_frames % temporal_patch_size)
138 + grid_t = max(padded_frames // temporal_patch_size, 1)
139 +
140 + num_tokens = grid_t * grid_h * grid_w // (merge_size**2)
141 + return ImageSize(width=new_w, height=new_h), num_tokens
142 +
143 + def get_num_image_tokens(
144 + self,
145 + *,
146 + image_width: int,
147 + image_height: int,
148 + image_processor,
149 + mm_kwargs: Mapping[str, object],
150 + ) -> int:
151 + _, n = self._get_vision_info(
152 + image_width=image_width,
153 + image_height=image_height,
154 + num_frames=1,
155 + image_processor=image_processor,
156 + )
157 + return n
158 +
159 + def get_num_video_tokens(
160 + self,
161 + *,
162 + image_width: int,
163 + image_height: int,
164 + num_frames: int,
165 + image_processor,
166 + mm_kwargs: Mapping[str, object],
167 + ) -> int:
168 + _, n = self._get_vision_info(
169 + image_width=image_width,
170 + image_height=image_height,
171 + num_frames=num_frames,
172 + image_processor=image_processor,
173 + )
174 + return n
175 +
176 + def get_image_size_with_most_features(self) -> ImageSize:
177 + # Largest square (a multiple of patch_size*merge_size) whose area is
178 + # within the image processor's bound — this yields the most vision
179 + # tokens for one image. With the long-side spec the square side is
180 + # capped by ``max_long_side_pixel`` (and the fixed ``max_total_pixels``);
181 + # otherwise it is bound by the ``max_pixels`` area.
182 + image_processor = self.get_image_processor()
183 + factor = image_processor.patch_size * image_processor.merge_size
184 + max_long_side_pixel = getattr(image_processor, "max_long_side_pixel", None)
185 + if max_long_side_pixel is not None:
186 + side_px = min(
187 + max_long_side_pixel,
188 + math.isqrt(image_processor.max_total_pixels),
189 + )
190 + else:
191 + side_px = math.isqrt(image_processor.max_pixels)
192 + side = max(factor, (side_px // factor) * factor)
193 + return ImageSize(width=side, height=side)
194 +
195 + def get_video_size_with_most_features(self) -> ImageSize:
196 + # Per-frame size that yields the most vision tokens, bound by the
197 + # *video* processor's ``max_pixels`` (which differs from the image
198 + # bound). Token count depends only on area, so maximize the area
199 + # achievable with both sides a multiple of patch_size*merge_size rather
200 + # than picking the largest square — a square (e.g. 756x756 for M3's
201 + # 602,112 bound) leaves area on the table, undercounting frames.
202 + video_processor = self.get_video_processor()
203 + factor = video_processor.patch_size * video_processor.merge_size
204 + per_frame_pixels = video_processor.max_pixels
205 + max_long_side_pixel = getattr(video_processor, "max_long_side_pixel", None)
206 + if max_long_side_pixel is not None:
207 + # Long-side spec: a frame's worst case is a square capped by
208 + # ``max_long_side_pixel`` (per-frame area, not the volumetric cap).
209 + per_frame_pixels = min(per_frame_pixels, max_long_side_pixel**2)
210 + units = per_frame_pixels // (factor * factor) # h_u * w_u
211 + h_u = math.isqrt(units)
212 + while units % h_u:
213 + h_u -= 1
214 + return ImageSize(width=(units // h_u) * factor, height=h_u * factor)
215 +
216 + def get_max_image_tokens(self) -> int:
217 + image_processor = self.get_image_processor()
218 + size = self.get_image_size_with_most_features()
219 + return self.get_num_image_tokens(
220 + image_width=size.width,
221 + image_height=size.height,
222 + image_processor=image_processor,
223 + mm_kwargs={},
224 + )
225 +
226 + def _get_max_video_frames(self, max_tokens: int) -> int:
227 + video_processor = self.get_video_processor()
228 + size = self.get_video_size_with_most_features()
229 + num_frames = 1
230 + while True:
231 + next_n = self.get_num_video_tokens(
232 + image_width=size.width,
233 + image_height=size.height,
234 + num_frames=num_frames + 1,
235 + image_processor=video_processor,
236 + mm_kwargs={},
237 + )
238 + if next_n > max_tokens:
239 + break
240 + num_frames += 1
241 + return num_frames
242 +
243 + def get_num_frames_with_most_features(
244 + self,
245 + seq_len: int,
246 + mm_counts: Mapping[str, int],
247 + max_frames_per_video: int = _MAX_FRAMES_PER_VIDEO,
248 + ) -> int:
249 + max_videos = mm_counts.get("video", 0)
250 + max_total_frames = self._get_max_video_frames(seq_len)
251 + max_frames_per_video = min(
252 + max_total_frames // max(max_videos, 1), max_frames_per_video
253 + )
254 + return max(max_frames_per_video, 1)
255 +
256 + def get_max_video_tokens(
257 + self,
258 + seq_len: int,
259 + mm_counts: Mapping[str, int],
260 + ) -> int:
261 + video_processor = self.get_video_processor()
262 + size = self.get_video_size_with_most_features()
263 + return self.get_num_video_tokens(
264 + image_width=size.width,
265 + image_height=size.height,
266 + num_frames=self.get_num_frames_with_most_features(seq_len, mm_counts),
267 + image_processor=video_processor,
268 + mm_kwargs={},
269 + )
270 +
271 +
272 +class MiniMaxM3VLDummyInputsBuilder(BaseDummyInputsBuilder[MiniMaxM3VLProcessingInfo]):
273 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
274 + num_images = mm_counts.get("image", 0)
275 + num_videos = mm_counts.get("video", 0)
276 + image_token: str = self.info.IMAGE_TOKEN
277 + video_token: str = self.info.VIDEO_TOKEN
278 + return image_token * num_images + video_token * num_videos
279 +
280 + def get_dummy_mm_data(
281 + self,
282 + seq_len: int,
283 + mm_counts: Mapping[str, int],
284 + mm_options: Mapping[str, BaseDummyOptions],
285 + ) -> MultiModalDataDict:
286 + size = self.info.get_image_size_with_most_features()
287 + video_size = self.info.get_video_size_with_most_features()
288 + num_frames = self.info.get_num_frames_with_most_features(seq_len, mm_counts)
289 + return {
290 + "image": self._get_dummy_images(
291 + width=size.width,
292 + height=size.height,
293 + num_images=mm_counts.get("image", 0),
294 + overrides=cast(ImageDummyOptions | None, mm_options.get("image")),
295 + ),
296 + "video": self._get_dummy_videos(
297 + width=video_size.width,
298 + height=video_size.height,
299 + num_frames=num_frames,
300 + num_videos=mm_counts.get("video", 0),
301 + overrides=cast(VideoDummyOptions | None, mm_options.get("video")),
302 + ),
303 + }
304 +
305 +
306 +class MiniMaxM3VLMultiModalProcessor(
307 + BaseMultiModalProcessor[MiniMaxM3VLProcessingInfo]
308 +):
309 + def _get_data_parser(self) -> MultiModalDataParser:
310 + # Request video metadata (fps + sampled frame indices) so the HF
311 + # processor can emit per-frame ``]<]X.X seconds[>[`` timestamp markers,
312 + # matching MiniMax's reference video token stream. ``_get_prompt_updates``
313 + # reconstructs the same markers from the metadata to keep the prompt
314 + # replacement aligned with the processor output.
315 + return MultiModalDataParser(video_needs_metadata=True)
316 +
317 + def _call_hf_processor(
318 + self,
319 + prompt: str,
320 + mm_data: Mapping[str, object],
321 + mm_kwargs: Mapping[str, object],
322 + tok_kwargs: Mapping[str, object],
323 + ) -> BatchFeature:
324 + mm_data = dict(mm_data)
325 + # With ``video_needs_metadata=True`` each video arrives as a
326 + # ``(frames, metadata)`` tuple. Split the frames back out and forward the
327 + # metadata as ``VideoMetadata`` so the processor emits timestamps.
328 + videos = cast(list | None, mm_data.get("videos"))
329 + video_metadata: list[VideoMetadata] | None = None
330 + if videos:
331 + frames_only = []
332 + video_metadata = []
333 + for item in videos:
334 + if isinstance(item, tuple) and len(item) == 2:
335 + frames, meta = item
336 + else:
337 + frames, meta = item, {}
338 + frames_only.append(frames)
339 + meta = {
340 + k: v for k, v in (meta or {}).items() if k != "do_sample_frames"
341 + }
342 + # VideoMetadata requires total_num_frames; derive it for
343 + # dummy/profiling videos whose metadata omits it. fps and
344 + # frames_indices default to None there → no timestamps, which
345 + # stays consistent with _get_prompt_updates.
346 + meta.setdefault("total_num_frames", len(frames))
347 + video_metadata.append(VideoMetadata(**meta))
348 + mm_data["videos"] = frames_only
349 +
350 + # Override the video processor's default do_resize=False (set for a
351 + # pre-resized pipeline) to True for vLLM's raw-frame inputs.
352 + merged = dict(do_resize=True, **mm_kwargs, **tok_kwargs)
353 + data = dict(text=prompt, **mm_data)
354 + if video_metadata is not None:
355 + data["video_metadata"] = video_metadata
356 + return self.info.ctx.call_hf_processor(
357 + self.info.get_hf_processor(**mm_kwargs),
358 + data,
359 + merged,
360 + )
361 +
362 + def _get_mm_fields_config(
363 + self,
364 + hf_inputs: BatchFeature,
365 + hf_processor_mm_kwargs: Mapping[str, object],
366 + ) -> Mapping[str, MultiModalFieldConfig]:
367 + image_grid_thw = hf_inputs.get("image_grid_thw")
368 + video_grid_thw = hf_inputs.get("video_grid_thw")
369 +
370 + # Total patches per item (grid_t * grid_h * grid_w)
371 + image_grid_sizes = (
372 + image_grid_thw.prod(-1)
373 + if image_grid_thw is not None
374 + else torch.empty(0, dtype=torch.long)
375 + )
376 + video_grid_sizes = (
377 + video_grid_thw.prod(-1)
378 + if video_grid_thw is not None
379 + else torch.empty(0, dtype=torch.long)
380 + )
381 +
382 + return {
383 + "pixel_values": MultiModalFieldConfig.flat_from_sizes(
384 + "image", image_grid_sizes
385 + ),
386 + "image_grid_thw": MultiModalFieldConfig.batched("image", keep_on_cpu=True),
387 + "pixel_values_videos": MultiModalFieldConfig.flat_from_sizes(
388 + "video", video_grid_sizes
389 + ),
390 + "video_grid_thw": MultiModalFieldConfig.batched("video", keep_on_cpu=True),
391 + }
392 +
393 + def _get_prompt_updates(
394 + self,
395 + mm_items: MultiModalDataItems,
396 + hf_processor_mm_kwargs: Mapping[str, object],
397 + out_mm_kwargs: MultiModalKwargsItems,
398 + ) -> Sequence[PromptUpdate]:
399 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
400 + tokenizer = self.info.get_tokenizer()
401 + vocab = tokenizer.get_vocab()
402 +
403 + image_token_id: int = vocab[self.info.IMAGE_TOKEN]
404 + video_token_id: int = vocab[self.info.VIDEO_TOKEN]
405 + start_token_id: int = vocab[self.info.VISION_START_TOKEN]
406 + end_token_id: int = vocab[self.info.VISION_END_TOKEN]
407 + merge_length: int = hf_processor.image_processor.merge_size**2
408 +
409 + def get_image_replacement(item_idx: int):
410 + grid_thw: torch.Tensor = out_mm_kwargs["image"][item_idx][
411 + "image_grid_thw"
412 + ].data
413 + # grid_thw shape: (3,) = [1, grid_h, grid_w]
414 + N = int(grid_thw.prod().item()) // merge_length
415 + full = [start_token_id] + [image_token_id] * N + [end_token_id]
416 + return PromptUpdateDetails.select_token_id(full, image_token_id)
417 +
418 + # Per-video metadata (fps + sampled frame indices) is carried on the
419 + # parsed video items; used to reproduce the HF processor's timestamps.
420 + video_items = mm_items.get("video")
421 + video_metadata = getattr(video_items, "metadata", None)
422 + temporal_patch_size: int = hf_processor.video_processor.temporal_patch_size
423 +
424 + def get_video_replacement(item_idx: int):
425 + grid_thw: torch.Tensor = out_mm_kwargs["video"][item_idx][
426 + "video_grid_thw"
427 + ].data
428 + # grid_thw shape: (3,) = [grid_t, grid_h, grid_w]
429 + # HF model uses VIDEO_TOKEN (not IMAGE_TOKEN) for video frame content:
430 + # processing_minimax.py L245: replace(placeholder, self.VIDEO_TOKEN)
431 + T = int(grid_thw[0].item())
432 + M = int(grid_thw[1].item() * grid_thw[2].item()) // merge_length
433 +
434 + # Reproduce the HF processor's per-frame timestamp markers
435 + # (processing_minimax.py: ts = frames_indices[frame*tps] / fps,
436 + # rendered as "]<]X.X seconds[>["). Falls back to no timestamps when
437 + # metadata is unavailable (keeping the replacement aligned with the
438 + # processor output in both cases).
439 + meta = (
440 + video_metadata[item_idx]
441 + if video_metadata is not None and item_idx < len(video_metadata)
442 + else None
443 + )
444 + fps = meta.get("fps") if meta else None
445 + frames_indices = meta.get("frames_indices") if meta else None
446 +
447 + full: list[int] = []
448 + for frame_idx in range(T):
449 + if fps is not None and frames_indices is not None:
450 + idx = min(frame_idx * temporal_patch_size, len(frames_indices) - 1)
451 + ts = frames_indices[idx] / fps
452 + full += tokenizer.encode(
453 + f"]<]{ts:.1f} seconds[>[", add_special_tokens=False
454 + )
455 + full += [start_token_id] + [video_token_id] * M + [end_token_id]
456 + return PromptUpdateDetails.select_token_id(full, video_token_id)
457 +
458 + return [
459 + PromptReplacement(
460 + modality="image",
461 + target=[image_token_id],
462 + replacement=get_image_replacement,
463 + ),
464 + PromptReplacement(
465 + modality="video",
466 + target=[video_token_id],
467 + replacement=get_video_replacement,
468 + ),
469 + ]
470 +
471 +
472 +# TODO(Isotr0py): Tie with MinimaxVideoProcessor
473 +# after https://github.com/vllm-project/vllm/pull/44126
474 +@VIDEO_LOADER_REGISTRY.register("minimax_m3_vl")
475 +class MiniMaxM3VideoBackend(VideoBackend):
476 + @classmethod
477 + def compute_frames_index_to_sample(
478 + cls,
479 + source: VideoSourceMetadata,
480 + target: VideoTargetMetadata,
481 + **kwargs,
482 + ) -> list[int]:
483 + total_frames = source.total_frames_num
484 + video_fps = source.original_fps
485 + fps = target.fps
486 +
487 + if total_frames <= 0 or video_fps <= 0 or fps <= 0:
488 + return [0] if total_frames > 0 else []
489 +
490 + read_time_interval = 1.0 / fps
491 + eps = 1e-4
492 +
493 + indices: list[int] = []
494 + prev_kept_ts = -float("inf")
495 + while True:
496 + if not indices:
497 + target_frame = 0
498 + else:
499 + target_ts = prev_kept_ts + read_time_interval - eps
500 + target_frame = math.ceil(target_ts * video_fps)
501 + target_frame = max(target_frame, indices[-1] + 1)
502 + if target_frame >= total_frames:
503 + break
504 + indices.append(target_frame)
505 + prev_kept_ts = target_frame / video_fps
506 +
507 + last_frame_idx = total_frames - 1
508 + last_ts = last_frame_idx / video_fps
509 + if indices and indices[-1] != last_frame_idx and last_ts - prev_kept_ts > eps:
510 + indices.append(last_frame_idx)
511 +
512 + if not indices:
513 + indices = [0]
514 + return indices
vllm/models/minimax_m3/common/ops/__init__.py new
+18
@@ -0,0 +1,18 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Cross-platform (Triton) kernels for MiniMax M3 sparse attention."""
4 +
5 +from .index_topk import (
6 + minimax_m3_index_decode,
7 + minimax_m3_index_score,
8 + minimax_m3_index_topk,
9 +)
10 +from .sparse_attn import minimax_m3_sparse_attn, minimax_m3_sparse_attn_decode
11 +
12 +__all__ = [
13 + "minimax_m3_index_decode",
14 + "minimax_m3_index_score",
15 + "minimax_m3_index_topk",
16 + "minimax_m3_sparse_attn",
17 + "minimax_m3_sparse_attn_decode",
18 +]
vllm/models/minimax_m3/common/ops/index_topk.py new
+898
@@ -0,0 +1,898 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Triton kernels for MiniMax M3 lightning-indexer block scoring + top-k.
4 +
5 +Index queries score each 128-token block of index keys (max over the block),
6 +then the top-k blocks (plus forced init/local blocks) are selected per query
7 +token. Adapted to vLLM's paged KV cache: the KV page size is forced to equal the
8 +sparse block size (128), so one sparse block maps to exactly one page.
9 +
10 +Index-K cache layout (vLLM): ``(num_blocks, 128, idx_head_dim)`` (single head).
11 +
12 +Only the paths MiniMax M3 uses are implemented: score_type="max", index value
13 +disabled (score-only indexer), single shared index head. The selected block ids
14 +feed the block-sparse attention kernels in ``sparse_attn``.
15 +"""
16 +
17 +import torch
18 +
19 +from vllm.platforms import current_platform
20 +from vllm.triton_utils import tl, triton
21 +from vllm.utils.math_utils import round_up
22 +
23 +# One sparse block == one KV page.
24 +SPARSE_BLOCK_SIZE = 128
25 +
26 +
27 +# ---------------------------------------------------------------------------
28 +# Bitonic top-k helpers (layout-agnostic).
29 +# ---------------------------------------------------------------------------
30 +@triton.jit
31 +def _compare_and_swap(x, ids, flip, i: tl.constexpr, n_dims: tl.constexpr):
32 + n_outer: tl.constexpr = x.numel >> n_dims
33 + shape: tl.constexpr = [n_outer * 2**i, 2, 2 ** (n_dims - i - 1)]
34 + y = tl.reshape(x, shape)
35 + mask = tl.arange(0, 2)[None, :, None]
36 + left = tl.broadcast_to(tl.sum(y * (1 - mask), 1)[:, None, :], shape).to(y.dtype)
37 + right = tl.broadcast_to(tl.sum(y * mask, 1)[:, None, :], shape).to(y.dtype)
38 + left = tl.reshape(left, x.shape)
39 + right = tl.reshape(right, x.shape)
40 + y_idx = tl.reshape(ids, shape)
41 + left_idx = tl.broadcast_to(tl.sum(y_idx * (1 - mask), 1)[:, None, :], shape)
42 + right_idx = tl.broadcast_to(tl.sum(y_idx * mask, 1)[:, None, :], shape)
43 + left_idx = tl.reshape(left_idx, x.shape).to(y_idx.dtype)
44 + right_idx = tl.reshape(right_idx, x.shape).to(y_idx.dtype)
45 + idtype = tl.core.get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True)
46 + ileft = left.to(idtype, bitcast=True)
47 + iright = right.to(idtype, bitcast=True)
48 + ix = x.to(idtype, bitcast=True)
49 + cond = (left > right) != flip
50 + ret = ix ^ tl.where(cond, ileft ^ iright, tl.zeros_like(ix))
51 + new_ids = ids ^ tl.where(cond, left_idx ^ right_idx, tl.zeros_like(ids))
52 + return ret.to(x.dtype, bitcast=True), new_ids
53 +
54 +
55 +@triton.jit
56 +def _bitonic_merge(
57 + x, ids, stage: tl.constexpr, order: tl.constexpr, n_dims: tl.constexpr
58 +):
59 + n_outer: tl.constexpr = x.numel >> n_dims
60 + tl.static_assert(stage <= n_dims)
61 + if order == 2:
62 + shape: tl.constexpr = [n_outer * 2 ** (n_dims - 1 - stage), 2, 2**stage]
63 + flip = tl.reshape(
64 + tl.broadcast_to(tl.arange(0, 2)[None, :, None], shape), x.shape
65 + )
66 + else:
67 + flip = order
68 + for i in tl.static_range(stage):
69 + x, ids = _compare_and_swap(x, ids, flip, i + (n_dims - stage), n_dims)
70 + return x, ids
71 +
72 +
73 +# ---------------------------------------------------------------------------
74 +# Index block-score kernel (paged). score[h, token, block] = max over the
75 +# 128-token block of (idx_q . index_k), causal-masked. BLOCK_SIZE_K == 128 so
76 +# each K-tile is exactly one page (BLOCKS_PER_K_BLOCK == 1).
77 +# ---------------------------------------------------------------------------
78 +# since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens
79 +# might lose pointer alignment, which trigger Triton recompiles. we don't actually
80 +# need pointer alignment for those tensors anyway because we do scalar load.
81 +@triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"])
82 +def _index_block_score_kernel(
83 + q_ptr, # idx_q: [total_q, num_idx_heads, head_dim]
84 + ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim]
85 + score_ptr, # [num_idx_heads, total_q, max_block]
86 + block_table_ptr, # [num_reqs, max_blocks]
87 + cu_seqlens, # [batch+1] query start offsets
88 + seq_lens, # [batch] total K length
89 + prefix_lens, # [batch] context length before this chunk's queries
90 + num_idx_heads,
91 + head_dim: tl.constexpr,
92 + sm_scale,
93 + stride_q_n,
94 + stride_q_h,
95 + stride_q_d,
96 + stride_ik_blk,
97 + stride_ik_pos,
98 + stride_ik_d,
99 + stride_s_h,
100 + stride_s_n,
101 + stride_s_k,
102 + stride_bt_b,
103 + BLOCK_SIZE_Q: tl.constexpr,
104 + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
105 +):
106 + sm_scale_log2e = sm_scale * 1.4426950409
107 + pid_q = tl.program_id(0)
108 + pid_bh = tl.program_id(1)
109 + pid_b = pid_bh // num_idx_heads
110 + pid_h = pid_bh % num_idx_heads
111 +
112 + seq_start = tl.load(cu_seqlens + pid_b)
113 + q_len = tl.load(cu_seqlens + pid_b + 1) - seq_start
114 + seq_len = tl.load(seq_lens + pid_b)
115 + prefix_len = tl.load(prefix_lens + pid_b)
116 + if BLOCK_SIZE_Q * pid_q >= q_len:
117 + return
118 +
119 + q_ptrs = tl.make_block_ptr(
120 + base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h,
121 + shape=(q_len, head_dim),
122 + strides=(stride_q_n, stride_q_d),
123 + offsets=(pid_q * BLOCK_SIZE_Q, 0),
124 + block_shape=(BLOCK_SIZE_Q, head_dim),
125 + order=(1, 0),
126 + )
127 + q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero")
128 + q_start = prefix_len + pid_q * BLOCK_SIZE_Q
129 +
130 + off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len
131 + off_k = tl.arange(0, BLOCK_SIZE_K)
132 + off_d = tl.arange(0, head_dim)
133 + # Block table row for this request.
134 + bt_row = block_table_ptr + pid_b * stride_bt_b
135 + # Causal window: only blocks up to the last query token's position.
136 + hi = min(seq_len, prefix_len + (pid_q + 1) * BLOCK_SIZE_Q)
137 + for i in tl.range(0, hi, BLOCK_SIZE_K):
138 + blk = i // BLOCK_SIZE_K
139 + page = tl.load(bt_row + blk).to(tl.int64)
140 + pos = i + off_k
141 + # index-K for this page: [BLOCK_SIZE_D, BLOCK_SIZE_K] (transposed)
142 + # we don't need masked load for K, because KV cache ensures
143 + # allocation is multiple of BLOCK_SIZE_K.
144 + # for tokens beyond seqlen, they will be masked in qk later.
145 + k = tl.load(
146 + ik_cache_ptr
147 + + page * stride_ik_blk
148 + + off_k[None, :] * stride_ik_pos
149 + + off_d[:, None] * stride_ik_d,
150 + )
151 + qk = tl.dot(q, k) * sm_scale_log2e
152 + # apply causal mask as needed
153 + if q_start < i + BLOCK_SIZE_K:
154 + qk = tl.where(off_q[:, None] >= pos[None, :], qk, float("-inf"))
155 + # one sparse block per K-tile -> max over the 128 positions
156 + score = tl.max(qk, axis=1) # [BLOCK_SIZE_Q]
157 + s_ptrs = (
158 + score_ptr
159 + + pid_h * stride_s_h
160 + + (seq_start + pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q))
161 + * stride_s_n
162 + + blk * stride_s_k
163 + )
164 + q_store_mask = (pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q)) < q_len
165 + tl.store(s_ptrs, score, mask=q_store_mask)
166 +
167 +
168 +# ---------------------------------------------------------------------------
169 +# Top-k selection over per-token block scores (layout-agnostic). block_size_q
170 +# is 1 for M3, so top-k is computed per query token.
171 +# ---------------------------------------------------------------------------
172 +# since prefill metadata is sliced from mixed batch metadata, prefix_lens
173 +# might lose pointer alignment, which trigger Triton recompiles. we don't actually
174 +# need pointer alignment for those tensors anyway because we do scalar load.
175 +@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])})
176 +@triton.autotune(
177 + configs=[
178 + triton.Config({"BLOCK_SIZE_K": 2048}, num_warps=8, num_stages=2),
179 + triton.Config({"BLOCK_SIZE_K": 1024}, num_warps=8, num_stages=2),
180 + triton.Config({"BLOCK_SIZE_K": 512}, num_warps=8, num_stages=2),
181 + triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2),
182 + triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2),
183 + triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2),
184 + ],
185 + key=["BLOCK_SIZE_T"],
186 +)
187 +@triton.jit(do_not_specialize_on_alignment=["prefix_lens"])
188 +def _topk_index_kernel(
189 + s_ptr, # [num_heads, total_q, max_block]
190 + ti_ptr, # [num_heads, total_q, topk]
191 + sample_interval: tl.constexpr, # block_size_q (1 for M3)
192 + block_size: tl.constexpr, # sparse block size (128)
193 + cu_seqlens,
194 + cu_seqblocks_q,
195 + prefix_lens,
196 + topk,
197 + init_blocks: tl.constexpr,
198 + local_blocks: tl.constexpr,
199 + stride_s_h,
200 + stride_s_n,
201 + stride_s_k,
202 + stride_ti_h,
203 + stride_ti_n,
204 + stride_ti_t,
205 + BLOCK_SIZE_K: tl.constexpr,
206 + BLOCK_SIZE_T: tl.constexpr,
207 + MASK_INIT: tl.constexpr,
208 + MASK_LOCAL: tl.constexpr,
209 +):
210 + tl.static_assert(BLOCK_SIZE_K > BLOCK_SIZE_T)
211 + pid_q = tl.program_id(0)
212 + pid_b = tl.program_id(1)
213 + pid_h = tl.program_id(2)
214 + seq_start = tl.load(cu_seqlens + pid_b)
215 + block_start = tl.load(cu_seqblocks_q + pid_b)
216 + block_num = tl.load(cu_seqblocks_q + pid_b + 1) - block_start
217 + prefix_len = tl.load(prefix_lens + pid_b)
218 + if pid_q >= block_num:
219 + return
220 + off_k = tl.arange(0, BLOCK_SIZE_K)
221 + off_t = tl.arange(0, BLOCK_SIZE_T)
222 + s_ptrs = (
223 + s_ptr
224 + + (seq_start + pid_q * sample_interval) * stride_s_n
225 + + pid_h * stride_s_h
226 + + off_k * stride_s_k
227 + )
228 + topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32)
229 + topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32)
230 + left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2
231 + valid_blocks = (prefix_len + pid_q * sample_interval + block_size) // block_size
232 + for i in tl.range(0, valid_blocks, BLOCK_SIZE_K):
233 + causal_mask = i + off_k < valid_blocks
234 + local_mask = i + off_k >= max(0, valid_blocks - local_blocks)
235 + init_mask = i + off_k < init_blocks
236 + score = tl.load(s_ptrs, mask=causal_mask, other=-1e30).to(tl.float32)
237 + score = tl.where(score != score, -1e30, score)
238 + s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K
239 + if MASK_INIT:
240 + score = tl.where(causal_mask & init_mask, score - 1e29, score)
241 + else:
242 + score = tl.where(causal_mask & init_mask, 1e30, score)
243 + if MASK_LOCAL:
244 + score = tl.where(causal_mask & local_mask, score - 1e28, score)
245 + else:
246 + score = tl.where(causal_mask & local_mask, 1e29, score)
247 + topk_score, last_topk_score = score, topk_score
248 + topk_idx, last_topk_idx = (tl.where(causal_mask, i + off_k + 1, 0), topk_idx)
249 + n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
250 + for j in tl.static_range(1, n_dims):
251 + topk_score, topk_idx = _bitonic_merge(
252 + topk_score, topk_idx.to(tl.int32), j, 2, n_dims
253 + )
254 + if i != 0:
255 + topk_score, topk_idx = _bitonic_merge(
256 + topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims
257 + )
258 + topk_score_new = last_topk_score * left_half_mask + topk_score * (
259 + 1 - left_half_mask
260 + )
261 + topk_idx_new = last_topk_idx * left_half_mask + topk_idx * (
262 + 1 - left_half_mask
263 + )
264 + topk_score, topk_idx = _bitonic_merge(
265 + topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims
266 + )
267 + else:
268 + topk_score, topk_idx = _bitonic_merge(
269 + topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims
270 + )
271 + topk_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
272 + topk_idx = tl.sum(
273 + topk_mask[:, None]
274 + * tl.reshape(topk_idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
275 + axis=0,
276 + )
277 + ti_ptrs = (
278 + ti_ptr
279 + + (block_start + pid_q) * stride_ti_n
280 + + pid_h * stride_ti_h
281 + + off_t * stride_ti_t
282 + )
283 + store_mask = off_t < topk
284 + valid_mask = off_t < valid_blocks
285 + topk_idx = tl.where(store_mask & valid_mask, topk_idx, -1)
286 + tl.store(ti_ptrs, topk_idx.to(ti_ptrs.dtype.element_ty), mask=store_mask)
287 +
288 +
289 +# ---------------------------------------------------------------------------
290 +# Decode index-score kernel (split-K over seq blocks). Decode batches are
291 +# flattened request-major, with a runtime query length used to map each query
292 +# token back to its request metadata. Chunk counts depend only on shape
293 +# constants so the grid is fixed within a cuda graph. Base-2 (exp2/log2)
294 +# softmax matches prefill.
295 +# ---------------------------------------------------------------------------
296 +@triton.jit(do_not_specialize=["num_kv_chunks", "decode_query_len"])
297 +def _decode_index_score_kernel(
298 + q_ptr, # idx_q: [total_q, num_idx_heads, head_dim]
299 + ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim]
300 + score_ptr, # [num_idx_heads, total_q, max_block]
301 + block_table_ptr, # [num_reqs, max_blocks]
302 + seq_lens, # [num_reqs]
303 + num_idx_heads: tl.constexpr,
304 + head_dim: tl.constexpr,
305 + init_blocks,
306 + local_blocks,
307 + sm_scale,
308 + decode_query_len,
309 + stride_q_n,
310 + stride_q_h,
311 + stride_q_d,
312 + stride_ik_blk,
313 + stride_ik_pos,
314 + stride_ik_d,
315 + stride_s_h,
316 + stride_s_n,
317 + stride_s_k,
318 + stride_bt_b,
319 + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
320 + num_kv_chunks,
321 + USE_PDL: tl.constexpr,
322 +):
323 + sm_scale_log2e = sm_scale * 1.4426950409
324 + pid_b = tl.program_id(0) # flattened query-token id
325 + pid_c = tl.program_id(1)
326 + req_id = pid_b // decode_query_len
327 + q_offset = pid_b - req_id * decode_query_len
328 +
329 + if USE_PDL:
330 + tl.extra.cuda.gdc_wait()
331 + tl.extra.cuda.gdc_launch_dependents()
332 +
333 + seq_len = tl.load(seq_lens + req_id)
334 + query_pos = seq_len - decode_query_len + q_offset
335 + # Full-CG padding uses zero-length request rows. Clamp to an empty
336 + # attention range instead of letting padded rows produce negative lengths.
337 + kv_len = tl.maximum(query_pos + 1, 0)
338 + num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
339 +
340 + # block-aligned fixed-count split: grid independent of seq_len (cuda graph).
341 + chunk_size_blocks = (num_blocks + num_kv_chunks - 1) // num_kv_chunks
342 + chunk_start_block = pid_c * chunk_size_blocks
343 + chunk_end_block = tl.minimum(chunk_start_block + chunk_size_blocks, num_blocks)
344 + if chunk_start_block >= chunk_end_block:
345 + return
346 + off_k = tl.arange(0, BLOCK_SIZE_K) # positions within a 128-block
347 + off_d = tl.arange(0, head_dim)
348 + bt_row = block_table_ptr + req_id * stride_bt_b
349 + # Force-select init (1e30) and local (1e29, higher priority) blocks.
350 + local_start = tl.maximum(0, num_blocks - local_blocks)
351 + # query vectors across all heads
352 + q = tl.load(
353 + q_ptr
354 + + pid_b * stride_q_n
355 + + tl.arange(0, num_idx_heads) * stride_q_h
356 + + off_d[:, None] * stride_q_d,
357 + ) # [D,H]
358 + for blk in tl.range(chunk_start_block, chunk_end_block):
359 + page = tl.load(bt_row + blk).to(tl.int64)
360 + pos = blk * BLOCK_SIZE_K + off_k
361 + pos_mask = pos < kv_len
362 + # we don't need masked load for K, because KV cache ensures
363 + # allocation is multiple of BLOCK_SIZE_K.
364 + # for tokens beyond seqlen, they will be masked in qk later.
365 + k = tl.load(
366 + ik_cache_ptr
367 + + page * stride_ik_blk
368 + + off_k[:, None] * stride_ik_pos
369 + + off_d * stride_ik_d,
370 + ) # [N,D]
371 + kq = tl.dot(k, q) * sm_scale_log2e # [N,H]
372 + kq = tl.where(pos_mask[:, None], kq, float("-inf"))
373 + score = tl.max(kq, axis=0) # [H]
374 + is_init = blk < init_blocks
375 + is_local = (blk >= local_start) & (blk < num_blocks)
376 + score = tl.where(is_local, 1e29, tl.where(is_init, 1e30, score))
377 + tl.store(
378 + score_ptr
379 + + tl.arange(0, num_idx_heads) * stride_s_h
380 + + pid_b * stride_s_n
381 + + blk * stride_s_k,
382 + score,
383 + )
384 +
385 +
386 +# ---------------------------------------------------------------------------
387 +# Decode top-k (split-K): per-chunk partial top-k + merge. Forced init/local
388 +# blocks are already encoded in the scores.
389 +# ---------------------------------------------------------------------------
390 +@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])})
391 +@triton.autotune(
392 + configs=[
393 + triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2),
394 + triton.Config({"BLOCK_SIZE_K": 256}, num_warps=4, num_stages=2),
395 + triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2),
396 + triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=3),
397 + triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2),
398 + ],
399 + key=["topk"],
400 +)
401 +@triton.jit(do_not_specialize=["chunk_blocks", "decode_query_len"])
402 +def _topk_index_partial_kernel(
403 + s_ptr, # score: [num_idx_heads, total_q, max_block]
404 + ts_partial_ptr, # partial scores out: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T]
405 + ti_partial_ptr, # partial idx out (1-indexed global, 0=invalid): same shape
406 + seq_lens, # [num_reqs]
407 + block_size: tl.constexpr, # sparse block size (128)
408 + topk: tl.constexpr,
409 + chunk_blocks, # how many score-blocks each chunk owns
410 + decode_query_len,
411 + stride_s_h,
412 + stride_s_b,
413 + stride_s_k,
414 + stride_ts_c,
415 + stride_ts_h,
416 + stride_ts_b,
417 + stride_ts_t,
418 + stride_ti_c,
419 + stride_ti_h,
420 + stride_ti_b,
421 + stride_ti_t,
422 + BLOCK_SIZE_K: tl.constexpr,
423 + BLOCK_SIZE_T: tl.constexpr,
424 + USE_PDL: tl.constexpr,
425 +):
426 + tl.static_assert(topk < BLOCK_SIZE_K)
427 + pid_b = tl.program_id(0) # flattened query-token id
428 + pid_h = tl.program_id(1)
429 + pid_chunk = tl.program_id(2)
430 + req_id = pid_b // decode_query_len
431 + q_offset = pid_b - req_id * decode_query_len
432 +
433 + if USE_PDL:
434 + tl.extra.cuda.gdc_wait()
435 +
436 + seq_len = tl.load(seq_lens + req_id)
437 + query_pos = seq_len - decode_query_len + q_offset
438 + # Full-CG padding uses zero-length request rows. Clamp to an empty
439 + # attention range instead of letting padded rows produce negative lengths.
440 + kv_len = tl.maximum(query_pos + 1, 0)
441 + num_blocks = (kv_len + block_size - 1) // block_size
442 +
443 + # Slice this chunk owns within [0, num_blocks).
444 + chunk_start = pid_chunk * chunk_blocks
445 + chunk_end = tl.minimum(chunk_start + chunk_blocks, num_blocks)
446 + chunk_actual = tl.maximum(chunk_end - chunk_start, 0)
447 +
448 + off_k = tl.arange(0, BLOCK_SIZE_K)
449 + off_t = tl.arange(0, BLOCK_SIZE_T)
450 +
451 + s_ptrs = (
452 + s_ptr
453 + + pid_b * stride_s_b
454 + + pid_h * stride_s_h
455 + + (chunk_start + off_k) * stride_s_k
456 + )
457 +
458 + topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32)
459 + topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32)
460 + left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2
461 +
462 + # Streaming top-K within this chunk. tl.range(0, 0) is a no-op so empty
463 + # chunks (chunk_actual == 0) skip the body and store sentinel -1e30 / 0.
464 + for i in tl.range(0, chunk_actual, BLOCK_SIZE_K):
465 + mask = off_k < chunk_actual - i
466 + score = tl.load(s_ptrs, mask=mask, other=-1e30).to(tl.float32)
467 + score = tl.where(score != score, -1e30, score)
468 + s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K
469 + topk_score, last_topk_score = score, topk_score
470 + topk_idx, last_topk_idx = (
471 + tl.where(mask, chunk_start + i + off_k + 1, 0), # 1-indexed global
472 + topk_idx,
473 + )
474 + n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
475 + for j in tl.static_range(1, n_dims):
476 + topk_score, topk_idx = _bitonic_merge(
477 + topk_score, topk_idx.to(tl.int32), j, 2, n_dims
478 + )
479 + if i != 0:
480 + topk_score, topk_idx = _bitonic_merge(
481 + topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims
482 + )
483 + topk_score_new = last_topk_score * left_half_mask + topk_score * (
484 + 1 - left_half_mask
485 + )
486 + topk_idx_new = last_topk_idx * left_half_mask + topk_idx * (
487 + 1 - left_half_mask
488 + )
489 + topk_score, topk_idx = _bitonic_merge(
490 + topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims
491 + )
492 + else:
493 + topk_score, topk_idx = _bitonic_merge(
494 + topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims
495 + )
496 +
497 + if USE_PDL:
498 + tl.extra.cuda.gdc_launch_dependents()
499 +
500 + # Extract first BLOCK_SIZE_T entries (top-K of this chunk after the sort).
501 + topk_mask_extract = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
502 + final_score = tl.sum(
503 + topk_mask_extract[:, None]
504 + * tl.reshape(topk_score, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
505 + axis=0,
506 + )
507 + final_idx = tl.sum(
508 + topk_mask_extract[:, None]
509 + * tl.reshape(topk_idx, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
510 + axis=0,
511 + )
512 +
513 + # Always write all BLOCK_SIZE_T slots — invalid slots carry -1e30 / 0
514 + # sentinels and lose to real scores in the merge stage.
515 + ts_ptrs = (
516 + ts_partial_ptr
517 + + pid_chunk * stride_ts_c
518 + + pid_b * stride_ts_b
519 + + pid_h * stride_ts_h
520 + + off_t * stride_ts_t
521 + )
522 + ti_ptrs = (
523 + ti_partial_ptr
524 + + pid_chunk * stride_ti_c
525 + + pid_b * stride_ti_b
526 + + pid_h * stride_ti_h
527 + + off_t * stride_ti_t
528 + )
529 + tl.store(ts_ptrs, final_score)
530 + tl.store(ti_ptrs, final_idx)
531 +
532 +
533 +@triton.heuristics(
534 + {
535 + "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"]),
536 + "BLOCK_SIZE_K": lambda args: triton.next_power_of_2(
537 + args["num_topk_chunks"] * triton.next_power_of_2(args["topk"])
538 + ),
539 + }
540 +)
541 +@triton.jit(do_not_specialize=["num_topk_chunks", "decode_query_len"])
542 +def _topk_index_merge_kernel(
543 + ts_partial_ptr, # partial scores: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T]
544 + ti_partial_ptr, # partial idx (1-indexed global, 0=invalid): same shape
545 + ti_final_ptr, # final idx (0-indexed, -1=invalid): [num_idx_heads, total_q, topk]
546 + seq_lens, # [num_reqs]
547 + block_size: tl.constexpr, # sparse block size (128)
548 + topk: tl.constexpr,
549 + decode_query_len,
550 + stride_ts_c,
551 + stride_ts_h,
552 + stride_ts_b,
553 + stride_ts_t,
554 + stride_ti_c,
555 + stride_ti_h,
556 + stride_ti_b,
557 + stride_ti_t,
558 + stride_tif_h,
559 + stride_tif_b,
560 + stride_tif_t,
561 + num_topk_chunks,
562 + BLOCK_SIZE_K: tl.constexpr,
563 + BLOCK_SIZE_T: tl.constexpr,
564 + USE_PDL: tl.constexpr,
565 +):
566 + pid_b = tl.program_id(0) # flattened query-token id
567 + pid_h = tl.program_id(1)
568 + req_id = pid_b // decode_query_len
569 + q_offset = pid_b - req_id * decode_query_len
570 +
571 + if USE_PDL:
572 + tl.extra.cuda.gdc_wait()
573 + tl.extra.cuda.gdc_launch_dependents()
574 +
575 + seq_len = tl.load(seq_lens + req_id)
576 + query_pos = seq_len - decode_query_len + q_offset
577 + # Full-CG padding uses zero-length request rows. Clamp to an empty
578 + # attention range instead of letting padded rows produce negative lengths.
579 + kv_len = tl.maximum(query_pos + 1, 0)
580 + num_blocks = (kv_len + block_size - 1) // block_size
581 +
582 + # Load NUM_TOPK_CHUNKS * BLOCK_SIZE_T candidates, padded to BLOCK_SIZE_K.
583 + # Candidate at flat position p comes from chunk = p // BLOCK_SIZE_T,
584 + # in_chunk = p % BLOCK_SIZE_T.
585 + off = tl.arange(0, BLOCK_SIZE_K)
586 + chunk_idx = off // BLOCK_SIZE_T
587 + in_chunk_idx = off % BLOCK_SIZE_T
588 + valid = chunk_idx < num_topk_chunks
589 +
590 + score_offset = (
591 + chunk_idx * stride_ts_c
592 + + pid_h * stride_ts_h
593 + + pid_b * stride_ts_b
594 + + in_chunk_idx * stride_ts_t
595 + )
596 + idx_offset = (
597 + chunk_idx * stride_ti_c
598 + + pid_h * stride_ti_h
599 + + pid_b * stride_ti_b
600 + + in_chunk_idx * stride_ti_t
601 + )
602 +
603 + score = tl.load(ts_partial_ptr + score_offset, mask=valid, other=-1e30).to(
604 + tl.float32
605 + )
606 + score = tl.where(score != score, -1e30, score)
607 + idx = tl.load(ti_partial_ptr + idx_offset, mask=valid, other=0).to(tl.int32)
608 +
609 + # Full bitonic descending sort of BLOCK_SIZE_K items.
610 + n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
611 + for j in tl.static_range(1, n_dims):
612 + score, idx = _bitonic_merge(score, idx.to(tl.int32), j, 2, n_dims)
613 + score, idx = _bitonic_merge(score, idx.to(tl.int32), n_dims, True, n_dims)
614 +
615 + # Extract first BLOCK_SIZE_T positions — these are the global top-K.
616 + extract_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
617 + topk_idx_final = tl.sum(
618 + extract_mask[:, None]
619 + * tl.reshape(idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
620 + axis=0,
621 + )
622 +
623 + off_t = tl.arange(0, BLOCK_SIZE_T)
624 + tif_ptrs = (
625 + ti_final_ptr
626 + + pid_h * stride_tif_h
627 + + pid_b * stride_tif_b
628 + + off_t * stride_tif_t
629 + )
630 + store_mask = off_t < topk
631 + topk_idx_final = tl.where(off_t < tl.minimum(topk, num_blocks), topk_idx_final, -1)
632 + tl.store(
633 + tif_ptrs, topk_idx_final.to(ti_final_ptr.dtype.element_ty), mask=store_mask
634 + )
635 +
636 +
637 +# ---------------------------------------------------------------------------
638 +# Python wrappers
639 +# ---------------------------------------------------------------------------
640 +@torch.no_grad()
641 +def minimax_m3_index_score(
642 + idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim]
643 + index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim]
644 + block_table: torch.Tensor, # [batch, max_blocks]
645 + cu_seqlens_q: torch.Tensor, # [batch+1] int32
646 + seq_lens: torch.Tensor, # [batch] int32
647 + prefix_lens: torch.Tensor, # [batch] int32
648 + max_query_len: int,
649 + max_seq_len: int,
650 + num_kv_heads: int,
651 + sm_scale: float,
652 +) -> torch.Tensor:
653 + """Compute per-token index scores for each visible sparse block.
654 +
655 + Returns score [num_kv_heads, total_q, max_block], where each score is the
656 + max over a 128-token index-K block. M3 has num_idx_heads == num_kv_heads.
657 + """
658 + total_q, num_idx_heads, head_dim = idx_q.shape
659 + assert num_idx_heads == num_kv_heads, (
660 + "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)"
661 + )
662 + batch = cu_seqlens_q.shape[0] - 1
663 + max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
664 +
665 + # Keep score strides 16-divisible to avoid Triton recompiles.
666 + score_block_stride = round_up(max_block, 16)
667 + score = torch.empty(
668 + (num_idx_heads, total_q, score_block_stride),
669 + dtype=torch.float32,
670 + device=idx_q.device,
671 + )
672 + BLOCK_SIZE_Q = 64
673 + grid_score = (triton.cdiv(max_query_len, BLOCK_SIZE_Q), batch * num_idx_heads)
674 + _index_block_score_kernel[grid_score](
675 + idx_q,
676 + index_kv_cache,
677 + score,
678 + block_table,
679 + cu_seqlens_q,
680 + seq_lens,
681 + prefix_lens,
682 + num_idx_heads,
683 + head_dim,
684 + sm_scale,
685 + idx_q.stride(0),
686 + idx_q.stride(1),
687 + idx_q.stride(2),
688 + index_kv_cache.stride(0),
689 + index_kv_cache.stride(1),
690 + index_kv_cache.stride(2),
691 + score.stride(0),
692 + score.stride(1),
693 + score.stride(2),
694 + block_table.stride(0),
695 + BLOCK_SIZE_Q=BLOCK_SIZE_Q,
696 + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
697 + )
698 + return score
699 +
700 +
701 +@torch.no_grad()
702 +def minimax_m3_index_topk(
703 + score: torch.Tensor, # [num_idx_heads, total_q, max_block]
704 + cu_seqlens_q: torch.Tensor, # [batch+1] int32
705 + prefix_lens: torch.Tensor, # [batch] int32
706 + max_query_len: int,
707 + topk: int,
708 + init_blocks: int,
709 + local_blocks: int,
710 +) -> torch.Tensor:
711 + """Select index top-k from a precomputed score tensor."""
712 + num_idx_heads = score.shape[0]
713 + batch = cu_seqlens_q.shape[0] - 1
714 + total_q = score.shape[1]
715 + topk_idx = torch.empty(
716 + (num_idx_heads, total_q, topk),
717 + dtype=torch.int32,
718 + device=score.device,
719 + )
720 + # block_size_q == 1 -> query blocks coincide with query tokens.
721 + grid_topk = (max_query_len, batch, num_idx_heads)
722 + _topk_index_kernel[grid_topk](
723 + score,
724 + topk_idx,
725 + 1, # sample_interval (block_size_q)
726 + SPARSE_BLOCK_SIZE,
727 + cu_seqlens_q,
728 + cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1
729 + prefix_lens,
730 + topk,
731 + init_blocks,
732 + local_blocks,
733 + score.stride(0),
734 + score.stride(1),
735 + score.stride(2),
736 + topk_idx.stride(0),
737 + topk_idx.stride(1),
738 + topk_idx.stride(2),
739 + MASK_INIT=False,
740 + MASK_LOCAL=False,
741 + )
742 + return topk_idx
743 +
744 +
745 +@torch.no_grad()
746 +def minimax_m3_index_decode(
747 + idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim]
748 + index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim]
749 + block_table: torch.Tensor, # [num_reqs, max_blocks]
750 + seq_lens: torch.Tensor, # [num_reqs] int32
751 + max_seq_len: int,
752 + topk: int,
753 + init_blocks: int,
754 + local_blocks: int,
755 + num_kv_heads: int,
756 + sm_scale: float,
757 + decode_query_len: int,
758 +) -> torch.Tensor:
759 + """Decode index block-score + top-k, both split-K (cudagraph-safe).
760 +
761 + Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad).
762 + """
763 + total_q, num_idx_heads, head_dim = idx_q.shape
764 + assert num_idx_heads == num_kv_heads, (
765 + "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)"
766 + )
767 + assert total_q == seq_lens.shape[0] * decode_query_len
768 + batch = total_q
769 + max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
770 + use_pdl = current_platform.is_arch_support_pdl()
771 + # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA
772 + # SM9+); this ROCm Triton rejects it even when False ("Keyword argument
773 + # launch_pdl was specified but unrecognised"). Only pass it when PDL is
774 + # actually supported -- on ROCm use_pdl is always False, so it's omitted.
775 + pdl_launch = {"launch_pdl": True} if use_pdl else {}
776 +
777 + # Keep score strides 16-divisible to avoid Triton recompiles.
778 + score_block_stride = round_up(max_block, 16)
779 + score = torch.empty(
780 + (num_idx_heads, total_q, score_block_stride),
781 + dtype=torch.float32,
782 + device=idx_q.device,
783 + )
784 + # split-K over seq blocks; chunk count depends only on shape constants so
785 + # the grid is fixed within a cuda graph.
786 + TARGET_GRID = 4096
787 + MAX_NUM_KV_CHUNKS = 256
788 + target = max(
789 + 1, min(MAX_NUM_KV_CHUNKS, TARGET_GRID // max(1, batch * num_idx_heads))
790 + )
791 + num_kv_chunks = 1 << (target.bit_length() - 1)
792 + grid_score = (batch, num_kv_chunks)
793 + _decode_index_score_kernel[grid_score](
794 + idx_q,
795 + index_kv_cache,
796 + score,
797 + block_table,
798 + seq_lens,
799 + num_idx_heads,
800 + head_dim,
801 + init_blocks,
802 + local_blocks,
803 + sm_scale,
804 + decode_query_len,
805 + idx_q.stride(0),
806 + idx_q.stride(1),
807 + idx_q.stride(2),
808 + index_kv_cache.stride(0),
809 + index_kv_cache.stride(1),
810 + index_kv_cache.stride(2),
811 + score.stride(0),
812 + score.stride(1),
813 + score.stride(2),
814 + block_table.stride(0),
815 + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
816 + num_kv_chunks=num_kv_chunks,
817 + USE_PDL=use_pdl,
818 + **pdl_launch,
819 + )
820 +
821 + topk_idx = torch.empty(
822 + (num_idx_heads, total_q, topk),
823 + dtype=torch.int32,
824 + device=idx_q.device,
825 + )
826 + # Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts
827 + # pow2(num_topk_chunks * pow2(topk)) candidates.
828 + TOPK_TARGET_GRID = 64
829 + MAX_NUM_TOPK_CHUNKS = 16
830 + topk_target = max(
831 + 1, min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, batch * num_idx_heads))
832 + )
833 + num_topk_chunks = 1 << (topk_target.bit_length() - 1)
834 + block_size_t = triton.next_power_of_2(topk)
835 + chunk_blocks = (max_block + num_topk_chunks - 1) // num_topk_chunks
836 + topk_score_partial = torch.empty(
837 + num_topk_chunks,
838 + num_idx_heads,
839 + batch,
840 + block_size_t,
841 + dtype=torch.float32,
842 + device=idx_q.device,
843 + )
844 + topk_idx_partial = torch.empty(
845 + num_topk_chunks,
846 + num_idx_heads,
847 + batch,
848 + block_size_t,
849 + dtype=torch.int32,
850 + device=idx_q.device,
851 + )
852 + _topk_index_partial_kernel[(batch, num_idx_heads, num_topk_chunks)](
853 + score,
854 + topk_score_partial,
855 + topk_idx_partial,
856 + seq_lens,
857 + SPARSE_BLOCK_SIZE,
858 + topk,
859 + chunk_blocks,
860 + decode_query_len,
861 + score.stride(0),
862 + score.stride(1),
863 + score.stride(2),
864 + topk_score_partial.stride(0),
865 + topk_score_partial.stride(1),
866 + topk_score_partial.stride(2),
867 + topk_score_partial.stride(3),
868 + topk_idx_partial.stride(0),
869 + topk_idx_partial.stride(1),
870 + topk_idx_partial.stride(2),
871 + topk_idx_partial.stride(3),
872 + USE_PDL=use_pdl,
873 + **pdl_launch,
874 + )
875 + _topk_index_merge_kernel[(batch, num_idx_heads)](
876 + topk_score_partial,
877 + topk_idx_partial,
878 + topk_idx,
879 + seq_lens,
880 + SPARSE_BLOCK_SIZE,
881 + topk,
882 + decode_query_len,
883 + topk_score_partial.stride(0),
884 + topk_score_partial.stride(1),
885 + topk_score_partial.stride(2),
886 + topk_score_partial.stride(3),
887 + topk_idx_partial.stride(0),
888 + topk_idx_partial.stride(1),
889 + topk_idx_partial.stride(2),
890 + topk_idx_partial.stride(3),
891 + topk_idx.stride(0),
892 + topk_idx.stride(1),
893 + topk_idx.stride(2),
894 + num_topk_chunks=num_topk_chunks,
895 + USE_PDL=use_pdl,
896 + **pdl_launch,
897 + )
898 + return topk_idx
vllm/models/minimax_m3/common/ops/sparse_attn.py new
+593
@@ -0,0 +1,593 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Triton kernels for MiniMax M3 block-sparse GQA attention.
4 +
5 +The main heads attend only to the blocks selected by the lightning indexer (see
6 +``index_topk``). Adapted to vLLM's paged KV cache: the KV page size is forced to
7 +equal the sparse block size (128), so one selected block maps to exactly one
8 +page.
9 +
10 +Main K/V cache layout (vLLM):
11 + ``(num_blocks, 2, 128, num_kv_heads, head_dim)`` K=[:,0] V=[:,1]
12 +
13 +Only the paths MiniMax M3 uses are implemented: no attention sink, base-2
14 +(exp2/log2) softmax. The decode kernels use split-K (flash-decoding) over the
15 +selected blocks with a separate merge step, since one query token per request
16 +leaves the prefill kernels (which parallelize over the query dim) idle.
17 +"""
18 +
19 +import torch
20 +
21 +from vllm.platforms import current_platform
22 +from vllm.triton_utils import tl, triton
23 +
24 +# One sparse block == one KV page.
25 +SPARSE_BLOCK_SIZE = 128
26 +
27 +
28 +_SPARSE_ATTN_NUM_STAGES_KWARG: dict | None = None
29 +
30 +
31 +def _sparse_attn_num_stages_kwarg() -> dict:
32 + """Triton ``num_stages`` override for the sparse-attn GEMM kernels.
33 +
34 + Forced only where required: CDNA3 (gfx942) caps LDS at
35 + 64 KB, and the default 2-stage pipeline double-buffers the 128x128 K/V tiles
36 + to ~66 KB ("out of resource: shared memory"), so pin gfx942 to a single
37 + stage (~32 KB, which fits). Everywhere else (NVIDIA, CDNA4 gfx950) return an
38 + empty kwarg and let Triton keep its own default -- don't second-guess it.
39 + Cached: the arch is fixed per process.
40 + """
41 + global _SPARSE_ATTN_NUM_STAGES_KWARG
42 + if _SPARSE_ATTN_NUM_STAGES_KWARG is None:
43 + kwarg: dict = {}
44 + if current_platform.is_rocm():
45 + from vllm.platforms.rocm import on_gfx942
46 +
47 + if on_gfx942():
48 + kwarg = {"num_stages": 1}
49 + _SPARSE_ATTN_NUM_STAGES_KWARG = kwarg
50 + return _SPARSE_ATTN_NUM_STAGES_KWARG
51 +
52 +
53 +# ---------------------------------------------------------------------------
54 +# GQA block-sparse attention (paged). Main heads attend only to the selected
55 +# blocks. BLOCK_SIZE_K == 128 so each selected block is one page.
56 +# ---------------------------------------------------------------------------
57 +# since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens
58 +# might lose pointer alignment, which trigger Triton recompiles. we don't actually
59 +# need pointer alignment for those tensors anyway because we do scalar load.
60 +@triton.heuristics(
61 + {
62 + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
63 + "BLOCK_SIZE_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]),
64 + "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
65 + "BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"]
66 + * triton.next_power_of_2(args["gqa_group_size"]),
67 + }
68 +)
69 +@triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"])
70 +def _gqa_sparse_fwd_kernel(
71 + q_ptr, # [total_q, num_heads, head_dim]
72 + kv_cache_ptr, # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim]
73 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk]
74 + o_ptr, # [total_q, num_heads, head_dim]
75 + block_table_ptr, # [num_reqs, max_blocks]
76 + cu_seqlens_q,
77 + cu_seqblocks_q,
78 + seq_lens,
79 + prefix_lens,
80 + num_kv_heads,
81 + gqa_group_size,
82 + head_dim,
83 + max_topk,
84 + num_q_loop,
85 + sm_scale,
86 + stride_qn,
87 + stride_qh,
88 + stride_qd,
89 + stride_kv_blk,
90 + stride_kv_kv,
91 + stride_kv_pos,
92 + stride_kv_h,
93 + stride_kv_d,
94 + stride_th,
95 + stride_tn,
96 + stride_tk,
97 + stride_on,
98 + stride_oh,
99 + stride_od,
100 + stride_bt_b,
101 + BLOCK_SIZE_Q: tl.constexpr,
102 + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
103 + BLOCK_SIZE_D: tl.constexpr,
104 + BLOCK_SIZE_H: tl.constexpr,
105 + BLOCK_SIZE_T: tl.constexpr,
106 + BLOCK_SIZE_QH: tl.constexpr,
107 + USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load
108 +):
109 + sm_scale_log2e = sm_scale * 1.4426950409
110 + pid_q = tl.program_id(0)
111 + pid_kh = tl.program_id(1)
112 + pid_b = tl.program_id(2)
113 + pid_h = pid_kh * gqa_group_size
114 + q_start = tl.load(cu_seqlens_q + pid_b)
115 + q_len = tl.load(cu_seqlens_q + pid_b + 1) - q_start
116 + q_block_start = tl.load(cu_seqblocks_q + pid_b)
117 + q_block_len = tl.load(cu_seqblocks_q + pid_b + 1) - q_block_start
118 + seq_len = tl.load(seq_lens + pid_b)
119 + prefix_len = tl.load(prefix_lens + pid_b)
120 + if pid_q * num_q_loop >= q_block_len:
121 + return
122 + real_q_loop = min(num_q_loop, q_block_len - pid_q * num_q_loop)
123 + bt_row = block_table_ptr + pid_b * stride_bt_b
124 + off_n = tl.arange(0, BLOCK_SIZE_K)
125 + off_d = tl.arange(0, BLOCK_SIZE_D)
126 + d_mask = off_d < head_dim
127 + for j in range(real_q_loop):
128 + pid_q_j = pid_q * num_q_loop + j
129 + t_ptr_j = t_ptr + (q_block_start + pid_q_j) * stride_tn + pid_kh * stride_th
130 + off_t = tl.arange(0, BLOCK_SIZE_T)
131 + topk_idx = tl.load(t_ptr_j + off_t * stride_tk, mask=off_t < max_topk, other=-1)
132 + real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0)
133 + q_ptrs = tl.make_block_ptr(
134 + base=q_ptr + q_start * stride_qn + pid_h * stride_qh,
135 + shape=(q_len, gqa_group_size, head_dim),
136 + strides=(stride_qn, stride_qh, stride_qd),
137 + offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
138 + block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
139 + order=(2, 1, 0),
140 + )
141 + q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero")
142 + off_q = (
143 + tl.arange(0, BLOCK_SIZE_Q)[:, None]
144 + + pid_q_j * BLOCK_SIZE_Q
145 + + prefix_len
146 + - tl.arange(0, BLOCK_SIZE_K)[None, :]
147 + )
148 + m_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
149 + lse_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
150 + acc_o = tl.zeros((BLOCK_SIZE_QH, BLOCK_SIZE_D), dtype=tl.float32)
151 + q = tl.reshape(q, BLOCK_SIZE_QH, BLOCK_SIZE_D)
152 + for _ in range(real_topk):
153 + blk = tl.load(t_ptr_j).to(tl.int32)
154 + t_ptr_j = t_ptr_j + stride_tk
155 + c = blk * BLOCK_SIZE_K
156 + page = tl.load(bt_row + blk).to(tl.int64)
157 + pos = c + off_n
158 + pos_mask = pos < seq_len
159 + k = tl.load(
160 + kv_cache_ptr
161 + + page * stride_kv_blk
162 + + 0 * stride_kv_kv
163 + + off_n[None, :] * stride_kv_pos
164 + + pid_kh * stride_kv_h
165 + + off_d[:, None] * stride_kv_d,
166 + mask=d_mask[:, None] & pos_mask[None, :],
167 + other=0.0,
168 + )
169 + if USE_FP8:
170 + k = k.to(q.dtype)
171 + qk = tl.zeros((BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_K), dtype=tl.float32)
172 + # causal: q_abs_pos - k_off >= block_start (c)
173 + qk += tl.where(off_q[:, None, :] >= c, 0, float("-inf"))
174 + qk = tl.reshape(qk, BLOCK_SIZE_QH, BLOCK_SIZE_K)
175 + qk += tl.dot(q, k) * sm_scale_log2e
176 + qk += tl.where(pos_mask[None, :], 0, float("-inf"))
177 + m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
178 + p = tl.exp2(qk - m_ij[:, None])
179 + l_ij = tl.sum(p, axis=1)
180 + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None]
181 + v = tl.load(
182 + kv_cache_ptr
183 + + page * stride_kv_blk
184 + + 1 * stride_kv_kv
185 + + off_n[:, None] * stride_kv_pos
186 + + pid_kh * stride_kv_h
187 + + off_d[None, :] * stride_kv_d,
188 + mask=pos_mask[:, None] & d_mask[None, :],
189 + other=0.0,
190 + )
191 + if USE_FP8:
192 + v = v.to(q.dtype)
193 + acc_o += tl.dot(p.to(v.dtype), v)
194 + m_i = m_ij
195 + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
196 + acc_o = acc_o * tl.exp2(m_i - lse_i)[:, None]
197 + acc_o = tl.reshape(acc_o, BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D)
198 + o_ptrs = tl.make_block_ptr(
199 + base=o_ptr + q_start * stride_on + pid_h * stride_oh,
200 + shape=(q_len, gqa_group_size, head_dim),
201 + strides=(stride_on, stride_oh, stride_od),
202 + offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
203 + block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
204 + order=(2, 1, 0),
205 + )
206 + tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2))
207 +
208 +
209 +# ---------------------------------------------------------------------------
210 +# Decode kernels (split-K). Decode batches are flattened request-major, with a
211 +# runtime query length used to map each query token back to its request metadata.
212 +# This parallelizes over the selected top-k blocks, producing partials that the
213 +# merge kernel combines (flash-decoding). All chunk counts depend only on shape
214 +# constants so the grid is fixed within a cuda graph. Base-2 (exp2/log2)
215 +# softmax matches the prefill kernel.
216 +# ---------------------------------------------------------------------------
217 +@triton.heuristics(
218 + {
219 + "BLOCK_SIZE_H": lambda args: max(
220 + 16, triton.next_power_of_2(args["gqa_group_size"])
221 + ),
222 + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
223 + "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
224 + }
225 +)
226 +@triton.jit(do_not_specialize=["decode_query_len"])
227 +def _gqa_sparse_decode_kernel(
228 + q_ptr, # [total_q, num_heads, head_dim]
229 + kv_cache_ptr, # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim]
230 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk]
231 + o_ptr, # partial out: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim]
232 + lse_ptr, # partial lse (log2): [NUM_TOPK_CHUNKS, total_q, num_heads]
233 + block_table_ptr, # [num_reqs, max_blocks]
234 + seq_lens, # [num_reqs]
235 + total_q,
236 + gqa_group_size,
237 + head_dim,
238 + max_topk,
239 + sm_scale,
240 + decode_query_len,
241 + stride_qn,
242 + stride_qh,
243 + stride_qd,
244 + stride_kv_blk,
245 + stride_kv_kv,
246 + stride_kv_pos,
247 + stride_kv_h,
248 + stride_kv_d,
249 + stride_th,
250 + stride_tn,
251 + stride_tk,
252 + stride_o_c,
253 + stride_o_b,
254 + stride_o_h,
255 + stride_o_d,
256 + stride_l_c,
257 + stride_l_b,
258 + stride_l_h,
259 + stride_bt_b,
260 + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
261 + NUM_TOPK_CHUNKS: tl.constexpr,
262 + BLOCK_SIZE_H: tl.constexpr,
263 + BLOCK_SIZE_D: tl.constexpr,
264 + BLOCK_SIZE_T: tl.constexpr,
265 + USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load
266 + USE_PDL: tl.constexpr,
267 +):
268 + sm_scale_log2e = sm_scale * 1.4426950409
269 + # split-K over the topk dimension: pid(0) folds (query-token, chunk).
270 + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1)
271 + pid_b = pid_bc % total_q
272 + pid_c = pid_bc // total_q
273 + req_id = pid_b // decode_query_len
274 + q_offset = pid_b - req_id * decode_query_len
275 + pid_h = pid_kh * gqa_group_size
276 + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS
277 + chunk_start_topk = pid_c * chunk_size_topk
278 + chunk_end_compiletime = chunk_start_topk + chunk_size_topk
279 +
280 + if USE_PDL:
281 + tl.extra.cuda.gdc_wait()
282 +
283 + seq_len = tl.load(seq_lens + req_id)
284 + query_pos = seq_len - decode_query_len + q_offset
285 + # Full-CG padding uses zero-length request rows. Clamp to an empty
286 + # attention range instead of letting padded rows produce negative lengths.
287 + kv_len = tl.maximum(query_pos + 1, 0)
288 +
289 + # number of valid (non-padded) selected blocks for this query token
290 + off_t = tl.arange(0, BLOCK_SIZE_T)
291 + idx_base = t_ptr + pid_kh * stride_th + pid_b * stride_tn
292 + topk_idx = tl.load(idx_base + off_t * stride_tk, mask=off_t < max_topk, other=-1)
293 + real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0)
294 + chunk_end_topk = tl.minimum(chunk_end_compiletime, real_topk)
295 +
296 + off_n = tl.arange(0, BLOCK_SIZE_K)
297 + off_d = tl.arange(0, BLOCK_SIZE_D)
298 + d_mask = off_d < head_dim
299 + bt_row = block_table_ptr + req_id * stride_bt_b
300 +
301 + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32)
302 + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32)
303 + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32)
304 + q_ptrs = tl.make_block_ptr(
305 + base=q_ptr + pid_b * stride_qn + pid_h * stride_qh,
306 + shape=(gqa_group_size, head_dim),
307 + strides=(stride_qh, stride_qd),
308 + offsets=(0, 0),
309 + block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D),
310 + order=(1, 0),
311 + )
312 + q = tl.load(q_ptrs, boundary_check=(0, 1), padding_option="zero")
313 +
314 + cur_idx_ptr = idx_base + chunk_start_topk * stride_tk
315 + for _ in tl.range(chunk_start_topk, chunk_end_topk):
316 + blk = tl.load(cur_idx_ptr).to(tl.int32)
317 + cur_idx_ptr = cur_idx_ptr + stride_tk
318 + c = blk * BLOCK_SIZE_K
319 + page = tl.load(bt_row + blk).to(tl.int64)
320 + pos = c + off_n
321 + pos_mask = pos < kv_len
322 + k = tl.load(
323 + kv_cache_ptr
324 + + page * stride_kv_blk
325 + + 0 * stride_kv_kv
326 + + off_n[None, :] * stride_kv_pos
327 + + pid_kh * stride_kv_h
328 + + off_d[:, None] * stride_kv_d,
329 + mask=d_mask[:, None] & pos_mask[None, :],
330 + other=0.0,
331 + )
332 + if USE_FP8:
333 + k = k.to(q.dtype)
334 + qk = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_K), dtype=tl.float32)
335 + qk += tl.where(pos_mask[None, :], 0, float("-inf"))
336 + qk += tl.dot(q, k) * sm_scale_log2e
337 + m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
338 + p = tl.exp2(qk - m_ij[:, None])
339 + l_ij = tl.sum(p, axis=1)
340 + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None]
341 + v = tl.load(
342 + kv_cache_ptr
343 + + page * stride_kv_blk
344 + + 1 * stride_kv_kv
345 + + off_n[:, None] * stride_kv_pos
346 + + pid_kh * stride_kv_h
347 + + off_d[None, :] * stride_kv_d,
348 + mask=pos_mask[:, None] & d_mask[None, :],
349 + other=0.0,
350 + )
351 + if USE_FP8:
352 + v = v.to(q.dtype)
353 + acc_o += tl.dot(p.to(v.dtype), v)
354 + m_i = m_ij
355 + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
356 +
357 + if USE_PDL:
358 + tl.extra.cuda.gdc_launch_dependents()
359 +
360 + # Empty chunks for active rows must store zero output; otherwise the merge
361 + # can hit 0 * NaN. All-empty padded rows may still produce NaNs in merge.
362 + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), tl.zeros_like(lse_i))
363 + acc_o = acc_o * scale[:, None]
364 + o_ptrs = tl.make_block_ptr(
365 + base=o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h,
366 + shape=(gqa_group_size, head_dim),
367 + strides=(stride_o_h, stride_o_d),
368 + offsets=(0, 0),
369 + block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D),
370 + order=(1, 0),
371 + )
372 + tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1))
373 + lse_ptrs = tl.make_block_ptr(
374 + base=lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h,
375 + shape=(gqa_group_size,),
376 + strides=(stride_l_h,),
377 + offsets=(0,),
378 + block_shape=(BLOCK_SIZE_H,),
379 + order=(0,),
380 + )
381 + tl.store(lse_ptrs, lse_i.to(lse_ptr.dtype.element_ty), boundary_check=(0,))
382 +
383 +
384 +@triton.heuristics(
385 + {"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}
386 +)
387 +@triton.jit
388 +def _merge_topk_attn_out_kernel(
389 + o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim]
390 + lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads]
391 + out_ptr, # merged out: [total_q, num_heads, head_dim]
392 + head_dim,
393 + stride_o_c,
394 + stride_o_b,
395 + stride_o_h,
396 + stride_o_d,
397 + stride_l_c,
398 + stride_l_b,
399 + stride_l_h,
400 + stride_out_n,
401 + stride_out_h,
402 + stride_out_d,
403 + NUM_TOPK_CHUNKS: tl.constexpr,
404 + BLOCK_SIZE_D: tl.constexpr,
405 + USE_PDL: tl.constexpr,
406 +):
407 + pid_b, pid_h = tl.program_id(0), tl.program_id(1)
408 +
409 + # NOTE: assume seq_lens is safe to load before gdc_wait()
410 + if USE_PDL:
411 + tl.extra.cuda.gdc_wait()
412 + tl.extra.cuda.gdc_launch_dependents()
413 +
414 + off_c = tl.arange(0, NUM_TOPK_CHUNKS)
415 + off_d = tl.arange(0, BLOCK_SIZE_D)
416 + o_ptrs = tl.make_block_ptr(
417 + base=o_ptr + pid_b * stride_o_b + pid_h * stride_o_h,
418 + shape=(NUM_TOPK_CHUNKS, head_dim),
419 + strides=(stride_o_c, stride_o_d),
420 + offsets=(0, 0),
421 + block_shape=(NUM_TOPK_CHUNKS, BLOCK_SIZE_D),
422 + order=(1, 0),
423 + )
424 + lse_ptrs = lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c
425 + o = tl.load(o_ptrs, boundary_check=(0, 1), padding_option="zero")
426 + lse = tl.load(lse_ptrs) # empty chunks contribute -inf -> weight 0
427 + lse_max = tl.max(lse, axis=0)
428 + weights = tl.exp2(lse - lse_max)
429 + weights = weights / tl.sum(weights, axis=0)
430 + o_merged = tl.sum(o * weights[:, None], axis=0)
431 + out_ptrs = (
432 + out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d
433 + )
434 + tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim)
435 +
436 +
437 +# ---------------------------------------------------------------------------
438 +# Python wrappers
439 +# ---------------------------------------------------------------------------
440 +@torch.no_grad()
441 +def minimax_m3_sparse_attn(
442 + q: torch.Tensor, # [total_q, num_heads, head_dim]
443 + kv_cache: torch.Tensor, # [num_blocks, 2, 128, num_kv_heads, head_dim]
444 + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk]
445 + block_table: torch.Tensor, # [batch, max_blocks]
446 + cu_seqlens_q: torch.Tensor, # [batch+1] int32
447 + seq_lens: torch.Tensor, # [batch] int32
448 + prefix_lens: torch.Tensor, # [batch] int32
449 + max_query_len: int,
450 + num_kv_heads: int,
451 + sm_scale: float,
452 + output: torch.Tensor, # [total_q, num_heads, head_dim]
453 +) -> None:
454 + """GQA block-sparse attention over the selected blocks. block_size_q == 1."""
455 + total_q, num_heads, head_dim = q.shape
456 + batch = cu_seqlens_q.shape[0] - 1
457 + topk = topk_idx.shape[-1]
458 + gqa_group_size = num_heads // num_kv_heads
459 + use_fp8 = kv_cache.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
460 + grid = (max_query_len, num_kv_heads, batch)
461 + _gqa_sparse_fwd_kernel[grid](
462 + q,
463 + kv_cache,
464 + topk_idx,
465 + output,
466 + block_table,
467 + cu_seqlens_q,
468 + cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1
469 + seq_lens,
470 + prefix_lens,
471 + num_kv_heads,
472 + gqa_group_size,
473 + head_dim,
474 + topk,
475 + 1, # num_q_loop
476 + sm_scale,
477 + q.stride(0),
478 + q.stride(1),
479 + q.stride(2),
480 + kv_cache.stride(0),
481 + kv_cache.stride(1),
482 + kv_cache.stride(2),
483 + kv_cache.stride(3),
484 + kv_cache.stride(4),
485 + topk_idx.stride(0),
486 + topk_idx.stride(1),
487 + topk_idx.stride(2),
488 + output.stride(0),
489 + output.stride(1),
490 + output.stride(2),
491 + block_table.stride(0),
492 + BLOCK_SIZE_Q=1,
493 + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
494 + USE_FP8=use_fp8,
495 + **_sparse_attn_num_stages_kwarg(),
496 + )
497 +
498 +
499 +@torch.no_grad()
500 +def minimax_m3_sparse_attn_decode(
501 + q: torch.Tensor, # [total_q, num_heads, head_dim]
502 + kv_cache: torch.Tensor, # [num_blocks, 2, 128, num_kv_heads, head_dim]
503 + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk]
504 + block_table: torch.Tensor, # [num_reqs, max_blocks]
505 + seq_lens: torch.Tensor, # [num_reqs] int32
506 + num_kv_heads: int,
507 + sm_scale: float,
508 + output: torch.Tensor, # [total_q, num_heads, head_dim]
509 + decode_query_len: int,
510 +) -> None:
511 + """GQA block-sparse attention for decode (split-K over the top-k blocks)."""
512 + total_q, num_heads, head_dim = q.shape
513 + assert total_q == seq_lens.shape[0] * decode_query_len
514 + max_topk = topk_idx.shape[-1]
515 + gqa_group_size = num_heads // num_kv_heads
516 + use_fp8 = kv_cache.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
517 + use_pdl = current_platform.is_arch_support_pdl()
518 + # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA
519 + # SM9+); this ROCm Triton rejects it even when False ("Keyword argument
520 + # launch_pdl was specified but unrecognised"). Only pass it when PDL is
521 + # actually supported -- on ROCm use_pdl is always False, so it's omitted.
522 + pdl_launch = {"launch_pdl": True} if use_pdl else {}
523 + # split-K over the selected blocks; chunk count is shape-constant (cuda graph).
524 + TARGET_GRID = 256
525 + target = max(1, min(max_topk, TARGET_GRID // max(1, total_q * num_kv_heads)))
526 + num_topk_chunks = 1 << (target.bit_length() - 1)
527 + o_partial = torch.empty(
528 + num_topk_chunks, total_q, num_heads, head_dim, dtype=q.dtype, device=q.device
529 + )
530 + lse_partial = torch.empty(
531 + num_topk_chunks, total_q, num_heads, dtype=torch.float32, device=q.device
532 + )
533 + grid = (total_q * num_topk_chunks, num_kv_heads)
534 + _gqa_sparse_decode_kernel[grid](
535 + q,
536 + kv_cache,
537 + topk_idx,
538 + o_partial,
539 + lse_partial,
540 + block_table,
541 + seq_lens,
542 + total_q,
543 + gqa_group_size,
544 + head_dim,
545 + max_topk,
546 + sm_scale,
547 + decode_query_len,
548 + q.stride(0),
549 + q.stride(1),
550 + q.stride(2),
551 + kv_cache.stride(0),
552 + kv_cache.stride(1),
553 + kv_cache.stride(2),
554 + kv_cache.stride(3),
555 + kv_cache.stride(4),
556 + topk_idx.stride(0),
557 + topk_idx.stride(1),
558 + topk_idx.stride(2),
559 + o_partial.stride(0),
560 + o_partial.stride(1),
561 + o_partial.stride(2),
562 + o_partial.stride(3),
563 + lse_partial.stride(0),
564 + lse_partial.stride(1),
565 + lse_partial.stride(2),
566 + block_table.stride(0),
567 + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
568 + NUM_TOPK_CHUNKS=num_topk_chunks,
569 + USE_FP8=use_fp8,
570 + USE_PDL=use_pdl,
571 + **_sparse_attn_num_stages_kwarg(),
572 + **pdl_launch,
573 + )
574 + merge_grid = (total_q, num_heads)
575 + _merge_topk_attn_out_kernel[merge_grid](
576 + o_partial,
577 + lse_partial,
578 + output,
579 + head_dim,
580 + o_partial.stride(0),
581 + o_partial.stride(1),
582 + o_partial.stride(2),
583 + o_partial.stride(3),
584 + lse_partial.stride(0),
585 + lse_partial.stride(1),
586 + lse_partial.stride(2),
587 + output.stride(0),
588 + output.stride(1),
589 + output.stride(2),
590 + NUM_TOPK_CHUNKS=num_topk_chunks,
591 + USE_PDL=use_pdl,
592 + **pdl_launch,
593 + )
vllm/models/minimax_m3/common/sparse_attention.py new
+398
@@ -0,0 +1,398 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Main block-sparse GQA attention for MiniMax M3 sparse layers.
4 +
5 +The lightning indexer (``indexer.py``) selects the top-k KV blocks; this module
6 +holds the main attention that attends only to those blocks: the paged K/V cache
7 +backend, its metadata + builder, and the impl that consumes the indexer's
8 +``topk_idx``. The Triton attend kernel lives here; the SM100 (MSA)
9 +``build_k2q_csr`` + ``sparse_atten_func`` attend lives in
10 +``nvidia/sparse_attention_msa.py``.
11 +
12 +``MiniMaxM3SparseBackend`` and ``MiniMaxM3SparseMetadata`` are referenced by the
13 +attention-backend registry (by dotted path) and by spec-decode, so they must
14 +keep these names and stay in this module.
15 +"""
16 +
17 +from dataclasses import dataclass
18 +from typing import ClassVar
19 +
20 +import torch
21 +
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.models.minimax_m3.common.ops.sparse_attn import (
26 + SPARSE_BLOCK_SIZE,
27 + minimax_m3_sparse_attn,
28 + minimax_m3_sparse_attn_decode,
29 +)
30 +from vllm.platforms import current_platform
31 +from vllm.v1.attention.backend import (
32 + AttentionBackend,
33 + AttentionCGSupport,
34 + AttentionImplBase,
35 + AttentionLayer,
36 + AttentionMetadata,
37 + AttentionMetadataBuilder,
38 + CommonAttentionMetadata,
39 + MultipleOf,
40 +)
41 +from vllm.v1.attention.backends.utils import (
42 + get_kv_cache_layout,
43 + split_decodes_and_prefills,
44 +)
45 +from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
46 +
47 +
48 +class MiniMaxM3SparseBackend(AttentionBackend):
49 + """Block-sparse GQA backend for MiniMax M3 sparse attention layers."""
50 +
51 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16, torch.float16]
52 + # bf16 or fp8 (e4m3/e5m2): the Triton kernels dequant fp8 before the dots.
53 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
54 + "bfloat16",
55 + "fp8",
56 + "fp8_e4m3",
57 + "fp8_e5m2",
58 + ]
59 +
60 + @staticmethod
61 + def get_name() -> str:
62 + return "MINIMAX_M3_SPARSE"
63 +
64 + @staticmethod
65 + def get_impl_cls() -> type["MiniMaxM3SparseImpl"]:
66 + # Concrete impl chosen by select_main_impl_cls; base for introspection.
67 + return MiniMaxM3SparseImpl
68 +
69 + @staticmethod
70 + def get_builder_cls() -> type["MiniMaxM3SparseMetadataBuilder"]:
71 + return MiniMaxM3SparseMetadataBuilder
72 +
73 + @classmethod
74 + def get_supported_head_sizes(cls) -> list[int]:
75 + return [128]
76 +
77 + @staticmethod
78 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
79 + # Page size == sparse block size (one sparse block per KV page).
80 + return [128]
81 +
82 + @classmethod
83 + def is_sparse(cls) -> bool:
84 + return True
85 +
86 + @staticmethod
87 + def get_kv_cache_shape(
88 + num_blocks: int,
89 + block_size: int,
90 + num_kv_heads: int,
91 + head_size: int,
92 + cache_dtype_str: str = "auto",
93 + ) -> tuple[int, ...]:
94 + return (num_blocks, 2, block_size, num_kv_heads, head_size)
95 +
96 + @staticmethod
97 + def get_kv_cache_stride_order(
98 + include_num_layers_dimension: bool = False,
99 + ) -> tuple[int, ...]:
100 + # Permutation from get_kv_cache_shape to the actual memory layout.
101 + if include_num_layers_dimension:
102 + raise NotImplementedError # no cross-layer KV blocks in M3
103 + cache_layout = get_kv_cache_layout()
104 + if cache_layout == "NHD":
105 + stride_order = (0, 1, 2, 3, 4)
106 + elif cache_layout == "HND":
107 + stride_order = (0, 1, 3, 2, 4)
108 + else:
109 + raise ValueError(f"Unknown cache layout format {cache_layout}.")
110 + return stride_order
111 +
112 +
113 +@dataclass
114 +class MiniMaxM3SparsePrefillMetadata:
115 + """Per-prefill state; ``cu_seqlens_k``/``total_kv_blocks`` feed the MSA CSR."""
116 +
117 + cu_seqlens_q: torch.Tensor # [num_prefills + 1] int32, rebased to 0
118 + cu_seqlens_k: torch.Tensor # [num_prefills + 1] int32, cumulative KV lengths
119 + seq_lens: torch.Tensor # [num_prefills] int32, total KV lengths
120 + context_lens: torch.Tensor # [num_prefills] int32 (cached/context tokens)
121 + block_table: torch.Tensor
122 + max_query_len: int
123 + max_seq_len: int
124 + total_kv_blocks: int
125 +
126 +
127 +@dataclass
128 +class MiniMaxM3SparseDecodeMetadata:
129 + """Per-decode state (cudagraph-safe). ``decode_query_len`` is the uniform
130 + per-request query length (1, or 1 + num_speculative_tokens)."""
131 +
132 + seq_lens: torch.Tensor # [num_decodes] int32
133 + block_table: torch.Tensor
134 + decode_query_len: int
135 +
136 +
137 +@dataclass
138 +class MiniMaxM3SparseMetadata(AttentionMetadata):
139 + """Sparse-attention metadata, split into prefill and decode sub-metadata."""
140 +
141 + seq_lens: torch.Tensor
142 + max_seq_len: int
143 + slot_mapping: torch.Tensor
144 +
145 + num_actual_tokens: int # total query tokens (decode-first batch)
146 +
147 + # Split counts (batch reordered decode-first).
148 + num_decodes: int
149 + num_decode_tokens: int
150 + num_prefills: int
151 + num_prefill_tokens: int
152 +
153 + prefill: MiniMaxM3SparsePrefillMetadata | None = None
154 + decode: MiniMaxM3SparseDecodeMetadata | None = None
155 +
156 +
157 +class MiniMaxM3SparseMetadataBuilder(AttentionMetadataBuilder[MiniMaxM3SparseMetadata]):
158 + # Full cudagraphs for uniform decode batches, incl. spec-decode verify
159 + # batches with >1 query token/request.
160 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
161 + # Raised to 1 + num_speculative_tokens by _init_reorder_batch_threshold when
162 + # spec decode is on; must match the indexer builder so the splits agree.
163 + reorder_batch_threshold: int = 1
164 +
165 + def __init__(
166 + self,
167 + kv_cache_spec: AttentionSpec,
168 + layer_names: list[str],
169 + vllm_config: VllmConfig,
170 + device: torch.device,
171 + ) -> None:
172 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
173 + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
174 + # Stable context-length buffer for decode cudagraph replays.
175 + self.context_len_buffer = torch.empty(
176 + vllm_config.scheduler_config.max_num_batched_tokens,
177 + dtype=torch.int32,
178 + device=device,
179 + )
180 +
181 + def build(
182 + self,
183 + common_prefix_len: int,
184 + common_attn_metadata: CommonAttentionMetadata,
185 + fast_build: bool = False,
186 + ) -> MiniMaxM3SparseMetadata:
187 + num_reqs = common_attn_metadata.num_reqs
188 + num_tokens = common_attn_metadata.num_actual_tokens
189 + query_start_loc = common_attn_metadata.query_start_loc
190 + seq_lens = common_attn_metadata.seq_lens
191 + block_table = common_attn_metadata.block_table_tensor
192 +
193 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
194 + split_decodes_and_prefills(
195 + common_attn_metadata,
196 + decode_threshold=self.reorder_batch_threshold,
197 + require_uniform=True,
198 + )
199 + )
200 + assert num_decodes + num_prefills == num_reqs
201 + assert num_decode_tokens + num_prefill_tokens == num_tokens
202 +
203 + # Decode-first batch: context lengths into the stable cudagraph buffer.
204 + context_lens = self.context_len_buffer[:num_reqs]
205 + context_lens.copy_(
206 + common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
207 + )
208 +
209 + prefill_metadata: MiniMaxM3SparsePrefillMetadata | None = None
210 + if num_prefills > 0:
211 + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
212 + assert seq_lens_cpu is not None
213 + prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:]
214 + prefill_total_kv_blocks = (
215 + ((prefill_seq_lens_cpu + SPARSE_BLOCK_SIZE - 1) // SPARSE_BLOCK_SIZE)
216 + .sum()
217 + .item()
218 + )
219 + prefill_kv_lens = seq_lens[num_decodes:]
220 + prefill_cu_seqlens_k = torch.empty(
221 + num_prefills + 1, dtype=torch.int32, device=seq_lens.device
222 + )
223 + prefill_cu_seqlens_k[0] = 0
224 + torch.cumsum(prefill_kv_lens, dim=0, out=prefill_cu_seqlens_k[1:])
225 + prefill_metadata = MiniMaxM3SparsePrefillMetadata(
226 + cu_seqlens_q=(query_start_loc[num_decodes:] - num_decode_tokens).to(
227 + torch.int32
228 + ),
229 + cu_seqlens_k=prefill_cu_seqlens_k,
230 + seq_lens=prefill_kv_lens,
231 + context_lens=context_lens[num_decodes:],
232 + block_table=block_table[num_decodes:],
233 + max_query_len=common_attn_metadata.max_query_len,
234 + max_seq_len=common_attn_metadata.max_seq_len,
235 + total_kv_blocks=prefill_total_kv_blocks,
236 + )
237 +
238 + decode_metadata: MiniMaxM3SparseDecodeMetadata | None = None
239 + if num_decodes > 0:
240 + qsl_cpu = common_attn_metadata.query_start_loc_cpu
241 + query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes]
242 + decode_query_len = int(query_lens_cpu[0].item())
243 + assert decode_query_len > 0
244 + assert torch.all(
245 + (query_lens_cpu == decode_query_len) | (query_lens_cpu == 0)
246 + )
247 + assert num_decode_tokens == num_decodes * decode_query_len
248 + decode_metadata = MiniMaxM3SparseDecodeMetadata(
249 + seq_lens=seq_lens[:num_decodes],
250 + block_table=block_table[:num_decodes],
251 + decode_query_len=decode_query_len,
252 + )
253 +
254 + return MiniMaxM3SparseMetadata(
255 + seq_lens=seq_lens,
256 + max_seq_len=common_attn_metadata.max_seq_len,
257 + slot_mapping=common_attn_metadata.slot_mapping,
258 + num_actual_tokens=num_tokens,
259 + num_decodes=num_decodes,
260 + num_decode_tokens=num_decode_tokens,
261 + num_prefills=num_prefills,
262 + num_prefill_tokens=num_prefill_tokens,
263 + prefill=prefill_metadata,
264 + decode=decode_metadata,
265 + )
266 +
267 +
268 +class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
269 + """Abstract base for block-sparse GQA over the indexer-selected blocks.
270 +
271 + Inherits ``AttentionImplBase`` for a custom forward signature (the layer
272 + pre-inserts K/V and runs the indexer, so forward takes the queries +
273 + ``topk_idx``). The Triton and MSA subclasses each own a full ``forward`` --
274 + no shared forward code.
275 + """
276 +
277 + def __init__(
278 + self,
279 + num_heads: int,
280 + head_size: int,
281 + scale: float,
282 + num_kv_heads: int | None = None,
283 + kv_cache_dtype: str = "auto",
284 + *,
285 + topk_blocks: int,
286 + sparse_block_size: int,
287 + ) -> None:
288 + self.num_heads = num_heads
289 + self.head_size = head_size
290 + self.scale = scale
291 + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
292 + self.kv_cache_dtype = kv_cache_dtype
293 + self.use_fp8_kv = is_quantized_kv_cache(kv_cache_dtype)
294 + self.kv_cache_fp8_dtype = (
295 + torch.float8_e5m2 if "e5m2" in kv_cache_dtype else torch.float8_e4m3fn
296 + )
297 + # Sparse selection parameters (block_size == page size == SPARSE_BLOCK_SIZE).
298 + self.topk_blocks = topk_blocks
299 + self.block_size = sparse_block_size
300 +
301 + def forward(
302 + self,
303 + layer: AttentionLayer,
304 + query: torch.Tensor,
305 + kv_cache: torch.Tensor,
306 + topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
307 + output: torch.Tensor,
308 + ) -> torch.Tensor:
309 + """Attend the queries to the indexer-selected blocks. Per kernel."""
310 + raise NotImplementedError
311 +
312 +
313 +class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
314 + """Triton block-sparse attend (``minimax_m3_sparse_attn``) + Triton decode."""
315 +
316 + def forward(
317 + self,
318 + layer: AttentionLayer,
319 + query: torch.Tensor,
320 + kv_cache: torch.Tensor,
321 + topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
322 + output: torch.Tensor,
323 + ) -> torch.Tensor:
324 + attn_metadata = get_forward_context().attn_metadata
325 + if not isinstance(attn_metadata, dict):
326 + return output # profiling run; caches unbound
327 + main_md = attn_metadata[layer.layer_name] # type: ignore[attr-defined]
328 + assert isinstance(main_md, MiniMaxM3SparseMetadata)
329 + decode_topk, prefill_topk = topk_idx
330 +
331 + nd = main_md.num_decode_tokens
332 + num_tokens = main_md.num_actual_tokens
333 + hd = self.head_size
334 + q = query[:num_tokens].view(-1, self.num_heads, hd)
335 + out = output[:num_tokens].view(-1, self.num_heads, hd)
336 + kv_cache = (
337 + kv_cache.view(self.kv_cache_fp8_dtype) if self.use_fp8_kv else kv_cache
338 + )
339 +
340 + # Decode [:nd]: split-K over the selected blocks (request-major chunks).
341 + if main_md.num_decodes > 0:
342 + d = main_md.decode
343 + assert d is not None and decode_topk is not None
344 + minimax_m3_sparse_attn_decode(
345 + q[:nd],
346 + kv_cache,
347 + decode_topk,
348 + d.block_table,
349 + d.seq_lens,
350 + self.num_kv_heads,
351 + self.scale,
352 + out[:nd],
353 + d.decode_query_len,
354 + )
355 +
356 + # Prefill [nd:]: cu_seqlens_q already rebased to 0.
357 + if main_md.num_prefills > 0:
358 + p = main_md.prefill
359 + assert p is not None and prefill_topk is not None
360 + minimax_m3_sparse_attn(
361 + q[nd:],
362 + kv_cache,
363 + prefill_topk,
364 + p.block_table,
365 + p.cu_seqlens_q,
366 + p.seq_lens,
367 + p.context_lens,
368 + p.max_query_len,
369 + self.num_kv_heads,
370 + self.scale,
371 + out[nd:],
372 + )
373 + return output
374 +
375 +
376 +def select_main_impl_cls(
377 + *,
378 + topk_blocks: int,
379 + kv_cache_dtype: str,
380 +) -> type[MiniMaxM3SparseImpl]:
381 + """Pick the main attend impl off the main KV-cache dtype.
382 +
383 + bf16 on Blackwell (SM100) uses the MSA attend; fp8 or non-Blackwell falls
384 + back to Triton. The MSA module is imported lazily so AMD/non-SM100 never
385 + import fmha_sm100.
386 + """
387 + if (
388 + current_platform.is_cuda()
389 + and current_platform.is_device_capability_family(100)
390 + and topk_blocks in (4, 8, 16, 32)
391 + and not is_quantized_kv_cache(kv_cache_dtype)
392 + ):
393 + from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
394 + MiniMaxM3SparseMSAImpl,
395 + )
396 +
397 + return MiniMaxM3SparseMSAImpl
398 + return MiniMaxM3SparseTritonImpl
vllm/models/minimax_m3/common/vision_tower.py new
+765
@@ -0,0 +1,765 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Iterable
4 +
5 +import numpy as np
6 +import torch
7 +import torch.nn as nn
8 +from einops import rearrange
9 +from transformers import PretrainedConfig
10 +
11 +from vllm.distributed import parallel_state
12 +from vllm.distributed import utils as dist_utils
13 +from vllm.model_executor.layers.activation import get_act_fn
14 +from vllm.model_executor.layers.attention.mm_encoder_attention import (
15 + MMEncoderAttention,
16 +)
17 +from vllm.model_executor.layers.linear import (
18 + ColumnParallelLinear,
19 + QKVParallelLinear,
20 + RowParallelLinear,
21 +)
22 +from vllm.model_executor.layers.quantization import QuantizationConfig
23 +from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb
24 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
25 +from vllm.model_executor.models.utils import maybe_prefix
26 +from vllm.model_executor.models.vision import (
27 + get_vit_attn_backend,
28 + is_vit_use_data_parallel,
29 +)
30 +from vllm.platforms import current_platform
31 +
32 +# ROCm caps a kernel-launch gridDim.y at 65536. The HIP flash-attn Triton
33 +# rotary kernel launches grid.y = cdiv(seqlen, BLOCK_M), so it fails with
34 +# hipErrorInvalidValue once cdiv(seqlen, BLOCK_M) > 65536. Used below to decide
35 +# when RoPE must be applied per video segment instead of in one launch.
36 +_HIP_MAX_GRID_DIM_Y = 65536
37 +
38 +
39 +class MiniMaxVLPatchEmbed(nn.Module):
40 + """Conv3d-based patch embedding.
41 +
42 + Takes flat tokens of shape (N, C * temporal_patch_size * patch_size²)
43 + and projects each to a hidden-size embedding.
44 + """
45 +
46 + def __init__(self, config: PretrainedConfig) -> None:
47 + super().__init__()
48 + compression = config.img_token_compression_config
49 + temporal_patch_size = compression.get("temporal_patch_size", 2)
50 + patch_size = config.patch_size
51 + num_channels = config.num_channels
52 +
53 + self.patch_size = patch_size
54 + self.temporal_patch_size = temporal_patch_size
55 + self.num_channels = num_channels
56 + self.hidden_size = config.hidden_size
57 +
58 + self.patch_embedding = nn.Conv3d(
59 + in_channels=num_channels,
60 + out_channels=config.hidden_size,
61 + kernel_size=(temporal_patch_size, patch_size, patch_size),
62 + stride=(temporal_patch_size, patch_size, patch_size),
63 + bias=False,
64 + )
65 +
66 + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
67 + # pixel_values: (N, C * temporal_patch_size * patch_size²)
68 + if self.patch_embedding.weight.dtype != pixel_values.dtype:
69 + self.patch_embedding = self.patch_embedding.to(pixel_values.dtype)
70 + x = pixel_values.reshape(
71 + pixel_values.shape[0],
72 + self.num_channels,
73 + self.temporal_patch_size,
74 + self.patch_size,
75 + self.patch_size,
76 + )
77 + return self.patch_embedding(x).reshape(x.shape[0], -1)
78 +
79 +
80 +class MiniMaxVLAttention(nn.Module):
81 + """Multi-head attention with MiniMax's partial 3D RoPE.
82 +
83 + Partial means only the first ``rot_dim`` (< head_dim) dimensions of
84 + Q and K are rotated; the remaining dims are passed through unchanged.
85 + """
86 +
87 + def __init__(
88 + self,
89 + embed_dim: int,
90 + num_heads: int,
91 + quant_config: QuantizationConfig | None = None,
92 + prefix: str = "",
93 + ) -> None:
94 + super().__init__()
95 + use_data_parallel = is_vit_use_data_parallel()
96 + self.tp_size = (
97 + 1
98 + if use_data_parallel
99 + else parallel_state.get_tensor_model_parallel_world_size()
100 + )
101 + self.head_dim = embed_dim // num_heads
102 + self.num_heads_per_partition = dist_utils.divide(num_heads, self.tp_size)
103 +
104 + self.qkv_proj = QKVParallelLinear(
105 + hidden_size=embed_dim,
106 + head_size=self.head_dim,
107 + total_num_heads=num_heads,
108 + total_num_kv_heads=num_heads,
109 + bias=True,
110 + quant_config=quant_config,
111 + prefix=f"{prefix}.qkv_proj",
112 + disable_tp=use_data_parallel,
113 + )
114 + self.out_proj = RowParallelLinear(
115 + input_size=embed_dim,
116 + output_size=embed_dim,
117 + bias=True,
118 + quant_config=quant_config,
119 + prefix=f"{prefix}.out_proj",
120 + disable_tp=use_data_parallel,
121 + )
122 + self.attn = MMEncoderAttention(
123 + num_heads=self.num_heads_per_partition,
124 + head_size=self.head_dim,
125 + prefix=f"{prefix}.attn",
126 + )
127 + # ApplyRotaryEmb handles the internal cos/sin repeat and partial
128 + # rotation (ro_dim = half_rot_dim * 2 < head_dim for MiniMax).
129 + # enable_fp32_compute=True runs the rotation in fp32 (q/k upcast,
130 + # fp32 cos/sin), matching the reference ``_minimax_rope_applier``.
131 + self.apply_rotary_emb = ApplyRotaryEmb(
132 + enforce_enable=True, enable_fp32_compute=True
133 + )
134 +
135 + def _apply_rotary_emb(
136 + self,
137 + qk_reshaped: torch.Tensor,
138 + rotary_cos: torch.Tensor,
139 + rotary_sin: torch.Tensor,
140 + seq_len: int,
141 + rotary_segment_lengths: list[int] | None,
142 + ) -> torch.Tensor:
143 + # Default fast path (all NVIDIA inputs, and ROCm short clips/images):
144 + # a single rotary kernel launch. ``rotary_segment_lengths`` is only
145 + # populated on ROCm (see ``MiniMaxVLVisionTransformer.forward``), so
146 + # the per-segment path below is ROCm-only and never touches the
147 + # NVIDIA/CUDA code path.
148 + if not current_platform.is_rocm() or rotary_segment_lengths is None:
149 + return self.apply_rotary_emb(qk_reshaped, rotary_cos, rotary_sin)
150 +
151 + # ROCm only: the HIP flash-attn Triton rotary kernel fails with
152 + # hipErrorInvalidValue once grid.y = cdiv(seqlen, BLOCK_M) exceeds
153 + # _HIP_MAX_GRID_DIM_Y (65536). BLOCK_M is 8 for rotary_dim <= 128
154 + # (MiniMax-M3 vision: rotary_dim=78), giving a hard limit of
155 + # 65536 * BLOCK_M tokens — measured exactly as 524288 OK / 524289 fail.
156 + # Only long videos cross it; since vision_segment_max_frames caps each
157 + # segment at a few frames (<< limit), applying RoPE per segment keeps
158 + # every sub-call in range. Splitting on segment boundaries is
159 + # mathematically exact because rotary_cos/sin are precomputed per token.
160 + # Images and short clips stay on the single-kernel fast path above.
161 + rotary_dim = rotary_cos.shape[-1] * 2
162 + block_m = 8 if rotary_dim <= 128 else 4
163 + hip_rotary_max_seqlen = _HIP_MAX_GRID_DIM_Y * block_m
164 + if seq_len <= hip_rotary_max_seqlen or len(rotary_segment_lengths) <= 1:
165 + return self.apply_rotary_emb(qk_reshaped, rotary_cos, rotary_sin)
166 +
167 + qk_segments = qk_reshaped.split(rotary_segment_lengths, dim=1)
168 + cos_segments = rotary_cos.split(rotary_segment_lengths, dim=0)
169 + sin_segments = rotary_sin.split(rotary_segment_lengths, dim=0)
170 + return torch.cat(
171 + [
172 + self.apply_rotary_emb(qk_s, cos_s, sin_s)
173 + for qk_s, cos_s, sin_s in zip(qk_segments, cos_segments, sin_segments)
174 + ],
175 + dim=1,
176 + )
177 +
178 + def forward(
179 + self,
180 + x: torch.Tensor,
181 + cu_seqlens: torch.Tensor,
182 + rotary_cos: torch.Tensor,
183 + rotary_sin: torch.Tensor,
184 + max_seqlen: torch.Tensor,
185 + rotary_segment_lengths: list[int] | None = None,
186 + sequence_lengths: torch.Tensor | None = None,
187 + ) -> torch.Tensor:
188 + # x: (N, 1, embed_dim) [seq=N, batch=1, chan=embed_dim]
189 + x_qkv, _ = self.qkv_proj(x) # (N, 1, 3 * heads_per_part * head_dim)
190 + seq_len, batch_size, _ = x_qkv.shape
191 +
192 + # Rearrange to (b=1, N, 3, heads, head_dim) — same as Qwen2_5_VisionAttention
193 + qkv = rearrange(
194 + x_qkv,
195 + "s b (three head d) -> b s three head d",
196 + three=3,
197 + head=self.num_heads_per_partition,
198 + )
199 + qk, v = qkv[:, :, :2], qkv[:, :, 2] # (b,N,2,h,d) and (b,N,h,d)
200 +
201 + # Stack q/k → (2*b, N, heads, head_dim) for joint RoPE application.
202 + # rotary_cos/sin: (N, half_rot_dim) — ApplyRotaryEmb expands internally
203 + # and rotates only the first 2*half_rot_dim dims, passing the rest through.
204 + qk_reshaped = rearrange(qk, "b s two h d -> (two b) s h d", two=2).contiguous()
205 + qk_rotated = self._apply_rotary_emb(
206 + qk_reshaped, rotary_cos, rotary_sin, seq_len, rotary_segment_lengths
207 + )
208 + qk_rotated = qk_rotated.view(
209 + 2, batch_size, seq_len, self.num_heads_per_partition, self.head_dim
210 + )
211 + q, k = qk_rotated.unbind(dim=0) # each (b=1, N, heads, head_dim)
212 +
213 + # Flash attention → (b, N, heads, head_dim)
214 + context = self.attn(
215 + query=q,
216 + key=k,
217 + value=v,
218 + cu_seqlens=cu_seqlens,
219 + max_seqlen=max_seqlen,
220 + sequence_lengths=sequence_lengths,
221 + )
222 +
223 + # Back to (N, 1, embed_dim)
224 + context = rearrange(context, "b s h d -> s b (h d)", b=batch_size)
225 + output, _ = self.out_proj(context)
226 + return output
227 +
228 +
229 +class MiniMaxVLEncoderLayer(nn.Module):
230 + """Single CLIP-style transformer block."""
231 +
232 + def __init__(
233 + self,
234 + config: PretrainedConfig,
235 + quant_config: QuantizationConfig | None = None,
236 + prefix: str = "",
237 + ) -> None:
238 + super().__init__()
239 + embed_dim = config.hidden_size
240 + self.layer_norm1 = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
241 + self.self_attn = MiniMaxVLAttention(
242 + embed_dim=embed_dim,
243 + num_heads=config.num_attention_heads,
244 + quant_config=quant_config,
245 + prefix=f"{prefix}.self_attn",
246 + )
247 + self.layer_norm2 = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
248 + use_data_parallel = is_vit_use_data_parallel()
249 + self.fc1 = ColumnParallelLinear(
250 + config.hidden_size,
251 + config.intermediate_size,
252 + bias=True,
253 + quant_config=quant_config,
254 + prefix=f"{prefix}.fc1",
255 + disable_tp=use_data_parallel,
256 + )
257 + self.act = get_act_fn(getattr(config, "hidden_act", "gelu"))
258 + self.fc2 = RowParallelLinear(
259 + config.intermediate_size,
260 + config.hidden_size,
261 + bias=True,
262 + quant_config=quant_config,
263 + prefix=f"{prefix}.fc2",
264 + disable_tp=use_data_parallel,
265 + )
266 +
267 + def forward(
268 + self,
269 + x: torch.Tensor,
270 + cu_seqlens: torch.Tensor,
271 + rotary_cos: torch.Tensor,
272 + rotary_sin: torch.Tensor,
273 + max_seqlen: torch.Tensor,
274 + rotary_segment_lengths: list[int] | None = None,
275 + sequence_lengths: torch.Tensor | None = None,
276 + ) -> torch.Tensor:
277 + # x: (N, 1, hidden_size)
278 + x = x + self.self_attn(
279 + self.layer_norm1(x),
280 + cu_seqlens,
281 + rotary_cos,
282 + rotary_sin,
283 + max_seqlen,
284 + rotary_segment_lengths,
285 + sequence_lengths,
286 + )
287 + residual = x
288 + x, _ = self.fc1(self.layer_norm2(x))
289 + x = self.act(x)
290 + x, _ = self.fc2(x)
291 + return residual + x
292 +
293 +
294 +class MiniMaxVLEncoder(nn.Module):
295 + def __init__(
296 + self,
297 + config: PretrainedConfig,
298 + num_hidden_layers_override: int | None = None,
299 + quant_config: QuantizationConfig | None = None,
300 + prefix: str = "",
301 + ) -> None:
302 + super().__init__()
303 + n = (
304 + config.num_hidden_layers
305 + if num_hidden_layers_override is None
306 + else num_hidden_layers_override
307 + )
308 + self.layers = nn.ModuleList(
309 + [
310 + MiniMaxVLEncoderLayer(
311 + config=config,
312 + quant_config=quant_config,
313 + prefix=f"{prefix}.layers.{i}",
314 + )
315 + for i in range(n)
316 + ]
317 + )
318 +
319 + def forward(
320 + self,
321 + x: torch.Tensor,
322 + cu_seqlens: torch.Tensor,
323 + rotary_cos: torch.Tensor,
324 + rotary_sin: torch.Tensor,
325 + max_seqlen: torch.Tensor,
326 + rotary_segment_lengths: list[int] | None = None,
327 + sequence_lengths: torch.Tensor | None = None,
328 + ) -> torch.Tensor:
329 + for layer in self.layers:
330 + x = layer(
331 + x,
332 + cu_seqlens,
333 + rotary_cos,
334 + rotary_sin,
335 + max_seqlen,
336 + rotary_segment_lengths,
337 + sequence_lengths,
338 + )
339 + return x
340 +
341 +
342 +class MiniMaxVLVisionTransformer(nn.Module):
343 + """CLIP-based ViT with 3D RoPE (t/h/w decomposed).
344 +
345 + Faithfully mirrors the reference ``MiniMaxVLVisionTransformer``.
346 + FLASHINFER backend is not supported; standard flash-attn is used.
347 + """
348 +
349 + def __init__(
350 + self,
351 + config: PretrainedConfig,
352 + num_hidden_layers_override: int | None = None,
353 + require_post_norm: bool | None = None,
354 + quant_config: QuantizationConfig | None = None,
355 + prefix: str = "",
356 + ) -> None:
357 + super().__init__()
358 + compression = config.img_token_compression_config
359 + self.spatial_merge_size: int = compression.get("spatial_merge_size", 2)
360 + self.temporal_patch_size: int = compression.get("temporal_patch_size", 2)
361 + self.vision_segment_max_frames: int | None = getattr(
362 + config, "vision_segment_max_frames", None
363 + )
364 + self.use_data_parallel = is_vit_use_data_parallel()
365 +
366 + embed_dim = config.hidden_size
367 + head_dim = embed_dim // config.num_attention_heads
368 + # Backend selection + sharding info for building encoder metadata.
369 + # Defaults to FLASH_ATTN on SM80+; --mm-encoder-attn-backend FLASHINFER
370 + # selects the cuDNN ViT prefill path.
371 + self.hidden_size = embed_dim
372 + self.tp_size = (
373 + 1
374 + if self.use_data_parallel
375 + else parallel_state.get_tensor_model_parallel_world_size()
376 + )
377 + self.attn_backend = get_vit_attn_backend(
378 + head_size=head_dim, dtype=torch.get_default_dtype()
379 + )
380 + rope_dims = 2 * (head_dim // 2)
381 +
382 + # Split rope dims evenly across t/h/w (same formula as the reference)
383 + self.t_dim = int(2 * ((rope_dims // 3) // 2))
384 + self.h_dim = int(2 * ((rope_dims // 3) // 2))
385 + self.w_dim = int(2 * ((rope_dims // 3) // 2))
386 + # rot_dim = t_dim + h_dim + w_dim (may be < head_dim)
387 +
388 + rope_theta: float = getattr(config, "rope_theta", 10000.0)
389 + inv_freq_t = 1.0 / (
390 + rope_theta
391 + ** (torch.arange(0, self.t_dim, 2, dtype=torch.float32) / self.t_dim)
392 + )
393 + inv_freq_h = 1.0 / (
394 + rope_theta
395 + ** (torch.arange(0, self.h_dim, 2, dtype=torch.float32) / self.h_dim)
396 + )
397 + inv_freq_w = 1.0 / (
398 + rope_theta
399 + ** (torch.arange(0, self.w_dim, 2, dtype=torch.float32) / self.w_dim)
400 + )
401 + self.register_buffer("inv_freq_t", inv_freq_t, persistent=False)
402 + self.register_buffer("inv_freq_h", inv_freq_h, persistent=False)
403 + self.register_buffer("inv_freq_w", inv_freq_w, persistent=False)
404 +
405 + self.embeddings = MiniMaxVLPatchEmbed(config)
406 + self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
407 +
408 + n_layers = config.num_hidden_layers
409 + if num_hidden_layers_override is None:
410 + num_hidden_layers_override = n_layers
411 + self.encoder = MiniMaxVLEncoder(
412 + config=config,
413 + num_hidden_layers_override=num_hidden_layers_override,
414 + quant_config=quant_config,
415 + prefix=f"{prefix}.encoder",
416 + )
417 +
418 + if require_post_norm is None:
419 + require_post_norm = num_hidden_layers_override == n_layers
420 + self.post_layernorm = (
421 + nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
422 + if require_post_norm
423 + else None
424 + )
425 +
426 + # out_hidden_size needed by run_dp_sharded_mrope_vision_model
427 + self.out_hidden_size = embed_dim
428 +
429 + # ── RoPE helpers ─────────────────────────────────────────────────────
430 +
431 + def _get_3d_rope_embed(
432 + self, grid_t: int, grid_h: int, grid_w: int, spatial_merge_size: int
433 + ) -> torch.Tensor:
434 + """Compute 3D RoPE frequencies for a single (T, H, W) grid.
435 +
436 + Returns (T*H*W, half_rot_dim) on the same device as inv_freq buffers.
437 + Mirrors the reference ``_get_3d_rope_embed`` exactly.
438 + """
439 + tokens_per_frame = grid_h * grid_w
440 +
441 + tpos_ids = (
442 + torch.arange(grid_t, device=self.inv_freq_t.device)
443 + .unsqueeze(1)
444 + .expand(-1, tokens_per_frame)
445 + .flatten()
446 + )
447 +
448 + hpos_ids = (
449 + torch.arange(grid_h, device=self.inv_freq_h.device)
450 + .unsqueeze(1)
451 + .expand(-1, grid_w)
452 + .reshape(
453 + grid_h // spatial_merge_size,
454 + spatial_merge_size,
455 + grid_w // spatial_merge_size,
456 + spatial_merge_size,
457 + )
458 + .permute(0, 2, 1, 3)
459 + .unsqueeze(0)
460 + .expand(grid_t, -1, -1, -1, -1)
461 + .flatten()
462 + )
463 + wpos_ids = (
464 + torch.arange(grid_w, device=self.inv_freq_w.device)
465 + .unsqueeze(0)
466 + .expand(grid_h, -1)
467 + .reshape(
468 + grid_h // spatial_merge_size,
469 + spatial_merge_size,
470 + grid_w // spatial_merge_size,
471 + spatial_merge_size,
472 + )
473 + .permute(0, 2, 1, 3)
474 + .unsqueeze(0)
475 + .expand(grid_t, -1, -1, -1, -1)
476 + .flatten()
477 + )
478 +
479 + max_t = max(grid_t, 1)
480 + max_hw = max(grid_h, grid_w)
481 +
482 + seq_t = torch.arange(
483 + max_t, device=self.inv_freq_t.device, dtype=self.inv_freq_t.dtype
484 + )
485 + seq_hw = torch.arange(
486 + max_hw, device=self.inv_freq_h.device, dtype=self.inv_freq_h.dtype
487 + )
488 +
489 + freqs_t = torch.outer(seq_t, self.inv_freq_t) # (max_t, t_dim/2)
490 + freqs_h = torch.outer(seq_hw, self.inv_freq_h) # (max_hw, h_dim/2)
491 + freqs_w = torch.outer(seq_hw, self.inv_freq_w) # (max_hw, w_dim/2)
492 +
493 + return torch.cat(
494 + [freqs_t[tpos_ids], freqs_h[hpos_ids], freqs_w[wpos_ids]], dim=-1
495 + ) # (T*H*W, half_rot_dim)
496 +
497 + def _get_rope_embed_3d(
498 + self, grid_thw: list[list[int]], spatial_merge_size: int
499 + ) -> torch.Tensor:
500 + embeds = [
501 + self._get_3d_rope_embed(t, h, w, spatial_merge_size) for t, h, w in grid_thw
502 + ]
503 + return torch.cat(embeds, dim=0) # (total_N, half_rot_dim)
504 +
505 + # ── Frame-limit helper (mirrors the reference) ───────────────────────
506 +
507 + def _apply_max_frames_limit(self, grid_thw: list[list[int]]) -> list[list[int]]:
508 + if self.vision_segment_max_frames is None:
509 + return grid_thw
510 + max_f = self.vision_segment_max_frames
511 + out: list[list[int]] = []
512 + for t, h, w in grid_thw:
513 + if t <= max_f:
514 + out.append([t, h, w])
515 + else:
516 + for i in range(0, t, max_f):
517 + out.append([min(max_f, t - i), h, w])
518 + return out
519 +
520 + # ── Forward ──────────────────────────────────────────────────────────
521 +
522 + def forward(
523 + self,
524 + pixel_values: torch.Tensor,
525 + grid_thw: list[list[int]],
526 + ) -> torch.Tensor:
527 + # pixel_values: (total_N, C * temporal_patch_size * patch_size²)
528 + # Output: (total_N, hidden_size)
529 +
530 + hidden = self.embeddings(pixel_values) # (total_N, hidden_size)
531 + hidden = self.pre_layrnorm(hidden)
532 +
533 + limited = self._apply_max_frames_limit(grid_thw)
534 +
535 + # Token-level cumulative sequence lengths (one segment per limited grid).
536 + lens = [t * h * w for t, h, w in limited]
537 + cu_seqlens_np = np.zeros(len(lens) + 1, dtype=np.int32)
538 + np.cumsum(np.array(lens, dtype=np.int32), out=cu_seqlens_np[1:])
539 +
540 + # Backend-specific encoder metadata. For FLASH_ATTN this returns the raw
541 + # token cu_seqlens, the max segment length, and sequence_lengths=None;
542 + # for FLASHINFER (cuDNN) it repacks cu_seqlens into element-offset
543 + # indptrs, buckets max_seqlen, and builds padded per-sequence lengths.
544 + sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
545 + self.attn_backend, cu_seqlens_np, hidden.device
546 + )
547 + max_seqlen = torch.tensor(
548 + MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens_np),
549 + dtype=torch.int32,
550 + )
551 + cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
552 + self.attn_backend,
553 + cu_seqlens_np,
554 + self.hidden_size,
555 + self.tp_size,
556 + hidden.device,
557 + )
558 +
559 + # 3D RoPE: (total_N, half_rot_dim); ApplyRotaryEmb expands internally
560 + freqs = self._get_rope_embed_3d(limited, self.spatial_merge_size)
561 + freqs = freqs.to(device=hidden.device)
562 + # Keep cos/sin in fp32; ApplyRotaryEmb(enable_fp32_compute=True) runs the
563 + # rotation in fp32 to match the reference precision.
564 + rotary_cos, rotary_sin = freqs.cos(), freqs.sin()
565 +
566 + # Encoder expects (N, 1, hidden_size) — add batch dim
567 + hidden = hidden.unsqueeze(1)
568 + # On ROCm, the flash_attn Triton rotary kernel can fail with
569 + # hipErrorInvalidValue when seqlen is very large, e.g. 192k video
570 + # tokens; pass per-segment lengths so RoPE can be applied in chunks.
571 + # On other platforms leave it None -> single-kernel fast path, so the
572 + # NVIDIA/CUDA code path is unchanged.
573 + rotary_segment_lengths = lens if current_platform.is_rocm() else None
574 +
575 + hidden = self.encoder(
576 + hidden,
577 + cu_seqlens,
578 + rotary_cos,
579 + rotary_sin,
580 + max_seqlen,
581 + rotary_segment_lengths,
582 + sequence_lengths=sequence_lengths,
583 + )
584 + hidden = hidden.squeeze(1) # back to (total_N, hidden_size)
585 +
586 + if self.post_layernorm is not None:
587 + hidden = self.post_layernorm(hidden)
588 +
589 + return hidden
590 +
591 +
592 +class MiniMaxVLMultiModalProjector(nn.Module):
593 + """Two-layer MLP projector: vision_hidden → text_hidden."""
594 +
595 + def __init__(
596 + self,
597 + vision_hidden_size: int,
598 + text_hidden_size: int,
599 + projector_hidden_size: int | None,
600 + multimodal_projector_bias: bool,
601 + projector_hidden_act: str = "gelu",
602 + quant_config: QuantizationConfig | None = None,
603 + prefix: str = "",
604 + ) -> None:
605 + super().__init__()
606 + mid = projector_hidden_size if projector_hidden_size else text_hidden_size
607 + use_dp = is_vit_use_data_parallel()
608 + self.linear_1 = ColumnParallelLinear(
609 + vision_hidden_size,
610 + mid,
611 + bias=multimodal_projector_bias,
612 + quant_config=quant_config,
613 + prefix=f"{prefix}.linear_1",
614 + disable_tp=use_dp,
615 + )
616 + self.act = get_act_fn(projector_hidden_act)
617 + self.linear_2 = RowParallelLinear(
618 + mid,
619 + text_hidden_size,
620 + bias=multimodal_projector_bias,
621 + quant_config=quant_config,
622 + prefix=f"{prefix}.linear_2",
623 + disable_tp=use_dp,
624 + )
625 +
626 + def forward(self, x: torch.Tensor) -> torch.Tensor:
627 + x, _ = self.linear_1(x)
628 + x = self.act(x)
629 + x, _ = self.linear_2(x)
630 + return x
631 +
632 +
633 +class MiniMaxVLPatchMerger(nn.Module):
634 + def __init__(
635 + self,
636 + spatial_merge_size: int,
637 + text_hidden_size: int,
638 + projector_hidden_size: int | None,
639 + patch_merge_bias: bool,
640 + projector_hidden_act: str = "gelu",
641 + quant_config: QuantizationConfig | None = None,
642 + prefix: str = "",
643 + ) -> None:
644 + super().__init__()
645 + self.spatial_merge_size = spatial_merge_size
646 + mid = projector_hidden_size if projector_hidden_size else text_hidden_size
647 + merge_in = text_hidden_size * spatial_merge_size**2
648 + use_dp = is_vit_use_data_parallel()
649 + self.linear_1 = ColumnParallelLinear(
650 + merge_in,
651 + mid,
652 + bias=patch_merge_bias,
653 + quant_config=quant_config,
654 + prefix=f"{prefix}.linear_1",
655 + disable_tp=use_dp,
656 + )
657 + self.act = get_act_fn(projector_hidden_act)
658 + self.linear_2 = RowParallelLinear(
659 + mid,
660 + text_hidden_size,
661 + bias=patch_merge_bias,
662 + quant_config=quant_config,
663 + prefix=f"{prefix}.linear_2",
664 + disable_tp=use_dp,
665 + )
666 +
667 + def forward(self, x: torch.Tensor) -> torch.Tensor:
668 + # x: (N, text_hidden_size) → (N // merge_size², text_hidden_size)
669 + x = x.reshape(x.shape[0] // (self.spatial_merge_size**2), -1)
670 + x, _ = self.linear_1(x)
671 + x = self.act(x)
672 + x, _ = self.linear_2(x)
673 + return x
674 +
675 +
676 +class MiniMaxVLVisionModel(nn.Module):
677 + """Full vision model: ViT → projector → patch merger."""
678 +
679 + def __init__(
680 + self,
681 + config: PretrainedConfig,
682 + text_hidden_size: int,
683 + projector_hidden_size: int | None = None,
684 + quant_config: QuantizationConfig | None = None,
685 + prefix: str = "",
686 + ) -> None:
687 + super().__init__()
688 + compression = config.img_token_compression_config
689 + spatial_merge_size: int = compression.get("spatial_merge_size", 2)
690 + self.spatial_merge_size = spatial_merge_size
691 + self.use_data_parallel = is_vit_use_data_parallel()
692 +
693 + # The released checkpoint ships no ``post_layernorm`` weights and
694 + # uses ``vision_feature_layer=-1`` with ``vision_feature_select_strategy
695 + # ="full"``, i.e. the raw last encoder hidden state (CLIP's
696 + # ``last_hidden_state`` is taken before the post layernorm). Applying an
697 + # untrained post layernorm here would corrupt the visual features.
698 + self.vision_model = MiniMaxVLVisionTransformer(
699 + config=config,
700 + require_post_norm=False,
701 + quant_config=quant_config,
702 + prefix=maybe_prefix(prefix, "vision_model"),
703 + )
704 + self.multi_modal_projector = MiniMaxVLMultiModalProjector(
705 + vision_hidden_size=config.hidden_size,
706 + text_hidden_size=text_hidden_size,
707 + projector_hidden_size=projector_hidden_size,
708 + multimodal_projector_bias=getattr(
709 + config, "multimodal_projector_bias", True
710 + ),
711 + projector_hidden_act=getattr(config, "projector_hidden_act", "gelu"),
712 + quant_config=quant_config,
713 + prefix=maybe_prefix(prefix, "multi_modal_projector"),
714 + )
715 + self.patch_merge_mlp = MiniMaxVLPatchMerger(
716 + spatial_merge_size=spatial_merge_size,
717 + text_hidden_size=text_hidden_size,
718 + projector_hidden_size=projector_hidden_size,
719 + patch_merge_bias=getattr(config, "patch_merge_bias", True),
720 + projector_hidden_act=getattr(config, "projector_hidden_act", "gelu"),
721 + quant_config=quant_config,
722 + prefix=maybe_prefix(prefix, "patch_merge_mlp"),
723 + )
724 +
725 + self.dtype = self.vision_model.embeddings.patch_embedding.weight.dtype
726 + self.out_hidden_size = text_hidden_size
727 +
728 + def forward(
729 + self,
730 + pixel_values: torch.Tensor,
731 + grid_thw: list[list[int]],
732 + ) -> torch.Tensor:
733 + hidden = self.vision_model(pixel_values=pixel_values, grid_thw=grid_thw)
734 + if hidden.dim() == 3:
735 + hidden = hidden.squeeze(0)
736 + hidden = self.multi_modal_projector(hidden)
737 + hidden = self.patch_merge_mlp(hidden)
738 + return hidden
739 +
740 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
741 + stacked_params_mapping = [
742 + # (param_name, shard_name, shard_id)
743 + ("qkv_proj.", "q_proj.", "q"),
744 + ("qkv_proj.", "k_proj.", "k"),
745 + ("qkv_proj.", "v_proj.", "v"),
746 + ]
747 + params_dict = dict(self.named_parameters(remove_duplicate=False))
748 + loaded_params: set[str] = set()
749 +
750 + for name, loaded_weight in weights:
751 + for param_name, weight_name, shard_id in stacked_params_mapping:
752 + if weight_name not in name:
753 + continue
754 + name = name.replace(weight_name, param_name)
755 +
756 + param = params_dict[name]
757 + weight_loader = param.weight_loader
758 + weight_loader(param, loaded_weight, shard_id)
759 + break
760 + else:
761 + param = params_dict[name]
762 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
763 + weight_loader(param, loaded_weight)
764 + loaded_params.add(name)
765 + return loaded_params
vllm/models/minimax_m3/nvidia/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/minimax_m3/nvidia/model.py new
+1177
@@ -0,0 +1,1177 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only MiniMax M3 (text backbone) model.
4 +
5 +The MiniMax-M3-preview config selects a single set of branches:
6 + * qk_norm_type == "per_head"
7 + * hidden_act == "swigluoai"
8 + * use_gemma_norm == True -> Gemma-style RMSNorm everywhere
9 + * attention_output_gate == False
10 + * scoring_func == "sigmoid" with a routing-bias correction term
11 + * sparse_attention_config present -> a subset of layers run the extra
12 + "index" attention branch.
13 +"""
14 +
15 +from collections.abc import Iterable
16 +
17 +import torch
18 +from torch import nn
19 +from transformers import PretrainedConfig
20 +
21 +from vllm import _custom_ops as ops
22 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
23 +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
24 +from vllm.distributed import 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_layer_base import AttentionLayerBase
29 +from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
30 + fused_allreduce_gemma_rms_norm,
31 +)
32 +from vllm.model_executor.layers.fused_moe import (
33 + FusedMoE,
34 + GateLinear,
35 + fused_moe_make_expert_params_mapping,
36 +)
37 +from vllm.model_executor.layers.linear import (
38 + MergedColumnParallelLinear,
39 + MinimaxM3QKVParallelLinearWithIndexer,
40 + QKVParallelLinear,
41 + RowParallelLinear,
42 +)
43 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
44 +from vllm.model_executor.layers.quantization import QuantizationConfig
45 +from vllm.model_executor.layers.rotary_embedding import get_rope
46 +from vllm.model_executor.layers.vocab_parallel_embedding import (
47 + ParallelLMHead,
48 + VocabParallelEmbedding,
49 +)
50 +from vllm.model_executor.model_loader.weight_utils import (
51 + default_weight_loader,
52 + maybe_remap_kv_scale_name,
53 +)
54 +from vllm.model_executor.models.interfaces import (
55 + EagleModelMixin,
56 + MultiModalEmbeddings,
57 + SupportsEagle3,
58 + SupportsMultiModal,
59 +)
60 +from vllm.model_executor.models.utils import (
61 + AutoWeightsLoader,
62 + WeightsMapper,
63 + init_vllm_registered_model,
64 + is_pp_missing_parameter,
65 + make_layers,
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 +)
73 +from vllm.models.minimax_m3.common.mm_preprocess import (
74 + MiniMaxM3VLDummyInputsBuilder,
75 + MiniMaxM3VLMultiModalProcessor,
76 + MiniMaxM3VLProcessingInfo,
77 +)
78 +from vllm.models.minimax_m3.common.sparse_attention import (
79 + MiniMaxM3SparseBackend,
80 + MiniMaxM3SparseImpl,
81 + MiniMaxM3SparseMetadata,
82 + select_main_impl_cls,
83 +)
84 +from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel
85 +from vllm.multimodal import MULTIMODAL_REGISTRY
86 +from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
87 +from vllm.v1.kv_cache_interface import (
88 + FullAttentionSpec,
89 + KVCacheSpec,
90 + get_kv_quant_mode,
91 +)
92 +
93 +
94 +def _sparse_attention_layer_ids(config: PretrainedConfig) -> set[int]:
95 + """Layer ids whose attention runs the extra sparse "index" branch."""
96 + cfg = getattr(config, "sparse_attention_config", None)
97 + if not cfg:
98 + return set()
99 + freq = cfg.get("sparse_attention_freq")
100 + if freq is None:
101 + return set()
102 + return {i for i, f in enumerate(freq) if f != 0}
103 +
104 +
105 +def _is_moe_layer(config: PretrainedConfig, layer_id: int) -> bool:
106 + """Whether this layer's MLP is a sparse MoE block (vs a dense MLP)."""
107 + moe_layer_freq = getattr(config, "moe_layer_freq", None)
108 + if moe_layer_freq is None:
109 + return True
110 + return moe_layer_freq[layer_id] != 0
111 +
112 +
113 +class MiniMAXGemmaRMSNorm(nn.Module):
114 + """Gemma-style RMS normalization backed by FlashInfer kernels.
115 +
116 + When ``residual`` is given, the fused add + norm runs in place and the
117 + updated ``(x, residual)`` pair is returned.
118 + """
119 +
120 + def __init__(
121 + self,
122 + hidden_size: int,
123 + eps: float = 1e-6,
124 + ) -> None:
125 + super().__init__()
126 + self.weight = nn.Parameter(torch.zeros(hidden_size))
127 + self.variance_epsilon = eps
128 +
129 + def forward(
130 + self,
131 + x: torch.Tensor,
132 + residual: torch.Tensor | None = None,
133 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
134 + from flashinfer.norm import gemma_fused_add_rmsnorm, gemma_rmsnorm
135 +
136 + if residual is None:
137 + return gemma_rmsnorm(x, self.weight, self.variance_epsilon)
138 +
139 + # gemma_fused_add_rmsnorm mutates x and residual in place.
140 + gemma_fused_add_rmsnorm(x, residual, self.weight, self.variance_epsilon)
141 + return x, residual
142 +
143 +
144 +class MiniMaxM3MLP(nn.Module):
145 + """Dense SwiGLU-OAI MLP (used by the leading dense layers)."""
146 +
147 + def __init__(
148 + self,
149 + config: PretrainedConfig,
150 + intermediate_size: int,
151 + quant_config: QuantizationConfig | None = None,
152 + reduce_results: bool = True,
153 + prefix: str = "",
154 + ) -> None:
155 + super().__init__()
156 + self.gate_up_proj = MergedColumnParallelLinear(
157 + config.hidden_size,
158 + [intermediate_size] * 2,
159 + bias=False,
160 + quant_config=quant_config,
161 + prefix=f"{prefix}.gate_up_proj",
162 + )
163 + self.down_proj = RowParallelLinear(
164 + intermediate_size,
165 + config.hidden_size,
166 + bias=False,
167 + quant_config=quant_config,
168 + reduce_results=reduce_results,
169 + prefix=f"{prefix}.down_proj",
170 + )
171 + if config.hidden_act != "swigluoai":
172 + raise ValueError(
173 + f"Unsupported activation: {config.hidden_act}. "
174 + "Only swigluoai is supported."
175 + )
176 + # gate * sigmoid(alpha * gate) * (up + beta), with both halves clamped.
177 + self.act_fn = SiluAndMulWithClamp(
178 + swiglu_limit=config.swiglu_limit,
179 + alpha=config.swiglu_alpha,
180 + beta=config.swiglu_beta,
181 + )
182 +
183 + def forward(self, x: torch.Tensor) -> torch.Tensor:
184 + gate_up, _ = self.gate_up_proj(x)
185 + x = self.act_fn(gate_up)
186 + x, _ = self.down_proj(x)
187 + return x
188 +
189 +
190 +class MiniMaxM3MoE(nn.Module):
191 + """Sigmoid-routed MoE block with a routing-bias correction and a shared
192 + expert."""
193 +
194 + def __init__(
195 + self,
196 + config: PretrainedConfig,
197 + layer_id: int,
198 + quant_config: QuantizationConfig | None = None,
199 + prefix: str = "",
200 + ) -> None:
201 + super().__init__()
202 + self.tp_size = get_tensor_model_parallel_world_size()
203 + if self.tp_size > config.num_local_experts:
204 + raise ValueError(
205 + f"Tensor parallel size {self.tp_size} is greater than "
206 + f"the number of experts {config.num_local_experts}."
207 + )
208 +
209 + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
210 + self.n_shared_experts = getattr(config, "n_shared_experts", None)
211 +
212 + # Sigmoid routing uses a per-expert score-correction bias for selection.
213 + self.use_routing_bias = getattr(config, "use_routing_bias", False)
214 + if self.use_routing_bias:
215 + self.e_score_correction_bias = nn.Parameter(
216 + torch.empty(config.num_local_experts, dtype=torch.float32)
217 + )
218 + self.e_score_correction_bias.weight_loader = (
219 + MiniMaxM3MoE.ebias_weight_loader
220 + )
221 + else:
222 + self.e_score_correction_bias = None
223 +
224 + # Router weights are stored in fp32; GateLinear upcasts the bf16
225 + # activations and computes the gate in fp32 (fp32 router logits).
226 + self.gate = GateLinear(
227 + config.hidden_size,
228 + config.num_local_experts,
229 + bias=False,
230 + params_dtype=torch.float32,
231 + out_dtype=torch.float32,
232 + prefix=f"{prefix}.gate",
233 + )
234 +
235 + self.shared_experts: MiniMaxM3MLP | None = None
236 + if self.n_shared_experts:
237 + self.shared_experts = MiniMaxM3MLP(
238 + config=config,
239 + intermediate_size=config.intermediate_size * self.n_shared_experts,
240 + quant_config=quant_config,
241 + reduce_results=False,
242 + prefix=f"{prefix}.shared_experts",
243 + )
244 +
245 + self.experts = FusedMoE(
246 + num_experts=config.num_local_experts,
247 + top_k=config.num_experts_per_tok,
248 + hidden_size=config.hidden_size,
249 + intermediate_size=config.intermediate_size,
250 + scoring_func=config.scoring_func,
251 + e_score_correction_bias=self.e_score_correction_bias,
252 + renormalize=True,
253 + # w13 (gate_up_proj) is loaded packed via MergedColumnParallelLinear
254 + # ([all gates; all ups]), so use the uninterleaved SwiGLU-OAI variant
255 + # rather than the interleaved gpt-oss layout.
256 + activation="swigluoai_uninterleave",
257 + swiglu_limit=config.swiglu_limit,
258 + swiglu_alpha=config.swiglu_alpha,
259 + swiglu_beta=config.swiglu_beta,
260 + routed_scaling_factor=self.routed_scaling_factor,
261 + apply_routed_scale_to_output=True,
262 + router_logits_dtype=self.gate.out_dtype,
263 + shared_experts=self.shared_experts,
264 + quant_config=quant_config,
265 + prefix=f"{prefix}.experts",
266 + )
267 +
268 + @staticmethod
269 + def ebias_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None:
270 + assert param.size() == loaded_weight.size()
271 + param.data.copy_(loaded_weight.to(torch.float32))
272 +
273 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
274 + num_tokens, hidden_dim = hidden_states.shape
275 + hidden_states = hidden_states.view(-1, hidden_dim)
276 +
277 + # router_logits: (num_tokens, n_experts); GateLinear casts to fp32.
278 + router_logits, _ = self.gate(hidden_states)
279 + final_hidden_states = self.experts(
280 + hidden_states=hidden_states, router_logits=router_logits
281 + )
282 +
283 + return final_hidden_states.view(num_tokens, hidden_dim)
284 +
285 +
286 +class MiniMaxM3Attention(nn.Module):
287 + """Dense attention with per-head QK norm and partial RoPE."""
288 +
289 + def __init__(
290 + self,
291 + config: PretrainedConfig,
292 + layer_id: int,
293 + quant_config: QuantizationConfig | None = None,
294 + prefix: str = "",
295 + cache_config: CacheConfig | None = None,
296 + ) -> None:
297 + super().__init__()
298 + self.hidden_size = config.hidden_size
299 + tp_size = get_tensor_model_parallel_world_size()
300 +
301 + self.total_num_heads = config.num_attention_heads
302 + assert self.total_num_heads % tp_size == 0
303 + self.num_heads = self.total_num_heads // tp_size
304 + self.total_num_kv_heads = config.num_key_value_heads
305 + if self.total_num_kv_heads >= tp_size:
306 + assert self.total_num_kv_heads % tp_size == 0
307 + else:
308 + assert tp_size % self.total_num_kv_heads == 0
309 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
310 + self.head_dim = config.head_dim
311 + self.q_size = self.num_heads * self.head_dim
312 + self.kv_size = self.num_kv_heads * self.head_dim
313 + self.scaling = self.head_dim**-0.5
314 +
315 + self.qkv_proj = QKVParallelLinear(
316 + self.hidden_size,
317 + self.head_dim,
318 + self.total_num_heads,
319 + self.total_num_kv_heads,
320 + bias=False,
321 + quant_config=quant_config,
322 + prefix=f"{prefix}.qkv_proj",
323 + )
324 + # reduce_results=False: the attention all-reduce is fused with the
325 + # following post_attention_layernorm (GemmaRMSNorm) in the decoder layer
326 + # via fused_allreduce_gemma_rms_norm.
327 + self.o_proj = RowParallelLinear(
328 + self.total_num_heads * self.head_dim,
329 + self.hidden_size,
330 + bias=False,
331 + reduce_results=False,
332 + quant_config=quant_config,
333 + prefix=f"{prefix}.o_proj",
334 + )
335 +
336 + # Per-head QK norm (qk_norm_type == "per_head", use_gemma_norm == True).
337 + self.q_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
338 + self.k_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
339 +
340 + # Partial RoPE: rotary_dim == head_dim * partial_rotary_factor.
341 + self.rotary_emb = get_rope(
342 + self.head_dim,
343 + max_position=config.max_position_embeddings,
344 + rope_parameters={
345 + "rope_theta": config.rope_theta,
346 + "partial_rotary_factor": config.partial_rotary_factor,
347 + },
348 + )
349 +
350 + self.attn = Attention(
351 + self.num_heads,
352 + self.head_dim,
353 + self.scaling,
354 + num_kv_heads=self.num_kv_heads,
355 + cache_config=cache_config,
356 + quant_config=quant_config,
357 + prefix=f"{prefix}.attn",
358 + )
359 +
360 + def forward(
361 + self,
362 + positions: torch.Tensor,
363 + hidden_states: torch.Tensor,
364 + ) -> torch.Tensor:
365 + qkv, _ = self.qkv_proj(hidden_states)
366 + # Fused per-head Gemma QK-norm + partial NeoX RoPE on q/k, in place.
367 +
368 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
369 + qkv,
370 + self.q_norm.weight,
371 + self.k_norm.weight,
372 + self.rotary_emb.cos_sin_cache,
373 + positions,
374 + self.num_heads,
375 + self.num_kv_heads,
376 + self.rotary_emb.rotary_dim,
377 + self.q_norm.variance_epsilon,
378 + )
379 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
380 + attn_output = self.attn(q, k, v)
381 + output, _ = self.o_proj(attn_output)
382 + return output
383 +
384 +
385 +class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
386 + """Block-sparse attention layer with the lightning-indexer branch.
387 +
388 + This is a merged attention layer: it owns the projections (qkv + index
389 + q/k), per-head QK norms and RoPE, *and* the attention-backend wiring that a
390 + generic ``Attention`` layer would normally provide — it binds the
391 + ``MiniMaxM3SparseBackend`` + main impl, registers the main paged K/V cache,
392 + and owns the lightning indexer (``MiniMaxM3Indexer``), which holds the
393 + index-key side cache.
394 +
395 + The index branch (index_{q,k}_proj + index_{q,k}_norm) feeds the sparse
396 + top-k block selection. M3 always disables the index value/output
397 + projections (``sparse_disable_index_value`` set for every sparse layer), so
398 + ``index_{v,o}_proj`` are never created.
399 + """
400 +
401 + def __init__(
402 + self,
403 + config: PretrainedConfig,
404 + layer_id: int,
405 + quant_config: QuantizationConfig | None = None,
406 + prefix: str = "",
407 + cache_config: CacheConfig | None = None,
408 + ) -> None:
409 + super().__init__()
410 + self.hidden_size = config.hidden_size
411 + tp_size = get_tensor_model_parallel_world_size()
412 +
413 + self.total_num_heads = config.num_attention_heads
414 + assert self.total_num_heads % tp_size == 0
415 + self.num_heads = self.total_num_heads // tp_size
416 + self.total_num_kv_heads = config.num_key_value_heads
417 + if self.total_num_kv_heads >= tp_size:
418 + assert self.total_num_kv_heads % tp_size == 0
419 + else:
420 + assert tp_size % self.total_num_kv_heads == 0
421 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
422 + self.head_dim = config.head_dim
423 + self.q_size = self.num_heads * self.head_dim
424 + self.kv_size = self.num_kv_heads * self.head_dim
425 + self.scaling = self.head_dim**-0.5
426 +
427 + # Sparse "index" branch dims. index_q has the same head count as the KV
428 + # heads (sparse_num_index_heads == num_key_value_heads), so it shards
429 + # identically -- including replication when tp_size > num_key_value_heads.
430 + sparse_cfg = config.sparse_attention_config
431 + self.total_idx_heads = sparse_cfg["sparse_num_index_heads"]
432 + self.num_idx_heads = self.num_kv_heads
433 + self.idx_head_dim = sparse_cfg["sparse_index_dim"]
434 + self.index_q_size = self.num_idx_heads * self.idx_head_dim
435 +
436 + # Single fused projection: q, k, v, index_q, index_k in one GEMM.
437 + self.qkv_proj = MinimaxM3QKVParallelLinearWithIndexer(
438 + self.hidden_size,
439 + self.head_dim,
440 + self.total_num_heads,
441 + self.total_num_kv_heads,
442 + self.total_idx_heads,
443 + self.idx_head_dim,
444 + bias=False,
445 + quant_config=quant_config,
446 + prefix=f"{prefix}.qkv_proj",
447 + )
448 + # reduce_results=False: the attention all-reduce is fused with the
449 + # following post_attention_layernorm (GemmaRMSNorm) in the decoder layer
450 + # via fused_allreduce_gemma_rms_norm.
451 + self.o_proj = RowParallelLinear(
452 + self.total_num_heads * self.head_dim,
453 + self.hidden_size,
454 + bias=False,
455 + reduce_results=False,
456 + quant_config=quant_config,
457 + prefix=f"{prefix}.o_proj",
458 + )
459 +
460 + # Per-head QK norm (qk_norm_type == "per_head", use_gemma_norm == True).
461 + self.q_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
462 + self.k_norm = MiniMAXGemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
463 +
464 + # Partial RoPE: rotary_dim == head_dim * partial_rotary_factor.
465 + self.rotary_emb = get_rope(
466 + self.head_dim,
467 + max_position=config.max_position_embeddings,
468 + rope_parameters={
469 + "rope_theta": config.rope_theta,
470 + "partial_rotary_factor": config.partial_rotary_factor,
471 + },
472 + )
473 +
474 + self.index_q_norm = MiniMAXGemmaRMSNorm(
475 + self.idx_head_dim, eps=config.rms_norm_eps
476 + )
477 + self.index_k_norm = MiniMAXGemmaRMSNorm(
478 + self.idx_head_dim, eps=config.rms_norm_eps
479 + )
480 + self.index_rotary_emb = self.rotary_emb
481 +
482 + # Attention-backend wiring.
483 + vllm_config = get_current_vllm_config()
484 + self.layer_name = f"{prefix}.attn"
485 + self.kv_cache_dtype = (
486 + cache_config.cache_dtype if cache_config is not None else "auto"
487 + )
488 + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
489 + self.kv_cache_dtype, vllm_config.model_config
490 + )
491 + # Indexer side-cache dtype, mirroring --kv-cache-dtype for the main
492 + # cache (--attention-config '{"indexer_kv_dtype": ...}').
493 + self.indexer_kv_dtype = vllm_config.attention_config.indexer_kv_dtype
494 +
495 + self.attn_backend = MiniMaxM3SparseBackend
496 + # Indexer (top-k selection) and main attention are separate impls, each
497 + # picking Triton vs MSA off its cache dtype. impl is AttentionImplBase
498 + # (broader than the AttentionImpl that AttentionLayerBase annotates).
499 + self.impl: MiniMaxM3SparseImpl = select_main_impl_cls( # type: ignore[assignment]
500 + topk_blocks=sparse_cfg["sparse_topk_blocks"],
501 + kv_cache_dtype=self.kv_cache_dtype,
502 + )(
503 + self.num_heads,
504 + self.head_dim,
505 + self.scaling,
506 + self.num_kv_heads,
507 + kv_cache_dtype=self.kv_cache_dtype,
508 + topk_blocks=sparse_cfg["sparse_topk_blocks"],
509 + sparse_block_size=sparse_cfg["sparse_block_size"],
510 + )
511 + # Self-contained nn.Module: owns its side cache, selects its impl in init.
512 + self.indexer = MiniMaxM3Indexer(
513 + num_kv_heads=self.num_kv_heads,
514 + scale=self.scaling,
515 + topk_blocks=sparse_cfg["sparse_topk_blocks"],
516 + sparse_block_size=sparse_cfg["sparse_block_size"],
517 + num_index_heads=self.num_idx_heads,
518 + index_head_dim=self.idx_head_dim,
519 + prefix=self.layer_name,
520 + init_blocks=sparse_cfg.get("sparse_init_block", 0),
521 + local_blocks=sparse_cfg.get("sparse_local_block", 0),
522 + score_type=sparse_cfg.get("sparse_score_type", "max"),
523 + cache_config=cache_config,
524 + indexer_kv_dtype=self.indexer_kv_dtype,
525 + )
526 +
527 + # Register the main K/V cache so the KV-cache manager allocates it.
528 + compilation_config = vllm_config.compilation_config
529 + if self.layer_name in compilation_config.static_forward_context:
530 + raise ValueError(f"Duplicate layer name: {self.layer_name}")
531 + compilation_config.static_forward_context[self.layer_name] = self
532 + self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache
533 +
534 + def get_attn_backend(self) -> type[MiniMaxM3SparseBackend]:
535 + return self.attn_backend
536 +
537 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
538 + # Main GQA K/V cache. Block size may change after load, refresh it.
539 + return FullAttentionSpec(
540 + block_size=vllm_config.cache_config.block_size,
541 + num_kv_heads=self.num_kv_heads,
542 + head_size=self.head_dim,
543 + head_size_v=self.head_dim,
544 + dtype=self.kv_cache_torch_dtype,
545 + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
546 + )
547 +
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 +
582 + def forward(
583 + self,
584 + positions: torch.Tensor,
585 + hidden_states: torch.Tensor,
586 + ) -> torch.Tensor:
587 + # Single fused projection emitting [q | k | v | index_q | index_k].
588 + qkv, _ = self.qkv_proj(hidden_states)
589 +
590 + # Horizontally-fused per-head Gemma QK-norm + partial NeoX RoPE on the
591 + # main (q/k) and index (index_q/index_k) branches, all read straight out
592 + # of the single fused ``qkv`` tensor (the "5 results"). Once the paged
593 + # caches are bound the kernel also inserts k/v and the index key into
594 + # 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
599 + # straight into the dedicated contiguous ``q``/``index_q`` buffers below.
600 +
601 + cos_sin_cache = self.rotary_emb.cos_sin_cache
602 + rotary_dim = self.rotary_emb.rotary_dim
603 + eps = self.q_norm.variance_epsilon
604 + num_tokens = qkv.shape[0]
605 +
606 + fwd_slot_mapping = get_forward_context().slot_mapping
607 + if (
608 + not isinstance(fwd_slot_mapping, dict)
609 + or self.layer_name not in fwd_slot_mapping
610 + ):
611 + # Memory-profiling run: caches not yet bound, slot_mapping is empty.
612 + return qkv.new_zeros((num_tokens, self.hidden_size))
613 +
614 + main_slot_mapping = fwd_slot_mapping[self.layer_name]
615 + index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
616 + q = qkv.new_empty((num_tokens, self.q_size))
617 + index_q = qkv.new_empty((num_tokens, self.index_q_size))
618 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
619 + qkv,
620 + self.q_norm.weight,
621 + self.k_norm.weight,
622 + cos_sin_cache,
623 + positions,
624 + self.num_heads,
625 + self.num_kv_heads,
626 + rotary_dim,
627 + eps,
628 + self.index_q_norm.weight,
629 + self.index_k_norm.weight,
630 + self.num_idx_heads,
631 + main_slot_mapping,
632 + index_slot_mapping,
633 + self.kv_cache,
634 + self.indexer.index_cache.kv_cache,
635 + self.kv_cache.size(2), # paged-cache block size
636 + q,
637 + index_q,
638 + )
639 +
640 + output = torch.empty_like(q)
641 + attn_output = self._run_attention(q, index_q, output)
642 + output, _ = self.o_proj(attn_output)
643 + return output
644 +
645 + @eager_break_during_capture
646 + def _run_attention(
647 + self,
648 + query: torch.Tensor,
649 + index_query: torch.Tensor,
650 + output: torch.Tensor,
651 + ) -> torch.Tensor:
652 + # Single eager break around both: their split-K kernels read per-request
653 + # metadata and can't be captured into a cudagraph.
654 + topk_idx = self.indexer(index_query)
655 + return self.impl.forward(self, query, self.kv_cache, topk_idx, output)
656 +
657 +
658 +class MiniMaxM3DecoderLayer(nn.Module):
659 + def __init__(
660 + self,
661 + *,
662 + vllm_config: VllmConfig,
663 + prefix: str,
664 + force_sparse_attn: bool = False,
665 + force_moe: bool = False,
666 + is_mtp_block: bool = False,
667 + ) -> None:
668 + super().__init__()
669 + if is_mtp_block:
670 + assert vllm_config.speculative_config is not None
671 + config = vllm_config.speculative_config.draft_model_config.hf_config
672 + else:
673 + config = vllm_config.model_config.hf_text_config
674 + cache_config = vllm_config.cache_config
675 + quant_config = vllm_config.quant_config
676 + self.hidden_size = config.hidden_size
677 + # DecoderLayers are created with `make_layers` which passes the prefix
678 + # with the layer's index.
679 + layer_id = int(prefix.split(sep=".")[-1])
680 + self.layer_id = layer_id
681 +
682 + # Complete the preceding dense MLP's deferred all-reduce
683 + # (reduce_results=False), fused into this layer's input_layernorm.
684 + # Disable this fusion when PP is set
685 + self.fuse_input_allreduce = (
686 + layer_id > 0
687 + and not _is_moe_layer(config, layer_id - 1)
688 + and vllm_config.parallel_config.pipeline_parallel_size == 1
689 + )
690 +
691 + is_sparse_attention_layer = (
692 + force_sparse_attn or layer_id in _sparse_attention_layer_ids(config)
693 + )
694 +
695 + if is_sparse_attention_layer:
696 + self.self_attn = MiniMaxM3SparseAttention(
697 + config=config,
698 + layer_id=layer_id,
699 + quant_config=quant_config,
700 + prefix=f"{prefix}.self_attn",
701 + cache_config=cache_config,
702 + )
703 + else:
704 + self.self_attn = MiniMaxM3Attention(
705 + config=config,
706 + layer_id=layer_id,
707 + quant_config=quant_config,
708 + prefix=f"{prefix}.self_attn",
709 + cache_config=cache_config,
710 + )
711 +
712 + # Dense layers store the FFN under `mlp`; MoE layers under
713 + # `block_sparse_moe` -- matching the checkpoint's naming.
714 + self.is_moe_layer = force_moe or _is_moe_layer(config, layer_id)
715 + if self.is_moe_layer:
716 + self.block_sparse_moe = MiniMaxM3MoE(
717 + config=config,
718 + layer_id=layer_id,
719 + quant_config=quant_config,
720 + prefix=f"{prefix}.block_sparse_moe",
721 + )
722 + else:
723 + self.mlp = MiniMaxM3MLP(
724 + config=config,
725 + intermediate_size=config.dense_intermediate_size,
726 + quant_config=quant_config,
727 + prefix=f"{prefix}.mlp",
728 + reduce_results=vllm_config.parallel_config.pipeline_parallel_size > 1,
729 + )
730 +
731 + # config.use_gemma_norm is True for M3 -> Gemma-style RMSNorm.
732 + self.input_layernorm = MiniMAXGemmaRMSNorm(
733 + config.hidden_size, eps=config.rms_norm_eps
734 + )
735 + self.post_attention_layernorm = MiniMAXGemmaRMSNorm(
736 + config.hidden_size, eps=config.rms_norm_eps
737 + )
738 +
739 + def forward(
740 + self,
741 + positions: torch.Tensor,
742 + hidden_states: torch.Tensor,
743 + residual: torch.Tensor | None,
744 + ) -> tuple[torch.Tensor, torch.Tensor]:
745 + if self.fuse_input_allreduce and residual is not None:
746 + hidden_states, residual = fused_allreduce_gemma_rms_norm(
747 + hidden_states, residual, self.input_layernorm
748 + )
749 + else:
750 + if residual is None:
751 + residual = hidden_states
752 + hidden_states = self.input_layernorm(hidden_states)
753 + else:
754 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
755 + hidden_states = self.self_attn(
756 + positions=positions,
757 + hidden_states=hidden_states,
758 + )
759 +
760 + hidden_states, residual = fused_allreduce_gemma_rms_norm(
761 + hidden_states, residual, self.post_attention_layernorm
762 + )
763 + ffn = self.block_sparse_moe if self.is_moe_layer else self.mlp
764 + hidden_states = ffn(hidden_states)
765 + return hidden_states, residual
766 +
767 +
768 +class MiniMaxM3Model(nn.Module, EagleModelMixin):
769 + fall_back_to_pt_during_load = False
770 +
771 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
772 + super().__init__()
773 +
774 + config = vllm_config.model_config.hf_text_config
775 + quant_config = vllm_config.quant_config
776 + self.config = config
777 +
778 + self.vocab_size = config.vocab_size
779 +
780 + self.embed_tokens = VocabParallelEmbedding(
781 + config.vocab_size,
782 + config.hidden_size,
783 + quant_config=quant_config,
784 + prefix=f"{prefix}.embed_tokens",
785 + )
786 +
787 + self.start_layer, self.end_layer, self.layers = make_layers(
788 + config.num_hidden_layers,
789 + lambda prefix: MiniMaxM3DecoderLayer(
790 + vllm_config=vllm_config,
791 + prefix=prefix,
792 + ),
793 + prefix=f"{prefix}.layers",
794 + )
795 +
796 + self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
797 +
798 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
799 + return self.embed_tokens(input_ids)
800 +
801 + def forward(
802 + self,
803 + input_ids: torch.Tensor | None,
804 + positions: torch.Tensor,
805 + inputs_embeds: torch.Tensor | None = None,
806 + ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
807 + if inputs_embeds is not None:
808 + hidden_states = inputs_embeds
809 + else:
810 + hidden_states = self.embed_input_ids(input_ids)
811 + residual = None
812 +
813 + # EAGLE3 is not yet compatible with pipeline parallel
814 + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
815 + for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]):
816 + hidden_states, residual = layer(positions, hidden_states, residual)
817 + self._maybe_add_hidden_state(
818 + aux_hidden_states, idx + 1, hidden_states, residual
819 + )
820 +
821 + hidden_states, _ = self.norm(hidden_states, residual)
822 +
823 + if len(aux_hidden_states) > 0:
824 + return hidden_states, aux_hidden_states
825 + return hidden_states
826 +
827 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
828 + # Checkpoint experts use w1=gate, w2=down, w3=up.
829 + return fused_moe_make_expert_params_mapping(
830 + self,
831 + ckpt_gate_proj_name="w1",
832 + ckpt_down_proj_name="w2",
833 + ckpt_up_proj_name="w3",
834 + num_experts=self.config.num_local_experts,
835 + )
836 +
837 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
838 + # q/k/v_proj -> fused qkv_proj; gate_proj/up_proj -> fused gate_up_proj
839 + # (dense MLP and shared expert). On sparse layers the indexer
840 + # index_q/index_k_proj fold into the same fused qkv_proj
841 + # (MinimaxM3QKVParallelLinearWithIndexer); these entries simply never match on
842 + # dense layers, whose checkpoints have no index_*_proj weights. Leading
843 + # dots keep `q_proj`/`k_proj` from matching `index_q_proj`/`index_k_proj`
844 + # (preceded by `_`, not `.`).
845 + stacked_params_mapping: list[tuple[str, str, int | str]] = [
846 + # (param_name, shard_name, shard_id)
847 + (".qkv_proj", ".q_proj", "q"),
848 + (".qkv_proj", ".k_proj", "k"),
849 + (".qkv_proj", ".v_proj", "v"),
850 + (".qkv_proj", ".index_q_proj", "index_q"),
851 + (".qkv_proj", ".index_k_proj", "index_k"),
852 + (".gate_up_proj", ".gate_proj", 0),
853 + (".gate_up_proj", ".up_proj", 1),
854 + ]
855 +
856 + # (param_name, weight_name, expert_id, shard_id)
857 + expert_params_mapping = self.get_expert_mapping()
858 +
859 + params_dict = dict(self.named_parameters())
860 + loaded_params: set[str] = set()
861 + for name, loaded_weight in weights:
862 + # The MTP module is not modeled yet.
863 + if "mtp." in name:
864 + continue
865 +
866 + # The checkpoint stores block scales as ``weight_scale_inv``; the
867 + # ModelOpt MXFP8 layers expose them as ``weight_scale``.
868 + if "weight_scale_inv" in name:
869 + name = name.replace("weight_scale_inv", "weight_scale")
870 +
871 + for param_name, weight_name, shard_id in stacked_params_mapping:
872 + if weight_name not in name:
873 + continue
874 + # Routed experts (w1/w2/w3) are handled below; don't let the
875 + # stacked mapping rewrite them.
876 + if ("block_sparse_moe.experts." in name) and name not in params_dict:
877 + continue
878 + name = name.replace(weight_name, param_name)
879 + if name.endswith(".bias") and name not in params_dict:
880 + continue
881 + if is_pp_missing_parameter(name, self):
882 + continue
883 + if name not in params_dict:
884 + continue
885 + param = params_dict[name]
886 + weight_loader = param.weight_loader
887 + weight_loader(param, loaded_weight, shard_id)
888 + break
889 + else:
890 + for (
891 + param_name,
892 + weight_name,
893 + expert_id,
894 + expert_shard_id,
895 + ) in expert_params_mapping:
896 + if weight_name not in name:
897 + continue
898 + name = name.replace(weight_name, param_name)
899 + if is_pp_missing_parameter(name, self):
900 + continue
901 + if name not in params_dict:
902 + continue
903 + param = params_dict[name]
904 + weight_loader = param.weight_loader
905 + weight_loader(
906 + param,
907 + loaded_weight,
908 + name,
909 + shard_id=expert_shard_id,
910 + expert_id=expert_id,
911 + )
912 + break
913 + else:
914 + if name.endswith(".bias") and name not in params_dict:
915 + continue
916 + remapped = maybe_remap_kv_scale_name(name, params_dict)
917 + if remapped is None:
918 + continue
919 + name = remapped
920 + if is_pp_missing_parameter(name, self):
921 + continue
922 + # Modules not modeled yet (e.g. attention) are skipped until
923 + # they are ported.
924 + if name not in params_dict:
925 + continue
926 + param = params_dict[name]
927 + weight_loader = getattr(
928 + param, "weight_loader", default_weight_loader
929 + )
930 + weight_loader(param, loaded_weight)
931 + loaded_params.add(name)
932 + return loaded_params
933 +
934 +
935 +class MiniMaxM3SparseForCausalLM(nn.Module, SupportsEagle3):
936 + """MiniMax M3 (sparse/dense backbone) for causal language modeling."""
937 +
938 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
939 + super().__init__()
940 + config = vllm_config.model_config.hf_text_config
941 + quant_config = vllm_config.quant_config
942 + self.config = config
943 + self.quant_config = quant_config
944 + self.model = MiniMaxM3Model(
945 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
946 + )
947 + self.lm_head = ParallelLMHead(
948 + config.vocab_size,
949 + config.hidden_size,
950 + quant_config=quant_config,
951 + prefix=maybe_prefix(prefix, "lm_head"),
952 + )
953 + self.logits_processor = LogitsProcessor(config.vocab_size)
954 +
955 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
956 + return self.model.embed_input_ids(input_ids)
957 +
958 + def forward(
959 + self,
960 + input_ids: torch.Tensor | None,
961 + positions: torch.Tensor,
962 + inputs_embeds: torch.Tensor | None = None,
963 + **kwargs,
964 + ) -> torch.Tensor:
965 + return self.model(input_ids, positions, inputs_embeds)
966 +
967 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
968 + return self.logits_processor(self.lm_head, hidden_states)
969 +
970 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
971 + return self.model.get_expert_mapping()
972 +
973 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
974 + loader = AutoWeightsLoader(self)
975 + return loader.load_weights(weights)
976 +
977 +
978 +@MULTIMODAL_REGISTRY.register_processor(
979 + MiniMaxM3VLMultiModalProcessor,
980 + info=MiniMaxM3VLProcessingInfo,
981 + dummy_inputs=MiniMaxM3VLDummyInputsBuilder,
982 +)
983 +class MiniMaxM3SparseForConditionalGeneration(
984 + nn.Module, SupportsMultiModal, SupportsEagle3
985 +):
986 + """Top-level (VL) entry point for MiniMax M3.
987 +
988 + The vision tower is not modeled yet; this wrapper routes the text
989 + backbone by constructing ``MiniMaxM3SparseForCausalLM`` from the nested
990 + ``text_config`` and delegating generation to it.
991 + """
992 +
993 + # The vision tower runs replicated per rank under ``--mm-encoder-tp-mode
994 + # data``; ``run_dp_sharded_mrope_vision_model`` shards the work across
995 + # ranks (see ``_process_image_input`` / ``_process_video_input``).
996 + supports_encoder_tp_data = True
997 +
998 + hf_to_vllm_mapper = WeightsMapper(
999 + orig_to_new_prefix={
1000 + "multi_modal_projector.": "vision_tower.multi_modal_projector.",
1001 + "patch_merge_mlp.": "vision_tower.patch_merge_mlp.",
1002 + },
1003 + orig_to_new_substr={
1004 + ".mlp.fc1.": ".fc1.",
1005 + ".mlp.fc2.": ".fc2.",
1006 + },
1007 + )
1008 +
1009 + @classmethod
1010 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
1011 + if modality == "image":
1012 + return MiniMaxM3VLProcessingInfo.IMAGE_TOKEN
1013 + if modality == "video":
1014 + return MiniMaxM3VLProcessingInfo.VIDEO_TOKEN
1015 + raise ValueError(f"Unsupported modality: {modality!r}")
1016 +
1017 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1018 + super().__init__()
1019 + config = vllm_config.model_config.hf_config
1020 + self.config = config
1021 + self.quant_config = vllm_config.quant_config
1022 + self.multimodal_config = vllm_config.model_config.multimodal_config
1023 + assert self.multimodal_config is not None
1024 + self.use_data_parallel = self.multimodal_config.mm_encoder_tp_mode == "data"
1025 +
1026 + text_hidden_size = getattr(config.text_config, "hidden_size", None)
1027 + assert text_hidden_size is not None, "text_config.hidden_size is required"
1028 + projector_hidden_size = getattr(config, "projector_hidden_size", None)
1029 +
1030 + with self._mark_tower_model(vllm_config, {"image", "video"}):
1031 + vision_config = config.vision_config
1032 + self.vision_tower = MiniMaxVLVisionModel(
1033 + config=PretrainedConfig.from_dict(vision_config),
1034 + text_hidden_size=text_hidden_size,
1035 + projector_hidden_size=projector_hidden_size,
1036 + quant_config=self.quant_config,
1037 + prefix=maybe_prefix(prefix, "vision_tower"),
1038 + )
1039 +
1040 + self.language_model = init_vllm_registered_model(
1041 + vllm_config=vllm_config,
1042 + hf_config=config.text_config,
1043 + prefix=maybe_prefix(prefix, "language_model"),
1044 + architectures=["MiniMaxM3SparseForCausalLM"],
1045 + )
1046 +
1047 + # Expose language model / lm_head for EAGLE3 spec decode.
1048 + @property
1049 + def model(self) -> nn.Module:
1050 + return self.language_model.model
1051 +
1052 + @property
1053 + def lm_head(self) -> nn.Module:
1054 + return self.language_model.lm_head
1055 +
1056 + def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None:
1057 + pixel_values = kwargs.pop("pixel_values", None)
1058 + image_grid_thw = kwargs.pop("image_grid_thw", None)
1059 + if pixel_values is None:
1060 + return None
1061 + return {"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}
1062 +
1063 + def _parse_and_validate_video_input(self, **kwargs: object) -> dict | None:
1064 + pixel_values_videos = kwargs.pop("pixel_values_videos", None)
1065 + video_grid_thw = kwargs.pop("video_grid_thw", None)
1066 + if pixel_values_videos is None:
1067 + return None
1068 + return {
1069 + "pixel_values_videos": pixel_values_videos,
1070 + "video_grid_thw": video_grid_thw,
1071 + }
1072 +
1073 + def _process_image_input(self, image_input: dict) -> tuple[torch.Tensor, ...]:
1074 + pixel_values: torch.Tensor = image_input["pixel_values"].type(
1075 + self.vision_tower.dtype
1076 + )
1077 + grid_thw: torch.Tensor = image_input["image_grid_thw"]
1078 + assert grid_thw.ndim == 2
1079 +
1080 + if self.use_data_parallel:
1081 + # Already returns a per-item tuple of embeddings.
1082 + return run_dp_sharded_mrope_vision_model(
1083 + self.vision_tower,
1084 + pixel_values,
1085 + grid_thw.tolist(),
1086 + rope_type="rope_3d",
1087 + )
1088 +
1089 + image_embeds = self.vision_tower(
1090 + pixel_values=pixel_values,
1091 + grid_thw=grid_thw.tolist(),
1092 + )
1093 +
1094 + # Split the concatenated output into one tensor per image item.
1095 + merge_size = self.vision_tower.spatial_merge_size
1096 + sizes = (grid_thw.prod(-1) // (merge_size * merge_size)).tolist()
1097 + return image_embeds.split(sizes)
1098 +
1099 + def _process_video_input(self, video_input: dict) -> tuple[torch.Tensor, ...]:
1100 + pixel_values: torch.Tensor = video_input["pixel_values_videos"].type(
1101 + self.vision_tower.dtype
1102 + )
1103 + grid_thw: torch.Tensor = video_input["video_grid_thw"]
1104 + assert grid_thw.ndim == 2
1105 +
1106 + if self.use_data_parallel:
1107 + # Already returns a per-item tuple of embeddings.
1108 + return run_dp_sharded_mrope_vision_model(
1109 + self.vision_tower,
1110 + pixel_values,
1111 + grid_thw.tolist(),
1112 + rope_type="rope_3d",
1113 + )
1114 +
1115 + video_embeds = self.vision_tower(
1116 + pixel_values=pixel_values,
1117 + grid_thw=grid_thw.tolist(),
1118 + )
1119 +
1120 + # Split the concatenated output into one tensor per video item.
1121 + merge_size = self.vision_tower.spatial_merge_size
1122 + sizes = (grid_thw.prod(-1) // (merge_size * merge_size)).tolist()
1123 + return video_embeds.split(sizes)
1124 +
1125 + def _parse_and_validate_multimodal_inputs(
1126 + self, **kwargs: object
1127 + ) -> dict[str, dict]:
1128 + mm_input_by_modality: dict[str, dict] = {}
1129 + for input_key in kwargs:
1130 + if input_key == "pixel_values" and "image" not in mm_input_by_modality:
1131 + image_input = self._parse_and_validate_image_input(**kwargs)
1132 + if image_input is not None:
1133 + mm_input_by_modality["image"] = image_input
1134 + if (
1135 + input_key == "pixel_values_videos"
1136 + and "video" not in mm_input_by_modality
1137 + ):
1138 + video_input = self._parse_and_validate_video_input(**kwargs)
1139 + if video_input is not None:
1140 + mm_input_by_modality["video"] = video_input
1141 + return mm_input_by_modality
1142 +
1143 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
1144 + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
1145 + if not mm_input_by_modality:
1146 + return []
1147 +
1148 + multimodal_embeddings: list[torch.Tensor] = []
1149 + for modality in mm_input_by_modality:
1150 + multimodal_input = mm_input_by_modality[modality]
1151 + if modality == "image":
1152 + image_embeddings = self._process_image_input(multimodal_input)
1153 + multimodal_embeddings.extend(image_embeddings)
1154 + if modality == "video":
1155 + video_embeddings = self._process_video_input(multimodal_input)
1156 + multimodal_embeddings.extend(video_embeddings)
1157 +
1158 + return tuple(multimodal_embeddings)
1159 +
1160 + def forward(
1161 + self,
1162 + input_ids: torch.Tensor | None,
1163 + positions: torch.Tensor,
1164 + inputs_embeds: torch.Tensor | None = None,
1165 + **kwargs,
1166 + ) -> torch.Tensor:
1167 + return self.language_model(input_ids, positions, inputs_embeds)
1168 +
1169 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
1170 + return self.language_model.compute_logits(hidden_states)
1171 +
1172 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1173 + return self.language_model.get_expert_mapping()
1174 +
1175 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1176 + loader = AutoWeightsLoader(self)
1177 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/models/minimax_m3/nvidia/mtp.py new
+312
@@ -0,0 +1,312 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Iterable
5 +
6 +import regex as re
7 +import torch
8 +import torch.nn as nn
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.model_executor.layers.fused_moe import (
12 + fused_moe_make_expert_params_mapping,
13 +)
14 +from vllm.model_executor.layers.linear import (
15 + ReplicatedLinear,
16 +)
17 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
18 +from vllm.model_executor.layers.vocab_parallel_embedding import (
19 + ParallelLMHead,
20 + VocabParallelEmbedding,
21 +)
22 +from vllm.model_executor.model_loader.weight_utils import (
23 + default_weight_loader,
24 + maybe_remap_kv_scale_name,
25 +)
26 +from vllm.model_executor.models.utils import (
27 + maybe_prefix,
28 +)
29 +from vllm.sequence import IntermediateTensors
30 +
31 +from .model import (
32 + MiniMAXGemmaRMSNorm,
33 + MiniMaxM3DecoderLayer,
34 +)
35 +
36 +
37 +class MiniMaxM3MultiTokenPredictorLayer(nn.Module):
38 + def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
39 + super().__init__()
40 +
41 + assert vllm_config.speculative_config is not None
42 + config = vllm_config.speculative_config.draft_model_config.hf_config
43 + quant_config = vllm_config.quant_config
44 +
45 + self.enorm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
46 + self.hnorm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
47 + self.eh_proj = ReplicatedLinear(
48 + config.hidden_size * 2,
49 + config.hidden_size,
50 + bias=False,
51 + quant_config=quant_config,
52 + prefix=f"{prefix}.eh_proj",
53 + )
54 + self.transformer_layer = MiniMaxM3DecoderLayer(
55 + vllm_config=vllm_config,
56 + prefix=prefix,
57 + force_sparse_attn=True,
58 + force_moe=True,
59 + is_mtp_block=True,
60 + )
61 + self.final_layernorm = MiniMAXGemmaRMSNorm(
62 + config.hidden_size, eps=config.rms_norm_eps
63 + )
64 +
65 + def forward(
66 + self,
67 + input_ids: torch.Tensor,
68 + positions: torch.Tensor,
69 + previous_hidden_states: torch.Tensor,
70 + inputs_embeds: torch.Tensor | None = None,
71 + spec_step_index: int = 0,
72 + ) -> torch.Tensor:
73 + assert inputs_embeds is not None
74 + # Mask out inputs at position 0, as not needed by MTP.
75 + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
76 +
77 + # Combine the normalized token embeddings with the normalized
78 + # previous hidden states.
79 + inputs_embeds = self.enorm(inputs_embeds)
80 + previous_hidden_states = self.hnorm(previous_hidden_states)
81 + hidden_states, _ = self.eh_proj(
82 + torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
83 + )
84 +
85 + # Apply transformer layer.
86 + hidden_states, residual = self.transformer_layer(
87 + positions=positions,
88 + hidden_states=hidden_states,
89 + residual=None,
90 + )
91 +
92 + hidden_states += residual
93 + return hidden_states
94 +
95 +
96 +class MiniMaxM3MultiTokenPredictor(nn.Module):
97 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
98 + super().__init__()
99 +
100 + assert vllm_config.speculative_config is not None
101 + # Use the draft (MTP) config, not the target model's. This is flat for a
102 + # standalone checkpoint, and the promoted text_config for a bundled one.
103 + config = vllm_config.speculative_config.draft_model_config.hf_config
104 + self.num_mtp_layers = config.num_mtp_modules
105 + self.layers = torch.nn.ModuleDict(
106 + {
107 + str(idx): MiniMaxM3MultiTokenPredictorLayer(
108 + vllm_config, f"{prefix}.layers.{idx}"
109 + )
110 + for idx in range(self.num_mtp_layers)
111 + }
112 + )
113 + self.embed_tokens = VocabParallelEmbedding(
114 + config.vocab_size,
115 + config.hidden_size,
116 + prefix=maybe_prefix(prefix, "embed_tokens"),
117 + )
118 +
119 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
120 + return self.embed_tokens(input_ids)
121 +
122 + def forward(
123 + self,
124 + input_ids: torch.Tensor,
125 + positions: torch.Tensor,
126 + previous_hidden_states: torch.Tensor,
127 + inputs_embeds: torch.Tensor | None = None,
128 + spec_step_idx: int = 0,
129 + ) -> torch.Tensor:
130 + if inputs_embeds is None:
131 + inputs_embeds = self.embed_tokens(input_ids)
132 + current_step_idx = spec_step_idx % self.num_mtp_layers
133 + return self.layers[str(current_step_idx)](
134 + input_ids,
135 + positions,
136 + previous_hidden_states,
137 + inputs_embeds,
138 + current_step_idx,
139 + )
140 +
141 +
142 +class MiniMaxM3MTP(nn.Module):
143 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
144 + super().__init__()
145 +
146 + assert vllm_config.speculative_config is not None
147 + self.config = vllm_config.speculative_config.draft_model_config.hf_config
148 + self.quant_config = vllm_config.quant_config
149 + self.model = MiniMaxM3MultiTokenPredictor(
150 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
151 + )
152 + self.lm_head = ParallelLMHead(
153 + self.config.vocab_size,
154 + self.config.hidden_size,
155 + quant_config=self.quant_config,
156 + prefix=maybe_prefix(prefix, "lm_head"),
157 + )
158 + self.logits_processor = LogitsProcessor(self.config.vocab_size)
159 +
160 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
161 + return self.model.embed_input_ids(input_ids)
162 +
163 + def forward(
164 + self,
165 + input_ids: torch.Tensor | None,
166 + positions: torch.Tensor,
167 + hidden_states: torch.Tensor,
168 + intermediate_tensors: IntermediateTensors | None = None,
169 + inputs_embeds: torch.Tensor | None = None,
170 + spec_step_idx: int = 0,
171 + ) -> torch.Tensor:
172 + return self.model(
173 + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
174 + )
175 +
176 + def compute_logits(
177 + self,
178 + hidden_states: torch.Tensor,
179 + spec_step_idx: int = 0,
180 + ) -> torch.Tensor | None:
181 + current_step_idx = spec_step_idx % self.model.num_mtp_layers
182 + mtp_layer = self.model.layers[str(current_step_idx)]
183 + return self.logits_processor(
184 + self.lm_head, mtp_layer.final_layernorm(hidden_states)
185 + )
186 +
187 + def _get_mtp_layer_idx_from_weight_name(self, name: str) -> int | None:
188 + """Return the MTP layer index in *.mtp.layers.{idx}.*, else None."""
189 + match = re.search(r"\.mtp\.layers\.(\d+)\.", name)
190 + return int(match.group(1)) if match else None
191 +
192 + def _map_checkpoint_name(self, name: str) -> str | None:
193 + """Map a full checkpoint key to this MTP module's parameter name.
194 +
195 + The MTP module only owns the *.mtp.layers.* weights plus the token
196 + embedding and LM head, which the checkpoint shares with the main model.
197 + Everything else belongs to other modules and is ignored here by returning
198 + None.
199 + """
200 + # In the bundled checkpoint, the MTP weights are prefixed with
201 + # "language_model". The standalone MTP checkpoint has no such prefix.
202 + # Strip it if present.
203 + name = name.removeprefix("language_model.")
204 +
205 + if name == "model.embed_tokens.weight":
206 + return "model.embed_tokens.weight"
207 + if name == "lm_head.weight":
208 + return "lm_head.weight"
209 + if "model.mtp.layers" in name:
210 + if "weight_scale_inv" in name:
211 + # The checkpoint stores block scales as "weight_scale_inv".
212 + # The ModelOpt MXFP8 layers expose them as "weight_scale".
213 + name = name.replace("weight_scale_inv", "weight_scale")
214 + # Strip "mtp" from prefix.
215 + return name.replace(".mtp.", ".")
216 + return None
217 +
218 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
219 + # Map q/k/v projections to qkv_proj, and gate/up projections to gate_up_proj.
220 + stacked_params_mapping: list[tuple[str, str, int | str]] = [
221 + # (param_name, shard_name, shard_id)
222 + (".qkv_proj", ".q_proj", "q"),
223 + (".qkv_proj", ".k_proj", "k"),
224 + (".qkv_proj", ".v_proj", "v"),
225 + (".qkv_proj", ".index_q_proj", "index_q"),
226 + (".qkv_proj", ".index_k_proj", "index_k"),
227 + (".gate_up_proj", ".gate_proj", 0),
228 + (".gate_up_proj", ".up_proj", 1),
229 + ]
230 +
231 + # Map expert weights w1/w2/w3 to gate/down/up.
232 + # (param_name, weight_name, expert_id, shard_id)
233 + expert_params_mapping = fused_moe_make_expert_params_mapping(
234 + self,
235 + ckpt_gate_proj_name="w1",
236 + ckpt_down_proj_name="w2",
237 + ckpt_up_proj_name="w3",
238 + num_experts=self.config.num_local_experts,
239 + )
240 +
241 + params_dict = dict(self.named_parameters())
242 + loaded_params: set[str] = set()
243 + loaded_mtp_layers: set[int] = set()
244 + for name, loaded_weight in weights:
245 + mtp_layer = self._get_mtp_layer_idx_from_weight_name(name)
246 + mapped_name = self._map_checkpoint_name(name)
247 + if mapped_name is None:
248 + # This weight does not belong to the MTP module, so skip it.
249 + continue
250 + name = mapped_name
251 +
252 + for param_name, weight_name, shard_id in stacked_params_mapping:
253 + if weight_name not in name:
254 + continue
255 +
256 + # Routed experts (w1/w2/w3) are handled below. Don't let the
257 + # stacked mapping rewrite them.
258 + if ("block_sparse_moe.experts." in name) and name not in params_dict:
259 + continue
260 + name = name.replace(weight_name, param_name)
261 + if name not in params_dict:
262 + continue
263 + param = params_dict[name]
264 + weight_loader = param.weight_loader
265 + weight_loader(param, loaded_weight, shard_id)
266 + break
267 + else:
268 + for (
269 + param_name,
270 + weight_name,
271 + expert_id,
272 + expert_shard_id,
273 + ) in expert_params_mapping:
274 + if weight_name not in name:
275 + continue
276 +
277 + name = name.replace(weight_name, param_name)
278 + if name not in params_dict:
279 + continue
280 + param = params_dict[name]
281 + weight_loader = param.weight_loader
282 + weight_loader(
283 + param,
284 + loaded_weight,
285 + name,
286 + shard_id=expert_shard_id,
287 + expert_id=expert_id,
288 + )
289 + break
290 + else:
291 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)
292 + if remapped_name is None or remapped_name not in params_dict:
293 + continue
294 + name = remapped_name
295 + param = params_dict[name]
296 + weight_loader = getattr(
297 + param, "weight_loader", default_weight_loader
298 + )
299 + weight_loader(param, loaded_weight)
300 +
301 + loaded_params.add(name)
302 + if mtp_layer is not None:
303 + loaded_mtp_layers.add(mtp_layer)
304 +
305 + # Validate that weights were loaded for each MTP layer.
306 + for layer_idx in range(self.model.num_mtp_layers):
307 + if layer_idx not in loaded_mtp_layers:
308 + raise ValueError(
309 + f"Failed to load MTP layer {layer_idx} weights from checkpoint."
310 + )
311 +
312 + return loaded_params
vllm/models/minimax_m3/nvidia/sparse_attention_msa.py new
+110
@@ -0,0 +1,110 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MSA (SM100/Blackwell) block-sparse attend for MiniMax M3.
4 +
5 +Prefill attends with ``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``);
6 +decode falls back to the Triton split-K kernel (no MSA decode yet). ``fmha_sm100``
7 +imports are function-local, so this module is import-safe on AMD/non-SM100.
8 +"""
9 +
10 +import torch
11 +
12 +from vllm.forward_context import get_forward_context
13 +from vllm.models.minimax_m3.common.ops.sparse_attn import (
14 + SPARSE_BLOCK_SIZE,
15 + minimax_m3_sparse_attn_decode,
16 +)
17 +from vllm.models.minimax_m3.common.sparse_attention import (
18 + MiniMaxM3SparseImpl,
19 + MiniMaxM3SparseMetadata,
20 +)
21 +from vllm.v1.attention.backend import AttentionLayer
22 +
23 +
24 +class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
25 + """MSA block-sparse attend (``fmha_sm100``); Triton split-K decode."""
26 +
27 + def forward(
28 + self,
29 + layer: AttentionLayer,
30 + query: torch.Tensor,
31 + kv_cache: torch.Tensor,
32 + topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
33 + output: torch.Tensor,
34 + ) -> torch.Tensor:
35 + attn_metadata = get_forward_context().attn_metadata
36 + if not isinstance(attn_metadata, dict):
37 + return output # profiling run; caches unbound
38 + main_md = attn_metadata[layer.layer_name] # type: ignore[attr-defined]
39 + assert isinstance(main_md, MiniMaxM3SparseMetadata)
40 + decode_topk, prefill_topk = topk_idx
41 +
42 + nd = main_md.num_decode_tokens
43 + num_tokens = main_md.num_actual_tokens
44 + hd = self.head_size
45 + q = query[:num_tokens].view(-1, self.num_heads, hd)
46 + out = output[:num_tokens].view(-1, self.num_heads, hd)
47 + kv_cache = (
48 + kv_cache.view(self.kv_cache_fp8_dtype) if self.use_fp8_kv else kv_cache
49 + )
50 +
51 + # Decode [:nd]: Triton split-K placeholder (no MSA decode yet).
52 + if main_md.num_decodes > 0:
53 + d = main_md.decode
54 + assert d is not None and decode_topk is not None
55 + minimax_m3_sparse_attn_decode(
56 + q[:nd],
57 + kv_cache,
58 + decode_topk,
59 + d.block_table,
60 + d.seq_lens,
61 + self.num_kv_heads,
62 + self.scale,
63 + out[:nd],
64 + d.decode_query_len,
65 + )
66 +
67 + # Prefill [nd:]: MSA sparse FMHA over the selected blocks.
68 + if main_md.num_prefills > 0:
69 + from vllm.third_party.fmha_sm100.sparse import (
70 + build_k2q_csr,
71 + sparse_atten_func,
72 + )
73 +
74 + p = main_md.prefill
75 + assert p is not None and prefill_topk is not None
76 + qp = q[nd:]
77 + k_cache = kv_cache[:, 0].transpose(1, 2)
78 + v_cache = kv_cache[:, 1].transpose(1, 2)
79 + k2q_row_ptr, k2q_q_indices, schedule = build_k2q_csr(
80 + prefill_topk,
81 + p.cu_seqlens_q,
82 + p.cu_seqlens_k,
83 + SPARSE_BLOCK_SIZE,
84 + total_k=0,
85 + max_seqlen_k=p.max_seq_len,
86 + max_seqlen_q=p.max_query_len,
87 + total_rows=p.total_kv_blocks,
88 + qhead_per_kv=qp.shape[1] // self.num_kv_heads,
89 + return_schedule=True,
90 + )
91 + sparse_atten_func(
92 + qp,
93 + k_cache,
94 + v_cache,
95 + k2q_row_ptr,
96 + k2q_q_indices,
97 + topK=self.topk_blocks,
98 + blk_kv=SPARSE_BLOCK_SIZE,
99 + causal=True,
100 + softmax_scale=self.scale,
101 + cu_seqlens_q=p.cu_seqlens_q,
102 + cu_seqlens_k=p.cu_seqlens_k,
103 + max_seqlen_q=p.max_query_len,
104 + max_seqlen_k=p.max_seq_len,
105 + page_table=p.block_table,
106 + seqused_k=p.seq_lens,
107 + schedule=schedule,
108 + out=out[nd:],
109 + )
110 + return output
vllm/reasoning/__init__.py
+4
@@ -92,6 +92,10 @@ _REASONING_PARSERS_TO_REGISTER = {
92 "minimax_m2_reasoning_parser",
93 "MiniMaxM2AppendThinkReasoningParser",
94 ),
95 + "minimax_m3": (
96 + "minimax_m3_reasoning_parser",
97 + "MiniMaxM3ReasoningParser",
98 + ),
99 "mistral": (
100 "mistral_reasoning_parser",
101 "MistralReasoningParser",
vllm/reasoning/minimax_m3_reasoning_parser.py new
+171
@@ -0,0 +1,171 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Iterable, Sequence
5 +from typing import TYPE_CHECKING
6 +
7 +from vllm.entrypoints.openai.engine.protocol import DeltaMessage
8 +from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
9 +
10 +if TYPE_CHECKING:
11 + from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
12 + from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
13 +
14 +
15 +class MiniMaxM3ReasoningParser(BaseThinkingReasoningParser):
16 + """Reasoning parser for MiniMax M3 explicit thinking blocks.
17 +
18 + MiniMax M3 emits reasoning as:
19 +
20 + <mm:think>reasoning text</mm:think>assistant content
21 +
22 + The M3 tokenizer exposes both markers as complete vocabulary tokens. The
23 + chat template may also prefill the start marker when
24 + ``thinking_mode="enabled"``, so generated text can begin directly inside a
25 + reasoning block without emitting ``<mm:think>`` again.
26 + """
27 +
28 + @property
29 + def start_token(self) -> str:
30 + return "<mm:think>"
31 +
32 + @property
33 + def end_token(self) -> str:
34 + return "</mm:think>"
35 +
36 + def __init__(self, tokenizer, *args, **kwargs):
37 + super().__init__(tokenizer, *args, **kwargs)
38 + chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
39 + self._initial_in_reasoning = chat_kwargs.get("thinking_mode") == "enabled"
40 + self._at_response_start = True
41 +
42 + def extract_reasoning(
43 + self,
44 + model_output: str,
45 + request: "ChatCompletionRequest | ResponsesRequest",
46 + ) -> tuple[str | None, str | None]:
47 + # MiniMax M3 can start a response with a stray closer. Drop that first
48 + # token only; later unmatched closers stay visible as content.
49 + if not self._initial_in_reasoning and model_output.startswith(self.end_token):
50 + content = model_output[len(self.end_token) :]
51 + return None, content or None
52 +
53 + if self._initial_in_reasoning and self.start_token not in model_output:
54 + reasoning, end, content = model_output.partition(self.end_token)
55 + if not end:
56 + return model_output, None
57 + return reasoning, content or None
58 +
59 + if self.start_token not in model_output:
60 + return None, model_output
61 +
62 + content_before, _, after_start = model_output.partition(self.start_token)
63 + reasoning, end, content_after = after_start.partition(self.end_token)
64 + if not end:
65 + return reasoning, content_before or None
66 +
67 + return reasoning, (content_before + content_after) or None
68 +
69 + def is_reasoning_end_streaming(
70 + self, input_ids: Sequence[int], delta_ids: Iterable[int]
71 + ) -> bool:
72 + delta_ids = tuple(delta_ids)
73 + if self.end_token_id in delta_ids:
74 + return True
75 + if self.end_token_id in input_ids:
76 + return True
77 + if self._initial_in_reasoning:
78 + return False
79 + if self.start_token_id not in input_ids:
80 + return bool(input_ids)
81 + return False
82 +
83 + def extract_content_ids(self, input_ids: list[int]) -> list[int]:
84 + if self.end_token_id in input_ids:
85 + end_index = len(input_ids) - 1 - input_ids[::-1].index(self.end_token_id)
86 + return input_ids[end_index + 1 :]
87 +
88 + if self._initial_in_reasoning and self.start_token_id not in input_ids:
89 + return []
90 +
91 + if self.start_token_id not in input_ids:
92 + return input_ids
93 + return []
94 +
95 + def extract_reasoning_streaming(
96 + self,
97 + previous_text: str,
98 + current_text: str,
99 + delta_text: str,
100 + previous_token_ids: Sequence[int],
101 + current_token_ids: Sequence[int],
102 + delta_token_ids: Sequence[int],
103 + ) -> DeltaMessage | None:
104 + if not delta_text:
105 + return None
106 +
107 + if self._at_response_start and not self._initial_in_reasoning:
108 + # Apply the leading-closer tolerance once. Later unmatched closers
109 + # stay visible as content.
110 + self._at_response_start = False
111 + if delta_text.startswith(self.end_token):
112 + delta_text = delta_text[len(self.end_token) :]
113 + if not delta_text:
114 + return None
115 + if delta_token_ids and delta_token_ids[0] == self.end_token_id:
116 + delta_token_ids = delta_token_ids[1:]
117 +
118 + if self.end_token_id in previous_token_ids:
119 + return DeltaMessage(content=delta_text)
120 +
121 + if (
122 + self._initial_in_reasoning
123 + and self.start_token_id not in previous_token_ids
124 + and self.start_token_id not in delta_token_ids
125 + ):
126 + if self.end_token_id in delta_token_ids:
127 + reasoning, _, content = delta_text.partition(self.end_token)
128 + return DeltaMessage(
129 + reasoning=reasoning or None,
130 + content=content or None,
131 + )
132 + return DeltaMessage(reasoning=delta_text)
133 +
134 + if (
135 + self.start_token_id not in previous_token_ids
136 + and self.start_token_id not in delta_token_ids
137 + ):
138 + return DeltaMessage(content=delta_text)
139 +
140 + if self.end_token_id in delta_token_ids:
141 + reasoning_text, _, content = delta_text.partition(self.end_token)
142 + if self.start_token_id in delta_token_ids:
143 + _, _, reasoning_text = reasoning_text.partition(self.start_token)
144 + return DeltaMessage(
145 + reasoning=reasoning_text or None,
146 + content=content or None,
147 + )
148 +
149 + if self.start_token_id in delta_token_ids:
150 + _, _, reasoning = delta_text.partition(self.start_token)
151 + return DeltaMessage(reasoning=reasoning) if reasoning else None
152 +
153 + return DeltaMessage(reasoning=delta_text)
154 +
155 + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
156 + if not self._initial_in_reasoning:
157 + return super().count_reasoning_tokens(token_ids)
158 +
159 + count = 0
160 + depth = 1
161 + for token_id in token_ids:
162 + if token_id == self.start_token_id:
163 + depth += 1
164 + continue
165 + if token_id == self.end_token_id:
166 + if depth > 0:
167 + depth -= 1
168 + continue
169 + if depth > 0:
170 + count += 1
171 + return count
vllm/tool_parsers/__init__.py
+4
@@ -126,6 +126,10 @@ _TOOL_PARSERS_TO_REGISTER = {
126 "minimax_m2_tool_parser",
127 "MinimaxM2ToolParser",
128 ),
129 + "minimax_m3": (
130 + "minimax_m3_tool_parser",
131 + "MinimaxM3ToolParser",
132 + ),
133 "minimax": (
134 "minimax_tool_parser",
135 "MinimaxToolParser",
vllm/tool_parsers/minimax_m3_tool_parser.py new
+19
@@ -0,0 +1,19 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.tool_parsers.rust_tool_parser import RustToolParser
5 +
6 +
7 +class MinimaxM3ToolParser(RustToolParser):
8 + """Adapter from the Rust MiniMax M3 parser to vLLM ToolParser.
9 +
10 + The real M3 grammar lives in the Rust tool-parser crate. This class only
11 + configures the generic Rust bridge with the MiniMax M3 parser name.
12 +
13 + M3 is not M2 with renamed tags: it prefixes each structural tag with the
14 + MiniMax namespace marker, allows multiple ``<invoke>`` tags in one wrapper,
15 + and represents nested arguments with parameter-name XML tags.
16 + """
17 +
18 + rust_parser_name = "MinimaxM3ToolParser"
19 + tool_call_start_token = "]<]minimax[>[<tool_call>"
vllm/transformers_utils/config.py
+2
@@ -103,6 +103,8 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
103 medusa="MedusaConfig",
104 mellum="MellumConfig",
105 midashenglm="MiDashengLMConfig",
106 + minimax_m3_vl="MiniMaxM3Config",
107 + minimax_m3_mtp="MiniMaxM3MTPConfig",
108 moondream3="Moondream3Config",
109 eagle="EAGLEConfig",
110 speculators="SpeculatorsConfig",
vllm/transformers_utils/configs/__init__.py
+6
@@ -53,6 +53,9 @@ _CLASS_TO_MODULE: dict[str, str] = {
53 "MedusaConfig": "vllm.transformers_utils.configs.medusa",
54 "MellumConfig": "vllm.transformers_utils.configs.mellum",
55 "MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm",
56 + "MiniMaxM3Config": "vllm.transformers_utils.configs.minimax_m3",
57 + "MiniMaxM3MTPConfig": "vllm.transformers_utils.configs.minimax_m3",
58 + "MiniMaxM3TextConfig": "vllm.transformers_utils.configs.minimax_m3",
59 "MLPSpeculatorConfig": "vllm.transformers_utils.configs.mlp_speculator",
60 "Moondream3Config": "vllm.transformers_utils.configs.moondream3",
61 "Moondream3TextConfig": "vllm.transformers_utils.configs.moondream3",
@@ -124,6 +127,9 @@ __all__ = [
127 "MedusaConfig",
128 "MellumConfig",
129 "MiDashengLMConfig",
130 + "MiniMaxM3Config",
131 + "MiniMaxM3MTPConfig",
132 + "MiniMaxM3TextConfig",
133 "MLPSpeculatorConfig",
134 "Moondream3Config",
135 "Moondream3TextConfig",
vllm/transformers_utils/configs/minimax_m3.py new
+149
@@ -0,0 +1,149 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any
4 +
5 +from transformers import PretrainedConfig
6 +
7 +
8 +class MiniMaxM3TextConfig(PretrainedConfig):
9 + """Config for the MiniMax M3 text backbone (MiniMaxM3SparseForCausalLM).
10 +
11 + Defaults mirror the ``text_config`` of the MiniMax-M3-preview checkpoint.
12 + """
13 +
14 + model_type = "minimax_m3_text"
15 + architectures = ["MiniMaxM3SparseForCausalLM"]
16 +
17 + def __init__(
18 + self,
19 + vocab_size: int = 200064,
20 + hidden_size: int = 6144,
21 + intermediate_size: int = 3072,
22 + dense_intermediate_size: int = 12288,
23 + shared_intermediate_size: int = 3072,
24 + num_hidden_layers: int = 60,
25 + num_attention_heads: int = 64,
26 + num_key_value_heads: int = 4,
27 + head_dim: int = 128,
28 + max_position_embeddings: int = 524288,
29 + rms_norm_eps: float = 1e-6,
30 + use_gemma_norm: bool = True,
31 + attention_output_gate: bool = False,
32 + rope_theta: float = 5000000,
33 + rotary_dim: int = 64,
34 + partial_rotary_factor: float = 0.5,
35 + hidden_act: str = "swigluoai",
36 + swiglu_alpha: float = 1.702,
37 + # SwiGLU-OAI uses the (up + 1) bias, i.e. beta=1.0 (matches the
38 + # reference: gate * sigmoid(gate * alpha) * (up + 1)). The checkpoint
39 + # config omits swiglu_beta, so this default must stay 1.0.
40 + swiglu_beta: float = 1.0,
41 + swiglu_limit: float = 7.0,
42 + use_qk_norm: bool = True,
43 + qk_norm_type: str = "per_head",
44 + num_local_experts: int = 128,
45 + num_experts_per_tok: int = 4,
46 + n_shared_experts: int = 1,
47 + scoring_func: str = "sigmoid",
48 + use_routing_bias: bool = True,
49 + routed_scaling_factor: float = 2.0,
50 + num_mtp_modules: int = 1,
51 + moe_layer_freq: list[int] | None = None,
52 + sparse_attention_config: dict[str, Any] | None = None,
53 + tie_word_embeddings: bool = False,
54 + **kwargs,
55 + ):
56 + self.vocab_size = vocab_size
57 + self.hidden_size = hidden_size
58 + self.intermediate_size = intermediate_size
59 + self.dense_intermediate_size = dense_intermediate_size
60 + self.shared_intermediate_size = shared_intermediate_size
61 + self.num_hidden_layers = num_hidden_layers
62 + self.num_attention_heads = num_attention_heads
63 + self.num_key_value_heads = num_key_value_heads
64 + self.head_dim = head_dim
65 + self.max_position_embeddings = max_position_embeddings
66 + self.rms_norm_eps = rms_norm_eps
67 + self.use_gemma_norm = use_gemma_norm
68 + self.attention_output_gate = attention_output_gate
69 + self.rope_theta = rope_theta
70 + self.rotary_dim = rotary_dim
71 + self.partial_rotary_factor = partial_rotary_factor
72 + self.hidden_act = hidden_act
73 + self.swiglu_alpha = swiglu_alpha
74 + self.swiglu_beta = swiglu_beta
75 + self.swiglu_limit = swiglu_limit
76 + self.use_qk_norm = use_qk_norm
77 + self.qk_norm_type = qk_norm_type
78 + self.num_local_experts = num_local_experts
79 + self.num_experts_per_tok = num_experts_per_tok
80 + self.n_shared_experts = n_shared_experts
81 + self.scoring_func = scoring_func
82 + self.use_routing_bias = use_routing_bias
83 + self.routed_scaling_factor = routed_scaling_factor
84 + self.num_mtp_modules = num_mtp_modules
85 + # First 3 layers are dense; the remaining 57 are sparse MoE.
86 + self.moe_layer_freq = (
87 + moe_layer_freq if moe_layer_freq is not None else [0] * 3 + [1] * 57
88 + )
89 + self.sparse_attention_config = (
90 + sparse_attention_config
91 + if sparse_attention_config is not None
92 + else {
93 + "use_sparse_attention": True,
94 + "sparse_index_dim": 128,
95 + "sparse_num_index_heads": 4,
96 + "sparse_topk_blocks": 16,
97 + "sparse_block_size": 128,
98 + "sparse_disable_index_value": [0] * 3 + [1] * 57,
99 + "sparse_score_type": "max",
100 + "sparse_init_block": 0,
101 + "sparse_local_block": 1,
102 + "sparse_attention_freq": [0] * 3 + [1] * 57,
103 + }
104 + )
105 + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
106 +
107 +
108 +class MiniMaxM3MTPConfig(MiniMaxM3TextConfig):
109 + """Config for a standalone MiniMax M3 MTP (multi-token prediction) head.
110 +
111 + The MTP transformer layer is structurally a single MiniMax M3 decoder
112 + layer, so this reuses the text backbone schema. Standalone MTP checkpoints
113 + use ``model_type='minimax_m3_mtp'`` and a single hidden layer.
114 + """
115 +
116 + model_type = "minimax_m3_mtp"
117 + architectures = ["MiniMaxM3MTP"]
118 +
119 + def __init__(self, num_hidden_layers: int = 1, **kwargs):
120 + super().__init__(num_hidden_layers=num_hidden_layers, **kwargs)
121 +
122 +
123 +class MiniMaxM3Config(PretrainedConfig):
124 + """Top-level MiniMax M3 (VL) config.
125 +
126 + Holds the text backbone as ``text_config`` so that
127 + ``config.get_text_config()`` extracts the MiniMaxM3SparseForCausalLM
128 + backbone. Vision components are kept as a raw dict passthrough and are
129 + not modeled here.
130 + """
131 +
132 + model_type = "minimax_m3_vl"
133 +
134 + def __init__(
135 + self,
136 + text_config: dict | MiniMaxM3TextConfig | None = None,
137 + vision_config: dict | None = None,
138 + **kwargs,
139 + ):
140 + if text_config is None:
141 + text_config = MiniMaxM3TextConfig()
142 + elif isinstance(text_config, dict):
143 + text_config = MiniMaxM3TextConfig(**text_config)
144 + self.text_config = text_config
145 + self.vision_config = vision_config
146 +
147 + self.hidden_size = text_config.hidden_size
148 +
149 + super().__init__(**kwargs)
vllm/transformers_utils/processors/__init__.py
+6
@@ -31,6 +31,9 @@ __all__ = [
31 "MiMoOmniProcessor",
32 "MiniCPMOProcessor",
33 "MiniCPMVProcessor",
34 + "MiniMaxM3VLImageProcessor",
35 + "MiniMaxM3VLVideoProcessor",
36 + "MiniMaxVLProcessor",
37 "MistralCommonPixtralProcessor",
38 "MistralCommonVoxtralProcessor",
39 "NanoNemotronVLProcessor",
@@ -64,6 +67,9 @@ _CLASS_TO_MODULE: dict[str, str] = {
67 "MiMoOmniProcessor": "vllm.transformers_utils.processors.mimo_v2_omni",
68 "MiniCPMOProcessor": "vllm.transformers_utils.processors.minicpmo",
69 "MiniCPMVProcessor": "vllm.transformers_utils.processors.minicpmv",
70 + "MiniMaxM3VLImageProcessor": "vllm.transformers_utils.processors.minimax_m3",
71 + "MiniMaxM3VLVideoProcessor": "vllm.transformers_utils.processors.minimax_m3",
72 + "MiniMaxVLProcessor": "vllm.transformers_utils.processors.minimax_m3",
73 "MistralCommonPixtralProcessor": "vllm.transformers_utils.processors.pixtral",
74 "MistralCommonVoxtralProcessor": "vllm.transformers_utils.processors.voxtral",
75 "Moondream3Processor": "vllm.transformers_utils.processors.moondream3",
vllm/transformers_utils/processors/minimax_m3.py new
+736
@@ -0,0 +1,736 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MiniMax M3 VL HuggingFace-compatible Processor / ImageProcessor /
4 +VideoProcessor, vendored into vLLM so the model loads without
5 +``--trust-remote-code`` (the released checkpoint only ships these classes as
6 +remote code via ``auto_map``).
7 +
8 +Adapted verbatim from the ``MiniMaxAI/Minimax-M3-preview`` repository files
9 +``image_processor.py``, ``video_processor.py`` and ``processing_minimax.py``
10 +(revision ``db01c0fe``). Both image and video processors use Qwen-style
11 +``smart_resize`` (bound by total pixels). The original async frame-sampling
12 +helpers are intentionally omitted: vLLM performs its own frame loading and
13 +feeds decoded frames to the processor.
14 +"""
15 +
16 +import math
17 +
18 +import regex as re
19 +import torch
20 +from torchvision.transforms import InterpolationMode
21 +from transformers import AutoTokenizer, BatchFeature
22 +from transformers.image_processing_utils_fast import (
23 + BaseImageProcessorFast,
24 + group_images_by_shape,
25 + reorder_images,
26 +)
27 +from transformers.image_utils import PILImageResampling, SizeDict
28 +from transformers.processing_utils import (
29 + ImagesKwargs,
30 + ProcessingKwargs,
31 + ProcessorMixin,
32 + Unpack,
33 + VideosKwargs,
34 +)
35 +from transformers.utils import TensorType
36 +from transformers.video_processing_utils import BaseVideoProcessor
37 +from transformers.video_utils import group_videos_by_shape, reorder_videos
38 +
39 +# Maximum allowed aspect ratio before smart_resize rejects the input.
40 +MAX_RATIO = 200
41 +
42 +# Fixed (non-configurable) bounds for the long-side resize logic, per the
43 +# MiniMax-M3 size spec. ``min_short_side_pixel`` is the floor the short edge is
44 +# enlarged to; ``*_MAX_TOTAL_PIXELS`` is the hard area cap that, once exceeded,
45 +# aborts processing instead of downscaling.
46 +MIN_SHORT_SIDE_PIXEL = 112
47 +IMAGE_MAX_TOTAL_PIXELS = 12_845_056 # 3584 ** 2 (width * height)
48 +VIDEO_MAX_TOTAL_PIXELS = 301_056_000 # width * height * frames
49 +
50 +
51 +def round_by_factor(number: int | float, factor: int) -> int:
52 + return round(number / factor) * factor
53 +
54 +
55 +def ceil_by_factor(number: int | float, factor: int) -> int:
56 + return math.ceil(number / factor) * factor
57 +
58 +
59 +def floor_by_factor(number: int | float, factor: int) -> int:
60 + return math.floor(number / factor) * factor
61 +
62 +
63 +def _smart_resize_by_long_side(
64 + height: int,
65 + width: int,
66 + factor: int,
67 + max_long_side_pixel: int,
68 + min_short_side_pixel: int,
69 + max_total_pixels: int | None,
70 +) -> tuple[int, int]:
71 + """Long-side based resize (MiniMax-M3 size spec).
72 +
73 + (a) if the long side exceeds ``max_long_side_pixel`` → shrink so the long
74 + side equals ``max_long_side_pixel``;
75 + (b) else if the short side is below ``min_short_side_pixel`` → enlarge so the
76 + short side equals ``min_short_side_pixel``;
77 + (c) if the resulting area still exceeds ``max_total_pixels`` → raise.
78 +
79 + (a) and (b) are mutually exclusive (they branch on the *original* long side).
80 + Both sides are then rounded to a multiple of ``factor``. For videos the
81 + ``max_total_pixels`` cap is volumetric (width * height * frames) and is
82 + enforced by the caller, so pass ``max_total_pixels=None`` here.
83 + """
84 + long_side = max(height, width)
85 + short_side = min(height, width)
86 +
87 + scaled_height: float = height
88 + scaled_width: float = width
89 + if long_side > max_long_side_pixel:
90 + beta = max_long_side_pixel / long_side
91 + scaled_height = height * beta
92 + scaled_width = width * beta
93 + elif short_side < min_short_side_pixel:
94 + beta = min_short_side_pixel / short_side
95 + scaled_height = height * beta
96 + scaled_width = width * beta
97 +
98 + h_bar = max(factor, round_by_factor(scaled_height, factor))
99 + w_bar = max(factor, round_by_factor(scaled_width, factor))
100 +
101 + if max_total_pixels is not None and h_bar * w_bar > max_total_pixels:
102 + raise ValueError(
103 + f"image area {h_bar * w_bar} exceeds max_total_pixels "
104 + f"{max_total_pixels} after resizing"
105 + )
106 + return h_bar, w_bar
107 +
108 +
109 +def smart_resize(
110 + height: int,
111 + width: int,
112 + factor: int = 28,
113 + min_pixels: int = 4 * 28 * 28,
114 + max_pixels: int = 451584,
115 + max_long_side_pixel: int | None = None,
116 + min_short_side_pixel: int = MIN_SHORT_SIDE_PIXEL,
117 + max_total_pixels: int | None = None,
118 +) -> tuple[int, int]:
119 + """Rescale (height, width) so each side is a multiple of ``factor``.
120 +
121 + When ``max_long_side_pixel`` is set, use the MiniMax-M3 long-side resize
122 + spec (see :func:`_smart_resize_by_long_side`). Otherwise fall back to the
123 + Qwen-VL area bound, keeping the total area within ``[min_pixels, max_pixels]``.
124 + """
125 + if max(height, width) / min(height, width) > MAX_RATIO:
126 + raise ValueError(
127 + f"absolute aspect ratio must be smaller than {MAX_RATIO}, "
128 + f"got {max(height, width) / min(height, width)}"
129 + )
130 + if max_long_side_pixel is not None:
131 + return _smart_resize_by_long_side(
132 + height,
133 + width,
134 + factor=factor,
135 + max_long_side_pixel=max_long_side_pixel,
136 + min_short_side_pixel=min_short_side_pixel,
137 + max_total_pixels=max_total_pixels,
138 + )
139 + h_bar = max(factor, round_by_factor(height, factor))
140 + w_bar = max(factor, round_by_factor(width, factor))
141 + if h_bar * w_bar > max_pixels:
142 + beta = math.sqrt((height * width) / max_pixels)
143 + h_bar = floor_by_factor(height / beta, factor)
144 + w_bar = floor_by_factor(width / beta, factor)
145 + elif h_bar * w_bar < min_pixels:
146 + beta = math.sqrt(min_pixels / (height * width))
147 + h_bar = ceil_by_factor(height * beta, factor)
148 + w_bar = ceil_by_factor(width * beta, factor)
149 + return h_bar, w_bar
150 +
151 +
152 +class MiniMaxM3VLImageProcessorKwargs(ImagesKwargs, total=False): # type: ignore[call-arg]
153 + patch_size: int
154 + temporal_patch_size: int
155 + merge_size: int
156 + max_pixels: int
157 + max_long_side_pixel: int
158 +
159 +
160 +class MiniMaxM3VLImageProcessor(BaseImageProcessorFast):
161 + do_resize = True
162 + resample = PILImageResampling.BICUBIC
163 + # required by base-class validation, not used as the resize bound
164 + size = {"height": 672, "width": 672}
165 + default_to_square = False
166 + do_rescale = True
167 + rescale_factor = 1 / 255
168 + do_normalize = True
169 + image_mean = [0.48145466, 0.4578275, 0.40821073]
170 + image_std = [0.26862954, 0.26130258, 0.27577711]
171 + do_convert_rgb = True
172 + patch_size = 14
173 + temporal_patch_size = 2
174 + merge_size = 2
175 + max_pixels = 451584 # 672 * 672
176 + # Long-side resize spec (opt-in via ``max_long_side_pixel``). The latter two
177 + # are fixed per the spec and are not exposed as configurable kwargs.
178 + max_long_side_pixel = None
179 + min_short_side_pixel = MIN_SHORT_SIDE_PIXEL
180 + max_total_pixels = IMAGE_MAX_TOTAL_PIXELS
181 + valid_kwargs = MiniMaxM3VLImageProcessorKwargs
182 + model_input_names = ["pixel_values", "image_grid_thw"]
183 +
184 + def __init__(self, **kwargs: Unpack[MiniMaxM3VLImageProcessorKwargs]):
185 + super().__init__(**kwargs)
186 +
187 + def preprocess(
188 + self, images, **kwargs: Unpack[MiniMaxM3VLImageProcessorKwargs]
189 + ) -> BatchFeature:
190 + return super().preprocess(images, **kwargs)
191 +
192 + def _preprocess(
193 + self,
194 + images: list[torch.Tensor],
195 + do_resize: bool,
196 + size: SizeDict,
197 + resample: "PILImageResampling | InterpolationMode | int | None",
198 + do_rescale: bool,
199 + rescale_factor: float,
200 + do_normalize: bool,
201 + image_mean: "float | list[float] | None",
202 + image_std: "float | list[float] | None",
203 + patch_size: int,
204 + temporal_patch_size: int,
205 + merge_size: int,
206 + max_pixels: int,
207 + max_long_side_pixel: "int | None",
208 + disable_grouping: "bool | None",
209 + return_tensors: "str | TensorType | None",
210 + **kwargs,
211 + ) -> BatchFeature:
212 + grouped_images, grouped_images_index = group_images_by_shape(
213 + images, disable_grouping=disable_grouping
214 + )
215 + resized_images_grouped = {}
216 + factor = patch_size * merge_size
217 + for shape, stacked_images in grouped_images.items():
218 + height, width = stacked_images.shape[-2:]
219 + if do_resize:
220 + resized_height, resized_width = smart_resize(
221 + height,
222 + width,
223 + factor=factor,
224 + max_pixels=max_pixels,
225 + max_long_side_pixel=max_long_side_pixel,
226 + min_short_side_pixel=self.min_short_side_pixel,
227 + max_total_pixels=self.max_total_pixels,
228 + )
229 + stacked_images = self.resize(
230 + stacked_images,
231 + size=SizeDict(height=resized_height, width=resized_width),
232 + resample=resample,
233 + )
234 + resized_images_grouped[shape] = stacked_images
235 +
236 + resized_images = reorder_images(resized_images_grouped, grouped_images_index)
237 +
238 + grouped_images, grouped_images_index = group_images_by_shape(
239 + resized_images, disable_grouping=disable_grouping
240 + )
241 + processed_images_grouped = {}
242 + processed_grids = {}
243 +
244 + for shape, stacked_images in grouped_images.items():
245 + resized_height, resized_width = stacked_images.shape[-2:]
246 +
247 + patches = self.rescale_and_normalize(
248 + stacked_images,
249 + do_rescale,
250 + rescale_factor,
251 + do_normalize,
252 + image_mean,
253 + image_std,
254 + )
255 + if patches.ndim == 4:
256 + patches = patches.unsqueeze(1)
257 +
258 + if patches.shape[1] % temporal_patch_size != 0:
259 + repeats = patches[:, -1:].repeat(
260 + 1,
261 + temporal_patch_size - (patches.shape[1] % temporal_patch_size),
262 + 1,
263 + 1,
264 + 1,
265 + )
266 + patches = torch.cat([patches, repeats], dim=1)
267 +
268 + batch_size, grid_t, channel = patches.shape[:3]
269 + grid_t = grid_t // temporal_patch_size
270 + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
271 +
272 + patches = patches.view(
273 + batch_size,
274 + grid_t,
275 + temporal_patch_size,
276 + channel,
277 + grid_h // merge_size,
278 + merge_size,
279 + patch_size,
280 + grid_w // merge_size,
281 + merge_size,
282 + patch_size,
283 + )
284 + patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)
285 +
286 + flatten_patches = patches.reshape(
287 + batch_size,
288 + grid_t * grid_h * grid_w,
289 + channel * temporal_patch_size * patch_size * patch_size,
290 + )
291 +
292 + processed_images_grouped[shape] = flatten_patches
293 + processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
294 +
295 + processed_images = reorder_images(
296 + processed_images_grouped, grouped_images_index
297 + )
298 + processed_grids = reorder_images(processed_grids, grouped_images_index)
299 +
300 + pixel_values = torch.cat(processed_images, dim=0)
301 + image_grid_thw = torch.tensor(processed_grids, dtype=torch.long)
302 +
303 + return BatchFeature(
304 + data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw},
305 + tensor_type=return_tensors,
306 + )
307 +
308 + def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None):
309 + images_kwargs = images_kwargs or {}
310 + patch_size = images_kwargs.get("patch_size", self.patch_size)
311 + merge_size = images_kwargs.get("merge_size", self.merge_size)
312 + max_pixels = images_kwargs.get("max_pixels", self.max_pixels)
313 + max_long_side_pixel = images_kwargs.get(
314 + "max_long_side_pixel", self.max_long_side_pixel
315 + )
316 +
317 + resized_height, resized_width = smart_resize(
318 + height,
319 + width,
320 + factor=patch_size * merge_size,
321 + max_pixels=max_pixels,
322 + max_long_side_pixel=max_long_side_pixel,
323 + min_short_side_pixel=self.min_short_side_pixel,
324 + max_total_pixels=self.max_total_pixels,
325 + )
326 + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
327 + return grid_h * grid_w
328 +
329 +
330 +class MiniMaxM3VLVideoProcessorKwargs(VideosKwargs, total=False): # type: ignore[call-arg]
331 + patch_size: int
332 + temporal_patch_size: int
333 + merge_size: int
334 + min_pixels: int
335 + max_pixels: int
336 + max_long_side_pixel: int
337 + total_pixels: int
338 + min_frames: int
339 + max_frames: int
340 + fps: "float | int"
341 +
342 +
343 +class MiniMaxM3VLVideoProcessor(BaseVideoProcessor):
344 + do_resize = True
345 + resample = PILImageResampling.BICUBIC
346 + size = {"height": 672, "width": 672}
347 + default_to_square = False
348 + do_rescale = True
349 + rescale_factor = 1 / 255
350 + do_normalize = True
351 + image_mean = [0.48145466, 0.4578275, 0.40821073]
352 + image_std = [0.26862954, 0.26130258, 0.27577711]
353 + do_convert_rgb = True
354 + do_sample_frames = False
355 + patch_size = 14
356 + temporal_patch_size = 2
357 + merge_size = 2
358 + min_pixels = 4 * 28 * 28
359 + max_pixels = 768 * 28 * 28 # 602,112
360 + total_pixels = int(64000 * 28 * 28 * 0.9) # ~45M, ~64k tokens budget
361 + # Long-side resize spec (opt-in via ``max_long_side_pixel``). The video
362 + # ``max_total_pixels`` cap is volumetric (width * height * frames) and is
363 + # enforced in ``_preprocess`` once the frame count is known.
364 + max_long_side_pixel = None
365 + min_short_side_pixel = MIN_SHORT_SIDE_PIXEL
366 + max_total_pixels = VIDEO_MAX_TOTAL_PIXELS
367 + fps = 1.0
368 + min_frames = 4
369 + max_frames = 768
370 + valid_kwargs = MiniMaxM3VLVideoProcessorKwargs
371 + model_input_names = ["pixel_values_videos", "video_grid_thw"]
372 +
373 + def __init__(self, **kwargs: Unpack[MiniMaxM3VLVideoProcessorKwargs]):
374 + super().__init__(**kwargs)
375 +
376 + def _preprocess(
377 + self,
378 + videos: list[torch.Tensor],
379 + do_convert_rgb: bool,
380 + do_resize: bool,
381 + size: SizeDict,
382 + resample: "PILImageResampling | InterpolationMode | int | None",
383 + do_rescale: bool,
384 + rescale_factor: float,
385 + do_normalize: bool,
386 + image_mean: "float | list[float] | None",
387 + image_std: "float | list[float] | None",
388 + patch_size: int,
389 + temporal_patch_size: int,
390 + merge_size: int,
391 + min_pixels: int,
392 + max_pixels: int,
393 + max_long_side_pixel: "int | None" = None,
394 + return_tensors: "str | TensorType | None" = None,
395 + **kwargs,
396 + ) -> BatchFeature:
397 + grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
398 + resized_videos_grouped = {}
399 + factor = patch_size * merge_size
400 + for shape, stacked_videos in grouped_videos.items():
401 + batch_size, num_frames, channels, height, width = stacked_videos.shape
402 + resized_height, resized_width = height, width
403 + if do_resize:
404 + resized_height, resized_width = smart_resize(
405 + height,
406 + width,
407 + factor=factor,
408 + min_pixels=min_pixels,
409 + max_pixels=max_pixels,
410 + max_long_side_pixel=max_long_side_pixel,
411 + min_short_side_pixel=self.min_short_side_pixel,
412 + # Per-frame raise disabled; the video cap is volumetric and
413 + # is enforced below once num_frames is known.
414 + max_total_pixels=None,
415 + )
416 + if (
417 + max_long_side_pixel is not None
418 + and resized_height * resized_width * num_frames
419 + > self.max_total_pixels
420 + ):
421 + raise ValueError(
422 + f"video area {resized_height * resized_width * num_frames} "
423 + f"(width * height * frames) exceeds max_total_pixels "
424 + f"{self.max_total_pixels} after resizing"
425 + )
426 + stacked_videos = stacked_videos.view(
427 + batch_size * num_frames, channels, height, width
428 + )
429 + stacked_videos = self.resize(
430 + stacked_videos,
431 + size=SizeDict(height=resized_height, width=resized_width),
432 + resample=resample,
433 + )
434 + stacked_videos = stacked_videos.view(
435 + batch_size,
436 + num_frames,
437 + channels,
438 + resized_height,
439 + resized_width,
440 + )
441 + resized_videos_grouped[shape] = stacked_videos
442 + resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)
443 +
444 + grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos)
445 + processed_videos_grouped = {}
446 + processed_grids = {}
447 + for shape, stacked_videos in grouped_videos.items():
448 + resized_height, resized_width = stacked_videos.shape[-2:]
449 + patches = self.rescale_and_normalize(
450 + stacked_videos,
451 + do_rescale,
452 + rescale_factor,
453 + do_normalize,
454 + image_mean,
455 + image_std,
456 + )
457 +
458 + if pad := -patches.shape[1] % temporal_patch_size:
459 + repeats = patches[:, -1:].expand(-1, pad, -1, -1, -1)
460 + patches = torch.cat([patches, repeats], dim=1)
461 +
462 + batch_size, grid_t, channels = patches.shape[:3]
463 + grid_t = grid_t // temporal_patch_size
464 + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
465 +
466 + patches = patches.view(
467 + batch_size,
468 + grid_t,
469 + temporal_patch_size,
470 + channels,
471 + grid_h // merge_size,
472 + merge_size,
473 + patch_size,
474 + grid_w // merge_size,
475 + merge_size,
476 + patch_size,
477 + )
478 + patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)
479 + flatten_patches = patches.reshape(
480 + batch_size,
481 + grid_t * grid_h * grid_w,
482 + channels * temporal_patch_size * patch_size * patch_size,
483 + )
484 +
485 + processed_videos_grouped[shape] = flatten_patches
486 + processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
487 +
488 + processed_videos = reorder_videos(
489 + processed_videos_grouped, grouped_videos_index
490 + )
491 + processed_grids = reorder_videos(processed_grids, grouped_videos_index)
492 + pixel_values_videos = torch.cat(processed_videos, dim=0)
493 + video_grid_thw = torch.tensor(processed_grids, dtype=torch.long)
494 +
495 + return BatchFeature(
496 + data={
497 + "pixel_values_videos": pixel_values_videos,
498 + "video_grid_thw": video_grid_thw,
499 + },
500 + tensor_type=return_tensors,
501 + )
502 +
503 +
504 +class MiniMaxVLProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[call-arg]
505 + _defaults = {
506 + "videos_kwargs": {
507 + "do_resize": False,
508 + "return_metadata": True,
509 + },
510 + }
511 +
512 +
513 +class MiniMaxVLProcessor(ProcessorMixin):
514 + IMAGE_TOKEN = "]<]image[>["
515 + VIDEO_TOKEN = "]<]video[>["
516 + VISION_START_TOKEN = "]<]start of image[>["
517 + VISION_END_TOKEN = "]<]end of image[>["
518 +
519 + @classmethod
520 + def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
521 + # Bypass ProcessorMixin's dynamic module lookup, which breaks in
522 + # transformers >= 5.9 when image_processor_class is a string: the
523 + # register() API now stores classes as {"pil": cls} dicts in
524 + # _extra_content, but get_possibly_dynamic_module() still calls
525 + # .__name__ on the raw value, crashing with AttributeError on dicts.
526 + tokenizer = AutoTokenizer.from_pretrained(
527 + pretrained_model_name_or_path, **kwargs
528 + )
529 + image_processor = MiniMaxM3VLImageProcessor.from_pretrained(
530 + pretrained_model_name_or_path, **kwargs
531 + )
532 + video_processor = MiniMaxM3VLVideoProcessor.from_pretrained(
533 + pretrained_model_name_or_path, **kwargs
534 + )
535 + return cls(
536 + image_processor=image_processor,
537 + tokenizer=tokenizer,
538 + video_processor=video_processor,
539 + )
540 +
541 + def __init__(
542 + self, image_processor=None, tokenizer=None, video_processor=None, **kwargs
543 + ):
544 + self.image_token_id = tokenizer.convert_tokens_to_ids(self.IMAGE_TOKEN)
545 + self.video_token_id = tokenizer.convert_tokens_to_ids(self.VIDEO_TOKEN)
546 + super().__init__(image_processor, tokenizer, video_processor)
547 + # Video expansion also uses image start/end tokens. Separate video
548 + # start/end tokens exist in the tokenizer, but the original MiniMax
549 + # serving path did not use them; keep that behavior for compatibility.
550 + self.vision_start_token_id = tokenizer.convert_tokens_to_ids(
551 + self.VISION_START_TOKEN
552 + )
553 + self.vision_end_token_id = tokenizer.convert_tokens_to_ids(
554 + self.VISION_END_TOKEN
555 + )
556 +
557 + def _prune_video_tokens(
558 + self,
559 + input_text: str,
560 + video_segments: list[int],
561 + video_token: str,
562 + ) -> str:
563 + """Prune video tokens by temporal_patch_size (e.g., 2:1).
564 +
565 + Expects the prompt to carry exactly sum(video_segments) video tokens
566 + — i.e. one token per *sampled* frame — then drops tokens.
567 + """
568 + # If no videos or temporal_patch_size <= 1, no pruning needed
569 + if not video_segments or self.video_processor.temporal_patch_size <= 1:
570 + return input_text
571 +
572 + # Split while keeping delimiters
573 + special_tokens = [video_token]
574 + pattern = "|".join(map(re.escape, special_tokens))
575 + parts = re.split(f"({pattern})", input_text)
576 +
577 + def is_timestamp(text: str) -> bool:
578 + """Check if text ends with timestamp format like ']<]0.0 seconds[>['"""
579 + return (
580 + text.endswith("seconds[>[")
581 + or text.endswith("seconds[>[ ")
582 + or text.endswith("seconds [>[")
583 + or text.endswith("seconds [>[ ")
584 + )
585 +
586 + def extract_timestamp(text: str) -> str:
587 + """Extract timestamp text from the end, starting from ']<]'"""
588 + start_index = text.rfind("]<]")
589 + if start_index == -1:
590 + raise ValueError(f"Failed to extract timestamp: {text}")
591 + return text[start_index:]
592 +
593 + # Build new text with pruned video tokens
594 + final_parts = []
595 + current_seg_idx = 0 # Which video segment we're in
596 + frame_in_seg = 0 # Frame index within current segment
597 + last_timestamp_len = 0 # Length of timestamp to potentially remove
598 +
599 + for part in parts:
600 + if part == video_token:
601 + if current_seg_idx < len(video_segments):
602 + if frame_in_seg % self.video_processor.temporal_patch_size == 0:
603 + # Keep this video token
604 + final_parts.append(part)
605 + frame_in_seg += 1
606 + if frame_in_seg >= video_segments[current_seg_idx]:
607 + current_seg_idx += 1
608 + frame_in_seg = 0
609 + last_timestamp_len = 0
610 + else:
611 + # Skip this video token
612 + frame_in_seg += 1
613 + if frame_in_seg >= video_segments[current_seg_idx]:
614 + current_seg_idx += 1
615 + frame_in_seg = 0
616 + # Remove the timestamp that was already appended
617 + if last_timestamp_len > 0:
618 + assert len(final_parts) > 0
619 + final_parts[-1] = final_parts[-1][:-last_timestamp_len]
620 + last_timestamp_len = 0
621 + else:
622 + # No more video segments, keep as is
623 + final_parts.append(part)
624 + last_timestamp_len = 0
625 + else:
626 + # Text part
627 + final_parts.append(part)
628 + # Check if this text ends with a timestamp
629 + if is_timestamp(part):
630 + last_timestamp_len = len(extract_timestamp(part))
631 + else:
632 + last_timestamp_len = 0
633 +
634 + return "".join(final_parts)
635 +
636 + def __call__(
637 + self,
638 + images=None,
639 + text=None,
640 + videos=None,
641 + **kwargs: Unpack[MiniMaxVLProcessorKwargs],
642 + ) -> BatchFeature:
643 + output_kwargs = self._merge_kwargs(
644 + MiniMaxVLProcessorKwargs,
645 + tokenizer_init_kwargs=self.tokenizer.init_kwargs,
646 + **kwargs,
647 + )
648 +
649 + if images is not None:
650 + images_kwargs = output_kwargs["images_kwargs"]
651 + image_inputs = self.image_processor(images=images, **images_kwargs)
652 + image_grid_thw = image_inputs["image_grid_thw"]
653 + else:
654 + image_inputs = {}
655 + image_grid_thw = None
656 +
657 + if videos is not None:
658 + videos_kwargs = output_kwargs["videos_kwargs"]
659 + video_inputs = self.video_processor(videos=videos, **videos_kwargs)
660 + video_grid_thw = video_inputs["video_grid_thw"]
661 + if not kwargs.get("return_metadata"):
662 + video_metadata = video_inputs.pop("video_metadata")
663 + else:
664 + video_metadata = video_inputs["video_metadata"]
665 + else:
666 + video_inputs = {}
667 + video_grid_thw = None
668 +
669 + if not isinstance(text, list):
670 + text = [text]
671 + text = text.copy()
672 +
673 + # Expand image tokens
674 + if image_grid_thw is not None:
675 + merge_length = self.image_processor.merge_size**2
676 + placeholder = "]<]placeholder[>["
677 + index = 0
678 + for i in range(len(text)):
679 + while self.IMAGE_TOKEN in text[i]:
680 + num_tokens = image_grid_thw[index].prod() // merge_length
681 + text[i] = text[i].replace(
682 + self.IMAGE_TOKEN,
683 + self.VISION_START_TOKEN
684 + + placeholder * num_tokens
685 + + self.VISION_END_TOKEN,
686 + 1,
687 + )
688 + index += 1
689 + text[i] = text[i].replace(placeholder, self.IMAGE_TOKEN)
690 +
691 + # Expand video tokens
692 + if video_grid_thw is not None:
693 + merge_length = self.image_processor.merge_size**2
694 + placeholder = "]<]placeholder[>["
695 + index = 0
696 + for i in range(len(text)):
697 + while self.VIDEO_TOKEN in text[i]:
698 + metadata = video_metadata[index]
699 + grid_t = video_grid_thw[index][0]
700 + frame_seqlen = video_grid_thw[index][1:].prod() // merge_length
701 +
702 + video_placeholder = ""
703 + for frame_idx in range(grid_t):
704 + if (
705 + metadata.fps is not None
706 + and metadata.frames_indices is not None
707 + ):
708 + ts = (
709 + metadata.frames_indices[
710 + min(
711 + frame_idx
712 + * self.video_processor.temporal_patch_size,
713 + len(metadata.frames_indices) - 1,
714 + )
715 + ]
716 + / metadata.fps
717 + )
718 + video_placeholder += f"]<]{ts:.1f} seconds[>["
719 + video_placeholder += (
720 + self.VISION_START_TOKEN
721 + + placeholder * frame_seqlen
722 + + self.VISION_END_TOKEN
723 + )
724 +
725 + text[i] = text[i].replace(self.VIDEO_TOKEN, video_placeholder, 1)
726 + index += 1
727 + text[i] = text[i].replace(placeholder, self.VIDEO_TOKEN)
728 +
729 + # Tokenize
730 + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
731 + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
732 +
733 + return BatchFeature(
734 + data={**text_inputs, **image_inputs, **video_inputs},
735 + tensor_type=return_tensors,
736 + )
vllm/v1/attention/backends/flashinfer.py
+29 -4
@@ -336,9 +336,24 @@ class FlashInferBackend(AttentionBackend):
336
337 @staticmethod
338 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
339 - # Note: Not sure for all platforms, but on Blackwell,
340 - # only support a page size of 16, 32, 64.
341 - return [16, 32, 64]
339 + # Page sizes >= 128 only run on the trtllm-gen dynamic kernel (GQA/MQA
340 + # on Blackwell); advertise them only when usable so selection never
341 + # picks a large kernel block we cannot serve.
342 + use_large_pages = False
343 + vllm_config = get_current_vllm_config_or_none()
344 + if vllm_config is not None and vllm_config.model_config is not None:
345 + pc = vllm_config.parallel_config
346 + mc = vllm_config.model_config
347 + num_qo_heads = mc.get_num_attention_heads(pc)
348 + num_kv_heads = mc.get_num_kv_heads(pc)
349 + use_large_pages = (
350 + num_kv_heads > 0
351 + and num_qo_heads // num_kv_heads > 1
352 + and can_use_trtllm_attention(num_qo_heads, num_kv_heads)
353 + )
354 + if not use_large_pages:
355 + return [16, 32, 64]
356 + return [16, 32, 64, 128, 256, 512, 1024]
357
358 @staticmethod
359 def get_name() -> str:
@@ -647,6 +662,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
662 # if TRTLLM attention kernel is not used when building attn metadata
663 can_use_trtllm = can_use_trtllm_attention(self.num_qo_heads, self.num_kv_heads)
664
665 + # Page sizes >= 128 require the trtllm-gen GQA/MQA path (guaranteed by
666 + # get_supported_kernel_block_sizes).
667 + assert self.page_size <= 64 or (
668 + can_use_trtllm and self.num_qo_heads // self.num_kv_heads > 1
669 + ), f"Unexpected FlashInfer page size {self.page_size} without trtllm-gen GQA"
670 +
671 if (
672 can_use_trtllm
673 and not vllm_config.attention_config.disable_flashinfer_q_quantization
@@ -917,6 +938,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
938 # - Decode (FI native or TRTLLM)
939 use_cascade = common_prefix_len > 0
940 uses_spec_reorder = self.reorder_batch_threshold > 1
941 + # Page sizes >= 128 must use trtllm-gen; force it for prefill too.
942 + prefill_force_trtllm = (
943 + True if page_size >= 128 else self.attention_config.use_trtllm_attention
944 + )
945 prefill_use_trtllm = use_trtllm_attention(
946 self.num_qo_heads,
947 self.num_kv_heads,
@@ -926,7 +951,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
951 self.cache_dtype,
952 self.q_data_type,
953 is_prefill=True,
929 - force_use_trtllm=self.attention_config.use_trtllm_attention,
954 + force_use_trtllm=prefill_force_trtllm,
955 has_sinks=self.has_sinks,
956 has_spec=uses_spec_reorder,
957 )
vllm/v1/attention/backends/registry.py
+3
@@ -91,6 +91,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
91 "vllm.models.deepseek_v4.amd.rocm.DeepseekV4ROCMAiterMLASparseBackend"
92 )
93 FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"
94 + MINIMAX_M3_SPARSE = (
95 + "vllm.models.minimax_m3.common.sparse_attention.MiniMaxM3SparseBackend"
96 + )
97 NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend"
98 FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
99 ROCM_AITER_UNIFIED_ATTN = (
vllm/v1/spec_decode/llm_base_proposer.py
+13 -1
@@ -7,6 +7,7 @@ import numpy as np
7 import torch
8 import torch.nn as nn
9
10 +from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
11 from vllm.config import (
12 CUDAGraphMode,
13 VllmConfig,
@@ -250,6 +251,13 @@ class SpecDecodeBaseProposer:
251 DeepseekV4ROCMAiterMLASparseMetadata,
252 DeepseekV4ROCMAiterSparseSWAMetadata,
253 )
254 +
255 + # MiniMax-M3 sparse (lightning-indexer) attention. The multi-step
256 + # drafting machinery is shared code at num_speculative_tokens>1.
257 + # this just opts the metadata into the ROCm allowlist.
258 + from vllm.models.minimax_m3.common.sparse_attention import (
259 + MiniMaxM3SparseMetadata,
260 + )
261 from vllm.v1.attention.backends.mla.indexer import (
262 DeepseekV32IndexerMetadata,
263 )
@@ -265,6 +273,7 @@ class SpecDecodeBaseProposer:
273 DeepseekV4ROCMAiterMLASparseMetadata,
274 DeepseekV4ROCMAiterSparseSWAMetadata,
275 DeepseekV32IndexerMetadata,
276 + MiniMaxM3SparseMetadata,
277 ]
278 # ROCM_AITER_FA is an optional backend
279 # We check is_enabled() here to avoid importing the backend module during
@@ -457,8 +466,11 @@ class SpecDecodeBaseProposer:
466 batch_size = common_attn_metadata.batch_size()
467
468 if self.method in ("eagle3", "dflash"):
469 + model = self.model
470 + if isinstance(model, BreakableCUDAGraphWrapper):
471 + model = model.unwrap()
472 assert isinstance(
461 - self.model,
473 + model,
474 (
475 Eagle3LlamaForCausalLM,
476 Eagle3DeepseekV2ForCausalLM,
vllm/v1/worker/block_table.py
+1 -1
@@ -322,7 +322,7 @@ class MultiGroupBlockTable:
322 return self.block_tables[idx]
323
324
325 -@triton.jit
325 +@triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"])
326 def _compute_slot_mapping_kernel(
327 num_tokens,
328 max_num_tokens,