refactor & simplify

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

CMakeLists.txt

@@ -315,7 +315,8 @@ set(VLLM_EXT_SRC
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")
   list(APPEND VLLM_EXT_SRC
-    "csrc/minimax_reduce_rms_kernel.cu")
+    "csrc/minimax_reduce_rms_kernel.cu"
+    "csrc/minimax_m3_build_k2q_csr.cu")
 
   SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
 
@@ -637,6 +638,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
     "csrc/libtorch_stable/pos_encoding_kernels.cu"
     "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
+    "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu"
     "csrc/libtorch_stable/layernorm_kernels.cu"
     "csrc/libtorch_stable/layernorm_quant_kernels.cu"
     "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"

csrc/libtorch_stable/activation_kernels.cu

@@ -10,11 +10,20 @@
 
 namespace vllm {
 
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
+// `alpha` and `beta` are applied to opposite operands:
+//   - alpha lives INSIDE the activation (the activated half): the gated
+//     activation computes act_half * sigmoid(alpha * act_half).
+//   - beta is added to the OTHER (non-activated) half before the multiply.
+// So the result is always ACT(act_half, alpha) * (other_half + beta).
+// Which half is which depends on `act_first` (see below). Defaults
+// alpha=1.0, beta=0.0 reproduce the plain SwiGLU/GeGLU behavior.
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float),
           bool act_first, bool HAS_CLAMP>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
                                             const scalar_t& y,
-                                            const float limit) {
+                                            const float limit,
+                                            const float alpha,
+                                            const float beta) {
   if constexpr (act_first) {
     scalar_t gate = x;
     scalar_t up = y;
@@ -22,7 +31,9 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
       gate = (scalar_t)fminf((float)gate, limit);
       up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
     }
-    return ACT_FN(gate) * up;
+    // act_first: gate is the activated half -> alpha applies to gate;
+    // beta is added to up (the non-activated half).
+    return ACT_FN(gate, alpha) * (scalar_t)((float)up + beta);
   } else {
     scalar_t gate = x;
     scalar_t up = y;
@@ -30,55 +41,66 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
       gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
       up = (scalar_t)fminf((float)up, limit);
     }
-    return gate * ACT_FN(up);
+    // !act_first: up is the activated half -> alpha applies to up;
+    // beta is added to gate (the non-activated half).
+    return (scalar_t)((float)gate + beta) * ACT_FN(up, alpha);
   }
 }
 
-template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
+template <typename packed_t,
+          packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
           bool act_first, bool HAS_CLAMP>
 __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
                                                    const packed_t& y,
-                                                   const float limit) {
+                                                   const float limit,
+                                                   const float alpha,
+                                                   const float beta) {
   if constexpr (act_first) {
     packed_t gate = x;
     packed_t up = y;
+    float2 u = cast_to_float2(up);
     if constexpr (HAS_CLAMP) {
       float2 g = cast_to_float2(gate);
-      float2 u = cast_to_float2(up);
       g.x = fminf(g.x, limit);
       g.y = fminf(g.y, limit);
       u.x = fmaxf(fminf(u.x, limit), -limit);
       u.y = fmaxf(fminf(u.y, limit), -limit);
       gate = cast_to_packed<packed_t>(g);
-      up = cast_to_packed<packed_t>(u);
     }
-    return packed_mul(PACKED_ACT_FN(gate), up);
+    // act_first: gate is the activated half -> alpha applies to gate;
+    // beta is added to up (the non-activated half).
+    u.x += beta;
+    u.y += beta;
+    return packed_mul(PACKED_ACT_FN(gate, alpha), cast_to_packed<packed_t>(u));
   } else {
     packed_t gate = x;
     packed_t up = y;
+    float2 g = cast_to_float2(gate);
     if constexpr (HAS_CLAMP) {
-      float2 g = cast_to_float2(gate);
       float2 u = cast_to_float2(up);
       g.x = fmaxf(fminf(g.x, limit), -limit);
       g.y = fmaxf(fminf(g.y, limit), -limit);
       u.x = fminf(u.x, limit);
       u.y = fminf(u.y, limit);
-      gate = cast_to_packed<packed_t>(g);
       up = cast_to_packed<packed_t>(u);
     }
-    return packed_mul(gate, PACKED_ACT_FN(up));
+    // !act_first: up is the activated half -> alpha applies to up;
+    // beta is added to gate (the non-activated half).
+    g.x += beta;
+    g.y += beta;
+    return packed_mul(cast_to_packed<packed_t>(g), PACKED_ACT_FN(up, alpha));
   }
 }
 
 // Activation and gating kernel template.
 template <typename scalar_t, typename packed_t,
-          scalar_t (*ACT_FN)(const scalar_t&),
-          packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
-          bool use_vec, bool HAS_CLAMP, bool use_256b = false>
+          scalar_t (*ACT_FN)(const scalar_t&, const float),
+          packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
+          bool act_first, bool use_vec, bool HAS_CLAMP, bool use_256b = false>
 __global__ void act_and_mul_kernel(
     scalar_t* __restrict__ out,          // [..., d]
     const scalar_t* __restrict__ input,  // [..., 2, d]
-    const int d, const float limit) {
+    const int d, const float limit, const float alpha, const float beta) {
   const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
   const scalar_t* y_ptr = x_ptr + d;
   scalar_t* out_ptr = out + blockIdx.x * d;
@@ -105,7 +127,7 @@ __global__ void act_and_mul_kernel(
       for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
         x.elts[j] =
             packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
-                x.elts[j], y.elts[j], limit);
+                x.elts[j], y.elts[j], limit, alpha, beta);
       }
       if constexpr (use_256b) {
         st256(x, &out_vec[i]);
@@ -118,29 +140,34 @@ __global__ void act_and_mul_kernel(
     for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
       const scalar_t x = VLLM_LDG(&x_ptr[idx]);
       const scalar_t y = VLLM_LDG(&y_ptr[idx]);
-      out_ptr[idx] =
-          compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
+      out_ptr[idx] = compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(
+          x, y, limit, alpha, beta);
     }
   }
 }
 
+// Gated activations take an `alpha` argument that scales the sigmoid input
+// (`x * sigmoid(alpha * x)`). alpha defaults to 1.0 at all call sites, which
+// is exactly SiLU; only the clamp path (silu_and_mul_with_clamp) passes a
+// non-default alpha. Activations that do not use alpha simply ignore it.
 template <typename T>
-__device__ __forceinline__ T silu_kernel(const T& x) {
-  // x * sigmoid(x)
-  return (T)(((float)x) / (1.0f + expf((float)-x)));
+__device__ __forceinline__ T silu_kernel(const T& x, const float alpha) {
+  // x * sigmoid(alpha * x)
+  return (T)(((float)x) / (1.0f + expf((float)-x * alpha)));
 }
 
 template <typename packed_t>
-__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val) {
-  // x * sigmoid(x)
+__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val,
+                                                       const float alpha) {
+  // x * sigmoid(alpha * x)
   float2 fval = cast_to_float2(val);
-  fval.x = fval.x / (1.0f + expf(-fval.x));
-  fval.y = fval.y / (1.0f + expf(-fval.y));
+  fval.x = fval.x / (1.0f + expf(-fval.x * alpha));
+  fval.y = fval.y / (1.0f + expf(-fval.y * alpha));
   return cast_to_packed<packed_t>(fval);
 }
 
 template <typename T>
-__device__ __forceinline__ T gelu_kernel(const T& x) {
+__device__ __forceinline__ T gelu_kernel(const T& x, const float /*alpha*/) {
   // Equivalent to PyTorch GELU with 'none' approximation.
   // Refer to:
   // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
@@ -150,7 +177,8 @@ __device__ __forceinline__ T gelu_kernel(const T& x) {
 }
 
 template <typename packed_t>
-__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
+__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val,
+                                                       const float /*alpha*/) {
   // Equivalent to PyTorch GELU with 'none' approximation.
   // Refer to:
   // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
@@ -162,7 +190,8 @@ __device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
 }
 
 template <typename T>
-__device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
+__device__ __forceinline__ T gelu_tanh_kernel(const T& x,
+                                              const float /*alpha*/) {
   // Equivalent to PyTorch GELU with 'tanh' approximation.
   // Refer to:
   // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
@@ -176,7 +205,7 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
 
 template <typename packed_t>
 __device__ __forceinline__ packed_t
-packed_gelu_tanh_kernel(const packed_t& val) {
+packed_gelu_tanh_kernel(const packed_t& val, const float /*alpha*/) {
   // Equivalent to PyTorch GELU with 'tanh' approximation.
   // Refer to:
   // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
@@ -202,7 +231,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
 // clamped (max only) and up input is clamped (both sides) before the
 // activation function is applied.
 #define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST,        \
-                                      HAS_CLAMP, LIMIT)                        \
+                                      HAS_CLAMP, LIMIT, ALPHA, BETA)           \
   auto dtype = input.scalar_type();                                            \
   int d = input.size(-1) / 2;                                                  \
   int64_t num_tokens = input.numel() / input.size(-1);                         \
@@ -230,7 +259,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
             ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
             out.mutable_data_ptr<scalar_t>(),                                  \
-            input.const_data_ptr<scalar_t>(), d, LIMIT);                       \
+            input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA);          \
       });                                                                      \
     } else {                                                                   \
       VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {   \
@@ -240,7 +269,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
             ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
             out.mutable_data_ptr<scalar_t>(),                                  \
-            input.const_data_ptr<scalar_t>(), d, LIMIT);                       \
+            input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA);          \
       });                                                                      \
     }                                                                          \
   } else {                                                                     \
@@ -252,7 +281,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
           PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
           ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
           out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),  \
-          d, LIMIT);                                                           \
+          d, LIMIT, ALPHA, BETA);                                              \
     });                                                                        \
   }
 
@@ -260,14 +289,18 @@ void silu_and_mul(torch::stable::Tensor& out,    // [..., d]
                   torch::stable::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, false, 0.0f);
+                                true, false, 0.0f, 1.0f, 0.0f);
 }
 
 void silu_and_mul_clamp(torch::stable::Tensor& out,    // [..., d]
                         torch::stable::Tensor& input,  // [..., 2 * d]
-                        double limit) {
+                        double limit, double alpha, double beta) {
+  // out = (gate.clamp(max=limit) * sigmoid(alpha * gate.clamp(max=limit)))
+  //       * (up.clamp(+-limit) + beta)
+  // alpha=1.0, beta=0.0 reduce this to silu(gate) * up.
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, true, (float)limit);
+                                true, true, (float)limit, (float)alpha,
+                                (float)beta);
 }
 
 void mul_and_silu(torch::stable::Tensor& out,    // [..., d]
@@ -276,21 +309,22 @@ void mul_and_silu(torch::stable::Tensor& out,    // [..., d]
   // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
   // applies the silu to the latter half of the input.
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                false, false, 0.0f);
+                                false, false, 0.0f, 1.0f, 0.0f);
 }
 
 void gelu_and_mul(torch::stable::Tensor& out,    // [..., d]
                   torch::stable::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
-                                true, false, 0.0f);
+                                true, false, 0.0f, 1.0f, 0.0f);
 }
 
 void gelu_tanh_and_mul(torch::stable::Tensor& out,    // [..., d]
                        torch::stable::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(
-      vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
+                                vllm::packed_gelu_tanh_kernel, true, false,
+                                0.0f, 1.0f, 0.0f);
 }
 
 namespace vllm {

csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu

@@ -0,0 +1,615 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+ *
+ * Horizontally-fused MiniMax-M3 attention pre-processing kernel.
+ *
+ * Replaces the per-token Python sequence in
+ * ``MiniMaxM3SparseAttention.forward`` / ``MiniMaxM3Attention.forward``:
+ *
+ *     q  = q_norm(q);  k = k_norm(k);  q, k = rotary_emb(pos, q, k)
+ *     index_q = index_q_norm(index_q);  index_k = index_k_norm(index_k)
+ *     index_q, index_k = rotary_emb(pos, index_q, index_k)
+ *     _insert_kv(k, v, index_k)
+ *
+ * All branches share head_dim=128 and the *same* partial-NeoX RoPE table
+ * (``rotary_dim`` rotated, the trailing dims pass through).  The four norms
+ * are Gemma-style RMSNorm (``x * rsqrt(mean(x^2)+eps) * (1 + weight)``) with
+ * independent weights.
+ *
+ * Everything lives in a single fused ``qkv`` tensor.  The sparse layer's
+ * fused projection (MinimaxM3QKVParallelLinearWithIndexer) emits, per token::
+ *
+ *     [ q | k | v | index_q | index_k ]   (the "5 results")
+ *
+ * while the dense layer emits just ``[ q | k | v ]``.  The kernel reads the
+ * index branch straight out of that packed row -- no separate index tensors.
+ *
+ * One kernel, one grid; each warp owns one (token, head-slot) pair.  Slot
+ * enumeration per token:
+ *     [0, nq)                         Q  heads   -> norm(q_w)  + RoPE, write
+ * qkv [nq, nq+nkv)                    K  heads   -> norm(k_w)  + RoPE, write
+ * qkv
+ *                                                  (+ insert into key cache)
+ *     [nq+nkv, nq+2*nkv)              V  heads   -> insert into value cache
+ *     IQ heads (niq)                             -> norm(iq_w) + RoPE, write iq
+ *     IK       (1)                               -> norm(ik_w) + RoPE
+ *                                                  (+ insert into index cache)
+ *
+ * The IQ/IK warps address the index_q/index_k sub-blocks *inside* qkv at the
+ * fixed physical offsets (nq+2*nkv)*128 and (nq+2*nkv+niq)*128.
+ *
+ * Dense vs sparse is a compile-time choice via the ``kIsSparse``/``kInsertKV``
+ * template bools (3 instantiations: dense <false,false>, sparse-profiling
+ * <true,false>, sparse-serving <true,true>), so the index slots, the V slots
+ * and the cache inserts fold away entirely on paths that don't use them. The
+ * dense layer passes no caches/index: norm+RoPE happens in place and the
+ * generic ``Attention`` layer owns the cache write.
+ *
+ * Q/K and (sparse) index_q/index_k are all rewritten in place inside the fused
+ * ``qkv`` tensor.  Caches (bf16) are scatter-written by slot.
+ */
+
+#include <cmath>
+#include <cuda_runtime.h>
+#include <type_traits>
+
+#include "torch_utils.h"
+
+#include "../cuda_compat.h"
+#include "../type_convert.cuh"
+#include "dispatch_utils.h"
+
+#ifndef FINAL_MASK
+  #ifdef USE_ROCM
+    #define FINAL_MASK 0xffffffffffffffffULL
+  #else
+    #define FINAL_MASK 0xffffffffu
+  #endif
+#endif
+
+namespace vllm {
+namespace minimax_m3_fused_ops {
+
+namespace {
+inline int getSMVersion() {
+  auto* props = get_device_prop();
+  return props->major * 10 + props->minor;
+}
+}  // namespace
+
+// ────────────────────────────────────────────────────────────────────────────
+// Constants (hard-coded for MiniMax-M3-preview).
+// ────────────────────────────────────────────────────────────────────────────
+constexpr int kHeadDim = 128;
+constexpr int kNumLanes = 32;
+constexpr int kElemsPerLane = kHeadDim / kNumLanes;  // 4
+
+// ────────────────────────────────────────────────────────────────────────────
+// Helpers
+// ────────────────────────────────────────────────────────────────────────────
+__device__ __forceinline__ float warpReduceSum(float val) {
+#pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1) {
+    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
+  }
+  return val;
+}
+
+// Gemma RMSNorm over the full head (no-op when ``weight == nullptr``) followed
+// by partial NeoX RoPE on the leading ``rotary_dim`` dims, all in fp32.  Each
+// lane owns ``kElemsPerLane`` contiguous dims [laneId*4, laneId*4+4).
+template <typename scalar_t>
+__device__ __forceinline__ void normAndRope(
+    float (&elems)[kElemsPerLane], int const laneId, float const eps,
+    scalar_t const* __restrict__ weight,  // [kHeadDim] or nullptr (no norm)
+    bool const do_rope, int const rotary_dim,
+    scalar_t const* __restrict__ cos_ptr,  // cos_sin_cache + pos*rotary_dim
+    bool const apply_norm) {
+  // ── Gemma RMSNorm: x * rsqrt(mean(x^2)+eps) * (1 + w) ──────────────────
+  if (apply_norm) {
+    float sumsq = 0.0f;
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) sumsq += elems[i] * elems[i];
+    sumsq = warpReduceSum(sumsq);
+    float const rms_rcp = rsqrtf(sumsq / static_cast<float>(kHeadDim) + eps);
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      int const dim = laneId * kElemsPerLane + i;
+      float const w = 1.0f + static_cast<float>(weight[dim]);
+      elems[i] = elems[i] * rms_rcp * w;
+    }
+  }
+
+  // ── Partial NeoX RoPE on dims [0, rotary_dim) ──────────────────────────
+  // half = rotary_dim/2.  Pair (i, i+half) for i in [0, half).  Lane L owns
+  // dims [4L, 4L+4); since half is a multiple of 4, a lane lies wholly in the
+  // first half (own=x[i]) or second half (own=x[i+half]); its partner lives
+  // ``half/4`` lanes away (XOR with that distance).
+  if (do_rope) {
+    int const half = rotary_dim / 2;
+    int const dim0 = laneId * kElemsPerLane;
+    bool const in_rope = dim0 < rotary_dim;
+    int const lane_xor = half / kElemsPerLane;  // partner-lane distance
+
+    float partner[kElemsPerLane];
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      partner[i] = __shfl_xor_sync(FINAL_MASK, elems[i], lane_xor, 32);
+    }
+    if (in_rope) {
+      bool const first_half = dim0 < half;
+      int const i_base = first_half ? dim0 : (dim0 - half);  // cos/sin index
+      scalar_t const* sin_ptr = cos_ptr + half;
+#pragma unroll
+      for (int i = 0; i < kElemsPerLane; i++) {
+        float const c = static_cast<float>(cos_ptr[i_base + i]);
+        float const s = static_cast<float>(sin_ptr[i_base + i]);
+        if (first_half) {
+          elems[i] = elems[i] * c - partner[i] * s;
+        } else {
+          elems[i] = elems[i] * c + partner[i] * s;
+        }
+      }
+    }
+  }
+}
+
+// Load 4 contiguous bf16 -> 4 fp32 registers.
+template <typename scalar_t>
+__device__ __forceinline__ void loadElems(scalar_t const* __restrict__ src,
+                                          float (&elems)[kElemsPerLane]) {
+  using Converter = vllm::_typeConvert<scalar_t>;
+  uint2 v = *reinterpret_cast<uint2 const*>(src);
+  auto const* p =
+      reinterpret_cast<typename Converter::packed_hip_type const*>(&v);
+#pragma unroll
+  for (int i = 0; i < kElemsPerLane / 2; i++) {
+    float2 f2 = Converter::convert(p[i]);
+    elems[2 * i] = f2.x;
+    elems[2 * i + 1] = f2.y;
+  }
+}
+
+// Store 4 fp32 registers -> 4 contiguous bf16.
+template <typename scalar_t>
+__device__ __forceinline__ void storeElems(
+    scalar_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
+  using Converter = vllm::_typeConvert<scalar_t>;
+  uint2 v;
+  auto* p = reinterpret_cast<typename Converter::packed_hip_type*>(&v);
+#pragma unroll
+  for (int i = 0; i < kElemsPerLane / 2; i++) {
+    p[i] = Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1]));
+  }
+  *reinterpret_cast<uint2*>(dst) = v;
+}
+
+// ────────────────────────────────────────────────────────────────────────────
+// Kernel
+// ────────────────────────────────────────────────────────────────────────────
+// Grid: 1D, ceil(num_tokens * slots_per_token / warps_per_block).
+// Each warp = one (token, slot).
+//
+// `kIsSparse` and `kInsertKV` are compile-time template bools, so all the
+// branch decisions that distinguish the dense layer from the sparse layer
+// (index slots, KV/index inserts, V slots) fold away per instantiation.
+// Three instantiations are built: dense <false,false>, sparse-profiling
+// <true,false> and sparse-serving <true,true>.  Slots per token:
+//     Q : nq                            (always — norm+RoPE)
+//     K : nkv                           (always — norm+RoPE; +K-cache insert)
+//     V : nkv  only if kInsertKV        (V-cache insert; no warps in dense)
+//     IQ: niq  only if kIsSparse        (norm+RoPE)
+//     IK: 1    only if kIsSparse        (norm+RoPE; +index-cache insert)
+template <typename scalar_t, bool kIsSparse, bool kInsertKV>
+__global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
+    scalar_t* __restrict__ qkv,  // [N, qkv_row] in/out (packs index if sparse)
+    scalar_t* __restrict__ q_out,        // [N, nq*128] contiguous, or nullptr
+    scalar_t* __restrict__ index_q_out,  // [N, niq*128] contiguous, or nullptr
+    scalar_t const* __restrict__ q_norm_w,
+    scalar_t const* __restrict__ k_norm_w,
+    scalar_t const* __restrict__ iq_norm_w,
+    scalar_t const* __restrict__ ik_norm_w,
+    scalar_t const* __restrict__ cos_sin_cache,  // [max_pos, rotary_dim]
+    int64_t const* __restrict__ positions,       // [N] i64
+    int64_t const* __restrict__ slot_mapping,    // [N] i64 or nullptr
+    scalar_t* __restrict__ kv_cache,             // [nb,2,bs,nkv,128] or nullptr
+    scalar_t* __restrict__ index_cache,          // [nb*bs, 128] or nullptr
+    float const eps, int const rotary_dim, int const num_tokens, int const nq,
+    int const nkv, int const niq, int const block_size,
+    // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
+    // The head_dim (last) dim is always innermost-contiguous (stride 1), so the
+    // NHD/HND layout choice is fully captured by these four strides: NHD keeps
+    // s_token < s_head, HND swaps them. dim_base addresses head_dim directly.
+    int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
+    int64_t const kv_s_head) {
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  // _typeConvert<BFloat16> is unavailable on pre-Ampere; the M3 kernel only
+  // runs with bf16/fp16 inputs in practice.  Discard the bf16 body there.
+  if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
+    return;
+  } else {
+#endif
+    int const warpsPerBlock = blockDim.x / 32;
+    int const laneId = threadIdx.x % 32;
+    int const globalWarpIdx = blockIdx.x * warpsPerBlock + (threadIdx.x / 32);
+
+    // Slot layout (compile-time gated: dense has neither V nor index slots).
+    int const v_slots = kInsertKV ? nkv : 0;
+    int const idx_slots = kIsSparse ? niq + 1 : 0;
+    int const slots_per_token = nq + nkv + v_slots + idx_slots;
+
+    int const tokenIdx = globalWarpIdx / slots_per_token;
+    int const slot = globalWarpIdx % slots_per_token;
+    if (tokenIdx >= num_tokens) return;
+
+    // Slot boundaries.
+    int const k_begin = nq;
+    int const v_begin = nq + nkv;             // valid only when kInsertKV
+    int const iq_begin = nq + nkv + v_slots;  // index block start
+    int const ik_slot = iq_begin + niq;       // valid only when kIsSparse
+
+    bool const isQ = slot < k_begin;
+    bool const isK = slot >= k_begin && slot < v_begin;
+    bool isV = false;
+    if constexpr (kInsertKV) isV = slot >= v_begin && slot < v_begin + nkv;
+    bool isIQ = false, isIK = false;
+    if constexpr (kIsSparse) {
+      isIQ = slot >= iq_begin && slot < ik_slot;
+      isIK = slot == ik_slot;
+    }
+
+    int const dim_base = laneId * kElemsPerLane;
+    // Physical row width of qkv: the dense layer packs [q|k|v]; the sparse
+    // layer additionally packs [index_q (niq heads) | index_k (1 head)].
+    int const qkv_row = (nq + 2 * nkv + (kIsSparse ? (niq + 1) : 0)) * kHeadDim;
+
+    // ── Resolve source pointer + per-branch parameters. ────────────────────
+    scalar_t* row_ptr = nullptr;       // in-place output location
+    scalar_t const* norm_w = nullptr;  // nullptr -> skip norm (V)
+    bool do_rope = true;
+    int head = 0;  // kv head index for inserts
+
+    if (isQ) {
+      row_ptr =
+          qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
+      norm_w = q_norm_w;
+    } else if (isK) {
+      head = slot - k_begin;
+      row_ptr =
+          qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
+      norm_w = k_norm_w;
+    } else if (isV) {
+      // qkv V section starts at slot index (nq + nkv): slot * kHeadDim is the
+      // correct in-tensor offset.
+      head = slot - v_begin;
+      row_ptr =
+          qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
+      norm_w = nullptr;  // V: no norm, no rope
+      do_rope = false;
+    } else if (isIQ) {
+      // index_q sub-block lives at physical offset (nq+2*nkv)*128 in qkv.
+      int const ih = slot - iq_begin;
+      row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
+                (nq + 2 * nkv + ih) * kHeadDim;
+      norm_w = iq_norm_w;
+    } else {  // isIK -- single shared index key at (nq+2*nkv+niq)*128.
+      row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
+                (nq + 2 * nkv + niq) * kHeadDim;
+      norm_w = ik_norm_w;
+    }
+
+    // Store destination.  Q and index_q are gathered into dedicated contiguous
+    // output buffers (when provided) so the downstream SM100 sparse kernel's
+    // flat TMA descriptor can address them as [tokens*heads, head_dim]; this
+    // folds the de-interleaving into the store the kernel already does, instead
+    // of a separate q.contiguous() copy.  Everything else stays in place.
+    scalar_t* store_ptr = row_ptr;
+    if (isQ && q_out != nullptr) {
+      store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
+                  slot * kHeadDim;
+    } else if (isIQ && index_q_out != nullptr) {
+      store_ptr = index_q_out +
+                  static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
+                  (slot - iq_begin) * kHeadDim;
+    }
+
+    // PDL: wait for the predecessor kernel (the qkv-projection GEMM that
+    // produces ``qkv``) to finish before touching any global memory.  No-op
+    // when PDL is not enabled on the launch.  The CUDA runtime wrapper emits
+    // the griddepcontrol.wait PTX with the required memory clobber internally.
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaGridDependencySynchronize();
+#endif
+
+    // ── Load -> norm+rope (fp32) -> store back in place. ───────────────────
+    float elems[kElemsPerLane];
+    loadElems<scalar_t>(row_ptr + dim_base, elems);
+
+    if (!isV) {
+      int64_t const pos = positions[tokenIdx];
+      scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim;
+      normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
+                            cos_ptr, /*apply_norm=*/norm_w != nullptr);
+      storeElems<scalar_t>(store_ptr + dim_base, elems);
+    }
+
+    // ── Cache inserts (sparse serving only). ───────────────────────────────
+    if constexpr (kInsertKV) {
+      // Guard (not early-return) so every thread reaches the PDL trigger below.
+      int64_t const sm = (isK || isV || isIK) ? slot_mapping[tokenIdx] : -1;
+      if (sm >= 0) {  // skip padded / unscheduled tokens
+        if (isIK) {
+          scalar_t* dst = index_cache + sm * kHeadDim + dim_base;
+          storeElems<scalar_t>(dst, elems);
+        } else if (isK || isV) {
+          // kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
+          // Paging is logical (block = sm/block_size, token = sm%block_size);
+          // the physical NHD/HND layout is honoured via the passed strides.
+          int64_t const b = sm / block_size;
+          int64_t const t = sm % block_size;
+          int const kv = isK ? 0 : 1;
+          int64_t const off =
+              b * kv_s_block + kv * kv_s_kv + t * kv_s_token + head * kv_s_head;
+          storeElems<scalar_t>(kv_cache + off + dim_base, elems);
+        }
+      }
+    }
+
+    // PDL: signal that this kernel is done so a dependent successor may launch
+    // early.  No-op when PDL is not enabled on the launch.
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  }
+#endif
+}
+
+// ────────────────────────────────────────────────────────────────────────────
+// Launch wrapper
+// ────────────────────────────────────────────────────────────────────────────
+template <typename scalar_t>
+void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
+                          scalar_t const* q_norm_w, scalar_t const* k_norm_w,
+                          scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
+                          scalar_t const* cos_sin_cache,
+                          int64_t const* positions, int64_t const* slot_mapping,
+                          scalar_t* kv_cache, scalar_t* index_cache,
+                          float const eps, int const rotary_dim,
+                          int const num_tokens, int const nq, int const nkv,
+                          int const niq, int const block_size,
+                          int64_t const kv_s_block, int64_t const kv_s_kv,
+                          int64_t const kv_s_token, int64_t const kv_s_head,
+                          bool const has_index, bool const insert_kv,
+                          cudaStream_t stream) {
+  // Slot count must match the kernel's compile-time gating.
+  int const v_slots = insert_kv ? nkv : 0;
+  int const idx_slots = has_index ? niq + 1 : 0;
+  int const slots_per_token = nq + nkv + v_slots + idx_slots;
+
+  constexpr int kBlockSize = 256;
+  constexpr int kWarpsPerBlock = kBlockSize / 32;
+  int64_t const total_warps =
+      static_cast<int64_t>(num_tokens) * slots_per_token;
+  int const grid =
+      static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
+  if (grid == 0) return;
+
+#ifndef USE_ROCM
+  // PDL: enable programmatic stream serialization whenever the hardware
+  // supports it (SM90+).  On pre-Hopper GPUs the attribute is unavailable, so
+  // leave numAttrs = 0 and launch as a regular kernel via cudaLaunchKernelEx.
+  static int const sm_version = getSMVersion();
+  cudaLaunchConfig_t config;
+  config.gridDim = dim3(grid);
+  config.blockDim = dim3(kBlockSize);
+  config.dynamicSmemBytes = 0;
+  config.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = 1;
+  config.attrs = attrs;
+  config.numAttrs = (sm_version >= 90) ? 1 : 0;
+
+  #define LAUNCH(IS_SPARSE, INSERT)                                            \
+    cudaLaunchKernelEx(                                                        \
+        &config,                                                               \
+        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT>,    \
+        qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w,     \
+        cos_sin_cache, positions, slot_mapping, kv_cache, index_cache, eps,    \
+        rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv, \
+        kv_s_token, kv_s_head)
+#else
+  // ROCm: standard kernel launch syntax (no PDL/stream serialization).
+  // clang-format off
+  #define LAUNCH(IS_SPARSE, INSERT)                                          \
+    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT>       \
+        <<<grid, kBlockSize, 0, stream>>>(                                   \
+            qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w,          \
+            ik_norm_w, cos_sin_cache, positions, slot_mapping, kv_cache,     \
+            index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq,          \
+            block_size, kv_s_block, kv_s_kv, kv_s_token, kv_s_head)
+  // clang-format on
+#endif
+
+  if (has_index) {
+    if (insert_kv) {
+      LAUNCH(true, true);  // sparse serving
+    } else {
+      LAUNCH(true, false);  // sparse profiling
+    }
+  } else {
+    // Dense layer: never has an index branch and never inserts here (the
+    // generic Attention layer owns the KV insert).
+    LAUNCH(false, false);
+  }
+#undef LAUNCH
+}
+
+}  // namespace minimax_m3_fused_ops
+}  // namespace vllm
+
+// ────────────────────────────────────────────────────────────────────────────
+// Torch op wrapper
+// ────────────────────────────────────────────────────────────────────────────
+void fused_minimax_m3_qknorm_rope_kv_insert(
+    torch::stable::Tensor& qkv,  // [N, qkv_row] (packs index if sparse)
+    torch::stable::Tensor const& q_norm_weight,  // [128]
+    torch::stable::Tensor const& k_norm_weight,  // [128]
+    torch::stable::Tensor const& cos_sin_cache,  // [max_pos, rotary_dim]
+    torch::stable::Tensor const& positions,      // [N] i64
+    int64_t num_heads, int64_t num_kv_heads, int64_t rotary_dim, double eps,
+    std::optional<torch::stable::Tensor> index_q_norm_weight,  // [128]
+    std::optional<torch::stable::Tensor> index_k_norm_weight,  // [128]
+    int64_t num_index_heads,                            // niq; 0 => dense
+    std::optional<torch::stable::Tensor> slot_mapping,  // [N] i64
+    std::optional<torch::stable::Tensor> kv_cache,      // [nb,2,bs,nkv,128]
+    std::optional<torch::stable::Tensor> index_cache,   // [nb,bs,128]
+    int64_t block_size,
+    std::optional<torch::stable::Tensor> q_out,  // [N, nq*128] contiguous
+    std::optional<torch::stable::Tensor>
+        index_q_out) {  // [N, niq*128] contiguous
+  STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
+                  "qkv must be contiguous CUDA");
+  STD_TORCH_CHECK(
+      positions.is_cuda() &&
+          positions.scalar_type() == torch::headeronly::ScalarType::Long,
+      "positions must be int64 CUDA");
+  STD_TORCH_CHECK(cos_sin_cache.is_cuda() && cos_sin_cache.is_contiguous(),
+                  "cos_sin_cache must be contiguous CUDA");
+  STD_TORCH_CHECK(cos_sin_cache.scalar_type() == qkv.scalar_type(),
+                  "cos_sin_cache dtype must match qkv");
+  STD_TORCH_CHECK(
+      cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == rotary_dim,
+      "cos_sin_cache shape [max_pos, rotary_dim]");
+
+  STD_TORCH_CHECK(q_norm_weight.scalar_type() == qkv.scalar_type() &&
+                      k_norm_weight.scalar_type() == qkv.scalar_type(),
+                  "q/k norm weight dtype must match qkv");
+  STD_TORCH_CHECK(
+      q_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim &&
+          k_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim,
+      "q/k norm weight must have 128 elements");
+  STD_TORCH_CHECK(rotary_dim > 0 && rotary_dim % 8 == 0 &&
+                      rotary_dim <= vllm::minimax_m3_fused_ops::kHeadDim,
+                  "rotary_dim must be a positive multiple of 8 and <= 128");
+
+  int const num_tokens = static_cast<int>(qkv.size(0));
+  int const nq = static_cast<int>(num_heads);
+  int const nkv = static_cast<int>(num_kv_heads);
+  int const niq = static_cast<int>(num_index_heads);
+
+  // The sparse layer packs the index branch ([index_q (niq heads) | index_k
+  // (1 head)]) right after [q|k|v] in the same row; the dense layer does not.
+  bool const has_index = niq > 0;
+  bool const insert_kv = kv_cache.has_value();
+  int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
+  int const expected_row =
+      (nq + 2 * nkv + (has_index ? niq + 1 : 0)) * kHeadDim;
+  STD_TORCH_CHECK(qkv.size(1) == expected_row,
+                  "qkv last dim must be (num_heads + 2*num_kv_heads"
+                  " + num_index_heads + 1) * 128 for sparse, "
+                  "(num_heads + 2*num_kv_heads) * 128 for dense");
+
+  // Only the sparse layer inserts here (dense lets the generic Attention layer
+  // own the KV write); there is no dense+insert kernel instantiation.
+  STD_TORCH_CHECK(
+      !insert_kv || has_index,
+      "insert mode (kv_cache) requires the index branch (sparse layer)");
+  if (has_index) {
+    STD_TORCH_CHECK(
+        index_q_norm_weight.has_value() && index_k_norm_weight.has_value(),
+        "index branch requires both index norm weights");
+    STD_TORCH_CHECK(index_q_norm_weight->scalar_type() == qkv.scalar_type() &&
+                        index_k_norm_weight->scalar_type() == qkv.scalar_type(),
+                    "index norm weights dtype must match qkv");
+    STD_TORCH_CHECK(index_q_norm_weight->numel() == kHeadDim &&
+                        index_k_norm_weight->numel() == kHeadDim,
+                    "index norm weights must have 128 elements");
+  }
+  // kv_cache strides (logical shape [nb, 2, bs, nkv, head_dim]). Read straight
+  // off the tensor so the kernel honours whatever physical layout the attention
+  // backend allocated (NHD: stride order (0,1,2,3,4); HND: (0,1,3,2,4)). No new
+  // op argument is needed -- the strides ride along with the tensor itself.
+  int64_t kv_s_block = 0, kv_s_kv = 0, kv_s_token = 0, kv_s_head = 0;
+  if (insert_kv) {
+    STD_TORCH_CHECK(
+        slot_mapping.has_value() &&
+            slot_mapping->scalar_type() == torch::headeronly::ScalarType::Long,
+        "insert mode requires int64 slot_mapping");
+    STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
+                    "kv_cache dtype must match qkv (bf16 cache only)");
+    STD_TORCH_CHECK(index_cache.has_value() &&
+                        index_cache->scalar_type() == qkv.scalar_type(),
+                    "insert mode requires matching index_cache");
+    STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
+                    "kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
+                    "head_dim (stride(4)==1)");
+    kv_s_block = kv_cache->stride(0);
+    kv_s_kv = kv_cache->stride(1);
+    kv_s_token = kv_cache->stride(2);
+    kv_s_head = kv_cache->stride(3);
+  }
+  // Optional contiguous gather targets: when given, the normed/roped q (and
+  // index_q) are written here instead of in place, so callers avoid a separate
+  // .contiguous() copy.  index_q_out only makes sense on the sparse path.
+  if (q_out.has_value()) {
+    STD_TORCH_CHECK(
+        q_out->is_cuda() && q_out->is_contiguous() &&
+            q_out->scalar_type() == qkv.scalar_type(),
+        "q_out must be a contiguous CUDA tensor matching qkv dtype");
+    STD_TORCH_CHECK(
+        q_out->numel() == static_cast<int64_t>(num_tokens) * nq * kHeadDim,
+        "q_out must have num_tokens * num_heads * 128 elements");
+  }
+  if (index_q_out.has_value()) {
+    STD_TORCH_CHECK(
+        has_index,
+        "index_q_out requires the index branch (num_index_heads > 0)");
+    STD_TORCH_CHECK(
+        index_q_out->is_cuda() && index_q_out->is_contiguous() &&
+            index_q_out->scalar_type() == qkv.scalar_type(),
+        "index_q_out must be a contiguous CUDA tensor matching qkv dtype");
+    STD_TORCH_CHECK(index_q_out->numel() ==
+                        static_cast<int64_t>(num_tokens) * niq * kHeadDim,
+                    "index_q_out must have num_tokens * num_index_heads * 128 "
+                    "elements");
+  }
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      qkv.get_device_index());
+  auto stream = get_current_cuda_stream(qkv.get_device_index());
+
+  VLLM_STABLE_DISPATCH_HALF_TYPES(
+      qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert", [&] {
+        using st = scalar_t;
+        vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st>(
+            reinterpret_cast<st*>(qkv.data_ptr()),
+            q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr())
+                              : nullptr,
+            index_q_out.has_value()
+                ? reinterpret_cast<st*>(index_q_out->data_ptr())
+                : nullptr,
+            reinterpret_cast<st const*>(q_norm_weight.data_ptr()),
+            reinterpret_cast<st const*>(k_norm_weight.data_ptr()),
+            has_index
+                ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr())
+                : nullptr,
+            has_index
+                ? reinterpret_cast<st const*>(index_k_norm_weight->data_ptr())
+                : nullptr,
+            reinterpret_cast<st const*>(cos_sin_cache.data_ptr()),
+            reinterpret_cast<int64_t const*>(positions.data_ptr()),
+            insert_kv
+                ? reinterpret_cast<int64_t const*>(slot_mapping->data_ptr())
+                : nullptr,
+            insert_kv ? reinterpret_cast<st*>(kv_cache->data_ptr()) : nullptr,
+            (insert_kv && has_index)
+                ? reinterpret_cast<st*>(index_cache->data_ptr())
+                : nullptr,
+            static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens,
+            nq, nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv,
+            kv_s_token, kv_s_head, has_index, insert_kv, stream);
+      });
+}

csrc/libtorch_stable/ops.h

@@ -231,6 +231,23 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
                         torch::stable::Tensor& position_ids,
                         int64_t forced_token_heads_per_warp);
 
+// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV /
+// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs
+// the index branch and scatters k/v/index_k into their paged caches.
+void fused_minimax_m3_qknorm_rope_kv_insert(
+    torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight,
+    torch::stable::Tensor const& k_norm_weight,
+    torch::stable::Tensor const& cos_sin_cache,
+    torch::stable::Tensor const& positions, int64_t num_heads,
+    int64_t num_kv_heads, int64_t rotary_dim, double eps,
+    std::optional<torch::stable::Tensor> index_q_norm_weight,
+    std::optional<torch::stable::Tensor> index_k_norm_weight,
+    int64_t num_index_heads, std::optional<torch::stable::Tensor> slot_mapping,
+    std::optional<torch::stable::Tensor> kv_cache,
+    std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
+    std::optional<torch::stable::Tensor> q_out,
+    std::optional<torch::stable::Tensor> index_q_out);
+
 // Sampler kernels (shared CUDA/ROCm)
 void apply_repetition_penalties_(
     torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -276,7 +293,8 @@ void selective_scan_fwd(
 // Activation kernels (shared CUDA/ROCm)
 void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void silu_and_mul_clamp(torch::stable::Tensor& out,
-                        torch::stable::Tensor& input, double limit);
+                        torch::stable::Tensor& input, double limit,
+                        double alpha = 1.0, double beta = 0.0);
 void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_tanh_and_mul(torch::stable::Tensor& out,

csrc/libtorch_stable/torch_bindings.cpp

@@ -337,6 +337,17 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "bool is_neox, Tensor position_ids, "
       "int forced_token_heads_per_warp=-1) -> ()");
 
+  // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert.
+  ops.def(
+      "fused_minimax_m3_qknorm_rope_kv_insert("
+      "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, "
+      "Tensor cos_sin_cache, Tensor positions, int num_heads, "
+      "int num_kv_heads, int rotary_dim, float eps, "
+      "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, "
+      "int num_index_heads, "
+      "Tensor? slot_mapping, Tensor!? kv_cache, Tensor!? index_cache, "
+      "int block_size, Tensor!? q_out, Tensor!? index_q_out) -> ()");
+
   // Apply repetition penalties to logits in-place.
   ops.def(
       "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -364,9 +375,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
 
   // SwiGLU activation with input clamping.
+  // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to
+  // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up.
   ops.def(
-      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
-      "-> ()");
+      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
+      "float alpha=1.0, float beta=0.0) -> ()");
 
   // Activation function used in GeGLU with `none` approximation.
   ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
@@ -571,6 +584,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   // Positional encoding kernels (shared CUDA/ROCm)
   ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
   ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
+  ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
+           TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
 
   // Sampler kernels (shared CUDA/ROCm)
   ops.impl("apply_repetition_penalties_",

csrc/minimax_m3_build_k2q_csr.cu

@@ -0,0 +1,742 @@
+// CUDA C++ q2k -> k2q CSR builder.
+//
+// Five-stage pipeline. q-ascending order within each CSR row is preserved
+// by partitioning q across (CTA, warp_in_CTA) units; each unit owns a
+// contiguous q-sub-range and reserves a contiguous slot range per row via
+// a precomputed exclusive prefix scan.
+//
+//   M:  build_row_map      -- round-robin packing of rows across batches
+//   H:  histogram + tile_counts
+//   PR: row prefix         -- single block per head, row_counts -> row_ptr
+//   PT: tile prefix        -- multi-block, scan tile_counts along (c, w) axis
+//   S:  scatter (sorted)   -- per-warp slot range, q-sequential within warp
+//
+// Per-warp partitioning: each CTA has kWarps warps; warp w of CTA c owns
+// q-range [c*q_per_cta + w*q_per_warp, c*q_per_cta + (w+1)*q_per_warp).
+// tile_counts is shaped [G * kWarps, H, total_rows]; the "row" dimension
+// of the prefix scan is the flattened (c * kWarps + w) index, scanned in
+// lexicographic order so that warp-local slot ranges concatenate to the
+// global q-sorted output.
+
+#include <torch/all.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <algorithm>
+
+#define CHECK_CUDA(x) TORCH_CHECK((x).is_cuda(), #x " must be CUDA")
+#define CHECK_CONTIGUOUS(x) \
+  TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
+#define CHECK_INT(x) \
+  TORCH_CHECK((x).scalar_type() == at::kInt, #x " must be int32")
+#define CHECK_INPUT(x) \
+  CHECK_CUDA(x);       \
+  CHECK_CONTIGUOUS(x); \
+  CHECK_INT(x)
+
+namespace {
+
+constexpr int kWarpSize = 32;
+
+__device__ __forceinline__ void advance_batch_only(int const* __restrict__ cu_q,
+                                                   int B, int q_abs, int& bi) {
+  while (bi < B && cu_q[bi + 1] <= q_abs) ++bi;
+}
+
+// Atomic increment of a 16-bit half within a 32-bit SMEM word; returns the
+// OLD 16-bit value (slot). Per-warp count must stay < 32768 so the low
+// half does not carry into the high half.
+//   base_int32 : int32 pointer; element i holds rows 2*i (low) and 2*i+1
+//   (high).
+__device__ __forceinline__ int atomic_inc_int16_packed(int* base_int32,
+                                                       int row) {
+  int idx = row >> 1;
+  int shift = (row & 1) << 4;  // 0 or 16
+  int delta = 1 << shift;
+  int old = atomicAdd(&base_int32[idx], delta);
+  return (old >> shift) & 0xFFFF;
+}
+
+// Read 16-bit half from packed int32 storage.
+__device__ __forceinline__ int read_int16_packed(int const* base_int32,
+                                                 int row) {
+  int v = base_int32[row >> 1];
+  int shift = (row & 1) << 4;
+  return (v >> shift) & 0xFFFF;
+}
+
+// ---------------------------------------------------------------------------
+// M: round-robin row map.
+// ---------------------------------------------------------------------------
+template <int kBlockK>
+__global__ void k2q_build_row_map_kernel(int const* __restrict__ cu_k,
+                                         int* __restrict__ row_map,
+                                         int* __restrict__ row_coords, int B,
+                                         int max_kv_blocks) {
+  int level = blockIdx.x;
+  if (level >= max_kv_blocks) return;
+  if (threadIdx.x != 0) return;
+  int rows_before = 0;
+  for (int b = 0; b < B; ++b) {
+    int rb = (cu_k[b + 1] - cu_k[b] + kBlockK - 1) / kBlockK;
+    rows_before += (rb < level ? rb : level);
+  }
+  int active_before = 0;
+  for (int b = 0; b < B; ++b) {
+    int rb = (cu_k[b + 1] - cu_k[b] + kBlockK - 1) / kBlockK;
+    if (rb > level) {
+      int row_linear = rows_before + active_before;
+      row_map[(size_t)b * max_kv_blocks + level] = row_linear;
+      if (row_coords != nullptr) {
+        row_coords[(size_t)row_linear * 2] = b;
+        row_coords[(size_t)row_linear * 2 + 1] = level;
+      }
+      ++active_before;
+    } else {
+      row_map[(size_t)b * max_kv_blocks + level] = -1;
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// H: per-warp histogram + tile_counts.
+// kWarps warps per CTA, each owns q-sub-range = q_per_cta / kWarps.
+// SMEM hist[kWarps, total_rows] int32 (stored as packed int16 cursor:
+// 2 entries per int32 word). Each warp counts to its own row.
+// At end-of-CTA, write tile_counts[c*kWarps + w, h, r] = smem_hist[w, r]
+// and atomicAdd(row_counts[h, r], sum over w of smem_hist[w, r]).
+// ---------------------------------------------------------------------------
+template <int kTopK, int kBlockK, int kWarps>
+__global__ void k2q_hist_kernel(int const* __restrict__ q2k,
+                                int const* __restrict__ cu_q,
+                                int const* __restrict__ row_map,
+                                int* __restrict__ row_counts,
+                                int* __restrict__ tile_counts, int H, int B,
+                                int S_Q, int total_rows, int max_kv_blocks,
+                                int q_per_cta, int q_per_warp) {
+  constexpr int kThreads = kWarps * kWarpSize;
+  extern __shared__ int smem_hist_int[];
+  int* smem_hist = smem_hist_int;
+  int tid = threadIdx.x;
+  int warp_id = tid >> 5;
+  int lane = tid & 31;
+  int c = blockIdx.x;
+  int q_start_cta = c * q_per_cta;
+  int q_end_cta = min(q_start_cta + q_per_cta, S_Q);
+  int q_start_warp = min(q_start_cta + warp_id * q_per_warp, q_end_cta);
+  int q_end_warp = min(q_start_warp + q_per_warp, q_end_cta);
+
+  constexpr int kInt4PerToken = kTopK / 4;
+  int packed_per_warp = (total_rows + 1) >> 1;
+  int* my_hist = smem_hist + warp_id * packed_per_warp;
+
+  for (int h = 0; h < H; ++h) {
+    for (int i = lane; i < packed_per_warp; i += kWarpSize) my_hist[i] = 0;
+    __syncthreads();
+
+    if (q_start_warp < q_end_warp) {
+      int bi = 0;
+      int qi = q_start_warp + lane;
+      advance_batch_only(cu_q, B, qi, bi);
+
+      int4 const* head_topk4 =
+          reinterpret_cast<int4 const*>(q2k + (size_t)h * S_Q * kTopK);
+
+      for (; qi < q_end_warp; qi += kWarpSize) {
+        advance_batch_only(cu_q, B, qi, bi);
+        int const* my_row_map = row_map + (size_t)bi * max_kv_blocks;
+
+        int4 buf[kInt4PerToken];
+#pragma unroll
+        for (int v = 0; v < kInt4PerToken; ++v) {
+          buf[v] = head_topk4[(size_t)qi * kInt4PerToken + v];
+        }
+#pragma unroll
+        for (int t = 0; t < kTopK; ++t) {
+          int kvb_local = reinterpret_cast<int const*>(buf)[t];
+          if (kvb_local >= 0 && kvb_local < max_kv_blocks) {
+            int row = my_row_map[kvb_local];
+            if (row >= 0 && row < total_rows) {
+              atomic_inc_int16_packed(my_hist, row);
+            }
+          }
+        }
+      }
+    }
+    __syncthreads();
+
+    int* head_row_counts = row_counts + (size_t)h * total_rows;
+    // Each warp writes its own slice of tile_counts (full int32) by
+    // unpacking int16 entries from SMEM.
+    int* my_tile =
+        tile_counts + ((size_t)(c * kWarps + warp_id) * H + h) * total_rows;
+    for (int i = lane; i < total_rows; i += kWarpSize) {
+      my_tile[i] = read_int16_packed(my_hist, i);
+    }
+    __syncthreads();
+
+    // Sum across warps (int32 accumulator), atomicAdd to row_counts.
+    for (int i = tid; i < total_rows; i += kThreads) {
+      int sum = 0;
+#pragma unroll
+      for (int w = 0; w < kWarps; ++w) {
+        sum += read_int16_packed(smem_hist + w * packed_per_warp, i);
+      }
+      if (sum > 0) atomicAdd(&head_row_counts[i], sum);
+    }
+    if (h + 1 < H) __syncthreads();
+  }
+}
+
+// ---------------------------------------------------------------------------
+// PR: row prefix. One block per head.
+// ---------------------------------------------------------------------------
+template <int kThreads>
+__global__ void k2q_row_prefix_kernel(int const* __restrict__ row_counts,
+                                      int* __restrict__ row_ptr,
+                                      int const* __restrict__ row_coords,
+                                      int* __restrict__ scheduler_metadata,
+                                      int* __restrict__ work_count,
+                                      int total_rows, int target_q_per_cta,
+                                      int work_capacity) {
+  int h = blockIdx.x;
+  int tid = threadIdx.x;
+  __shared__ int scan_buf[kThreads];
+
+  int const* head_counts = row_counts + (size_t)h * total_rows;
+  int* head_rowptr = row_ptr + (size_t)h * (total_rows + 1);
+  int chunk = (total_rows + kThreads - 1) / kThreads;
+  int lo = tid * chunk;
+  int hi = min(lo + chunk, total_rows);
+
+  int local_sum = 0;
+  for (int i = lo; i < hi; ++i) local_sum += head_counts[i];
+  scan_buf[tid] = local_sum;
+  __syncthreads();
+
+  for (int off = 1; off < kThreads; off <<= 1) {
+    int add = (tid >= off) ? scan_buf[tid - off] : 0;
+    __syncthreads();
+    scan_buf[tid] += add;
+    __syncthreads();
+  }
+  int running = scan_buf[tid] - local_sum;
+  for (int i = lo; i < hi; ++i) {
+    int row_count = head_counts[i];
+    running += row_count;
+    head_rowptr[i + 1] = running;
+    if (scheduler_metadata != nullptr && work_count != nullptr &&
+        row_count > 0) {
+      int num_chunks = (row_count + target_q_per_cta - 1) / target_q_per_cta;
+      int base = atomicAdd(work_count, num_chunks);
+      int batch_idx = row_coords[(size_t)i * 2];
+      int kv_block_idx = row_coords[(size_t)i * 2 + 1];
+      for (int c = 0; c < num_chunks; ++c) {
+        int work_idx = base + c;
+        if (work_idx < work_capacity) {
+          int q_begin = c * target_q_per_cta;
+          int q_count = min(target_q_per_cta, row_count - q_begin);
+          int* meta = scheduler_metadata + (size_t)work_idx * 6;
+          meta[0] = h;
+          meta[1] = i;
+          meta[2] = q_begin;
+          meta[3] = q_count;
+          meta[4] = batch_idx;
+          meta[5] = kv_block_idx;
+        }
+      }
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// PT_smem: SMEM-staged tile prefix scan.
+// Each block handles kRowsPerBlock rows for one head h. Cooperative load
+// of tile_counts[*, h, base_r..base_r+M) into SMEM (better coalescing
+// than per-warp uncoalesced stride reads), then per-warp scan in SMEM,
+// then cooperative store back. Fuses row_ptr into the base.
+// ---------------------------------------------------------------------------
+template <int kThreads, int kRowsPerBlock>
+__global__ void k2q_tile_prefix_smem_kernel(int* __restrict__ tile_counts,
+                                            int const* __restrict__ row_ptr,
+                                            int H, int total_rows,
+                                            int G_total) {
+  static_assert(kRowsPerBlock > 0, "kRowsPerBlock must be positive");
+  extern __shared__ int smem_tprefix[];
+  // smem layout: smem[r_off][g] for r_off in [0, M), g in [0, G_total).
+
+  int tid = threadIdx.x;
+  int lane = tid & 31;
+  int warp_id = tid >> 5;
+
+  // Grid: H * blocks_per_h. Each block stays within a single head h
+  // and processes kRowsPerBlock contiguous rows starting at b_in_h *
+  // kRowsPerBlock. (Earlier flat-grid mapping `h = block_job /
+  // total_rows; base_r = block_job - h*total_rows` skipped rows when
+  // total_rows was not a multiple of kRowsPerBlock and H > 1, because
+  // the last partial block of head h-1 left blocks of head h starting
+  // at a non-zero row offset.)
+  int blocks_per_h = (total_rows + kRowsPerBlock - 1) / kRowsPerBlock;
+  int h = blockIdx.x / blocks_per_h;
+  int b_in_h = blockIdx.x - h * blocks_per_h;
+  if (h >= H) return;
+  int base_r = b_in_h * kRowsPerBlock;
+  if (base_r >= total_rows) return;
+  int actual_M = min(kRowsPerBlock, total_rows - base_r);
+
+  size_t stride_g = (size_t)H * total_rows;
+  int* base_ptr = tile_counts + (size_t)h * total_rows + base_r;
+  int total_elems = G_total * actual_M;
+
+  // Cooperative load. Pattern: thread tid -> (r_off=tid%M, g=tid/M),
+  // then strided. 32 lanes hit M r's × (32/M) g's, giving 32/M cache
+  // lines per warp (vs 32 in the naive stride-along-g pattern).
+  for (int i = tid; i < total_elems; i += kThreads) {
+    int r_off = i % actual_M;
+    int g = i / actual_M;
+    smem_tprefix[r_off * G_total + g] = base_ptr[g * stride_g + r_off];
+  }
+  __syncthreads();
+
+  // Per-warp scan: warp w scans row (base_r + w) if w < actual_M.
+  if (warp_id < actual_M) {
+    int abs_r = base_r + warp_id;
+    int rp = row_ptr[(size_t)h * (total_rows + 1) + abs_r];
+    int* my_smem = smem_tprefix + warp_id * G_total;
+    int running = rp;
+    for (int g0 = 0; g0 < G_total; g0 += kWarpSize) {
+      int g = g0 + lane;
+      int v = (g < G_total) ? my_smem[g] : 0;
+      int x = v;
+#pragma unroll
+      for (int off = 1; off < kWarpSize; off <<= 1) {
+        int nbr = __shfl_up_sync(0xFFFFFFFF, x, off);
+        if (lane >= off) x += nbr;
+      }
+      int excl = running + x - v;
+      if (g < G_total) my_smem[g] = excl;
+      int chunk_sum = __shfl_sync(0xFFFFFFFF, x, 31);
+      running += chunk_sum;
+    }
+  }
+  __syncthreads();
+
+  // Cooperative store back.
+  for (int i = tid; i < total_elems; i += kThreads) {
+    int r_off = i % actual_M;
+    int g = i / actual_M;
+    base_ptr[g * stride_g + r_off] = smem_tprefix[r_off * G_total + g];
+  }
+}
+
+// ---------------------------------------------------------------------------
+// S: scatter. kWarps warps per CTA, each owns q-sub-range. Per-warp SMEM
+// cursor and per-warp tile_offset slot range. Within a warp, q's are
+// processed sequentially; lanes 0..kTopK-1 handle the topK slots in
+// lockstep. Across distinct q's in the same warp, the lockstep ordering
+// guarantees q-monotonic atomicAdd on smem_cursor[r].
+// ---------------------------------------------------------------------------
+// kQPerIter * kTopK lanes are active per warp iter; remaining lanes idle.
+// For kTopK=16, kQPerIter=2 uses all 32 lanes; for kTopK=8, kQPerIter=4.
+// CORRECTNESS NOTE: relies on lane-ordered SMEM atomicAdd return values
+// within a single warp instruction (verified on B200; tests pass).
+//
+// SMEM cursor stored as packed int16 (two cursors per int32). Per-warp
+// row count must stay < 32768 (~q_per_warp * kTopK at max sink), which
+// holds for all task.md sizes up to 1024K.
+template <int kTopK, int kBlockK, int kWarps>
+__global__ void k2q_scatter_kernel(
+    int const* __restrict__ q2k, int const* __restrict__ cu_q,
+    int const* __restrict__ row_map, int const* __restrict__ abs_base,
+    int* __restrict__ q_idx, int* __restrict__ qsplit_idx,
+    int* __restrict__ split_counts, int H, int B, int S_Q, int total_rows,
+    int max_kv_blocks, int q_per_cta, int q_per_warp, int max_seqlen_q) {
+  constexpr int kQPerIter = kWarpSize / kTopK > 0 ? kWarpSize / kTopK : 1;
+  extern __shared__ int smem_cursor_int[];
+  int* smem_cursor = smem_cursor_int;
+  int tid = threadIdx.x;
+  int warp_id = tid >> 5;
+  int lane = tid & 31;
+  int c = blockIdx.x;
+  int q_start_cta = c * q_per_cta;
+  int q_end_cta = min(q_start_cta + q_per_cta, S_Q);
+  int q_start_warp = min(q_start_cta + warp_id * q_per_warp, q_end_cta);
+  int q_end_warp = min(q_start_warp + q_per_warp, q_end_cta);
+
+  int q_in_iter = lane / kTopK;
+  int slot_in_q = lane % kTopK;
+  bool lane_active = (lane < kQPerIter * kTopK);
+
+  // Per-warp packed cursor: total_rows int16 entries -> ceil(total_rows/2)
+  // int32.
+  int packed_per_warp = (total_rows + 1) >> 1;
+  int* my_cursor = smem_cursor + warp_id * packed_per_warp;
+
+  for (int h = 0; h < H; ++h) {
+    for (int i = lane; i < packed_per_warp; i += kWarpSize) my_cursor[i] = 0;
+    __syncwarp();
+
+    if (q_start_warp < q_end_warp) {
+      int bi = 0;
+      advance_batch_only(cu_q, B, q_start_warp, bi);
+
+      int const* head_q2k = q2k + (size_t)h * S_Q * kTopK;
+      int const* my_abs_base =
+          abs_base + ((size_t)(c * kWarps + warp_id) * H + h) * total_rows;
+      int* head_qidx = q_idx + (size_t)h * S_Q * kTopK;
+
+      // (Hot-row register cache experiment showed no measurable
+      // benefit; relying on L1 to keep row 0 / row total_rows-1
+      // hot since they're hit every iteration in sink workloads.)
+
+      constexpr int kUnroll = 16;
+      int qi_base = q_start_warp;
+      for (; qi_base + kUnroll * kQPerIter <= q_end_warp;
+           qi_base += kUnroll * kQPerIter) {
+        int kvb[kUnroll];
+        int qloc[kUnroll];
+        int batch[kUnroll];
+        int const* rmap[kUnroll];
+
+#pragma unroll
+        for (int u = 0; u < kUnroll; ++u) {
+          int qi_u = qi_base + u * kQPerIter + q_in_iter;
+          kvb[u] = -1;
+          qloc[u] = 0;
+          batch[u] = 0;
+          if (lane_active) {
+            advance_batch_only(cu_q, B, qi_u, bi);
+            qloc[u] = qi_u - cu_q[bi];
+            batch[u] = bi;
+            kvb[u] = head_q2k[(size_t)qi_u * kTopK + slot_in_q];
+          }
+          rmap[u] = row_map + (size_t)bi * max_kv_blocks;
+        }
+
+        int row[kUnroll];
+#pragma unroll
+        for (int u = 0; u < kUnroll; ++u) {
+          row[u] = -1;
+          if (lane_active && kvb[u] >= 0 && kvb[u] < max_kv_blocks)
+            row[u] = rmap[u][kvb[u]];
+        }
+
+        // Pre-issue all kUnroll abs_base loads in parallel before
+        // the atomic chain so memory pipeline runs concurrently
+        // with SMEM atomic-adds.
+        int abs_v[kUnroll];
+#pragma unroll
+        for (int u = 0; u < kUnroll; ++u) {
+          abs_v[u] =
+              (row[u] >= 0 && row[u] < total_rows) ? my_abs_base[row[u]] : 0;
+        }
+
+#pragma unroll
+        for (int u = 0; u < kUnroll; ++u) {
+          int r = row[u];
+          bool valid_edge = r >= 0 && r < total_rows;
+          unsigned int valid_mask = __ballot_sync(0xFFFFFFFFu, valid_edge);
+          unsigned int group_mask =
+              (kTopK == 32) ? 0xFFFFFFFFu
+                            : (((1u << kTopK) - 1u) << (q_in_iter * kTopK));
+          unsigned int lower_lane_mask = lane == 0 ? 0u : ((1u << lane) - 1u);
+          int split_slot = __popc(valid_mask & group_mask & lower_lane_mask);
+          int valid_count = __popc(valid_mask & group_mask);
+          if (split_counts != nullptr && slot_in_q == 0) {
+            split_counts[((size_t)batch[u] * max_seqlen_q + qloc[u]) * H + h] =
+                valid_count;
+          }
+          if (valid_edge) {
+            int slot = atomic_inc_int16_packed(my_cursor, r);
+            int out_pos = abs_v[u] + slot;
+            head_qidx[out_pos] = qloc[u];
+            if (qsplit_idx != nullptr) {
+              qsplit_idx[(size_t)h * S_Q * kTopK + out_pos] =
+                  qloc[u] | ((split_slot & 0xFF) << 24);
+            }
+          }
+        }
+      }
+      // Tail: 1-3 iters left.
+      for (; qi_base < q_end_warp; qi_base += kQPerIter) {
+        int my_qi = qi_base + q_in_iter;
+        bool valid_q = (my_qi < q_end_warp) && lane_active;
+        int kvb_local = -1;
+        int q_local = 0;
+        int batch_local = 0;
+        if (valid_q) {
+          advance_batch_only(cu_q, B, my_qi, bi);
+          batch_local = bi;
+          q_local = my_qi - cu_q[bi];
+          kvb_local = head_q2k[(size_t)my_qi * kTopK + slot_in_q];
+        }
+        int const* my_row_map = row_map + (size_t)bi * max_kv_blocks;
+        int row = -1;
+        if (valid_q && kvb_local >= 0 && kvb_local < max_kv_blocks) {
+          row = my_row_map[kvb_local];
+        }
+        bool valid_edge = row >= 0 && row < total_rows;
+        unsigned int valid_mask = __ballot_sync(0xFFFFFFFFu, valid_edge);
+        unsigned int group_mask =
+            (kTopK == 32) ? 0xFFFFFFFFu
+                          : (((1u << kTopK) - 1u) << (q_in_iter * kTopK));
+        unsigned int lower_lane_mask = lane == 0 ? 0u : ((1u << lane) - 1u);
+        int split_slot = __popc(valid_mask & group_mask & lower_lane_mask);
+        int valid_count = __popc(valid_mask & group_mask);
+        if (split_counts != nullptr && valid_q && slot_in_q == 0) {
+          split_counts[((size_t)batch_local * max_seqlen_q + q_local) * H + h] =
+              valid_count;
+        }
+        if (valid_edge) {
+          int slot = atomic_inc_int16_packed(my_cursor, row);
+          int out_pos = my_abs_base[row] + slot;
+          head_qidx[out_pos] = q_local;
+          if (qsplit_idx != nullptr) {
+            qsplit_idx[(size_t)h * S_Q * kTopK + out_pos] =
+                q_local | ((split_slot & 0xFF) << 24);
+          }
+        }
+      }
+    }
+    if (h + 1 < H) __syncthreads();
+  }
+}
+
+}  // anonymous namespace
+
+// ===========================================================================
+// Host orchestration
+// ===========================================================================
+
+template <int kTopK, int kBlockK>
+static void launch_pipeline(torch::Tensor q2k, torch::Tensor cu_q,
+                            torch::Tensor cu_k, torch::Tensor row_ptr,
+                            torch::Tensor q_idx, int total_rows,
+                            int max_kv_blocks,
+                            torch::Tensor scheduler_metadata = torch::Tensor(),
+                            torch::Tensor work_count = torch::Tensor(),
+                            torch::Tensor qsplit_idx = torch::Tensor(),
+                            torch::Tensor split_counts = torch::Tensor(),
+                            int target_q_per_cta = 1, int work_capacity = 0,
+                            int max_seqlen_q = 0) {
+  int H = (int)q2k.size(0);
+  int S_Q = (int)q2k.size(1);
+  int topK = (int)q2k.size(2);
+  TORCH_CHECK(topK == kTopK, "topK runtime != template kTopK");
+  int B = (int)cu_q.size(0) - 1;
+  auto device = q2k.device();
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  AT_CUDA_CHECK(cudaMemsetAsync(row_ptr.data_ptr<int>(), 0,
+                                (size_t)H * (total_rows + 1) * sizeof(int),
+                                stream));
+  AT_CUDA_CHECK(cudaMemsetAsync(q_idx.data_ptr<int>(), 0xFF,
+                                (size_t)H * S_Q * kTopK * sizeof(int), stream));
+
+  auto opts = torch::TensorOptions().dtype(torch::kInt32).device(device);
+  auto row_counts = torch::zeros({H, total_rows}, opts);
+  auto row_map = torch::empty({B, max_kv_blocks}, opts);
+  bool emit_schedule = scheduler_metadata.defined();
+  auto row_coords =
+      emit_schedule ? torch::empty({total_rows, 2}, opts) : torch::Tensor();
+  int* scheduler_metadata_ptr =
+      emit_schedule ? scheduler_metadata.data_ptr<int>() : nullptr;
+  int* work_count_ptr = emit_schedule ? work_count.data_ptr<int>() : nullptr;
+  int* qsplit_idx_ptr = emit_schedule ? qsplit_idx.data_ptr<int>() : nullptr;
+  int* split_counts_ptr =
+      emit_schedule ? split_counts.data_ptr<int>() : nullptr;
+  int* row_coords_ptr = emit_schedule ? row_coords.data_ptr<int>() : nullptr;
+  if (emit_schedule) {
+    AT_CUDA_CHECK(cudaMemsetAsync(work_count_ptr, 0, sizeof(int), stream));
+    AT_CUDA_CHECK(cudaMemsetAsync(scheduler_metadata_ptr, 0,
+                                  (size_t)work_capacity * 6 * sizeof(int),
+                                  stream));
+  }
+
+  int dev = q2k.get_device();
+  int num_sms = 0;
+  AT_CUDA_CHECK(
+      cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev));
+
+  // -- Pick kWarps per CTA based on SMEM budget for cursor/hist ---------
+  // SMEM per CTA = kWarps * total_rows * sizeof(int) (for both H and S).
+  // Want at least 2 CTAs/SM for memory parallelism. SM100 SMEM = 228KB.
+  // Pick the largest kWarps that fits two CTAs/SM, capped at 4.
+  // SMEM cursor packed as int16 (2 entries per int32 word):
+  int per_warp_smem = ((total_rows + 1) >> 1) * (int)sizeof(int);
+  int kWarps_pick = 4;
+  while (kWarps_pick > 1 && (kWarps_pick * per_warp_smem) * 2 > 228 * 1024) {
+    kWarps_pick >>= 1;
+  }
+  if (kWarps_pick < 1) kWarps_pick = 1;
+
+  // -- Pick G (CTAs) ----------------------------------------------------
+  // For each (kWarps, per_warp_smem) pair, the SMEM-bound occupancy is
+  // 228KB / (kWarps*per_warp_smem) CTAs/SM. We size G as
+  // num_sms * occupancy so a single resident wave covers all CTAs and
+  // the memory pipeline runs at peak.
+  int per_cta_smem_bytes = kWarps_pick * per_warp_smem;
+  int max_ctas_per_sm =
+      std::max(1, (228 * 1024) / std::max(1, per_cta_smem_bytes));
+  if (max_ctas_per_sm > 8) max_ctas_per_sm = 8;
+  constexpr int kMinQPerCta = 256;
+  // Cap target_g at num_sms * 3 — empirically this balances
+  // per-CTA work-size against parallelism. Higher caps regress
+  // mid-size cases due to row_counts atomicAdd contention and
+  // smaller q_per_cta. SMEM-bound configurations naturally cap
+  // lower if max_ctas_per_sm < 3.
+  int target_g = num_sms * std::min(max_ctas_per_sm, 3);
+  int max_g_for_q = (S_Q + kMinQPerCta - 1) / kMinQPerCta;
+  int G = std::min({target_g, max_g_for_q, S_Q});
+  if (G < 1) G = 1;
+  int q_per_cta = (S_Q + G - 1) / G;
+  G = (S_Q + q_per_cta - 1) / q_per_cta;
+  int q_per_warp = (q_per_cta + kWarps_pick - 1) / kWarps_pick;
+  int G_total = G * kWarps_pick;
+
+  auto tile_counts = torch::empty({G_total, H, total_rows}, opts);
+
+  // -- Compile-time switch on kWarps for the templated kernels ---------
+  auto rmap_fn = k2q_build_row_map_kernel<kBlockK>;
+  auto rprefix_fn = k2q_row_prefix_kernel<1024>;
+  constexpr int kPtRowsPerBlock = 8;
+  constexpr int kPtThreads = 256;
+  auto tprefix_smem_fn =
+      k2q_tile_prefix_smem_kernel<kPtThreads, kPtRowsPerBlock>;
+
+  if (max_kv_blocks > 0) {
+    rmap_fn<<<max_kv_blocks, 32, 0, stream>>>(cu_k.data_ptr<int>(),
+                                              row_map.data_ptr<int>(),
+                                              row_coords_ptr, B, max_kv_blocks);
+  }
+
+  auto launch_hist_scatter = [&](auto kWarps_const) {
+    constexpr int W = decltype(kWarps_const)::value;
+    size_t smem_bytes = (size_t)W * per_warp_smem;
+    auto hist_fn = k2q_hist_kernel<kTopK, kBlockK, W>;
+    auto scat_fn = k2q_scatter_kernel<kTopK, kBlockK, W>;
+    AT_CUDA_CHECK(cudaFuncSetAttribute(
+        hist_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_bytes));
+    AT_CUDA_CHECK(cudaFuncSetAttribute(
+        scat_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_bytes));
+
+    hist_fn<<<G, W * kWarpSize, smem_bytes, stream>>>(
+        q2k.data_ptr<int>(), cu_q.data_ptr<int>(), row_map.data_ptr<int>(),
+        row_counts.data_ptr<int>(), tile_counts.data_ptr<int>(), H, B, S_Q,
+        total_rows, max_kv_blocks, q_per_cta, q_per_warp);
+
+    rprefix_fn<<<H, 1024, 0, stream>>>(
+        row_counts.data_ptr<int>(), row_ptr.data_ptr<int>(),
+        emit_schedule ? row_coords.data_ptr<int>() : nullptr,
+        scheduler_metadata_ptr, work_count_ptr, total_rows, target_q_per_cta,
+        work_capacity);
+
+    // Grid is H * blocks_per_h so each block stays within a single
+    // head; flat (H*total_rows) grid would skip rows when total_rows
+    // is not a multiple of kPtRowsPerBlock.
+    int blocks_per_h = (total_rows + kPtRowsPerBlock - 1) / kPtRowsPerBlock;
+    int pt_grid = H * blocks_per_h;
+    if (pt_grid < 1) pt_grid = 1;
+    size_t pt_smem = (size_t)kPtRowsPerBlock * G_total * sizeof(int);
+    AT_CUDA_CHECK(cudaFuncSetAttribute(
+        tprefix_smem_fn, cudaFuncAttributeMaxDynamicSharedMemorySize,
+        (int)pt_smem));
+    tprefix_smem_fn<<<pt_grid, kPtThreads, pt_smem, stream>>>(
+        tile_counts.data_ptr<int>(), row_ptr.data_ptr<int>(), H, total_rows,
+        G_total);
+
+    scat_fn<<<G, W * kWarpSize, smem_bytes, stream>>>(
+        q2k.data_ptr<int>(), cu_q.data_ptr<int>(), row_map.data_ptr<int>(),
+        tile_counts.data_ptr<int>(), q_idx.data_ptr<int>(), qsplit_idx_ptr,
+        split_counts_ptr, H, B, S_Q, total_rows, max_kv_blocks, q_per_cta,
+        q_per_warp, max_seqlen_q);
+  };
+
+  if (kWarps_pick == 4) {
+    launch_hist_scatter(std::integral_constant<int, 4>{});
+  } else if (kWarps_pick == 2) {
+    launch_hist_scatter(std::integral_constant<int, 2>{});
+  } else {
+    launch_hist_scatter(std::integral_constant<int, 1>{});
+  }
+}
+
+void run_minimax_m3_build_k2q_csr_with_schedule(
+    torch::Tensor q2k, torch::Tensor cu_q, torch::Tensor cu_k,
+    torch::Tensor row_ptr, torch::Tensor q_idx,
+    torch::Tensor scheduler_metadata, torch::Tensor work_count,
+    torch::Tensor qsplit_idx, torch::Tensor split_counts, int64_t topk,
+    int64_t blk_kv, int64_t total_rows, int64_t max_kv_blocks,
+    int64_t target_q_per_cta, int64_t work_capacity, int64_t max_seqlen_q) {
+  CHECK_INPUT(q2k);
+  CHECK_INPUT(cu_q);
+  CHECK_INPUT(cu_k);
+  CHECK_INPUT(row_ptr);
+  CHECK_INPUT(q_idx);
+  CHECK_INPUT(scheduler_metadata);
+  CHECK_INPUT(work_count);
+  CHECK_INPUT(qsplit_idx);
+  CHECK_INPUT(split_counts);
+  TORCH_CHECK(blk_kv == 128, "build_k2q_csr only supports blk_kv == 128");
+  int H = (int)q2k.size(0);
+  int S_Q = (int)q2k.size(1);
+  int tr = (int)total_rows;
+  int mkv = (int)max_kv_blocks;
+  int target = (int)target_q_per_cta;
+  int capacity = (int)work_capacity;
+  int max_sq = (int)max_seqlen_q;
+  TORCH_CHECK(tr >= 0 && mkv >= 0 && target > 0 && capacity > 0 && max_sq >= 0,
+              "invalid schedule sizing arguments");
+  TORCH_CHECK(row_ptr.size(0) == H && row_ptr.size(1) == tr + 1,
+              "row_ptr shape mismatch");
+  TORCH_CHECK(q_idx.size(0) == H && q_idx.size(1) == (int64_t)S_Q * (int)topk,
+              "q_idx shape mismatch");
+  TORCH_CHECK(qsplit_idx.sizes() == q_idx.sizes(), "qsplit_idx shape mismatch");
+  TORCH_CHECK(
+      scheduler_metadata.size(0) == capacity && scheduler_metadata.size(1) == 6,
+      "scheduler_metadata shape mismatch");
+  TORCH_CHECK(work_count.numel() == 1,
+              "work_count must have one int32 element");
+  TORCH_CHECK(split_counts.dim() == 3 &&
+                  split_counts.size(0) == cu_q.size(0) - 1 &&
+                  split_counts.size(1) == max_sq && split_counts.size(2) == H,
+              "split_counts shape mismatch");
+  if (S_Q == 0 || tr == 0 || H == 0 || mkv == 0) {
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    AT_CUDA_CHECK(cudaMemsetAsync(row_ptr.data_ptr<int>(), 0,
+                                  (size_t)H * (tr + 1) * sizeof(int), stream));
+    AT_CUDA_CHECK(cudaMemsetAsync(q_idx.data_ptr<int>(), 0xFF,
+                                  (size_t)H * S_Q * (int)topk * sizeof(int),
+                                  stream));
+    AT_CUDA_CHECK(
+        cudaMemsetAsync(work_count.data_ptr<int>(), 0, sizeof(int), stream));
+    if (split_counts.numel() > 0) {
+      AT_CUDA_CHECK(cudaMemsetAsync(split_counts.data_ptr<int>(), 0,
+                                    (size_t)split_counts.numel() * sizeof(int),
+                                    stream));
+    }
+    return;
+  }
+
+  if (topk == 16) {
+    launch_pipeline<16, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
+                             scheduler_metadata, work_count, qsplit_idx,
+                             split_counts, target, capacity, max_sq);
+  } else if (topk == 8) {
+    launch_pipeline<8, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
+                            scheduler_metadata, work_count, qsplit_idx,
+                            split_counts, target, capacity, max_sq);
+  } else if (topk == 32) {
+    launch_pipeline<32, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
+                             scheduler_metadata, work_count, qsplit_idx,
+                             split_counts, target, capacity, max_sq);
+  } else if (topk == 4) {
+    launch_pipeline<4, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
+                            scheduler_metadata, work_count, qsplit_idx,
+                            split_counts, target, capacity, max_sq);
+  } else {
+    TORCH_CHECK(false, "unsupported topK ", topk,
+                " (expected 4, 8, 16, or 32)");
+  }
+}

csrc/ops.h

@@ -62,7 +62,8 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
 
 void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
 
-void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
+void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
+                        double alpha = 1.0, double beta = 0.0);
 
 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
                         torch::Tensor& scale);
@@ -146,4 +147,12 @@ std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
     torch::Tensor const& norm_weight_k, torch::Tensor workspace,
     int64_t const q_size, int64_t const kv_size, int64_t const rank,
     int64_t const nranks, double const eps);
+
+void run_minimax_m3_build_k2q_csr_with_schedule(
+    torch::Tensor q2k, torch::Tensor cu_q, torch::Tensor cu_k,
+    torch::Tensor row_ptr, torch::Tensor q_idx,
+    torch::Tensor scheduler_metadata, torch::Tensor work_count,
+    torch::Tensor qsplit_idx, torch::Tensor split_counts, int64_t topk,
+    int64_t blk_kv, int64_t total_rows, int64_t max_kv_blocks,
+    int64_t target_q_per_cta, int64_t work_capacity, int64_t max_seqlen_q);
 #endif

csrc/torch_bindings.cpp

@@ -187,6 +187,27 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "float eps) -> (Tensor, Tensor)");
   ops.impl("minimax_allreduce_rms_qk", torch::kCUDA, &minimax_allreduce_rms_qk);
 
+  ops.def(
+      "minimax_m3_build_k2q_csr_with_schedule("
+      "Tensor q2k,"
+      "Tensor cu_q,"
+      "Tensor cu_k,"
+      "Tensor! row_ptr,"
+      "Tensor! q_idx,"
+      "Tensor! scheduler_metadata,"
+      "Tensor! work_count,"
+      "Tensor! qsplit_idx,"
+      "Tensor! split_counts,"
+      "int topk,"
+      "int blk_kv,"
+      "int total_rows,"
+      "int max_kv_blocks,"
+      "int target_q_per_cta,"
+      "int work_capacity,"
+      "int max_seqlen_q) -> ()");
+  ops.impl("minimax_m3_build_k2q_csr_with_schedule", torch::kCUDA,
+           &run_minimax_m3_build_k2q_csr_with_schedule);
+
   //  conditionally compiled so impl in source file
 #endif
 }

docker/Dockerfile

@@ -757,7 +757,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.11.post2
+ARG FLASHINFER_VERSION=0.6.12
 RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
         --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

docker/Dockerfile.nightly_torch

@@ -256,13 +256,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
 
 
 # build flashinfer for torch nightly from source around 10 mins
-# release version: v0.6.11.post2
+# release version: v0.6.12
 # todo(elainewy): cache flashinfer build result for faster build
 ENV CCACHE_DIR=/root/.cache/ccache
 RUN --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=cache,target=/root/.cache/uv \
     echo "git clone flashinfer..." \
-    && git clone --depth 1 --branch v0.6.11.post2 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && git clone --depth 1 --branch v0.6.12 --recursive https://github.com/flashinfer-ai/flashinfer.git \
     && cd flashinfer \
     && git submodule update --init --recursive \
     && echo "finish git clone flashinfer..." \

docker/versions.json

@@ -68,7 +68,7 @@
       "default": "true"
     },
     "FLASHINFER_VERSION": {
-      "default": "0.6.11.post2"
+      "default": "0.6.12"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"

docs/design/attention_backends.md

@@ -170,8 +170,8 @@ Priority is **1 = highest** (tried first).
 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
 | `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 |
-| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
+| `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 |
+| `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 |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
 | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
@@ -187,6 +187,18 @@ Priority is **1 = highest** (tried first).
 >
 > **\*** 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.
 
+## MiniMax M3 Sparse Attention Backends
+
+Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")
+layers. It is wired in directly by the model and is not part of the
+automatic priority lists above. A lightning indexer scores KV blocks, the
+top-k blocks (plus fixed init/local blocks) are selected, and attention
+attends only to those blocks; index keys live in a separate side cache.
+
+| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
+| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
+| `MINIMAX_M3_SPARSE` | bf16, fp16 | `bfloat16` | 128 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
+
 ## MLA (Multi-head Latent Attention) Backends
 
 MLA uses separate backends for prefill and decode phases.

pyproject.toml

@@ -162,7 +162,12 @@ dout = "dout"
 Pn = "Pn"
 arange = "arange"
 thw = "thw"
+# temporal position ids (parallels hpos/wpos in vision RoPE)
+tpos = "tpos"
 subtile = "subtile"
+subtiles = "subtiles"
+reord = "reord"
+Ot = "Ot"
 HSA = "HSA"
 setp = "setp"
 CPY = "CPY"

requirements/common.txt

@@ -29,6 +29,7 @@ xgrammar >= 0.2.0, < 1.0.0; platform_machine == "x86_64" or platform_machine ==
 typing_extensions >= 4.10
 filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
 partial-json-parser # used for parsing partial JSON outputs
+jsonschema >= 4.23.0 # required for MiniMax M3 tool schema validation
 pyzmq >= 25.0.0
 msgspec
 gguf >= 0.17.0

requirements/cuda.txt

@@ -9,8 +9,8 @@ torchaudio==2.11.0
 # These must be updated alongside torch
 torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 # FlashInfer should be updated together with the Dockerfile
-flashinfer-python==0.6.11.post2
-flashinfer-cubin==0.6.11.post2
+flashinfer-python==0.6.12
+flashinfer-cubin==0.6.12
 apache-tvm-ffi==0.1.9
 tilelang==0.1.9
 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to

requirements/test/cuda.txt

@@ -360,6 +360,7 @@ jsonpointer==3.0.0
     # via jsonschema
 jsonschema==4.23.0
     # via
+    #   -c requirements/common.txt
     #   hypothesis-jsonschema
     #   mistral-common
     #   ray

requirements/test/rocm.txt

@@ -440,6 +440,8 @@ jsonpointer==3.1.0
     # via jsonschema
 jsonschema==4.26.0
     # via
+    #   -c requirements/common.txt
+    #   -r requirements/test/../common.txt
     #   hypothesis-jsonschema
     #   mcp
     #   mistral-common

requirements/test/xpu.txt

@@ -229,6 +229,7 @@ jsonlines==4.0.0
     # via lm-eval
 jsonschema==4.26.0
     # via
+    #   -c requirements/common.txt
     #   hypothesis-jsonschema
     #   mistral-common
     #   schemathesis

rust/Cargo.lock

@@ -3458,6 +3458,75 @@ version = "0.1.29"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f"
 
+[[package]]
+name = "pyo3"
+version = "0.28.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12"
+dependencies = [
+ "libc",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.28.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e"
+dependencies = [
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.28.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.28.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn 2.0.117",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.28.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "pyo3-build-config",
+ "quote",
+ "syn 2.0.117",
+]
+
+[[package]]
+name = "pythonize"
+version = "0.28.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b79f670c9626c8b651c0581011b57b6ba6970bb69faf01a7c4c0cfc81c43f95"
+dependencies = [
+ "pyo3",
+ "serde",
+ "serde_json",
+]
+
 [[package]]
 name = "qoi"
 version = "0.4.1"
@@ -4669,6 +4738,12 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "target-lexicon"
+version = "0.13.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
+
 [[package]]
 name = "task-local"
 version = "0.1.1"
@@ -5622,6 +5697,7 @@ dependencies = [
  "expect-test",
  "futures",
  "half",
+ "indexmap 2.13.0",
  "itertools 0.14.0",
  "llm-multimodal",
  "minijinja",
@@ -5900,6 +5976,17 @@ dependencies = [
  "winnow",
 ]
 
+[[package]]
+name = "vllm-tool-parser-py"
+version = "0.1.0"
+dependencies = [
+ "pyo3",
+ "pythonize",
+ "serde_json",
+ "thiserror-ext",
+ "vllm-tool-parser",
+]
+
 [[package]]
 name = "walkdir"
 version = "2.5.0"

rust/Cargo.toml

@@ -12,6 +12,7 @@ members = [
     "src/text",
     "src/tokenizer",
     "src/tool-parser",
+    "src/tool-parser/python",
 ]
 resolver = "3"
 
@@ -43,6 +44,7 @@ half = { version = "2.7.1", features = ["bytemuck"] }
 hex = "0.4.3"
 hf-hub = { version = "0.5.0", features = ["tokio"] }
 http-body = "1.0.1"
+indexmap = "2.13.0"
 itertools = "0.14.0"
 libc = "0.2.177"
 llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "5b558989844d1c7af3e43d0f604069ffd9c06320" }
@@ -59,6 +61,8 @@ prometheus-client = "0.24.0"
 prometheus-client-derive-encode = "0.5.0"
 prost = "0.14.3"
 prost-types = "0.14.3"
+pyo3 = "0.28.3"
+pythonize = "0.28.0"
 rand = "0.9.2"
 reasoning-parser = "1.2.2"
 reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] }
@@ -69,7 +73,7 @@ rustc-hash = "1.1.0"
 serde = { version = "1.0.228", features = ["derive"] }
 serde-json-fmt = "0.1.0"
 serde_default = "0.2.0"
-serde_json = { version = "1.0.145", features = ["arbitrary_precision", "preserve_order"] }
+serde_json = { version = "1.0.145", features = ["preserve_order"] }
 serde_repr = "0.1.20"
 serde_tuple = "1.1.3"
 serde_with = "3.18.0"

rust/src/chat/Cargo.toml

@@ -10,6 +10,7 @@ asynk-strim-attr.workspace = true
 easy-ext.workspace = true
 futures.workspace = true
 half.workspace = true
+indexmap.workspace = true
 itertools.workspace = true
 llm-multimodal.workspace = true
 minijinja.workspace = true

rust/src/chat/src/lib.rs

@@ -233,7 +233,7 @@ mod tests {
         )
         .unwrap_err();
 
-        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, hermes, hy_v3, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
+        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, hermes, hy_v3, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
     }
 
     #[test]

rust/src/chat/src/parser/reasoning/mod.rs

@@ -5,8 +5,9 @@ use std::sync::LazyLock;
 pub use vllm_reasoning_parser::{
     CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
     DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
-    KimiReasoningParser, MiniMaxM2ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser,
-    ReasoningDelta, ReasoningError, ReasoningParser, Step3ReasoningParser,
+    KimiReasoningParser, MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser,
+    NemotronV3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta, ReasoningError,
+    ReasoningParser, Step3ReasoningParser,
 };
 use vllm_tokenizer::DynTokenizer;
 
@@ -23,6 +24,7 @@ pub mod names {
     pub const KIMI: &str = "kimi";
     pub const KIMI_K2: &str = "kimi_k2";
     pub const MINIMAX_M2: &str = "minimax_m2";
+    pub const MINIMAX_M3: &str = "minimax_m3";
     pub const NEMOTRON_V3: &str = "nemotron_v3";
     pub const QWEN3: &str = "qwen3";
     pub const STEP3: &str = "step3";
@@ -59,6 +61,7 @@ impl ReasoningParserFactory {
             .register_parser::<KimiReasoningParser>(names::KIMI)
             .register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
             .register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
+            .register_parser::<MiniMaxM3ReasoningParser>(names::MINIMAX_M3)
             .register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)
             .register_parser::<Qwen3ReasoningParser>(names::QWEN3)
             .register_parser::<Step3ReasoningParser>(names::STEP3);
@@ -78,6 +81,8 @@ impl ReasoningParserFactory {
             .register_pattern("kimi-k2", names::KIMI_K2)
             .register_pattern("kimi", names::KIMI)
             .register_pattern("step3", names::STEP3)
+            .register_pattern("minimax-m3", names::MINIMAX_M3)
+            .register_pattern("mm-m3", names::MINIMAX_M3)
             .register_pattern("minimax", names::MINIMAX_M2)
             .register_pattern("mm-m2", names::MINIMAX_M2)
             .register_pattern("cohere", names::COHERE_CMD)

rust/src/chat/src/parser/reasoning/tests.rs

@@ -32,8 +32,10 @@ fn factory_contains_and_lists_registered_parsers() {
     let factory = ReasoningParserFactory::new();
     assert!(factory.contains(names::QWEN3));
     assert!(factory.contains(names::DEEPSEEK_V4));
+    assert!(factory.contains(names::MINIMAX_M3));
     assert!(factory.list().contains(&names::QWEN3.to_string()));
     assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string()));
+    assert!(factory.list().contains(&names::MINIMAX_M3.to_string()));
 }
 
 #[test]
@@ -49,6 +51,19 @@ fn factory_resolves_deepseek_v4_to_qwen3_alias() {
     );
 }
 
+#[test]
+fn factory_resolves_minimax_m3_before_generic_minimax() {
+    let factory = ReasoningParserFactory::new();
+    assert_eq!(
+        factory.resolve_name_for_model("MiniMaxAI/Minimax-M3-preview"),
+        Some(names::MINIMAX_M3)
+    );
+    assert_eq!(
+        factory.resolve_name_for_model("mm-m3"),
+        Some(names::MINIMAX_M3)
+    );
+}
+
 #[test]
 fn factory_rejects_unknown_parser_names() {
     let tokenizer = Arc::new(FakeTokenizer);

rust/src/chat/src/parser/tool/mod.rs

@@ -5,9 +5,9 @@ use std::sync::LazyLock;
 pub use vllm_tool_parser::{
     DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
     Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, HyV3ToolParser,
-    KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MistralToolParser,
-    Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError,
-    ToolParserOutput,
+    KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MinimaxM3ToolParser,
+    MistralToolParser, Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser,
+    ToolParserError, ToolParserOutput,
 };
 
 use crate::parser::ParserFactory;
@@ -28,6 +28,7 @@ pub mod names {
     pub const LLAMA3_JSON: &str = "llama3_json";
     pub const LLAMA4_JSON: &str = "llama4_json";
     pub const MINIMAX_M2: &str = "minimax_m2";
+    pub const MINIMAX_M3: &str = "minimax_m3";
     pub const MISTRAL: &str = "mistral";
     pub const QWEN3_CODER: &str = "qwen3_coder";
     pub const QWEN3_XML: &str = "qwen3_xml";
@@ -66,6 +67,7 @@ impl ToolParserFactory {
             .register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
             .register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
             .register_parser::<MinimaxM2ToolParser>(names::MINIMAX_M2)
+            .register_parser::<MinimaxM3ToolParser>(names::MINIMAX_M3)
             .register_parser::<MistralToolParser>(names::MISTRAL)
             .register_parser::<Qwen3XmlToolParser>(names::QWEN3_XML)
             .register_parser::<Qwen3CoderToolParser>(names::QWEN3_CODER);
@@ -96,6 +98,8 @@ impl ToolParserFactory {
             .register_pattern("gemma4", names::GEMMA4)
             .register_pattern("gemma-4", names::GEMMA4)
             .register_pattern("kimi-k2", names::KIMI_K2)
+            .register_pattern("minimax-m3", names::MINIMAX_M3)
+            .register_pattern("mm-m3", names::MINIMAX_M3)
             .register_pattern("minimax", names::MINIMAX_M2)
             .register_pattern("mm-m2", names::MINIMAX_M2);
 

rust/src/chat/src/parser/tool/tests.rs

@@ -153,6 +153,14 @@ fn factory_new_resolves_default_patterns() {
         factory.resolve_name_for_model("tencent/Hy3-preview"),
         Some(names::HY_V3)
     );
+    assert_eq!(
+        factory.resolve_name_for_model("MiniMax/MiniMax-M3-Text"),
+        Some(names::MINIMAX_M3)
+    );
+    assert_eq!(
+        factory.resolve_name_for_model("org/mm-m3-base"),
+        Some(names::MINIMAX_M3)
+    );
     assert_eq!(
         factory.resolve_name_for_model("MiniMax/MiniMax-M2-01"),
         Some(names::MINIMAX_M2)

rust/src/chat/src/renderer/hf/mod.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 
 use serde::Serialize;
-use serde_json::Value;
+use serde_json::Value as JsonValue;
 use thiserror_ext::AsReport as _;
 use tracing::{info, trace, warn};
 use vllm_text::Prompt;
@@ -13,6 +13,7 @@ use self::format::{
     ChatTemplateContentFormat, ChatTemplateContentFormatOption as ContentFormatOption,
 };
 use self::template::{CompiledChatTemplate, TemplateContext};
+use self::value::{TemplateValue, to_template_value};
 use super::{ChatRenderer, RenderedPrompt};
 use crate::error::Result;
 use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest};
@@ -24,6 +25,7 @@ mod error;
 mod format;
 mod template;
 mod tojson;
+mod value;
 
 pub use template::{load_chat_template, resolve_chat_template};
 
@@ -38,7 +40,7 @@ pub struct MultimodalRenderInfo {
 /// state.
 pub struct HfChatRenderer {
     default_template: Option<CompiledChatTemplate>,
-    default_template_kwargs: HashMap<String, Value>,
+    default_template_kwargs: HashMap<String, JsonValue>,
     content_format: ContentFormatOption,
     special_tokens: Option<HfSpecialTokens>,
     multimodal: Option<MultimodalRenderInfo>,
@@ -48,7 +50,7 @@ impl HfChatRenderer {
     /// Create a renderer from the given template string.
     pub fn new(
         template: Option<String>,
-        default_template_kwargs: HashMap<String, Value>,
+        default_template_kwargs: HashMap<String, JsonValue>,
         content_format: ContentFormatOption,
     ) -> Result<Self> {
         Ok(Self {
@@ -245,7 +247,7 @@ struct TemplateToolCall {
 #[derive(Debug, Serialize)]
 struct TemplateToolFunction {
     name: String,
-    arguments: Value,
+    arguments: TemplateValue,
 }
 
 #[derive(Debug, Serialize)]
@@ -259,7 +261,7 @@ pub(super) struct TemplateTool {
 struct TemplateToolDefinition {
     name: String,
     description: Option<String>,
-    parameters: Value,
+    parameters: TemplateValue,
     strict: Option<bool>,
 }
 
@@ -345,13 +347,14 @@ fn to_template_tool_calls(
     let mut tool_calls = Vec::new();
 
     for tool_call in content.tool_calls() {
-        let arguments = serde_json::from_str::<Value>(&tool_call.arguments).map_err(|error| {
+        let arguments = serde_json::from_str(&tool_call.arguments).map_err(|error| {
             Error::ChatTemplate(format!(
                 "assistant tool call `{}` has invalid JSON arguments: {}",
                 tool_call.id,
                 error.as_report()
             ))
         })?;
+        let arguments = to_template_value(arguments);
 
         tool_calls.push(TemplateToolCall {
             id: tool_call.id.clone(),
@@ -434,7 +437,7 @@ fn to_template_tools(tools: &[ChatTool]) -> Vec<TemplateTool> {
             function: TemplateToolDefinition {
                 name: tool.name.clone(),
                 description: tool.description.clone(),
-                parameters: tool.parameters.clone(),
+                parameters: to_template_value(tool.parameters.clone()),
                 strict: tool.strict,
             },
         })
@@ -909,6 +912,29 @@ mod tests {
         assert_eq!(rendered, "get_weather|Paris|call_1|Sunny");
     }
 
+    #[test]
+    fn chat_template_tool_call_argument_items_method_is_not_shadowed_by_field() {
+        let request = sample_request(vec![ChatMessage::assistant_blocks(vec![
+            AssistantContentBlock::ToolCall(crate::AssistantToolCall {
+                id: "call_1".to_string(),
+                name: "add".to_string(),
+                arguments: r#"{"items":"operands","x":2,"y":1.0}"#.to_string(),
+            }),
+        ])]);
+
+        let rendered = render(
+            Some(
+                "{%- set arguments = messages[0].tool_calls[0].function.arguments -%}
+{%- for key, value in arguments.items() -%}{{ key }}={{ value }};{%- endfor -%}
+|{{ arguments['items'] }}",
+            ),
+            &request,
+        )
+        .unwrap();
+
+        assert_eq!(rendered, "items=operands;x=2;y=1.0;|operands");
+    }
+
     #[test]
     fn qwen35_template_renders_prefilled_reasoning_start_when_thinking_enabled() {
         let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);

rust/src/chat/src/renderer/hf/tojson.rs

@@ -208,11 +208,27 @@ mod tests {
     }
 
     #[test]
-    fn tojson_preserves_arbitrary_precision_number_spelling() {
+    fn tojson_uses_standard_serde_json_number_spelling() {
         let payload = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap();
         let rendered = render("{{ payload|tojson }}", payload);
 
-        assert_eq!(rendered, "{\"x\": 2, \"y\": 1.00}");
+        // TODO: we cannot preserve the original number precision by enabling `serde_json`'s
+        // `arbitrary_precision` feature, otherwise the following test
+        // `serialized_json_numbers_do_not_leak_serde_private_representation` will fail.
+        // See issue: https://github.com/mitsuhiko/minijinja/issues/641
+        assert_eq!(rendered, "{\"x\": 2, \"y\": 1.0}");
+    }
+
+    #[test]
+    fn serialized_json_numbers_do_not_leak_serde_private_representation() {
+        let payload: serde_json::Value = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap();
+        let rendered = render("{{ payload }}", payload);
+
+        // TODO: we cannot preserve the original number precision by enabling `serde_json`'s
+        // `arbitrary_precision` feature, otherwise this will fail.
+        // See issue: https://github.com/mitsuhiko/minijinja/issues/641
+        assert!(!rendered.contains("$serde_json::private::Number"));
+        assert_eq!(rendered, r#"{"x": 2, "y": 1.0}"#);
     }
 
     #[test]

rust/src/chat/src/renderer/hf/value.rs

@@ -0,0 +1,77 @@
+use std::sync::Arc;
+
+use indexmap::IndexMap;
+use minijinja::value::{Enumerator, Object, ObjectExt, ObjectRepr};
+use minijinja::{Error as TemplateError, ErrorKind as TemplateErrorKind, State};
+use serde::Serialize;
+use serde_json::Value as JsonValue;
+
+/// A wrapper around `minijinja::Value` that can be constructed with `to_template_value` and used
+/// as a value in the chat template.
+#[derive(Debug, Serialize)]
+#[serde(transparent)]
+pub(super) struct TemplateValue(minijinja::Value);
+
+pub(super) fn to_template_value(value: JsonValue) -> TemplateValue {
+    TemplateValue(match value {
+        JsonValue::Array(values) => values
+            .into_iter()
+            .map(to_template_value)
+            .map(|value| value.0)
+            .collect::<minijinja::Value>(),
+        JsonValue::Object(values) => minijinja::Value::from_object(TemplateMap(
+            values
+                .into_iter()
+                .map(|(key, value)| (key, to_template_value(value).0))
+                .collect(),
+        )),
+        // For primitive values, directly convert them to `minijinja::Value` using `from_serialize`.
+        value => minijinja::Value::from_serialize(value),
+    })
+}
+
+/// A custom map type that always returns `UnknownMethod` for method calls, so that pycompat can
+/// always handle dict methods through the unknown-method callback.
+///
+/// Use `IndexMap` to preserve the original key order when iterating.
+///
+/// MiniJinja's default map can resolve a same-named field before Python dict methods. HF templates
+/// commonly call `dict.items()`, which would fail if the map had an `items` field.
+/// See issue: https://github.com/mitsuhiko/minijinja/issues/903
+#[derive(Debug)]
+struct TemplateMap(IndexMap<String, minijinja::Value>);
+
+impl Object for TemplateMap {
+    fn repr(self: &Arc<Self>) -> ObjectRepr {
+        ObjectRepr::Map
+    }
+
+    fn get_value(self: &Arc<Self>, key: &minijinja::Value) -> Option<minijinja::Value> {
+        self.0.get(key.as_str()?).cloned()
+    }
+
+    fn get_value_by_str(self: &Arc<Self>, key: &str) -> Option<minijinja::Value> {
+        self.0.get(key).cloned()
+    }
+
+    fn enumerate(self: &Arc<Self>) -> Enumerator {
+        self.mapped_rev_enumerator(|this| {
+            Box::new(this.0.keys().map(|key| minijinja::Value::from(key.as_str())))
+        })
+    }
+
+    fn enumerator_len(self: &Arc<Self>) -> Option<usize> {
+        Some(self.0.len())
+    }
+
+    fn call_method(
+        self: &Arc<Self>,
+        _state: &State<'_, '_>,
+        _method: &str,
+        _args: &[minijinja::Value],
+    ) -> std::result::Result<minijinja::Value, TemplateError> {
+        // Always return `UnknownMethod` for method calls,
+        // so that pycompat can handle dict methods through the unknown-method callback.
+        Err(TemplateError::from(TemplateErrorKind::UnknownMethod))
+    }
+}

rust/src/chat/tests/roundtrip.rs

@@ -183,7 +183,7 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
         "roundtrip-reasoning-tools",
         vec![ChatMessage::text(
             ChatRole::User,
-            "Check Shanghai weather and add 1.00 plus 2.",
+            "Check Shanghai weather and add 1.0 plus 2.",
         )],
         test_tools(),
     );
@@ -210,9 +210,10 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
                 AssistantContentBlock::ToolCall(AssistantToolCall {
                     id: "functions.add:1".to_string(),
                     name: "add".to_string(),
-                    // Intentionally use a non-lexical order of keys and a different number
-                    // formatting style to verify text-level fidelity of the roundtrip.
-                    arguments: r#"{"y":1.00,"x":2}"#.to_string(),
+                    // Intentionally use a non-lexical order of keys to verify text-level
+                    // fidelity of the roundtrip where JSON formatting remains stable. The
+                    // `items` key also exercises templates that call `arguments.items()`.
+                    arguments: r#"{"y":1.0,"x":2,"items":["left","right"]}"#.to_string(),
                 }),
             ],
         },
@@ -240,7 +241,7 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
     assert_eq!(tool_calls[1].name, "add");
     assert_eq!(
         tool_calls[1].arguments,
-        expected_arguments(&case, r#"{"y": 1.00, "x": 2}"#)?,
+        expected_arguments(&case, r#"{"y": 1.0, "x": 2, "items": ["left", "right"]}"#)?,
     );
 
     assert_eq!(
@@ -531,9 +532,13 @@ fn test_tools() -> Vec<ChatTool> {
                 "type": "object",
                 "properties": {
                     "y": { "type": "number" },
-                    "x": { "type": "number" }
+                    "x": { "type": "number" },
+                    "items": {
+                        "type": "array",
+                        "items": { "type": "string" }
+                    }
                 },
-                "required": ["y", "x"]
+                "required": ["y", "x", "items"]
             }),
             strict: None,
         },

rust/src/reasoning-parser/src/lib.rs

@@ -19,6 +19,7 @@ mod deepseek_r1;
 mod delimited;
 mod gemma4;
 mod kimi;
+mod minimax_m3;
 mod qwen3;
 
 use thiserror::Error;
@@ -29,6 +30,7 @@ pub use self::deepseek_r1::DeepSeekR1ReasoningParser;
 pub(crate) use self::delimited::DelimitedReasoningParser;
 pub use self::gemma4::Gemma4ReasoningParser;
 pub use self::kimi::KimiReasoningParser;
+pub use self::minimax_m3::MiniMaxM3ReasoningParser;
 pub use self::qwen3::Qwen3ReasoningParser;
 
 /// DeepSeek V3 currently shares the standard `<think>...</think>` parser.

rust/src/reasoning-parser/src/minimax_m3.rs

@@ -0,0 +1,43 @@
+use vllm_tokenizer::DynTokenizer;
+
+use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
+
+/// Reasoning parser for MiniMax M3 style outputs.
+///
+/// MiniMax M3 uses `<mm:think>...</mm:think>` delimiters. Its chat template may
+/// prefill either delimiter depending on the requested thinking mode, so the
+/// shared delimited parser derives the starting state from the rendered prompt.
+pub struct MiniMaxM3ReasoningParser {
+    inner: DelimitedReasoningParser,
+}
+
+impl MiniMaxM3ReasoningParser {
+    /// Create a MiniMax M3 parser backed by the shared delimited state machine.
+    pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
+        Ok(Self {
+            inner: DelimitedReasoningParser::new(tokenizer, "<mm:think>", "</mm:think>", false)?,
+        })
+    }
+}
+
+impl ReasoningParser for MiniMaxM3ReasoningParser {
+    fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
+    where
+        Self: Sized + 'static,
+    {
+        Ok(Box::new(Self::new(tokenizer)?))
+    }
+
+    fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
+        self.inner.initialize(prompt_token_ids);
+        Ok(())
+    }
+
+    fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
+        Ok(self.inner.push(delta))
+    }
+
+    fn finish(&mut self) -> Result<ReasoningDelta> {
+        Ok(self.inner.finish())
+    }
+}

rust/src/reasoning-parser/src/tests.rs

@@ -3,7 +3,8 @@ use std::sync::Arc;
 use vllm_tokenizer::Tokenizer;
 
 use super::{
-    DeepSeekR1ReasoningParser, DelimitedReasoningParser, Qwen3ReasoningParser, ReasoningParser,
+    DeepSeekR1ReasoningParser, DelimitedReasoningParser, MiniMaxM3ReasoningParser,
+    Qwen3ReasoningParser, ReasoningParser,
 };
 
 struct FakeTokenizer;
@@ -32,6 +33,8 @@ impl Tokenizer for FakeTokenizer {
             "<|END_THINKING|>" => Some(4),
             "◁think▷" => Some(5),
             "◁/think▷" => Some(6),
+            "<mm:think>" => Some(8),
+            "</mm:think>" => Some(9),
             _ => None,
         }
     }
@@ -159,3 +162,35 @@ fn deepseek_r1_stops_scanning_at_last_special_token() {
     assert_eq!(delta.reasoning.as_deref(), Some("reason"));
     assert_eq!(delta.content.as_deref(), Some("answer"));
 }
+
+#[test]
+fn minimax_m3_handles_explicit_think_delimiters() {
+    let tokenizer = Arc::new(FakeTokenizer);
+    let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
+
+    let delta = parser.push("<mm:think>reason</mm:think>answer").unwrap();
+    assert_eq!(delta.reasoning.as_deref(), Some("reason"));
+    assert_eq!(delta.content.as_deref(), Some("answer"));
+}
+
+#[test]
+fn minimax_m3_uses_prompt_prefilled_start_marker() {
+    let tokenizer = Arc::new(FakeTokenizer);
+    let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
+    parser.initialize(&[8]).unwrap();
+
+    let delta = parser.push("reason</mm:think>answer").unwrap();
+    assert_eq!(delta.reasoning.as_deref(), Some("reason"));
+    assert_eq!(delta.content.as_deref(), Some("answer"));
+}
+
+#[test]
+fn minimax_m3_uses_prompt_prefilled_end_marker() {
+    let tokenizer = Arc::new(FakeTokenizer);
+    let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
+    parser.initialize(&[9]).unwrap();
+
+    let delta = parser.push("answer").unwrap();
+    assert_eq!(delta.reasoning, None);
+    assert_eq!(delta.content.as_deref(), Some("answer"));
+}

rust/src/tool-parser/python/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "vllm-tool-parser-py"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+
+[lib]
+name = "_rust_tool_parser"
+crate-type = ["cdylib", "rlib"]
+
+[features]
+default = []
+extension-module = ["pyo3/extension-module"]
+
+[dependencies]
+pyo3.workspace = true
+pythonize = { workspace = true, features = ["serde_json"] }
+serde_json.workspace = true
+thiserror-ext.workspace = true
+vllm-tool-parser.workspace = true
+
+[lints]
+workspace = true

rust/src/tool-parser/python/src/lib.rs

@@ -0,0 +1,367 @@
+//! Thin PyO3 bindings for `vllm_tool_parser`.
+//!
+//! This crate exposes the Rust tool parser trait and data shapes to Python
+//! while keeping parser state, grammar, and schema-aware argument conversion in
+//! Rust. Python callers should use this module as a typed bridge and keep any
+//! vLLM protocol adaptation outside the binding.
+
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::types::{PyAny, PyModule};
+use pythonize::{depythonize, pythonize};
+use serde_json::Value;
+use thiserror_ext::AsReport as _;
+use vllm_tool_parser::{Tool, ToolCallDelta, ToolParser, ToolParserOutput};
+
+macro_rules! tool_parser_factory {
+    ($($parser:ident),+ $(,)?) => {
+        fn create_tool_parser(
+            name: &str,
+            tools: &[Tool],
+        ) -> PyResult<Box<dyn ToolParser>> {
+            match name {
+                $(
+                    stringify!($parser) => {
+                        <vllm_tool_parser::$parser as ToolParser>::create(tools)
+                    }
+                )+
+                _ => {
+                    return Err(PyValueError::new_err(format!(
+                        "unsupported tool parser `{name}`"
+                    )));
+                }
+            }
+            .map_err(|error| PyValueError::new_err(error.to_report_string()))
+        }
+    };
+}
+
+// Export a tool parser to Python by registering it here.
+tool_parser_factory! {
+    DeepSeekV4ToolParser,
+    MinimaxM3ToolParser,
+}
+
+#[pyclass(name = "Tool", module = "vllm._rust_tool_parser", skip_from_py_object)]
+#[derive(Clone)]
+struct PyTool(Tool);
+
+#[pymethods]
+impl PyTool {
+    #[new]
+    #[pyo3(signature = (name, description, parameters, strict=None))]
+    fn new(
+        name: String,
+        description: Option<String>,
+        parameters: &Bound<'_, PyAny>,
+        strict: Option<bool>,
+    ) -> PyResult<Self> {
+        let parameters = depythonize::<Value>(parameters).map_err(|error| {
+            PyValueError::new_err(format!(
+                "failed to convert tool parameters from Python to JSON: {error}"
+            ))
+        })?;
+        Ok(Self(Tool {
+            name,
+            description,
+            parameters,
+            strict,
+        }))
+    }
+
+    #[getter]
+    fn name(&self) -> &str {
+        &self.0.name
+    }
+
+    #[getter]
+    fn description(&self) -> Option<&str> {
+        self.0.description.as_deref()
+    }
+
+    #[getter]
+    fn parameters(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
+        pythonize(py, &self.0.parameters).map(Bound::unbind).map_err(|error| {
+            PyValueError::new_err(format!(
+                "failed to convert tool parameters from JSON to Python: {error}"
+            ))
+        })
+    }
+
+    #[getter]
+    fn strict(&self) -> Option<bool> {
+        self.0.strict
+    }
+}
+
+#[pyclass(
+    name = "ToolCallDelta",
+    module = "vllm._rust_tool_parser",
+    skip_from_py_object
+)]
+#[derive(Clone)]
+struct PyToolCallDelta(ToolCallDelta);
+
+#[pymethods]
+impl PyToolCallDelta {
+    #[new]
+    #[pyo3(signature = (tool_index, name, arguments))]
+    fn new(tool_index: usize, name: Option<String>, arguments: String) -> Self {
+        Self(ToolCallDelta {
+            tool_index,
+            name,
+            arguments,
+        })
+    }
+
+    #[getter]
+    fn tool_index(&self) -> usize {
+        self.0.tool_index
+    }
+
+    #[getter]
+    fn name(&self) -> Option<&str> {
+        self.0.name.as_deref()
+    }
+
+    #[getter]
+    fn arguments(&self) -> &str {
+        &self.0.arguments
+    }
+}
+
+#[pyclass(
+    name = "ToolParserOutput",
+    module = "vllm._rust_tool_parser",
+    skip_from_py_object
+)]
+#[derive(Clone)]
+struct PyToolParserOutput(ToolParserOutput);
+
+#[pymethods]
+impl PyToolParserOutput {
+    #[new]
+    #[pyo3(signature = (normal_text="", calls=None))]
+    fn new(py: Python<'_>, normal_text: &str, calls: Option<Vec<Py<PyToolCallDelta>>>) -> Self {
+        let calls =
+            calls.unwrap_or_default().iter().map(|call| call.borrow(py).0.clone()).collect();
+        Self(ToolParserOutput {
+            normal_text: normal_text.to_owned(),
+            calls,
+        })
+    }
+
+    #[getter]
+    fn normal_text(&self) -> &str {
+        &self.0.normal_text
+    }
+
+    #[getter]
+    fn calls(&self) -> Vec<PyToolCallDelta> {
+        self.0.calls.iter().cloned().map(PyToolCallDelta).collect()
+    }
+
+    fn append(&mut self, other: PyRef<'_, PyToolParserOutput>) {
+        self.0.append(other.0.clone());
+    }
+
+    fn coalesce_calls(&self) -> Self {
+        Self(self.0.clone().coalesce_calls())
+    }
+}
+
+#[pyclass(name = "ToolParser", module = "vllm._rust_tool_parser", unsendable)]
+struct PyToolParser(Box<dyn ToolParser>);
+
+impl PyToolParser {
+    fn parse_into_output(&mut self, chunk: &str, output: &mut PyToolParserOutput) -> PyResult<()> {
+        self.0
+            .parse_into(chunk, &mut output.0)
+            .map_err(|error| PyValueError::new_err(error.to_report_string()))
+    }
+}
+
+#[pymethods]
+impl PyToolParser {
+    #[new]
+    fn new(py: Python<'_>, parser_name: &str, tools: Vec<Py<PyTool>>) -> PyResult<Self> {
+        let tools = tools.iter().map(|tool| tool.borrow(py).0.clone()).collect::<Vec<_>>();
+        create_tool_parser(parser_name, &tools).map(Self)
+    }
+
+    fn parse_into(
+        &mut self,
+        chunk: &str,
+        mut output: PyRefMut<'_, PyToolParserOutput>,
+    ) -> PyResult<()> {
+        self.parse_into_output(chunk, &mut output)
+    }
+
+    fn finish(&mut self) -> PyResult<PyToolParserOutput> {
+        self.0
+            .finish()
+            .map(PyToolParserOutput)
+            .map_err(|error| PyValueError::new_err(error.to_report_string()))
+    }
+
+    fn reset(&mut self) -> String {
+        self.0.reset()
+    }
+
+    fn preserve_special_tokens(&self) -> bool {
+        self.0.preserve_special_tokens()
+    }
+}
+
+#[pymodule]
+fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<PyTool>()?;
+    m.add_class::<PyToolCallDelta>()?;
+    m.add_class::<PyToolParserOutput>()?;
+    m.add_class::<PyToolParser>()?;
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    fn with_python<R>(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R {
+        Python::initialize();
+        Python::attach(f)
+    }
+
+    fn tool_schema() -> Value {
+        json!({
+            "type": "object",
+            "properties": {
+                "user_id": {"type": "integer"},
+                "shipping": {
+                    "type": "object",
+                    "properties": {
+                        "city": {"type": "string"},
+                        "zip": {"type": "integer"}
+                    }
+                }
+            }
+        })
+    }
+
+    fn build_call() -> String {
+        r#"<｜DSML｜tool_calls>
+<｜DSML｜invoke name="create_order">
+<｜DSML｜parameter name="user_id" string="false">42</｜DSML｜parameter>
+<｜DSML｜parameter name="shipping" string="false">{"city":"Singapore","zip":18956}</｜DSML｜parameter>
+</｜DSML｜invoke>
+</｜DSML｜tool_calls>"#
+            .to_owned()
+    }
+
+    fn make_py_tool(py: Python<'_>) -> PyResult<Py<PyTool>> {
+        let parameters = pythonize(py, &tool_schema()).map_err(|error| {
+            PyValueError::new_err(format!(
+                "failed to convert test schema from JSON to Python: {error}"
+            ))
+        })?;
+        Py::new(
+            py,
+            PyTool::new(
+                "create_order".to_owned(),
+                Some("Create an order".to_owned()),
+                &parameters,
+                None,
+            )?,
+        )
+    }
+
+    #[test]
+    fn tool_round_trips_typed_fields() {
+        with_python(|py| {
+            let tool = make_py_tool(py)?;
+            let borrowed = tool.borrow(py);
+            assert_eq!(borrowed.name(), "create_order");
+            assert_eq!(borrowed.description(), Some("Create an order"));
+            assert_eq!(borrowed.strict(), None);
+
+            let parameters = borrowed.parameters(py)?;
+            let parameters = depythonize::<Value>(parameters.bind(py))?;
+            assert_eq!(parameters, tool_schema());
+            PyResult::Ok(())
+        })
+        .unwrap();
+    }
+
+    #[test]
+    fn output_append_and_coalesce_calls() {
+        with_python(|py| {
+            let first = Py::new(
+                py,
+                PyToolCallDelta::new(0, Some("create_order".to_owned()), "{\"a\"".to_owned()),
+            )?;
+            let second = Py::new(py, PyToolCallDelta::new(0, None, ":1}".to_owned()))?;
+            let mut output = PyToolParserOutput::new(py, "text", Some(vec![first]));
+            let other = Py::new(py, PyToolParserOutput::new(py, "", Some(vec![second])))?;
+            output.append(other.borrow(py));
+
+            let coalesced = output.coalesce_calls();
+            assert_eq!(coalesced.normal_text(), "text");
+            let calls = coalesced.calls();
+            assert_eq!(calls.len(), 1);
+            assert_eq!(calls[0].tool_index(), 0);
+            assert_eq!(calls[0].name(), Some("create_order"));
+            assert_eq!(calls[0].arguments(), "{\"a\":1}");
+            PyResult::Ok(())
+        })
+        .unwrap();
+    }
+
+    #[test]
+    fn parser_parse_finish_and_preserve_special_tokens() {
+        with_python(|py| {
+            let tool = make_py_tool(py)?;
+            let mut parser = PyToolParser::new(py, "DeepSeekV4ToolParser", vec![tool])?;
+            assert!(parser.preserve_special_tokens());
+
+            let mut output = PyToolParserOutput::new(py, "", None);
+            parser.parse_into_output(&build_call(), &mut output)?;
+            let finish = Py::new(py, parser.finish()?)?;
+            output.append(finish.borrow(py));
+            let output = output.coalesce_calls();
+
+            assert_eq!(output.normal_text(), "");
+            let calls = output.calls();
+            assert_eq!(calls.len(), 1);
+            assert_eq!(calls[0].name(), Some("create_order"));
+            assert_eq!(
+                serde_json::from_str::<Value>(calls[0].arguments()).unwrap(),
+                json!({
+                    "user_id": 42,
+                    "shipping": {
+                        "city": "Singapore",
+                        "zip": 18956
+                    }
+                })
+            );
+
+            assert_eq!(parser.reset(), "");
+            PyResult::Ok(())
+        })
+        .unwrap();
+    }
+
+    #[test]
+    fn parser_errors_for_unknown_name() {
+        with_python(|py| {
+            let tool = make_py_tool(py)?;
+            let error = match PyToolParser::new(py, "missing", vec![tool]) {
+                Ok(_) => panic!("missing parser name unexpectedly succeeded"),
+                Err(error) => error,
+            };
+            let message = format!("{error}");
+            assert!(message.contains("unsupported tool parser `missing`"));
+            PyResult::Ok(())
+        })
+        .unwrap();
+    }
+}

rust/src/tool-parser/src/deepseek_dsml/mod.rs

@@ -104,7 +104,7 @@ impl DeepSeekDsmlToolParser {
                         self.tool_parameters.convert_param_with_schema(
                             &name,
                             &param.name,
-                            &param.value,
+                            param.value,
                         )
                     };
                     arguments.insert(param.name, value);

rust/src/tool-parser/src/lib.rs

@@ -10,6 +10,7 @@ mod hy_v3;
 mod json;
 mod kimi_k2;
 mod minimax_m2;
+mod minimax_m3;
 mod parameters;
 mod qwen_coder;
 #[cfg(any(test, feature = "test-util"))]
@@ -27,6 +28,7 @@ pub use hy_v3::HyV3ToolParser;
 pub use json::{HermesToolParser, Llama3JsonToolParser, MistralToolParser, Qwen3XmlToolParser};
 pub use kimi_k2::KimiK2ToolParser;
 pub use minimax_m2::MinimaxM2ToolParser;
+pub use minimax_m3::MinimaxM3ToolParser;
 pub use qwen_coder::Qwen3CoderToolParser;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;

rust/src/tool-parser/src/minimax_m3.rs

@@ -0,0 +1,812 @@
+use winnow::ascii::{multispace0 as ws0, multispace1 as ws1};
+use winnow::combinator::{alt, delimited, eof, repeat, seq, terminated};
+use winnow::error::{ContextError, ErrMode};
+use winnow::prelude::*;
+use winnow::stream::Partial;
+use winnow::token::{literal, rest, take_until};
+
+use super::parameters::{ParamElement, ParamInput, ToolSchemas};
+use super::utils::{parse_buffered_event, safe_text_len};
+use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
+use crate::Tool;
+
+const NAMESPACE: &str = "]<]minimax[>[";
+const TOOL_CALL_START: &str = "]<]minimax[>[<tool_call>";
+const TOOL_CALL_END: &str = "]<]minimax[>[</tool_call>";
+const INVOKE_START: &str = "]<]minimax[>[<invoke";
+const INVOKE_END: &str = "]<]minimax[>[</invoke>";
+const ELEMENT_START: &str = "]<]minimax[>[<";
+const ELEMENT_END_START: &str = "]<]minimax[>[</";
+const MIXED_TEXT_FIELD: &str = "$text";
+
+type MinimaxM3Input<'i> = Partial<&'i str>;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum MinimaxM3Mode {
+    Text,
+    ToolBlock,
+    Done,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+enum MinimaxM3Event {
+    Text {
+        len: usize,
+    },
+    ToolBlockStart,
+    Invoke {
+        name: String,
+        params: Vec<(String, ParamInput)>,
+    },
+    ToolBlockEnd,
+    IgnoredRest,
+}
+
+/// Tool parser for MiniMax M3 namespace-delimited XML-style tool calls.
+///
+/// Example tool call content with recursive parameters:
+///
+/// ```text
+/// ]<]minimax[>[<tool_call>
+/// ]<]minimax[>[<invoke name="create_order">
+/// ]<]minimax[>[<user_id>42]<]minimax[>[</user_id>
+/// ]<]minimax[>[<shipping>
+/// ]<]minimax[>[<city>Singapore]<]minimax[>[</city>
+/// ]<]minimax[>[<zip>018956]<]minimax[>[</zip>
+/// ]<]minimax[>[</shipping>
+/// ]<]minimax[>[<items>
+/// ]<]minimax[>[<item>
+/// ]<]minimax[>[<sku>book-001]<]minimax[>[</sku>
+/// ]<]minimax[>[<qty>2]<]minimax[>[</qty>
+/// ]<]minimax[>[</item>
+/// ]<]minimax[>[</items>
+/// ]<]minimax[>[</invoke>
+/// ]<]minimax[>[</tool_call>
+/// ```
+///
+/// With a schema where `shipping` is an object and `items` is an array of
+/// objects, recursive parameter conversion produces:
+///
+/// ```json
+/// {
+///   "user_id": 42,
+///   "shipping": {
+///     "city": "Singapore",
+///     "zip": 18956
+///   },
+///   "items": [
+///     {
+///       "sku": "book-001",
+///       "qty": 2
+///     }
+///   ]
+/// }
+/// ```
+///
+/// MiniMax M3 emits the namespace marker `]<]minimax[>[` before each structural
+/// tag. Arguments are emitted only after a full `<invoke>` block is parsed.
+pub struct MinimaxM3ToolParser {
+    buffer: String,
+    mode: MinimaxM3Mode,
+    emitted_tool_count: usize,
+    tool_parameters: ToolSchemas,
+}
+
+impl MinimaxM3ToolParser {
+    /// Create a MiniMax M3 tool parser.
+    pub fn new(tools: &[Tool]) -> Self {
+        Self {
+            buffer: String::new(),
+            mode: MinimaxM3Mode::Text,
+            emitted_tool_count: 0,
+            tool_parameters: ToolSchemas::from_tools(tools),
+        }
+    }
+
+    /// Apply one parsed MiniMax M3 event to parser state and output.
+    fn apply_event(&mut self, event: MinimaxM3Event, output: &mut ToolParserOutput) -> Result<()> {
+        match event {
+            MinimaxM3Event::Text { len: consumed_len } => {
+                output.normal_text.push_str(&self.buffer[..consumed_len]);
+            }
+            MinimaxM3Event::ToolBlockStart => self.mode = MinimaxM3Mode::ToolBlock,
+            MinimaxM3Event::Invoke { name, params } => {
+                let arguments = self.tool_parameters.convert_params_with_schema(&name, params);
+                let arguments = serde_json::to_string(&arguments)
+                    .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
+
+                output.calls.push(ToolCallDelta {
+                    tool_index: self.emitted_tool_count,
+                    name: Some(name),
+                    arguments,
+                });
+                self.emitted_tool_count += 1;
+            }
+            MinimaxM3Event::ToolBlockEnd => self.mode = MinimaxM3Mode::Done,
+            MinimaxM3Event::IgnoredRest => {}
+        }
+        Ok(())
+    }
+}
+
+impl ToolParser for MinimaxM3ToolParser {
+    fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
+    where
+        Self: Sized + 'static,
+    {
+        Ok(Box::new(Self::new(tools)))
+    }
+
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+        self.buffer.push_str(chunk);
+
+        while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
+            parse_next_minimax_m3_event(input, self.mode)
+        })? {
+            self.apply_event(event, output)?;
+            self.buffer.drain(..consumed_len);
+        }
+
+        Ok(())
+    }
+
+    fn finish(&mut self) -> Result<ToolParserOutput> {
+        let mut output = ToolParserOutput::default();
+        match self.mode {
+            MinimaxM3Mode::Text => {
+                output.normal_text.push_str(&self.buffer);
+            }
+            MinimaxM3Mode::ToolBlock => {
+                return Err(parsing_failed!("incomplete MiniMax M3 tool call"));
+            }
+            MinimaxM3Mode::Done => {}
+        }
+        let _ = self.reset();
+        Ok(output)
+    }
+
+    fn reset(&mut self) -> String {
+        self.mode = MinimaxM3Mode::Text;
+        self.emitted_tool_count = 0;
+        std::mem::take(&mut self.buffer)
+    }
+}
+
+/// Parse a MiniMax M3 event for the current parser mode.
+fn parse_next_minimax_m3_event(
+    input: &mut MinimaxM3Input<'_>,
+    mode: MinimaxM3Mode,
+) -> ModalResult<MinimaxM3Event> {
+    match mode {
+        MinimaxM3Mode::Text => parse_text_event(input),
+        MinimaxM3Mode::ToolBlock => parse_tool_block_event(input),
+        MinimaxM3Mode::Done => ignored_rest_event(input),
+    }
+}
+
+/// Parse a text-mode MiniMax M3 event.
+fn parse_text_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
+    alt((tool_block_start_event, safe_text_event)).parse_next(input)
+}
+
+/// Parse a MiniMax M3 tool-block start marker.
+fn tool_block_start_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
+    literal(TOOL_CALL_START).value(MinimaxM3Event::ToolBlockStart).parse_next(input)
+}
+
+/// Parse a safe text run before the next MiniMax M3 marker.
+fn safe_text_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
+    safe_text_len(input, TOOL_CALL_START).map(|len| MinimaxM3Event::Text { len })
+}
+
+/// Parse one event inside a MiniMax M3 tool block.
+fn parse_tool_block_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
+    alt((tool_block_end_event, invoke_event)).parse_next(input)
+}
+
+/// Parse a MiniMax M3 tool-block end marker.
+fn tool_block_end_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
+    (ws0, literal(TOOL_CALL_END))
+        .value(MinimaxM3Event::ToolBlockEnd)
+        .parse_next(input)
+}
+
+/// Parse a complete MiniMax M3 invoke block.
+fn invoke_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
+    let (name, body) = seq!(
+        _: ws0,
+        _: literal(INVOKE_START),
+        _: (ws1, literal("name=")),
+        partial_attr_value,
+        _: literal(">"),
+        take_until(0.., INVOKE_END),
+        _: literal(INVOKE_END),
+    )
+    .parse_next(input)?;
+    let params = parse_invoke_params(body)?;
+
+    Ok(MinimaxM3Event::Invoke {
+        name: name.trim().to_string(),
+        params,
+    })
+}
+
+/// Parse all parameter elements inside a complete MiniMax M3 invoke body.
+fn parse_invoke_params(invoke_body: &str) -> ModalResult<Vec<(String, ParamInput)>> {
+    let mut input = invoke_body;
+    let elements: Vec<ParamElement> =
+        delimited(ws0, repeat(0.., terminated(parameter_element, ws0)), eof)
+            .parse_next(&mut input)?;
+
+    Ok(elements.into_iter().map(|element| (element.name, element.value)).collect())
+}
+
+/// Parse a MiniMax M3 parameter element.
+fn parameter_element(input: &mut &str) -> ModalResult<ParamElement> {
+    let name = open_element_tag(input)?.to_string();
+    let value = element_body(input, &name)?;
+    close_element_tag(input, &name)?;
+    Ok(ParamElement { name, value })
+}
+
+/// Parse a MiniMax M3 opening element tag.
+fn open_element_tag<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
+    let name = seq!(
+        _: literal(ELEMENT_START),
+        take_until(1.., ">"),
+        _: literal(">"),
+    )
+    .parse_next(input)?;
+
+    let name = name.0;
+    if name.starts_with('/') || name.trim().is_empty() {
+        return malformed();
+    }
+
+    Ok(name)
+}
+
+/// Parse a MiniMax M3 closing element tag.
+fn close_element_tag(input: &mut &str, name: &str) -> ModalResult<()> {
+    literal(ELEMENT_END_START).void().parse_next(input)?;
+    literal(name).void().parse_next(input)?;
+    literal(">").void().parse_next(input)
+}
+
+/// Parse the body of one MiniMax M3 element.
+fn element_body(input: &mut &str, closing_name: &str) -> ModalResult<ParamInput> {
+    let close_tag = format!("{ELEMENT_END_START}{closing_name}>");
+    let mut text = String::new();
+    let mut elements = Vec::new();
+
+    loop {
+        text.push_str(text_until_namespace(input)?);
+
+        if input.starts_with(&close_tag) {
+            // Close tag reached, end of element body.
+            break;
+        }
+        if input.starts_with(ELEMENT_START) {
+            // Child element start reached, parse child element recursively.
+            elements.push(parameter_element(input)?);
+            continue;
+        }
+        if input.starts_with(NAMESPACE) {
+            // Unexpected namespace marker.
+            return malformed();
+        }
+    }
+
+    if elements.is_empty() {
+        Ok(ParamInput::Text(text))
+    } else {
+        if !text.trim().is_empty() {
+            push_mixed_text_element(&mut elements, text);
+        }
+        Ok(ParamInput::Elements(elements))
+    }
+}
+
+/// Parse text until the next MiniMax M3 namespace marker.
+fn text_until_namespace<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
+    take_until(0.., NAMESPACE).parse_next(input)
+}
+
+/// Preserve mixed text content under a reserved object field.
+///
+/// By default, the field name is `$text`, but if that collides with an existing
+/// child element name, prepend `$` until there is no collision.
+fn push_mixed_text_element(elements: &mut Vec<ParamElement>, text: String) {
+    let mut name = MIXED_TEXT_FIELD.to_string();
+    while elements.iter().any(|element| element.name == name) {
+        name.insert(0, '$');
+    }
+    elements.push(ParamElement {
+        name,
+        value: ParamInput::Text(text),
+    });
+}
+
+/// Parse a quoted or unquoted XML attribute value from partial streaming input.
+fn partial_attr_value<'i>(input: &mut MinimaxM3Input<'i>) -> ModalResult<&'i str> {
+    alt((
+        delimited(literal("\""), take_until(1.., "\""), literal("\"")),
+        delimited(literal("'"), take_until(1.., "'"), literal("'")),
+        take_until(1.., ">"),
+    ))
+    .parse_next(input)
+}
+
+/// Parse ignored rest after the MiniMax M3 tool block ends.
+fn ignored_rest_event(input: &mut MinimaxM3Input<'_>) -> ModalResult<MinimaxM3Event> {
+    rest.value(MinimaxM3Event::IgnoredRest).parse_next(input)
+}
+
+fn malformed<T>() -> ModalResult<T> {
+    Err(ErrMode::Cut(ContextError::new()))
+}
+
+#[cfg(test)]
+mod tests {
+    use expect_test::expect;
+    use serde_json::{Value, json};
+    use thiserror_ext::AsReport;
+
+    use super::{
+        ELEMENT_END_START, ELEMENT_START, INVOKE_END, INVOKE_START, MinimaxM3ToolParser,
+        TOOL_CALL_END, TOOL_CALL_START, ToolParser,
+    };
+    use crate::test_utils::{collect_stream, split_by_chars, test_tools};
+    use crate::{Tool, ToolParserTestExt as _};
+
+    fn element(name: &str, body: &str) -> String {
+        format!("{ELEMENT_START}{name}>{body}{ELEMENT_END_START}{name}>")
+    }
+
+    fn invoke(function_name: &str, body: &str) -> String {
+        format!("{INVOKE_START} name=\"{function_name}\">{body}{INVOKE_END}")
+    }
+
+    fn build_tool_block(invokes: &[(&str, String)]) -> String {
+        let invokes = invokes
+            .iter()
+            .map(|(function_name, body)| invoke(function_name, body))
+            .collect::<Vec<_>>()
+            .join("\n");
+        format!("{TOOL_CALL_START}\n{invokes}\n{TOOL_CALL_END}")
+    }
+
+    fn m3_test_tools() -> Vec<Tool> {
+        let mut tools = test_tools();
+        tools.push(Tool {
+            name: "create_order".to_string(),
+            description: None,
+            parameters: json!({
+                "type": "object",
+                "properties": {
+                    "user_id": { "type": "integer" },
+                    "urgent": { "type": "boolean" },
+                    "note": { "type": "string" },
+                    "shipping": {
+                        "type": "object",
+                        "properties": {
+                            "city": { "type": "string" },
+                            "zip": { "type": "integer" }
+                        }
+                    },
+                    "items": {
+                        "type": "array",
+                        "items": {
+                            "type": "object",
+                            "properties": {
+                                "sku": { "type": "string" },
+                                "qty": { "type": "integer" }
+                            }
+                        }
+                    },
+                    "metadata": {
+                        "type": "object",
+                        "additionalProperties": { "type": "integer" }
+                    },
+                    "duplicate_demo": {
+                        "type": "object",
+                        "properties": {
+                            "tag": { "type": "string" }
+                        }
+                    },
+                    "schema_mismatch_array": {
+                        "type": "array",
+                        "items": { "type": "integer" }
+                    }
+                }
+            }),
+            strict: None,
+        });
+        tools
+    }
+
+    fn order_arguments() -> String {
+        let shipping = element(
+            "shipping",
+            &format!(
+                "{}{}",
+                element("city", "Singapore"),
+                element("zip", "018956")
+            ),
+        );
+        let first_item = element(
+            "item",
+            &format!("{}{}", element("sku", "book-001"), element("qty", "2")),
+        );
+        let second_item = element(
+            "item",
+            &format!("{}{}", element("sku", "pen-007"), element("qty", "5")),
+        );
+        let items = element("items", &format!("{first_item}{second_item}"));
+        let metadata = element(
+            "metadata",
+            &format!("{}{}", element("score", "42"), element("rank", "7")),
+        );
+        let duplicate_demo = element(
+            "duplicate_demo",
+            &format!("{}{}", element("tag", "a"), element("tag", "b")),
+        );
+        let schema_mismatch_array = element(
+            "schema_mismatch_array",
+            &format!("{}{}", element("x", "1"), element("x", "2")),
+        );
+
+        [
+            element("user_id", "42"),
+            element("urgent", "true"),
+            element("note", "Please leave at front desk."),
+            shipping,
+            items,
+            metadata,
+            duplicate_demo,
+            schema_mismatch_array,
+            element(
+                "unknown_struct",
+                &format!("{}{}", element("a", "1"), element("a", "2")),
+            ),
+        ]
+        .join("")
+    }
+
+    #[test]
+    fn minimax_m3_parse_complete_without_tool_call_keeps_text() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = parser.parse_complete("Hello, world!").unwrap();
+
+        assert_eq!(output.normal_text, "Hello, world!");
+        assert!(output.calls.is_empty());
+    }
+
+    #[test]
+    fn minimax_m3_parse_complete_extracts_single_tool_call() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = parser
+            .parse_complete(&build_tool_block(&[(
+                "get_weather",
+                format!("{}{}", element("city", "Seattle"), element("days", "5")),
+            )]))
+            .unwrap();
+
+        assert!(output.normal_text.is_empty());
+        assert_eq!(output.calls.len(), 1);
+        assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
+        assert_eq!(
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
+            json!({ "city": "Seattle", "days": 5 })
+        );
+    }
+
+    #[test]
+    fn minimax_m3_parse_complete_preserves_prefix_and_ignores_trailing_text() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = format!(
+            "Let me check. {} This trailing text is ignored.",
+            build_tool_block(&[("get_weather", element("city", "Seattle"))])
+        );
+        let output = parser.parse_complete(&output).unwrap();
+
+        assert_eq!(output.normal_text, "Let me check. ");
+        assert_eq!(output.calls.len(), 1);
+    }
+
+    #[test]
+    fn minimax_m3_parse_complete_extracts_multiple_invokes() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = parser
+            .parse_complete(&build_tool_block(&[
+                ("get_weather", element("city", "Seattle")),
+                ("get_weather", element("city", "NYC")),
+            ]))
+            .unwrap();
+
+        assert_eq!(output.calls.len(), 2);
+        assert_eq!(output.calls[0].tool_index, 0);
+        assert_eq!(output.calls[1].tool_index, 1);
+        assert_eq!(
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
+            json!({ "city": "Seattle" })
+        );
+        assert_eq!(
+            serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
+            json!({ "city": "NYC" })
+        );
+    }
+
+    #[test]
+    fn minimax_m3_parse_complete_converts_schema_types() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = parser
+            .parse_complete(&build_tool_block(&[(
+                "convert",
+                [
+                    element("whole", "5.0"),
+                    element("flag", "true"),
+                    element("payload", r#"{"nested":true}"#),
+                    element("items", "[1,2]"),
+                    element("empty", "42"),
+                ]
+                .join(""),
+            )]))
+            .unwrap();
+
+        assert_eq!(
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
+            json!({
+                "whole": 5.0,
+                "flag": true,
+                "payload": { "nested": true },
+                "items": [1, 2],
+                "empty": "42",
+            })
+        );
+    }
+
+    #[test]
+    fn minimax_m3_parse_complete_converts_nested_arguments() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = parser
+            .parse_complete(&build_tool_block(&[("create_order", order_arguments())]))
+            .unwrap();
+
+        assert_eq!(
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
+            json!({
+                "user_id": 42,
+                "urgent": true,
+                "note": "Please leave at front desk.",
+                "shipping": {
+                    "city": "Singapore",
+                    "zip": 18956
+                },
+                "items": [
+                    {
+                        "sku": "book-001",
+                        "qty": 2
+                    },
+                    {
+                        "sku": "pen-007",
+                        "qty": 5
+                    }
+                ],
+                "metadata": {
+                    "score": 42,
+                    "rank": 7
+                },
+                "duplicate_demo": {
+                    "tag": ["a", "b"]
+                },
+                "schema_mismatch_array": [1, 2],
+                "unknown_struct": {
+                    "a": ["1", "2"]
+                }
+            })
+        );
+    }
+
+    #[test]
+    fn minimax_m3_parse_complete_handles_multiline_leaf_parameters() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = parser
+            .parse_complete(&build_tool_block(&[(
+                "calculate_area",
+                [
+                    element("shape", "\nrectangle\n"),
+                    element("dimensions", r#"{"width":10,"height":20}"#),
+                    element("precision", "2"),
+                ]
+                .join(""),
+            )]))
+            .unwrap();
+
+        assert_eq!(
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
+            json!({
+                "shape": "\nrectangle\n",
+                "dimensions": { "width": 10, "height": 20 },
+                "precision": 2,
+            })
+        );
+    }
+
+    #[test]
+    fn minimax_m3_streaming_extracts_single_tool_call() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = collect_stream(
+            &mut parser,
+            &[
+                TOOL_CALL_START,
+                &invoke("get_weather", &element("city", "Seattle")),
+                TOOL_CALL_END,
+            ],
+        );
+
+        assert!(output.normal_text.is_empty());
+        assert_eq!(output.calls.len(), 1);
+        assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
+        assert_eq!(
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
+            json!({ "city": "Seattle" })
+        );
+    }
+
+    #[test]
+    fn minimax_m3_streaming_preserves_prefix_text() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = collect_stream(
+            &mut parser,
+            &[
+                "Let me check. ",
+                TOOL_CALL_START,
+                &invoke("get_weather", &element("city", "Seattle")),
+                TOOL_CALL_END,
+            ],
+        );
+
+        assert_eq!(output.normal_text, "Let me check. ");
+        assert_eq!(output.calls.len(), 1);
+    }
+
+    #[test]
+    fn minimax_m3_streaming_without_tool_call_emits_text_incrementally() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = collect_stream(&mut parser, &["Hello, ", "world!"]);
+
+        assert_eq!(output.normal_text, "Hello, world!");
+        assert!(output.calls.is_empty());
+    }
+
+    #[test]
+    fn minimax_m3_streaming_handles_marker_split_across_chunks() {
+        let text = build_tool_block(&[("get_weather", element("city", "Seattle"))]);
+        let chunks = split_by_chars(&text, 3);
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = collect_stream(&mut parser, &chunks);
+
+        assert_eq!(output.calls.len(), 1);
+        assert!(output.normal_text.is_empty());
+    }
+
+    #[test]
+    fn minimax_m3_streaming_extracts_multiple_invokes_in_order() {
+        let text = build_tool_block(&[
+            ("get_weather", element("city", "Seattle")),
+            ("get_weather", element("city", "NYC")),
+        ]);
+        let chunks = split_by_chars(&text, 7);
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = collect_stream(&mut parser, &chunks);
+
+        assert_eq!(output.calls.len(), 2);
+        assert_eq!(output.calls[0].tool_index, 0);
+        assert_eq!(output.calls[1].tool_index, 1);
+    }
+
+    #[test]
+    fn minimax_m3_streaming_does_not_emit_incomplete_tool_call() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = parser
+            .parse_chunk(&format!(
+                "{TOOL_CALL_START}{INVOKE_START} name=\"get_weather\">"
+            ))
+            .unwrap();
+
+        assert!(output.normal_text.is_empty());
+        assert!(output.calls.is_empty());
+    }
+
+    #[test]
+    fn minimax_m3_streaming_ignores_text_after_tool_block() {
+        let text = format!(
+            "{} ignored",
+            build_tool_block(&[("get_weather", element("city", "Seattle"))])
+        );
+        let chunks = split_by_chars(&text, 5);
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let output = collect_stream(&mut parser, &chunks);
+
+        assert!(output.normal_text.is_empty());
+        assert_eq!(output.calls.len(), 1);
+    }
+
+    #[test]
+    fn minimax_m3_finish_fails_incomplete_tool_call() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        parser
+            .parse_chunk(&format!(
+                "{TOOL_CALL_START}{INVOKE_START} name=\"get_weather\">"
+            ))
+            .unwrap();
+
+        assert!(parser.finish().is_err());
+    }
+
+    #[test]
+    fn minimax_m3_finish_fails_after_bare_tool_block_start() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        parser.parse_chunk(TOOL_CALL_START).unwrap();
+
+        assert!(parser.finish().is_err());
+    }
+
+    #[test]
+    fn minimax_m3_malformed_tool_call_fails_fast() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let error = parser
+            .parse_chunk(&format!(
+                "{TOOL_CALL_START}{ELEMENT_START}bad>{TOOL_CALL_END}"
+            ))
+            .unwrap_err();
+
+        expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
+    }
+
+    #[test]
+    fn minimax_m3_mixed_content_is_preserved_as_text_field() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let body = element(
+            "payload",
+            &format!("text before {} text after", element("child", "value")),
+        );
+        let output = parser.parse_complete(&build_tool_block(&[("convert", body)])).unwrap();
+
+        assert_eq!(
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
+            json!({
+                "payload": {
+                    "child": "value",
+                    "$text": "text before  text after"
+                }
+            })
+        );
+    }
+
+    #[test]
+    fn minimax_m3_mixed_text_field_avoids_child_name_collision() {
+        let mut parser = MinimaxM3ToolParser::new(&m3_test_tools());
+        let body = element(
+            "payload",
+            &format!(
+                "text{}{}",
+                element("$text", "child text"),
+                element("child", "value")
+            ),
+        );
+        let output = parser.parse_complete(&build_tool_block(&[("convert", body)])).unwrap();
+
+        assert_eq!(
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
+            json!({
+                "payload": {
+                    "$text": "child text",
+                    "$$text": "text",
+                    "child": "value"
+                }
+            })
+        );
+    }
+}

rust/src/tool-parser/src/parameters.rs

@@ -1,6 +1,6 @@
 use std::collections::BTreeMap;
 
-use serde_json::{Number, Value};
+use serde_json::{Map, Number, Value};
 
 use crate::Tool;
 
@@ -21,6 +21,29 @@ pub(super) struct ToolSchema {
     params: BTreeMap<String, JsonParamType>,
 }
 
+/// Parameter input for schema-aware conversion.
+///
+/// It can be either a raw text string, or a structured input with named child elements.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(super) enum ParamInput {
+    Text(String),
+    #[allow(dead_code)]
+    Elements(Vec<ParamElement>),
+}
+
+impl From<String> for ParamInput {
+    fn from(value: String) -> Self {
+        Self::Text(value)
+    }
+}
+
+/// One named structured parameter child.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(super) struct ParamElement {
+    pub name: String,
+    pub value: ParamInput,
+}
+
 /// Normalized JSON parameter type used for raw string coercion.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(super) enum JsonParamType {
@@ -28,8 +51,13 @@ pub(super) enum JsonParamType {
     Integer,
     Number,
     Boolean,
-    Object,
-    Array,
+    Object {
+        properties: BTreeMap<String, JsonParamType>,
+        additional_properties: Option<Box<JsonParamType>>,
+    },
+    Array {
+        items: Option<Box<JsonParamType>>,
+    },
     Null,
     OneOf(Vec<JsonParamType>),
 }
@@ -45,33 +73,39 @@ impl ToolSchemas {
         Self { tools }
     }
 
-    /// Convert raw string parameter values for one named tool.
+    /// Convert parameter values for one named tool.
     ///
     /// Unknown tool names use an empty schema, so all parameters fall back to
-    /// strings.
-    pub(super) fn convert_params_with_schema(
+    /// strings or object-like JSON for structured inputs.
+    pub(super) fn convert_params_with_schema<P>(
         &self,
         function_name: &str,
-        params: Vec<(String, String)>,
-    ) -> serde_json::Map<String, Value> {
+        params: Vec<(String, P)>,
+    ) -> Map<String, Value>
+    where
+        P: Into<ParamInput>,
+    {
         let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty());
-        let mut converted = serde_json::Map::with_capacity(params.len());
+        let mut converted = Map::with_capacity(params.len());
         for (name, value) in params {
-            let value = tool_schema.convert(&name, &value);
+            let value = tool_schema.convert(&name, value.into());
             converted.insert(name, value);
         }
         converted
     }
 
-    /// Convert one raw string parameter value for one named tool.
-    pub(super) fn convert_param_with_schema(
+    /// Convert one parameter value for one named tool.
+    pub(super) fn convert_param_with_schema<P>(
         &self,
         function_name: &str,
         name: &str,
-        value: &str,
-    ) -> Value {
+        value: P,
+    ) -> Value
+    where
+        P: Into<ParamInput>,
+    {
         let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty());
-        tool_schema.convert(name, value)
+        tool_schema.convert(name, value.into())
     }
 }
 
@@ -101,21 +135,13 @@ impl ToolSchema {
         Self { params }
     }
 
-    /// Convert one raw parameter value using its normalized schema type.
+    /// Convert one parameter value using its normalized schema type.
     ///
     /// If the parameter name is unknown, or we don't have a schema for it, or
     /// the value fails to convert, this falls back to returning the raw
-    /// string as a JSON string value.
-    fn convert(&self, name: &str, value: &str) -> Value {
-        if value.eq_ignore_ascii_case("null") {
-            return Value::Null;
-        }
-
-        let Some(param_type) = self.params.get(name) else {
-            return Value::String(value.to_string());
-        };
-
-        convert_value(param_type, value).unwrap_or_else(|| Value::String(value.to_string()))
+    /// string as a JSON string value, or object-like JSON for structured input.
+    fn convert(&self, name: &str, input: ParamInput) -> Value {
+        convert_with_optional_schema(self.params.get(name), &input)
     }
 }
 
@@ -125,7 +151,7 @@ impl JsonParamType {
         let schema = schema.as_object()?;
 
         if let Some(type_value) = schema.get("type") {
-            return Self::from_type_value(type_value);
+            return Self::from_type_value(type_value, schema);
         }
 
         if let Some(composite) = schema.get("anyOf").or_else(|| schema.get("oneOf")) {
@@ -134,32 +160,34 @@ impl JsonParamType {
                 .map(|schemas| schemas.iter().filter_map(Self::from_schema).collect::<Vec<_>>())
                 .filter(|types| !types.is_empty())
                 .map(Self::one_of)
-                .unwrap_or(Self::Object);
+                .unwrap_or_else(|| Self::object_from_schema(Some(schema)));
             return Some(param_type);
         }
 
+        // Typically, these types are already handled by checking the "type" field, but
+        // we can also infer them from their characteristic fields if "type" is missing.
         if schema.contains_key("enum") {
             return Some(Self::String);
         }
         if schema.contains_key("items") {
-            return Some(Self::Array);
+            return Some(Self::array_from_schema(Some(schema)));
         }
-        if schema.contains_key("properties") {
-            return Some(Self::Object);
+        if schema.contains_key("properties") || schema.contains_key("additionalProperties") {
+            return Some(Self::object_from_schema(Some(schema)));
         }
 
         None
     }
 
     /// Normalize a JSON schema `type` value.
-    fn from_type_value(type_value: &Value) -> Option<Self> {
+    fn from_type_value(type_value: &Value, schema: &Map<String, Value>) -> Option<Self> {
         match type_value {
-            Value::String(kind) => Self::from_type_name(kind),
+            Value::String(kind) => Self::from_type_name(kind, Some(schema)),
             Value::Array(kinds) => {
                 let types = kinds
                     .iter()
                     .filter_map(Value::as_str)
-                    .filter_map(Self::from_type_name)
+                    .filter_map(|kind| Self::from_type_name(kind, Some(schema)))
                     .collect::<Vec<_>>();
                 if types.is_empty() {
                     None
@@ -172,15 +200,15 @@ impl JsonParamType {
     }
 
     /// Normalize one JSON schema type name.
-    fn from_type_name(kind: &str) -> Option<Self> {
+    fn from_type_name(kind: &str, schema: Option<&Map<String, Value>>) -> Option<Self> {
         let kind = kind.trim().to_ascii_lowercase();
         match kind.as_str() {
             "string" | "str" | "text" | "varchar" | "char" | "enum" => Some(Self::String),
             "integer" | "int" => Some(Self::Integer),
             "number" | "float" | "double" => Some(Self::Number),
             "boolean" | "bool" | "binary" => Some(Self::Boolean),
-            "object" | "dict" | "map" => Some(Self::Object),
-            "array" | "arr" | "list" | "sequence" => Some(Self::Array),
+            "object" | "dict" | "map" => Some(Self::object_from_schema(schema)),
+            "array" | "arr" | "list" | "sequence" => Some(Self::array_from_schema(schema)),
             "null" => Some(Self::Null),
             _ if kind.starts_with("int")
                 || kind.starts_with("uint")
@@ -191,12 +219,52 @@ impl JsonParamType {
                 Some(Self::Integer)
             }
             _ if kind.starts_with("num") || kind.starts_with("float") => Some(Self::Number),
-            _ if kind.starts_with("dict") => Some(Self::Object),
-            _ if kind.starts_with("list") => Some(Self::Array),
+            _ if kind.starts_with("dict") => Some(Self::object_from_schema(schema)),
+            _ if kind.starts_with("list") => Some(Self::array_from_schema(schema)),
             _ => None,
         }
     }
 
+    /// Normalize object schema fields.
+    fn object_from_schema(schema: Option<&Map<String, Value>>) -> Self {
+        let properties = schema
+            .and_then(|schema| schema.get("properties"))
+            .and_then(Value::as_object)
+            .map(|properties| {
+                properties
+                    .iter()
+                    .filter_map(|(name, schema)| {
+                        Self::from_schema(schema).map(|param_type| (name.clone(), param_type))
+                    })
+                    .collect()
+            })
+            .unwrap_or_default();
+
+        let additional_properties =
+            schema.and_then(|schema| schema.get("additionalProperties")).and_then(|schema| {
+                if schema.is_object() {
+                    Self::from_schema(schema).map(Box::new)
+                } else {
+                    None
+                }
+            });
+
+        Self::Object {
+            properties,
+            additional_properties,
+        }
+    }
+
+    /// Normalize array schema fields.
+    fn array_from_schema(schema: Option<&Map<String, Value>>) -> Self {
+        let items = schema
+            .and_then(|schema| schema.get("items"))
+            .and_then(Self::from_schema)
+            .map(Box::new);
+
+        Self::Array { items }
+    }
+
     /// Collapse a candidate type list into one normalized type.
     fn one_of(mut types: Vec<Self>) -> Self {
         if types.len() == 1 {
@@ -207,23 +275,126 @@ impl JsonParamType {
     }
 }
 
-/// Convert one raw string value to a normalized JSON type.
-fn convert_value(param_type: &JsonParamType, value: &str) -> Option<Value> {
-    match param_type {
-        JsonParamType::String => Some(Value::String(value.to_string())),
-        JsonParamType::Integer => value.parse::<i64>().ok().map(Number::from).map(Value::Number),
-        JsonParamType::Number => convert_number(value),
-        JsonParamType::Boolean => convert_boolean(value),
-        JsonParamType::Object | JsonParamType::Array => serde_json::from_str(value).ok(),
-        JsonParamType::Null => value.eq_ignore_ascii_case("null").then_some(Value::Null),
-        JsonParamType::OneOf(types) => {
-            types.iter().find_map(|param_type| convert_value(param_type, value))
+/// Convert one parameter input to a normalized JSON value.
+fn convert_with_optional_schema(param_type: Option<&JsonParamType>, input: &ParamInput) -> Value {
+    // For literal `null`, always convert to JSON null value.
+    if let ParamInput::Text(value) = input
+        && value.eq_ignore_ascii_case("null")
+    {
+        return Value::Null;
+    }
+
+    // If we have a schema, try to convert the value using it.
+    if let Some(param_type) = param_type
+        && let Some(value) = try_convert_value(param_type, input)
+    {
+        return value;
+    }
+    // We don't have a schema, or conversion failed, use fallback logic.
+    match input {
+        ParamInput::Text(value) => Value::String(value.clone()),
+        ParamInput::Elements(elements) => {
+            // Convert structured input to object without a schema.
+            Value::Object(convert_elements_to_object(elements, &BTreeMap::new(), None))
         }
     }
 }
 
+/// Convert one parameter input to a normalized JSON type.
+fn try_convert_value(param_type: &JsonParamType, input: &ParamInput) -> Option<Value> {
+    match input {
+        ParamInput::Text(value) => try_convert_text_value(param_type, value),
+        ParamInput::Elements(elements) => try_convert_elements_value(param_type, elements),
+    }
+}
+
+/// Convert one raw string value to a normalized JSON type.
+fn try_convert_text_value(param_type: &JsonParamType, value: &str) -> Option<Value> {
+    match param_type {
+        JsonParamType::String => Some(Value::String(value.to_string())),
+        JsonParamType::Integer => value.parse::<i64>().ok().map(Number::from).map(Value::Number),
+        JsonParamType::Number => try_convert_number(value),
+        JsonParamType::Boolean => try_convert_boolean(value),
+        JsonParamType::Object { .. } if value.is_empty() => Some(Value::Object(Map::new())),
+        JsonParamType::Array { .. } if value.is_empty() => Some(Value::Array(Vec::new())),
+        JsonParamType::Object { .. } | JsonParamType::Array { .. } => {
+            // For composite types with string input, simply interpret the string as JSON.
+            serde_json::from_str(value).ok()
+        }
+        JsonParamType::Null => value.eq_ignore_ascii_case("null").then_some(Value::Null),
+        JsonParamType::OneOf(types) => {
+            types.iter().find_map(|param_type| try_convert_text_value(param_type, value))
+        }
+    }
+}
+
+/// Convert one structured parameter input to a normalized JSON type.
+fn try_convert_elements_value(
+    param_type: &JsonParamType,
+    elements: &[ParamElement],
+) -> Option<Value> {
+    match param_type {
+        JsonParamType::Object {
+            properties,
+            additional_properties,
+        } => Some(Value::Object(convert_elements_to_object(
+            elements,
+            properties,
+            additional_properties.as_deref(),
+        ))),
+        JsonParamType::Array { items } => Some(Value::Array(
+            // Collect all child elements into an array, regardless of their names.
+            elements
+                .iter()
+                .map(|element| convert_with_optional_schema(items.as_deref(), &element.value))
+                .collect(),
+        )),
+        JsonParamType::OneOf(types) => types
+            .iter()
+            .find_map(|param_type| try_convert_elements_value(param_type, elements)),
+
+        // Primitive types can't be converted from structured input.
+        JsonParamType::String
+        | JsonParamType::Integer
+        | JsonParamType::Number
+        | JsonParamType::Boolean
+        | JsonParamType::Null => None,
+    }
+}
+
+/// Convert structured elements to an object, using field schemas when present.
+fn convert_elements_to_object(
+    elements: &[ParamElement],
+    properties: &BTreeMap<String, JsonParamType>,
+    additional_properties: Option<&JsonParamType>,
+) -> Map<String, Value> {
+    let mut object = Map::with_capacity(elements.len());
+    for element in elements {
+        let param_type = properties.get(&element.name).or(additional_properties);
+        let value = convert_with_optional_schema(param_type, &element.value);
+        insert_object_value(&mut object, element.name.clone(), value);
+    }
+    object
+}
+
+/// Insert an object field while preserving duplicate keys as arrays.
+fn insert_object_value(object: &mut Map<String, Value>, key: String, value: Value) {
+    if let Some(existing) = object.get_mut(&key) {
+        match existing {
+            // Collect values under the same key into an array.
+            Value::Array(values) => values.push(value),
+            existing => {
+                let first = std::mem::replace(existing, Value::Null);
+                *existing = Value::Array(vec![first, value]);
+            }
+        }
+    } else {
+        object.insert(key, value);
+    }
+}
+
 /// Convert one raw string value to a JSON number.
-fn convert_number(value: &str) -> Option<Value> {
+fn try_convert_number(value: &str) -> Option<Value> {
     serde_json::from_str::<Number>(value)
         .or_else(|_| value.parse::<i64>().map(Number::from))
         .or_else(|_| value.parse::<f64>().ok().and_then(Number::from_f64).ok_or(()))
@@ -232,7 +403,7 @@ fn convert_number(value: &str) -> Option<Value> {
 }
 
 /// Convert one raw string value to a boolean.
-fn convert_boolean(value: &str) -> Option<Value> {
+fn try_convert_boolean(value: &str) -> Option<Value> {
     match value.trim().to_ascii_lowercase().as_str() {
         "true" | "1" => Some(Value::Bool(true)),
         "false" | "0" => Some(Value::Bool(false)),
@@ -242,9 +413,9 @@ fn convert_boolean(value: &str) -> Option<Value> {
 
 #[cfg(test)]
 mod tests {
-    use serde_json::json;
+    use serde_json::{Value, json};
 
-    use super::{ToolSchema, ToolSchemas};
+    use super::{ParamElement, ParamInput, ToolSchema, ToolSchemas};
     use crate::Tool;
 
     fn test_tool(name: &str, parameters: serde_json::Value) -> Tool {
@@ -260,8 +431,8 @@ mod tests {
     fn invalid_schema_converts_everything_as_string() {
         let params = ToolSchema::from_schema(&json!({ "type": "object" }));
 
-        assert_eq!(params.convert("count", "42"), json!("42"));
-        assert_eq!(params.convert("count", "null"), json!(null));
+        assert_eq!(params.convert("count", text("42")), json!("42"));
+        assert_eq!(params.convert("count", text("null")), json!(null));
     }
 
     #[test]
@@ -275,9 +446,9 @@ mod tests {
             }
         }));
 
-        assert_eq!(params.convert("unknown_schema", "42"), json!("42"));
-        assert_eq!(params.convert("unknown_type", "42"), json!("42"));
-        assert_eq!(params.convert("known", "42"), json!(42));
+        assert_eq!(params.convert("unknown_schema", text("42")), json!("42"));
+        assert_eq!(params.convert("unknown_type", text("42")), json!("42"));
+        assert_eq!(params.convert("known", text("42")), json!(42));
     }
 
     #[test]
@@ -298,16 +469,25 @@ mod tests {
             }
         }));
 
-        assert_eq!(params.convert("text", "42"), json!("42"));
-        assert_eq!(params.convert("count", "42"), json!(42));
-        assert_eq!(params.convert("size", "5.0"), json!(5.0));
-        assert_eq!(params.convert("ratio", "2.5"), json!(2.5));
-        assert_eq!(params.convert("enabled", "1"), json!(true));
-        assert_eq!(params.convert("payload", r#"{"k":1}"#), json!({ "k": 1 }));
-        assert_eq!(params.convert("mapping", r#"{"k":1}"#), json!({ "k": 1 }));
-        assert_eq!(params.convert("items", "[1,2]"), json!([1, 2]));
-        assert_eq!(params.convert("names", r#"["a","b"]"#), json!(["a", "b"]));
-        assert_eq!(params.convert("nothing", "null"), json!(null));
+        assert_eq!(params.convert("text", text("42")), json!("42"));
+        assert_eq!(params.convert("count", text("42")), json!(42));
+        assert_eq!(params.convert("size", text("5.0")), json!(5.0));
+        assert_eq!(params.convert("ratio", text("2.5")), json!(2.5));
+        assert_eq!(params.convert("enabled", text("1")), json!(true));
+        assert_eq!(
+            params.convert("payload", text(r#"{"k":1}"#)),
+            json!({ "k": 1 })
+        );
+        assert_eq!(
+            params.convert("mapping", text(r#"{"k":1}"#)),
+            json!({ "k": 1 })
+        );
+        assert_eq!(params.convert("items", text("[1,2]")), json!([1, 2]));
+        assert_eq!(
+            params.convert("names", text(r#"["a","b"]"#)),
+            json!(["a", "b"])
+        );
+        assert_eq!(params.convert("nothing", text("null")), json!(null));
     }
 
     #[test]
@@ -321,19 +501,40 @@ mod tests {
 
         assert_eq!(converted_number_text(&params, "5"), "5");
         assert_eq!(converted_number_text(&params, "5.0"), "5.0");
-        assert_eq!(converted_number_text(&params, "5.00"), "5.00");
-        assert_eq!(converted_number_text(&params, "1e0"), "1e+0");
         assert_eq!(converted_number_text(&params, "5."), "5.0");
         assert_eq!(converted_number_text(&params, "+1"), "1");
         assert_eq!(converted_number_text(&params, "+1.0"), "1.0");
-        assert_eq!(
-            converted_number_text(&params, "9223372036854775807.5"),
-            "9223372036854775807.5"
-        );
+
+        // TODO: we cannot preserve the original number precision by enabling `serde_json`'s
+        // `arbitrary_precision` feature, otherwise the test
+        // `serialized_json_numbers_do_not_leak_serde_private_representation` will fail.
+        // See issue: https://github.com/mitsuhiko/minijinja/issues/641
+
+        // assert_eq!(converted_number_text(&params, "5.00"), "5.00");
+        // assert_eq!(converted_number_text(&params, "1e0"), "1e+0");
+        // assert_eq!(
+        //     converted_number_text(&params, "9223372036854775807.5"),
+        //     "9223372036854775807.5"
+        // );
     }
 
     fn converted_number_text(params: &ToolSchema, value: &str) -> String {
-        serde_json::to_string(&params.convert("value", value)).unwrap()
+        serde_json::to_string(&params.convert("value", text(value))).unwrap()
+    }
+
+    fn text(value: &str) -> ParamInput {
+        ParamInput::Text(value.to_string())
+    }
+
+    fn elem(name: &str, value: ParamInput) -> ParamElement {
+        ParamElement {
+            name: name.to_string(),
+            value,
+        }
+    }
+
+    fn elements(elements: Vec<ParamElement>) -> ParamInput {
+        ParamInput::Elements(elements)
     }
 
     #[test]
@@ -350,12 +551,12 @@ mod tests {
             }
         }));
 
-        assert_eq!(params.convert("s", "x"), json!("x"));
-        assert_eq!(params.convert("i", "7"), json!(7));
-        assert_eq!(params.convert("n", "7.5"), json!(7.5));
-        assert_eq!(params.convert("b", "true"), json!(true));
-        assert_eq!(params.convert("a", "[1]"), json!([1]));
-        assert_eq!(params.convert("o", r#"{"x":1}"#), json!({ "x": 1 }));
+        assert_eq!(params.convert("s", text("x")), json!("x"));
+        assert_eq!(params.convert("i", text("7")), json!(7));
+        assert_eq!(params.convert("n", text("7.5")), json!(7.5));
+        assert_eq!(params.convert("b", text("true")), json!(true));
+        assert_eq!(params.convert("a", text("[1]")), json!([1]));
+        assert_eq!(params.convert("o", text(r#"{"x":1}"#)), json!({ "x": 1 }));
     }
 
     #[test]
@@ -373,8 +574,8 @@ mod tests {
             }
         }));
 
-        assert_eq!(integer_first.convert("value", "42"), json!(42));
-        assert_eq!(string_first.convert("value", "42"), json!("42"));
+        assert_eq!(integer_first.convert("value", text("42")), json!(42));
+        assert_eq!(string_first.convert("value", text("42")), json!("42"));
     }
 
     #[test]
@@ -396,9 +597,9 @@ mod tests {
             }
         }));
 
-        assert_eq!(params.convert("choice", "42"), json!(42));
+        assert_eq!(params.convert("choice", text("42")), json!(42));
         assert_eq!(
-            params.convert("fallback_object", r#"{"x":1}"#),
+            params.convert("fallback_object", text(r#"{"x":1}"#)),
             json!({ "x": 1 })
         );
     }
@@ -414,9 +615,12 @@ mod tests {
             }
         }));
 
-        assert_eq!(params.convert("choice", "a"), json!("a"));
-        assert_eq!(params.convert("items", "[1,2]"), json!([1, 2]));
-        assert_eq!(params.convert("payload", r#"{"x":1}"#), json!({ "x": 1 }));
+        assert_eq!(params.convert("choice", text("a")), json!("a"));
+        assert_eq!(params.convert("items", text("[1,2]")), json!([1, 2]));
+        assert_eq!(
+            params.convert("payload", text(r#"{"x":1}"#)),
+            json!({ "x": 1 })
+        );
     }
 
     #[test]
@@ -518,4 +722,162 @@ mod tests {
         assert_eq!(converted.get("topn"), Some(&json!("5")));
         assert_eq!(converted.get("nullish"), Some(&json!(null)));
     }
+
+    #[test]
+    fn converts_structured_inputs_with_recursive_schema() {
+        let schemas = ToolSchemas::from_tools(&[test_tool(
+            "create_order",
+            json!({
+                "type": "object",
+                "properties": {
+                    "user_id": { "type": "integer" },
+                    "urgent": { "type": "boolean" },
+                    "note": { "type": "string" },
+                    "nil": { "type": "string" },
+                    "shipping": {
+                        "type": "object",
+                        "properties": {
+                            "city": { "type": "string" },
+                            "zip": { "type": "integer" }
+                        }
+                    },
+                    "items": {
+                        "type": "array",
+                        "items": {
+                            "type": "object",
+                            "properties": {
+                                "sku": { "type": "string" },
+                                "qty": { "type": "integer" }
+                            }
+                        }
+                    },
+                    "metadata": {
+                        "type": "object",
+                        "additionalProperties": { "type": "integer" }
+                    },
+                    "duplicate_demo": {
+                        "type": "object",
+                        "properties": {
+                            "tag": { "type": "string" }
+                        }
+                    },
+                    "schema_mismatch_array": {
+                        "type": "array",
+                        "items": { "type": "integer" }
+                    },
+                    "closed_object": {
+                        "type": "object",
+                        "additionalProperties": false
+                    },
+                    "open_object": {
+                        "type": "object",
+                        "additionalProperties": true
+                    },
+                    "payload_text": { "type": "object" },
+                    "items_text": { "type": "array" }
+                }
+            }),
+        )]);
+
+        let converted = schemas.convert_params_with_schema(
+            "create_order",
+            vec![
+                ("user_id".to_string(), text("42")),
+                ("urgent".to_string(), text("true")),
+                ("note".to_string(), text("Please leave at front desk.")),
+                ("nil".to_string(), text("NULL")),
+                (
+                    "shipping".to_string(),
+                    elements(vec![
+                        elem("city", text("Singapore")),
+                        elem("zip", text("018956")),
+                    ]),
+                ),
+                (
+                    "items".to_string(),
+                    elements(vec![
+                        elem(
+                            "item1",
+                            elements(vec![elem("sku", text("book-001")), elem("qty", text("2"))]),
+                        ),
+                        elem(
+                            "item2",
+                            elements(vec![elem("sku", text("pen-007")), elem("qty", text("5"))]),
+                        ),
+                    ]),
+                ),
+                (
+                    "metadata".to_string(),
+                    elements(vec![elem("score", text("42")), elem("rank", text("7"))]),
+                ),
+                (
+                    "duplicate_demo".to_string(),
+                    elements(vec![elem("tag", text("a")), elem("tag", text("b"))]),
+                ),
+                (
+                    "closed_object".to_string(),
+                    elements(vec![elem("unknown", text("x"))]),
+                ),
+                (
+                    "open_object".to_string(),
+                    elements(vec![elem("unknown", text("y"))]),
+                ),
+                ("payload_text".to_string(), text(r#"{"x":1}"#)),
+                ("items_text".to_string(), text("[1,2]")),
+                (
+                    "unknown_struct".to_string(),
+                    elements(vec![
+                        elem("a", text("1")),
+                        elem("a", text("2")),
+                        elem("nil", text("null")),
+                    ]),
+                ),
+            ],
+        );
+
+        assert_eq!(
+            Value::Object(converted),
+            json!({
+                "user_id": 42,
+                "urgent": true,
+                "note": "Please leave at front desk.",
+                "nil": null,
+                "shipping": {
+                    "city": "Singapore",
+                    "zip": 18956
+                },
+                "items": [
+                    {
+                        "sku": "book-001",
+                        "qty": 2
+                    },
+                    {
+                        "sku": "pen-007",
+                        "qty": 5
+                    }
+                ],
+                "metadata": {
+                    "score": 42,
+                    "rank": 7
+                },
+                "duplicate_demo": {
+                    "tag": ["a", "b"]
+                },
+                "closed_object": {
+                    "unknown": "x"
+                },
+                "open_object": {
+                    "unknown": "y"
+                },
+                "payload_text": {
+                    "x": 1
+                },
+                "items_text": [1, 2],
+                "unknown_struct": {
+                    "a": ["1", "2"],
+                    "nil": null
+                }
+            })
+        );
+    }
 }

setup.py

@@ -36,6 +36,8 @@ ROOT_DIR = Path(__file__).parent
 logger = logging.getLogger(__name__)
 
 PRECOMPILED_RUST_FRONTEND_PATH = ROOT_DIR / "vllm" / "vllm-rs"
+PRECOMPILED_RUST_EXTENSION_GLOB = "_rust_*.so"
+PRECOMPILED_RUST_EXTENSION_MEMBER_REGEX = re.compile(r"vllm/_rust_[^/]*\.so$")
 
 # cannot import envs directly because it depends on vllm,
 #  which is not installed yet
@@ -54,6 +56,59 @@ def should_require_rust_frontend() -> bool:
     return value.lower() not in ("", "0", "false", "no")
 
 
+# Rust frontend binary, built via setuptools-rust and installed into the
+# package directory alongside the Python modules.
+# TODO: we may use `RustBin` to directly install it into `bin` directory, but this
+# requires extra work on using precompiled binaries.
+rust_extensions = [
+    RustExtension(
+        target="vllm.vllm-rs",
+        path="rust/src/cmd/Cargo.toml",
+        args=["--bin", "vllm-rs"],
+        features=["native-tls-vendored"],
+        binding=Binding.Exec,
+        optional=not should_require_rust_frontend(),
+    ),
+    RustExtension(
+        target="vllm._rust_tool_parser",
+        path="rust/src/tool-parser/python/Cargo.toml",
+        features=["extension-module"],
+        binding=Binding.PyO3,
+        optional=not should_require_rust_frontend(),
+    ),
+]
+
+
+def get_precompiled_rust_extension_paths() -> list[Path]:
+    return sorted((ROOT_DIR / "vllm").glob(PRECOMPILED_RUST_EXTENSION_GLOB))
+
+
+def get_expected_rust_extension_module_names() -> list[str]:
+    """Return configured PyO3 Rust extension module names under ``vllm``."""
+    module_names = []
+    for rust_extension in rust_extensions:
+        if rust_extension.binding != Binding.PyO3:
+            continue
+
+        for target_name in rust_extension.target.values():
+            if target_name.startswith("vllm._rust_"):
+                module_names.append(target_name.rsplit(".", 1)[-1])
+
+    return module_names
+
+
+def get_missing_precompiled_rust_extension_modules() -> list[str]:
+    missing = []
+    for module_name in get_expected_rust_extension_module_names():
+        if not list((ROOT_DIR / "vllm").glob(f"{module_name}*.so")):
+            missing.append(module_name)
+    return missing
+
+
+def has_precompiled_rust_extensions() -> bool:
+    return not get_missing_precompiled_rust_extension_modules()
+
+
 if sys.platform.startswith("darwin") and VLLM_TARGET_DEVICE != "cpu":
     logger.warning("VLLM_TARGET_DEVICE automatically set to `cpu` due to macOS")
     VLLM_TARGET_DEVICE = "cpu"
@@ -421,19 +476,33 @@ class precompiled_build_ext(build_ext):
 
 
 class precompiled_build_rust(build_rust):
-    """Skips local Rust builds when the precompiled wheel already ships vllm-rs."""
+    """Skips local Rust builds when all precompiled Rust artifacts are present."""
 
     def run(self) -> None:
-        if PRECOMPILED_RUST_FRONTEND_PATH.exists():
+        if (
+            PRECOMPILED_RUST_FRONTEND_PATH.exists()
+            and has_precompiled_rust_extensions()
+        ):
             logger.info(
-                "Skipping local Rust build: using precompiled %s",
+                "Skipping local Rust build: using precompiled %s and %s",
                 PRECOMPILED_RUST_FRONTEND_PATH,
+                get_precompiled_rust_extension_paths(),
             )
             return
 
+        missing = []
+        if not PRECOMPILED_RUST_FRONTEND_PATH.exists():
+            missing.append(str(PRECOMPILED_RUST_FRONTEND_PATH))
+        missing_rust_extensions = get_missing_precompiled_rust_extension_modules()
+        if missing_rust_extensions:
+            missing.extend(
+                str(ROOT_DIR / "vllm" / f"{module_name}*.so")
+                for module_name in missing_rust_extensions
+            )
         logger.warning(
-            "Precompiled wheel did not provide %s; falling back to local Rust build.",
-            PRECOMPILED_RUST_FRONTEND_PATH,
+            "Precompiled wheel did not provide all Rust artifacts (%s); "
+            "falling back to local Rust build.",
+            ", ".join(missing),
         )
         super().run()
 
@@ -756,6 +825,14 @@ class precompiled_wheel_utils:
                     if member.filename in exact_members:
                         file_members.append(member)
                         continue
+                    if (
+                        extract_rust_frontend
+                        and PRECOMPILED_RUST_EXTENSION_MEMBER_REGEX.match(
+                            member.filename
+                        )
+                    ):
+                        file_members.append(member)
+                        continue
 
                     if not extract_extensions:
                         continue
@@ -1127,6 +1204,10 @@ if PRECOMPILED_RUST_FRONTEND_PATH.exists():
     vllm_files = package_data.setdefault("vllm", [])
     if "vllm-rs" not in vllm_files:
         vllm_files.append("vllm-rs")
+vllm_files = package_data.setdefault("vllm", [])
+for rust_extension_path in get_precompiled_rust_extension_paths():
+    if rust_extension_path.name not in vllm_files:
+        vllm_files.append(rust_extension_path.name)
 
 if _no_device():
     ext_modules = []
@@ -1139,24 +1220,13 @@ else:
         if USE_PRECOMPILED_EXTENSIONS
         else cmake_build_ext,
     }
-if USE_PRECOMPILED_RUST_FRONTEND or PRECOMPILED_RUST_FRONTEND_PATH.exists():
+if (
+    USE_PRECOMPILED_RUST_FRONTEND
+    or PRECOMPILED_RUST_FRONTEND_PATH.exists()
+    or has_precompiled_rust_extensions()
+):
     cmdclass["build_rust"] = precompiled_build_rust
 
-# Rust frontend binary, built via setuptools-rust and installed into the
-# package directory alongside the Python modules.
-# TODO: we may use `RustBin` to directly install it into `bin` directory, but this
-# requires extra work on using precompiled binaries.
-rust_extensions = [
-    RustExtension(
-        target="vllm.vllm-rs",
-        path="rust/src/cmd/Cargo.toml",
-        args=["--bin", "vllm-rs"],
-        features=["native-tls-vendored"],
-        binding=Binding.Exec,
-        optional=not should_require_rust_frontend(),
-    ),
-]
-
 setup(
     # static metadata should rather go in pyproject.toml
     version=get_vllm_version(),

tests/kernels/attention/test_minimax_m3.py

@@ -0,0 +1,810 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Correctness tests for MiniMax M3 sparse prefill attention kernels."""
+
+import pytest
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.models.minimax_m3.common.ops.index_topk import (
+    minimax_m3_index_topk,
+    minimax_m3_index_topk_decode,
+)
+from vllm.models.minimax_m3.common.ops.sparse_attn import (
+    minimax_m3_sparse_attn,
+    minimax_m3_sparse_attn_decode,
+)
+from vllm.models.minimax_m3.common.sparse_attention import (
+    MiniMaxM3IndexerBackend,
+    MiniMaxM3SparseBackend,
+)
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_cutedsl
+from vllm.v1.attention.backends.utils import set_kv_cache_layout
+from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
+from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache
+from vllm.v1.worker.utils import AttentionGroup
+
+if not current_platform.is_cuda():
+    pytest.skip("MiniMax M3 attention kernels require CUDA.", allow_module_level=True)
+
+
+@pytest.fixture
+def kv_layout(request):
+    """Set the global KV cache layout for one test and restore it after."""
+    set_kv_cache_layout(request.param)
+    try:
+        yield request.param
+    finally:
+        set_kv_cache_layout(None)
+
+
+def _stride_order_for(backend: type[MiniMaxM3SparseBackend], ndim: int) -> tuple:
+    """Mirror the allocator's stride-order resolution (identity fallback)."""
+    try:
+        stride_order = backend.get_kv_cache_stride_order()
+        assert len(stride_order) == ndim
+    except (AttributeError, NotImplementedError):
+        stride_order = tuple(range(ndim))
+    return stride_order
+
+
+def _allocate_main_kv_via_contract(
+    num_pages: int, device: torch.device | str = "cuda"
+) -> torch.Tensor:
+    """Build the main KV cache exactly as the production allocator does for the
+    currently active layout: allocate the physical (permuted) tensor, then
+    expose the inverse-permuted logical-NHD view the backend sees."""
+    logical_shape = MiniMaxM3SparseBackend.get_kv_cache_shape(
+        num_pages, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
+    )
+    stride_order = _stride_order_for(MiniMaxM3SparseBackend, len(logical_shape))
+    physical_shape = tuple(logical_shape[i] for i in stride_order)
+    inv_order = [stride_order.index(i) for i in range(len(stride_order))]
+    raw = torch.randn(physical_shape, device=device, dtype=DTYPE)
+    return raw.permute(*inv_order)
+
+
+NUM_Q_HEADS = 32
+NUM_KV_HEADS = 2
+HEAD_DIM = 128
+BLOCK_SIZE = 128
+DTYPE = torch.bfloat16
+SM_SCALE = HEAD_DIM**-0.5
+TOPK = 16
+
+
+# Index top-k kernels.
+def _reference_index_topk(
+    idx_q: torch.Tensor,
+    index_kv_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    q_lens: torch.Tensor,
+    seq_lens: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    topk: int,
+    init_blocks: int,
+    local_blocks: int,
+    sm_scale: float,
+) -> torch.Tensor:
+    total_q, num_idx_heads, _ = idx_q.shape
+    out = torch.full(
+        (num_idx_heads, total_q, topk), -1, device=idx_q.device, dtype=torch.int32
+    )
+
+    q_start = 0
+    for req_id, (q_len, seq_len, prefix_len) in enumerate(
+        zip(q_lens.tolist(), seq_lens.tolist(), prefix_lens.tolist())
+    ):
+        q_end = q_start + q_len
+        q = idx_q[q_start:q_end]
+        num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
+        pages = block_table[req_id, :num_blocks]
+        k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1)
+        score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale
+
+        q_pos = prefix_len + torch.arange(q_len, device=idx_q.device)
+        k_pos = torch.arange(k.shape[0], device=idx_q.device)
+        score.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf"))
+        score = score.reshape(num_idx_heads, q_len, num_blocks, BLOCK_SIZE)
+        score_tensor = score.max(dim=3).values
+
+        valid_blocks = (q_pos + BLOCK_SIZE) // BLOCK_SIZE
+        for local_q, num_valid_blocks in enumerate(valid_blocks.tolist()):
+            end = min(init_blocks, num_valid_blocks)
+            score_tensor[:, local_q, :end] = 1e30
+            start = max(0, num_valid_blocks - local_blocks)
+            score_tensor[:, local_q, start:num_valid_blocks] = 1e29
+
+            k = min(topk, num_valid_blocks)
+            topk_idx = score_tensor[:, local_q].topk(k, dim=1).indices
+            out[:, q_start + local_q, :k] = topk_idx
+        q_start = q_end
+
+    return out
+
+
+def test_prefill_index_topk_correctness():
+    topk = 6
+    init_blocks = 0
+    local_blocks = 1
+    num_idx_heads = 2
+    head_dim = 16
+    q_lens = torch.tensor((4, 3), device="cuda", dtype=torch.int32)
+    prefix_lens = torch.tensor((0, 1024), device="cuda", dtype=torch.int32)
+    seq_lens = prefix_lens + q_lens
+    batch = q_lens.numel()
+    max_seq_len = seq_lens.max().item()
+    max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
+    num_pages = batch * max_blocks
+
+    cu_seqlens = torch.zeros(batch + 1, device="cuda", dtype=torch.int32)
+    cu_seqlens[1:] = q_lens.cumsum(0)
+    block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(
+        batch, max_blocks
+    )
+    idx_q = torch.ones(q_lens.sum().item(), num_idx_heads, head_dim, device="cuda")
+    index_kv_cache = torch.empty(num_pages, BLOCK_SIZE, head_dim, device="cuda")
+    for req_id in range(batch):
+        for block_id in range(max_blocks):
+            page = block_table[req_id, block_id]
+            index_kv_cache[page].fill_(block_id + 1)
+
+    actual = minimax_m3_index_topk(
+        idx_q,
+        index_kv_cache,
+        block_table,
+        cu_seqlens,
+        seq_lens,
+        prefix_lens,
+        max_query_len=q_lens.max().item(),
+        max_seq_len=max_seq_len,
+        topk=topk,
+        init_blocks=init_blocks,
+        local_blocks=local_blocks,
+        num_kv_heads=num_idx_heads,
+        sm_scale=head_dim**-0.5,
+    )
+    expected = _reference_index_topk(
+        idx_q,
+        index_kv_cache,
+        block_table,
+        q_lens,
+        seq_lens,
+        prefix_lens,
+        topk,
+        init_blocks,
+        local_blocks,
+        head_dim**-0.5,
+    )
+    assert torch.equal(actual, expected)
+
+
+def test_decode_index_topk_correctness():
+    topk = 6
+    init_blocks = 0
+    local_blocks = 1
+    num_idx_heads = 2
+    head_dim = 16
+    seq_lens = torch.tensor((7, 129, 1025), device="cuda", dtype=torch.int32)
+    q_lens = torch.ones_like(seq_lens)
+    prefix_lens = seq_lens - 1
+    batch = seq_lens.numel()
+    max_seq_len = seq_lens.max().item()
+    max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
+    num_pages = batch * max_blocks
+
+    block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(
+        batch, max_blocks
+    )
+    idx_q = torch.ones(batch, num_idx_heads, head_dim, device="cuda")
+    index_kv_cache = torch.empty(num_pages, BLOCK_SIZE, head_dim, device="cuda")
+    for req_id in range(batch):
+        for block_id in range(max_blocks):
+            page = block_table[req_id, block_id]
+            index_kv_cache[page].fill_(block_id + 1)
+
+    actual = minimax_m3_index_topk_decode(
+        idx_q,
+        index_kv_cache,
+        block_table,
+        seq_lens,
+        max_seq_len=max_seq_len,
+        topk=topk,
+        init_blocks=init_blocks,
+        local_blocks=local_blocks,
+        num_kv_heads=num_idx_heads,
+        sm_scale=head_dim**-0.5,
+    )
+    expected = _reference_index_topk(
+        idx_q,
+        index_kv_cache,
+        block_table,
+        q_lens,
+        seq_lens,
+        prefix_lens,
+        topk,
+        init_blocks,
+        local_blocks,
+        head_dim**-0.5,
+    )
+    assert torch.equal(actual, expected)
+
+
+# Sparse attention kernels.
+def _reference_sparse_attn(
+    q: torch.Tensor,
+    kv_cache: torch.Tensor,
+    topk_idx: torch.Tensor,
+    block_table: torch.Tensor,
+    q_lens: torch.Tensor,
+    seq_lens: torch.Tensor,
+    prefix_lens: torch.Tensor,
+) -> torch.Tensor:
+    out = torch.empty_like(q, dtype=torch.float32)
+    gqa_group_size = NUM_Q_HEADS // NUM_KV_HEADS
+    q_start = 0
+    for req_id, (q_len, seq_len, prefix_len) in enumerate(
+        zip(q_lens.tolist(), seq_lens.tolist(), prefix_lens.tolist())
+    ):
+        q_end = q_start + q_len
+        q_req = q[q_start:q_end]
+        positions = torch.arange(seq_len, device="cuda")
+        pages = block_table[req_id, positions // BLOCK_SIZE]
+        rows = positions % BLOCK_SIZE
+        k_req = kv_cache[pages, 0, rows]
+        v_req = kv_cache[pages, 1, rows].float()
+
+        q_pos = prefix_len + torch.arange(q_len, device="cuda")
+        key_blocks = positions // BLOCK_SIZE
+        causal_mask = positions.unsqueeze(0) <= q_pos.unsqueeze(1)
+
+        for kv_head in range(NUM_KV_HEADS):
+            selected = topk_idx[kv_head, q_start:q_end]
+            selected_mask = (key_blocks[None, :, None] == selected[:, None, :]).any(-1)
+            mask = causal_mask & selected_mask
+            head_start = kv_head * gqa_group_size
+            head_end = head_start + gqa_group_size
+
+            q_heads = q_req[:, head_start:head_end].transpose(0, 1)
+            k_head = k_req[:, kv_head].T.expand(gqa_group_size, -1, -1)
+            scores = torch.bmm(q_heads, k_head, out_dtype=torch.float32)
+            scores = scores.transpose(0, 1) * SM_SCALE
+            probs = torch.softmax(
+                scores.masked_fill(~mask[:, None, :], -float("inf")), -1
+            )
+            out[q_start:q_end, head_start:head_end] = torch.einsum(
+                "qhk,kd->qhd", probs, v_req[:, kv_head]
+            )
+        q_start += q_len
+    return out.to(q.dtype)
+
+
+@pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
+@pytest.mark.parametrize("backend", ["triton", "cutedsl"])
+@pytest.mark.parametrize(
+    ("q_lens", "kv_lens"),
+    [
+        ((129, 257), (129, 257)),
+        ((65, 129, 257), (129, 257, 385)),
+    ],
+)
+def test_prefill_sparse_attention_correctness(
+    kv_layout: str,
+    backend: str,
+    q_lens: tuple[int, ...],
+    kv_lens: tuple[int, ...],
+):
+    if backend == "cutedsl":
+        if not current_platform.is_device_capability_family(100):
+            pytest.skip("MiniMax M3 CuteDSL prefill requires CUDA SM10x.")
+        if not has_cutedsl():
+            pytest.skip("cutedsl (cutlass) is not installed")
+
+    assert len(q_lens) == len(kv_lens)
+    assert all(kv_len >= q_len for q_len, kv_len in zip(q_lens, kv_lens))
+
+    # Build paged-KV metadata, including a non-identity page order.
+    batch = len(q_lens)
+    pages_per_req = [(kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE for kv_len in kv_lens]
+    max_blocks = max(pages_per_req)
+    num_pages = sum(pages_per_req)
+    physical_pages = torch.randperm(num_pages, device="cuda", dtype=torch.int32)
+    block_table = torch.zeros(batch, max_blocks, device="cuda", dtype=torch.int32)
+    base_page = 0
+    for req_id, num_req_pages in enumerate(pages_per_req):
+        block_table[req_id, :num_req_pages] = physical_pages[
+            base_page : base_page + num_req_pages
+        ]
+        base_page += num_req_pages
+
+    q_lens_t = torch.tensor(q_lens, device="cuda", dtype=torch.int32)
+    seq_lens = torch.tensor(kv_lens, device="cuda", dtype=torch.int32)
+    prefix_lens = seq_lens - q_lens_t
+    cu_seqlens = torch.zeros(batch + 1, device="cuda", dtype=torch.int32)
+    cu_seqlens[1:] = q_lens_t.cumsum(0)
+    cu_seqlens_k = torch.zeros(batch + 1, device="cuda", dtype=torch.int32)
+    cu_seqlens_k[1:] = seq_lens.cumsum(0)
+    total_q = sum(q_lens)
+    max_seqlen_q = max(q_lens)
+    max_seqlen_k = max(kv_lens)
+
+    q_shape = (total_q, NUM_Q_HEADS, HEAD_DIM)
+    q = torch.randn(q_shape, device="cuda", dtype=DTYPE)
+    # Allocate the main KV cache through the backend layout contract so the
+    # physical storage matches the active layout (contiguous NHD or strided
+    # HND), while the kernels and reference see the logical-NHD view.
+    kv_cache = _allocate_main_kv_via_contract(num_pages)
+
+    # Build sparse block indices with the same contract as the real M3 indexer:
+    # one forced local block, then score-selected older causal blocks.
+    topk_shape = (NUM_KV_HEADS, total_q, TOPK)
+    topk_idx = torch.full(topk_shape, -1, device="cuda", dtype=torch.int32)
+    q_start = 0
+    for q_len, prefix_len in zip(q_lens_t.tolist(), prefix_lens.tolist()):
+        for local_q in range(q_len):
+            current_block = (prefix_len + local_q) // BLOCK_SIZE
+            older_blocks = torch.randperm(
+                current_block, device="cuda", dtype=torch.int32
+            )
+            selected = torch.cat(
+                [
+                    torch.tensor([current_block], device="cuda", dtype=torch.int32),
+                    older_blocks[: TOPK - 1],
+                ]
+            )
+            topk_idx[:, q_start + local_q, : selected.numel()] = selected
+        q_start += q_len
+
+    actual = torch.empty_like(q)
+    if backend == "triton":
+        minimax_m3_sparse_attn(
+            q,
+            kv_cache,
+            topk_idx,
+            block_table,
+            cu_seqlens,
+            seq_lens,
+            prefix_lens,
+            max_seqlen_q,
+            NUM_KV_HEADS,
+            SM_SCALE,
+            actual,
+        )
+    else:
+        from vllm.models.minimax_m3.nvidia.ops.prefill_gqa_sparse import (
+            minimax_m3_sparse_attn_cutedsl,
+        )
+
+        minimax_m3_sparse_attn_cutedsl(
+            q,
+            kv_cache,
+            topk_idx,
+            block_table,
+            cu_seqlens,
+            cu_seqlens_k,
+            seq_lens,
+            max_seqlen_q,
+            max_seqlen_k,
+            NUM_KV_HEADS,
+            SM_SCALE,
+            actual,
+            total_kv_blocks=num_pages,
+        )
+
+    expected = _reference_sparse_attn(
+        q,
+        kv_cache,
+        topk_idx,
+        block_table,
+        q_lens_t,
+        seq_lens,
+        prefix_lens,
+    )
+    torch.accelerator.synchronize()
+
+    error = (actual.float() - expected.float()).abs()
+    assert error.mean().item() < 2.5e-4
+    assert error.max().item() < 1.7e-2
+
+
+def test_main_backend_layout_contract():
+    """The main sparse backend exposes the logical-NHD shape and the
+    flash_attn-style stride order for each layout."""
+    nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
+    logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
+    assert logical == (nb, 2, bs, h, d)
+    # The old HND-ordered shape is no longer the logical shape.
+    assert logical != (nb, 2, h, bs, d)
+
+    try:
+        set_kv_cache_layout("HND")
+        assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 3, 2, 4)
+        set_kv_cache_layout("NHD")
+        assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 2, 3, 4)
+    finally:
+        set_kv_cache_layout(None)
+
+    for layout in ("NHD", "HND"):
+        try:
+            set_kv_cache_layout(layout)
+            order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+        finally:
+            set_kv_cache_layout(None)
+        # Valid permutation: no duplicates, covers every axis.
+        assert set(order) == set(range(len(order)))
+
+    # M3 has no cross-layer KV blocks.
+    with pytest.raises(NotImplementedError):
+        MiniMaxM3SparseBackend.get_kv_cache_stride_order(
+            include_num_layers_dimension=True
+        )
+
+
+def test_main_backend_unknown_layout_raises(monkeypatch):
+    """An unrecognized layout (injected past env-var validation) is rejected."""
+    import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
+
+    monkeypatch.setattr(sparse_attn_mod, "get_kv_cache_layout", lambda: "BOGUS")
+    with pytest.raises(ValueError, match="Unknown cache layout format"):
+        MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+
+
+def test_indexer_backend_stride_order_is_identity():
+    """The 3-dim indexer cache must not inherit the parent's 5-element stride
+    order; it overrides to the 3-element identity so the allocator keeps the
+    contiguous layout."""
+    assert MiniMaxM3IndexerBackend.get_kv_cache_stride_order() == (0, 1, 2)
+
+    # Cross-layer (per-layer-stacked) KV blocks are not supported.
+    with pytest.raises(NotImplementedError):
+        MiniMaxM3IndexerBackend.get_kv_cache_stride_order(
+            include_num_layers_dimension=True
+        )
+
+    # The stride order matches the 3-dim indexer shape rank.
+    indexer_shape = MiniMaxM3IndexerBackend.get_kv_cache_shape(
+        5, BLOCK_SIZE, 1, HEAD_DIM
+    )
+    assert len(indexer_shape) == 3
+    assert _stride_order_for(MiniMaxM3IndexerBackend, len(indexer_shape)) == (0, 1, 2)
+
+
+def test_hnd_allocation_is_byte_identical_to_transpose():
+    """Under HND the backend-visible logical view is byte-identical to the
+    pre-change allocate-HND-then-transpose(2, 3) workaround."""
+    nb, bs, h, d = 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
+    logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
+    try:
+        set_kv_cache_layout("HND")
+        stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+    finally:
+        set_kv_cache_layout(None)
+
+    physical_shape = tuple(logical[i] for i in stride_order)
+    # The physical (permuted) shape equals the old hardcoded HND shape.
+    assert physical_shape == (nb, 2, h, bs, d)
+
+    inv_order = [stride_order.index(i) for i in range(len(stride_order))]
+    raw = torch.empty(physical_shape, device="cuda", dtype=DTYPE)
+    view = raw.permute(*inv_order)
+    expected = raw.view((nb, 2, h, bs, d)).transpose(2, 3)
+
+    assert view.shape == expected.shape
+    assert view.stride() == expected.stride()
+    assert view.storage_offset() == expected.storage_offset()
+
+    # Negative: the identity (wrong) stride order under HND does not reproduce
+    # the transpose view.
+    wrong_view = raw.view(logical)
+    assert wrong_view.stride() != expected.stride()
+
+
+def test_main_cache_is_block_first_and_unpadded():
+    """The allocator's contiguous-view branch (not the padded-strided branch)
+    is used for the main GQA cache: its spec is unpadded and the physical
+    layout keeps num_blocks as the first dimension under both layouts."""
+    from vllm.v1.kv_cache_interface import FullAttentionSpec
+
+    spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_DIM,
+        head_size_v=HEAD_DIM,
+        dtype=DTYPE,
+    )
+    # Unpadded -> allocator uses kv_tensor.view(...) rather than as_strided().
+    assert spec.page_size_padded is None
+
+    logical = MiniMaxM3SparseBackend.get_kv_cache_shape(
+        4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
+    )
+    for layout in ("NHD", "HND"):
+        try:
+            set_kv_cache_layout(layout)
+            order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+        finally:
+            set_kv_cache_layout(None)
+        inv_order = [order.index(i) for i in range(len(order))]
+        # Physical first dim is num_blocks (block-first); required by the
+        # padded-strided branch's block-first assumption if it were ever taken.
+        assert inv_order[0] == 0
+        assert logical[order[0]] == logical[0]
+
+
+def _build_decode_inputs(seq_lens_list: tuple[int, ...]):
+    """Shared decode setup: one query token per request at position seq_len-1,
+    a non-identity block table, and topk indices selecting the current block
+    plus older causal blocks."""
+    batch = len(seq_lens_list)
+    pages_per_req = [(s + BLOCK_SIZE - 1) // BLOCK_SIZE for s in seq_lens_list]
+    max_blocks = max(pages_per_req)
+    num_pages = sum(pages_per_req)
+    physical_pages = torch.randperm(num_pages, device="cuda", dtype=torch.int32)
+    block_table = torch.zeros(batch, max_blocks, device="cuda", dtype=torch.int32)
+    base_page = 0
+    for req_id, num_req_pages in enumerate(pages_per_req):
+        block_table[req_id, :num_req_pages] = physical_pages[
+            base_page : base_page + num_req_pages
+        ]
+        base_page += num_req_pages
+
+    seq_lens = torch.tensor(seq_lens_list, device="cuda", dtype=torch.int32)
+    q = torch.randn(batch, NUM_Q_HEADS, HEAD_DIM, device="cuda", dtype=DTYPE)
+
+    topk_idx = torch.full(
+        (NUM_KV_HEADS, batch, TOPK), -1, device="cuda", dtype=torch.int32
+    )
+    for req_id, seq_len in enumerate(seq_lens_list):
+        current_block = (seq_len - 1) // BLOCK_SIZE
+        older_blocks = torch.randperm(current_block, device="cuda", dtype=torch.int32)
+        selected = torch.cat(
+            [
+                torch.tensor([current_block], device="cuda", dtype=torch.int32),
+                older_blocks[: TOPK - 1],
+            ]
+        )
+        topk_idx[:, req_id, : selected.numel()] = selected
+
+    return q, block_table, seq_lens, topk_idx, num_pages
+
+
+@pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
+@pytest.mark.parametrize(
+    "seq_lens_list",
+    [(130, 257), (129, 200, 384)],
+)
+def test_decode_sparse_attention_correctness(
+    kv_layout: str,
+    seq_lens_list: tuple[int, ...],
+):
+    """Decode (split-K) parity under both layouts: this is the only coverage of
+    the decode-site cache feed, and the strided HND case fails if the kernel
+    ignores the cache strides."""
+    torch.manual_seed(0)
+    q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(seq_lens_list)
+    kv_cache = _allocate_main_kv_via_contract(num_pages)
+
+    actual = torch.empty_like(q)
+    minimax_m3_sparse_attn_decode(
+        q,
+        kv_cache,
+        topk_idx,
+        block_table,
+        seq_lens,
+        NUM_KV_HEADS,
+        SM_SCALE,
+        actual,
+    )
+
+    # Reuse the prefill reference: each request is a single query token at
+    # position seq_len-1 (q_len == 1, prefix_len == seq_len-1).
+    q_lens_t = torch.ones(len(seq_lens_list), device="cuda", dtype=torch.int32)
+    prefix_lens = seq_lens - q_lens_t
+    expected = _reference_sparse_attn(
+        q, kv_cache, topk_idx, block_table, q_lens_t, seq_lens, prefix_lens
+    )
+    torch.accelerator.synchronize()
+
+    error = (actual.float() - expected.float()).abs()
+    assert error.mean().item() < 2.5e-4
+    assert error.max().item() < 1.7e-2
+
+
+def test_decode_wrong_layout_breaks_parity():
+    """Negative (AC-3/AC-5): consuming the physical HND buffer as if it were
+    already contiguous-NHD (i.e. skipping the allocator's inverse permute)
+    reorders the K/V content, so the decode output no longer matches the
+    reference computed on the correct logical view. The mislabeled tensor keeps
+    the same shape as the correct view, so the kernel stays in bounds."""
+    torch.manual_seed(0)
+    seq_lens_list = (130, 257)
+    q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(seq_lens_list)
+
+    # Physical HND storage [blocks, 2, heads, block, dim].
+    phys = torch.randn(
+        (num_pages, 2, NUM_KV_HEADS, BLOCK_SIZE, HEAD_DIM), device="cuda", dtype=DTYPE
+    )
+    # Correct logical-NHD view (strided) vs. the same bytes mislabeled as a
+    # contiguous-NHD cache — same shape, different content mapping.
+    correct = phys.permute(0, 1, 3, 2, 4)
+    wrong = phys.reshape(num_pages, 2, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM)
+
+    q_lens_t = torch.ones(len(seq_lens_list), device="cuda", dtype=torch.int32)
+    prefix_lens = seq_lens - q_lens_t
+    expected = _reference_sparse_attn(
+        q, correct, topk_idx, block_table, q_lens_t, seq_lens, prefix_lens
+    )
+
+    actual = torch.empty_like(q)
+    minimax_m3_sparse_attn_decode(
+        q, wrong, topk_idx, block_table, seq_lens, NUM_KV_HEADS, SM_SCALE, actual
+    )
+    torch.accelerator.synchronize()
+    assert (actual.float() - expected.float()).abs().max().item() > 1.7e-2
+
+
+def _make_attn_group(backend, spec):
+    return AttentionGroup(
+        backend=backend,
+        layer_names=["main"],
+        kv_cache_spec=spec,
+        kv_cache_group_id=0,
+    )
+
+
+def test_main_cache_byte_identical_through_production_allocator():
+    """AC-2: drive the real allocator (`_reshape_kv_cache`) for the M3 main
+    `FullAttentionSpec` under HND and assert the backend-visible view has the
+    same shape, stride, and storage offset as the pre-change
+    allocate-HND-then-transpose path; the indexer `MLAAttentionSpec` allocates
+    through the same path to its 3-dim shape."""
+    nb = 4
+    spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_DIM,
+        head_size_v=HEAD_DIM,
+        dtype=DTYPE,
+    )
+    raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
+    group = _make_attn_group(MiniMaxM3SparseBackend, spec)
+    try:
+        set_kv_cache_layout("HND")
+        kv_caches = _reshape_kv_cache([group], {"main": raw}, "auto", [BLOCK_SIZE], {})
+    finally:
+        set_kv_cache_layout(None)
+    view = kv_caches["main"]
+
+    oracle = raw.view(DTYPE).view((nb, 2, NUM_KV_HEADS, BLOCK_SIZE, HEAD_DIM))
+    oracle = oracle.transpose(2, 3)
+    assert tuple(view.shape) == tuple(oracle.shape)
+    assert view.stride() == oracle.stride()
+    assert view.storage_offset() == oracle.storage_offset()
+
+    # Indexer cache allocates through the same path under both layouts.
+    ispec = MLAAttentionSpec(
+        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
+    )
+    for layout in ("NHD", "HND"):
+        iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
+        igroup = AttentionGroup(
+            backend=MiniMaxM3IndexerBackend,
+            layer_names=["idx"],
+            kv_cache_spec=ispec,
+            kv_cache_group_id=0,
+        )
+        try:
+            set_kv_cache_layout(layout)
+            iout = _reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {})
+        finally:
+            set_kv_cache_layout(None)
+        assert tuple(iout["idx"].shape) == (nb, BLOCK_SIZE, HEAD_DIM)
+
+
+def test_indexer_inherited_stride_order_trips_allocator_assert():
+    """AC-4 negative: without the indexer override, the inherited 5-element
+    stride order trips the allocator's `len(stride_order) == len(shape)` assert
+    for the 3-dim indexer shape; the `AssertionError` is NOT swallowed by the
+    allocator's `(AttributeError, NotImplementedError)` fallback."""
+
+    class _BrokenIndexerBackend(MiniMaxM3IndexerBackend):
+        # Simulate inheriting the parent's 5-element stride order.
+        get_kv_cache_stride_order = staticmethod(
+            MiniMaxM3SparseBackend.get_kv_cache_stride_order
+        )
+
+    nb = 4
+    ispec = MLAAttentionSpec(
+        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
+    )
+    iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
+    igroup = AttentionGroup(
+        backend=_BrokenIndexerBackend,
+        layer_names=["idx"],
+        kv_cache_spec=ispec,
+        kv_cache_group_id=0,
+    )
+    try:
+        set_kv_cache_layout("HND")
+        with pytest.raises(AssertionError):
+            _reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {})
+    finally:
+        set_kv_cache_layout(None)
+
+
+def test_padded_main_cache_is_flagged():
+    """AC-2.1 negative: the M3 main cache relies on the allocator's
+    contiguous-view branch (`page_size_padded is None`). A spec that sets
+    `page_size_padded` is explicitly flagged rather than silently wrong-strided."""
+
+    def _require_unpadded_block_first(spec, stride_order):
+        inv_order = [stride_order.index(i) for i in range(len(stride_order))]
+        assert spec.page_size_padded is None, (
+            "main GQA cache must be unpadded to use the contiguous-view "
+            "allocator branch"
+        )
+        assert inv_order[0] == 0, "main GQA cache must remain block-first"
+
+    try:
+        set_kv_cache_layout("HND")
+        stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+    finally:
+        set_kv_cache_layout(None)
+
+    good = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_DIM,
+        head_size_v=HEAD_DIM,
+        dtype=DTYPE,
+    )
+    _require_unpadded_block_first(good, stride_order)  # passes
+
+    padded = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_DIM,
+        head_size_v=HEAD_DIM,
+        dtype=DTYPE,
+        page_size_padded=good.page_size_bytes + 128,
+    )
+    with pytest.raises(AssertionError):
+        _require_unpadded_block_first(padded, stride_order)
+
+
+@pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
+def test_reshape_and_cache_flash_write_persists(kv_layout: str):
+    """AC-5 write path: the `reshape_and_cache_flash` write site now consumes
+    `self.kv_cache.unbind(1)` directly. Writing through those views must persist
+    into the bound storage (read back through an independent logical view) under
+    both layouts — a `.contiguous()` copy of the unbind slice would leave the
+    bound storage unchanged."""
+    torch.manual_seed(0)
+    num_pages = 4
+    kv_cache = _allocate_main_kv_via_contract(num_pages)
+    with torch.no_grad():
+        kv_cache.zero_()
+
+    # Exactly the production write-site code under test.
+    key_cache, value_cache = kv_cache.unbind(1)
+
+    num_tokens = 12
+    slot_mapping = torch.randperm(num_pages * BLOCK_SIZE, device="cuda")[
+        :num_tokens
+    ].to(torch.int64)
+    key = torch.randn(num_tokens, NUM_KV_HEADS, HEAD_DIM, device="cuda", dtype=DTYPE)
+    value = torch.randn(num_tokens, NUM_KV_HEADS, HEAD_DIM, device="cuda", dtype=DTYPE)
+    scale = torch.ones((), device="cuda")
+    ops.reshape_and_cache_flash(
+        key, value, key_cache, value_cache, slot_mapping, "auto", scale, scale
+    )
+    torch.accelerator.synchronize()
+
+    # Read back through the independent logical view; proves the writes landed
+    # in the engine-bound storage, not a detached copy.
+    for t in range(num_tokens):
+        slot = int(slot_mapping[t].item())
+        blk, intra = divmod(slot, BLOCK_SIZE)
+        torch.testing.assert_close(kv_cache[blk, 0, intra], key[t])
+        torch.testing.assert_close(kv_cache[blk, 1, intra], value[t])

tests/kernels/core/test_fused_allreduce_gemma_rms_norm.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the manual AllReduce + GemmaRMSNorm fusion used by MiniMax M3.
+
+``fused_allreduce_gemma_rms_norm`` must match the unfused model path, i.e.
+``GemmaRMSNorm(all_reduce(partial), residual)``, both on the flashinfer fast
+path (TP>1 with flashinfer + NVSwitch) and on the eager fallback (TP==1, or when
+flashinfer is unavailable / the GPU has no NVSwitch).
+"""
+
+import pytest
+import torch
+from torch.multiprocessing import spawn
+
+from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
+from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
+from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
+    fused_allreduce_gemma_rms_norm,
+)
+from vllm.model_executor.layers.layernorm import GemmaRMSNorm
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
+from vllm.utils.torch_utils import set_random_seed
+
+
+@ensure_current_vllm_config()
+def _worker_fused_ar_norm(
+    local_rank,
+    world_size,
+    port,
+    num_tokens,
+    hidden_size,
+    dtype,
+    seed,
+    eps,
+):
+    """Per-rank worker: compare the fused helper vs all_reduce + GemmaRMSNorm."""
+    device = torch.device(f"cuda:{local_rank}")
+    torch.accelerator.set_device_index(device)
+    init_test_distributed_environment(
+        world_size, 1, local_rank, port, local_rank=local_rank
+    )
+
+    # Norm weights are identical across ranks (replicated GemmaRMSNorm).
+    set_random_seed(seed)
+    norm = GemmaRMSNorm(hidden_size, eps=eps).cuda().to(dtype)
+    with torch.no_grad():
+        norm.weight.normal_(mean=0.0, std=0.1)
+
+    # Residual is shared across ranks; the partial o_proj output differs per rank
+    # (each rank holds a partial sum that all_reduce combines).
+    torch.manual_seed(seed + 7)
+    residual = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
+    torch.manual_seed(seed + 1000 + local_rank)
+    partial = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
+
+    # Reference: the unfused model path.
+    reduced = tensor_model_parallel_all_reduce(partial.clone())
+    ref_out, ref_res = norm(reduced, residual.clone())
+
+    # Fused helper (flashinfer fast path when available, else fallback).
+    out, res = fused_allreduce_gemma_rms_norm(partial.clone(), residual.clone(), norm)
+    torch.accelerator.synchronize()
+
+    torch.testing.assert_close(out, ref_out, atol=2e-2, rtol=2e-2)
+    torch.testing.assert_close(res, ref_res, atol=2e-2, rtol=2e-2)
+
+    cleanup_dist_env_and_memory()
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="CUDA required",
+)
+# world_size=1 exercises the TP==1 identity branch on a single GPU; >1 exercises
+# the all_reduce + GemmaRMSNorm equivalence (flashinfer kernel or fallback).
+@pytest.mark.parametrize("world_size", [1, 2, 4])
+@pytest.mark.parametrize("num_tokens", [1, 128, 333])
+@pytest.mark.parametrize("hidden_size", [2048, 4096])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("eps", [1e-6])
+@pytest.mark.parametrize("seed", [42])
+def test_fused_allreduce_gemma_rms_norm(
+    world_size,
+    num_tokens,
+    hidden_size,
+    dtype,
+    eps,
+    seed,
+):
+    num_gpus = current_platform.device_count()
+    if num_gpus < world_size:
+        pytest.skip(f"Need >= {world_size} GPUs, have {num_gpus}")
+    port = str(get_open_port())
+    spawn(
+        _worker_fused_ar_norm,
+        args=(
+            world_size,
+            port,
+            num_tokens,
+            hidden_size,
+            dtype,
+            seed,
+            eps,
+        ),
+        nprocs=world_size,
+        join=True,
+    )

tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py

@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit test for the horizontally-fused MiniMax-M3 attention pre-processing
+kernel:
+
+  fused_minimax_m3_qknorm_rope_kv_insert
+    - q / k / index_q / index_k: Gemma RMSNorm + partial NeoX RoPE (in place)
+    - sparse (insert) mode: scatter k/v into the paged bf16 KV cache and the
+      index key into the index cache by slot_mapping.
+
+Reference: PyTorch GemmaRMSNorm + RotaryEmbedding.forward_static (neox style).
+"""
+
+import pytest
+import torch
+
+import vllm._custom_ops as ops
+
+HEAD_DIM = 128
+ROTARY_DIM = 64
+
+
+def _op_available() -> bool:
+    return hasattr(torch.ops._C, "fused_minimax_m3_qknorm_rope_kv_insert")
+
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or not _op_available(),
+    reason="CUDA not available or fused MiniMax-M3 op not built in",
+)
+
+
+def make_cos_sin_cache(max_pos, rotary_dim, base, dtype, device):
+    inv_freq = 1.0 / (
+        base
+        ** (
+            torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device)
+            / rotary_dim
+        )
+    )
+    t = torch.arange(max_pos, dtype=torch.float32, device=device)
+    freqs = torch.einsum("i,j->ij", t, inv_freq)  # [max_pos, rotary_dim/2]
+    cache = torch.cat((freqs.cos(), freqs.sin()), dim=-1)  # [max_pos, rotary_dim]
+    return cache.to(dtype)
+
+
+def gemma_rmsnorm(x, weight, eps):
+    """x: [..., 128] fp32; weight: [128]. Returns fp32 (one round happens in
+    the caller, matching the kernel's single final cast)."""
+    xf = x.float()
+    var = xf.pow(2).mean(dim=-1, keepdim=True)
+    out = xf * torch.rsqrt(var + eps)
+    return out * (1.0 + weight.float())
+
+
+def apply_rope_neox_partial(x, positions, cos_sin_cache, rotary_dim):
+    """NeoX-style RoPE on the leading rotary_dim dims; rest pass through.
+
+    x: [num_tokens, num_heads, head_dim] fp32
+    cos_sin_cache: [max_pos, rotary_dim] (cos||sin), read as float (matches the
+    kernel, which loads the bf16 cache and converts to fp32).
+    """
+    half = rotary_dim // 2
+    cs = cos_sin_cache[positions].float()  # [num_tokens, rotary_dim]
+    cos = cs[..., :half].unsqueeze(1)  # [nt, 1, half]
+    sin = cs[..., half:].unsqueeze(1)
+
+    rot = x[..., :rotary_dim]
+    x1 = rot[..., :half]
+    x2 = rot[..., half:]
+    o1 = x1 * cos - x2 * sin
+    o2 = x2 * cos + x1 * sin
+    out = x.clone()
+    out[..., :half] = o1
+    out[..., half:rotary_dim] = o2
+    return out
+
+
+def norm_rope_ref(x, weight, positions, cos_sin_cache, eps, dtype):
+    """[nt, nheads, 128] -> Gemma norm + neox partial rope, rounded once."""
+    normed = gemma_rmsnorm(x.float(), weight, eps)
+    roped = apply_rope_neox_partial(normed, positions, cos_sin_cache, ROTARY_DIM)
+    return roped.to(dtype)
+
+
+# ── Test 1: dense mode (norm+rope only, no index, no insert) ─────────────────
+
+
+@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
+@pytest.mark.parametrize("num_heads,num_kv_heads", [(8, 2), (16, 4), (64, 4)])
+def test_dense_norm_rope(num_tokens, num_heads, num_kv_heads):
+    torch.manual_seed(0)
+    device, dtype, eps = "cuda", torch.bfloat16, 1e-6
+    base, max_pos = 5_000_000.0, 4096
+
+    q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
+    positions = torch.randint(
+        0, max_pos, (num_tokens,), dtype=torch.int64, device=device
+    )
+
+    qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
+    qkv = torch.randn(num_tokens, qsz + 2 * kvsz, dtype=dtype, device=device)
+    qkv_orig = qkv.clone()
+
+    ops.fused_minimax_m3_qknorm_rope_kv_insert(
+        qkv, q_w, k_w, cos_sin, positions, num_heads, num_kv_heads, ROTARY_DIM, eps
+    )
+    q_out, k_out, v_out = qkv.split([qsz, kvsz, kvsz], dim=-1)
+
+    q_in, k_in, v_in = qkv_orig.split([qsz, kvsz, kvsz], dim=-1)
+    q_ref = norm_rope_ref(
+        q_in.view(num_tokens, num_heads, HEAD_DIM), q_w, positions, cos_sin, eps, dtype
+    ).view(num_tokens, qsz)
+    k_ref = norm_rope_ref(
+        k_in.view(num_tokens, num_kv_heads, HEAD_DIM),
+        k_w,
+        positions,
+        cos_sin,
+        eps,
+        dtype,
+    ).view(num_tokens, kvsz)
+
+    torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2)
+    # V is untouched.
+    torch.testing.assert_close(v_out, v_in, rtol=0, atol=0)
+
+
+# ── Test 2: sparse mode (full: index branch + cache inserts) ─────────────────
+
+
+@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
+@pytest.mark.parametrize("block_size", [16, 64])
+def test_sparse_full(num_tokens, block_size):
+    torch.manual_seed(1)
+    device, dtype, eps = "cuda", torch.bfloat16, 1e-6
+    base, max_pos = 5_000_000.0, 4096
+    num_heads, num_kv_heads, num_idx_heads = 16, 4, 4
+
+    q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    iq_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    ik_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
+    positions = torch.randint(
+        0, max_pos, (num_tokens,), dtype=torch.int64, device=device
+    )
+
+    qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
+    iqsz, iksz = num_idx_heads * HEAD_DIM, HEAD_DIM
+    # Single fused tensor packing [q | k | v | index_q | index_k].
+    qkv = torch.randn(
+        num_tokens, qsz + 2 * kvsz + iqsz + iksz, dtype=dtype, device=device
+    )
+    qkv_orig = qkv.clone()
+    splits = [qsz, kvsz, kvsz, iqsz, iksz]
+
+    num_blocks = (num_tokens + block_size - 1) // block_size + 1
+    kv_cache = torch.zeros(
+        num_blocks, 2, block_size, num_kv_heads, HEAD_DIM, dtype=dtype, device=device
+    )
+    index_cache = torch.zeros(
+        num_blocks, block_size, HEAD_DIM, dtype=dtype, device=device
+    )
+    slot_mapping = torch.randperm(
+        num_blocks * block_size, dtype=torch.int64, device=device
+    )[:num_tokens]
+
+    # Contiguous gather targets: the kernel writes the normed/roped q and
+    # index_q here (de-interleaved from the packed qkv); k/v/index_k stay in
+    # place inside qkv and are scatter-inserted into the caches.
+    q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device)
+    index_q = torch.empty(num_tokens, iqsz, dtype=dtype, device=device)
+
+    ops.fused_minimax_m3_qknorm_rope_kv_insert(
+        qkv,
+        q_w,
+        k_w,
+        cos_sin,
+        positions,
+        num_heads,
+        num_kv_heads,
+        ROTARY_DIM,
+        eps,
+        iq_w,
+        ik_w,
+        num_idx_heads,
+        slot_mapping,
+        kv_cache,
+        index_cache,
+        block_size,
+        q_out,
+        index_q,
+    )
+
+    # ── norm+rope parity. q/index_q land in their gather buffers; k/index_k are
+    # rewritten in place inside qkv. ──
+    _, k_out, _, _, index_k = qkv.split(splits, dim=-1)
+    q_in, k_in, v_in, iq_orig, ik_orig = qkv_orig.split(splits, dim=-1)
+    q_ref = norm_rope_ref(
+        q_in.view(num_tokens, num_heads, HEAD_DIM), q_w, positions, cos_sin, eps, dtype
+    ).view(num_tokens, qsz)
+    k_ref = norm_rope_ref(
+        k_in.view(num_tokens, num_kv_heads, HEAD_DIM),
+        k_w,
+        positions,
+        cos_sin,
+        eps,
+        dtype,
+    ).view(num_tokens, kvsz)
+    iq_ref = norm_rope_ref(
+        iq_orig.view(num_tokens, num_idx_heads, HEAD_DIM),
+        iq_w,
+        positions,
+        cos_sin,
+        eps,
+        dtype,
+    ).view(num_tokens, num_idx_heads * HEAD_DIM)
+    ik_ref = norm_rope_ref(
+        ik_orig.view(num_tokens, 1, HEAD_DIM), ik_w, positions, cos_sin, eps, dtype
+    ).view(num_tokens, HEAD_DIM)
+
+    torch.testing.assert_close(q_out, q_ref, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(k_out, k_ref, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(index_q, iq_ref, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(index_k, ik_ref, rtol=1e-2, atol=1e-2)
+
+    # ── Cache inserts. ──
+    # Main cache layout is [num_blocks, 2, block_size, num_kv_heads, head_dim]
+    # (the K/V axis sits *before* block_size); index cache is [nb, bs, head_dim].
+    idx_flat = index_cache.view(num_blocks * block_size, HEAD_DIM)
+    k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
+    v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM)  # v is raw (no norm/rope)
+    for t in range(num_tokens):
+        s = slot_mapping[t].item()
+        b, pos = s // block_size, s % block_size
+        torch.testing.assert_close(
+            kv_cache[b, 0, pos], k_ref_h[t], rtol=1e-2, atol=1e-2
+        )
+        torch.testing.assert_close(kv_cache[b, 1, pos], v_ref_h[t], rtol=0, atol=0)
+        torch.testing.assert_close(idx_flat[s], ik_ref[t], rtol=1e-2, atol=1e-2)

tests/kernels/test_minimax_m3_amd_ops.py

@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Reference-vs-optimized unit tests for the MiniMax-M3 AMD/ROCm fused kernels.
+
+Each optimized kernel added for the ROCm port has a slow PyTorch reference; the
+tests assert the two agree within tolerance:
+
+  * Gemma RMSNorm (plain + fused-add-residual)  -> fp32 PyTorch normalize
+  * SwiGLU-OAI (split layout)                    -> fp32 PyTorch elementwise
+  * Fused MXFP8 activation quant (Triton)        -> _mxfp8_e4m3_quantize_torch
+  * Native MXFP8 linear (dot_scaled)             -> dequant-to-bf16 @ matmul
+  * Native MXFP8 MoE (dot_scaled grouped GEMM)   -> dequant-to-bf16 MoE math
+
+The native MXFP8 GEMMs also guard the ``dot_scaled`` rhs-scale orientation: the
+scale is loaded ``[N, K//32]`` and passed WITHOUT transpose; a stray ``.T``
+makes the shape ``[K//32, N]`` and Triton raises before producing output, so any
+regression there fails these tests loudly.
+
+Hardware scope: the whole module is ROCm-only (these are the AMD path; NVIDIA
+uses the FlashInfer kernels). The norm/activation/quant kernels run on any ROCm
+arch; the native MXFP8 ``dot_scaled`` linear/MoE tests are additionally gated to
+CDNA4 gfx95x (``@requires_gfx950``) since gfx942 uses the BF16 emulation path.
+
+Run:  pytest tests/kernels/test_minimax_m3_amd_ops.py -v
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_rocm():
+    pytest.skip(
+        "MiniMax-M3 AMD fused ops require ROCm.", allow_module_level=True
+    )
+if not torch.cuda.is_available():
+    pytest.skip("Requires a GPU.", allow_module_level=True)
+
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (  # noqa: E402
+    _mxfp8_e4m3_quantize_torch,
+    _mxfp8_e4m3_quantize_triton,
+    dequant_mxfp8_to_bf16,
+)
+from vllm.models.minimax_m3.amd.ops import (  # noqa: E402
+    gemma_fused_add_rmsnorm,
+    gemma_rmsnorm,
+    swiglu_oai_split,
+)
+from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import _num_warps  # noqa: E402
+
+DEVICE = "cuda"
+EPS = 1e-6
+
+
+def _gcn_arch() -> str:
+    try:
+        return torch.cuda.get_device_properties(0).gcnArchName
+    except Exception:  # pragma: no cover - no device / non-AMD
+        return ""
+
+
+# The pure-Triton norm/activation/quant kernels run on any ROCm arch (CDNA3
+# gfx942 and CDNA4 gfx950). The native MXFP8 ``dot_scaled`` GEMMs (linear + MoE)
+# use CDNA4 hardware microscaling and are gated to gfx95x in the source
+# (``RocmDotScaledMxfp8LinearKernel.is_supported``; the MoE oracle routes gfx942
+# to the BF16 emulation path instead) — so those tests are gfx950-only.
+requires_gfx950 = pytest.mark.skipif(
+    "gfx95" not in _gcn_arch(),
+    reason="native MXFP8 dot_scaled is a CDNA4 (gfx95x) feature; "
+    "gfx942 uses the BF16 emulation path instead.",
+)
+
+
+def _relerr(a: torch.Tensor, b: torch.Tensor) -> float:
+    a = a.float()
+    b = b.float()
+    return ((a - b).norm() / (b.norm() + 1e-8)).item()
+
+
+# --------------------------------------------------------------------------- #
+# Gemma RMSNorm
+# --------------------------------------------------------------------------- #
+def _ref_gemma_rmsnorm(x, w, eps, residual=None):
+    orig_dtype = x.dtype
+    xf = x.float()
+    res_out = None
+    if residual is not None:
+        xf = xf + residual.float()
+        res_out = xf.to(orig_dtype)
+    xf = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps)
+    xf = xf * (1.0 + w.float())
+    out = xf.to(orig_dtype)
+    return out if residual is None else (out, res_out)
+
+
+@pytest.mark.parametrize("shape", [(1, 4096), (37, 6144), (128, 2048)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("seed", [0, 1234])
+@torch.inference_mode()
+def test_gemma_rmsnorm(shape, dtype, seed):
+    torch.manual_seed(seed)
+    x = torch.randn(*shape, device=DEVICE, dtype=dtype)
+    w = (torch.randn(shape[-1], device=DEVICE, dtype=dtype) * 0.1)
+    got = gemma_rmsnorm(x, w, EPS)
+    ref = _ref_gemma_rmsnorm(x, w, EPS)
+    assert got.shape == x.shape
+    assert _relerr(got, ref) < 5e-3
+
+
+@pytest.mark.parametrize("shape", [(1, 6144), (64, 4096)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@torch.inference_mode()
+def test_gemma_fused_add_rmsnorm(shape, dtype):
+    torch.manual_seed(0)
+    x = torch.randn(*shape, device=DEVICE, dtype=dtype)
+    res = torch.randn(*shape, device=DEVICE, dtype=dtype)
+    w = torch.randn(shape[-1], device=DEVICE, dtype=dtype) * 0.1
+    got_out, got_res = gemma_fused_add_rmsnorm(x, res, w, EPS)
+    ref_out, ref_res = _ref_gemma_rmsnorm(x, w, EPS, residual=res)
+    assert _relerr(got_out, ref_out) < 5e-3
+    # residual_out is the pre-norm sum (x + res): bit-for-bit identical cast.
+    assert torch.equal(got_res, ref_res)
+
+
+@torch.inference_mode()
+def test_gemma_rmsnorm_per_head_strided():
+    """q_norm/k_norm normalize a non-contiguous ``qkv.split`` slice over head_dim."""
+    torch.manual_seed(0)
+    T, H, D, kv = 7, 48, 128, 8
+    total = (H + 2 * kv) * D
+    qkv = torch.randn(T, total, device=DEVICE, dtype=torch.bfloat16)
+    q = qkv[..., : H * D]  # non-contiguous view (row stride == total)
+    q_by_head = q.view(T, H, D)
+    assert not q_by_head.is_contiguous()
+    w = torch.randn(D, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    got = gemma_rmsnorm(q_by_head, w, EPS)
+    ref = _ref_gemma_rmsnorm(q_by_head, w, EPS)
+    assert got.shape == q_by_head.shape
+    assert _relerr(got, ref) < 5e-3
+
+
+def test_num_warps_monotonic():
+    assert _num_warps(128) <= _num_warps(2048) <= _num_warps(8192)
+
+
+# --------------------------------------------------------------------------- #
+# SwiGLU-OAI (split layout)
+# --------------------------------------------------------------------------- #
+def _ref_swiglu(gate_up, alpha, beta, limit):
+    d = gate_up.shape[-1] // 2
+    gate = gate_up[..., :d].float()
+    up = gate_up[..., d:].float()
+    if limit is not None:
+        gate = gate.clamp(max=limit)
+        up = up.clamp(min=-limit, max=limit)
+    return (gate * torch.sigmoid(alpha * gate) * (up + beta)).to(gate_up.dtype)
+
+
+@pytest.mark.parametrize("m,inter", [(1, 768), (64, 1536), (128, 1024)])
+@pytest.mark.parametrize("limit", [7.0, None])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@torch.inference_mode()
+def test_swiglu_oai_split(m, inter, limit, dtype):
+    torch.manual_seed(0)
+    gate_up = torch.randn(m, 2 * inter, device=DEVICE, dtype=dtype)
+    got = swiglu_oai_split(gate_up, alpha=1.702, beta=1.0, limit=limit)
+    ref = _ref_swiglu(gate_up, 1.702, 1.0, limit)
+    assert got.shape == (m, inter)
+    assert _relerr(got, ref) < 5e-3
+
+
+# --------------------------------------------------------------------------- #
+# Fused MXFP8 activation quant (Triton vs torch reference)
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize("shape", [(64, 4096), (1, 6144), (333, 2048)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@torch.inference_mode()
+def test_mxfp8_quant_triton_matches_torch(shape, dtype):
+    torch.manual_seed(0)
+    x = torch.randn(*shape, device=DEVICE, dtype=dtype)
+    xq_t, s_t = _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout=False)
+    xq_k, s_k = _mxfp8_e4m3_quantize_triton(x)
+    assert s_k.shape == s_t.shape == (shape[0], shape[1] // 32)
+    # E8M0 block exponents share the floor(log2(amax))+127 algorithm; allow at
+    # most a 1-step difference at exact powers of two.
+    assert (s_k.int() - s_t.int()).abs().max().item() <= 1
+    # Dequantized values agree to fp8 granularity.
+    deq_t = dequant_mxfp8_to_bf16(xq_t, s_t)
+    deq_k = dequant_mxfp8_to_bf16(xq_k, s_k)
+    assert _relerr(deq_k, deq_t) < 1e-2
+
+
+# --------------------------------------------------------------------------- #
+# Native MXFP8 linear (dot_scaled) vs dequant-to-bf16 matmul
+# --------------------------------------------------------------------------- #
+@requires_gfx950
+@pytest.mark.parametrize("m,n,k", [(64, 256, 128), (37, 512, 256), (1, 6144, 4096)])
+@torch.inference_mode()
+def test_mxfp8_native_linear(m, n, k):
+    from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
+        _mxfp8_dot_scaled_linear,
+    )
+
+    torch.manual_seed(0)
+    w_bf16 = torch.randn(n, k, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
+    x = torch.randn(m, k, device=DEVICE, dtype=torch.bfloat16) * 0.5
+
+    got = _mxfp8_dot_scaled_linear(x, w_fp8, w_scale)
+    # Reference: consume the SAME quantized weights (isolates activation-quant
+    # noise) -> dequant to bf16, plain matmul.
+    w_deq = dequant_mxfp8_to_bf16(w_fp8, w_scale)
+    ref = torch.nn.functional.linear(x, w_deq).to(x.dtype)
+    assert got.shape == (m, n)
+    # Only the activation is re-quantized inside the kernel -> small MX noise.
+    assert _relerr(got, ref) < 5e-2
+
+
+# --------------------------------------------------------------------------- #
+# Native MXFP8 MoE (dot_scaled grouped GEMM) vs dequant-to-bf16 MoE math
+# --------------------------------------------------------------------------- #
+def _ref_moe(x, w13, w2, topk_weights, topk_ids, alpha, beta, limit):
+    T, H = x.shape
+    inter = w2.shape[-1]
+    top_k = topk_ids.shape[1]
+    out = torch.zeros(T, H, device=x.device, dtype=torch.float32)
+    for t in range(T):
+        for j in range(top_k):
+            e = int(topk_ids[t, j].item())
+            g1 = x[t].float() @ w13[e].float().T  # [2I]
+            gate = g1[:inter]
+            up = g1[inter:]
+            if limit is not None:
+                gate = gate.clamp(max=limit)
+                up = up.clamp(min=-limit, max=limit)
+            act = gate * torch.sigmoid(alpha * gate) * (up + beta)
+            g2 = act @ w2[e].float().T  # [H]
+            out[t] += topk_weights[t, j].float() * g2
+    return out.to(x.dtype)
+
+
+@requires_gfx950
+@pytest.mark.parametrize(
+    "T,H,inter,E,top_k", [(8, 256, 512, 8, 2), (1, 512, 256, 16, 4)]
+)
+@torch.inference_mode()
+def test_mxfp8_native_moe(T, H, inter, E, top_k):
+    from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
+        fused_moe_mxfp8_native,
+    )
+
+    torch.manual_seed(0)
+    alpha, beta, limit = 1.702, 1.0, 7.0
+    w13_bf16 = torch.randn(E, 2 * inter, H, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    w2_bf16 = torch.randn(E, H, inter, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    w13_fp8, w13_scale = _mxfp8_e4m3_quantize_torch(
+        w13_bf16, is_sf_swizzled_layout=False
+    )
+    w2_fp8, w2_scale = _mxfp8_e4m3_quantize_torch(
+        w2_bf16, is_sf_swizzled_layout=False
+    )
+
+    x = torch.randn(T, H, device=DEVICE, dtype=torch.bfloat16) * 0.5
+    logits = torch.randn(T, E, device=DEVICE, dtype=torch.float32)
+    topk_weights, topk_ids = logits.softmax(dim=-1).topk(top_k, dim=-1)
+    topk_weights = topk_weights.to(torch.float32)
+    topk_ids = topk_ids.to(torch.int32)
+
+    got = fused_moe_mxfp8_native(
+        x, w13_fp8, w13_scale, w2_fp8, w2_scale, topk_weights, topk_ids,
+        alpha=alpha, beta=beta, limit=limit,
+        global_num_experts=E, expert_map=None,
+    )
+    # Reference consumes the dequantized weights (same bits the kernel reads).
+    w13_deq = dequant_mxfp8_to_bf16(w13_fp8, w13_scale)
+    w2_deq = dequant_mxfp8_to_bf16(w2_fp8, w2_scale)
+    ref = _ref_moe(x, w13_deq, w2_deq, topk_weights, topk_ids, alpha, beta, limit)
+    assert got.shape == (T, H)
+    assert _relerr(got, ref) < 5e-2

tests/reasoning/test_minimax_m3_reasoning_parser.py

@@ -0,0 +1,278 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import string
+from collections.abc import Sequence
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.reasoning import ReasoningParserManager
+from vllm.reasoning.minimax_m3_reasoning_parser import MiniMaxM3ReasoningParser
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+class MiniMaxM3Tokenizer:
+    """Small tokenizer with MiniMax M3 reasoning tags as special tokens."""
+
+    special_tokens = ("<mm:think>", "</mm:think>")
+
+    def __init__(self):
+        self._token_to_id: dict[str, int] = {}
+        self._id_to_token: dict[int, str] = {}
+        for token in self.special_tokens:
+            self._add_token(token)
+        for char in string.printable:
+            self._add_token(char)
+
+    def _add_token(self, token: str) -> int:
+        token_id = self._token_to_id.get(token)
+        if token_id is None:
+            token_id = len(self._token_to_id) + 1
+            self._token_to_id[token] = token_id
+            self._id_to_token[token_id] = token
+        return token_id
+
+    def get_vocab(self) -> dict[str, int]:
+        return dict(self._token_to_id)
+
+    def encode(
+        self,
+        text: str,
+        truncation: bool | None = None,
+        max_length: int | None = None,
+        add_special_tokens: bool = True,
+    ) -> list[int]:
+        return [self._add_token(token) for token in self.tokenize(text)]
+
+    def decode(
+        self, ids: Sequence[int] | int, skip_special_tokens: bool = False
+    ) -> str:
+        if isinstance(ids, int):
+            ids = [ids]
+        return "".join(self._id_to_token[token_id] for token_id in ids)
+
+    def tokenize(self, text: str) -> list[str]:
+        tokens: list[str] = []
+        pos = 0
+        while pos < len(text):
+            for special_token in self.special_tokens:
+                if text.startswith(special_token, pos):
+                    tokens.append(special_token)
+                    pos += len(special_token)
+                    break
+            else:
+                tokens.append(text[pos])
+                pos += 1
+        return tokens
+
+    def convert_ids_to_tokens(
+        self,
+        ids: Sequence[int],
+        skip_special_tokens: bool = False,
+    ) -> list[str]:
+        return [self._id_to_token[token_id] for token_id in ids]
+
+    def convert_tokens_to_ids(self, tokens: str | list[str]) -> int | list[int]:
+        if isinstance(tokens, str):
+            return self._add_token(tokens)
+        return [self._add_token(token) for token in tokens]
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        return "".join(tokens)
+
+
+def make_parser(
+    chat_template_kwargs: dict[str, str] | None = None,
+) -> tuple[MiniMaxM3ReasoningParser, MiniMaxM3Tokenizer]:
+    tokenizer = MiniMaxM3Tokenizer()
+    return (
+        MiniMaxM3ReasoningParser(
+            tokenizer, chat_template_kwargs=chat_template_kwargs
+        ),
+        tokenizer,
+    )
+
+
+def run_streaming(
+    parser: MiniMaxM3ReasoningParser,
+    tokenizer: MiniMaxM3Tokenizer,
+    chunks: list[str],
+) -> tuple[str | None, str | None, list[bool]]:
+    previous_text = ""
+    previous_token_ids: list[int] = []
+    reasoning_parts: list[str] = []
+    content_parts: list[str] = []
+    reasoning_end_states: list[bool] = []
+
+    for chunk in chunks:
+        delta_token_ids = tokenizer.encode(chunk, add_special_tokens=False)
+        current_text = previous_text + chunk
+        current_token_ids = previous_token_ids + delta_token_ids
+        delta = parser.extract_reasoning_streaming(
+            previous_text=previous_text,
+            current_text=current_text,
+            delta_text=chunk,
+            previous_token_ids=previous_token_ids,
+            current_token_ids=current_token_ids,
+            delta_token_ids=delta_token_ids,
+        )
+        reasoning_end_states.append(
+            parser.is_reasoning_end_streaming(current_token_ids, delta_token_ids)
+        )
+
+        if delta is not None:
+            if delta.reasoning is not None:
+                reasoning_parts.append(delta.reasoning)
+            if delta.content is not None:
+                content_parts.append(delta.content)
+
+        previous_text = current_text
+        previous_token_ids = current_token_ids
+
+    return (
+        "".join(reasoning_parts) or None,
+        "".join(content_parts) or None,
+        reasoning_end_states,
+    )
+
+
+def test_parser_registration():
+    parser_cls = ReasoningParserManager.get_reasoning_parser("minimax_m3")
+
+    assert parser_cls is MiniMaxM3ReasoningParser
+
+
+def test_nonstreaming_extracts_explicit_reasoning_block():
+    parser, _ = make_parser()
+    request = ChatCompletionRequest(messages=[], model="test-model")
+
+    reasoning, content = parser.extract_reasoning(
+        "<mm:think>plan</mm:think>answer", request
+    )
+
+    assert reasoning == "plan"
+    assert content == "answer"
+
+
+def test_nonstreaming_without_start_tag_is_content():
+    parser, _ = make_parser()
+    request = ChatCompletionRequest(messages=[], model="test-model")
+
+    reasoning, content = parser.extract_reasoning("plain answer", request)
+
+    assert reasoning is None
+    assert content == "plain answer"
+
+
+def test_nonstreaming_enabled_mode_starts_in_reasoning():
+    parser, _ = make_parser(chat_template_kwargs={"thinking_mode": "enabled"})
+    request = ChatCompletionRequest(messages=[], model="test-model")
+
+    reasoning, content = parser.extract_reasoning("plan</mm:think>answer", request)
+
+    assert reasoning == "plan"
+    assert content == "answer"
+
+
+def test_nonstreaming_open_reasoning_block():
+    parser, _ = make_parser()
+    request = ChatCompletionRequest(messages=[], model="test-model")
+
+    reasoning, content = parser.extract_reasoning("<mm:think>still thinking", request)
+
+    assert reasoning == "still thinking"
+    assert content is None
+
+
+def test_streaming_reasoning_tags_are_not_returned():
+    parser, tokenizer = make_parser()
+
+    reasoning, content, end_states = run_streaming(
+        parser,
+        tokenizer,
+        ["<mm:think>", "plan", "</mm:think>", "answer"],
+    )
+
+    assert reasoning == "plan"
+    assert content == "answer"
+    assert end_states == [False, False, True, True]
+
+
+def test_streaming_boundary_can_emit_reasoning_and_content():
+    parser, tokenizer = make_parser()
+
+    reasoning, content, end_states = run_streaming(
+        parser,
+        tokenizer,
+        ["<mm:think>plan</mm:think>answer"],
+    )
+
+    assert reasoning == "plan"
+    assert content == "answer"
+    assert end_states == [True]
+
+
+def test_streaming_enabled_mode_starts_in_reasoning():
+    parser, tokenizer = make_parser(
+        chat_template_kwargs={"thinking_mode": "enabled"}
+    )
+
+    reasoning, content, end_states = run_streaming(
+        parser,
+        tokenizer,
+        ["plan", "</mm:think>", "answer"],
+    )
+
+    assert reasoning == "plan"
+    assert content == "answer"
+    assert end_states == [False, True, True]
+
+
+def test_streaming_plain_content_ends_reasoning_phase():
+    parser, tokenizer = make_parser()
+
+    reasoning, content, end_states = run_streaming(
+        parser,
+        tokenizer,
+        ["plain ", "answer"],
+    )
+
+    assert reasoning is None
+    assert content == "plain answer"
+    assert end_states == [True, True]
+
+
+def test_token_id_helpers():
+    parser, tokenizer = make_parser()
+    output_ids = tokenizer.encode(
+        "<mm:think>abc</mm:think>def", add_special_tokens=False
+    )
+    open_reasoning_ids = tokenizer.encode("<mm:think>abc", add_special_tokens=False)
+    content_ids = tokenizer.encode("plain", add_special_tokens=False)
+
+    assert parser.is_reasoning_end(output_ids)
+    assert not parser.is_reasoning_end(open_reasoning_ids)
+    assert not parser.is_reasoning_end(content_ids)
+    assert tokenizer.decode(parser.extract_content_ids(output_ids)) == "def"
+    assert parser.extract_content_ids(open_reasoning_ids) == []
+    assert parser.extract_content_ids(content_ids) == content_ids
+    assert parser.count_reasoning_tokens(output_ids) == len(tokenizer.encode("abc"))
+
+
+def test_token_id_helpers_enabled_mode():
+    parser, tokenizer = make_parser(
+        chat_template_kwargs={"thinking_mode": "enabled"}
+    )
+    output_ids = tokenizer.encode("abc</mm:think>def", add_special_tokens=False)
+    open_reasoning_ids = tokenizer.encode("abc", add_special_tokens=False)
+
+    assert parser.is_reasoning_end(output_ids)
+    assert not parser.is_reasoning_end(open_reasoning_ids)
+    assert tokenizer.decode(parser.extract_content_ids(output_ids)) == "def"
+    assert parser.extract_content_ids(open_reasoning_ids) == []
+    assert parser.count_reasoning_tokens(output_ids) == len(tokenizer.encode("abc"))
+    assert parser.count_reasoning_tokens(open_reasoning_ids) == len(
+        tokenizer.encode("abc")
+    )

tests/tool_parsers/test_minimax_m3_tool_parser.py

@@ -0,0 +1,261 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+from typing import Any
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionToolsParam,
+    FunctionDefinition,
+)
+from vllm.entrypoints.openai.engine.protocol import DeltaMessage
+from vllm.tool_parsers import ToolParserManager
+from vllm.tool_parsers.minimax_m3_tool_parser import MinimaxM3ToolParser
+
+pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
+
+NS = "]<]minimax[>["
+EOS_ID = 99
+
+
+class FakeTokenizer:
+    """Minimal fake tokenizer for unit tests."""
+
+    def __init__(self):
+        self.model_tokenizer = True
+        self.vocab: dict[str, int] = {}
+
+    def get_vocab(self) -> dict[str, int]:
+        return self.vocab
+
+
+def sample_tools() -> list[ChatCompletionToolsParam]:
+    return [
+        ChatCompletionToolsParam(
+            function=FunctionDefinition(
+                name="create_order",
+                parameters={
+                    "type": "object",
+                    "properties": {
+                        "user_id": {"type": "integer"},
+                        "urgent": {"type": "boolean"},
+                        "note": {"type": "string"},
+                        "shipping": {
+                            "type": "object",
+                            "properties": {
+                                "city": {"type": "string"},
+                                "zip": {"type": "integer"},
+                            },
+                        },
+                        "items": {
+                            "type": "array",
+                            "items": {
+                                "type": "object",
+                                "properties": {
+                                    "sku": {"type": "string"},
+                                    "qty": {"type": "integer"},
+                                },
+                            },
+                        },
+                        "metadata": {
+                            "type": "object",
+                            "additionalProperties": {"type": "string"},
+                        },
+                        "duplicate_demo": {"type": "object"},
+                    },
+                },
+            ),
+        )
+    ]
+
+
+@pytest.fixture
+def parser() -> MinimaxM3ToolParser:
+    return MinimaxM3ToolParser(FakeTokenizer(), tools=sample_tools())
+
+
+def build_order_call() -> str:
+    return (
+        f"{NS}<tool_call>\n"
+        f'{NS}<invoke name="create_order">'
+        f"{NS}<user_id>42{NS}</user_id>"
+        f"{NS}<urgent>true{NS}</urgent>"
+        f"{NS}<note>Please leave at front desk.{NS}</note>"
+        f"{NS}<shipping>"
+        f"{NS}<city>Singapore{NS}</city>"
+        f"{NS}<zip>018956{NS}</zip>"
+        f"{NS}</shipping>"
+        f"{NS}<items>"
+        f"{NS}<item>{NS}<sku>book-001{NS}</sku>{NS}<qty>2{NS}</qty>{NS}</item>"
+        f"{NS}<item>{NS}<sku>pen-007{NS}</sku>{NS}<qty>5{NS}</qty>{NS}</item>"
+        f"{NS}</items>"
+        f"{NS}<metadata>"
+        f"{NS}<source>mobile{NS}</source>"
+        f"{NS}<campaign>may-launch{NS}</campaign>"
+        f"{NS}</metadata>"
+        f"{NS}<duplicate_demo>"
+        f"{NS}<tag>a{NS}</tag>"
+        f"{NS}<tag>b{NS}</tag>"
+        f"{NS}</duplicate_demo>"
+        f"{NS}</invoke>\n"
+        f"{NS}</tool_call>"
+    )
+
+
+def build_order_invocation(user_id: int) -> str:
+    return (
+        f'{NS}<invoke name="create_order">'
+        f"{NS}<user_id>{user_id}{NS}</user_id>"
+        f"{NS}</invoke>"
+    )
+
+
+def build_multiple_order_call() -> str:
+    return (
+        f"{NS}<tool_call>\n"
+        f"{build_order_invocation(1)}\n"
+        f"{build_order_invocation(2)}\n"
+        f"{NS}</tool_call>"
+    )
+
+
+def _feed(
+    parser: MinimaxM3ToolParser, chunks: list[str | tuple[str, list[int]]]
+) -> list[DeltaMessage]:
+    previous = ""
+    results: list[DeltaMessage] = []
+    for chunk in chunks:
+        if isinstance(chunk, tuple):
+            delta, delta_ids = chunk
+        else:
+            delta = chunk
+            delta_ids = []
+
+        current = previous + delta
+        result = parser.extract_tool_calls_streaming(
+            previous_text=previous,
+            current_text=current,
+            delta_text=delta,
+            previous_token_ids=[],
+            current_token_ids=[],
+            delta_token_ids=delta_ids,
+            request=None,
+        )
+        if result is not None:
+            results.append(result)
+        previous = current
+    return results
+
+
+def _collect_content(results: list[DeltaMessage]) -> str:
+    return "".join(result.content for result in results if result.content)
+
+
+def _collect_tool_calls(results: list[DeltaMessage]) -> dict[int, dict[str, Any]]:
+    tool_calls: dict[int, dict[str, Any]] = {}
+    for result in results:
+        for tool_call in result.tool_calls or []:
+            tool_calls.setdefault(
+                tool_call.index,
+                {"id": None, "name": "", "arguments": ""},
+            )
+            if tool_call.id:
+                tool_calls[tool_call.index]["id"] = tool_call.id
+            if tool_call.function:
+                if tool_call.function.name:
+                    tool_calls[tool_call.index]["name"] += tool_call.function.name
+                if tool_call.function.arguments:
+                    tool_calls[tool_call.index]["arguments"] += (
+                        tool_call.function.arguments
+                    )
+    return tool_calls
+
+
+def test_minimax_m3_parser_registered():
+    assert ToolParserManager.get_tool_parser("minimax_m3") is MinimaxM3ToolParser
+
+
+def test_non_streaming_nested_tool_call(parser):
+    result = parser.extract_tool_calls(
+        "I will create it.\n" + build_order_call(),
+        request=None,
+    )
+
+    assert result.tools_called
+    assert result.content == "I will create it.\n"
+    assert len(result.tool_calls) == 1
+    tool_call = result.tool_calls[0]
+    assert tool_call.function.name == "create_order"
+    assert json.loads(tool_call.function.arguments) == {
+        "user_id": 42,
+        "urgent": True,
+        "note": "Please leave at front desk.",
+        "shipping": {"city": "Singapore", "zip": 18956},
+        "items": [
+            {"sku": "book-001", "qty": 2},
+            {"sku": "pen-007", "qty": 5},
+        ],
+        "metadata": {
+            "source": "mobile",
+            "campaign": "may-launch",
+        },
+        "duplicate_demo": {"tag": ["a", "b"]},
+    }
+
+
+def test_non_streaming_without_tool_call_keeps_content(parser):
+    result = parser.extract_tool_calls("plain response", request=None)
+
+    assert not result.tools_called
+    assert result.tool_calls == []
+    assert result.content == "plain response"
+
+
+def test_non_streaming_multiple_tool_calls(parser):
+    result = parser.extract_tool_calls(build_multiple_order_call(), request=None)
+
+    assert result.tools_called
+    assert result.content is None
+    assert [tool_call.function.name for tool_call in result.tool_calls] == [
+        "create_order",
+        "create_order",
+    ]
+    assert [
+        json.loads(tool_call.function.arguments)["user_id"]
+        for tool_call in result.tool_calls
+    ] == [1, 2]
+
+
+def test_streaming_without_tool_call_emits_text(parser):
+    results = _feed(parser, ["plain ", "response"])
+
+    assert _collect_content(results) == "plain response"
+    assert _collect_tool_calls(results) == {}
+
+
+def test_streaming_nested_tool_call(parser):
+    tool_call_text = build_order_call()
+    results = _feed(
+        parser,
+        [
+            "I will create it.\n",
+            tool_call_text[:5],
+            tool_call_text[5:17],
+            tool_call_text[17:120],
+            tool_call_text[120:],
+            ("", [EOS_ID]),
+        ],
+    )
+
+    assert _collect_content(results) == "I will create it.\n"
+    tool_calls = _collect_tool_calls(results)
+    assert len(tool_calls) == 1
+    assert tool_calls[0]["name"] == "create_order"
+    assert tool_calls[0]["id"] is not None
+    assert json.loads(tool_calls[0]["arguments"]) == json.loads(
+        parser.streamed_args_for_tool[0]
+    )
+    assert json.loads(parser.prev_tool_call_arr[0]["arguments"])["items"][1]["qty"] == 5
+    assert results[-1].content == ""

tools/pre_commit/generate_attention_backend_docs.py

@@ -30,6 +30,7 @@ REPO_ROOT = Path(__file__).parent.parent.parent
 RELEVANT_PATTERNS = [
     "vllm/v1/attention/backends/*.py",
     "vllm/v1/attention/backends/**/*.py",
+    "vllm/models/minimax_m3/common/sparse_attention.py",
     "vllm/model_executor/layers/attention/mla_attention.py",
     "vllm/platforms/cuda.py",
     "tools/pre_commit/generate_attention_backend_docs.py",
@@ -1615,6 +1616,24 @@ def generate_mla_section(
     return "\n".join(lines)
 
 
+def generate_minimax_section(backends: list[dict[str, Any]]) -> str:
+    """Generate the MiniMax M3 sparse attention section."""
+    lines = [
+        "## MiniMax M3 Sparse Attention Backends",
+        "",
+        'Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")',
+        "layers. It is wired in directly by the model and is not part of the",
+        "automatic priority lists above. A lightning indexer scores KV blocks, the",
+        "top-k blocks (plus fixed init/local blocks) are selected, and attention",
+        "attends only to those blocks; index keys live in a separate side cache.",
+        "",
+    ]
+    columns = _build_columns(is_mla=False, has_versions=False)
+    lines.extend(_render_table(columns, backends))
+    lines.append("")
+    return "\n".join(lines)
+
+
 # ---------------------------------------------------------------------------
 # Top-level orchestration
 # ---------------------------------------------------------------------------
@@ -1651,9 +1670,16 @@ def generate_docs() -> str:
     if fi_features:
         all_backends = _expand_flashinfer_variants(all_backends, fi_features)
 
-    # Split into MLA and non-MLA
+    # Split into MLA, MiniMax M3 sparse, and standard (MHA/MQA/GQA) backends.
     mla_backends = [b for b in all_backends if b["is_mla"]]
-    non_mla_backends = [b for b in all_backends if not b["is_mla"]]
+    minimax_backends = [
+        b for b in all_backends if not b["is_mla"] and b["name"].startswith("MINIMAX")
+    ]
+    non_mla_backends = [
+        b
+        for b in all_backends
+        if not b["is_mla"] and not b["name"].startswith("MINIMAX")
+    ]
 
     # Generate documentation
     script_path = "tools/pre_commit/generate_attention_backend_docs.py"
@@ -1702,6 +1728,10 @@ def generate_docs() -> str:
     if footnotes:
         doc_lines.append("\n>\n".join(footnotes) + "\n")
 
+    # Add MiniMax M3 sparse section (separate category after standard GQA)
+    if minimax_backends:
+        doc_lines.append(generate_minimax_section(minimax_backends))
+
     # Add MLA section with prefill and decode backends
     doc_lines.append(generate_mla_section(mla_prefill_backends, mla_backends))
 

vllm/_custom_ops.py

@@ -2680,6 +2680,67 @@ def reshape_and_cache_flash(
     )
 
 
+def fused_minimax_m3_qknorm_rope_kv_insert(
+    qkv: torch.Tensor,
+    q_norm_weight: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    positions: torch.Tensor,
+    num_heads: int,
+    num_kv_heads: int,
+    rotary_dim: int,
+    eps: float,
+    index_q_norm_weight: torch.Tensor | None = None,
+    index_k_norm_weight: torch.Tensor | None = None,
+    num_index_heads: int = 0,
+    slot_mapping: torch.Tensor | None = None,
+    kv_cache: torch.Tensor | None = None,
+    index_cache: torch.Tensor | None = None,
+    block_size: int = 0,
+    q_out: torch.Tensor | None = None,
+    index_q_out: torch.Tensor | None = None,
+) -> None:
+    """Fused MiniMax-M3 attention pre-processing (in-place).
+
+    Applies Gemma RMSNorm + partial NeoX RoPE to ``qkv`` in place. ``qkv`` is a
+    single fused tensor:
+
+    - dense layer (``num_index_heads == 0``): ``[q | k | v]``;
+    - sparse layer (``num_index_heads > 0``): ``[q | k | v | index_q |
+      index_k]`` — the index branch is read straight out of ``qkv``.
+
+    When ``kv_cache`` is given (sparse serving), also scatter-inserts the
+    normed/roped k & v into the paged bf16 KV cache and the index key into
+    ``index_cache`` by ``slot_mapping``.
+
+    If ``q_out`` / ``index_q_out`` (contiguous ``[N, nq*128]`` / ``[N,
+    niq*128]``) are given, the normed/roped q / index_q are written there
+    instead of in place — folding the de-interleave into this kernel's store so
+    callers skip a separate ``.contiguous()`` copy before the SM100 sparse
+    attention's flat TMA descriptor.
+    """
+    torch.ops._C.fused_minimax_m3_qknorm_rope_kv_insert(
+        qkv,
+        q_norm_weight,
+        k_norm_weight,
+        cos_sin_cache,
+        positions,
+        num_heads,
+        num_kv_heads,
+        rotary_dim,
+        eps,
+        index_q_norm_weight,
+        index_k_norm_weight,
+        num_index_heads,
+        slot_mapping,
+        kv_cache,
+        index_cache,
+        block_size,
+        q_out,
+        index_q_out,
+    )
+
+
 def concat_and_cache_mla(
     kv_c: torch.Tensor,
     k_pe: torch.Tensor,

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

@@ -132,6 +132,7 @@ if flashinfer_comm is not None:
         quant_out: torch.Tensor | None = None,
         scale_out: torch.Tensor | None = None,
         scale_factor: torch.Tensor | None = None,
+        weight_bias: float = 0.0,
     ) -> None:
         num_tokens, hidden_size = allreduce_in.shape
         element_size = allreduce_in.element_size()
@@ -209,6 +210,7 @@ if flashinfer_comm is not None:
             use_oneshot=use_oneshot,
             fp32_acc=fp32_acc,
             trigger_completion_at_end=num_tokens > PDL_ADVANCE_LAUNCH_TOKENS,
+            weight_bias=weight_bias,
         )
 
     def call_trtllm_fused_allreduce_norm_fake(
@@ -225,6 +227,7 @@ if flashinfer_comm is not None:
         quant_out: torch.Tensor | None = None,
         scale_out: torch.Tensor | None = None,
         scale_factor: torch.Tensor | None = None,
+        weight_bias: float = 0.0,
     ) -> None:
         pass
 

vllm/config/vllm.py

@@ -1051,20 +1051,26 @@ class VllmConfig:
             )
             self.compilation_config.mode = CompilationMode.NONE
 
-        # DeepSeek V4's model classes don't carry @support_torch_compile —
+        # For model classes don't carry @support_torch_compile —
         # the breakable cudagraph is the supported PIECEWISE path. Auto-enable
         # it unless the user has explicitly opted out via the env var.
         if (
             self.model_config is not None
             and "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ
             and any(
-                a in ("DeepseekV4ForCausalLM", "DeepSeekV4MTPModel")
+                a
+                in (
+                    "DeepseekV4ForCausalLM",
+                    "DeepSeekV4MTPModel",
+                    "MiniMaxM3SparseForCausalLM",
+                    "MiniMaxM3SparseForConditionalGeneration",
+                )
                 for a in self.model_config.architectures
             )
         ):
             os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
             logger.info_once(
-                "Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1 for DeepSeek V4. "
+                "Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1. "
                 "Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out."
             )
 

vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py

@@ -922,12 +922,9 @@ class NixlConnectorWorker:
                         f"{self.transfer_topo.is_kv_layout_blocks_first}"
                     )
 
-                if not self.use_mla:
-                    # Different kv cache shape is not supported by HeteroTP.
-                    # This must also hold true for Mamba-like models.
-                    assert tensor_size_bytes == curr_tensor_size_bytes, (
-                        "All kv cache tensors must have the same size"
-                    )
+                # Allow heterogeneous per-layer KV tensor sizes (non-MLA), e.g.
+                # MiniMax-M3 full-attn + MLA indexer; per-layer sizes live in
+                # block_len_per_layer. Equal-TP enforced at handshake.
                 # Need to make sure the device ID is non-negative for NIXL,
                 # Torch uses -1 to indicate CPU tensors.
                 self.device_id = max(cache.get_device(), 0)
@@ -1524,6 +1521,34 @@ class NixlConnectorWorker:
                         self.block_len_per_layer[i] // block_size_ratio
                         == nixl_agent_meta.block_lens[i]
                     ), "KV cache sizes must match between P and D when replicated"
+        elif (
+            len(set(self.block_len_per_layer)) > 1
+            or len(set(nixl_agent_meta.block_lens)) > 1
+        ):
+            # Non-MLA, non-replicated, HETEROGENEOUS per-layer block lengths
+            # (e.g. MiniMax-M3: full-attn K/V layers + smaller MLA lightning-
+            # indexer layers grouped together). Check either side so a P/D pair
+            # with one homogeneous side still validates per-layer. Only equal-TP
+            # is supported: the linear tp_ratio scaling assumes a uniform
+            # block_len, which does not hold across heterogeneous layers.
+            assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), (
+                "Number of KV layers must match between prefill and decode"
+            )
+            if abs(tp_ratio) != 1 or block_size_ratio != 1:
+                raise NotImplementedError(
+                    "Non-MLA heterogeneous KV cache (mixed full-attention + MLA "
+                    "layers) requires equal tensor-parallel and block size "
+                    "between prefill and decode; got tp_ratio="
+                    f"{tp_ratio}, block_size_ratio={block_size_ratio}."
+                )
+            if not self._has_mamba:
+                # Validate each layer independently (like the MLA/replicated
+                # path); the descriptor builders index block_lens[i] per layer.
+                for i in range(len(self.block_len_per_layer)):
+                    assert (
+                        self.block_len_per_layer[i] // block_size_ratio
+                        == nixl_agent_meta.block_lens[i]
+                    ), "Per-layer KV block_len mismatch between P and D"
         else:
             # When MLA is not used, this is a list of the same block length
             for block_len in nixl_agent_meta.block_lens:

vllm/model_executor/kernels/linear/__init__.py

@@ -90,6 +90,9 @@ from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
 from vllm.model_executor.kernels.linear.mxfp8.marlin import (
     MarlinMxfp8LinearKernel,
 )
+from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
+    RocmDotScaledMxfp8LinearKernel,
+)
 from vllm.model_executor.kernels.linear.mxfp8.xpu import (
     XPUMxFp8LinearKernel,
 )
@@ -372,6 +375,9 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
         EmulationMxfp8LinearKernel,
     ],
     PlatformEnum.ROCM: [
+        # Native CDNA4 (gfx950) MX linear; is_supported() gates to gfx95x and
+        # falls through to BF16 emulation (hipBLASLt) elsewhere / on regression.
+        RocmDotScaledMxfp8LinearKernel,
         EmulationMxfp8LinearKernel,
     ],
     PlatformEnum.XPU: [

vllm/model_executor/kernels/linear/mxfp8/flashinfer.py

@@ -56,8 +56,6 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
 
         input_shape = x.shape
         input_2d = x.view(-1, K)
-        M_orig = input_2d.shape[0]
-
         min_dim = 128
 
         assert min_dim <= K, (
@@ -72,11 +70,6 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
             f"out_features is too small for mm_mxfp8."
         )
 
-        M_padded = ((M_orig + min_dim - 1) // min_dim) * min_dim
-        if M_padded != M_orig:
-            pad_rows = M_padded - M_orig
-            input_2d = torch.nn.functional.pad(input_2d, (0, 0, 0, pad_rows))
-
         input_mxfp8, input_scale = mxfp8_e4m3_quantize(
             input_2d, is_sf_swizzled_layout=True
         )
@@ -93,9 +86,6 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
             backend="cutlass",
         )
 
-        if M_padded != M_orig:
-            output = output[:M_orig, :]
-
         if bias is not None:
             output = output + bias
 

vllm/model_executor/kernels/linear/mxfp8/rocm_native.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Native MXFP8 linear GEMM for AMD CDNA4 (gfx950) via Triton ``tl.dot_scaled``.
+
+Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16);
+activations are MXFP8-quantized per token. Uses the CDNA4 hardware microscaling
+matrix cores. Falls back (via the kernel selector) to the BF16
+``EmulationMxfp8LinearKernel`` on archs without native MX or for shapes with
+``K % 128 != 0``.
+"""
+
+import torch
+from torch.nn.parameter import Parameter
+
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    MXFP8_BLOCK_SIZE,
+    MXFP8_SCALE_DTYPE,
+    dequant_mxfp8_to_bf16,
+    mxfp8_e4m3_quantize,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
+
+
+@triton.jit
+def _mxfp8_linear_kernel(
+    x_ptr, xs_ptr, w_ptr, ws_ptr, out_ptr,
+    M, N, K,
+    stride_xm, stride_xk, stride_xsm, stride_xsk,
+    stride_wn, stride_wk, stride_wsn, stride_wsk,
+    stride_om, stride_on,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_k = tl.arange(0, BLOCK_K)
+    offs_sk = tl.arange(0, BLOCK_K // 32)
+    m_mask = offs_m < M
+    n_mask = offs_n < N
+
+    x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk
+    xs_ptrs = xs_ptr + offs_m[:, None] * stride_xsm + offs_sk[None, :] * stride_xsk
+    w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk
+    ws_ptrs = ws_ptr + offs_n[:, None] * stride_wsn + offs_sk[None, :] * stride_wsk
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for _ in range(0, tl.cdiv(K, BLOCK_K)):
+        x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0)
+        w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0)
+        xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=0)
+        ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0)
+        acc += tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3")
+        x_ptrs += BLOCK_K * stride_xk
+        w_ptrs += BLOCK_K * stride_wk
+        xs_ptrs += (BLOCK_K // 32) * stride_xsk
+        ws_ptrs += (BLOCK_K // 32) * stride_wsk
+
+    o_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
+    tl.store(o_ptrs, acc.to(out_ptr.dtype.element_ty),
+             mask=m_mask[:, None] & n_mask[None, :])
+
+
+def _mxfp8_dot_scaled_linear(
+    x: torch.Tensor,        # [M, K] bf16/fp16
+    w: torch.Tensor,        # [N, K] fp8 e4m3
+    w_scale: torch.Tensor,  # [N, K//32] uint8 (E8M0)
+) -> torch.Tensor:
+    M, K = x.shape
+    N = w.shape[0]
+    x_q, x_scale = mxfp8_e4m3_quantize(x)
+    out = torch.empty((M, N), dtype=x.dtype, device=x.device)
+    BLOCK_M, BLOCK_N, BLOCK_K = 64, 128, 128
+    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
+    _mxfp8_linear_kernel[grid](
+        x_q, x_scale, w, w_scale, out,
+        M, N, K,
+        x_q.stride(0), x_q.stride(1), x_scale.stride(0), x_scale.stride(1),
+        w.stride(0), w.stride(1), w_scale.stride(0), w_scale.stride(1),
+        out.stride(0), out.stride(1),
+        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+        num_warps=8,
+    )
+    return out
+
+
+class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
+    """Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_rocm():
+            return False, "not ROCm"
+        # supports_mx() == gfx95x (CDNA4 native microscaling hardware). On other
+        # archs dot_scaled would upcast to BF16, so the kernel selector falls
+        # through to the BF16 emulation (hipBLASLt) path instead.
+        if not current_platform.supports_mx():
+            return False, "native MX requires CDNA4 (gfx95x)"
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = layer.weight.data  # [N, K] fp8
+        N, K = weight.shape
+        scale_k = K // MXFP8_BLOCK_SIZE
+        weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
+        layer.weight = Parameter(weight.contiguous(), requires_grad=False)
+        layer.weight_scale = Parameter(weight_scale, requires_grad=False)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if layer.weight_scale.dtype != MXFP8_SCALE_DTYPE:
+            raise ValueError(
+                f"Expected {MXFP8_SCALE_DTYPE} weight_scale, got "
+                f"{layer.weight_scale.dtype}."
+            )
+        out_shape = (*x.shape[:-1], layer.weight.shape[0])
+        x2d = x.reshape(-1, x.shape[-1])
+        if x2d.shape[-1] % 128 == 0:
+            out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale)
+        else:
+            # dot_scaled tiling needs K % 128 == 0; dequantize fallback otherwise.
+            w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale)
+            out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype)
+        out = out.reshape(out_shape)
+        if bias is not None:
+            out = out + bias
+        return out

vllm/model_executor/layers/activation.py

@@ -158,17 +158,28 @@ class SiluAndMulWithClamp(CustomOp):
     Computes:
         gate = clamp(x[..., :d], max=swiglu_limit)
         up   = clamp(x[..., d:], min=-swiglu_limit, max=swiglu_limit)
-        out  = silu(gate) * up
-    where d = x.shape[-1] // 2.
+        out  = gate * sigmoid(alpha * gate) * (up + beta)
+    where d = x.shape[-1] // 2. The defaults alpha=1.0, beta=0.0 reduce this to
+    ``silu(gate) * up``; SwiGLU-OAI style models pass alpha (sigmoid scale) and
+    beta=1.0 (up bias).
 
     Shapes:
         x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
         return: (num_tokens, d) or (batch_size, seq_len, d)
     """
 
-    def __init__(self, swiglu_limit: float, *, compile_native: bool = True):
+    def __init__(
+        self,
+        swiglu_limit: float,
+        alpha: float = 1.0,
+        beta: float = 0.0,
+        *,
+        compile_native: bool = True,
+    ):
         super().__init__(compile_native=compile_native)
         self.swiglu_limit = float(swiglu_limit)
+        self.alpha = float(alpha)
+        self.beta = float(beta)
         if current_platform.is_rocm() or current_platform.is_xpu():
             self._forward_method = self.forward_native
         elif current_platform.is_cuda_alike():
@@ -180,18 +191,24 @@ class SiluAndMulWithClamp(CustomOp):
         d = x.shape[-1] // 2
         gate = torch.clamp(x[..., :d], max=self.swiglu_limit)
         up = torch.clamp(x[..., d:], min=-self.swiglu_limit, max=self.swiglu_limit)
-        return F.silu(gate) * up
+        return gate * torch.sigmoid(self.alpha * gate) * (up + self.beta)
 
     def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
         d = x.shape[-1] // 2
         output_shape = x.shape[:-1] + (d,)
         out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        self.op(out, x, self.swiglu_limit)
+        self.op(out, x, self.swiglu_limit, self.alpha, self.beta)
         return out
 
     def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
         return self.forward_native(x)
 
+    def extra_repr(self) -> str:
+        return (
+            f"swiglu_limit={self.swiglu_limit!r}, "
+            f"alpha={self.alpha!r}, beta={self.beta!r}"
+        )
+
 
 # --8<-- [start:mul_and_silu]
 @CustomOp.register("mul_and_silu")

vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Manual fusion of tensor-parallel all-reduce with the following GemmaRMSNorm.
+
+Under tensor parallelism a ``RowParallelLinear`` (e.g. attention ``o_proj``)
+produces a per-rank partial sum that is all-reduced, and the result is then fed
+into a ``GemmaRMSNorm`` that adds the residual and normalizes. flashinfer ships a
+kernel that fuses all-reduce + residual-add + RMSNorm into a single launch; this
+helper drives it directly (no torch.compile pass) for models that run eager.
+
+Scope: attention output only, no quantization. When the flashinfer fast path is
+not applicable (TP==1, flashinfer/NVSwitch unavailable, unsupported dtype, or an
+oversize batch) it falls back to ``all_reduce`` + ``GemmaRMSNorm``, which is
+numerically identical to the unfused model path.
+"""
+
+import torch
+
+from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
+from vllm.distributed.parallel_state import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    get_tp_group,
+)
+from vllm.model_executor.layers.layernorm import GemmaRMSNorm
+
+MiB = 1024 * 1024
+
+# flashinfer fused all-reduce + RMSNorm is wired as a registered custom op in
+# allreduce_rms_fusion; both that op and the workspace helpers only exist when
+# flashinfer.comm.allreduce_fusion is importable.
+try:
+    from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
+        flashinfer_trtllm_fused_allreduce_norm,
+    )
+    from vllm.distributed.device_communicators.flashinfer_all_reduce import (
+        flashinfer_comm,
+        get_fi_ar_workspace,
+    )
+
+    _AR_RESIDUAL_RMS_NORM = (
+        flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm
+        if flashinfer_comm is not None
+        else None
+    )
+except ImportError:
+    flashinfer_trtllm_fused_allreduce_norm = None  # type: ignore[assignment]
+    get_fi_ar_workspace = None  # type: ignore[assignment]
+    _AR_RESIDUAL_RMS_NORM = None
+
+
+_FI_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16)
+
+
+def _max_token_num(tp_size: int, hidden_size: int, dtype: torch.dtype) -> int | None:
+    """Workspace token budget for flashinfer fused all-reduce, or None if the
+    current world size / device is unsupported. Mirrors ``FlashInferAllReduce``."""
+    from vllm.config.compilation import PassConfig
+
+    max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(tp_size)
+    if not max_size_mb:
+        return None
+    element_size = torch.tensor([], dtype=dtype).element_size()
+    return int(max_size_mb * MiB) // (hidden_size * element_size)
+
+
+def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool, int]:
+    """Whether the flashinfer fused path applies; returns (ok, max_token_num)."""
+    if (
+        flashinfer_trtllm_fused_allreduce_norm is None
+        or get_fi_ar_workspace is None
+        or _AR_RESIDUAL_RMS_NORM is None
+    ):
+        return False, 0
+    if (
+        not hidden_states.is_cuda
+        or hidden_states.dim() != 2
+        or not hidden_states.is_contiguous()
+        or hidden_states.dtype not in _FI_SUPPORTED_DTYPES
+    ):
+        return False, 0
+
+    num_tokens, hidden_size = hidden_states.shape
+    max_token_num = _max_token_num(tp_size, hidden_size, hidden_states.dtype)
+    if max_token_num is None or num_tokens > max_token_num:
+        return False, 0
+
+    # Lazily create / fetch the (globally cached) workspace; returns None on
+    # GPUs without NVSwitch, in which case we fall back gracefully.
+    workspace = get_fi_ar_workspace(
+        world_size=tp_size,
+        rank=get_tensor_model_parallel_rank(),
+        max_token_num=max_token_num,
+        hidden_dim=hidden_size,
+        dtype=hidden_states.dtype,
+        group=get_tp_group().device_group,
+    )
+    if workspace is None:
+        return False, 0
+    return True, max_token_num
+
+
+def fused_allreduce_gemma_rms_norm(
+    hidden_states: torch.Tensor,
+    residual: torch.Tensor,
+    norm: GemmaRMSNorm,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """All-reduce ``hidden_states`` + add ``residual`` + GemmaRMSNorm, fused.
+
+    ``hidden_states`` is the per-rank *partial* (un-reduced) output of a
+    row-parallel linear; ``norm`` is the GemmaRMSNorm applied right after.
+    Returns ``(normed_output, new_residual)``, equivalent to
+    ``norm(all_reduce(hidden_states), residual)``.
+    """
+    tp_size = get_tensor_model_parallel_world_size()
+    if tp_size == 1:
+        # No all-reduce needed; identical to the unfused path.
+        return norm(hidden_states, residual)
+
+    ok, max_token_num = _can_use_flashinfer(hidden_states, tp_size)
+    if ok:
+        norm_out = torch.empty_like(hidden_states)
+        # With norm_out provided, the kernel writes the new residual
+        # (all_reduce(hidden_states) + residual) into the hidden_states buffer
+        # and the normalized result into norm_out, leaving `residual` untouched.
+        flashinfer_trtllm_fused_allreduce_norm(
+            allreduce_in=hidden_states,
+            residual=residual,
+            rms_gamma=norm.weight,
+            rms_eps=norm.variance_epsilon,
+            world_size=tp_size,
+            weight_bias=1.0,  # GemmaRMSNorm-style
+            launch_with_pdl=True,
+            fp32_acc=True,
+            max_token_num=max_token_num,
+            pattern_code=_AR_RESIDUAL_RMS_NORM,
+            norm_out=norm_out,
+        )
+        return norm_out, hidden_states
+
+    # Fallback: explicit all-reduce + GemmaRMSNorm (matches the unfused model).
+    reduced = tensor_model_parallel_all_reduce(hidden_states)
+    return norm(reduced, residual)

vllm/model_executor/layers/fused_moe/activation.py

@@ -17,7 +17,12 @@ class MoEActivation(Enum):
     GELU = "gelu"
     GELU_TANH = "gelu_tanh"
     RELU2 = "relu2"
+    # SWIGLUOAI expects gate/up *interleaved* in w13 ([gate0, up0, gate1, ...]),
+    # as in gpt-oss checkpoints. SWIGLUOAI_UNINTERLEAVE has identical math but
+    # expects the *packed* layout ([all gates; all ups]), as produced by a
+    # MergedColumnParallelLinear gate_up_proj (e.g. MiniMax-M3).
     SWIGLUOAI = "swigluoai"
+    SWIGLUOAI_UNINTERLEAVE = "swigluoai_uninterleave"
     SWIGLUSTEP = "swiglustep"
 
     # Non-gated activations (no mul with gate) expect input of shape [..., d]
@@ -73,6 +78,7 @@ _CUSTOM_OP_NAMES: dict[MoEActivation, str] = {
     MoEActivation.GELU: "gelu_and_mul",
     MoEActivation.GELU_TANH: "gelu_tanh_and_mul",
     MoEActivation.SWIGLUOAI: "swigluoai_and_mul",
+    MoEActivation.SWIGLUOAI_UNINTERLEAVE: "silu_and_mul_with_clamp",
     MoEActivation.SWIGLUSTEP: "swiglustep_and_mul",
     MoEActivation.RELU2: "relu2",
     MoEActivation.SILU_NO_MUL: "silu_and_mul",
@@ -105,8 +111,17 @@ def apply_moe_activation(
     activation: MoEActivation,
     output: torch.Tensor,
     input: torch.Tensor,
+    *,
+    clamp_limit: float | None = None,
+    alpha: float = 1.0,
+    beta: float = 0.0,
 ) -> torch.Tensor:
-    """Apply MoE activation function."""
+    """Apply MoE activation function.
+
+    ``clamp_limit``/``alpha``/``beta`` (from the quant config) drive the clamped
+    SwiGLU kernels: ``SILU`` + ``clamp_limit`` and ``SWIGLUOAI_UNINTERLEAVE`` both
+    map to ``silu_and_mul_with_clamp``. Other activations ignore them.
+    """
     assert input.dim() == 2, "Input must be 2D"
     assert output.dim() == 2, "Output must be 2D"
     if activation.is_gated:
@@ -122,13 +137,21 @@ def apply_moe_activation(
 
     # Activations with gated multiplication (gate × activation(up))
     if activation == MoEActivation.SILU:
-        torch.ops._C.silu_and_mul(output, input)
+        if clamp_limit is not None:
+            # Fused silu(clamp(gate)) * clamp(up); equivalent to swiglu_limit_func.
+            torch.ops._C.silu_and_mul_with_clamp(output, input, clamp_limit, 1.0, 0.0)
+        else:
+            torch.ops._C.silu_and_mul(output, input)
     elif activation == MoEActivation.GELU:
         torch.ops._C.gelu_and_mul(output, input)
     elif activation == MoEActivation.GELU_TANH:
         torch.ops._C.gelu_tanh_and_mul(output, input)
     elif activation == MoEActivation.SWIGLUOAI:
         torch.ops._C.swigluoai_and_mul(output, input)
+    elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
+        # SwiGLU-OAI on packed w13 (gate = first half, up = second half).
+        assert clamp_limit is not None, "SWIGLUOAI_UNINTERLEAVE requires clamp_limit"
+        torch.ops._C.silu_and_mul_with_clamp(output, input, clamp_limit, alpha, beta)
     elif activation == MoEActivation.SWIGLUSTEP:
         from vllm.model_executor.layers.activation import swiglustep_and_mul_triton
 

vllm/model_executor/layers/fused_moe/config.py

@@ -895,6 +895,9 @@ def fp8_w8a16_moe_quant_config(
     w1_bias: torch.Tensor | None = None,
     w2_bias: torch.Tensor | None = None,
     block_shape: list[int] | None = None,
+    gemm1_alpha: float | None = None,
+    gemm1_beta: float | None = None,
+    gemm1_clamp_limit: float | None = None,
 ) -> FusedMoEQuantConfig:
     """
     Construct a quant config for 16-bit float activations and fp8 weights.
@@ -920,6 +923,9 @@ def fp8_w8a16_moe_quant_config(
             None,
             w2_bias,
         ),
+        gemm1_alpha=gemm1_alpha,
+        gemm1_beta=gemm1_beta,
+        gemm1_clamp_limit=gemm1_clamp_limit,
     )
 
 

vllm/model_executor/layers/fused_moe/deep_gemm_utils.py

@@ -130,6 +130,9 @@ def _fwd_kernel_ep_scatter_2(
     HIDDEN_SIZE_PAD: tl.constexpr,
     SCALE_HIDDEN_SIZE: tl.constexpr,
     SCALE_HIDDEN_SIZE_PAD: tl.constexpr,
+    PACK_UE8M0: tl.constexpr,
+    SCALE_PACKED_SIZE: tl.constexpr,
+    SCALE_PACKED_SIZE_PAD: tl.constexpr,
 ):
     start_token_id = tl.program_id(0)
     grid_num = tl.num_programs(0)
@@ -137,16 +140,47 @@ def _fwd_kernel_ep_scatter_2(
     offset_in = tl.arange(0, HIDDEN_SIZE_PAD)
     mask = offset_in < HIDDEN_SIZE
 
-    offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
-    mask_s = offset_in_s < SCALE_HIDDEN_SIZE
-
     output_tensor_stride0 = output_tensor_stride0.to(tl.int64)
 
+    if PACK_UE8M0:
+        # One int32 per 4 consecutive 32-wide UE8M0 groups, stored MN-major.
+        offs_pk = tl.arange(0, SCALE_PACKED_SIZE_PAD)
+        mask_pk = offs_pk < SCALE_PACKED_SIZE
+    else:
+        offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
+        mask_s = offset_in_s < SCALE_HIDDEN_SIZE
+
     for token_id in range(start_token_id, total_token_num, grid_num):
         to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask)
-        to_copy_s = tl.load(
-            recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s, mask=mask_s
-        )
+
+        if PACK_UE8M0:
+            # Pack 4 UE8M0 bytes into one int32 (byte j = group 4*pk+j).
+            base_s = recv_x_scale + token_id * recv_x_scale_stride0
+            g0, g1 = offs_pk * 4, offs_pk * 4 + 1
+            g2, g3 = offs_pk * 4 + 2, offs_pk * 4 + 3
+            b0 = tl.load(
+                base_s + g0 * recv_x_scale_stride1, mask=g0 < SCALE_HIDDEN_SIZE
+            )
+            b1 = tl.load(
+                base_s + g1 * recv_x_scale_stride1, mask=g1 < SCALE_HIDDEN_SIZE
+            )
+            b2 = tl.load(
+                base_s + g2 * recv_x_scale_stride1, mask=g2 < SCALE_HIDDEN_SIZE
+            )
+            b3 = tl.load(
+                base_s + g3 * recv_x_scale_stride1, mask=g3 < SCALE_HIDDEN_SIZE
+            )
+            packed_s = (
+                b0.to(tl.int32)
+                | (b1.to(tl.int32) << 8)
+                | (b2.to(tl.int32) << 16)
+                | (b3.to(tl.int32) << 24)
+            )
+        else:
+            to_copy_s = tl.load(
+                recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s,
+                mask=mask_s,
+            )
 
         for topk_index in tl.range(0, topk_num, 1, num_stages=4):
             expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index)
@@ -164,11 +198,21 @@ def _fwd_kernel_ep_scatter_2(
                 output_tensor_ptr = (
                     output_tensor + dest_token_index_i64 * output_tensor_stride0
                 )
+                tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
+
                 output_tensor_scale_ptr = (
                     output_tensor_scale + dest_token_index * output_tensor_scale_stride0
                 )
-                tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask)
-                tl.store(output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s)
+                if PACK_UE8M0:
+                    tl.store(
+                        output_tensor_scale_ptr + offs_pk * output_tensor_scale_stride1,
+                        packed_s,
+                        mask=mask_pk,
+                    )
+                else:
+                    tl.store(
+                        output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s
+                    )
 
 
 @torch.no_grad()
@@ -183,9 +227,11 @@ def ep_scatter(
     output_tensor_scale: torch.Tensor,
     m_indices: torch.Tensor,
     output_index: torch.Tensor,
+    block_size: int = 128,
+    pack_ue8m0: bool = False,
 ):
     BLOCK_E = 128  # token num of per expert is aligned to 128
-    BLOCK_D = 128  # block size of quantization
+    BLOCK_D = block_size  # block size of activation-scale quantization
     num_warps = 8
     num_experts = num_recv_tokens_per_expert.shape[0]
     hidden_size = recv_x.shape[1]
@@ -195,6 +241,10 @@ def ep_scatter(
     assert m_indices.shape[0] % BLOCK_E == 0
     assert expert_start_loc.shape[0] == num_experts
 
+    # pack_ue8m0: scatter packs 4 UE8M0 bytes per int32; else copies scales as-is.
+    scale_hidden_size = hidden_size // BLOCK_D
+    scale_packed_size = (scale_hidden_size + 3) // 4 if pack_ue8m0 else 1
+
     _fwd_kernel_ep_scatter_1[(grid,)](
         num_recv_tokens_per_expert,
         expert_start_loc,
@@ -234,8 +284,11 @@ def ep_scatter(
         num_warps=num_warps,
         HIDDEN_SIZE=hidden_size,
         HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size),
-        SCALE_HIDDEN_SIZE=hidden_size // BLOCK_D,
-        SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size // BLOCK_D),
+        SCALE_HIDDEN_SIZE=scale_hidden_size,
+        SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size),
+        PACK_UE8M0=pack_ue8m0,
+        SCALE_PACKED_SIZE=scale_packed_size,
+        SCALE_PACKED_SIZE_PAD=triton.next_power_of_2(scale_packed_size),
     )
     return
 
@@ -352,6 +405,7 @@ def deepgemm_moe_permute(
     expert_map: torch.Tensor | None,
     expert_tokens_meta: mk.ExpertTokensMetadata | None,
     aq_out: torch.Tensor | None = None,
+    block_size: int | None = None,
 ):
     assert aq.ndim == 2
     assert topk_ids.dtype.is_signed, "The kernel uses -1 to represent invalid topk_ids"
@@ -359,6 +413,10 @@ def deepgemm_moe_permute(
     device = aq.device
 
     block_m, block_k = get_mk_alignment_for_contiguous_layout()
+    # The activation-scale group size may differ from the M/K tile alignment
+    # (e.g. MXFP8 uses a 32-element scale group while block_k stays 128).
+    if block_size is not None:
+        block_k = block_size
 
     M_sum = compute_aligned_M(
         M=topk_ids.size(0),
@@ -376,9 +434,21 @@ def deepgemm_moe_permute(
     if aq_out is None:
         aq_out = torch.empty((M_sum, H), device=device, dtype=aq.dtype)
 
-    aq_scale_out = torch.empty(
-        (M_sum, H // block_k), device=device, dtype=torch.float32
-    )
+    # uint8 UE8M0 (MXFP8) -> scatter packs into DeepGEMM's int32 MN-major
+    # TMA-aligned layout; float32 (FP8/FP4) scattered row-major as-is.
+    pack_ue8m0 = aq_scale.dtype == torch.uint8
+    sf_k = H // block_k
+    if pack_ue8m0:
+        packed_sf_k = (sf_k + 3) // 4
+        tma_aligned_mn = round_up(M_sum, 4)
+        aq_scale_out = torch.empty_strided(
+            (M_sum, packed_sf_k),
+            (1, tma_aligned_mn),
+            device=device,
+            dtype=torch.int32,
+        )
+    else:
+        aq_scale_out = torch.empty((M_sum, sf_k), device=device, dtype=torch.float32)
 
     # DeepGEMM uses negative values in m_indices (here expert_ids) to mark
     # completely invalid / padded blocks that should be skipped. We always
@@ -412,6 +482,8 @@ def deepgemm_moe_permute(
         output_tensor_scale=aq_scale_out,
         m_indices=expert_ids,
         output_index=inv_perm,
+        block_size=block_k,
+        pack_ue8m0=pack_ue8m0,
     )
 
     return aq_out, aq_scale_out, expert_ids, inv_perm

vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py

@@ -33,7 +33,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8Dynamic128Sym,
     kFp8Static128BlockSym,
     kMxfp4Static,
+    kMxfp8Dynamic,
+    kMxfp8Static,
 )
+from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import (
     DeepGemmQuantScaleFMT,
     get_mk_alignment_for_contiguous_layout,
@@ -123,12 +126,26 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
 
     def __init__(self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig):
         super().__init__(moe_config=moe_config, quant_config=quant_config)
-        assert quant_config.block_shape == get_mk_alignment_for_contiguous_layout()
-        assert quant_config.quant_dtype == torch.float8_e4m3fn
+        # MXFP8: FP8 e4m3 values + UE8M0 1x32 block scales (Blackwell). Reuses
+        # the same grouped GEMM (aliased to fp8_fp4) with recipe (1, 32).
+        self.mxfp8 = quant_config.block_shape == [1, 32]
+        if self.mxfp8:
+            assert quant_config.quant_dtype == "mxfp8"
+        else:
+            assert quant_config.block_shape == get_mk_alignment_for_contiguous_layout()
+            assert quant_config.quant_dtype == torch.float8_e4m3fn
         assert not quant_config.per_act_token_quant
         assert not quant_config.per_out_ch_quant
 
         self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
+        # Gated-activation params: silu == swigluoai with alpha=1, beta=0.
+        # FP8 (silu) configs leave these None, reproducing plain silu.
+        self.gemm1_alpha = (
+            quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
+        )
+        self.gemm1_beta = (
+            quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0
+        )
 
     @staticmethod
     def activation_format() -> mk.FusedMoEActivationFormat:
@@ -147,14 +164,25 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
         weight_key: QuantKey | None,
         activation_key: QuantKey | None,
     ) -> bool:
-        SUPPORTED_W_A = [
-            (kFp8Static128BlockSym, kFp8Dynamic128Sym),
-        ]
-        return (weight_key, activation_key) in SUPPORTED_W_A
+        if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym):
+            return True
+        # MXFP8 1x32 uses the fp8_fp4 grouped GEMM with recipe (1, 32) — only
+        # available on Blackwell (SM100).
+        if (weight_key, activation_key) == (kMxfp8Static, kMxfp8Dynamic):
+            return current_platform.is_device_capability_family(100)
+        return False
 
     @staticmethod
     def _supports_activation(activation: MoEActivation) -> bool:
-        return activation in [MoEActivation.SILU, MoEActivation.SWIGLUSTEP]
+        # silu/swigluoai go through the fused alpha/beta kernel; swiglustep
+        # uses the unfused activation path. The fused kernel reads packed w13
+        # (gate = first half, up = second half), so it implements the
+        # *uninterleaved* SwiGLU-OAI variant.
+        return activation in [
+            MoEActivation.SILU,
+            MoEActivation.SWIGLUSTEP,
+            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
+        ]
 
     @staticmethod
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -179,7 +207,9 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
         activation: MoEActivation,
     ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
         assert self.block_shape is not None
-        block_m = self.block_shape[0]
+        # Use the contiguous-layout M alignment (matches apply()); block_shape[0]
+        # is the quant block (1 for MXFP8) and would under-size the workspace.
+        block_m = get_mk_alignment_for_contiguous_layout()[0]
         M_sum = compute_aligned_M(
             M, topk, local_num_experts, block_m, expert_tokens_meta
         )
@@ -201,14 +231,24 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
         M_sum, N = input.size()
         activation_out_dim = self.adjust_N_for_activation(N, activation)
 
-        # 1. DeepGemm UE8M0: fused SiLU+mul+clamp+quant+pack
+        # silu and swigluoai are both expressible by the fused gated kernel via
+        # (alpha, beta): silu uses alpha=1, beta=0; swigluoai uses config values.
+        # The fused kernel reads packed w13, hence SWIGLUOAI_UNINTERLEAVE.
+        fused_gated = activation in (
+            MoEActivation.SILU,
+            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
+        )
+
+        # 1. DeepGemm UE8M0: fused gate+mul+clamp+quant+pack
         if scale_fmt == DeepGemmQuantScaleFMT.UE8M0:
-            if activation == MoEActivation.SILU:
+            if fused_gated:
                 return fused_silu_mul_fp8_quant_packed(
                     input=input,
                     output_q=output,
                     group_size=block_k,
                     clamp_limit=self.gemm1_clamp_limit,
+                    alpha=self.gemm1_alpha,
+                    beta=self.gemm1_beta,
                 )
             act_out = torch.empty(
                 (M_sum, activation_out_dim), dtype=input.dtype, device=input.device
@@ -221,14 +261,17 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
             )
             return a2q, a2q_scale
 
-        # 2. Hopper / non‑E8M0: prefer the fused SiLU+mul+quant kernel
-        if activation == MoEActivation.SILU:
+        # 2. Hopper / non‑E8M0: prefer the fused gate+mul+quant kernel
+        if fused_gated:
             use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0
             return silu_mul_per_token_group_quant_fp8_colmajor(
                 input=input,
                 output=output,
                 use_ue8m0=use_ue8m0,
                 clamp_limit=self.gemm1_clamp_limit,
+                group_size=block_k,
+                alpha=self.gemm1_alpha,
+                beta=self.gemm1_beta,
             )
 
         # 3. fallback path for non-SiLU activations in non‑UE8M0 cases.
@@ -292,12 +335,23 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
             expert_map=expert_map,
             expert_tokens_meta=expert_tokens_meta,
             aq_out=a1q_perm,
+            # MXFP8 uses a 32-element activation-scale group (block_shape[1]);
+            # FP8-block keeps the default (128) alignment.
+            block_size=self.block_shape[1] if self.mxfp8 else None,
         )
         assert a1q.size(0) == M_sum
 
+        # MXFP8 (1x32) drives the fp8_fp4-aliased grouped GEMM with recipe
+        # (1, 32); the FP8 block path keeps the default (128) recipe.
+        gemm_kwargs = (
+            {"recipe_a": (1, self.block_shape[1]), "recipe_b": (1, self.block_shape[1])}
+            if self.mxfp8
+            else {}
+        )
+
         mm1_out = _resize_cache(workspace2, (M_sum, N))
         m_grouped_fp8_gemm_nt_contiguous(
-            (a1q, a1q_scale), (w1, self.w1_scale), mm1_out, expert_ids
+            (a1q, a1q_scale), (w1, self.w1_scale), mm1_out, expert_ids, **gemm_kwargs
         )
 
         activation_out_dim = self.adjust_N_for_activation(N, activation)
@@ -310,7 +364,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
 
         mm2_out = _resize_cache(workspace2, (M_sum, K))
         m_grouped_fp8_gemm_nt_contiguous(
-            (a2q, a2q_scale), (w2, self.w2_scale), mm2_out, expert_ids
+            (a2q, a2q_scale), (w2, self.w2_scale), mm2_out, expert_ids, **gemm_kwargs
         )
 
         if apply_router_weight_on_input:

vllm/model_executor/layers/fused_moe/experts/marlin_moe.py

@@ -28,10 +28,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceDelegate,
     TopKWeightAndReduceNoOP,
 )
-from vllm.model_executor.layers.fused_moe.utils import (
-    _resize_cache,
-    swiglu_limit_func,
-)
+from vllm.model_executor.layers.fused_moe.utils import _resize_cache
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     get_marlin_input_dtype,
     marlin_make_workspace_new,
@@ -72,9 +69,7 @@ def _fused_marlin_moe(
     expert_ids: torch.Tensor,
     num_tokens_post_padded: torch.Tensor,
     activation: MoEActivation = MoEActivation.SILU,
-    activation_func: Callable[
-        [MoEActivation, torch.Tensor, torch.Tensor], None
-    ] = apply_moe_activation,
+    activation_func: Callable[..., None] = apply_moe_activation,
     input_global_scale1: torch.Tensor | None = None,
     input_global_scale2: torch.Tensor | None = None,
     global_scale1: torch.Tensor | None = None,
@@ -92,6 +87,8 @@ def _fused_marlin_moe(
     input_dtype: torch.dtype | None = None,
     is_k_full: bool = True,
     clamp_limit: float | None = None,
+    gemm1_alpha: float = 1.0,
+    gemm1_beta: float = 0.0,
 ) -> torch.Tensor:
     assert hidden_states.ndim == 2
     M, K = hidden_states.size()
@@ -159,18 +156,16 @@ def _fused_marlin_moe(
         use_fp32_reduce=True,
         is_zp_float=False,
     )
-    if clamp_limit is not None and activation == MoEActivation.SILU:
-        swiglu_limit_func(
-            intermediate_cache2,
-            intermediate_cache1.view(-1, w13_num_shards * N),
-            clamp_limit,
-        )
-    else:
-        activation_func(
-            activation,
-            intermediate_cache2,
-            intermediate_cache1.view(-1, w13_num_shards * N),
-        )
+    # apply_moe_activation fuses the clamp/gate params: SILU + clamp_limit and
+    # SWIGLUOAI_UNINTERLEAVE both map to the silu_and_mul_with_clamp kernel.
+    activation_func(
+        activation,
+        intermediate_cache2,
+        intermediate_cache1.view(-1, w13_num_shards * N),
+        clamp_limit=clamp_limit,
+        alpha=gemm1_alpha,
+        beta=gemm1_beta,
+    )
 
     if output is None:
         output = intermediate_cache3
@@ -236,9 +231,7 @@ def fused_marlin_moe(
     apply_router_weight_on_input: bool = False,
     global_num_experts: int = -1,
     activation: MoEActivation = MoEActivation.SILU,
-    activation_func: Callable[
-        [MoEActivation, torch.Tensor, torch.Tensor], None
-    ] = apply_moe_activation,
+    activation_func: Callable[..., None] = apply_moe_activation,
     moe_sum: Callable[[torch.Tensor, torch.Tensor], None] | None = None,
     expert_map: torch.Tensor | None = None,
     input_global_scale1: torch.Tensor | None = None,
@@ -258,6 +251,8 @@ def fused_marlin_moe(
     output: torch.Tensor | None = None,
     input_dtype: torch.dtype | None = None,
     clamp_limit: float | None = None,
+    gemm1_alpha: float = 1.0,
+    gemm1_beta: float = 0.0,
 ) -> torch.Tensor:
     """
     This function computes a Mixture of Experts (MoE) layer using two sets of
@@ -371,6 +366,8 @@ def fused_marlin_moe(
         input_dtype=input_dtype,
         is_k_full=is_k_full,
         clamp_limit=clamp_limit,
+        gemm1_alpha=gemm1_alpha,
+        gemm1_beta=gemm1_beta,
     ).view(-1, topk, K)
 
     if output is None:
@@ -413,6 +410,8 @@ def batched_fused_marlin_moe(
     output: torch.Tensor | None = None,
     input_dtype: torch.dtype | None = None,
     clamp_limit: float | None = None,
+    gemm1_alpha: float = 1.0,
+    gemm1_beta: float = 0.0,
 ) -> torch.Tensor:
     """
     This function massages the inputs so the batched hidden_states can be
@@ -542,6 +541,8 @@ def batched_fused_marlin_moe(
         input_dtype=input_dtype,
         is_k_full=is_k_full,
         clamp_limit=clamp_limit,
+        gemm1_alpha=gemm1_alpha,
+        gemm1_beta=gemm1_beta,
     )
 
     output = output.view(B, BATCH_TOKENS_MAX, K)
@@ -577,6 +578,15 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
         self.is_k_full = is_k_full
         self.input_dtype = get_marlin_input_dtype()
         self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
+        # Gated-activation params (used by SWIGLUOAI_UNINTERLEAVE on packed w13).
+        # silu == swigluoai with alpha=1, beta=0; configs that don't set these
+        # (plain silu) fall back to the silu identity.
+        self.gemm1_alpha = (
+            quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
+        )
+        self.gemm1_beta = (
+            quant_config.gemm1_beta if quant_config.gemm1_beta is not None else 0.0
+        )
 
         super().__init__(
             moe_config=moe_config,
@@ -623,6 +633,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
             MoEActivation.GELU,
             MoEActivation.GELU_TANH,
             MoEActivation.SWIGLUOAI,
+            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
             MoEActivation.SWIGLUSTEP,
             MoEActivation.SILU_NO_MUL,
             MoEActivation.GELU_NO_MUL,
@@ -783,6 +794,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
                 is_k_full=self.is_k_full,
                 input_dtype=self.input_dtype,
                 clamp_limit=self.gemm1_clamp_limit,
+                gemm1_alpha=self.gemm1_alpha,
+                gemm1_beta=self.gemm1_beta,
             )
             return
 
@@ -801,6 +814,10 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
             act_enum: MoEActivation,
             act_output: torch.Tensor,
             act_input: torch.Tensor,
+            *,
+            clamp_limit: float | None = None,
+            alpha: float = 1.0,
+            beta: float = 0.0,
         ) -> None:
             # act_input  = intermediate_cache1 (M*topk, 2N for gated)
             # act_output = intermediate_cache2 (M*topk, N)
@@ -830,7 +847,14 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
                     "tlm": token_lora_mapping,
                 }
             )
-            self.activation(act_enum, act_output, act_input)
+            self.activation(
+                act_enum,
+                act_output,
+                act_input,
+                clamp_limit=clamp_limit,
+                alpha=alpha,
+                beta=beta,
+            )
             lora_state["cache2"] = act_output
 
         def moe_sum_with_lora(moe_out: torch.Tensor, out: torch.Tensor) -> None:
@@ -884,6 +908,8 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
             is_k_full=self.is_k_full,
             input_dtype=self.input_dtype,
             clamp_limit=self.gemm1_clamp_limit,
+            gemm1_alpha=self.gemm1_alpha,
+            gemm1_beta=self.gemm1_beta,
         )
 
     def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
@@ -992,4 +1018,6 @@ class BatchedMarlinExperts(MarlinExpertsBase):
             input_dtype=self.input_dtype,
             is_k_full=self.is_k_full,
             clamp_limit=self.gemm1_clamp_limit,
+            gemm1_alpha=self.gemm1_alpha,
+            gemm1_beta=self.gemm1_beta,
         )

vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py

@@ -0,0 +1,158 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MXFP8 (1x32 block, E8M0 scale) MoE experts on Triton.
+
+``Mxfp8TritonExpertsBase`` stashes E8M0 weight scales for checkpoint layout.
+``Mxfp8EmulationTritonExperts`` dequantizes to BF16 and runs ``TritonExperts``
+for devices without a native MXFP8 MoE kernel (e.g. ROCm gfx942 / MI300).
+"""
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    dequant_mxfp8_to_bf16,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kMxfp8Dynamic,
+    kMxfp8Static,
+)
+
+logger = init_logger(__name__)
+
+
+class Mxfp8TritonExpertsBase(TritonExperts):
+    """Shared MXFP8 MoE setup: stash E8M0 scales, clear scales on ``quant_config``."""
+
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+    ):
+        super().__init__(moe_config, quant_config)
+        self.w1_scale_val = self.quant_config.w1_scale
+        self.w2_scale_val = self.quant_config.w2_scale
+        self.quant_config._w1.scale = None
+        self.quant_config._w2.scale = None
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        return (weight_key, activation_key) == (kMxfp8Static, kMxfp8Dynamic)
+
+    @staticmethod
+    def _supports_activation(activation) -> bool:
+        from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+
+        if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
+            return True
+        return TritonExperts._supports_activation(activation)
+
+
+class Mxfp8EmulationTritonExperts(Mxfp8TritonExpertsBase):
+    """Dequantize MXFP8 weights to BF16 on the fly and run ``TritonExperts``."""
+
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+    ):
+        super().__init__(moe_config, quant_config)
+        logger.warning_once(
+            "Using Mxfp8EmulationTritonExperts MoE backend. Weights are "
+            "dequantized to BF16 on the fly; this is slower than a native "
+            "MXFP8 MoE kernel and is intended for devices without one."
+        )
+
+    @property
+    def quant_dtype(self) -> torch.dtype | str | None:
+        # BF16 fallback: do not MXFP8-quantize activations in ``TritonExperts``.
+        return None
+
+    @property
+    def block_shape(self) -> list[int] | None:
+        return None
+
+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        return True
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        return True
+
+    def activation(self, activation, output: torch.Tensor, input: torch.Tensor):
+        """Apply GEMM1 activation with quant-config alpha/beta/clamp."""
+        from vllm.model_executor.layers.fused_moe.activation import (
+            MoEActivation,
+            apply_moe_activation,
+        )
+
+        if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
+            limit = self.quant_config.gemm1_clamp_limit
+            if limit is None:
+                raise ValueError(
+                    "SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit"
+                )
+            alpha = self.quant_config.gemm1_alpha
+            alpha = 1.702 if alpha is None else float(alpha)
+            beta = self.quant_config.gemm1_beta
+            beta = 1.0 if beta is None else float(beta)
+            apply_moe_activation(
+                activation,
+                output,
+                input,
+                clamp_limit=float(limit),
+                alpha=alpha,
+                beta=beta,
+            )
+            return
+        super().activation(activation, output, input)
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ):
+        w1_bf16 = dequant_mxfp8_to_bf16(w1, self.w1_scale_val).to(hidden_states.dtype)
+        w2_bf16 = dequant_mxfp8_to_bf16(w2, self.w2_scale_val).to(hidden_states.dtype)
+
+        super().apply(
+            output=output,
+            hidden_states=hidden_states,
+            w1=w1_bf16,
+            w2=w2_bf16,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            global_num_experts=global_num_experts,
+            expert_map=expert_map,
+            a1q_scale=None,
+            a2_scale=None,
+            workspace13=workspace13,
+            workspace2=workspace2,
+            expert_tokens_meta=expert_tokens_meta,
+            apply_router_weight_on_input=apply_router_weight_on_input,
+        )

vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py

@@ -0,0 +1,264 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Native MXFP8 (1x32 block, E8M0 scale) MoE for AMD CDNA4 (gfx950) via Triton
+``tl.dot_scaled`` (hardware microscaling matmul).
+
+The expert GEMMs consume the FP8 E4M3 weights and their E8M0 block scales
+directly (no dequant-to-BF16), and activations are MXFP8-quantized per token.
+On CDNA4 ``dot_scaled`` maps to the native MX matrix-core ops; on other archs
+Triton upcasts to BF16 (so this stays correct, just not faster) — but the
+oracle only selects this path on gfx950 and routes everything else to the
+BF16 ``Mxfp8EmulationTritonExperts`` fallback.
+
+Structure mirrors vLLM's ``fused_moe_kernel``: tokens are sorted by expert
+(``moe_align_block_size``); each program computes a ``[BLOCK_M, BLOCK_N]`` tile
+for one expert, accumulating over K with ``dot_scaled``. SwiGLU-OAI activation
+and the top-k weighted reduction run in PyTorch between/after the two GEMMs.
+"""
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
+    Mxfp8TritonExpertsBase,
+)
+from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
+    moe_align_block_size,
+)
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    MXFP8_BLOCK_SIZE,
+    mxfp8_e4m3_quantize,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+
+logger = init_logger(__name__)
+
+
+@triton.jit
+def _mxfp8_grouped_gemm_kernel(
+    a_ptr,
+    a_scale_ptr,
+    b_ptr,
+    b_scale_ptr,
+    c_ptr,
+    topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    N,
+    K,
+    num_valid_tokens,
+    top_k,
+    stride_am,
+    stride_ak,
+    stride_asm,
+    stride_ask,
+    stride_be,
+    stride_bn,
+    stride_bk,
+    stride_bse,
+    stride_bsn,
+    stride_bsk,
+    stride_cm,
+    stride_cn,
+    A_DIV: tl.constexpr,
+    MUL_WEIGHT: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+    num_post = tl.load(num_tokens_post_padded_ptr)
+    if pid_m * BLOCK_M >= num_post:
+        return
+
+    offs_tid = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_token = tl.load(sorted_token_ids_ptr + offs_tid).to(tl.int64)
+    token_mask = offs_token < num_valid_tokens
+    off_e = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_k = tl.arange(0, BLOCK_K)
+    offs_sk = tl.arange(0, BLOCK_K // 32)
+    a_row = offs_token // A_DIV
+
+    a_ptrs = a_ptr + a_row[:, None] * stride_am + offs_k[None, :] * stride_ak
+    as_ptrs = a_scale_ptr + a_row[:, None] * stride_asm + offs_sk[None, :] * stride_ask
+    b_ptrs = b_ptr + off_e * stride_be + offs_n[:, None] * stride_bn + \
+        offs_k[None, :] * stride_bk
+    bs_ptrs = b_scale_ptr + off_e * stride_bse + offs_n[:, None] * stride_bsn + \
+        offs_sk[None, :] * stride_bsk
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    n_mask = offs_n < N
+    for _ in range(0, tl.cdiv(K, BLOCK_K)):
+        a = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0)
+        b = tl.load(b_ptrs, mask=n_mask[:, None], other=0.0)
+        asc = tl.load(as_ptrs, mask=token_mask[:, None], other=0)
+        bsc = tl.load(bs_ptrs, mask=n_mask[:, None], other=0)
+        acc += tl.dot_scaled(a, asc, "e4m3", b.T, bsc, "e4m3")
+
+        a_ptrs += BLOCK_K * stride_ak
+        b_ptrs += BLOCK_K * stride_bk
+        as_ptrs += (BLOCK_K // 32) * stride_ask
+        bs_ptrs += (BLOCK_K // 32) * stride_bsk
+
+    if MUL_WEIGHT:
+        w = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0)
+        acc = acc * w[:, None]
+
+    c_ptrs = c_ptr + offs_token[:, None] * stride_cm + offs_n[None, :] * stride_cn
+    tl.store(c_ptrs, acc.to(c_ptr.dtype.element_ty),
+             mask=token_mask[:, None] & n_mask[None, :])
+
+
+def _grouped_gemm_mxfp8(
+    a_q: torch.Tensor,        # [M, K] fp8 e4m3
+    a_scale: torch.Tensor,    # [M, K//32] uint8 (E8M0)
+    w: torch.Tensor,          # [E, N, K] fp8 e4m3
+    w_scale: torch.Tensor,    # [E, N, K//32] uint8 (E8M0)
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    num_valid_tokens: int,
+    top_k: int,
+    block_m: int,
+    out_dtype: torch.dtype,
+    a_div: int,
+    mul_weight_by: torch.Tensor | None = None,
+) -> torch.Tensor:
+    M_routed = num_valid_tokens
+    E, N, K = w.shape
+    assert K % 128 == 0, f"MXFP8 native MoE requires K%128==0, got K={K}"
+    out = torch.zeros((M_routed, N), dtype=out_dtype, device=a_q.device)
+    BLOCK_N = 128
+    BLOCK_K = 128
+    grid = (triton.cdiv(sorted_token_ids.shape[0], block_m), triton.cdiv(N, BLOCK_N))
+    _mxfp8_grouped_gemm_kernel[grid](
+        a_q, a_scale, w, w_scale, out,
+        mul_weight_by if mul_weight_by is not None else a_q,
+        sorted_token_ids, expert_ids, num_tokens_post_padded,
+        N, K, num_valid_tokens, top_k,
+        a_q.stride(0), a_q.stride(1), a_scale.stride(0), a_scale.stride(1),
+        w.stride(0), w.stride(1), w.stride(2),
+        w_scale.stride(0), w_scale.stride(1), w_scale.stride(2),
+        out.stride(0), out.stride(1),
+        A_DIV=a_div,
+        MUL_WEIGHT=mul_weight_by is not None,
+        BLOCK_M=block_m, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+        num_warps=8,
+    )
+    return out
+
+
+def fused_moe_mxfp8_native(
+    hidden_states: torch.Tensor,   # [T, H] bf16
+    w13: torch.Tensor,             # [E, 2I, H] fp8
+    w13_scale: torch.Tensor,       # [E, 2I, H//32] uint8
+    w2: torch.Tensor,              # [E, H, I] fp8
+    w2_scale: torch.Tensor,        # [E, H, I//32] uint8
+    topk_weights: torch.Tensor,    # [T, top_k]
+    topk_ids: torch.Tensor,        # [T, top_k] (global expert ids)
+    *,
+    alpha: float,
+    beta: float,
+    limit: float | None,
+    global_num_experts: int,
+    expert_map: torch.Tensor | None,
+) -> torch.Tensor:
+    T, H = hidden_states.shape
+    top_k = topk_ids.shape[1]
+    M = T * top_k
+
+    block_m = 64
+    sorted_ids, expert_ids, num_post = moe_align_block_size(
+        topk_ids, block_m, global_num_experts, expert_map
+    )
+
+    # GEMM1: x (mxfp8) @ w13^T -> [M, 2I]
+    a_q, a_s = mxfp8_e4m3_quantize(hidden_states)
+    g1 = _grouped_gemm_mxfp8(
+        a_q, a_s, w13, w13_scale, sorted_ids, expert_ids, num_post,
+        M, top_k, block_m, hidden_states.dtype, a_div=top_k,
+    )  # [M, 2I]
+
+    # SwiGLU-OAI (split layout): gate=g1[:, :I], up=g1[:, I:] — fp32 Triton pass.
+    # Not the #22 ``silu_and_mul_with_clamp`` op: it is built on ROCm but rounds
+    # intermediates to bf16 (rel ~3e-3 vs ~1e-6 here), and this feeds the GEMM2
+    # MXFP8 quant. Lazy import: the amd.ops package pulls in the minimax_m3
+    # platform dispatch, only resolvable after the model module finishes loading.
+    from vllm.models.minimax_m3.amd.ops import swiglu_oai_split
+
+    act = swiglu_oai_split(
+        g1, alpha=alpha, beta=beta, limit=limit, out_dtype=hidden_states.dtype
+    )  # [M, I]
+
+    # GEMM2: act (mxfp8) @ w2^T -> [M, H], weighted by topk_weights, then reduce.
+    act_q, act_s = mxfp8_e4m3_quantize(act)
+    g2 = _grouped_gemm_mxfp8(
+        act_q, act_s, w2, w2_scale, sorted_ids, expert_ids, num_post,
+        M, top_k, block_m, torch.float32, a_div=1,
+        mul_weight_by=topk_weights.reshape(-1).to(torch.float32),
+    )  # [M, H] == [T*top_k, H]
+
+    return g2.view(T, top_k, H).sum(dim=1).to(hidden_states.dtype)
+
+
+class Mxfp8NativeTritonExperts(Mxfp8TritonExpertsBase):
+    """Native MXFP8 MoE (CDNA4 ``dot_scaled``) on gfx950."""
+
+    @property
+    def quant_dtype(self) -> torch.dtype | str | None:
+        return self.quant_config.quant_dtype
+
+    @property
+    def block_shape(self) -> list[int] | None:
+        return self.quant_config.block_shape
+
+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        # Activations are MXFP8-quantized inside ``fused_moe_mxfp8_native``.
+        return True
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        return current_platform.is_rocm() and current_platform.supports_mx()
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ):
+        alpha = self.quant_config.gemm1_alpha
+        alpha = 1.702 if alpha is None else float(alpha)
+        beta = self.quant_config.gemm1_beta
+        beta = 1.0 if beta is None else float(beta)
+        limit = self.quant_config.gemm1_clamp_limit
+        limit = None if limit is None else float(limit)
+        out = fused_moe_mxfp8_native(
+            hidden_states, w1, self.w1_scale_val, w2, self.w2_scale_val,
+            topk_weights, topk_ids,
+            alpha=alpha, beta=beta, limit=limit,
+            global_num_experts=global_num_experts, expert_map=expert_map,
+        )
+        output.copy_(out)

vllm/model_executor/layers/fused_moe/layer.py

@@ -120,6 +120,8 @@ class FusedMoE(PluggableLayer):
         scoring_func: str = "softmax",
         routed_scaling_factor: float = 1.0,
         swiglu_limit: float | None = None,
+        swiglu_alpha: float | None = None,
+        swiglu_beta: float | None = None,
         e_score_correction_bias: torch.Tensor | None = None,
         apply_router_weight_on_input: bool = False,
         activation: str = "silu",
@@ -149,6 +151,10 @@ class FusedMoE(PluggableLayer):
         vllm_config = get_current_vllm_config()
         self.vllm_config = vllm_config
         self.swiglu_limit = swiglu_limit
+        # SwiGLU/swigluoai gate params (alpha, beta); read by quant configs
+        # (e.g. ModelOptMxFp8FusedMoE) to plumb into the fused activation kernel.
+        self.swiglu_alpha = swiglu_alpha
+        self.swiglu_beta = swiglu_beta
 
         # FIXME (varun): We should have a better way of inferring the activation
         # datatype. This works for now as the tensor datatype entering the MoE

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -880,9 +880,18 @@ class FusedMoEExpertsModular(FusedMoEExperts):
         return N if not activation.is_gated else N // 2
 
     def activation(
-        self, activation: MoEActivation, output: torch.Tensor, input: torch.Tensor
+        self,
+        activation: MoEActivation,
+        output: torch.Tensor,
+        input: torch.Tensor,
+        *,
+        clamp_limit: float | None = None,
+        alpha: float = 1.0,
+        beta: float = 0.0,
     ) -> None:
-        apply_moe_activation(activation, output, input)
+        apply_moe_activation(
+            activation, output, input, clamp_limit=clamp_limit, alpha=alpha, beta=beta
+        )
 
     @abstractmethod
     def finalize_weight_and_reduce_impl(self) -> TopKWeightAndReduce:

vllm/model_executor/layers/fused_moe/oracle/fp8.py

@@ -53,6 +53,13 @@ class Fp8MoeBackend(Enum):
     BATCHED_VLLM_CUTLASS = "BATCHED_VLLM_CUTLASS"
     XPU = "XPU"
     CPU = "CPU"
+    # Dequantize-to-BF16 emulation for MXFP8 on devices without a native
+    # MXFP8 MoE kernel (e.g. ROCm). Weights pass through unchanged here.
+    EMULATION = "EMULATION"
+    # MXFP8 MoE via a Triton ``dot_scaled`` kernel that lowers to CDNA4
+    # (gfx950) native MX matrix-core ops. Weights stay in MXFP8 (no load-time
+    # format conversion); the FP8 values + E8M0 scales are consumed directly.
+    NATIVE_MXFP8 = "NATIVE_MXFP8"
 
 
 def _get_priority_backends(
@@ -499,6 +506,10 @@ def convert_to_fp8_moe_kernel_format(
             Fp8MoeBackend.VLLM_CUTLASS,
             Fp8MoeBackend.BATCHED_VLLM_CUTLASS,
             Fp8MoeBackend.XPU,
+            # EMULATION dequantizes weights at runtime; NATIVE_MXFP8 consumes
+            # the MXFP8 weights as-is — neither needs a load-time layout change.
+            Fp8MoeBackend.EMULATION,
+            Fp8MoeBackend.NATIVE_MXFP8,
         ]:
             raise ValueError(f"Unsupported FP8 MoE backend: {fp8_backend.value}")
 
@@ -517,6 +528,8 @@ def make_fp8_moe_quant_config(
     per_act_token_quant: bool = False,
     per_out_ch_quant: bool = False,
     swiglu_limit: float | None = None,
+    gemm1_alpha: float | None = None,
+    gemm1_beta: float | None = None,
 ) -> FusedMoEQuantConfig:
     """
     Create FusedMoEQuantConfig for the specified FP8 Backend.
@@ -539,6 +552,9 @@ def make_fp8_moe_quant_config(
             w1_bias=w1_bias,
             w2_bias=w2_bias,
             block_shape=block_shape,
+            gemm1_alpha=gemm1_alpha,
+            gemm1_beta=gemm1_beta,
+            gemm1_clamp_limit=swiglu_limit,
         )
 
     # Flashinfer CUTLASS per-tensor uses single dq scale
@@ -558,10 +574,9 @@ def make_fp8_moe_quant_config(
             g2_alphas=(w2_scale * a2_scale).squeeze(),
             gemm1_clamp_limit=swiglu_limit,
         )
-    # MXFP8 uses "mxfp8" quant_dtype so the prepare step dispatches to
-    # _mxfp8_e4m3_quantize rather than standard FP8 block quantization.
-    # Non-swizzled layout is required since the TRTLLM kernel expects
-    # scales in (num_tokens, hidden_dim // 32) format.
+    # MXFP8 (block [1, 32]) dispatches to the mxfp8 activation quant. Scales are
+    # the non-swizzled (num_tokens, hidden_dim // 32) uint8 UE8M0 layout for all
+    # backends; the DeepGEMM expert permute repacks them for the grouped GEMM.
     if block_shape == [1, 32]:
         return FusedMoEQuantConfig.make(
             "mxfp8",
@@ -573,6 +588,8 @@ def make_fp8_moe_quant_config(
             a2_scale=a2_scale,
             block_shape=block_shape,
             is_scale_swizzled=False,
+            gemm1_alpha=gemm1_alpha,
+            gemm1_beta=gemm1_beta,
             gemm1_clamp_limit=swiglu_limit,
         )
 

vllm/model_executor/layers/fused_moe/oracle/mxfp8.py

@@ -12,22 +12,43 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kMxfp8Dynamic,
     kMxfp8Static,
 )
+from vllm.platforms import current_platform
 
 logger = init_logger(__name__)
 
 _SUPPORTED_BACKENDS = (
     Fp8MoeBackend.FLASHINFER_TRTLLM,
+    Fp8MoeBackend.DEEPGEMM,
     Fp8MoeBackend.MARLIN,
     Fp8MoeBackend.XPU,
 )
 
 _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = {
     "flashinfer_trtllm": Fp8MoeBackend.FLASHINFER_TRTLLM,
+    "deep_gemm": Fp8MoeBackend.DEEPGEMM,
     "marlin": Fp8MoeBackend.MARLIN,
     "xpu": Fp8MoeBackend.XPU,
 }
 
 
+def _mxfp8_backend_to_kernel_cls(
+    backend: Fp8MoeBackend,
+) -> list[type[mk.FusedMoEExperts]]:
+    """Resolve the MXFP8 expert classes for a backend.
+
+    DeepGEMM resolves directly to ``DeepGemmExperts`` (not the
+    ``TritonOrDeepGemmExperts`` wrapper, whose Triton fallback cannot handle the
+    MXFP8 1x32 scheme); all other backends defer to the FP8 resolver.
+    """
+    if backend == Fp8MoeBackend.DEEPGEMM:
+        from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
+            DeepGemmExperts,
+        )
+
+        return [DeepGemmExperts]
+    return backend_to_kernel_cls(backend)
+
+
 def _select_kernel_cls(
     backend: Fp8MoeBackend,
     config: FusedMoEConfig,
@@ -39,7 +60,7 @@ def _select_kernel_cls(
         else mk.FusedMoEActivationFormat.Standard
     )
     last_reason: str | None = None
-    for cls in backend_to_kernel_cls(backend):
+    for cls in _mxfp8_backend_to_kernel_cls(backend):
         supported, reason = cls.is_supported_config(
             cls,
             config,
@@ -55,6 +76,29 @@ def _select_kernel_cls(
     )
 
 
+def _select_rocm_mxfp8_backend() -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:
+    """ROCm fallback when vendor MXFP8 backends are unavailable."""
+
+    if current_platform.supports_mx():
+        from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
+            Mxfp8NativeTritonExperts,
+        )
+
+        logger.info_once(
+            "Using native CDNA4 (gfx950) MXFP8 dot_scaled MoE backend."
+        )
+        return Fp8MoeBackend.NATIVE_MXFP8, Mxfp8NativeTritonExperts
+
+    from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
+        Mxfp8EmulationTritonExperts,
+    )
+
+    logger.info_once(
+        "No native MXFP8 MoE backend available; using BF16 dequant emulation."
+    )
+    return Fp8MoeBackend.EMULATION, Mxfp8EmulationTritonExperts
+
+
 def select_mxfp8_moe_backend(
     config: FusedMoEConfig,
 ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:
@@ -88,4 +132,8 @@ def select_mxfp8_moe_backend(
         logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value)
         return backend, experts_cls
 
-    raise ValueError("No MXFP8 MoE backends available.")
+    # simplify the logic for rocm, refactor later when more backends are supported
+    if current_platform.is_rocm():
+        return _select_rocm_mxfp8_backend()
+
+    raise ValueError("No MXFP8 MoE backends available.")

vllm/model_executor/layers/fused_moe/utils.py

@@ -313,6 +313,8 @@ def moe_kernel_quantize_input(
                 "moe_kernel_quantize_input does not support quant_dtype='mxfp8' MOE "
                 "quantization emulation. Please open an issue."
             )
+        # Non-swizzled (M, K/32) uint8 UE8M0 scales; deepgemm_moe_permute packs
+        # them for DeepGEMM, TRTLLM takes them as-is.
         return _mxfp8_e4m3_quantize(
             A,
             A_scale,

vllm/model_executor/layers/linear.py

@@ -1163,6 +1163,7 @@ class QKVParallelLinear(ColumnParallelLinear):
 
         shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
         shard_size = self._get_shard_size_mapping(loaded_shard_id)
+        assert shard_offset is not None and shard_size is not None
 
         if isinstance(param, BlockQuantScaleParameter):
             weight_block_size = getattr(self, "weight_block_size", None)
@@ -1384,6 +1385,191 @@ class QKVParallelLinear(ColumnParallelLinear):
         param_data.copy_(loaded_weight)
 
 
+class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear):
+    """QKV projection fused with a lightning-indexer's index_q/index_k.
+
+    NOTE: MiniMax-M3-specific. This is tailored to the M3 sparse-attention
+    layers (it assumes the indexer's head count equals the KV head count and
+    shares the main head_dim); it is not a general-purpose linear layer. It
+    lives here only to sit alongside QKVParallelLinear, whose sharding /
+    weight-loading machinery it reuses.
+
+    A single column-parallel GEMM emits, per rank::
+
+        [q | k | v | index_q | index_k]
+
+    ``index_q`` must have the same head count as the KV heads
+    (``total_num_index_heads == total_num_kv_heads``) and ``index_head_size ==
+    head_size``, so it shards exactly like K/V -- including the KV-head
+    *replication* path when ``tp_size > total_num_kv_heads`` (this is what makes
+    a TP size greater than the KV-head count work). ``index_k`` is a single
+    shared head, replicated to every rank.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        head_size: int,
+        total_num_heads: int,
+        total_num_kv_heads: int,
+        total_num_index_heads: int,
+        index_head_size: int,
+        bias: bool = False,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        # index_q rides the KV-head sharding/replication path, so its head count
+        # must match the KV heads.
+        assert total_num_index_heads == total_num_kv_heads, (
+            "MinimaxM3QKVParallelLinearWithIndexer requires "
+            "total_num_index_heads == total_num_kv_heads"
+        )
+        self.hidden_size = hidden_size
+        self.head_size = head_size
+        self.v_head_size = head_size
+        self.total_num_heads = total_num_heads
+        self.total_num_kv_heads = total_num_kv_heads
+        self.total_num_index_heads = total_num_index_heads
+        self.index_head_size = index_head_size
+
+        tp_size = get_tensor_model_parallel_world_size()
+        self.num_heads = divide(self.total_num_heads, tp_size)
+        if tp_size >= self.total_num_kv_heads:
+            self.num_kv_heads = 1
+            self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads)
+        else:
+            self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
+            self.num_kv_head_replicas = 1
+        # index_q shards identically to the KV heads.
+        self.num_index_heads = self.num_kv_heads
+
+        # Global per-group sizes (replicated groups counted x tp_size, matching
+        # the QKVParallelLinear convention). index_k is a single replicated head.
+        q = self.num_heads * self.head_size
+        kv = self.num_kv_heads * self.head_size
+        iq = self.num_index_heads * self.index_head_size
+        ik = self.index_head_size
+        self.output_sizes = [
+            q * tp_size,  # q
+            kv * tp_size,  # k
+            kv * tp_size,  # v
+            iq * tp_size,  # index_q
+            ik * tp_size,  # index_k (replicated)
+        ]
+
+        # Skip QKVParallelLinear.__init__ (3-group layout); build the 5-group
+        # column-parallel weight directly.
+        ColumnParallelLinear.__init__(
+            self,
+            input_size=self.hidden_size,
+            output_size=sum(self.output_sizes),
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+
+    def validate_shard_id(self, loaded_shard_id: str | None) -> None:
+        if loaded_shard_id is None:
+            return
+        if loaded_shard_id not in ("q", "k", "v", "index_q", "index_k"):
+            raise ValueError(
+                "Shard id for MinimaxM3QKVParallelLinearWithIndexer must be one of "
+                "'q', 'k', 'v', 'index_q', 'index_k'; got "
+                f"{loaded_shard_id}."
+            )
+
+    def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int | None:
+        h = self.head_size
+        nq, nkv, nidx = self.num_heads, self.num_kv_heads, self.num_index_heads
+        return {
+            "q": 0,
+            "k": nq * h,
+            "v": (nq + nkv) * h,
+            "index_q": (nq + 2 * nkv) * h,
+            "index_k": (nq + 2 * nkv + nidx) * h,
+        }.get(loaded_shard_id)
+
+    def _get_shard_size_mapping(self, loaded_shard_id: str) -> int | None:
+        h = self.head_size
+        return {
+            "q": self.num_heads * h,
+            "k": self.num_kv_heads * h,
+            "v": self.num_kv_heads * h,
+            "index_q": self.num_index_heads * h,
+            "index_k": self.index_head_size,
+        }.get(loaded_shard_id)
+
+    def weight_loader_v2(
+        self,
+        param: BasevLLMParameter,
+        loaded_weight: torch.Tensor,
+        loaded_shard_id: str | None = None,
+    ) -> None:
+        self.validate_shard_id(loaded_shard_id)
+        # Index checkpoints are never pre-fused on disk; a shard id is always given.
+        assert loaded_shard_id in ("q", "k", "v", "index_q", "index_k")
+
+        shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
+        shard_size = self._get_shard_size_mapping(loaded_shard_id)
+        assert shard_offset is not None and shard_size is not None
+        if isinstance(param, BlockQuantScaleParameter):
+            weight_block_size = getattr(self, "weight_block_size", None)
+            shard_size, shard_offset = adjust_block_scale_shard(
+                weight_block_size, shard_size, shard_offset
+            )
+
+        # index_k is fully replicated: num_heads == tp_size makes
+        # load_qkv_weight pick shard_id_int == 0 on every rank. q/k/v/index_q ride
+        # the KV-head replication factor.
+        num_heads = (
+            self.tp_size if loaded_shard_id == "index_k" else self.num_kv_head_replicas
+        )
+        param.load_qkv_weight(
+            loaded_weight=loaded_weight,
+            num_heads=num_heads,
+            shard_id=loaded_shard_id,
+            shard_offset=shard_offset,
+            shard_size=shard_size,
+            tp_rank=self.tp_rank,
+        )
+
+    def weight_loader(
+        self,
+        param: Parameter,
+        loaded_weight: torch.Tensor,
+        loaded_shard_id: str | None = None,
+    ) -> None:
+        # Unquantized (bf16) path. MXFP8 checkpoints use weight_loader_v2; this
+        # keeps an unquantized load correct too.
+        self.validate_shard_id(loaded_shard_id)
+        assert loaded_shard_id in ("q", "k", "v", "index_q", "index_k")
+        output_dim = getattr(param, "output_dim", None)
+        assert output_dim is not None
+
+        shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
+        shard_size = self._get_shard_size_mapping(loaded_shard_id)
+        assert shard_offset is not None and shard_size is not None
+        if isinstance(param, BlockQuantScaleParameter):
+            weight_block_size = getattr(self, "weight_block_size", None)
+            shard_size, shard_offset = adjust_block_scale_shard(
+                weight_block_size, shard_size, shard_offset
+            )
+
+        param_data = param.data.narrow(output_dim, shard_offset, shard_size)
+        if loaded_shard_id == "q":
+            shard_rank = self.tp_rank
+        elif loaded_shard_id == "index_k":
+            shard_rank = 0  # replicated to every rank
+        else:
+            shard_rank = self.tp_rank // self.num_kv_head_replicas
+        loaded_weight = loaded_weight.narrow(
+            output_dim, shard_rank * shard_size, shard_size
+        )
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+
+
 # --8<-- [start:row_parallel_linear]
 @PluggableLayer.register("row_parallel_linear")
 class RowParallelLinear(LinearBase):

vllm/model_executor/layers/quantization/__init__.py

@@ -165,17 +165,18 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
         "deepseek_v4_fp8": DeepseekV4FP8Config,
         "humming": HummingConfig,
         "online": OnlineQuantizationConfig,
+        # MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the
+        # ModelOpt MXFP8 config (same format). The "mxfp8" online shorthand
+        # below only applies to the `--quantization mxfp8` CLI path.
+        "mxfp8": ModelOptMxFp8Config,
     }
 
-    # Register online shorthands as quantization methods so the user can
-    # specify "LLM(..., quantization='fp8_per_tensor')" as shorthand for
-    # creating a more complicated online quant config object.
+    # Register online shorthands (e.g. "fp8_per_tensor") as quant methods.
+    # setdefault so a shorthand that is also a checkpoint method (e.g. "mxfp8")
+    # keeps its checkpoint config; the shorthand still works via the
+    # `--quantization` CLI path in `resolve_quantization_config`.
     for shorthand in _ONLINE_SHORTHANDS:
-        assert shorthand not in method_to_config, (
-            f"Online quant shorthand {shorthand!r} conflicts with an "
-            f"existing quantization method"
-        )
-        method_to_config[shorthand] = OnlineQuantizationConfig
+        method_to_config.setdefault(shorthand, OnlineQuantizationConfig)
 
     # Update the `method_to_config` with customized quantization methods.
     method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)

vllm/model_executor/layers/quantization/modelopt.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from fnmatch import fnmatch
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 import torch
 from torch.nn.parameter import Parameter
@@ -1714,6 +1714,22 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase):
             return "modelopt_mxfp8"
         return None
 
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "ModelOptMxFp8Config":
+        # MiniMax-style checkpoints tag `quant_method: "mxfp8"` + `ignored_layers`
+        # (same on-disk format as ModelOpt MXFP8); normalize to the ModelOpt
+        # schema and reuse the shared parser.
+        if "quantization" not in config and not config.get("quant_algo"):
+            config = {
+                "quant_method": "modelopt",
+                "quantization": {
+                    "quant_algo": "MXFP8",
+                    "kv_cache_quant_algo": config.get("kv_cache_quant_algo"),
+                    "exclude_modules": config.get("ignored_layers", []) or [],
+                },
+            }
+        return cast("ModelOptMxFp8Config", super().from_config(config))
+
     @classmethod
     def _from_config(
         cls,
@@ -2125,6 +2141,9 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
             a1_scale=None,
             a2_scale=None,
             block_shape=self.weight_block_size,
+            swiglu_limit=getattr(layer, "swiglu_limit", None),
+            gemm1_alpha=getattr(layer, "swiglu_alpha", None),
+            gemm1_beta=getattr(layer, "swiglu_beta", None),
         )
 
     def apply_monolithic(

vllm/model_executor/layers/quantization/utils/fp8_utils.py

@@ -159,75 +159,104 @@ def _silu_mul_quant_fp8_packed_kernel(
     output_q_stride_m,
     output_scale_stride_k,
     clamp_limit,
+    alpha,
+    beta,
     N: tl.constexpr,
-    NUM_GROUPS: tl.constexpr,
+    GROUPS_PER_ROW: tl.constexpr,
+    PACKS_PER_ROW: tl.constexpr,
     fp8_min: tl.constexpr,
     fp8_max: tl.constexpr,
     GROUP_SIZE: tl.constexpr,
+    PACKS_PER_CTA: tl.constexpr,
     BLOCK_M: tl.constexpr,
     HAS_CLAMP: tl.constexpr,
 ):
-    N_2: tl.constexpr = N // 2
+    GROUPS_PER_PACK: tl.constexpr = 4
+    hidden_size: tl.constexpr = N // 2
 
-    pid_pack = tl.program_id(0)
-    pid_m = tl.program_id(1)
-    m_offset = pid_m.to(tl.int64) * BLOCK_M
+    pack_tile = tl.program_id(0)
+    row_start = tl.program_id(1).to(tl.int64) * BLOCK_M
+    row_step = tl.num_programs(1).to(tl.int64) * BLOCK_M
 
-    if m_offset >= M:
-        return
+    groups_per_cta: tl.constexpr = PACKS_PER_CTA * GROUPS_PER_PACK
+    elems_per_cta: tl.constexpr = groups_per_cta * GROUP_SIZE
+    col_start = pack_tile * elems_per_cta
+    col_offsets = tl.arange(0, elems_per_cta)
+    row_offsets = tl.arange(0, BLOCK_M)
+    pack_offsets = tl.arange(0, PACKS_PER_CTA)
 
-    offs_m = tl.arange(0, BLOCK_M)
-    offs_n = tl.arange(0, GROUP_SIZE)
-    row_mask = (m_offset + offs_m) < M
+    col_mask = (col_start + col_offsets) < (GROUPS_PER_ROW * GROUP_SIZE)
 
-    base_row_offset = (m_offset + offs_m[:, None]) * input_stride_m
-    base_out_offset = (m_offset + offs_m[:, None]) * output_q_stride_m
+    # persistent with grid_m-stride loop
+    while row_start < M:
+        rows = row_start + row_offsets
+        row_mask = rows < M
+        input_row_start = rows[:, None] * input_stride_m
+        output_row_start = rows[:, None] * output_q_stride_m
 
-    packed_scale = tl.zeros((BLOCK_M,), dtype=tl.int32)
+        gate_flat = tl.load(
+            input_ptr + input_row_start + col_start + col_offsets[None, :],
+            mask=row_mask[:, None] & col_mask[None, :],
+            other=0.0,
+        )
+        up_flat = tl.load(
+            input_ptr
+            + input_row_start
+            + hidden_size
+            + col_start
+            + col_offsets[None, :],
+            mask=row_mask[:, None] & col_mask[None, :],
+            other=0.0,
+        )
 
-    for pack_idx in tl.static_range(4):
-        group_id = pid_pack * 4 + pack_idx
+        gate = tl.reshape(gate_flat, (BLOCK_M, groups_per_cta, GROUP_SIZE)).to(
+            tl.float32
+        )
+        up = tl.reshape(up_flat, (BLOCK_M, groups_per_cta, GROUP_SIZE)).to(tl.float32)
 
-        if group_id < NUM_GROUPS:
-            n_offset = group_id * GROUP_SIZE
+        if HAS_CLAMP:
+            gate = tl.minimum(gate, clamp_limit)
+            up = tl.clamp(up, -clamp_limit, clamp_limit)
 
-            act_ptrs = input_ptr + base_row_offset + n_offset + offs_n[None, :]
-            act_in = tl.load(act_ptrs, mask=row_mask[:, None], other=0.0)
+        # Unified gated activation: silu == swigluoai with alpha=1, beta=0.
+        #   glu = gate * sigmoid(alpha * gate); y = (up + beta) * glu
+        glu = gate / (1.0 + tl.exp(-gate * alpha))
+        y = (up + beta) * glu
+        # Round through bf16 to match unfused precision path
+        y = y.to(tl.bfloat16).to(tl.float32)
 
-            mul_ptrs = act_ptrs + N_2
-            mul_in = tl.load(mul_ptrs, mask=row_mask[:, None], other=0.0)
+        absmax = tl.max(tl.abs(y), axis=2)
+        scale_raw = tl.maximum(absmax / fp8_max, 1e-10)
+        exponent = tl.ceil(tl.log2(scale_raw))
+        scale = tl.math.exp2(exponent)
 
-            act_f32 = act_in.to(tl.float32)
-            mul_f32 = mul_in.to(tl.float32)
+        y_q = tl.clamp(y / scale[:, :, None], fp8_min, fp8_max)
 
-            if HAS_CLAMP:
-                act_f32 = tl.minimum(act_f32, clamp_limit)
-                mul_f32 = tl.clamp(mul_f32, -clamp_limit, clamp_limit)
+        y_q_flat = tl.reshape(y_q, (BLOCK_M, elems_per_cta))
+        tl.store(
+            output_q_ptr + output_row_start + col_start + col_offsets[None, :],
+            y_q_flat.to(output_q_ptr.dtype.element_ty),
+            mask=row_mask[:, None] & col_mask[None, :],
+        )
 
-            y = (act_f32 / (1.0 + tl.exp(-act_f32))) * mul_f32
-            # Round through bf16 to match unfused precision path
-            y = y.to(tl.bfloat16).to(tl.float32)
+        scale_byte = tl.clamp(exponent + 127.0, 0.0, 255.0).to(tl.int32)
+        scale_bytes = tl.reshape(scale_byte, (BLOCK_M, PACKS_PER_CTA, GROUPS_PER_PACK))
+        shifts = tl.arange(0, GROUPS_PER_PACK) * 8
+        packed_scale = tl.sum(scale_bytes << shifts[None, None, :], axis=2)
 
-            absmax = tl.max(tl.abs(y), axis=1)
+        scale_pack = pack_tile * PACKS_PER_CTA + pack_offsets
+        scale_ptrs = (
+            output_scale_ptr
+            + scale_pack[None, :] * output_scale_stride_k
+            + rows[:, None]
+        )
+        tl.store(
+            scale_ptrs,
+            packed_scale,
+            mask=row_mask[:, None] & (scale_pack[None, :] < PACKS_PER_ROW),
+        )
 
-            scale_raw = tl.maximum(absmax / fp8_max, 1e-10)
-            exponent = tl.ceil(tl.log2(scale_raw))
-            scale = tl.math.exp2(exponent)
-
-            y_q = tl.clamp(y / scale[:, None], fp8_min, fp8_max)
-
-            out_q_ptrs = output_q_ptr + base_out_offset + n_offset + offs_n[None, :]
-            tl.store(
-                out_q_ptrs,
-                y_q.to(output_q_ptr.dtype.element_ty),
-                mask=row_mask[:, None],
-            )
-
-            exponent_biased = tl.clamp(exponent + 127.0, 0.0, 255.0).to(tl.int32)
-            packed_scale = packed_scale | (exponent_biased << (pack_idx * 8))
-
-    scale_ptrs = output_scale_ptr + pid_pack * output_scale_stride_k + m_offset + offs_m
-    tl.store(scale_ptrs, packed_scale, mask=row_mask)
+        row_start += row_step
 
 
 def silu_mul_quant_fp8_packed_triton(
@@ -235,37 +264,49 @@ def silu_mul_quant_fp8_packed_triton(
     group_size: int = 128,
     output_q: torch.Tensor | None = None,
     clamp_limit: float | None = None,
+    alpha: float = 1.0,
+    beta: float = 0.0,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert input.dim() == 2
     assert input.is_contiguous()
 
     M, N = input.shape
-    N_2 = N // 2
+    hidden_size = N // 2
 
-    assert N_2 % group_size == 0
+    assert hidden_size % group_size == 0
 
     fp8_dtype = torch.float8_e4m3fn
     finfo = torch.finfo(fp8_dtype)
     fp8_min, fp8_max = finfo.min, finfo.max
 
-    num_groups_per_row = N_2 // group_size
-    num_packed_groups = (num_groups_per_row + 3) // 4
-    tma_aligned_M = ((M + 3) // 4) * 4
+    groups_per_row = hidden_size // group_size
+    groups_per_pack = 4  # pack 4 UE8M0 scales to a single INT32
+    assert groups_per_row % groups_per_pack == 0
+    packs_per_row = groups_per_row // groups_per_pack
 
     if output_q is None:
-        output_q = torch.empty((M, N_2), dtype=fp8_dtype, device=input.device)
+        output_q = torch.empty((M, hidden_size), dtype=fp8_dtype, device=input.device)
 
+    aligned_m = triton.cdiv(M, 4) * 4
     output_scale_packed = torch.empty(
-        (num_packed_groups, tma_aligned_M),
+        (packs_per_row, aligned_m),
         dtype=torch.int32,
         device=input.device,
     ).T[:M, :]
 
-    BLOCK_M = 8
-    grid = (num_packed_groups, (M + BLOCK_M - 1) // BLOCK_M)
-
-    num_warps = max(4, group_size // 32)
+    # Tuned for group_size=32 (MXFP8) and group_size=128 (DeepSeek-V4)
+    num_warps = 4
     num_stages = 2
+    if group_size < 128:
+        BM = 1
+        packs_per_cta = 8
+    else:
+        BM = 1 if M < 512 else 4
+        packs_per_cta = 2 if M < 512 else 1
+
+    grid_n = triton.cdiv(packs_per_row, packs_per_cta)
+    grid_m = min(triton.cdiv(M, BM), 4096)
+    grid = (grid_n, grid_m)
 
     has_clamp = clamp_limit is not None
     _silu_mul_quant_fp8_packed_kernel[grid](
@@ -277,12 +318,16 @@ def silu_mul_quant_fp8_packed_triton(
         output_q.stride(0),
         output_scale_packed.stride(1),
         clamp_limit if has_clamp else 0.0,
+        alpha,
+        beta,
         N=N,
-        NUM_GROUPS=num_groups_per_row,
+        GROUPS_PER_ROW=groups_per_row,
+        PACKS_PER_ROW=packs_per_row,
         fp8_min=fp8_min,
         fp8_max=fp8_max,
         GROUP_SIZE=group_size,
-        BLOCK_M=BLOCK_M,
+        PACKS_PER_CTA=packs_per_cta,
+        BLOCK_M=BM,
         HAS_CLAMP=has_clamp,
         num_warps=num_warps,
         num_stages=num_stages,
@@ -303,6 +348,8 @@ def _silu_mul_per_token_group_quant_fp8_colmajor(
     # Information for float8
     eps,
     clamp_limit,
+    alpha,
+    beta,
     fp8_min: tl.constexpr,
     fp8_max: tl.constexpr,
     use_ue8m0: tl.constexpr,
@@ -348,10 +395,14 @@ def _silu_mul_per_token_group_quant_fp8_colmajor(
         mul_in = tl.clamp(mul_in.to(tl.float32), -clamp_limit, clamp_limit).to(
             y_ptr.dtype.element_ty
         )
+    # Unified gated activation: silu == swigluoai with alpha=1, beta=0.
+    #   glu = gate * sigmoid(alpha * gate); y = (up + beta) * glu
+    # Keep glu/up at input precision (narrow before the mul) so the alpha=1,
+    # beta=0 defaults match the C++ silu_and_mul path bit-for-bit.
     act_in = act_in.to(tl.float32)
-    one_f32 = tl.cast(1, tl.float32)
-    silu_out = (act_in / (one_f32 + tl.exp(-act_in))).to(y_ptr.dtype.element_ty)
-    y = (silu_out * mul_in).to(tl.float32)
+    glu = (act_in / (1.0 + tl.exp(-act_in * alpha))).to(y_ptr.dtype.element_ty)
+    up = (mul_in.to(tl.float32) + beta).to(y_ptr.dtype.element_ty)
+    y = (glu * up).to(tl.float32)
 
     # quant
     _absmax = tl.maximum(tl.max(tl.abs(y), axis=1), eps)
@@ -379,11 +430,15 @@ def silu_mul_per_token_group_quant_fp8_colmajor(
     use_ue8m0: bool | None = None,
     eps: float = 1e-10,
     clamp_limit: float | None = None,
+    group_size: int = 128,
+    alpha: float = 1.0,
+    beta: float = 0.0,
 ):
     """
-    silu+mul + block-fp8 quant with group size 128.
+    Gated activation + block-fp8 quant. ``alpha``/``beta`` select the gate
+    (silu: alpha=1, beta=0; swigluoai: alpha, beta from config).
     """
-    GROUP_SIZE = 128
+    GROUP_SIZE = group_size
     assert input.ndim == 2
     if output is not None:
         assert output.ndim == 2
@@ -431,6 +486,8 @@ def silu_mul_per_token_group_quant_fp8_colmajor(
         output_scales.stride(-1),
         eps,
         clamp_limit if has_clamp else 0.0,
+        alpha,
+        beta,
         fp8_min,
         fp8_max,
         use_ue8m0,
@@ -1020,9 +1077,10 @@ def deepgemm_post_process_fp8_weight_block(
         f"to be torch.float8_e4m3fn, got {wq.dtype} instead."
     )
 
-    if ws.dtype == torch.float8_e8m0fnu:
-        # Scales already in E8M0 from checkpoint — upcast to fp32
-        # and skip requantization (weights already have power-of-two scales).
+    if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
+        # Scales already in E8M0 from checkpoint (float8_e8m0fnu, or raw E8M0
+        # bits as uint8 for MXFP8) — upcast to fp32 and skip requantization
+        # (weights already have power-of-two scales).
         ws = _upcast_e8m0_to_fp32(ws)
     else:
         assert ws.dtype == torch.float32, (
@@ -1062,7 +1120,8 @@ def deepgemm_post_process_fp8_weight_block(
         ws = ws.unsqueeze(0)
 
     # From https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/utils/layout.hpp#L46
-    recipe = (1, 128, 128)
+    # (1, block_n, block_k): (1, 128, 128) for FP8 block, (1, 1, 32) for MXFP8.
+    recipe = (1, quant_block_shape[0], quant_block_shape[1])
 
     # Ref : https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/apis/gemm.hpp
     # DeepGemm uses the `transform_sf_into_required_layout` function to

vllm/model_executor/layers/quantization/utils/mxfp8_utils.py

@@ -84,6 +84,73 @@ def _mxfp8_e4m3_quantize_torch(
     return x_fp8, scales_uint8
 
 
+def _mxfp8_quant_triton_kernel():
+    """Lazily-built Triton kernel: per-32-block E8M0 scale + FP8-E4M3 quant.
+
+    Fuses what ``_mxfp8_e4m3_quantize_torch`` does in several elementwise passes
+    into one launch. Each program handles ``[BLOCK_M, 32]`` (one MX block).
+    """
+    from vllm.triton_utils import tl, triton
+
+    @triton.jit
+    def _kernel(
+        x_ptr, xq_ptr, s_ptr, M, K,
+        sxm, sxk, sqm, sqk, ssm, ssk,
+        BLOCK_M: tl.constexpr,
+    ):
+        pid_m = tl.program_id(0)
+        pid_b = tl.program_id(1)  # which 32-element block along K
+        offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        offs_k = pid_b * 32 + tl.arange(0, 32)
+        m_mask = offs_m < M
+        x = tl.load(
+            x_ptr + offs_m[:, None] * sxm + offs_k[None, :] * sxk,
+            mask=m_mask[:, None], other=0.0,
+        ).to(tl.float32)
+        amax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-30)  # [BLOCK_M]
+        sb = tl.floor(tl.log2(amax)) + 127.0
+        sb = tl.minimum(tl.maximum(sb, 0.0), 254.0)
+        descale = tl.exp2(sb - 127.0)
+        xq = (x / descale[:, None]).to(xq_ptr.dtype.element_ty)
+        tl.store(
+            xq_ptr + offs_m[:, None] * sqm + offs_k[None, :] * sqk, xq,
+            mask=m_mask[:, None],
+        )
+        tl.store(s_ptr + offs_m * ssm + pid_b * ssk, sb.to(tl.uint8), mask=m_mask)
+
+    return _kernel
+
+
+_MXFP8_QUANT_KERNEL = None
+
+
+def _mxfp8_e4m3_quantize_triton(
+    x: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fused 2D MXFP8 quant (non-swizzled, row-major [M, K//32] scales)."""
+    from vllm.triton_utils import triton
+
+    global _MXFP8_QUANT_KERNEL
+    if _MXFP8_QUANT_KERNEL is None:
+        _MXFP8_QUANT_KERNEL = _mxfp8_quant_triton_kernel()
+
+    M, K = x.shape
+    x = x.contiguous()
+    xq = torch.empty((M, K), dtype=MXFP8_VALUE_DTYPE, device=x.device)
+    scales = torch.empty(
+        (M, K // MXFP8_BLOCK_SIZE), dtype=MXFP8_SCALE_DTYPE, device=x.device
+    )
+    BLOCK_M = 64
+    grid = (triton.cdiv(M, BLOCK_M), K // MXFP8_BLOCK_SIZE)
+    _MXFP8_QUANT_KERNEL[grid](
+        x, xq, scales, M, K,
+        x.stride(0), x.stride(1), xq.stride(0), xq.stride(1),
+        scales.stride(0), scales.stride(1),
+        BLOCK_M=BLOCK_M,
+    )
+    return xq, scales
+
+
 def _mxfp8_e4m3_quantize_impl(
     x: torch.Tensor,
     is_sf_swizzled_layout: bool = False,
@@ -103,6 +170,17 @@ def _mxfp8_e4m3_quantize_impl(
             x_scales = x_scales.view(x.size(0), -1)
         return x_q, x_scales
 
+    # ROCm: a single fused Triton kernel beats the multi-pass torch path for the
+    # common 2D, non-swizzled activation-quant case (used by the native MX
+    # linear/MoE). Falls back to torch otherwise (3D weights, swizzled layout).
+    if (
+        current_platform.is_rocm()
+        and not is_sf_swizzled_layout
+        and x.ndim == 2
+        and x.shape[-1] % MXFP8_BLOCK_SIZE == 0
+    ):
+        return _mxfp8_e4m3_quantize_triton(x)
+
     return _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout)
 
 

vllm/model_executor/models/registry.py

@@ -165,6 +165,10 @@ _TEXT_GENERATION_MODELS = {
     "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
     "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
     "MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"),
+    "MiniMaxM3SparseForCausalLM": (
+        "vllm.models.minimax_m3",
+        "MiniMaxM3SparseForCausalLM",
+    ),
     "Ministral3ForCausalLM": ("mistral", "MistralForCausalLM"),
     "MistralForCausalLM": ("mistral", "MistralForCausalLM"),
     "MistralLarge3ForCausalLM": ("mistral_large_3", "MistralLarge3ForCausalLM"),
@@ -477,6 +481,10 @@ _MULTIMODAL_MODELS = {
     "MantisForConditionalGeneration": ("llava", "MantisForConditionalGeneration"),
     "MiDashengLMModel": ("midashenglm", "MiDashengLMModel"),
     "MiMoV2OmniForCausalLM": ("mimo_v2_omni", "MiMoV2OmniForCausalLM"),
+    "MiniMaxM3SparseForConditionalGeneration": (
+        "vllm.models.minimax_m3",
+        "MiniMaxM3SparseForConditionalGeneration",
+    ),
     "MiniMaxVL01ForConditionalGeneration": (
         "minimax_vl_01",
         "MiniMaxVL01ForConditionalGeneration",

vllm/models/minimax_m3/__init__.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax M3 model — hardware-isolated entry point.
+
+The implementation lives under ``nvidia/`` and ``amd/``; this module picks the
+right one for the current platform and re-exports the public classes used by
+the model registry. (Mirrors ``vllm.models.deepseek_v4``.)
+"""
+
+from typing import TYPE_CHECKING
+
+from vllm.platforms import current_platform
+
+# The NVIDIA branch is the static default that type-checkers see; the ROCm
+# branch overrides it at runtime (kept type-compatible via type: ignore).
+if TYPE_CHECKING or not current_platform.is_rocm():
+    from .nvidia.model import (
+        MiniMaxM3SparseForCausalLM,
+        MiniMaxM3SparseForConditionalGeneration,
+    )
+else:
+    from .amd.model import (  # type: ignore[assignment]
+        MiniMaxM3SparseForCausalLM,
+        MiniMaxM3SparseForConditionalGeneration,
+    )
+
+__all__ = [
+    "MiniMaxM3SparseForCausalLM",
+    "MiniMaxM3SparseForConditionalGeneration",
+]

vllm/models/minimax_m3/amd/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/minimax_m3/amd/ops/__init__.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""AMD/ROCm fused Triton ops for MiniMax-M3.
+
+These replace per-element PyTorch fallbacks (FlashInfer / fused HIP kernels are
+unavailable on ROCm) with single-pass Triton kernels to cut launch overhead and
+intermediate-tensor traffic during decode.
+"""
+
+from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import (
+    gemma_fused_add_rmsnorm,
+    gemma_rmsnorm,
+)
+from vllm.models.minimax_m3.amd.ops.swiglu_oai import swiglu_oai_split
+
+__all__ = [
+    "gemma_rmsnorm",
+    "gemma_fused_add_rmsnorm",
+    "swiglu_oai_split",
+]

vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused Gemma-style RMSNorm for AMD ROCm via Triton.
+
+Gemma RMSNorm = normalize(x) * (1 + weight), computed in fp32. FlashInfer's
+``gemma_rmsnorm`` / ``gemma_fused_add_rmsnorm`` CUDA kernels are unavailable on
+ROCm, so the AMD path previously used a ~8-op PyTorch sequence (float cast, add,
+pow, mean, rsqrt, two muls, cast) — each a separate kernel launch materializing
+fp32 intermediates. These kernels collapse that into a single pass per row.
+
+Two entry points:
+  * ``gemma_rmsnorm(x, w, eps)``                 -> normalized tensor
+  * ``gemma_fused_add_rmsnorm(x, res, w, eps)``  -> (normalized, x + res)
+
+Both normalize over the last dim and broadcast ``weight`` (shape [N]) over it,
+so they serve both the full-hidden norms (input/post-attn/final) and the
+per-head q_norm/k_norm (N == head_dim). Inputs may be non-contiguous views
+(e.g. ``qkv.split`` slices); strides are passed through and outputs are written
+contiguous.
+"""
+
+import torch
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _gemma_rmsnorm_kernel(
+    x_ptr,
+    w_ptr,
+    out_ptr,
+    n_cols,
+    stride_row,
+    stride_col,
+    eps,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < n_cols
+    x = tl.load(
+        x_ptr + row * stride_row + cols * stride_col, mask=mask, other=0.0
+    ).to(tl.float32)
+    var = tl.sum(x * x, axis=0) / n_cols
+    rstd = 1.0 / tl.sqrt(var + eps)
+    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    out = x * rstd * (1.0 + w)
+    tl.store(
+        out_ptr + row * n_cols + cols,
+        out.to(out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+@triton.jit
+def _gemma_fused_add_rmsnorm_kernel(
+    x_ptr,
+    res_ptr,
+    w_ptr,
+    out_ptr,
+    res_out_ptr,
+    n_cols,
+    stride_xrow,
+    stride_xcol,
+    stride_rrow,
+    stride_rcol,
+    eps,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < n_cols
+    x = tl.load(
+        x_ptr + row * stride_xrow + cols * stride_xcol, mask=mask, other=0.0
+    ).to(tl.float32)
+    r = tl.load(
+        res_ptr + row * stride_rrow + cols * stride_rcol, mask=mask, other=0.0
+    ).to(tl.float32)
+    s = x + r
+    # residual_out is the pre-norm sum (consumed by the next layer's add).
+    tl.store(
+        res_out_ptr + row * n_cols + cols,
+        s.to(res_out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+    var = tl.sum(s * s, axis=0) / n_cols
+    rstd = 1.0 / tl.sqrt(var + eps)
+    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    out = s * rstd * (1.0 + w)
+    tl.store(
+        out_ptr + row * n_cols + cols,
+        out.to(out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+def _num_warps(block_n: int) -> int:
+    if block_n >= 4096:
+        return 16
+    if block_n >= 1024:
+        return 8
+    return 4
+
+
+def gemma_rmsnorm(
+    x: torch.Tensor, weight: torch.Tensor, eps: float
+) -> torch.Tensor:
+    orig_shape = x.shape
+    n = orig_shape[-1]
+    x2 = x.reshape(-1, n)
+    m = x2.shape[0]
+    out = torch.empty((m, n), dtype=x.dtype, device=x.device)
+    block_n = triton.next_power_of_2(n)
+    _gemma_rmsnorm_kernel[(m,)](
+        x2,
+        weight,
+        out,
+        n,
+        x2.stride(0),
+        x2.stride(1),
+        eps,
+        BLOCK_N=block_n,
+        num_warps=_num_warps(block_n),
+    )
+    return out.reshape(orig_shape)
+
+
+def gemma_fused_add_rmsnorm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    orig_shape = x.shape
+    n = orig_shape[-1]
+    x2 = x.reshape(-1, n)
+    r2 = residual.reshape(-1, n)
+    m = x2.shape[0]
+    out = torch.empty((m, n), dtype=x.dtype, device=x.device)
+    res_out = torch.empty((m, n), dtype=x.dtype, device=x.device)
+    block_n = triton.next_power_of_2(n)
+    _gemma_fused_add_rmsnorm_kernel[(m,)](
+        x2,
+        r2,
+        weight,
+        out,
+        res_out,
+        n,
+        x2.stride(0),
+        x2.stride(1),
+        r2.stride(0),
+        r2.stride(1),
+        eps,
+        BLOCK_N=block_n,
+        num_warps=_num_warps(block_n),
+    )
+    return out.reshape(orig_shape), res_out.reshape(orig_shape)

vllm/models/minimax_m3/amd/ops/swiglu_oai.py

@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused SwiGLU-OAI activation (split layout) for AMD ROCm via Triton.
+
+SwiGLU-OAI on a ``[*, 2I]`` split-layout input (gate = first half, up = second
+half):
+
+    gate = clamp(gate, max=limit)
+    up   = clamp(up, -limit, +limit)
+    out  = gate * sigmoid(alpha * gate) * (up + beta)
+
+On ROCm the dense MLP and the native MXFP8 MoE (between its two GEMMs) fell back
+to a chain of elementwise PyTorch ops with fp32 intermediates: vLLM's shared
+``SiluAndMulWithClamp`` blanket-routes ROCm to ``forward_native``, and the MoE
+applies the activation inline in PyTorch. This Triton kernel collapses that into
+a single pass producing the ``[*, I]`` output directly, and computes in fp32
+(rel ~1e-6 vs reference).
+
+Note: the vectorized ``torch.ops._C.silu_and_mul_with_clamp`` op IS built on
+ROCm and is ~1.2-2.2x faster in isolation, but the win is launch overhead that
+HIP graphs already eliminate — measured end-to-end throughput is identical
+(within noise), so we keep the fp32-accurate Triton kernel.
+"""
+
+import torch
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _swiglu_oai_kernel(
+    g_ptr,
+    out_ptr,
+    n_inter,
+    stride_gm,
+    stride_gn,
+    stride_om,
+    stride_on,
+    alpha,
+    beta,
+    limit,
+    HAS_LIMIT: tl.constexpr,
+    BLOCK_I: tl.constexpr,
+):
+    row = tl.program_id(0)
+    pid_i = tl.program_id(1)
+    cols = pid_i * BLOCK_I + tl.arange(0, BLOCK_I)
+    mask = cols < n_inter
+    gate = tl.load(
+        g_ptr + row * stride_gm + cols * stride_gn, mask=mask, other=0.0
+    ).to(tl.float32)
+    up = tl.load(
+        g_ptr + row * stride_gm + (n_inter + cols) * stride_gn,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    if HAS_LIMIT:
+        gate = tl.minimum(gate, limit)
+        up = tl.minimum(tl.maximum(up, -limit), limit)
+    out = gate * tl.sigmoid(alpha * gate) * (up + beta)
+    tl.store(
+        out_ptr + row * stride_om + cols * stride_on,
+        out.to(out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+def swiglu_oai_split(
+    gate_up: torch.Tensor,
+    alpha: float,
+    beta: float,
+    limit: float | None,
+    out_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """SwiGLU-OAI on a split-layout ``[*, 2I]`` tensor -> ``[*, I]``."""
+    orig_shape = gate_up.shape
+    two_i = orig_shape[-1]
+    n_inter = two_i // 2
+    x2 = gate_up.reshape(-1, two_i)
+    m = x2.shape[0]
+    dt = out_dtype if out_dtype is not None else gate_up.dtype
+    out = torch.empty((m, n_inter), dtype=dt, device=gate_up.device)
+    # Tile tuned on gfx950. The SwiGLU intermediate is sharded across tensor
+    # parallel ranks (per-rank n_inter = I / tp: dense I=12288, MoE I=3072), and
+    # a 512-wide tile (4 warps, ~2 elems/lane) only helps once the per-rank slice
+    # is large enough to be bandwidth-bound — at TP=1 prefill that is ~1.25-1.35x
+    # faster than 256. For small sharded slices (high TP) the kernel is launch-
+    # bound (~12us) and a wide tile can slightly regress, so fall back to 256.
+    # Decode is launch-bound at every TP. num_warps=8 underfills this tile, so it
+    # is pinned to 4.
+    block_i = 512 if n_inter >= 2048 else 256
+    grid = (m, triton.cdiv(n_inter, block_i))
+    _swiglu_oai_kernel[grid](
+        x2,
+        out,
+        n_inter,
+        x2.stride(0),
+        x2.stride(1),
+        out.stride(0),
+        out.stride(1),
+        float(alpha),
+        float(beta),
+        0.0 if limit is None else float(limit),
+        HAS_LIMIT=limit is not None,
+        BLOCK_I=block_i,
+        num_warps=4,
+    )
+    return out.reshape(*orig_shape[:-1], n_inter)

vllm/models/minimax_m3/common/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/minimax_m3/common/mm_preprocess.py

@@ -0,0 +1,462 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+from collections.abc import Mapping, Sequence
+from typing import cast
+
+import torch
+from transformers import BatchFeature
+from transformers.video_utils import VideoMetadata
+
+from vllm.config.multimodal import (
+    BaseDummyOptions,
+    ImageDummyOptions,
+    VideoDummyOptions,
+)
+from vllm.inputs import MultiModalDataDict
+from vllm.multimodal.inputs import (
+    MultiModalFieldConfig,
+    MultiModalKwargsItems,
+)
+from vllm.multimodal.parse import (
+    ImageSize,
+    MultiModalDataItems,
+    MultiModalDataParser,
+)
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    PromptReplacement,
+    PromptUpdate,
+    PromptUpdateDetails,
+)
+from vllm.multimodal.video import (
+    VIDEO_LOADER_REGISTRY,
+    VideoBackend,
+    VideoSourceMetadata,
+    VideoTargetMetadata,
+)
+from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3Config
+from vllm.transformers_utils.processors.minimax_m3 import (
+    MiniMaxM3VLImageProcessor,
+    MiniMaxM3VLVideoProcessor,
+    MiniMaxVLProcessor,
+    smart_resize,
+)
+
+# Upper bound on the number of frames used to build the dummy video during
+# memory profiling. Without this cap, ``_get_max_video_frames(seq_len)`` with
+# M3's large ``max_model_len`` yields ~1800 frames, producing a multi-GB dummy
+# tensor that overflows the multimodal processor cache. Mirrors Qwen2-VL's
+# ``_MAX_FRAMES_PER_VIDEO``.
+_MAX_FRAMES_PER_VIDEO = 14
+
+
+class MiniMaxM3VLProcessingInfo(BaseProcessingInfo):
+    IMAGE_TOKEN = "]<]image[>["
+    VIDEO_TOKEN = "]<]video[>["
+    VISION_START_TOKEN = "]<]start of image[>["
+    VISION_END_TOKEN = "]<]end of image[>["
+
+    def get_hf_config(self) -> MiniMaxM3Config:
+        return self.ctx.get_hf_config(MiniMaxM3Config)
+
+    def get_hf_processor(self, **kwargs: object) -> MiniMaxVLProcessor:
+        # The released checkpoint only ships the processor as remote code
+        # (via ``auto_map``). Construct the vendored processor directly so the
+        # model loads without ``--trust-remote-code``.
+        return self.ctx.get_hf_processor(MiniMaxVLProcessor, **kwargs)
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"image": None, "video": None}
+
+    def get_mm_max_tokens_per_item(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> Mapping[str, int]:
+        return {
+            "image": self.get_max_image_tokens(),
+            "video": self.get_max_video_tokens(seq_len, mm_counts),
+        }
+
+    def get_image_processor(self, **kwargs: object) -> MiniMaxM3VLImageProcessor:
+        return self.get_hf_processor(**kwargs).image_processor
+
+    def get_video_processor(self, **kwargs: object) -> MiniMaxM3VLVideoProcessor:
+        return self.get_hf_processor(**kwargs).video_processor
+
+    def _get_vision_info(
+        self,
+        *,
+        image_width: int,
+        image_height: int,
+        num_frames: int,
+        image_processor,
+    ) -> tuple[ImageSize, int]:
+        """Compute resized image size and number of vision tokens.
+
+        Mirrors the processor's Qwen-style ``smart_resize`` (area bound by
+        ``max_pixels``) so token counts match the actual processor output.
+        """
+        patch_size: int = image_processor.patch_size
+        merge_size: int = image_processor.merge_size
+        temporal_patch_size: int = image_processor.temporal_patch_size
+        factor = patch_size * merge_size
+        max_pixels: int = image_processor.max_pixels
+
+        new_h, new_w = smart_resize(
+            image_height, image_width, factor=factor, max_pixels=max_pixels
+        )
+        grid_h = new_h // patch_size
+        grid_w = new_w // patch_size
+
+        # Pad frames to be divisible by temporal_patch_size
+        padded_frames = num_frames + (-num_frames % temporal_patch_size)
+        grid_t = max(padded_frames // temporal_patch_size, 1)
+
+        num_tokens = grid_t * grid_h * grid_w // (merge_size**2)
+        return ImageSize(width=new_w, height=new_h), num_tokens
+
+    def get_num_image_tokens(
+        self,
+        *,
+        image_width: int,
+        image_height: int,
+        image_processor,
+        mm_kwargs: Mapping[str, object],
+    ) -> int:
+        _, n = self._get_vision_info(
+            image_width=image_width,
+            image_height=image_height,
+            num_frames=1,
+            image_processor=image_processor,
+        )
+        return n
+
+    def get_num_video_tokens(
+        self,
+        *,
+        image_width: int,
+        image_height: int,
+        num_frames: int,
+        image_processor,
+        mm_kwargs: Mapping[str, object],
+    ) -> int:
+        _, n = self._get_vision_info(
+            image_width=image_width,
+            image_height=image_height,
+            num_frames=num_frames,
+            image_processor=image_processor,
+        )
+        return n
+
+    def get_image_size_with_most_features(self) -> ImageSize:
+        # Largest square (a multiple of patch_size*merge_size) whose area is
+        # within the image processor's ``max_pixels`` bound — this yields the
+        # most vision tokens for one image.
+        image_processor = self.get_image_processor()
+        factor = image_processor.patch_size * image_processor.merge_size
+        side = max(factor, (math.isqrt(image_processor.max_pixels) // factor) * factor)
+        return ImageSize(width=side, height=side)
+
+    def get_max_image_tokens(self) -> int:
+        image_processor = self.get_image_processor()
+        size = self.get_image_size_with_most_features()
+        return self.get_num_image_tokens(
+            image_width=size.width,
+            image_height=size.height,
+            image_processor=image_processor,
+            mm_kwargs={},
+        )
+
+    def _get_max_video_frames(self, max_tokens: int) -> int:
+        image_processor = self.get_image_processor()
+        size = self.get_image_size_with_most_features()
+        num_frames = 1
+        while True:
+            next_n = self.get_num_video_tokens(
+                image_width=size.width,
+                image_height=size.height,
+                num_frames=num_frames + 1,
+                image_processor=image_processor,
+                mm_kwargs={},
+            )
+            if next_n > max_tokens:
+                break
+            num_frames += 1
+        return num_frames
+
+    def get_num_frames_with_most_features(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        max_frames_per_video: int = _MAX_FRAMES_PER_VIDEO,
+    ) -> int:
+        max_videos = mm_counts.get("video", 0)
+        max_total_frames = self._get_max_video_frames(seq_len)
+        max_frames_per_video = min(
+            max_total_frames // max(max_videos, 1), max_frames_per_video
+        )
+        return max(max_frames_per_video, 1)
+
+    def get_max_video_tokens(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> int:
+        image_processor = self.get_image_processor()
+        size = self.get_image_size_with_most_features()
+        return self.get_num_video_tokens(
+            image_width=size.width,
+            image_height=size.height,
+            num_frames=self.get_num_frames_with_most_features(seq_len, mm_counts),
+            image_processor=image_processor,
+            mm_kwargs={},
+        )
+
+
+class MiniMaxM3VLDummyInputsBuilder(BaseDummyInputsBuilder[MiniMaxM3VLProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        num_images = mm_counts.get("image", 0)
+        num_videos = mm_counts.get("video", 0)
+        image_token: str = self.info.IMAGE_TOKEN
+        video_token: str = self.info.VIDEO_TOKEN
+        return image_token * num_images + video_token * num_videos
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions],
+    ) -> MultiModalDataDict:
+        size = self.info.get_image_size_with_most_features()
+        num_frames = self.info.get_num_frames_with_most_features(seq_len, mm_counts)
+        return {
+            "image": self._get_dummy_images(
+                width=size.width,
+                height=size.height,
+                num_images=mm_counts.get("image", 0),
+                overrides=cast(ImageDummyOptions | None, mm_options.get("image")),
+            ),
+            "video": self._get_dummy_videos(
+                width=size.width,
+                height=size.height,
+                num_frames=num_frames,
+                num_videos=mm_counts.get("video", 0),
+                overrides=cast(VideoDummyOptions | None, mm_options.get("video")),
+            ),
+        }
+
+
+class MiniMaxM3VLMultiModalProcessor(
+    BaseMultiModalProcessor[MiniMaxM3VLProcessingInfo]
+):
+    def _get_data_parser(self) -> MultiModalDataParser:
+        # Request video metadata (fps + sampled frame indices) so the HF
+        # processor can emit per-frame ``]<]X.X seconds[>[`` timestamp markers,
+        # matching MiniMax's reference video token stream. ``_get_prompt_updates``
+        # reconstructs the same markers from the metadata to keep the prompt
+        # replacement aligned with the processor output.
+        return MultiModalDataParser(video_needs_metadata=True)
+
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object],
+        tok_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        mm_data = dict(mm_data)
+        # With ``video_needs_metadata=True`` each video arrives as a
+        # ``(frames, metadata)`` tuple. Split the frames back out and forward the
+        # metadata as ``VideoMetadata`` so the processor emits timestamps.
+        videos = cast(list | None, mm_data.get("videos"))
+        video_metadata: list[VideoMetadata] | None = None
+        if videos:
+            frames_only = []
+            video_metadata = []
+            for item in videos:
+                if isinstance(item, tuple) and len(item) == 2:
+                    frames, meta = item
+                else:
+                    frames, meta = item, {}
+                frames_only.append(frames)
+                meta = {
+                    k: v for k, v in (meta or {}).items() if k != "do_sample_frames"
+                }
+                # VideoMetadata requires total_num_frames; derive it for
+                # dummy/profiling videos whose metadata omits it. fps and
+                # frames_indices default to None there → no timestamps, which
+                # stays consistent with _get_prompt_updates.
+                meta.setdefault("total_num_frames", len(frames))
+                video_metadata.append(VideoMetadata(**meta))
+            mm_data["videos"] = frames_only
+
+        # Override the video processor's default do_resize=False (set for
+        # SGLang's pre-resized pipeline) to True for vLLM's raw-frame inputs.
+        merged = dict(do_resize=True, **mm_kwargs, **tok_kwargs)
+        data = dict(text=prompt, **mm_data)
+        if video_metadata is not None:
+            data["video_metadata"] = video_metadata
+        return self.info.ctx.call_hf_processor(
+            self.info.get_hf_processor(**mm_kwargs),
+            data,
+            merged,
+        )
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        image_grid_thw = hf_inputs.get("image_grid_thw")
+        video_grid_thw = hf_inputs.get("video_grid_thw")
+
+        # Total patches per item (grid_t * grid_h * grid_w)
+        image_grid_sizes = (
+            image_grid_thw.prod(-1)
+            if image_grid_thw is not None
+            else torch.empty(0, dtype=torch.long)
+        )
+        video_grid_sizes = (
+            video_grid_thw.prod(-1)
+            if video_grid_thw is not None
+            else torch.empty(0, dtype=torch.long)
+        )
+
+        return {
+            "pixel_values": MultiModalFieldConfig.flat_from_sizes(
+                "image", image_grid_sizes
+            ),
+            "image_grid_thw": MultiModalFieldConfig.batched("image", keep_on_cpu=True),
+            "pixel_values_videos": MultiModalFieldConfig.flat_from_sizes(
+                "video", video_grid_sizes
+            ),
+            "video_grid_thw": MultiModalFieldConfig.batched("video", keep_on_cpu=True),
+        }
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+        tokenizer = self.info.get_tokenizer()
+        vocab = tokenizer.get_vocab()
+
+        image_token_id: int = vocab[self.info.IMAGE_TOKEN]
+        video_token_id: int = vocab[self.info.VIDEO_TOKEN]
+        start_token_id: int = vocab[self.info.VISION_START_TOKEN]
+        end_token_id: int = vocab[self.info.VISION_END_TOKEN]
+        merge_length: int = hf_processor.image_processor.merge_size**2
+
+        def get_image_replacement(item_idx: int):
+            grid_thw: torch.Tensor = out_mm_kwargs["image"][item_idx][
+                "image_grid_thw"
+            ].data
+            # grid_thw shape: (3,) = [1, grid_h, grid_w]
+            N = int(grid_thw.prod().item()) // merge_length
+            full = [start_token_id] + [image_token_id] * N + [end_token_id]
+            return PromptUpdateDetails.select_token_id(full, image_token_id)
+
+        # Per-video metadata (fps + sampled frame indices) is carried on the
+        # parsed video items; used to reproduce the HF processor's timestamps.
+        video_items = mm_items.get("video")
+        video_metadata = getattr(video_items, "metadata", None)
+        temporal_patch_size: int = hf_processor.video_processor.temporal_patch_size
+
+        def get_video_replacement(item_idx: int):
+            grid_thw: torch.Tensor = out_mm_kwargs["video"][item_idx][
+                "video_grid_thw"
+            ].data
+            # grid_thw shape: (3,) = [grid_t, grid_h, grid_w]
+            # HF model uses VIDEO_TOKEN (not IMAGE_TOKEN) for video frame content:
+            # processing_minimax.py L245: replace(placeholder, self.VIDEO_TOKEN)
+            T = int(grid_thw[0].item())
+            M = int(grid_thw[1].item() * grid_thw[2].item()) // merge_length
+
+            # Reproduce the HF processor's per-frame timestamp markers
+            # (processing_minimax.py: ts = frames_indices[frame*tps] / fps,
+            # rendered as "]<]X.X seconds[>["). Falls back to no timestamps when
+            # metadata is unavailable (keeping the replacement aligned with the
+            # processor output in both cases).
+            meta = (
+                video_metadata[item_idx]
+                if video_metadata is not None and item_idx < len(video_metadata)
+                else None
+            )
+            fps = meta.get("fps") if meta else None
+            frames_indices = meta.get("frames_indices") if meta else None
+
+            full: list[int] = []
+            for frame_idx in range(T):
+                if fps is not None and frames_indices is not None:
+                    idx = min(frame_idx * temporal_patch_size, len(frames_indices) - 1)
+                    ts = frames_indices[idx] / fps
+                    full += tokenizer.encode(
+                        f"]<]{ts:.1f} seconds[>[", add_special_tokens=False
+                    )
+                full += [start_token_id] + [video_token_id] * M + [end_token_id]
+            return PromptUpdateDetails.select_token_id(full, video_token_id)
+
+        return [
+            PromptReplacement(
+                modality="image",
+                target=[image_token_id],
+                replacement=get_image_replacement,
+            ),
+            PromptReplacement(
+                modality="video",
+                target=[video_token_id],
+                replacement=get_video_replacement,
+            ),
+        ]
+
+
+# TODO(Isotr0py): Tie with MinimaxVideoProcessor
+# after https://github.com/vllm-project/vllm/pull/44126
+@VIDEO_LOADER_REGISTRY.register("minimax_m3_vl")
+class MiniMaxM3VideoBackend(VideoBackend):
+    @classmethod
+    def compute_frames_index_to_sample(
+        cls,
+        source: VideoSourceMetadata,
+        target: VideoTargetMetadata,
+        **kwargs,
+    ) -> list[int]:
+        total_frames = source.total_frames_num
+        video_fps = source.original_fps
+        fps = target.fps
+
+        if total_frames <= 0 or video_fps <= 0 or fps <= 0:
+            return [0] if total_frames > 0 else []
+
+        read_time_interval = 1.0 / fps
+        eps = 1e-4
+
+        indices: list[int] = []
+        prev_kept_ts = -float("inf")
+        while True:
+            if not indices:
+                target_frame = 0
+            else:
+                target_ts = prev_kept_ts + read_time_interval - eps
+                target_frame = math.ceil(target_ts * video_fps)
+                target_frame = max(target_frame, indices[-1] + 1)
+            if target_frame >= total_frames:
+                break
+            indices.append(target_frame)
+            prev_kept_ts = target_frame / video_fps
+
+        last_frame_idx = total_frames - 1
+        last_ts = last_frame_idx / video_fps
+        if indices and indices[-1] != last_frame_idx and last_ts - prev_kept_ts > eps:
+            indices.append(last_frame_idx)
+
+        if not indices:
+            indices = [0]
+        return indices

vllm/models/minimax_m3/common/ops/__init__.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Cross-platform (Triton) kernels for MiniMax M3 sparse attention."""
+
+from .index_topk import minimax_m3_index_topk, minimax_m3_index_topk_decode
+from .sparse_attn import minimax_m3_sparse_attn, minimax_m3_sparse_attn_decode
+
+__all__ = [
+    "minimax_m3_index_topk",
+    "minimax_m3_index_topk_decode",
+    "minimax_m3_sparse_attn",
+    "minimax_m3_sparse_attn_decode",
+]

vllm/models/minimax_m3/common/ops/index_topk.py

@@ -0,0 +1,819 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Triton kernels for MiniMax M3 lightning-indexer block scoring + top-k.
+
+Index queries score each 128-token block of index keys (max over the block),
+then the top-k blocks (plus forced init/local blocks) are selected per query
+token. Ported from the sglang reference (minimax_sparse_ops), adapted to vLLM's
+paged KV cache: the KV page size is forced to equal the sparse block size (128),
+so one sparse block maps to exactly one page.
+
+Index-K cache layout (vLLM): ``(num_blocks, 128, idx_head_dim)`` (single head).
+
+Only the paths MiniMax M3 uses are implemented: score_type="max", index value
+disabled (score-only indexer), single shared index head. The selected block ids
+feed the block-sparse attention kernels in ``sparse_attn``.
+"""
+
+import torch
+
+from vllm.triton_utils import tl, triton
+
+# One sparse block == one KV page.
+SPARSE_BLOCK_SIZE = 128
+
+
+# ---------------------------------------------------------------------------
+# Bitonic top-k helpers (layout-agnostic; ported verbatim from sglang).
+# ---------------------------------------------------------------------------
+@triton.jit
+def _compare_and_swap(x, ids, flip, i: tl.constexpr, n_dims: tl.constexpr):
+    n_outer: tl.constexpr = x.numel >> n_dims
+    shape: tl.constexpr = [n_outer * 2**i, 2, 2 ** (n_dims - i - 1)]
+    y = tl.reshape(x, shape)
+    mask = tl.arange(0, 2)[None, :, None]
+    left = tl.broadcast_to(tl.sum(y * (1 - mask), 1)[:, None, :], shape).to(y.dtype)
+    right = tl.broadcast_to(tl.sum(y * mask, 1)[:, None, :], shape).to(y.dtype)
+    left = tl.reshape(left, x.shape)
+    right = tl.reshape(right, x.shape)
+    y_idx = tl.reshape(ids, shape)
+    left_idx = tl.broadcast_to(tl.sum(y_idx * (1 - mask), 1)[:, None, :], shape)
+    right_idx = tl.broadcast_to(tl.sum(y_idx * mask, 1)[:, None, :], shape)
+    left_idx = tl.reshape(left_idx, x.shape).to(y_idx.dtype)
+    right_idx = tl.reshape(right_idx, x.shape).to(y_idx.dtype)
+    idtype = tl.core.get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True)
+    ileft = left.to(idtype, bitcast=True)
+    iright = right.to(idtype, bitcast=True)
+    ix = x.to(idtype, bitcast=True)
+    cond = (left > right) != flip
+    ret = ix ^ tl.where(cond, ileft ^ iright, tl.zeros_like(ix))
+    new_ids = ids ^ tl.where(cond, left_idx ^ right_idx, tl.zeros_like(ids))
+    return ret.to(x.dtype, bitcast=True), new_ids
+
+
+@triton.jit
+def _bitonic_merge(
+    x, ids, stage: tl.constexpr, order: tl.constexpr, n_dims: tl.constexpr
+):
+    n_outer: tl.constexpr = x.numel >> n_dims
+    tl.static_assert(stage <= n_dims)
+    if order == 2:
+        shape: tl.constexpr = [n_outer * 2 ** (n_dims - 1 - stage), 2, 2**stage]
+        flip = tl.reshape(
+            tl.broadcast_to(tl.arange(0, 2)[None, :, None], shape), x.shape
+        )
+    else:
+        flip = order
+    for i in tl.static_range(stage):
+        x, ids = _compare_and_swap(x, ids, flip, i + (n_dims - stage), n_dims)
+    return x, ids
+
+
+# ---------------------------------------------------------------------------
+# Index block-score kernel (paged). score[h, token, block] = max over the
+# 128-token block of (idx_q . index_k), causal-masked. BLOCK_SIZE_K == 128 so
+# each K-tile is exactly one page (BLOCKS_PER_K_BLOCK == 1).
+# ---------------------------------------------------------------------------
+@triton.jit
+def _index_block_score_kernel(
+    q_ptr,  # idx_q: [total_q, num_idx_heads, head_dim]
+    ik_cache_ptr,  # index-K cache: [num_blocks, 128, head_dim]
+    score_ptr,  # [num_idx_heads, total_q, max_block]
+    block_table_ptr,  # [num_reqs, max_blocks]
+    cu_seqlens,  # [batch+1] query start offsets
+    seq_lens,  # [batch] total K length
+    prefix_lens,  # [batch] context length before this chunk's queries
+    num_idx_heads,
+    head_dim: tl.constexpr,
+    sm_scale,
+    stride_q_n,
+    stride_q_h,
+    stride_q_d,
+    stride_ik_blk,
+    stride_ik_pos,
+    stride_ik_d,
+    stride_s_h,
+    stride_s_n,
+    stride_s_k,
+    stride_bt_b,
+    BLOCK_SIZE_Q: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,  # == SPARSE_BLOCK_SIZE (128)
+):
+    sm_scale_log2e = sm_scale * 1.4426950409
+    pid_q = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    pid_b = pid_bh // num_idx_heads
+    pid_h = pid_bh % num_idx_heads
+
+    seq_start = tl.load(cu_seqlens + pid_b)
+    q_len = tl.load(cu_seqlens + pid_b + 1) - seq_start
+    seq_len = tl.load(seq_lens + pid_b)
+    prefix_len = tl.load(prefix_lens + pid_b)
+    if BLOCK_SIZE_Q * pid_q >= q_len:
+        return
+
+    q_ptrs = tl.make_block_ptr(
+        base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h,
+        shape=(q_len, head_dim),
+        strides=(stride_q_n, stride_q_d),
+        offsets=(pid_q * BLOCK_SIZE_Q, 0),
+        block_shape=(BLOCK_SIZE_Q, head_dim),
+        order=(1, 0),
+    )
+    q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero")
+    q_start = prefix_len + pid_q * BLOCK_SIZE_Q
+
+    off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len
+    off_k = tl.arange(0, BLOCK_SIZE_K)
+    off_d = tl.arange(0, head_dim)
+    # Block table row for this request.
+    bt_row = block_table_ptr + pid_b * stride_bt_b
+    # Causal window: only blocks up to the last query token's position.
+    hi = min(seq_len, prefix_len + (pid_q + 1) * BLOCK_SIZE_Q)
+    for i in tl.range(0, hi, BLOCK_SIZE_K):
+        blk = i // BLOCK_SIZE_K
+        page = tl.load(bt_row + blk).to(tl.int64)
+        pos = i + off_k
+        # index-K for this page: [BLOCK_SIZE_D, BLOCK_SIZE_K] (transposed)
+        # we don't need masked load for K, because KV cache ensures
+        # allocation is multiple of BLOCK_SIZE_K.
+        # for tokens beyond seqlen, they will be masked in qk later.
+        k = tl.load(
+            ik_cache_ptr
+            + page * stride_ik_blk
+            + off_k[None, :] * stride_ik_pos
+            + off_d[:, None] * stride_ik_d,
+        )
+        qk = tl.dot(q, k) * sm_scale_log2e
+        # apply causal mask as needed
+        if q_start < i + BLOCK_SIZE_K:
+            qk = tl.where(off_q[:, None] >= pos[None, :], qk, float("-inf"))
+        # one sparse block per K-tile -> max over the 128 positions
+        score = tl.max(qk, axis=1)  # [BLOCK_SIZE_Q]
+        s_ptrs = (
+            score_ptr
+            + pid_h * stride_s_h
+            + (seq_start + pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q))
+            * stride_s_n
+            + blk * stride_s_k
+        )
+        q_store_mask = (pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q)) < q_len
+        tl.store(s_ptrs, score, mask=q_store_mask)
+
+
+# ---------------------------------------------------------------------------
+# Top-k selection over per-token block scores (layout-agnostic). block_size_q
+# is 1 for M3, so top-k is computed per query token.
+# ---------------------------------------------------------------------------
+@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])})
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_SIZE_K": 2048}, num_warps=8, num_stages=2),
+        triton.Config({"BLOCK_SIZE_K": 1024}, num_warps=8, num_stages=2),
+        triton.Config({"BLOCK_SIZE_K": 512}, num_warps=8, num_stages=2),
+        triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2),
+        triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2),
+    ],
+    key=["BLOCK_SIZE_T"],
+)
+@triton.jit
+def _topk_index_kernel(
+    s_ptr,  # [num_heads, total_q, max_block]
+    ti_ptr,  # [num_heads, total_q, topk]
+    sample_interval: tl.constexpr,  # block_size_q (1 for M3)
+    block_size: tl.constexpr,  # sparse block size (128)
+    cu_seqlens,
+    cu_seqblocks_q,
+    prefix_lens,
+    topk,
+    init_blocks: tl.constexpr,
+    local_blocks: tl.constexpr,
+    stride_s_h,
+    stride_s_n,
+    stride_s_k,
+    stride_ti_h,
+    stride_ti_n,
+    stride_ti_t,
+    BLOCK_SIZE_K: tl.constexpr,
+    BLOCK_SIZE_T: tl.constexpr,
+    MASK_INIT: tl.constexpr,
+    MASK_LOCAL: tl.constexpr,
+):
+    tl.static_assert(BLOCK_SIZE_K > BLOCK_SIZE_T)
+    pid_q = tl.program_id(0)
+    pid_b = tl.program_id(1)
+    pid_h = tl.program_id(2)
+    seq_start = tl.load(cu_seqlens + pid_b)
+    block_start = tl.load(cu_seqblocks_q + pid_b)
+    block_num = tl.load(cu_seqblocks_q + pid_b + 1) - block_start
+    prefix_len = tl.load(prefix_lens + pid_b)
+    if pid_q >= block_num:
+        return
+    off_k = tl.arange(0, BLOCK_SIZE_K)
+    off_t = tl.arange(0, BLOCK_SIZE_T)
+    s_ptrs = (
+        s_ptr
+        + (seq_start + pid_q * sample_interval) * stride_s_n
+        + pid_h * stride_s_h
+        + off_k * stride_s_k
+    )
+    topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32)
+    topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32)
+    left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2
+    valid_blocks = (prefix_len + pid_q * sample_interval + block_size) // block_size
+    for i in tl.range(0, valid_blocks, BLOCK_SIZE_K):
+        causal_mask = i + off_k < valid_blocks
+        local_mask = i + off_k >= max(0, valid_blocks - local_blocks)
+        init_mask = i + off_k < init_blocks
+        score = tl.load(s_ptrs, mask=causal_mask, other=-1e30).to(tl.float32)
+        score = tl.where(score != score, -1e30, score)
+        s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K
+        if MASK_INIT:
+            score = tl.where(causal_mask & init_mask, score - 1e29, score)
+        else:
+            score = tl.where(causal_mask & init_mask, 1e30, score)
+        if MASK_LOCAL:
+            score = tl.where(causal_mask & local_mask, score - 1e28, score)
+        else:
+            score = tl.where(causal_mask & local_mask, 1e29, score)
+        topk_score, last_topk_score = score, topk_score
+        topk_idx, last_topk_idx = (tl.where(causal_mask, i + off_k + 1, 0), topk_idx)
+        n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
+        for j in tl.static_range(1, n_dims):
+            topk_score, topk_idx = _bitonic_merge(
+                topk_score, topk_idx.to(tl.int32), j, 2, n_dims
+            )
+        if i != 0:
+            topk_score, topk_idx = _bitonic_merge(
+                topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims
+            )
+            topk_score_new = last_topk_score * left_half_mask + topk_score * (
+                1 - left_half_mask
+            )
+            topk_idx_new = last_topk_idx * left_half_mask + topk_idx * (
+                1 - left_half_mask
+            )
+            topk_score, topk_idx = _bitonic_merge(
+                topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims
+            )
+        else:
+            topk_score, topk_idx = _bitonic_merge(
+                topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims
+            )
+    topk_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
+    topk_idx = tl.sum(
+        topk_mask[:, None]
+        * tl.reshape(topk_idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
+        axis=0,
+    )
+    ti_ptrs = (
+        ti_ptr
+        + (block_start + pid_q) * stride_ti_n
+        + pid_h * stride_ti_h
+        + off_t * stride_ti_t
+    )
+    store_mask = off_t < topk
+    valid_mask = off_t < valid_blocks
+    topk_idx = tl.where(store_mask & valid_mask, topk_idx, -1)
+    tl.store(ti_ptrs, topk_idx.to(ti_ptrs.dtype.element_ty), mask=store_mask)
+
+
+# ---------------------------------------------------------------------------
+# Decode index-score kernel (split-K over seq blocks). Decode == one query
+# token per request, so this parallelizes over the KV dimension instead of the
+# query dimension. Chunk counts depend only on shape constants so the grid is
+# fixed within a cuda graph. Base-2 (exp2/log2) softmax matches prefill.
+# ---------------------------------------------------------------------------
+@triton.heuristics(
+    {"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}
+)
+@triton.jit
+def _decode_index_score_kernel(
+    q_ptr,  # idx_q: [total_q (== batch), num_idx_heads, head_dim]
+    ik_cache_ptr,  # index-K cache: [num_blocks, 128, head_dim]
+    score_ptr,  # [num_idx_heads, total_q, max_block]
+    block_table_ptr,  # [num_reqs, max_blocks]
+    seq_lens,  # [batch]
+    num_idx_heads,
+    batch_size,
+    head_dim,
+    init_blocks,
+    local_blocks,
+    sm_scale,
+    stride_q_n,
+    stride_q_h,
+    stride_q_d,
+    stride_ik_blk,
+    stride_ik_pos,
+    stride_ik_d,
+    stride_s_h,
+    stride_s_n,
+    stride_s_k,
+    stride_bt_b,
+    BLOCK_SIZE_K: tl.constexpr,  # == SPARSE_BLOCK_SIZE (128)
+    NUM_KV_CHUNKS: tl.constexpr,
+    BLOCK_SIZE_D: tl.constexpr,
+):
+    sm_scale_log2e = sm_scale * 1.4426950409
+    pid_bc, pid_h = tl.program_id(0), tl.program_id(1)
+    pid_b = pid_bc % batch_size
+    pid_c = pid_bc // batch_size
+    seq_len = tl.load(seq_lens + pid_b)
+    num_blocks = (seq_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
+    # block-aligned fixed-count split: grid independent of seq_len (cuda graph).
+    chunk_size_blocks = (num_blocks + NUM_KV_CHUNKS - 1) // NUM_KV_CHUNKS
+    chunk_start_block = pid_c * chunk_size_blocks
+    chunk_end_block = tl.minimum(chunk_start_block + chunk_size_blocks, num_blocks)
+    if chunk_start_block >= chunk_end_block:
+        return
+    off_k = tl.arange(0, BLOCK_SIZE_K)  # positions within a 128-block
+    off_d = tl.arange(0, BLOCK_SIZE_D)
+    d_mask = off_d < head_dim
+    bt_row = block_table_ptr + pid_b * stride_bt_b
+    # Force-select init (1e30) and local (1e29, higher priority) blocks.
+    local_start = tl.maximum(0, num_blocks - local_blocks)
+    # single query vector for this (token, index head)
+    q = tl.load(
+        q_ptr + pid_b * stride_q_n + pid_h * stride_q_h + off_d * stride_q_d,
+        mask=d_mask,
+        other=0.0,
+    ).to(tl.float32)  # [D]
+    for blk in tl.range(chunk_start_block, chunk_end_block):
+        page = tl.load(bt_row + blk).to(tl.int64)
+        pos = blk * BLOCK_SIZE_K + off_k
+        pos_mask = pos < seq_len
+        k = tl.load(
+            ik_cache_ptr
+            + page * stride_ik_blk
+            + off_k[None, :] * stride_ik_pos
+            + off_d[:, None] * stride_ik_d,
+            mask=d_mask[:, None] & pos_mask[None, :],
+            other=0.0,
+        ).to(tl.float32)  # [D, N]
+        qk = tl.sum(q[:, None] * k, axis=0) * sm_scale_log2e  # [N]
+        qk = tl.where(pos_mask, qk, float("-inf"))
+        score = tl.max(qk, axis=0)  # one score for this 128-block
+        is_init = blk < init_blocks
+        is_local = (blk >= local_start) & (blk < num_blocks)
+        score = tl.where(is_local, 1e29, tl.where(is_init, 1e30, score))
+        tl.store(
+            score_ptr + pid_h * stride_s_h + pid_b * stride_s_n + blk * stride_s_k,
+            score,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Decode top-k (split-K): per-chunk partial top-k + merge. Forced init/local
+# blocks are already encoded in the scores. Ported from the sglang reference.
+# ---------------------------------------------------------------------------
+@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])})
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2),
+        triton.Config({"BLOCK_SIZE_K": 256}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2),
+    ],
+    key=["topk"],
+)
+@triton.jit
+def _topk_index_partial_kernel(
+    s_ptr,  # score: [num_idx_heads, batch, max_block]
+    ts_partial_ptr,  # partial scores out: [NUM_TOPK_CHUNKS, num_idx_heads, batch, T]
+    ti_partial_ptr,  # partial idx out (1-indexed global, 0=invalid): same shape
+    seq_lens,  # [batch]
+    block_size: tl.constexpr,  # sparse block size (128)
+    topk: tl.constexpr,
+    chunk_blocks: tl.constexpr,  # how many score-blocks each chunk owns
+    stride_s_h,
+    stride_s_b,
+    stride_s_k,
+    stride_ts_c,
+    stride_ts_h,
+    stride_ts_b,
+    stride_ts_t,
+    stride_ti_c,
+    stride_ti_h,
+    stride_ti_b,
+    stride_ti_t,
+    BLOCK_SIZE_K: tl.constexpr,
+    BLOCK_SIZE_T: tl.constexpr,
+):
+    tl.static_assert(topk < BLOCK_SIZE_K)
+    pid_b = tl.program_id(0)
+    pid_h = tl.program_id(1)
+    pid_chunk = tl.program_id(2)
+
+    seq_len = tl.load(seq_lens + pid_b)
+    num_blocks = (seq_len + block_size - 1) // block_size
+
+    # Slice this chunk owns within [0, num_blocks).
+    chunk_start = pid_chunk * chunk_blocks
+    chunk_end = tl.minimum(chunk_start + chunk_blocks, num_blocks)
+    chunk_actual = tl.maximum(chunk_end - chunk_start, 0)
+
+    off_k = tl.arange(0, BLOCK_SIZE_K)
+    off_t = tl.arange(0, BLOCK_SIZE_T)
+
+    s_ptrs = (
+        s_ptr
+        + pid_b * stride_s_b
+        + pid_h * stride_s_h
+        + (chunk_start + off_k) * stride_s_k
+    )
+
+    topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32)
+    topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32)
+    left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2
+
+    # Streaming top-K within this chunk. tl.range(0, 0) is a no-op so empty
+    # chunks (chunk_actual == 0) skip the body and store sentinel -1e30 / 0.
+    for i in tl.range(0, chunk_actual, BLOCK_SIZE_K):
+        mask = off_k < chunk_actual - i
+        score = tl.load(s_ptrs, mask=mask, other=-1e30).to(tl.float32)
+        score = tl.where(score != score, -1e30, score)
+        s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K
+        topk_score, last_topk_score = score, topk_score
+        topk_idx, last_topk_idx = (
+            tl.where(mask, chunk_start + i + off_k + 1, 0),  # 1-indexed global
+            topk_idx,
+        )
+        n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
+        for j in tl.static_range(1, n_dims):
+            topk_score, topk_idx = _bitonic_merge(
+                topk_score, topk_idx.to(tl.int32), j, 2, n_dims
+            )
+        if i != 0:
+            topk_score, topk_idx = _bitonic_merge(
+                topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims
+            )
+            topk_score_new = last_topk_score * left_half_mask + topk_score * (
+                1 - left_half_mask
+            )
+            topk_idx_new = last_topk_idx * left_half_mask + topk_idx * (
+                1 - left_half_mask
+            )
+            topk_score, topk_idx = _bitonic_merge(
+                topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims
+            )
+        else:
+            topk_score, topk_idx = _bitonic_merge(
+                topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims
+            )
+
+    # Extract first BLOCK_SIZE_T entries (top-K of this chunk after the sort).
+    topk_mask_extract = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
+    final_score = tl.sum(
+        topk_mask_extract[:, None]
+        * tl.reshape(topk_score, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
+        axis=0,
+    )
+    final_idx = tl.sum(
+        topk_mask_extract[:, None]
+        * tl.reshape(topk_idx, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
+        axis=0,
+    )
+
+    # Always write all BLOCK_SIZE_T slots — invalid slots carry -1e30 / 0
+    # sentinels and lose to real scores in the merge stage.
+    ts_ptrs = (
+        ts_partial_ptr
+        + pid_chunk * stride_ts_c
+        + pid_b * stride_ts_b
+        + pid_h * stride_ts_h
+        + off_t * stride_ts_t
+    )
+    ti_ptrs = (
+        ti_partial_ptr
+        + pid_chunk * stride_ti_c
+        + pid_b * stride_ti_b
+        + pid_h * stride_ti_h
+        + off_t * stride_ti_t
+    )
+    tl.store(ts_ptrs, final_score)
+    tl.store(ti_ptrs, final_idx)
+
+
+@triton.heuristics(
+    {
+        "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"]),
+        "BLOCK_SIZE_K": lambda args: triton.next_power_of_2(
+            args["NUM_TOPK_CHUNKS"] * triton.next_power_of_2(args["topk"])
+        ),
+    }
+)
+@triton.jit
+def _topk_index_merge_kernel(
+    ts_partial_ptr,  # partial scores: [NUM_TOPK_CHUNKS, num_idx_heads, batch, T]
+    ti_partial_ptr,  # partial idx (1-indexed global, 0=invalid): same shape
+    ti_final_ptr,  # final idx (0-indexed, -1=invalid): [num_idx_heads, batch, topk]
+    seq_lens,  # [batch]
+    block_size: tl.constexpr,  # sparse block size (128)
+    topk: tl.constexpr,
+    stride_ts_c,
+    stride_ts_h,
+    stride_ts_b,
+    stride_ts_t,
+    stride_ti_c,
+    stride_ti_h,
+    stride_ti_b,
+    stride_ti_t,
+    stride_tif_h,
+    stride_tif_b,
+    stride_tif_t,
+    NUM_TOPK_CHUNKS: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    BLOCK_SIZE_T: tl.constexpr,
+):
+    pid_b = tl.program_id(0)
+    pid_h = tl.program_id(1)
+
+    seq_len = tl.load(seq_lens + pid_b)
+    num_blocks = (seq_len + block_size - 1) // block_size
+
+    # Load NUM_TOPK_CHUNKS * BLOCK_SIZE_T candidates, padded to BLOCK_SIZE_K.
+    # Candidate at flat position p comes from chunk = p // BLOCK_SIZE_T,
+    # in_chunk = p % BLOCK_SIZE_T.
+    off = tl.arange(0, BLOCK_SIZE_K)
+    chunk_idx = off // BLOCK_SIZE_T
+    in_chunk_idx = off % BLOCK_SIZE_T
+    valid = chunk_idx < NUM_TOPK_CHUNKS
+
+    score_offset = (
+        chunk_idx * stride_ts_c
+        + pid_h * stride_ts_h
+        + pid_b * stride_ts_b
+        + in_chunk_idx * stride_ts_t
+    )
+    idx_offset = (
+        chunk_idx * stride_ti_c
+        + pid_h * stride_ti_h
+        + pid_b * stride_ti_b
+        + in_chunk_idx * stride_ti_t
+    )
+
+    score = tl.load(ts_partial_ptr + score_offset, mask=valid, other=-1e30).to(
+        tl.float32
+    )
+    score = tl.where(score != score, -1e30, score)
+    idx = tl.load(ti_partial_ptr + idx_offset, mask=valid, other=0).to(tl.int32)
+
+    # Full bitonic descending sort of BLOCK_SIZE_K items.
+    n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K)
+    for j in tl.static_range(1, n_dims):
+        score, idx = _bitonic_merge(score, idx.to(tl.int32), j, 2, n_dims)
+    score, idx = _bitonic_merge(score, idx.to(tl.int32), n_dims, True, n_dims)
+
+    # Extract first BLOCK_SIZE_T positions — these are the global top-K.
+    extract_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0
+    topk_idx_final = tl.sum(
+        extract_mask[:, None]
+        * tl.reshape(idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]),
+        axis=0,
+    )
+
+    off_t = tl.arange(0, BLOCK_SIZE_T)
+    tif_ptrs = (
+        ti_final_ptr
+        + pid_h * stride_tif_h
+        + pid_b * stride_tif_b
+        + off_t * stride_tif_t
+    )
+    store_mask = off_t < topk
+    topk_idx_final = tl.where(off_t < tl.minimum(topk, num_blocks), topk_idx_final, -1)
+    tl.store(
+        tif_ptrs, topk_idx_final.to(ti_final_ptr.dtype.element_ty), mask=store_mask
+    )
+
+
+# ---------------------------------------------------------------------------
+# Python wrappers
+# ---------------------------------------------------------------------------
+@torch.no_grad()
+def minimax_m3_index_topk(
+    idx_q: torch.Tensor,  # [total_q, num_idx_heads, head_dim]
+    index_kv_cache: torch.Tensor,  # [num_blocks, 128, head_dim]
+    block_table: torch.Tensor,  # [batch, max_blocks]
+    cu_seqlens_q: torch.Tensor,  # [batch+1] int32
+    seq_lens: torch.Tensor,  # [batch] int32
+    prefix_lens: torch.Tensor,  # [batch] int32
+    max_query_len: int,
+    max_seq_len: int,
+    topk: int,
+    init_blocks: int,
+    local_blocks: int,
+    num_kv_heads: int,
+    sm_scale: float,
+) -> torch.Tensor:
+    """Index block-score + top-k selection. block_size_q == 1 (per-token).
+
+    Returns topk_idx [num_kv_heads, total_q, topk] of 0-indexed block ids
+    (right-padded with -1). M3 has num_idx_heads == num_kv_heads, so the
+    per-index-head top-k maps 1:1 to kv heads (no index-head reduction needed).
+    """
+    total_q, num_idx_heads, head_dim = idx_q.shape
+    assert num_idx_heads == num_kv_heads, (
+        "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)"
+    )
+    batch = cu_seqlens_q.shape[0] - 1
+    max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
+
+    score = torch.empty(
+        (num_idx_heads, total_q, max_block),
+        dtype=torch.float32,
+        device=idx_q.device,
+    )
+    BLOCK_SIZE_Q = 64
+    grid_score = (triton.cdiv(max_query_len, BLOCK_SIZE_Q), batch * num_idx_heads)
+    _index_block_score_kernel[grid_score](
+        idx_q,
+        index_kv_cache,
+        score,
+        block_table,
+        cu_seqlens_q,
+        seq_lens,
+        prefix_lens,
+        num_idx_heads,
+        head_dim,
+        sm_scale,
+        idx_q.stride(0),
+        idx_q.stride(1),
+        idx_q.stride(2),
+        index_kv_cache.stride(0),
+        index_kv_cache.stride(1),
+        index_kv_cache.stride(2),
+        score.stride(0),
+        score.stride(1),
+        score.stride(2),
+        block_table.stride(0),
+        BLOCK_SIZE_Q=BLOCK_SIZE_Q,
+        BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
+    )
+
+    topk_idx = torch.empty(
+        (num_idx_heads, total_q, topk),
+        dtype=torch.int32,
+        device=idx_q.device,
+    )
+    # block_size_q == 1 -> query blocks coincide with query tokens.
+    grid_topk = (max_query_len, batch, num_idx_heads)
+    _topk_index_kernel[grid_topk](
+        score,
+        topk_idx,
+        1,  # sample_interval (block_size_q)
+        SPARSE_BLOCK_SIZE,
+        cu_seqlens_q,
+        cu_seqlens_q,  # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1
+        prefix_lens,
+        topk,
+        init_blocks,
+        local_blocks,
+        score.stride(0),
+        score.stride(1),
+        score.stride(2),
+        topk_idx.stride(0),
+        topk_idx.stride(1),
+        topk_idx.stride(2),
+        MASK_INIT=False,
+        MASK_LOCAL=False,
+    )
+    return topk_idx
+
+
+@torch.no_grad()
+def minimax_m3_index_topk_decode(
+    idx_q: torch.Tensor,  # [batch, num_idx_heads, head_dim]
+    index_kv_cache: torch.Tensor,  # [num_blocks, 128, head_dim]
+    block_table: torch.Tensor,  # [batch, max_blocks]
+    seq_lens: torch.Tensor,  # [batch] int32
+    max_seq_len: int,
+    topk: int,
+    init_blocks: int,
+    local_blocks: int,
+    num_kv_heads: int,
+    sm_scale: float,
+) -> torch.Tensor:
+    """Decode index block-score + top-k, both split-K (cudagraph-safe).
+
+    Returns topk_idx [num_kv_heads, batch, topk] (0-indexed block ids, -1 pad).
+    """
+    total_q, num_idx_heads, head_dim = idx_q.shape
+    assert num_idx_heads == num_kv_heads, (
+        "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)"
+    )
+    batch = total_q
+    max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
+    score = torch.empty(
+        (num_idx_heads, total_q, max_block),
+        dtype=torch.float32,
+        device=idx_q.device,
+    )
+    # split-K over seq blocks; chunk count depends only on shape constants so
+    # the grid is fixed within a cuda graph.
+    TARGET_GRID = 4096
+    MAX_NUM_KV_CHUNKS = 256
+    target = max(
+        1, min(MAX_NUM_KV_CHUNKS, TARGET_GRID // max(1, batch * num_idx_heads))
+    )
+    num_kv_chunks = 1 << (target.bit_length() - 1)
+    grid_score = (batch * num_kv_chunks, num_idx_heads)
+    _decode_index_score_kernel[grid_score](
+        idx_q,
+        index_kv_cache,
+        score,
+        block_table,
+        seq_lens,
+        num_idx_heads,
+        batch,
+        head_dim,
+        init_blocks,
+        local_blocks,
+        sm_scale,
+        idx_q.stride(0),
+        idx_q.stride(1),
+        idx_q.stride(2),
+        index_kv_cache.stride(0),
+        index_kv_cache.stride(1),
+        index_kv_cache.stride(2),
+        score.stride(0),
+        score.stride(1),
+        score.stride(2),
+        block_table.stride(0),
+        BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
+        NUM_KV_CHUNKS=num_kv_chunks,
+    )
+
+    topk_idx = torch.empty(
+        (num_idx_heads, total_q, topk),
+        dtype=torch.int32,
+        device=idx_q.device,
+    )
+    # Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts
+    # pow2(num_topk_chunks * pow2(topk)) candidates.
+    TOPK_TARGET_GRID = 64
+    MAX_NUM_TOPK_CHUNKS = 16
+    topk_target = max(
+        1, min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, batch * num_idx_heads))
+    )
+    num_topk_chunks = 1 << (topk_target.bit_length() - 1)
+    block_size_t = triton.next_power_of_2(topk)
+    chunk_blocks = (max_block + num_topk_chunks - 1) // num_topk_chunks
+    topk_score_partial = torch.empty(
+        num_topk_chunks,
+        num_idx_heads,
+        batch,
+        block_size_t,
+        dtype=torch.float32,
+        device=idx_q.device,
+    )
+    topk_idx_partial = torch.empty(
+        num_topk_chunks,
+        num_idx_heads,
+        batch,
+        block_size_t,
+        dtype=torch.int32,
+        device=idx_q.device,
+    )
+    _topk_index_partial_kernel[(batch, num_idx_heads, num_topk_chunks)](
+        score,
+        topk_score_partial,
+        topk_idx_partial,
+        seq_lens,
+        SPARSE_BLOCK_SIZE,
+        topk,
+        chunk_blocks,
+        score.stride(0),
+        score.stride(1),
+        score.stride(2),
+        topk_score_partial.stride(0),
+        topk_score_partial.stride(1),
+        topk_score_partial.stride(2),
+        topk_score_partial.stride(3),
+        topk_idx_partial.stride(0),
+        topk_idx_partial.stride(1),
+        topk_idx_partial.stride(2),
+        topk_idx_partial.stride(3),
+    )
+    _topk_index_merge_kernel[(batch, num_idx_heads)](
+        topk_score_partial,
+        topk_idx_partial,
+        topk_idx,
+        seq_lens,
+        SPARSE_BLOCK_SIZE,
+        topk,
+        topk_score_partial.stride(0),
+        topk_score_partial.stride(1),
+        topk_score_partial.stride(2),
+        topk_score_partial.stride(3),
+        topk_idx_partial.stride(0),
+        topk_idx_partial.stride(1),
+        topk_idx_partial.stride(2),
+        topk_idx_partial.stride(3),
+        topk_idx.stride(0),
+        topk_idx.stride(1),
+        topk_idx.stride(2),
+        NUM_TOPK_CHUNKS=num_topk_chunks,
+    )
+    return topk_idx

vllm/models/minimax_m3/common/ops/sparse_attn.py

@@ -0,0 +1,510 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Triton kernels for MiniMax M3 block-sparse GQA attention.
+
+The main heads attend only to the blocks selected by the lightning indexer (see
+``index_topk``). Ported from the sglang reference (minimax_sparse_ops), adapted
+to vLLM's paged KV cache: the KV page size is forced to equal the sparse block
+size (128), so one selected block maps to exactly one page.
+
+Main K/V cache layout (vLLM):
+  ``(num_blocks, 2, 128, num_kv_heads, head_dim)``  K=[:,0] V=[:,1]
+
+Only the paths MiniMax M3 uses are implemented: no attention sink, base-2
+(exp2/log2) softmax. The decode kernels use split-K (flash-decoding) over the
+selected blocks with a separate merge step, since one query token per request
+leaves the prefill kernels (which parallelize over the query dim) idle.
+"""
+
+import torch
+
+from vllm.triton_utils import tl, triton
+
+# One sparse block == one KV page.
+SPARSE_BLOCK_SIZE = 128
+
+
+# ---------------------------------------------------------------------------
+# GQA block-sparse attention (paged). Main heads attend only to the selected
+# blocks. BLOCK_SIZE_K == 128 so each selected block is one page.
+# ---------------------------------------------------------------------------
+@triton.heuristics(
+    {
+        "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
+        "BLOCK_SIZE_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]),
+        "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
+        "BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"]
+        * triton.next_power_of_2(args["gqa_group_size"]),
+    }
+)
+@triton.jit
+def _gqa_sparse_fwd_kernel(
+    q_ptr,  # [total_q, num_heads, head_dim]
+    kv_cache_ptr,  # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim]
+    t_ptr,  # topk_idx: [num_kv_heads, total_q, topk]
+    o_ptr,  # [total_q, num_heads, head_dim]
+    block_table_ptr,  # [num_reqs, max_blocks]
+    cu_seqlens_q,
+    cu_seqblocks_q,
+    seq_lens,
+    prefix_lens,
+    num_kv_heads,
+    gqa_group_size,
+    head_dim,
+    max_topk,
+    num_q_loop,
+    sm_scale,
+    stride_qn,
+    stride_qh,
+    stride_qd,
+    stride_kv_blk,
+    stride_kv_kv,
+    stride_kv_pos,
+    stride_kv_h,
+    stride_kv_d,
+    stride_th,
+    stride_tn,
+    stride_tk,
+    stride_on,
+    stride_oh,
+    stride_od,
+    stride_bt_b,
+    BLOCK_SIZE_Q: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,  # == SPARSE_BLOCK_SIZE (128)
+    BLOCK_SIZE_D: tl.constexpr,
+    BLOCK_SIZE_H: tl.constexpr,
+    BLOCK_SIZE_T: tl.constexpr,
+    BLOCK_SIZE_QH: tl.constexpr,
+):
+    sm_scale_log2e = sm_scale * 1.4426950409
+    pid_q = tl.program_id(0)
+    pid_kh = tl.program_id(1)
+    pid_b = tl.program_id(2)
+    pid_h = pid_kh * gqa_group_size
+    q_start = tl.load(cu_seqlens_q + pid_b)
+    q_len = tl.load(cu_seqlens_q + pid_b + 1) - q_start
+    q_block_start = tl.load(cu_seqblocks_q + pid_b)
+    q_block_len = tl.load(cu_seqblocks_q + pid_b + 1) - q_block_start
+    seq_len = tl.load(seq_lens + pid_b)
+    prefix_len = tl.load(prefix_lens + pid_b)
+    if pid_q * num_q_loop >= q_block_len:
+        return
+    real_q_loop = min(num_q_loop, q_block_len - pid_q * num_q_loop)
+    bt_row = block_table_ptr + pid_b * stride_bt_b
+    off_n = tl.arange(0, BLOCK_SIZE_K)
+    off_d = tl.arange(0, BLOCK_SIZE_D)
+    d_mask = off_d < head_dim
+    for j in range(real_q_loop):
+        pid_q_j = pid_q * num_q_loop + j
+        t_ptr_j = t_ptr + (q_block_start + pid_q_j) * stride_tn + pid_kh * stride_th
+        off_t = tl.arange(0, BLOCK_SIZE_T)
+        topk_idx = tl.load(t_ptr_j + off_t * stride_tk, mask=off_t < max_topk, other=-1)
+        real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0)
+        q_ptrs = tl.make_block_ptr(
+            base=q_ptr + q_start * stride_qn + pid_h * stride_qh,
+            shape=(q_len, gqa_group_size, head_dim),
+            strides=(stride_qn, stride_qh, stride_qd),
+            offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
+            block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
+            order=(2, 1, 0),
+        )
+        q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero")
+        off_q = (
+            tl.arange(0, BLOCK_SIZE_Q)[:, None]
+            + pid_q_j * BLOCK_SIZE_Q
+            + prefix_len
+            - tl.arange(0, BLOCK_SIZE_K)[None, :]
+        )
+        m_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
+        lse_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
+        acc_o = tl.zeros((BLOCK_SIZE_QH, BLOCK_SIZE_D), dtype=tl.float32)
+        q = tl.reshape(q, BLOCK_SIZE_QH, BLOCK_SIZE_D)
+        for _ in range(real_topk):
+            blk = tl.load(t_ptr_j).to(tl.int32)
+            t_ptr_j = t_ptr_j + stride_tk
+            c = blk * BLOCK_SIZE_K
+            page = tl.load(bt_row + blk).to(tl.int64)
+            pos = c + off_n
+            pos_mask = pos < seq_len
+            k = tl.load(
+                kv_cache_ptr
+                + page * stride_kv_blk
+                + 0 * stride_kv_kv
+                + off_n[None, :] * stride_kv_pos
+                + pid_kh * stride_kv_h
+                + off_d[:, None] * stride_kv_d,
+                mask=d_mask[:, None] & pos_mask[None, :],
+                other=0.0,
+            )
+            qk = tl.zeros((BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_K), dtype=tl.float32)
+            # causal: q_abs_pos - k_off >= block_start (c)
+            qk += tl.where(off_q[:, None, :] >= c, 0, float("-inf"))
+            qk = tl.reshape(qk, BLOCK_SIZE_QH, BLOCK_SIZE_K)
+            qk += tl.dot(q, k) * sm_scale_log2e
+            qk += tl.where(pos_mask[None, :], 0, float("-inf"))
+            m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
+            p = tl.exp2(qk - m_ij[:, None])
+            l_ij = tl.sum(p, axis=1)
+            acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None]
+            v = tl.load(
+                kv_cache_ptr
+                + page * stride_kv_blk
+                + 1 * stride_kv_kv
+                + off_n[:, None] * stride_kv_pos
+                + pid_kh * stride_kv_h
+                + off_d[None, :] * stride_kv_d,
+                mask=pos_mask[:, None] & d_mask[None, :],
+                other=0.0,
+            )
+            acc_o += tl.dot(p.to(v.dtype), v)
+            m_i = m_ij
+            lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
+        acc_o = acc_o * tl.exp2(m_i - lse_i)[:, None]
+        acc_o = tl.reshape(acc_o, BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D)
+        o_ptrs = tl.make_block_ptr(
+            base=o_ptr + q_start * stride_on + pid_h * stride_oh,
+            shape=(q_len, gqa_group_size, head_dim),
+            strides=(stride_on, stride_oh, stride_od),
+            offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
+            block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
+            order=(2, 1, 0),
+        )
+        tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2))
+
+
+# ---------------------------------------------------------------------------
+# Decode kernels (split-K). Decode == one query token per request, so the
+# prefill kernel (which parallelizes over the query dim) leaves the GPU idle.
+# This instead parallelizes over the selected top-k blocks, producing partials
+# that the merge kernel combines (flash-decoding). All chunk counts depend only
+# on shape constants so the grid is fixed within a cuda graph. Base-2
+# (exp2/log2) softmax matches the prefill kernel.
+# ---------------------------------------------------------------------------
+@triton.heuristics(
+    {
+        "BLOCK_SIZE_H": lambda args: max(
+            16, triton.next_power_of_2(args["gqa_group_size"])
+        ),
+        "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
+        "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
+    }
+)
+@triton.jit
+def _gqa_sparse_decode_kernel(
+    q_ptr,  # [total_q (== batch), num_heads, head_dim]
+    kv_cache_ptr,  # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim]
+    t_ptr,  # topk_idx: [num_kv_heads, batch, topk]
+    o_ptr,  # partial out: [NUM_TOPK_CHUNKS, batch, num_heads, head_dim]
+    lse_ptr,  # partial lse (log2): [NUM_TOPK_CHUNKS, batch, num_heads]
+    block_table_ptr,  # [num_reqs, max_blocks]
+    seq_lens,  # [batch]
+    batch_size,
+    gqa_group_size,
+    head_dim,
+    max_topk,
+    sm_scale,
+    stride_qn,
+    stride_qh,
+    stride_qd,
+    stride_kv_blk,
+    stride_kv_kv,
+    stride_kv_pos,
+    stride_kv_h,
+    stride_kv_d,
+    stride_th,
+    stride_tn,
+    stride_tk,
+    stride_o_c,
+    stride_o_b,
+    stride_o_h,
+    stride_o_d,
+    stride_l_c,
+    stride_l_b,
+    stride_l_h,
+    stride_bt_b,
+    BLOCK_SIZE_K: tl.constexpr,  # == SPARSE_BLOCK_SIZE (128)
+    NUM_TOPK_CHUNKS: tl.constexpr,
+    BLOCK_SIZE_H: tl.constexpr,
+    BLOCK_SIZE_D: tl.constexpr,
+    BLOCK_SIZE_T: tl.constexpr,
+):
+    sm_scale_log2e = sm_scale * 1.4426950409
+    # split-K over the topk dimension: pid(0) folds (batch, chunk) together.
+    pid_bc, pid_kh = tl.program_id(0), tl.program_id(1)
+    pid_b = pid_bc % batch_size
+    pid_c = pid_bc // batch_size
+    pid_h = pid_kh * gqa_group_size
+    chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS
+    chunk_start_topk = pid_c * chunk_size_topk
+    chunk_end_compiletime = chunk_start_topk + chunk_size_topk
+    seq_len = tl.load(seq_lens + pid_b)
+    # number of valid (non-padded) selected blocks for this request
+    off_t = tl.arange(0, BLOCK_SIZE_T)
+    idx_base = t_ptr + pid_kh * stride_th + pid_b * stride_tn
+    topk_idx = tl.load(idx_base + off_t * stride_tk, mask=off_t < max_topk, other=-1)
+    real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0)
+    chunk_end_topk = tl.minimum(chunk_end_compiletime, real_topk)
+
+    off_n = tl.arange(0, BLOCK_SIZE_K)
+    off_d = tl.arange(0, BLOCK_SIZE_D)
+    d_mask = off_d < head_dim
+    bt_row = block_table_ptr + pid_b * stride_bt_b
+
+    m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32)
+    lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32)
+    acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32)
+    q_ptrs = tl.make_block_ptr(
+        base=q_ptr + pid_b * stride_qn + pid_h * stride_qh,
+        shape=(gqa_group_size, head_dim),
+        strides=(stride_qh, stride_qd),
+        offsets=(0, 0),
+        block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D),
+        order=(1, 0),
+    )
+    q = tl.load(q_ptrs, boundary_check=(0, 1), padding_option="zero")
+
+    cur_idx_ptr = idx_base + chunk_start_topk * stride_tk
+    for _ in tl.range(chunk_start_topk, chunk_end_topk):
+        blk = tl.load(cur_idx_ptr).to(tl.int32)
+        cur_idx_ptr = cur_idx_ptr + stride_tk
+        c = blk * BLOCK_SIZE_K
+        page = tl.load(bt_row + blk).to(tl.int64)
+        pos = c + off_n
+        pos_mask = pos < seq_len  # decode query is the last token: attend all valid
+        k = tl.load(
+            kv_cache_ptr
+            + page * stride_kv_blk
+            + 0 * stride_kv_kv
+            + off_n[None, :] * stride_kv_pos
+            + pid_kh * stride_kv_h
+            + off_d[:, None] * stride_kv_d,
+            mask=d_mask[:, None] & pos_mask[None, :],
+            other=0.0,
+        )
+        qk = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_K), dtype=tl.float32)
+        qk += tl.where(pos_mask[None, :], 0, float("-inf"))
+        qk += tl.dot(q, k) * sm_scale_log2e
+        m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
+        p = tl.exp2(qk - m_ij[:, None])
+        l_ij = tl.sum(p, axis=1)
+        acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None]
+        v = tl.load(
+            kv_cache_ptr
+            + page * stride_kv_blk
+            + 1 * stride_kv_kv
+            + off_n[:, None] * stride_kv_pos
+            + pid_kh * stride_kv_h
+            + off_d[None, :] * stride_kv_d,
+            mask=pos_mask[:, None] & d_mask[None, :],
+            other=0.0,
+        )
+        acc_o += tl.dot(p.to(v.dtype), v)
+        m_i = m_ij
+        lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
+    # empty chunks (chunk_start >= real_topk) keep lse_i = -inf -> weight 0 in merge
+    scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), tl.zeros_like(lse_i))
+    acc_o = acc_o * scale[:, None]
+    o_ptrs = tl.make_block_ptr(
+        base=o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h,
+        shape=(gqa_group_size, head_dim),
+        strides=(stride_o_h, stride_o_d),
+        offsets=(0, 0),
+        block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D),
+        order=(1, 0),
+    )
+    tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1))
+    lse_ptrs = tl.make_block_ptr(
+        base=lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h,
+        shape=(gqa_group_size,),
+        strides=(stride_l_h,),
+        offsets=(0,),
+        block_shape=(BLOCK_SIZE_H,),
+        order=(0,),
+    )
+    tl.store(lse_ptrs, lse_i.to(lse_ptr.dtype.element_ty), boundary_check=(0,))
+
+
+@triton.heuristics(
+    {"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}
+)
+@triton.jit
+def _merge_topk_attn_out_kernel(
+    o_ptr,  # partials: [NUM_TOPK_CHUNKS, batch, num_heads, head_dim]
+    lse_ptr,  # partials (log2): [NUM_TOPK_CHUNKS, batch, num_heads]
+    out_ptr,  # merged out: [total_q (== batch), num_heads, head_dim]
+    head_dim,
+    stride_o_c,
+    stride_o_b,
+    stride_o_h,
+    stride_o_d,
+    stride_l_c,
+    stride_l_b,
+    stride_l_h,
+    stride_out_n,
+    stride_out_h,
+    stride_out_d,
+    NUM_TOPK_CHUNKS: tl.constexpr,
+    BLOCK_SIZE_D: tl.constexpr,
+):
+    pid_b, pid_h = tl.program_id(0), tl.program_id(1)
+    off_c = tl.arange(0, NUM_TOPK_CHUNKS)
+    off_d = tl.arange(0, BLOCK_SIZE_D)
+    o_ptrs = tl.make_block_ptr(
+        base=o_ptr + pid_b * stride_o_b + pid_h * stride_o_h,
+        shape=(NUM_TOPK_CHUNKS, head_dim),
+        strides=(stride_o_c, stride_o_d),
+        offsets=(0, 0),
+        block_shape=(NUM_TOPK_CHUNKS, BLOCK_SIZE_D),
+        order=(1, 0),
+    )
+    lse_ptrs = lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c
+    o = tl.load(o_ptrs, boundary_check=(0, 1), padding_option="zero")
+    lse = tl.load(lse_ptrs)  # empty chunks contribute -inf -> weight 0
+    lse_max = tl.max(lse, axis=0)
+    weights = tl.exp2(lse - lse_max)
+    weights = weights / tl.sum(weights, axis=0)
+    o_merged = tl.sum(o * weights[:, None], axis=0)
+    out_ptrs = (
+        out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d
+    )
+    tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim)
+
+
+# ---------------------------------------------------------------------------
+# Python wrappers
+# ---------------------------------------------------------------------------
+@torch.no_grad()
+def minimax_m3_sparse_attn(
+    q: torch.Tensor,  # [total_q, num_heads, head_dim]
+    kv_cache: torch.Tensor,  # [num_blocks, 2, 128, num_kv_heads, head_dim]
+    topk_idx: torch.Tensor,  # [num_kv_heads, total_q, topk]
+    block_table: torch.Tensor,  # [batch, max_blocks]
+    cu_seqlens_q: torch.Tensor,  # [batch+1] int32
+    seq_lens: torch.Tensor,  # [batch] int32
+    prefix_lens: torch.Tensor,  # [batch] int32
+    max_query_len: int,
+    num_kv_heads: int,
+    sm_scale: float,
+    output: torch.Tensor,  # [total_q, num_heads, head_dim]
+) -> None:
+    """GQA block-sparse attention over the selected blocks. block_size_q == 1."""
+    total_q, num_heads, head_dim = q.shape
+    batch = cu_seqlens_q.shape[0] - 1
+    topk = topk_idx.shape[-1]
+    gqa_group_size = num_heads // num_kv_heads
+    grid = (max_query_len, num_kv_heads, batch)
+    _gqa_sparse_fwd_kernel[grid](
+        q,
+        kv_cache,
+        topk_idx,
+        output,
+        block_table,
+        cu_seqlens_q,
+        cu_seqlens_q,  # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1
+        seq_lens,
+        prefix_lens,
+        num_kv_heads,
+        gqa_group_size,
+        head_dim,
+        topk,
+        1,  # num_q_loop
+        sm_scale,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        kv_cache.stride(0),
+        kv_cache.stride(1),
+        kv_cache.stride(2),
+        kv_cache.stride(3),
+        kv_cache.stride(4),
+        topk_idx.stride(0),
+        topk_idx.stride(1),
+        topk_idx.stride(2),
+        output.stride(0),
+        output.stride(1),
+        output.stride(2),
+        block_table.stride(0),
+        BLOCK_SIZE_Q=1,
+        BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
+    )
+
+
+@torch.no_grad()
+def minimax_m3_sparse_attn_decode(
+    q: torch.Tensor,  # [batch, num_heads, head_dim]
+    kv_cache: torch.Tensor,  # [num_blocks, 2, 128, num_kv_heads, head_dim]
+    topk_idx: torch.Tensor,  # [num_kv_heads, batch, topk]
+    block_table: torch.Tensor,  # [batch, max_blocks]
+    seq_lens: torch.Tensor,  # [batch] int32
+    num_kv_heads: int,
+    sm_scale: float,
+    output: torch.Tensor,  # [batch, num_heads, head_dim]
+) -> None:
+    """GQA block-sparse attention for decode (split-K over the top-k blocks)."""
+    batch, num_heads, head_dim = q.shape
+    max_topk = topk_idx.shape[-1]
+    gqa_group_size = num_heads // num_kv_heads
+    # split-K over the selected blocks; chunk count is shape-constant (cuda graph).
+    TARGET_GRID = 256
+    target = max(1, min(max_topk, TARGET_GRID // max(1, batch * num_kv_heads)))
+    num_topk_chunks = 1 << (target.bit_length() - 1)
+    o_partial = torch.empty(
+        num_topk_chunks, batch, num_heads, head_dim, dtype=q.dtype, device=q.device
+    )
+    lse_partial = torch.empty(
+        num_topk_chunks, batch, num_heads, dtype=torch.float32, device=q.device
+    )
+    grid = (batch * num_topk_chunks, num_kv_heads)
+    _gqa_sparse_decode_kernel[grid](
+        q,
+        kv_cache,
+        topk_idx,
+        o_partial,
+        lse_partial,
+        block_table,
+        seq_lens,
+        batch,
+        gqa_group_size,
+        head_dim,
+        max_topk,
+        sm_scale,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        kv_cache.stride(0),
+        kv_cache.stride(1),
+        kv_cache.stride(2),
+        kv_cache.stride(3),
+        kv_cache.stride(4),
+        topk_idx.stride(0),
+        topk_idx.stride(1),
+        topk_idx.stride(2),
+        o_partial.stride(0),
+        o_partial.stride(1),
+        o_partial.stride(2),
+        o_partial.stride(3),
+        lse_partial.stride(0),
+        lse_partial.stride(1),
+        lse_partial.stride(2),
+        block_table.stride(0),
+        BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
+        NUM_TOPK_CHUNKS=num_topk_chunks,
+    )
+    merge_grid = (batch, num_heads)
+    _merge_topk_attn_out_kernel[merge_grid](
+        o_partial,
+        lse_partial,
+        output,
+        head_dim,
+        o_partial.stride(0),
+        o_partial.stride(1),
+        o_partial.stride(2),
+        o_partial.stride(3),
+        lse_partial.stride(0),
+        lse_partial.stride(1),
+        lse_partial.stride(2),
+        output.stride(0),
+        output.stride(1),
+        output.stride(2),
+        NUM_TOPK_CHUNKS=num_topk_chunks,
+    )

vllm/models/minimax_m3/common/sparse_attention.py

@@ -0,0 +1,605 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Attention backend for MiniMax M3 sparse ("lightning indexer") attention.
+
+MiniMax M3 sparse layers run GQA attention restricted to a small set of KV
+blocks chosen by a lightning indexer: index heads score KV blocks, the top-k
+blocks (plus fixed init/local blocks) are selected, and the main attention
+attends only to those blocks. Index keys live in a separate side cache
+(``MiniMaxM3IndexerCache``).
+"""
+
+from dataclasses import dataclass
+from typing import ClassVar
+
+import torch
+from torch import nn
+
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+from vllm.config.cache import CacheDType
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.models.minimax_m3.common.ops.index_topk import (
+    minimax_m3_index_topk,
+    minimax_m3_index_topk_decode,
+)
+from vllm.models.minimax_m3.common.ops.sparse_attn import (
+    SPARSE_BLOCK_SIZE,
+    minimax_m3_sparse_attn,
+    minimax_m3_sparse_attn_decode,
+)
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_cutedsl
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    AttentionImplBase,
+    AttentionLayer,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+    CommonAttentionMetadata,
+    MultipleOf,
+)
+from vllm.v1.attention.backends.utils import (
+    get_kv_cache_layout,
+    split_decodes_and_prefills,
+)
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    KVCacheSpec,
+    MLAAttentionSpec,
+)
+
+logger = init_logger(__name__)
+
+
+class MiniMaxM3SparseBackend(AttentionBackend):
+    """Block-sparse GQA backend for MiniMax M3 sparse attention layers."""
+
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16, torch.float16]
+    # Sparse kernels operate on a bf16 KV cache only.
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["bfloat16"]
+
+    @staticmethod
+    def get_name() -> str:
+        return "MINIMAX_M3_SPARSE"
+
+    @staticmethod
+    def get_impl_cls() -> type["MiniMaxM3SparseImpl"]:
+        return MiniMaxM3SparseImpl
+
+    @staticmethod
+    def get_builder_cls() -> type["MiniMaxM3SparseMetadataBuilder"]:
+        return MiniMaxM3SparseMetadataBuilder
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return [128]
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        # Page size must equal the sparse block size so one sparse block maps
+        # to one KV page (see common.ops.sparse_attn).
+        return [128]
+
+    @classmethod
+    def is_sparse(cls) -> bool:
+        return True
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        return (num_blocks, 2, block_size, num_kv_heads, head_size)
+
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        # `stride_order` indicates the permutation that gets us from
+        # `get_kv_cache_shape` to the actual memory layout we want.
+        if include_num_layers_dimension:
+            # M3 does not use cross-layer (per-layer-stacked) KV blocks for now.
+            raise NotImplementedError
+        cache_layout = get_kv_cache_layout()
+        if cache_layout == "NHD":
+            stride_order = (0, 1, 2, 3, 4)
+        elif cache_layout == "HND":
+            stride_order = (0, 1, 3, 2, 4)
+        else:
+            raise ValueError(f"Unknown cache layout format {cache_layout}.")
+        return stride_order
+
+
+class MiniMaxM3IndexerBackend(MiniMaxM3SparseBackend):
+    """Backend for the lightning-indexer side cache (key-only, single vector).
+
+    Shares the impl/builder/metadata with the main sparse backend but stores a
+    single index-key vector per token, matching the MLAAttentionSpec page.
+    """
+
+    @staticmethod
+    def get_name() -> str:
+        return "MINIMAX_M3_SPARSE_INDEXER"
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        return (num_blocks, block_size, head_size)
+
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        if include_num_layers_dimension:
+            # M3 does not use cross-layer (per-layer-stacked) KV blocks.
+            raise NotImplementedError
+        return (0, 1, 2)
+
+
+class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase):
+    """Side KV cache for the lightning indexer's per-token index keys.
+
+    Stores a single index-key vector per token (no value: M3 disables the index
+    value projection). Modeled on ``DeepseekV32IndexerCache``; registers itself
+    in the static forward context so the KV-cache manager allocates its cache.
+    """
+
+    def __init__(
+        self,
+        head_dim: int,
+        dtype: torch.dtype,
+        prefix: str,
+        cache_config: CacheConfig | None = None,
+    ) -> None:
+        super().__init__()
+        self.kv_cache = torch.tensor([])
+        self.head_dim = head_dim
+        self.dtype = dtype
+        self.prefix = prefix
+        self.cache_config = cache_config
+        compilation_config = get_current_vllm_config().compilation_config
+        if prefix in compilation_config.static_forward_context:
+            raise ValueError(f"Duplicate layer name: {prefix}")
+        compilation_config.static_forward_context[prefix] = self
+
+    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
+        # Key-only cache: one vector per token. MLAAttentionSpec budgets a
+        # single vector; FullAttentionSpec would reserve 2x for K+V.
+        return MLAAttentionSpec(
+            block_size=vllm_config.cache_config.block_size,
+            num_kv_heads=1,
+            head_size=self.head_dim,
+            dtype=self.dtype,
+        )
+
+    def forward(self) -> None: ...
+
+    def get_attn_backend(self) -> type[AttentionBackend]:
+        return MiniMaxM3IndexerBackend
+
+
+@dataclass
+class MiniMaxM3SparsePrefillMetadata:
+    """Per-prefill state for index scoring + block-sparse attention.
+
+    ``cu_seqlens_q``/``context_lens`` are precomputed in the builder so the
+    forward path stays free of host syncs and per-step tensor allocations.
+    """
+
+    cu_seqlens_q: torch.Tensor  # [num_prefills + 1] int32, rebased to 0
+    cu_seqlens_k: torch.Tensor  # [num_prefills + 1] int32, cumulative KV lengths
+    seq_lens: torch.Tensor  # [num_prefills] int32, total KV lengths
+    context_lens: torch.Tensor  # [num_prefills] int32 (cached/context tokens)
+    block_table: torch.Tensor
+    max_query_len: int
+    max_seq_len: int
+    total_kv_blocks: int
+
+
+@dataclass
+class MiniMaxM3SparseDecodeMetadata:
+    """Per-decode state (cudagraph-safe); split-K kernels key only on seq_lens."""
+
+    seq_lens: torch.Tensor  # [num_decodes] int32
+    block_table: torch.Tensor
+
+
+@dataclass
+class MiniMaxM3SparseMetadata(AttentionMetadata):
+    """Sparse-attention metadata, split into prefill and decode sub-metadata."""
+
+    seq_lens: torch.Tensor
+    max_seq_len: int
+    slot_mapping: torch.Tensor
+
+    # Total query tokens in the (decode-first) batch.
+    num_actual_tokens: int
+
+    # Split counts (batch is reordered decode-first).
+    num_decodes: int
+    num_decode_tokens: int
+    num_prefills: int
+    num_prefill_tokens: int
+
+    prefill: MiniMaxM3SparsePrefillMetadata | None = None
+    decode: MiniMaxM3SparseDecodeMetadata | None = None
+
+
+class MiniMaxM3SparseMetadataBuilder(AttentionMetadataBuilder[MiniMaxM3SparseMetadata]):
+    # Full cudagraphs for uniform single-query decode batches.
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+    # The split-K decode kernel doesn't support spec decode yet (it handles one
+    # query token per request only). Keep the threshold at 1 so multi-query
+    # verify batches route to the prefill kernels instead of the decode kernel.
+    reorder_batch_threshold: int = 1
+
+    def __init__(
+        self,
+        kv_cache_spec: AttentionSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+    ) -> None:
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+
+        # Stable per-request context-length buffer for decode cudagraph replays.
+        # Sized to max_num_batched_tokens (>= num_reqs).
+        self.context_len_buffer = torch.empty(
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            dtype=torch.int32,
+            device=device,
+        )
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> MiniMaxM3SparseMetadata:
+        num_reqs = common_attn_metadata.num_reqs
+        num_tokens = common_attn_metadata.num_actual_tokens
+        query_start_loc = common_attn_metadata.query_start_loc
+        seq_lens = common_attn_metadata.seq_lens
+        block_table = common_attn_metadata.block_table_tensor
+
+        num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
+            split_decodes_and_prefills(
+                common_attn_metadata,
+                decode_threshold=self.reorder_batch_threshold,
+            )
+        )
+        assert num_decodes + num_prefills == num_reqs
+        assert num_decode_tokens + num_prefill_tokens == num_tokens
+
+        # Per-request context tokens (seq_len - query_len) into the stable
+        # buffer; batch is decode-first ([:num_decodes] decode, rest prefill).
+        context_lens = self.context_len_buffer[:num_reqs]
+        context_lens.copy_(
+            common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
+        )
+
+        prefill_metadata: MiniMaxM3SparsePrefillMetadata | None = None
+        if num_prefills > 0:
+            seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
+            assert seq_lens_cpu is not None
+            prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:]
+            prefill_total_kv_blocks = (
+                ((prefill_seq_lens_cpu + SPARSE_BLOCK_SIZE - 1) // SPARSE_BLOCK_SIZE)
+                .sum()
+                .item()
+            )
+            prefill_kv_lens = seq_lens[num_decodes:]
+            prefill_cu_seqlens_k = torch.empty(
+                num_prefills + 1, dtype=torch.int32, device=seq_lens.device
+            )
+            prefill_cu_seqlens_k[0] = 0
+            torch.cumsum(prefill_kv_lens, dim=0, out=prefill_cu_seqlens_k[1:])
+            prefill_metadata = MiniMaxM3SparsePrefillMetadata(
+                cu_seqlens_q=(query_start_loc[num_decodes:] - num_decode_tokens).to(
+                    torch.int32
+                ),
+                cu_seqlens_k=prefill_cu_seqlens_k,
+                seq_lens=prefill_kv_lens,
+                context_lens=context_lens[num_decodes:],
+                block_table=block_table[num_decodes:],
+                max_query_len=common_attn_metadata.max_query_len,
+                max_seq_len=common_attn_metadata.max_seq_len,
+                total_kv_blocks=prefill_total_kv_blocks,
+            )
+
+        decode_metadata: MiniMaxM3SparseDecodeMetadata | None = None
+        if num_decodes > 0:
+            decode_metadata = MiniMaxM3SparseDecodeMetadata(
+                seq_lens=seq_lens[:num_decodes],
+                block_table=block_table[:num_decodes],
+            )
+
+        return MiniMaxM3SparseMetadata(
+            seq_lens=seq_lens,
+            max_seq_len=common_attn_metadata.max_seq_len,
+            slot_mapping=common_attn_metadata.slot_mapping,
+            num_actual_tokens=num_tokens,
+            num_decodes=num_decodes,
+            num_decode_tokens=num_decode_tokens,
+            num_prefills=num_prefills,
+            num_prefill_tokens=num_prefill_tokens,
+            prefill=prefill_metadata,
+            decode=decode_metadata,
+        )
+
+
+class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
+    """Block-sparse GQA attention for MiniMax M3 (forward not yet implemented).
+
+    Inherits ``AttentionImplBase`` (not ``AttentionImpl``) because the sparse
+    path needs a custom forward signature: the owning layer pre-inserts K/V and
+    index-K into their caches, so forward takes only the queries.
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int | None = None,
+        kv_cache_dtype: str = "auto",
+        *,
+        topk_blocks: int,
+        sparse_block_size: int,
+        num_index_heads: int,
+        index_head_dim: int,
+        init_blocks: int = 0,
+        local_blocks: int = 0,
+        score_type: str = "max",
+    ) -> None:
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = scale
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        self.kv_cache_dtype = kv_cache_dtype
+
+        # Sparse selection parameters and index-branch dims.
+        self.topk_blocks = topk_blocks
+        self.block_size = sparse_block_size
+        self.init_blocks = init_blocks
+        self.local_blocks = local_blocks
+        self.score_type = score_type
+        self.num_index_heads = num_index_heads
+        self.index_head_dim = index_head_dim
+        can_run_prefill_cutedsl = (
+            current_platform.is_cuda()
+            and current_platform.is_device_capability_family(100)
+            and has_cutedsl()
+            and self.head_size == 128
+            and self.block_size == 128
+            and self.topk_blocks in (4, 8, 16, 32)
+        )
+        self._prefill_gqa_sparse = (
+            self._prefill_gqa_sparse_cutedsl
+            if can_run_prefill_cutedsl
+            else self._prefill_gqa_sparse_triton
+        )
+
+    def _run_prefill(
+        self,
+        q: torch.Tensor,  # [tot, num_heads, head_dim]
+        iq: torch.Tensor,  # [tot, num_idx_heads, head_dim]
+        out: torch.Tensor,  # [tot, num_heads, head_dim]
+        cu_seqlens_q: torch.Tensor,
+        cu_seqlens_k: torch.Tensor,
+        seq_lens: torch.Tensor,
+        context_lens: torch.Tensor,
+        main_block_table: torch.Tensor,
+        index_block_table: torch.Tensor,
+        max_query_len: int,
+        max_seq_len: int,
+        total_kv_blocks: int,
+    ) -> None:
+        # 1. Index block-score + top-k (reads the index-K cache).
+        topk_idx = minimax_m3_index_topk(
+            iq,
+            self._index_kv_cache,
+            index_block_table,
+            cu_seqlens_q,
+            seq_lens,
+            context_lens,
+            max_query_len,
+            max_seq_len,
+            self.topk_blocks,
+            self.init_blocks,
+            self.local_blocks,
+            self.num_kv_heads,
+            self.scale,
+        )
+        # 2. GQA block-sparse attention over the selected blocks (main cache).
+        self._prefill_gqa_sparse(
+            q,
+            out,
+            topk_idx,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seq_lens,
+            context_lens,
+            main_block_table,
+            max_query_len,
+            max_seq_len,
+            total_kv_blocks,
+        )
+
+    def _prefill_gqa_sparse_triton(
+        self,
+        q: torch.Tensor,
+        out: torch.Tensor,
+        topk_idx: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        cu_seqlens_k: torch.Tensor,
+        seq_lens: torch.Tensor,
+        context_lens: torch.Tensor,
+        main_block_table: torch.Tensor,
+        max_query_len: int,
+        max_seq_len: int,
+        total_kv_blocks: int,
+    ) -> None:
+        minimax_m3_sparse_attn(
+            q,
+            self._kv_cache,
+            topk_idx,
+            main_block_table,
+            cu_seqlens_q,
+            seq_lens,
+            context_lens,
+            max_query_len,
+            self.num_kv_heads,
+            self.scale,
+            out,
+        )
+
+    def _prefill_gqa_sparse_cutedsl(
+        self,
+        q: torch.Tensor,
+        out: torch.Tensor,
+        topk_idx: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        cu_seqlens_k: torch.Tensor,
+        seq_lens: torch.Tensor,
+        context_lens: torch.Tensor,
+        main_block_table: torch.Tensor,
+        max_query_len: int,
+        max_seq_len: int,
+        total_kv_blocks: int,
+    ) -> None:
+        from vllm.models.minimax_m3.nvidia.ops.prefill_gqa_sparse import (
+            minimax_m3_sparse_attn_cutedsl,
+        )
+
+        minimax_m3_sparse_attn_cutedsl(
+            q,
+            self._kv_cache,
+            topk_idx,
+            main_block_table,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seq_lens,
+            max_query_len,
+            max_seq_len,
+            self.num_kv_heads,
+            self.scale,
+            out,
+            total_kv_blocks=total_kv_blocks,
+        )
+
+    def _run_decode(
+        self,
+        q: torch.Tensor,  # [batch, num_heads, head_dim]
+        iq: torch.Tensor,  # [batch, num_idx_heads, head_dim]
+        out: torch.Tensor,  # [batch, num_heads, head_dim]
+        seq_lens: torch.Tensor,
+        main_block_table: torch.Tensor,
+        index_block_table: torch.Tensor,
+        max_seq_len: int,
+    ) -> None:
+        # Split-K decode kernels (parallelize over KV; one query token/request).
+        # 1. Index block-score + top-k.
+        topk_idx = minimax_m3_index_topk_decode(
+            iq,
+            self._index_kv_cache,
+            index_block_table,
+            seq_lens,
+            max_seq_len,
+            self.topk_blocks,
+            self.init_blocks,
+            self.local_blocks,
+            self.num_kv_heads,
+            self.scale,
+        )
+        # 2. GQA block-sparse attention (split-K over the selected blocks).
+        minimax_m3_sparse_attn_decode(
+            q,
+            self._kv_cache,
+            topk_idx,
+            main_block_table,
+            seq_lens,
+            self.num_kv_heads,
+            self.scale,
+            out,
+        )
+
+    @eager_break_during_capture
+    def forward(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        index_query: torch.Tensor,
+        kv_cache: torch.Tensor,
+        index_kv_cache: torch.Tensor,
+        output: torch.Tensor,
+    ) -> torch.Tensor:
+        # K/V and index-K are pre-inserted into their caches; per-request
+        # metadata comes from the forward context. Only queries are passed in.
+        attn_metadata = get_forward_context().attn_metadata
+        if not isinstance(attn_metadata, dict):
+            # Profiling run: caches unbound, output left as-is.
+            return output
+        main_md = attn_metadata[layer.layer_name]  # type: ignore[attr-defined]
+        index_md = attn_metadata[layer.index_cache.prefix]  # type: ignore[attr-defined]
+        assert isinstance(main_md, MiniMaxM3SparseMetadata)
+        assert isinstance(index_md, MiniMaxM3SparseMetadata)
+
+        nd = main_md.num_decode_tokens
+        num_tokens = main_md.num_actual_tokens
+        hd = self.head_size
+        q = query[:num_tokens].view(-1, self.num_heads, hd)
+        iq = index_query[:num_tokens].view(
+            -1, self.num_index_heads, self.index_head_dim
+        )
+        out = output[:num_tokens].view(-1, self.num_heads, hd)
+        # Stash caches for _run_phase (avoid threading through every call).
+        self._kv_cache = kv_cache
+        self._index_kv_cache = index_kv_cache
+
+        # Decode slice [:nd]: each token is a 1-token "prefill" at seq_len-1.
+        # All kernel args are precomputed in the builder (cudagraph-safe).
+        if main_md.num_decodes > 0:
+            d, idx_d = main_md.decode, index_md.decode
+            assert d is not None and idx_d is not None
+            self._run_decode(
+                q[:nd],
+                iq[:nd],
+                out[:nd],
+                d.seq_lens,
+                d.block_table,
+                idx_d.block_table,
+                main_md.max_seq_len,
+            )
+
+        # Prefill slice [nd:]: cu_seqlens_q already rebased to 0.
+        if main_md.num_prefills > 0:
+            p, idx_p = main_md.prefill, index_md.prefill
+            assert p is not None and idx_p is not None
+            self._run_prefill(
+                q[nd:],
+                iq[nd:],
+                out[nd:],
+                p.cu_seqlens_q,
+                p.cu_seqlens_k,
+                p.seq_lens,
+                p.context_lens,
+                p.block_table,
+                idx_p.block_table,
+                p.max_query_len,
+                p.max_seq_len,
+                p.total_kv_blocks,
+            )
+        return output

vllm/models/minimax_m3/common/vision_tower.py

@@ -0,0 +1,695 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterable
+
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange
+from transformers import PretrainedConfig
+
+from vllm.distributed import parallel_state
+from vllm.distributed import utils as dist_utils
+from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.attention.mm_encoder_attention import (
+    MMEncoderAttention,
+)
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.utils import maybe_prefix
+from vllm.model_executor.models.vision import (
+    get_vit_attn_backend,
+    is_vit_use_data_parallel,
+)
+
+
+class MiniMaxVLPatchEmbed(nn.Module):
+    """Conv3d-based patch embedding.
+
+    Takes flat tokens of shape (N, C * temporal_patch_size * patch_size²)
+    and projects each to a hidden-size embedding.
+    """
+
+    def __init__(self, config: PretrainedConfig) -> None:
+        super().__init__()
+        compression = config.img_token_compression_config
+        temporal_patch_size = compression.get("temporal_patch_size", 2)
+        patch_size = config.patch_size
+        num_channels = config.num_channels
+
+        self.patch_size = patch_size
+        self.temporal_patch_size = temporal_patch_size
+        self.num_channels = num_channels
+        self.hidden_size = config.hidden_size
+
+        self.patch_embedding = nn.Conv3d(
+            in_channels=num_channels,
+            out_channels=config.hidden_size,
+            kernel_size=(temporal_patch_size, patch_size, patch_size),
+            stride=(temporal_patch_size, patch_size, patch_size),
+            bias=False,
+        )
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        # pixel_values: (N, C * temporal_patch_size * patch_size²)
+        if self.patch_embedding.weight.dtype != pixel_values.dtype:
+            self.patch_embedding = self.patch_embedding.to(pixel_values.dtype)
+        x = pixel_values.reshape(
+            pixel_values.shape[0],
+            self.num_channels,
+            self.temporal_patch_size,
+            self.patch_size,
+            self.patch_size,
+        )
+        return self.patch_embedding(x).reshape(x.shape[0], -1)
+
+
+class MiniMaxVLAttention(nn.Module):
+    """Multi-head attention with MiniMax's partial 3D RoPE.
+
+    Partial means only the first ``rot_dim`` (< head_dim) dimensions of
+    Q and K are rotated; the remaining dims are passed through unchanged.
+    """
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        use_data_parallel = is_vit_use_data_parallel()
+        self.tp_size = (
+            1
+            if use_data_parallel
+            else parallel_state.get_tensor_model_parallel_world_size()
+        )
+        self.head_dim = embed_dim // num_heads
+        self.num_heads_per_partition = dist_utils.divide(num_heads, self.tp_size)
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=embed_dim,
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            total_num_kv_heads=num_heads,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+            disable_tp=use_data_parallel,
+        )
+        self.out_proj = RowParallelLinear(
+            input_size=embed_dim,
+            output_size=embed_dim,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.out_proj",
+            disable_tp=use_data_parallel,
+        )
+        self.attn = MMEncoderAttention(
+            num_heads=self.num_heads_per_partition,
+            head_size=self.head_dim,
+            prefix=f"{prefix}.attn",
+        )
+        # ApplyRotaryEmb handles the internal cos/sin repeat and partial
+        # rotation (ro_dim = half_rot_dim * 2 < head_dim for MiniMax).
+        # enable_fp32_compute=True runs the rotation in fp32 (q/k upcast,
+        # fp32 cos/sin), matching the SGLang reference's _minimax_rope_applier.
+        self.apply_rotary_emb = ApplyRotaryEmb(
+            enforce_enable=True, enable_fp32_compute=True
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        rotary_cos: torch.Tensor,
+        rotary_sin: torch.Tensor,
+        max_seqlen: torch.Tensor,
+        sequence_lengths: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # x: (N, 1, embed_dim)  [seq=N, batch=1, chan=embed_dim]
+        x_qkv, _ = self.qkv_proj(x)  # (N, 1, 3 * heads_per_part * head_dim)
+        seq_len, batch_size, _ = x_qkv.shape
+
+        # Rearrange to (b=1, N, 3, heads, head_dim) — same as Qwen2_5_VisionAttention
+        qkv = rearrange(
+            x_qkv,
+            "s b (three head d) -> b s three head d",
+            three=3,
+            head=self.num_heads_per_partition,
+        )
+        qk, v = qkv[:, :, :2], qkv[:, :, 2]  # (b,N,2,h,d) and (b,N,h,d)
+
+        # Stack q/k → (2*b, N, heads, head_dim) for joint RoPE application.
+        # rotary_cos/sin: (N, half_rot_dim) — ApplyRotaryEmb expands internally
+        # and rotates only the first 2*half_rot_dim dims, passing the rest through.
+        qk_reshaped = rearrange(qk, "b s two h d -> (two b) s h d", two=2).contiguous()
+        qk_rotated = self.apply_rotary_emb(qk_reshaped, rotary_cos, rotary_sin)
+        qk_rotated = qk_rotated.view(
+            2, batch_size, seq_len, self.num_heads_per_partition, self.head_dim
+        )
+        q, k = qk_rotated.unbind(dim=0)  # each (b=1, N, heads, head_dim)
+
+        # Flash attention → (b, N, heads, head_dim)
+        context = self.attn(
+            query=q,
+            key=k,
+            value=v,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            sequence_lengths=sequence_lengths,
+        )
+
+        # Back to (N, 1, embed_dim)
+        context = rearrange(context, "b s h d -> s b (h d)", b=batch_size)
+        output, _ = self.out_proj(context)
+        return output
+
+
+class MiniMaxVLEncoderLayer(nn.Module):
+    """Single CLIP-style transformer block."""
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        embed_dim = config.hidden_size
+        self.layer_norm1 = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+        self.self_attn = MiniMaxVLAttention(
+            embed_dim=embed_dim,
+            num_heads=config.num_attention_heads,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.layer_norm2 = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+        use_data_parallel = is_vit_use_data_parallel()
+        self.fc1 = ColumnParallelLinear(
+            config.hidden_size,
+            config.intermediate_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc1",
+            disable_tp=use_data_parallel,
+        )
+        self.act = get_act_fn(getattr(config, "hidden_act", "gelu"))
+        self.fc2 = RowParallelLinear(
+            config.intermediate_size,
+            config.hidden_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc2",
+            disable_tp=use_data_parallel,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        rotary_cos: torch.Tensor,
+        rotary_sin: torch.Tensor,
+        max_seqlen: torch.Tensor,
+        sequence_lengths: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # x: (N, 1, hidden_size)
+        x = x + self.self_attn(
+            self.layer_norm1(x),
+            cu_seqlens,
+            rotary_cos,
+            rotary_sin,
+            max_seqlen,
+            sequence_lengths,
+        )
+        residual = x
+        x, _ = self.fc1(self.layer_norm2(x))
+        x = self.act(x)
+        x, _ = self.fc2(x)
+        return residual + x
+
+
+class MiniMaxVLEncoder(nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        num_hidden_layers_override: int | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        n = (
+            config.num_hidden_layers
+            if num_hidden_layers_override is None
+            else num_hidden_layers_override
+        )
+        self.layers = nn.ModuleList(
+            [
+                MiniMaxVLEncoderLayer(
+                    config=config,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.layers.{i}",
+                )
+                for i in range(n)
+            ]
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        rotary_cos: torch.Tensor,
+        rotary_sin: torch.Tensor,
+        max_seqlen: torch.Tensor,
+        sequence_lengths: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        for layer in self.layers:
+            x = layer(
+                x, cu_seqlens, rotary_cos, rotary_sin, max_seqlen, sequence_lengths
+            )
+        return x
+
+
+class MiniMaxVLVisionTransformer(nn.Module):
+    """CLIP-based ViT with 3D RoPE (t/h/w decomposed).
+
+    Faithfully mirrors SGLang's ``MiniMaxVLVisionTransformer``.
+    FLASHINFER backend is not supported; standard flash-attn is used.
+    """
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        num_hidden_layers_override: int | None = None,
+        require_post_norm: bool | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        compression = config.img_token_compression_config
+        self.spatial_merge_size: int = compression.get("spatial_merge_size", 2)
+        self.temporal_patch_size: int = compression.get("temporal_patch_size", 2)
+        self.vision_segment_max_frames: int | None = getattr(
+            config, "vision_segment_max_frames", None
+        )
+        self.use_data_parallel = is_vit_use_data_parallel()
+
+        embed_dim = config.hidden_size
+        head_dim = embed_dim // config.num_attention_heads
+        # Backend selection + sharding info for building encoder metadata.
+        # Defaults to FLASH_ATTN on SM80+; --mm-encoder-attn-backend FLASHINFER
+        # selects the cuDNN ViT prefill path.
+        self.hidden_size = embed_dim
+        self.tp_size = (
+            1
+            if self.use_data_parallel
+            else parallel_state.get_tensor_model_parallel_world_size()
+        )
+        self.attn_backend = get_vit_attn_backend(
+            head_size=head_dim, dtype=torch.get_default_dtype()
+        )
+        rope_dims = 2 * (head_dim // 2)
+
+        # Split rope dims evenly across t/h/w (same formula as SGLang)
+        self.t_dim = int(2 * ((rope_dims // 3) // 2))
+        self.h_dim = int(2 * ((rope_dims // 3) // 2))
+        self.w_dim = int(2 * ((rope_dims // 3) // 2))
+        # rot_dim = t_dim + h_dim + w_dim (may be < head_dim)
+
+        rope_theta: float = getattr(config, "rope_theta", 10000.0)
+        inv_freq_t = 1.0 / (
+            rope_theta
+            ** (torch.arange(0, self.t_dim, 2, dtype=torch.float32) / self.t_dim)
+        )
+        inv_freq_h = 1.0 / (
+            rope_theta
+            ** (torch.arange(0, self.h_dim, 2, dtype=torch.float32) / self.h_dim)
+        )
+        inv_freq_w = 1.0 / (
+            rope_theta
+            ** (torch.arange(0, self.w_dim, 2, dtype=torch.float32) / self.w_dim)
+        )
+        self.register_buffer("inv_freq_t", inv_freq_t, persistent=False)
+        self.register_buffer("inv_freq_h", inv_freq_h, persistent=False)
+        self.register_buffer("inv_freq_w", inv_freq_w, persistent=False)
+
+        self.embeddings = MiniMaxVLPatchEmbed(config)
+        self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+        n_layers = config.num_hidden_layers
+        if num_hidden_layers_override is None:
+            num_hidden_layers_override = n_layers
+        self.encoder = MiniMaxVLEncoder(
+            config=config,
+            num_hidden_layers_override=num_hidden_layers_override,
+            quant_config=quant_config,
+            prefix=f"{prefix}.encoder",
+        )
+
+        if require_post_norm is None:
+            require_post_norm = num_hidden_layers_override == n_layers
+        self.post_layernorm = (
+            nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+            if require_post_norm
+            else None
+        )
+
+        # out_hidden_size needed by run_dp_sharded_mrope_vision_model
+        self.out_hidden_size = embed_dim
+
+    # ── RoPE helpers ─────────────────────────────────────────────────────
+
+    def _get_3d_rope_embed(
+        self, grid_t: int, grid_h: int, grid_w: int, spatial_merge_size: int
+    ) -> torch.Tensor:
+        """Compute 3D RoPE frequencies for a single (T, H, W) grid.
+
+        Returns (T*H*W, half_rot_dim) on the same device as inv_freq buffers.
+        Mirrors SGLang's ``_get_3d_rope_embed`` exactly.
+        """
+        tokens_per_frame = grid_h * grid_w
+
+        tpos_ids = (
+            torch.arange(grid_t, device=self.inv_freq_t.device)
+            .unsqueeze(1)
+            .expand(-1, tokens_per_frame)
+            .flatten()
+        )
+
+        hpos_ids = (
+            torch.arange(grid_h, device=self.inv_freq_h.device)
+            .unsqueeze(1)
+            .expand(-1, grid_w)
+            .reshape(
+                grid_h // spatial_merge_size,
+                spatial_merge_size,
+                grid_w // spatial_merge_size,
+                spatial_merge_size,
+            )
+            .permute(0, 2, 1, 3)
+            .unsqueeze(0)
+            .expand(grid_t, -1, -1, -1, -1)
+            .flatten()
+        )
+        wpos_ids = (
+            torch.arange(grid_w, device=self.inv_freq_w.device)
+            .unsqueeze(0)
+            .expand(grid_h, -1)
+            .reshape(
+                grid_h // spatial_merge_size,
+                spatial_merge_size,
+                grid_w // spatial_merge_size,
+                spatial_merge_size,
+            )
+            .permute(0, 2, 1, 3)
+            .unsqueeze(0)
+            .expand(grid_t, -1, -1, -1, -1)
+            .flatten()
+        )
+
+        max_t = max(grid_t, 1)
+        max_hw = max(grid_h, grid_w)
+
+        seq_t = torch.arange(
+            max_t, device=self.inv_freq_t.device, dtype=self.inv_freq_t.dtype
+        )
+        seq_hw = torch.arange(
+            max_hw, device=self.inv_freq_h.device, dtype=self.inv_freq_h.dtype
+        )
+
+        freqs_t = torch.outer(seq_t, self.inv_freq_t)  # (max_t, t_dim/2)
+        freqs_h = torch.outer(seq_hw, self.inv_freq_h)  # (max_hw, h_dim/2)
+        freqs_w = torch.outer(seq_hw, self.inv_freq_w)  # (max_hw, w_dim/2)
+
+        return torch.cat(
+            [freqs_t[tpos_ids], freqs_h[hpos_ids], freqs_w[wpos_ids]], dim=-1
+        )  # (T*H*W, half_rot_dim)
+
+    def _get_rope_embed_3d(
+        self, grid_thw: list[list[int]], spatial_merge_size: int
+    ) -> torch.Tensor:
+        embeds = [
+            self._get_3d_rope_embed(t, h, w, spatial_merge_size) for t, h, w in grid_thw
+        ]
+        return torch.cat(embeds, dim=0)  # (total_N, half_rot_dim)
+
+    # ── Frame-limit helper (mirrors SGLang) ──────────────────────────────
+
+    def _apply_max_frames_limit(self, grid_thw: list[list[int]]) -> list[list[int]]:
+        if self.vision_segment_max_frames is None:
+            return grid_thw
+        max_f = self.vision_segment_max_frames
+        out: list[list[int]] = []
+        for t, h, w in grid_thw:
+            if t <= max_f:
+                out.append([t, h, w])
+            else:
+                for i in range(0, t, max_f):
+                    out.append([min(max_f, t - i), h, w])
+        return out
+
+    # ── Forward ──────────────────────────────────────────────────────────
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        grid_thw: list[list[int]],
+    ) -> torch.Tensor:
+        # pixel_values: (total_N, C * temporal_patch_size * patch_size²)
+        # Output:       (total_N, hidden_size)
+
+        hidden = self.embeddings(pixel_values)  # (total_N, hidden_size)
+        hidden = self.pre_layrnorm(hidden)
+
+        limited = self._apply_max_frames_limit(grid_thw)
+
+        # Token-level cumulative sequence lengths (one segment per limited grid).
+        lens = [t * h * w for t, h, w in limited]
+        cu_seqlens_np = np.zeros(len(lens) + 1, dtype=np.int32)
+        np.cumsum(np.array(lens, dtype=np.int32), out=cu_seqlens_np[1:])
+
+        # Backend-specific encoder metadata. For FLASH_ATTN this returns the raw
+        # token cu_seqlens, the max segment length, and sequence_lengths=None;
+        # for FLASHINFER (cuDNN) it repacks cu_seqlens into element-offset
+        # indptrs, buckets max_seqlen, and builds padded per-sequence lengths.
+        sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
+            self.attn_backend, cu_seqlens_np, hidden.device
+        )
+        max_seqlen = torch.tensor(
+            MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens_np),
+            dtype=torch.int32,
+        )
+        cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
+            self.attn_backend,
+            cu_seqlens_np,
+            self.hidden_size,
+            self.tp_size,
+            hidden.device,
+        )
+
+        # 3D RoPE: (total_N, half_rot_dim); ApplyRotaryEmb expands internally
+        freqs = self._get_rope_embed_3d(limited, self.spatial_merge_size)
+        freqs = freqs.to(device=hidden.device)
+        # Keep cos/sin in fp32; ApplyRotaryEmb(enable_fp32_compute=True) runs the
+        # rotation in fp32 to match the SGLang reference precision.
+        rotary_cos, rotary_sin = freqs.cos(), freqs.sin()
+
+        # Encoder expects (N, 1, hidden_size) — add batch dim
+        hidden = hidden.unsqueeze(1)
+        hidden = self.encoder(
+            hidden,
+            cu_seqlens,
+            rotary_cos,
+            rotary_sin,
+            max_seqlen,
+            sequence_lengths=sequence_lengths,
+        )
+        hidden = hidden.squeeze(1)  # back to (total_N, hidden_size)
+
+        if self.post_layernorm is not None:
+            hidden = self.post_layernorm(hidden)
+
+        return hidden
+
+
+class MiniMaxVLMultiModalProjector(nn.Module):
+    """Two-layer MLP projector: vision_hidden → text_hidden."""
+
+    def __init__(
+        self,
+        vision_hidden_size: int,
+        text_hidden_size: int,
+        projector_hidden_size: int | None,
+        multimodal_projector_bias: bool,
+        projector_hidden_act: str = "gelu",
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        mid = projector_hidden_size if projector_hidden_size else text_hidden_size
+        use_dp = is_vit_use_data_parallel()
+        self.linear_1 = ColumnParallelLinear(
+            vision_hidden_size,
+            mid,
+            bias=multimodal_projector_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.linear_1",
+            disable_tp=use_dp,
+        )
+        self.act = get_act_fn(projector_hidden_act)
+        self.linear_2 = RowParallelLinear(
+            mid,
+            text_hidden_size,
+            bias=multimodal_projector_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.linear_2",
+            disable_tp=use_dp,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.linear_1(x)
+        x = self.act(x)
+        x, _ = self.linear_2(x)
+        return x
+
+
+class MiniMaxVLPatchMerger(nn.Module):
+    def __init__(
+        self,
+        spatial_merge_size: int,
+        text_hidden_size: int,
+        projector_hidden_size: int | None,
+        patch_merge_bias: bool,
+        projector_hidden_act: str = "gelu",
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.spatial_merge_size = spatial_merge_size
+        mid = projector_hidden_size if projector_hidden_size else text_hidden_size
+        merge_in = text_hidden_size * spatial_merge_size**2
+        use_dp = is_vit_use_data_parallel()
+        self.linear_1 = ColumnParallelLinear(
+            merge_in,
+            mid,
+            bias=patch_merge_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.linear_1",
+            disable_tp=use_dp,
+        )
+        self.act = get_act_fn(projector_hidden_act)
+        self.linear_2 = RowParallelLinear(
+            mid,
+            text_hidden_size,
+            bias=patch_merge_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.linear_2",
+            disable_tp=use_dp,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (N, text_hidden_size) → (N // merge_size², text_hidden_size)
+        x = x.reshape(x.shape[0] // (self.spatial_merge_size**2), -1)
+        x, _ = self.linear_1(x)
+        x = self.act(x)
+        x, _ = self.linear_2(x)
+        return x
+
+
+class MiniMaxVLVisionModel(nn.Module):
+    """Full vision model: ViT → projector → patch merger."""
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        text_hidden_size: int,
+        projector_hidden_size: int | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        compression = config.img_token_compression_config
+        spatial_merge_size: int = compression.get("spatial_merge_size", 2)
+        self.spatial_merge_size = spatial_merge_size
+        self.use_data_parallel = is_vit_use_data_parallel()
+
+        # The released checkpoint ships no ``post_layernorm`` weights and
+        # uses ``vision_feature_layer=-1`` with ``vision_feature_select_strategy
+        # ="full"``, i.e. the raw last encoder hidden state (CLIP's
+        # ``last_hidden_state`` is taken before the post layernorm). Applying an
+        # untrained post layernorm here would corrupt the visual features.
+        self.vision_model = MiniMaxVLVisionTransformer(
+            config=config,
+            require_post_norm=False,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "vision_model"),
+        )
+        self.multi_modal_projector = MiniMaxVLMultiModalProjector(
+            vision_hidden_size=config.hidden_size,
+            text_hidden_size=text_hidden_size,
+            projector_hidden_size=projector_hidden_size,
+            multimodal_projector_bias=getattr(
+                config, "multimodal_projector_bias", True
+            ),
+            projector_hidden_act=getattr(config, "projector_hidden_act", "gelu"),
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "multi_modal_projector"),
+        )
+        self.patch_merge_mlp = MiniMaxVLPatchMerger(
+            spatial_merge_size=spatial_merge_size,
+            text_hidden_size=text_hidden_size,
+            projector_hidden_size=projector_hidden_size,
+            patch_merge_bias=getattr(config, "patch_merge_bias", True),
+            projector_hidden_act=getattr(config, "projector_hidden_act", "gelu"),
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "patch_merge_mlp"),
+        )
+
+        self.dtype = self.vision_model.embeddings.patch_embedding.weight.dtype
+        self.out_hidden_size = text_hidden_size
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        grid_thw: list[list[int]],
+    ) -> torch.Tensor:
+        hidden = self.vision_model(pixel_values=pixel_values, grid_thw=grid_thw)
+        if hidden.dim() == 3:
+            hidden = hidden.squeeze(0)
+        hidden = self.multi_modal_projector(hidden)
+        hidden = self.patch_merge_mlp(hidden)
+        return hidden
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj.", "q_proj.", "q"),
+            ("qkv_proj.", "k_proj.", "k"),
+            ("qkv_proj.", "v_proj.", "v"),
+        ]
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params

vllm/models/minimax_m3/nvidia/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/minimax_m3/nvidia/ops/prefill_gqa_sparse/__init__.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax M3 CuteDSL sparse prefill attention wrapper."""
+
+import torch
+
+from .interface import sparse_atten_func
+from .sm100.prepare_k2q_csr import build_k2q_csr_with_schedule_sm100
+
+SPARSE_BLOCK_SIZE = 128
+
+
+def minimax_m3_sparse_attn_cutedsl(
+    q: torch.Tensor,
+    kv_cache: torch.Tensor,
+    topk_idx: torch.Tensor,
+    block_table: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    seq_lens: torch.Tensor,
+    max_query_len: int,
+    max_seq_len: int,
+    num_kv_heads: int,
+    sm_scale: float,
+    out: torch.Tensor,
+    *,
+    total_kv_blocks: int,
+) -> None:
+    """Run CuteDSL sparse attention directly over vLLM's paged KV cache."""
+    if kv_cache.shape[2] != SPARSE_BLOCK_SIZE:
+        raise ValueError("MiniMax M3 CuteDSL path requires block_size == 128")
+    topk = topk_idx.shape[-1]
+    k_cache = kv_cache[:, 0]
+    v_cache = kv_cache[:, 1]
+    k2q_row_ptr, k2q_q_indices, schedule = build_k2q_csr_with_schedule_sm100(
+        topk_idx,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        blk_kv=SPARSE_BLOCK_SIZE,
+        max_seqlen_k=max_seq_len,
+        max_seqlen_q=max_query_len,
+        total_rows=total_kv_blocks,
+        qhead_per_kv=q.shape[1] // num_kv_heads,
+    )
+    sparse_atten_func(
+        q,
+        k_cache,
+        v_cache,
+        k2q_row_ptr,
+        k2q_q_indices,
+        topK=topk,
+        blk_kv=SPARSE_BLOCK_SIZE,
+        causal=True,
+        softmax_scale=sm_scale,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_query_len,
+        page_table=block_table,
+        seqused_k=seq_lens,
+        schedule=schedule,
+        out=out,
+    )

vllm/models/minimax_m3/nvidia/ops/prefill_gqa_sparse/common/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/minimax_m3/nvidia/ops/prefill_gqa_sparse/common/cute_dsl_utils.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa
+# mypy: ignore-errors
+import os
+import pathlib
+from dataclasses import dataclass, fields
+from functools import lru_cache, partial
+
+import torch
+
+try:
+    from triton.tools.disasm import extract
+except ImportError:
+    extract = None
+
+import cutlass
+import cutlass.cute as cute
+from cutlass.base_dsl.typing import JitArgument
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import NumericMeta
+
+StaticTypes = (cutlass.Constexpr, NumericMeta, int, bool, str, float, type(None))
+
+
+load_cubin_module_data_og = cutlass.base_dsl.runtime.cuda.load_cubin_module_data
+cute_compile_og = cute.compile
+
+
+torch2cute_dtype_map = {
+    torch.float16: cutlass.Float16,
+    torch.bfloat16: cutlass.BFloat16,
+    torch.float32: cutlass.Float32,
+}
+
+
+@lru_cache
+def get_max_active_clusters(cluster_size):
+    return cutlass.utils.HardwareInfo().get_max_active_clusters(
+        cluster_size=cluster_size
+    )
+
+
+@lru_cache
+def get_device_capacity(device: torch.device = None) -> tuple[int, int]:
+    return torch.cuda.get_device_capability(device)
+
+
+@dataclass
+class ArgumentsBase(JitArgument):
+    def __c_pointers__(self):
+        all_fields = [getattr(self, field.name) for field in fields(self)]
+        non_constexpr_fields = [f for f in all_fields if not isinstance(f, StaticTypes)]
+        c_ptrs = []
+        for obj in non_constexpr_fields:
+            if hasattr(obj, "__c_pointers__"):
+                c_ptrs.extend(obj.__c_pointers__())
+        return c_ptrs
+
+    def __get_mlir_types__(self):
+        all_fields = [getattr(self, field.name) for field in fields(self)]
+        non_constexpr_fields = [f for f in all_fields if not isinstance(f, StaticTypes)]
+        types, self._values_pos = [], []
+        for obj in non_constexpr_fields:
+            if hasattr(obj, "__get_mlir_types__"):
+                obj_types = obj.__get_mlir_types__()
+                types.extend(obj_types)
+                self._values_pos.append(len(obj_types))
+            else:
+                self._values_pos.append(0)
+        return types
+
+    def __new_from_mlir_values__(self, values):
+        all_fields = {field.name: getattr(self, field.name) for field in fields(self)}
+        constexpr_fields = {
+            n: f for n, f in all_fields.items() if isinstance(f, StaticTypes)
+        }
+        non_constexpr_fields = {
+            n: f for n, f in all_fields.items() if not isinstance(f, StaticTypes)
+        }
+        for (name, field), n_items in zip(
+            non_constexpr_fields.items(), self._values_pos
+        ):
+            non_constexpr_fields[name] = cutlass.new_from_mlir_values(
+                field, values[:n_items]
+            )
+            values = values[n_items:]
+        return self.__class__(**non_constexpr_fields, **constexpr_fields)
+
+
+def load_cubin_module_data_patched(cubin_data, filepath):
+    pathlib.Path(filepath).write_bytes(cubin_data)
+    return load_cubin_module_data_og(cubin_data)
+
+
+def cute_compile_patched(*args, **kwargs):
+    """A patched version of cute.compile that dump the SASS to a file if CUTE_CUBIN_PATH is set."""
+    cubin_path = os.getenv("CUTE_CUBIN_PATH", None)
+    if cubin_path is not None:
+        cutlass.base_dsl.runtime.cuda.load_cubin_module_data = partial(
+            load_cubin_module_data_patched, filepath=cubin_path
+        )
+    output = cute_compile_og(*args, **kwargs)
+    if cubin_path is not None:
+        cutlass.base_dsl.runtime.cuda.load_cubin_module_data = load_cubin_module_data_og
+        if extract is not None:
+            sass = extract(cubin_path, None)
+            pathlib.Path(cubin_path).with_suffix(".annotated.sass").write_text(sass)
+    return output
+
+
+def assume_strides_aligned(t):
+    """Assume all strides except the last are divisible by 128 bits.
+
+    Python int strides (e.g., stride=0 from GQA expand) are kept as-is
+    since they're static and don't need alignment assumptions.
+    """
+    divby = 128 // t.element_type.width
+    strides = tuple(
+        s if isinstance(s, int) else cute.assume(s, divby=divby) for s in t.stride[:-1]
+    )
+    return (*strides, t.stride[-1])
+
+
+def assume_tensor_aligned(t):
+    """Rebuild a tensor with 128-bit aligned stride assumptions. Passes through None."""
+    if t is None:
+        return None
+    return cute.make_tensor(
+        t.iterator, cute.make_layout(t.shape, stride=assume_strides_aligned(t))
+    )
+
+
+def to_cute_tensor(
+    t, assumed_align=16, leading_dim=-1, fully_dynamic=False, enable_tvm_ffi=True
+):
+    """Convert torch tensor to cute tensor for TVM FFI. leading_dim=-1 defaults to t.ndim-1."""
+    tensor = from_dlpack(
+        t.detach(), assumed_align=assumed_align, enable_tvm_ffi=enable_tvm_ffi
+    )
+    if fully_dynamic:
+        return tensor.mark_layout_dynamic()
+    if leading_dim == -1:
+        leading_dim = t.ndim - 1
+    return tensor.mark_layout_dynamic(leading_dim=leading_dim)
+
+
+def to_cute_aux_tensor(t, enable_tvm_ffi=True):
+    """Convert torch tensor to cute tensor for TVM FFI, tailored to FlexAttention aux tensors.
+    This allows the user to specify alignment and leading dimension for aux tensors used in
+    custom score_mod callables.
+    """
+    assumed_align: int = getattr(t, "__assumed_align__", None)
+    leading_dim: int = getattr(t, "__leading_dim__", None)
+    fully_dynamic: bool = leading_dim is None
+
+    return to_cute_tensor(
+        t,
+        assumed_align=assumed_align,
+        leading_dim=leading_dim,
+        fully_dynamic=fully_dynamic,
+        enable_tvm_ffi=enable_tvm_ffi,
+    )
+
+
+def get_broadcast_dims(tensor: torch.Tensor) -> tuple[bool, ...]:
+    """Return tuple of bools indicating which dims have stride=0 (broadcast).
+
+    This is useful for compile keys since CuTe's mark_layout_dynamic() keeps
+    stride=0 as static, meaning kernels compiled with different broadcast
+    patterns are not interchangeable.
+    """
+    return tuple(s == 0 for s in tensor.stride())

vllm/models/minimax_m3/nvidia/ops/prefill_gqa_sparse/common/mask.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import TypeAlias
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, Uint32, const_expr
+
+from . import utils
+from .seqlen_info import SeqlenInfoQK
+
+MaskGenFn: TypeAlias = Callable[[int], Uint32]
+MASK_R2P_CHUNK_SIZE: int = 32
+
+
+@cute.jit
+def r2p_bitmask_below(limit: Int32, s: int) -> Uint32:
+    m = max((s + 1) * MASK_R2P_CHUNK_SIZE - limit, 0)
+    return utils.shr_u32(Uint32(0xFFFFFFFF), Uint32(m))
+
+
+@cute.jit
+def r2p_bitmask_above(limit: Int32, s: int) -> Uint32:
+    n = max(limit - s * MASK_R2P_CHUNK_SIZE, 0)
+    return utils.shl_u32(Uint32(0xFFFFFFFF), Uint32(n))
+
+
+@cute.jit
+def mask_r2p_lambda(
+    X: cute.Tensor,
+    mask_gen_fn: cutlass.Constexpr[MaskGenFn],
+    rank1: bool = False,
+) -> None:
+    ncol = const_expr(
+        cute.size(X.shape[cute.rank(X) - 1]) if not rank1 else cute.size(X.shape)
+    )
+    for s in cutlass.range_constexpr(cute.ceil_div(ncol, MASK_R2P_CHUNK_SIZE)):
+        mask = mask_gen_fn(s)
+        for i in cutlass.range_constexpr(
+            min(MASK_R2P_CHUNK_SIZE, ncol - s * MASK_R2P_CHUNK_SIZE)
+        ):
+            in_bound = cutlass.Boolean(mask & (Uint32(1) << i))
+            c = s * MASK_R2P_CHUNK_SIZE + i
+            if const_expr(rank1):
+                X[c] = X[c] if in_bound else -Float32.inf
+            else:
+                for r in cutlass.range_constexpr(cute.size(X.shape[0])):
+                    X[r, c] = X[r, c] if in_bound else -Float32.inf
+
+
+@cute.jit
+def row_to_r2p_idx(x: Int32, num_rep: int, num_wg: int) -> Int32:
+    return x // (num_rep * num_wg) * num_rep + min(x % (num_rep * num_wg), num_rep)
+
+
+@dataclass(frozen=True)
+class AttentionMask:
+    tile_m: cutlass.Constexpr[int]
+    tile_n: cutlass.Constexpr[int]
+    seqlen_info: SeqlenInfoQK
+    qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1
+    swap_AB: cutlass.Constexpr[bool] = False
+
+    @property
+    def seqlen_q(self) -> Int32:
+        return self.seqlen_info.seqlen_q
+
+    @property
+    def seqlen_k(self) -> Int32:
+        return self.seqlen_info.seqlen_k
+
+    @cute.jit
+    def apply_mask_sm100_transposed(
+        self,
+        acc_S: cute.Tensor,
+        tScS_t2r: cute.Tensor,
+        t0ScS_t2r: cute.Tensor,
+        m_block: cutlass.Int32,
+        n_block: cutlass.Int32,
+        mask_seqlen: cutlass.Constexpr,
+        mask_causal: cutlass.Constexpr,
+        is_full_block: bool = False,
+        check_m_boundary: bool = True,
+    ) -> None:
+        del is_full_block, check_m_boundary
+        del t0ScS_t2r
+        row_axis = 0 if const_expr(not self.swap_AB) else 1
+        col_axis = 1 if const_expr(not self.swap_AB) else 0
+        thr_col_offset = tScS_t2r[0][col_axis]
+        seqlenk_col_limit = self.seqlen_k - n_block * self.tile_n - thr_col_offset
+
+        if const_expr(not mask_causal):
+            if const_expr(mask_seqlen) and seqlenk_col_limit <= 0:
+                for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True):
+                    acc_S[i] = -cutlass.Float32.inf
+            return
+
+        thr_row_offset = tScS_t2r[0][row_axis]
+        seqlenq_row_limit = self.seqlen_q - m_block * self.tile_m - thr_row_offset
+        row_limit_top = seqlenq_row_limit - seqlenk_col_limit
+        if const_expr(mask_seqlen) and seqlenk_col_limit <= 0:
+            row_limit_top = self.tile_m
+        num_rep = cute.size(tScS_t2r, mode=[0])
+        row_limit = row_to_r2p_idx(row_limit_top, num_rep, 2)
+        mask_r2p_lambda(
+            acc_S,
+            lambda s: r2p_bitmask_above(row_limit, s),
+            rank1=True,
+        )

vllm/models/minimax_m3/nvidia/ops/prefill_gqa_sparse/common/mma_sm100_desc.py

@@ -0,0 +1,311 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa
+from enum import IntEnum
+
+import cutlass
+import cutlass.cute as cute
+
+# ---------------------------------------------------------------------------
+# Enumerations that match the HW encodings (values MUST stay identical)
+# ---------------------------------------------------------------------------
+
+
+class Major(IntEnum):  # matrix "layout" in the ISA docs
+    K = 0
+    MN = 1
+
+
+class ScaleIn(IntEnum):  # negate flags
+    One = 0
+    Neg = 1
+
+
+class Saturate(IntEnum):
+    False_ = 0
+    True_ = 1
+
+
+class CFormat(IntEnum):  # 2-bit field (bits 4-5)
+    F16 = 0
+    F32 = 1
+    S32 = 2
+
+
+class F16F32Format(IntEnum):  # 3-bit field (A/B element type)
+    F16 = 0
+    BF16 = 1
+    TF32 = 2
+
+
+class S8Format(IntEnum):
+    UINT8 = 0
+    INT8 = 1
+
+
+class MXF8F6F4Format(IntEnum):
+    E4M3 = 0
+    E5M2 = 1
+    E2M3 = 3
+    E3M2 = 4
+    E2M1 = 5
+
+
+class MaxShift(IntEnum):
+    NoShift = 0
+    MaxShift8 = 1
+    MaxShift16 = 2
+    MaxShift32 = 3
+
+
+# ---------------------------------------------------------------------------
+# CUTLASS-type -> encoding helpers
+# ---------------------------------------------------------------------------
+
+
+def to_UMMA_format(cutlass_type) -> int:
+    """
+    Map a CUTLASS scalar class to the 3-bit encoding for Matrix A/B.
+    """
+    if cutlass_type is cutlass.Int8:
+        return S8Format.INT8
+    # Unsigned 8-bit (if available in your CUTLASS build)
+    if cutlass_type is cutlass.Uint8:
+        return S8Format.UINT8
+    # FP-16 / BF-16
+    if cutlass_type is cutlass.Float16:
+        return F16F32Format.F16
+    if cutlass_type is cutlass.BFloat16:
+        return F16F32Format.BF16
+    # TensorFloat-32 (8-bit exponent, 10-bit mantissa packed in 19 bits)
+    if cutlass_type is cutlass.TFloat32:
+        return F16F32Format.TF32
+    # Float-8 / Float-6 / Float-4 – add whenever CUTLASS exposes them
+    if cutlass_type is cutlass.FloatE4M3FN:
+        return MXF8F6F4Format.E4M3
+    if cutlass_type is cutlass.FloatE5M2:
+        return MXF8F6F4Format.E5M2
+    raise TypeError(f"Unsupported CUTLASS scalar type for A/B: {cutlass_type!r}")
+
+
+def to_C_format(cutlass_type) -> int:
+    """
+    Map a CUTLASS scalar class to the 2-bit accumulator encoding.
+    """
+    if cutlass_type is cutlass.Float16:
+        return CFormat.F16
+    if cutlass_type is cutlass.Float32:
+        return CFormat.F32
+    if cutlass_type is cutlass.Int32:
+        return CFormat.S32
+    raise TypeError(
+        f"Unsupported CUTLASS scalar type for accumulator: {cutlass_type!r}"
+    )
+
+
+# ---------------------------------------------------------------------------
+# The constructor – accepts only CUTLASS scalar classes
+# ---------------------------------------------------------------------------
+
+
+def make_instr_desc(
+    a_type,  # CUTLASS scalar class, e.g. cutlass.Int8
+    b_type,
+    c_type,
+    M: int,  # 64, 128 or 256
+    N: int,  # 8 … 256 (multiple of 8)
+    a_major: Major,
+    b_major: Major,
+    a_neg: ScaleIn = ScaleIn.One,
+    b_neg: ScaleIn = ScaleIn.One,
+    c_sat: Saturate = Saturate.False_,
+    is_sparse: bool = False,
+    max_shift: MaxShift = MaxShift.NoShift,
+) -> int:
+    """
+    Build the 32-bit instruction descriptor for Blackwell MMA.
+    All matrix/accumulator **types must be CUTLASS scalar classes** –
+    passing integers is forbidden.
+    """
+    # --- encode element formats -------------------------------------------------
+    a_fmt = int(to_UMMA_format(a_type))
+    b_fmt = int(to_UMMA_format(b_type))
+    c_fmt = int(to_C_format(c_type))
+
+    # --- range checks on M/N -----------------------------------------------------
+    if M not in (64, 128, 256):
+        raise ValueError("M must be 64, 128 or 256")
+    if N < 8 or N > 256 or (N & 7):
+        raise ValueError("N must be a multiple of 8 in the range 8…256")
+
+    m_dim = M >> 4  # 5-bit field
+    n_dim = N >> 3  # 6-bit field
+
+    # fmt: off
+    # --- pack the bit-fields -----------------------------------------------------
+    desc = 0
+    desc |= (0                 & 0x3) << 0        # sparse_id2 (always 0 here)
+    desc |= (int(is_sparse)    & 0x1) << 2        # sparse_flag
+    desc |= (int(c_sat)        & 0x1) << 3        # saturate
+    desc |= (c_fmt             & 0x3) << 4        # c_format
+    desc |= (a_fmt             & 0x7) << 7        # a_format
+    desc |= (b_fmt             & 0x7) << 10       # b_format
+    desc |= (int(a_neg)        & 0x1) << 13       # a_negate
+    desc |= (int(b_neg)        & 0x1) << 14       # b_negate
+    desc |= (int(a_major)      & 0x1) << 15       # a_major
+    desc |= (int(b_major)      & 0x1) << 16       # b_major
+    desc |= (n_dim             & 0x3F) << 17      # n_dim (6 bits)
+    desc |= (m_dim             & 0x1F) << 24      # m_dim (5 bits)
+    desc |= (int(max_shift)    & 0x3) << 30       # max_shift (2 bits)
+    # fmt: on
+
+    return desc & 0xFFFF_FFFF  # ensure 32-bit result
+
+
+def mma_op_to_idesc(op: cute.nvgpu.tcgen05.mma.MmaOp):
+    return make_instr_desc(
+        op.a_dtype,
+        op.b_dtype,
+        op.acc_dtype,
+        op.shape_mnk[0],
+        op.shape_mnk[1],
+        Major.K
+        if op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K
+        else Major.MN,
+        Major.K
+        if op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K
+        else Major.MN,
+    )
+
+
+class LayoutType(IntEnum):  # occupies the top-3 bits [61:64)
+    SWIZZLE_NONE = 0  # (a.k.a. "INTERLEAVE" in older docs)
+    SWIZZLE_128B_BASE32B = 1
+    SWIZZLE_128B = 2
+    SWIZZLE_64B = 4
+    SWIZZLE_32B = 6
+    # values 3,5,7 are reserved / illegal for UMMA
+
+
+# ---------------------------------------------------------------------------
+#  Helpers – figure out the SWIZZLE_* family from the tensor layout
+# ---------------------------------------------------------------------------
+
+
+def _layout_type(swizzle: cute.Swizzle) -> LayoutType:
+    B, M, S = swizzle.num_bits, swizzle.num_base, swizzle.num_shift
+
+    if M == 4:  # Swizzle<*,4,3>
+        if S != 3:
+            raise ValueError("Unexpected swizzle shift – want S==3 for M==4")
+        return {
+            0: LayoutType.SWIZZLE_NONE,
+            1: LayoutType.SWIZZLE_32B,
+            2: LayoutType.SWIZZLE_64B,
+            3: LayoutType.SWIZZLE_128B,
+        }[B]  # KeyError ⇒ invalid B→ raise
+    if M == 5:  # Swizzle<2,5,2> (the only legal triple for M==5)
+        if (B, S) != (2, 2):
+            raise ValueError("Only Swizzle<2,5,2> supported for 128B_BASE32B")
+        return LayoutType.SWIZZLE_128B_BASE32B
+
+    # Any other (M,B,S) triple is not a UMMA-legal shared-memory layout
+    raise ValueError("Unsupported swizzle triple for UMMA smem descriptor")
+
+
+def make_smem_desc_base(
+    layout: cute.Layout, swizzle: cute.Swizzle, major: Major
+) -> int:
+    """
+    Convert a 2-D *shared-memory* Cute layout into the Blackwell 64-bit
+    smem-descriptor, without the smem start address.
+    layout must correspond to layout of an uint128 tensor.
+    """
+    # ------------------------------------------------------------------ meta
+    layout_type = _layout_type(swizzle)  # resolve SWIZZLE_* family
+
+    VERSION = 1  # bits 46–47
+    LBO_MODE = 0  # bit  52
+    BASE_OFFSET = 0  # bits 49–51   (CUTLASS always 0)
+
+    # ---------------------------------------------------------- strides  (units: uint128_t = 16 B)
+    swizzle_atom_mn_size = {
+        LayoutType.SWIZZLE_NONE: 1,
+        LayoutType.SWIZZLE_32B: 2,
+        LayoutType.SWIZZLE_64B: 4,
+        LayoutType.SWIZZLE_128B: 8,
+        LayoutType.SWIZZLE_128B_BASE32B: 8,
+    }[layout_type]
+
+    if major is Major.MN:
+        swizzle_atom_k_size = 4 if layout_type is LayoutType.SWIZZLE_128B_BASE32B else 8
+        canonical_layout = cute.logical_divide(
+            layout, (swizzle_atom_mn_size, swizzle_atom_k_size)
+        )
+        if not cute.is_congruent(canonical_layout, ((1, 1), (1, 1))):
+            raise ValueError(
+                "Not a canonical UMMA_MN Layout: Expected profile failure."
+            )
+        stride_00 = canonical_layout.stride[0][0]
+        if layout_type is not LayoutType.SWIZZLE_NONE and stride_00 != 1:
+            raise ValueError("Not a canonical UMMA_MN Layout: Expected stride failure.")
+        stride_10 = canonical_layout.stride[1][0]
+        if stride_10 != swizzle_atom_mn_size:
+            raise ValueError("Not a canonical UMMA_MN Layout: Expected stride failure.")
+        stride_01, stride_11 = (
+            canonical_layout.stride[0][1],
+            canonical_layout.stride[1][1],
+        )
+        if layout_type is LayoutType.SWIZZLE_NONE:
+            stride_byte_offset, leading_byte_offset = stride_01, stride_11
+        else:
+            stride_byte_offset, leading_byte_offset = stride_11, stride_01
+    else:
+        if layout_type == LayoutType.SWIZZLE_128B_BASE32B:
+            raise ValueError("SWIZZLE_128B_BASE32B is invalid for Major-K")
+        if not cute.size(layout.shape[0]) % 8 == 0:
+            raise ValueError(
+                "Not a canonical UMMA_K Layout: Expected MN-size multiple of 8."
+            )
+        canonical_layout = cute.logical_divide(layout, (8, 2))
+        if not cute.is_congruent(canonical_layout, ((1, 1), (1, 1))):
+            raise ValueError("Not a canonical UMMA_K Layout: Expected profile failure.")
+        stride_00 = canonical_layout.stride[0][0]
+        if stride_00 != swizzle_atom_mn_size:
+            raise ValueError("Not a canonical UMMA_K Layout: Expected stride failure.")
+        stride_10 = canonical_layout.stride[1][0]
+        if layout_type is not LayoutType.SWIZZLE_NONE and stride_10 != 1:
+            raise ValueError("Not a canonical UMMA_K Layout: Expected stride failure.")
+        stride_01 = canonical_layout.stride[0][1]
+        stride_byte_offset, leading_byte_offset = stride_01, stride_10
+
+    # ------------------------------------------------------------------ pack
+    desc = 0
+    # leading_byte_offset_  [16:30)
+    desc |= (leading_byte_offset & 0x3FFF) << 16
+    # stride_byte_offset_   [32:46)
+    desc |= (stride_byte_offset & 0x3FFF) << 32
+    # version_             [46:48)
+    desc |= (VERSION & 0x3) << 46
+    # base_offset_         [49:52)
+    desc |= (BASE_OFFSET & 0x7) << 49
+    # lbo_mode_            [52:53)
+    desc |= (LBO_MODE & 0x1) << 52
+    # layout_type_         [61:64)
+    desc |= (int(layout_type) & 0x7) << 61
+
+    return desc & 0xFFFF_FFFF_FFFF_FFFF  # force 64-bit width
+
+
+def make_smem_desc_start_addr(start_addr: cute.Pointer) -> cutlass.Int32:
+    # 14 bits, remove 4 LSB (bits 0-13 in desc)
+    return (start_addr.toint() & 0x3FFFF) >> 4
+
+
+def smem_desc_base_from_tensor(sA: cute.Tensor, major: Major) -> int:
+    sA_swizzle = sA.iterator.type.swizzle_type
+    return make_smem_desc_base(
+        cute.recast_layout(128, sA.element_type.width, sA.layout[0]),
+        sA_swizzle,
+        major,
+    )

vllm/models/minimax_m3/nvidia/ops/prefill_gqa_sparse/common/named_barrier.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa
+import enum
+
+
+class NamedBarrierFwdSm100(enum.IntEnum):
+    Epilogue = enum.auto()  # starts from 1 as barrier 0 is reserved for sync_threads()
+    TmemPtr = enum.auto()
+    SoftmaxStatsW0 = enum.auto()
+    SoftmaxStatsW1 = enum.auto()
+    SoftmaxStatsW2 = enum.auto()
+    SoftmaxStatsW3 = enum.auto()
+    SoftmaxStatsW4 = enum.auto()
+    SoftmaxStatsW5 = enum.auto()
+    SoftmaxStatsW6 = enum.auto()
+    SoftmaxStatsW7 = enum.auto()
+    LoadWG = enum.auto()
+    StoreEpilogue = enum.auto()
+
+
+class NamedBarrierBwdSm100(enum.IntEnum):
+    EpilogueWG1 = enum.auto()
+    EpilogueWG2 = enum.auto()
+    Compute = enum.auto()
+    dQaccReduce = enum.auto()
+    TmemPtr = enum.auto()

vllm/models/minimax_m3/nvidia/ops/prefill_gqa_sparse/common/pack_gqa.py

@@ -0,0 +1,350 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa
+# mypy: ignore-errors
+"""PackGQA primitives for GQA (grouped-query attention) tile layouts.
+
+Contains:
+- ``pack_gqa_layout`` / ``unpack_gqa_layout``: fold/unfold ``qhead_per_kvhead``
+  into the seqlen dimension of a tensor layout (zero-copy view).
+- ``PackGQA``: base class with ``compute_ptr`` / ``load_Q`` / ``store_LSE`` /
+  ``store_O`` helpers for kernels that treat ``(qhead_per_kvhead × seqlen_q)``
+  as a single packed row dimension.
+- ``PackGQAComb``: subclass used by the K2 combine kernel; adds ``load_LSE``
+  for coalesced GMEM→SMEM async copies when LSE_partial is laid out with H_q
+  innermost (stride-1).
+"""
+
+from dataclasses import dataclass
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, const_expr
+from cutlass.cute import FastDivmodDivisor
+from quack import layout_utils
+
+from . import utils
+
+
+def pack_gqa_layout(T, qhead_per_kvhead, nheads_kv, head_idx):
+    """Reshape a tensor to fold qhead_per_kvhead into the seqlen dimension (mode 0).
+
+    The head dimension is at mode ``head_idx``.  Modes before it (1..head_idx-1)
+    are kept as-is (e.g. headdim for Q/O tensors), and modes after it are kept
+    as-is (e.g. batch).
+
+    For Q/O tensors (head_idx=2):
+        (seqlen_q, headdim, nheads, batch, ...) -> ((qhead_per_kvhead, seqlen_q), headdim, nheads_kv, batch, ...)
+    For LSE tensors (head_idx=1):
+        (seqlen_q, nheads, batch, ...) -> ((qhead_per_kvhead, seqlen_q), nheads_kv, batch, ...)
+    """
+    head_stride = T.stride[head_idx]
+    shape_packed = (
+        (qhead_per_kvhead, T.shape[0]),
+        *[T.shape[i] for i in range(1, head_idx)],
+        nheads_kv,
+        *[T.shape[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    stride_packed = (
+        (head_stride, T.stride[0]),
+        *[T.stride[i] for i in range(1, head_idx)],
+        head_stride * qhead_per_kvhead,
+        *[T.stride[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    return cute.make_tensor(
+        T.iterator, cute.make_layout(shape_packed, stride=stride_packed)
+    )
+
+
+def unpack_gqa_layout(T, qhead_per_kvhead, head_idx):
+    """Reverse of pack_gqa_layout: unfold qhead_per_kvhead from the seqlen dimension (mode 0).
+
+    The head dimension is at mode ``head_idx``.  Modes before it (1..head_idx-1)
+    are kept as-is (e.g. headdim for Q/O tensors), and modes after it are kept
+    as-is (e.g. batch).
+
+    For Q/O tensors (head_idx=2):
+        ((qhead_per_kvhead, seqlen_q), headdim, nheads_kv, batch, ...) -> (seqlen_q, headdim, nheads, batch, ...)
+    For LSE tensors (head_idx=1):
+        ((qhead_per_kvhead, seqlen_q), nheads_kv, batch, ...) -> (seqlen_q, nheads, batch, ...)
+    """
+    seqlen_stride = T.stride[0][1]
+    head_stride = T.stride[0][0]
+    shape_unpacked = (
+        T.shape[0][1],
+        *[T.shape[i] for i in range(1, head_idx)],
+        T.shape[head_idx] * qhead_per_kvhead,
+        *[T.shape[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    stride_unpacked = (
+        seqlen_stride,
+        *[T.stride[i] for i in range(1, head_idx)],
+        head_stride,
+        *[T.stride[i] for i in range(head_idx + 1, len(T.shape))],
+    )
+    return cute.make_tensor(
+        T.iterator, cute.make_layout(shape_unpacked, stride=stride_unpacked)
+    )
+
+
+@dataclass
+class PackGQA:
+    m_block_size: cutlass.Constexpr[int]
+    head_dim_padded: cutlass.Constexpr[int]
+    check_hdim_oob: cutlass.Constexpr[bool]
+    qhead_per_kvhead: cutlass.Constexpr[bool]
+
+    @cute.jit
+    def compute_ptr(
+        self,
+        tensor: cute.Tensor,
+        cRows: cute.Tensor,
+        tidx: cutlass.Int32,
+        block: cutlass.Int32,
+        threads_per_row: cutlass.Constexpr[int],
+        num_threads: cutlass.Constexpr[int],
+    ):
+        num_ptr_per_thread = cute.ceil_div(cute.size(cRows), threads_per_row)
+        tPrPtr = cute.make_rmem_tensor(num_ptr_per_thread, cutlass.Int64)
+        for i in cutlass.range_constexpr(num_ptr_per_thread):
+            row = i * num_threads + cRows[tidx % threads_per_row][0]
+            idx = block * self.m_block_size + row
+            m_idx = idx // self.qhead_per_kvhead
+            h_idx = idx - m_idx * self.qhead_per_kvhead
+            tPrPtr[i] = utils.elem_pointer(tensor, ((h_idx, m_idx),)).toint()
+        return tPrPtr
+
+    @cute.jit
+    def load_Q(
+        self,
+        mQ: cute.Tensor,  # ((qhead_per_kvhead, seqlen_q), headdim)
+        sQ: cute.Tensor,  # (m_block_size, head_dim_padded)
+        gmem_tiled_copy: cute.TiledCopy,
+        tidx: cutlass.Int32,
+        block: cutlass.Int32,
+        seqlen: cutlass.Int32,
+    ):
+        gmem_thr_copy = gmem_tiled_copy.get_slice(tidx)
+        cQ = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded))
+        tQsQ = gmem_thr_copy.partition_D(sQ)
+        tQcQ = gmem_thr_copy.partition_S(cQ)
+        t0QcQ = gmem_thr_copy.get_slice(0).partition_S(cQ)
+        tQpQ = utils.predicate_k(tQcQ, limit=mQ.shape[1])
+        tQcQ_row = tQcQ[0, None, 0]
+        threads_per_row = gmem_tiled_copy.layout_tv_tiled.shape[0][0]
+        assert cute.arch.WARP_SIZE % threads_per_row == 0, (
+            "threads_per_row must divide WARP_SIZE"
+        )
+        num_threads = gmem_tiled_copy.size
+        tPrQPtr = self.compute_ptr(
+            mQ[None, 0], tQcQ_row, tidx, block, threads_per_row, num_threads
+        )
+        for m in cutlass.range_constexpr(cute.size(tQsQ.shape[1])):
+            q_ptr_i64 = utils.shuffle_sync(
+                tPrQPtr[m // threads_per_row],
+                m % threads_per_row,
+                width=threads_per_row,
+            )
+            q_gmem_ptr = cute.make_ptr(
+                mQ.element_type, q_ptr_i64, cute.AddressSpace.gmem, assumed_align=16
+            )
+            if (
+                t0QcQ[0, m, 0][0]
+                < seqlen * self.qhead_per_kvhead
+                - block * self.m_block_size
+                - tQcQ_row[0][0]
+            ):
+                mQ_cur = cute.make_tensor(q_gmem_ptr, (self.head_dim_padded,))
+                elems_per_load = cute.size(tQsQ.shape[0][0])
+                mQ_cur_copy = cute.tiled_divide(mQ_cur, (elems_per_load,))
+                for k in cutlass.range_constexpr(cute.size(tQsQ.shape[2])):
+                    ki = tQcQ[0, 0, k][1] // elems_per_load
+                    cute.copy(
+                        gmem_thr_copy,
+                        mQ_cur_copy[None, ki],
+                        tQsQ[None, m, k],
+                        pred=tQpQ[None, m, k]
+                        if cutlass.const_expr(self.check_hdim_oob)
+                        else None,
+                    )
+
+    @cute.jit
+    def store_LSE(
+        self,
+        mLSE: cute.Tensor,  # (qhead_per_kvhead, seqlen_q)
+        tLSErLSE: cute.Tensor,  # (m_block_size, head_dim_padded)
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        block: cutlass.Int32,
+        seqlen: cutlass.Int32,
+    ):
+        thr_mma = tiled_mma.get_slice(tidx)
+        caccO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded))
+        taccOcO = thr_mma.partition_C(caccO)
+        taccOcO_row = layout_utils.reshape_acc_to_mn(taccOcO)[None, 0]
+        assert cute.size(tLSErLSE) == cute.size(taccOcO_row)
+        threads_per_row = tiled_mma.tv_layout_C.shape[0][0]
+        assert cute.arch.WARP_SIZE % threads_per_row == 0, (
+            "threads_per_row must divide WARP_SIZE"
+        )
+        assert cute.size(tLSErLSE) <= threads_per_row
+        num_threads = tiled_mma.size
+        tPrLSEPtr = self.compute_ptr(
+            mLSE, taccOcO_row, tidx, block, threads_per_row, num_threads
+        )
+        for m in cutlass.range_constexpr(cute.size(tLSErLSE)):
+            lse_ptr_i64 = utils.shuffle_sync(
+                tPrLSEPtr[m // threads_per_row],
+                m % threads_per_row,
+                width=threads_per_row,
+            )
+            lse_gmem_ptr = cute.make_ptr(
+                mLSE.element_type, lse_ptr_i64, cute.AddressSpace.gmem, assumed_align=4
+            )
+            row = block * self.m_block_size + taccOcO_row[m][0]
+            # Only the thread corresponding to column 0 writes out the lse to gmem
+            if taccOcO[0][1] == 0 and row < seqlen * self.qhead_per_kvhead:
+                mLSE_copy = cute.make_tensor(lse_gmem_ptr, (1,))
+                mLSE_copy[0] = tLSErLSE[m]
+
+    @cute.jit
+    def store_O(
+        self,
+        mO: cute.Tensor,  # ((qhead_per_kvhead, seqlen_q), headdim)
+        tOrO: cute.Tensor,  # (m_block_size, head_dim_padded) split across threads according to gmem_tiled_copy
+        gmem_tiled_copy: cute.TiledCopy,
+        tidx: cutlass.Int32,
+        block: cutlass.Int32,
+        seqlen: cutlass.Int32,
+    ):
+        gmem_thr_copy = gmem_tiled_copy.get_slice(tidx)
+        cO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded))
+        tOcO = gmem_thr_copy.partition_S(cO)
+        t0OcO = gmem_thr_copy.get_slice(0).partition_S(cO)
+        tOpO = utils.predicate_k(tOcO, limit=mO.shape[1])
+        tOcO_row = tOcO[0, None, 0]
+        threads_per_row = gmem_tiled_copy.layout_tv_tiled.shape[0][0]
+        assert cute.arch.WARP_SIZE % threads_per_row == 0, (
+            "threads_per_row must divide WARP_SIZE"
+        )
+        num_threads = gmem_tiled_copy.size
+        tPrOPtr = self.compute_ptr(
+            mO[None, 0], tOcO_row, tidx, block, threads_per_row, num_threads
+        )
+        for m in cutlass.range_constexpr(cute.size(tOrO.shape[1])):
+            o_ptr_i64 = utils.shuffle_sync(
+                tPrOPtr[m // threads_per_row],
+                m % threads_per_row,
+                width=threads_per_row,
+            )
+            o_gmem_ptr = cute.make_ptr(
+                mO.element_type, o_ptr_i64, cute.AddressSpace.gmem, assumed_align=16
+            )
+            if (
+                t0OcO[0, m, 0][0]
+                < seqlen * self.qhead_per_kvhead
+                - block * self.m_block_size
+                - tOcO_row[0][0]
+            ):
+                mO_cur = cute.make_tensor(o_gmem_ptr, (self.head_dim_padded,))
+                elems_per_load = cute.size(tOrO.shape[0][0])
+                mO_cur_copy = cute.tiled_divide(mO_cur, (elems_per_load,))
+                for k in cutlass.range_constexpr(cute.size(tOrO.shape[2])):
+                    ki = tOcO[0, 0, k][1] // elems_per_load
+                    cute.copy(
+                        gmem_thr_copy,
+                        tOrO[None, m, k],
+                        mO_cur_copy[None, ki],
+                        pred=tOpO[None, m, k]
+                        if cutlass.const_expr(self.check_hdim_oob)
+                        else None,
+                    )
+
+
+@dataclass
+class PackGQAComb(PackGQA):
+    """PackGQA subclass for the K2 combine kernel.
+
+    Inherits ``compute_ptr`` / ``load_Q`` / ``store_LSE`` / ``store_O`` from
+    ``PackGQA``. Adds ``load_LSE`` for coalesced GMEM→SMEM async copies when
+    LSE_partial is laid out with H_q innermost.
+
+    K2 combine treats each query head independently (no GQA grouping in combine
+    itself), so ``qhead_per_kvhead`` is set to ``num_heads_q`` by the caller —
+    all heads are folded into one "group" per Sq position.
+    """
+
+    @cute.jit
+    def load_LSE(
+        self,
+        mLSE_partial: cute.Tensor,
+        # Packed layout after caller-side reshape:
+        #   shape  ((qhead_per_kvhead, seqlen_q), num_splits)
+        #   stride ((1, qhead_per_kvhead), ...)
+        # — H_q is the innermost (stride-1) element of the packed first dim.
+        sLSE: cute.Tensor,
+        # SMEM destination: ``(max_splits, m_block_size)`` fp32.
+        max_splits: cutlass.Constexpr[int],
+        # Explicit max_splits so the identity tensor shape is a plain int,
+        # avoiding compound-shape traps from sLSE.shape[0] after tile_to_shape.
+        gmem_tiled_copy: cute.TiledCopy,
+        tidx: Int32,
+        block: Int32,
+        num_splits: Int32,
+        seqlen: Int32,
+        num_heads_divmod: FastDivmodDivisor,
+        mCounter: cute.Tensor | None = None,
+        batch_idx: Int32 | None = None,
+        qhead_per_kvhead: Int32 = Int32(1),
+        # divmod for ``m_pos = idx // qhead_per_kvhead``; passed explicitly so
+        # caller controls whether the divisor is constexpr or a runtime value.
+    ):
+        """Coalesced GMEM→SMEM async load of LSE_partial for one tile.
+
+        For each (split, row) slot this thread owns in the tile, compute the
+        GMEM coordinate ``(h_pos, m_pos)`` via PackGQA divmod and copy one fp32.
+        Out-of-bounds rows (``m_pos >= seqlen``) and splits (``si >= num_splits``)
+        are filled with ``-inf`` so they flow cleanly through downstream reductions.
+
+        Coalescing: adjacent thread rows correspond to adjacent ``h_pos`` values
+        (head varies fast under ``divmod(idx, qhead_per_kvhead)``), which map to
+        adjacent GMEM addresses when H_q is stride-1 — one sector per warp.
+        """
+        gmem_thr_copy = gmem_tiled_copy.get_slice(tidx)
+        cLSE = cute.make_identity_tensor((max_splits, self.m_block_size))
+        tLSEcLSE = gmem_thr_copy.partition_S(cLSE)
+        tLSEsLSE = gmem_thr_copy.partition_D(sLSE)
+
+        for m in cutlass.range(cute.size(tLSEcLSE, mode=[2]), unroll_full=True):
+            mi = tLSEcLSE[0, 0, m][1]
+            idx = block * self.m_block_size + mi
+            m_pos, h_pos = divmod(idx, num_heads_divmod)
+
+            if m_pos < seqlen:
+                row_count = (
+                    mCounter[batch_idx, m_pos, h_pos // qhead_per_kvhead]
+                    if const_expr(mCounter is not None)
+                    else num_splits
+                )
+                for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True):
+                    si = tLSEcLSE[0, s, 0][0]
+                    if si < num_splits and si < row_count:
+                        # Build a 1-element GMEM tensor at ((h_pos, m_pos), si),
+                        # matching PackGQA.store_LSE's ptr pattern so cute.copy
+                        # receives a proper Tensor, not a scalar.
+                        src_ptr_i64 = utils.elem_pointer(
+                            mLSE_partial, ((h_pos, m_pos), si)
+                        ).toint()
+                        src_ptr = cute.make_ptr(
+                            Float32,
+                            src_ptr_i64,
+                            cute.AddressSpace.gmem,
+                            assumed_align=4,
+                        )
+                        src_t = cute.make_tensor(src_ptr, (1,))
+                        cute.copy(gmem_thr_copy, src_t, tLSEsLSE[None, s, m])
+                    else:
+                        tLSEsLSE[None, s, m].fill(-Float32.inf)
+            else:
+                for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True):
+                    tLSEsLSE[None, s, m].fill(-Float32.inf)

vllm/models/minimax_m3/nvidia/ops/prefill_gqa_sparse/common/paged_kv.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa
+from dataclasses import dataclass
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, const_expr
+
+
+@dataclass(frozen=True)
+class PagedKVManager:
+    mPageTable: cute.Tensor
+    page_size: cutlass.Constexpr[int]
+    n_block_size: cutlass.Constexpr[int]
+    segment_rows: cutlass.Constexpr[int]
+    segments_per_block: cutlass.Constexpr[int]
+    blocks_per_page: cutlass.Constexpr[int]
+
+    @staticmethod
+    def create(
+        mPageTable: cute.Tensor,
+        *,
+        page_size: int,
+        n_block_size: int,
+    ):
+        if page_size < 8:
+            raise ValueError(
+                f"page_size must be >= 8 for TMA segmented load, got {page_size}"
+            )
+        if page_size % n_block_size == 0:
+            segment_rows = n_block_size
+            segments_per_block = 1
+            blocks_per_page = page_size // n_block_size
+        elif n_block_size % page_size == 0:
+            segment_rows = page_size
+            segments_per_block = n_block_size // page_size
+            blocks_per_page = 1
+        else:
+            raise ValueError(
+                f"page_size ({page_size}) must divide blk_kv ({n_block_size}) "
+                f"or be divisible by it"
+            )
+        return PagedKVManager(
+            mPageTable,
+            page_size=page_size,
+            n_block_size=n_block_size,
+            segment_rows=segment_rows,
+            segments_per_block=segments_per_block,
+            blocks_per_page=blocks_per_page,
+        )
+
+    @cute.jit
+    def logical_length(
+        self,
+        batch_idx: Int32,
+        num_kv_blocks: Int32,
+        mSeqUsedK=None,
+    ) -> Int32:
+        if const_expr(mSeqUsedK is not None):
+            return mSeqUsedK[batch_idx]
+        return num_kv_blocks * Int32(self.n_block_size)
+
+    @cute.jit
+    def valid_cols_in_block(
+        self,
+        batch_idx: Int32,
+        kv_block_idx: Int32,
+        num_kv_blocks: Int32,
+        mSeqUsedK=None,
+    ) -> Int32:
+        seqlen_k = self.logical_length(batch_idx, num_kv_blocks, mSeqUsedK)
+        block_start = kv_block_idx * Int32(self.n_block_size)
+        remaining = seqlen_k - block_start
+        remaining = cutlass.max(remaining, Int32(0))
+        return cutlass.min(remaining, Int32(self.n_block_size))
+
+    @cute.jit
+    def physical_page_and_tile(
+        self,
+        batch_idx: Int32,
+        kv_block_idx: Int32,
+        seg_idx: Int32,
+    ) -> tuple[Int32, Int32]:
+        if const_expr(self.page_size >= self.n_block_size):
+            logical_page = kv_block_idx // Int32(self.blocks_per_page)
+            page_tile_idx = kv_block_idx % Int32(self.blocks_per_page)
+            physical_page = self.mPageTable[batch_idx, logical_page]
+            return physical_page, page_tile_idx
+        logical_page = kv_block_idx * Int32(self.segments_per_block) + seg_idx
+        physical_page = self.mPageTable[batch_idx, logical_page]
+        return physical_page, Int32(0)
+
+    @cute.jit
+    def physical_block_index(
+        self,
+        batch_idx: Int32,
+        kv_block_idx: Int32,
+    ) -> Int32:
+        physical_page, page_tile_idx = self.physical_page_and_tile(
+            batch_idx, kv_block_idx, Int32(0)
+        )
+        return physical_page * Int32(self.blocks_per_page) + page_tile_idx
+
+    @cute.jit
+    def physical_page_index(
+        self,
+        batch_idx: Int32,
+        kv_block_idx: Int32,
+        seg_idx: Int32,
+    ) -> Int32:
+        physical_page, _ = self.physical_page_and_tile(batch_idx, kv_block_idx, seg_idx)
+        return physical_page
+
+    @cute.jit
+    def physical_row_start(
+        self,
+        batch_idx: Int32,
+        kv_block_idx: Int32,
+        seg_idx: Int32,
+    ) -> Int32:
+        physical_page, page_tile_idx = self.physical_page_and_tile(
+            batch_idx, kv_block_idx, seg_idx
+        )
+        return physical_page * Int32(self.page_size) + page_tile_idx * Int32(
+            self.n_block_size
+        )
+
+
+__all__ = ["PagedKVManager"]