refactor & simplify

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

.buildkite/image_build/image_build.yaml

@@ -14,7 +14,7 @@ steps:
           limit: 2
 
   - label: ":docker: :smoking: Non-root smoke tests"
-    key: image-smoke-test
+    key: image-build-smoke-test
     depends_on:
       - image-build
     commands:

.buildkite/test-amd.yaml

@@ -1261,7 +1261,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
+  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
 
 - label: Entrypoints Integration (API Server openai - Part 2) # TBD
   timeout_in_minutes: 180
@@ -2768,7 +2768,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
+  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
 
 - label: Entrypoints Integration (API Server openai - Part 2) # TBD
   timeout_in_minutes: 180

.buildkite/test_areas/entrypoints.yaml

@@ -43,7 +43,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
+  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
   mirror:
     amd:
       device: mi325_1

.buildkite/test_areas/rust_frontend.yaml

@@ -16,7 +16,7 @@ steps:
   - tests/benchmarks/test_serve_cli.py
   - tests/entrypoints/openai/chat_completion/test_chat_completion.py
   # - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
-  # - tests/entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py
+
   # - tests/entrypoints/openai/completion/test_prompt_validation.py
   - tests/entrypoints/openai/completion/test_shutdown.py
   # - tests/entrypoints/openai/test_return_token_ids.py
@@ -28,7 +28,7 @@ steps:
   - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
   - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py
   # - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
-  # - pytest -v -s entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py
+
   # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
   - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
   # - pytest -v -s entrypoints/openai/test_return_token_ids.py

.github/CODEOWNERS

@@ -40,6 +40,12 @@
 /vllm/entrypoints/chat_utils.py @DarkLight1337
 /vllm/entrypoints/llm.py @DarkLight1337
 
+# Rust Frontend
+/rust/ @BugenZhao @njhill
+/build_rust.sh @BugenZhao @njhill
+/rust-toolchain.toml @BugenZhao @njhill
+/.buildkite/test_areas/rust* @BugenZhao @njhill
+
 # Input/Output Processing
 /vllm/sampling_params.py @njhill @NickLucche
 /vllm/pooling_params.py @noooop @DarkLight1337
@@ -72,11 +78,13 @@
 /vllm/v1/worker/gpu/kv_connector.py @orozery
 
 # CI & building
-/.buildkite @Harry-Chen
-/docker/Dockerfile @Harry-Chen
+/.buildkite @Harry-Chen @khluu
+/docker/Dockerfile @Harry-Chen @khluu
+/pyproject.toml @khluu
+/setup.py @khluu
 
 # Test ownership
-/.buildkite/lm-eval-harness @mgoin 
+/.buildkite/lm-eval-harness @mgoin
 /tests/distributed/test_multi_node_assignment.py @youkaichao
 /tests/distributed/test_pipeline_parallel.py @youkaichao
 /tests/distributed/test_same_node.py @youkaichao

CMakeLists.txt

@@ -144,14 +144,14 @@ endif()
 # Set up GPU language and check the torch version and warn if it isn't
 # what is expected.
 #
-if (NOT HIP_FOUND AND CUDA_FOUND)
+if (NOT HIP_FOUND AND NOT PYTORCH_FOUND_HIP AND CUDA_FOUND)
   set(VLLM_GPU_LANG "CUDA")
 
   if (NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_CUDA})
     message(WARNING "Pytorch version ${TORCH_SUPPORTED_VERSION_CUDA} "
       "expected for CUDA build, saw ${Torch_VERSION} instead.")
   endif()
-elseif(HIP_FOUND)
+elseif(HIP_FOUND OR PYTORCH_FOUND_HIP)
   set(VLLM_GPU_LANG "HIP")
 
   # Importing torch recognizes and sets up some HIP/ROCm configuration but does
@@ -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"
@@ -683,6 +685,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
                    "in CUDA target architectures.")
   endif()
 
+  # FP32 router GEMM (H=3072, E=256, M<=32). Requires SM90+ and CUDA >= 12.0.
+  cuda_archs_sm90plus(FP32_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND FP32_ROUTER_GEMM_ARCHS)
+    set(SRCS
+      "csrc/libtorch_stable/fp32_router_gemm_entry.cu"
+      "csrc/libtorch_stable/fp32_router_gemm.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${SRCS}"
+      CUDA_ARCHS "${FP32_ROUTER_GEMM_ARCHS}")
+    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    message(STATUS "Building fp32_router_gemm for archs: ${FP32_ROUTER_GEMM_ARCHS}")
+  else()
+    message(STATUS "Not building fp32_router_gemm as no compatible archs found "
+                   "(requires SM90+ and CUDA >= 12.0).")
+  endif()
+
   # Only build AllSpark kernels if we are building for at least some compatible archs.
   cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
   if (ALLSPARK_ARCHS)
@@ -1240,24 +1258,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                    " in CUDA target architectures")
   endif()
 
-  # DeepSeek V3 router GEMM kernel - requires SM90+
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_ROUTER_GEMM_ARCHS)
+  # DeepSeek V3 router GEMM kernel requires SM90+ and CUDA >= 12.0.
+  # (fp32_router_gemm has been migrated to _C_stable_libtorch above.)
+  cuda_archs_sm90plus(SM90PLUS_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SM90PLUS_ROUTER_GEMM_ARCHS)
     set(DSV3_ROUTER_GEMM_SRC
       "csrc/moe/dsv3_router_gemm_entry.cu"
       "csrc/moe/dsv3_router_gemm_float_out.cu"
       "csrc/moe/dsv3_router_gemm_bf16_out.cu")
     set_gencode_flags_for_srcs(
       SRCS "${DSV3_ROUTER_GEMM_SRC}"
-      CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
+      CUDA_ARCHS "${SM90PLUS_ROUTER_GEMM_ARCHS}")
     list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
-    message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
+
+    message(STATUS "Building DSV3 router GEMM kernels for archs: ${SM90PLUS_ROUTER_GEMM_ARCHS}")
   else()
-    message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found"
+    message(STATUS "Not building DSV3 router GEMM kernels as no compatible archs found"
                    " (requires SM90+ and CUDA >= 12.0)")
   endif()
 endif()

benchmarks/kernels/benchmark_router_gemm.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn.functional as F
+
+from vllm import _custom_ops as ops
+from vllm.platforms import current_platform
+from vllm.transformers_utils.config import get_config
+from vllm.triton_utils import triton
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+# Dimensions supported by the DSV3 specialized kernel
+DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
+DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
+
+# Dimensions supported by the gpt-oss specialized kernel
+GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128]
+GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880]
+
+# Dimensions supported by the fp32 specialized kernel (MiniMax-M2)
+FP32_SUPPORTED_NUM_EXPERTS = [256]
+FP32_SUPPORTED_HIDDEN_SIZES = [3072]
+FP32_MAX_TOKENS = 32
+
+
+def get_batch_size_range(max_batch_size):
+    return [2**x for x in range(14) if 2**x <= max_batch_size]
+
+
+def get_model_params(config):
+    if config.architectures[0] in (
+        "DeepseekV2ForCausalLM",
+        "DeepseekV3ForCausalLM",
+        "DeepseekV32ForCausalLM",
+    ):
+        num_experts = config.n_routed_experts
+        hidden_size = config.hidden_size
+    elif config.architectures[0] in ("GptOssForCausalLM",) or config.architectures[
+        0
+    ] in ("MiniMaxM2ForCausalLM",):
+        num_experts = config.num_local_experts
+        hidden_size = config.hidden_size
+    else:
+        raise ValueError(f"Unsupported architecture: {config.architectures}")
+    return num_experts, hidden_size
+
+
+def get_benchmark(model, max_batch_size, trust_remote_code):
+    @triton.testing.perf_report(
+        triton.testing.Benchmark(
+            x_names=["batch_size"],
+            x_vals=get_batch_size_range(max_batch_size),
+            x_log=False,
+            line_arg="provider",
+            line_vals=[
+                "torch",
+                "vllm",
+            ],
+            line_names=["PyTorch", "vLLM"],
+            styles=([("blue", "-"), ("red", "-")]),
+            ylabel="TFLOPs",
+            plot_name=f"{model} router gemm throughput",
+            args={},
+        )
+    )
+    def benchmark(batch_size, provider):
+        config = get_config(model=model, trust_remote_code=trust_remote_code)
+        num_experts, hidden_size = get_model_params(config)
+
+        is_hopper_or_blackwell = current_platform.is_device_capability(
+            90
+        ) or current_platform.is_device_capability_family(100)
+        allow_dsv3_router_gemm = (
+            is_hopper_or_blackwell
+            and num_experts in DSV3_SUPPORTED_NUM_EXPERTS
+            and hidden_size in DSV3_SUPPORTED_HIDDEN_SIZES
+        )
+        allow_gpt_oss_router_gemm = (
+            is_hopper_or_blackwell
+            and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS
+            and hidden_size in GPT_OSS_SUPPORTED_HIDDEN_SIZES
+        )
+        is_fp32_router_model = (
+            is_hopper_or_blackwell
+            and num_experts in FP32_SUPPORTED_NUM_EXPERTS
+            and hidden_size in FP32_SUPPORTED_HIDDEN_SIZES
+        )
+        allow_fp32_router_gemm = is_fp32_router_model and batch_size <= FP32_MAX_TOKENS
+
+        # Weight dtype: fp32 kernel requires fp32 weights; others use bf16.
+        weight_dtype = torch.float32 if is_fp32_router_model else torch.bfloat16
+        mat_a = torch.randn(
+            (batch_size, hidden_size), dtype=torch.bfloat16, device="cuda"
+        ).contiguous()
+        mat_b = torch.randn(
+            (num_experts, hidden_size), dtype=weight_dtype, device="cuda"
+        ).contiguous()
+        bias = torch.randn(
+            num_experts, dtype=torch.bfloat16, device="cuda"
+        ).contiguous()
+
+        has_bias = allow_gpt_oss_router_gemm
+
+        quantiles = [0.5, 0.2, 0.8]
+
+        if provider == "torch":
+
+            def runner():
+                if allow_fp32_router_gemm:
+                    F.linear(mat_a.float(), mat_b)
+                elif has_bias:
+                    F.linear(mat_a, mat_b, bias)
+                else:
+                    F.linear(mat_a, mat_b)
+        elif provider == "vllm":
+
+            def runner():
+                if allow_dsv3_router_gemm:
+                    ops.dsv3_router_gemm(mat_a, mat_b, torch.bfloat16)
+                elif allow_fp32_router_gemm:
+                    ops.fp32_router_gemm(mat_a, mat_b)
+                elif allow_gpt_oss_router_gemm:
+                    ops.gpt_oss_router_gemm(mat_a, mat_b, bias)
+                elif is_fp32_router_model:
+                    # batch_size > FP32_MAX_TOKENS: fall back to F.linear
+                    F.linear(mat_a.float(), mat_b)
+                else:
+                    F.linear(mat_a, mat_b)
+
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            runner, quantiles=quantiles
+        )
+
+        def tflops(t_ms):
+            flops = 2 * batch_size * hidden_size * num_experts
+            return flops / (t_ms * 1e-3) / 1e12
+
+        return tflops(ms), tflops(max_ms), tflops(min_ms)
+
+    return benchmark
+
+
+if __name__ == "__main__":
+    parser = FlexibleArgumentParser()
+    parser.add_argument("--model", type=str, default="openai/gpt-oss-20b")
+    parser.add_argument("--max-batch-size", default=16, type=int)
+    parser.add_argument("--trust-remote-code", action="store_true")
+    args = parser.parse_args()
+
+    # Get the benchmark function
+    benchmark = get_benchmark(args.model, args.max_batch_size, args.trust_remote_code)
+    # Run performance benchmark
+    benchmark.run(print_data=True)

cmake/utils.cmake

@@ -476,6 +476,16 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
   set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE)
 endfunction()
 
+
+function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS)
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f" "${TGT_CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a" "${TGT_CUDA_ARCHS}")
+  endif()
+  set(${OUT_CUDA_ARCHS} ${_archs} PARENT_SCOPE)
+endfunction()
+
 #
 # Override the GPU architectures detected by cmake/torch and filter them by
 # `GPU_SUPPORTED_ARCHES`. Sets the final set of architectures in

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/fp32_router_gemm.cu

@@ -0,0 +1,223 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Router GEMM: activation(T) x weight(fp32) -> fp32, H=3072, E=256, M<=32.
+// Supports bf16 or fp32 activation; weight is always fp32.
+// Adapted from dsv3_router_gemm_float_out.cu.
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+// ---------------------------------------------------------------------------
+// Load helpers
+// ---------------------------------------------------------------------------
+
+// Load VPT fp32 values from the weight matrix (always fp32).
+//   VPT=4 when activation is fp32 (one float4 load)
+//   VPT=8 when activation is bf16 (two float4 loads)
+template <int VPT>
+__device__ __forceinline__ void load_weight(float const* ptr, float* dst);
+
+template <>
+__device__ __forceinline__ void load_weight<4>(float const* ptr, float* dst) {
+  float4 v = *reinterpret_cast<float4 const*>(ptr);
+  dst[0] = v.x;
+  dst[1] = v.y;
+  dst[2] = v.z;
+  dst[3] = v.w;
+}
+
+template <>
+__device__ __forceinline__ void load_weight<8>(float const* ptr, float* dst) {
+  float4 v0 = *reinterpret_cast<float4 const*>(ptr);
+  float4 v1 = *reinterpret_cast<float4 const*>(ptr + 4);
+  dst[0] = v0.x;
+  dst[1] = v0.y;
+  dst[2] = v0.z;
+  dst[3] = v0.w;
+  dst[4] = v1.x;
+  dst[5] = v1.y;
+  dst[6] = v1.z;
+  dst[7] = v1.w;
+}
+
+// Load VPT activation values and convert to fp32.
+template <typename T, int VPT>
+__device__ __forceinline__ void load_activation(T const* ptr, float* dst);
+
+// fp32 activation: one float4 load, no conversion needed.
+template <>
+__device__ __forceinline__ void load_activation<float, 4>(float const* ptr,
+                                                          float* dst) {
+  float4 v = *reinterpret_cast<float4 const*>(ptr);
+  dst[0] = v.x;
+  dst[1] = v.y;
+  dst[2] = v.z;
+  dst[3] = v.w;
+}
+
+// bf16 activation: one uint4 load (8 × bf16) + element-wise conversion.
+template <>
+__device__ __forceinline__ void load_activation<__nv_bfloat16, 8>(
+    __nv_bfloat16 const* ptr, float* dst) {
+  uint4 v = *reinterpret_cast<uint4 const*>(ptr);
+  __nv_bfloat16 const* bf16_ptr = reinterpret_cast<__nv_bfloat16 const*>(&v);
+#pragma unroll
+  for (int i = 0; i < 8; i++) dst[i] = __bfloat162float(bf16_ptr[i]);
+}
+
+// ---------------------------------------------------------------------------
+// Kernel
+// ---------------------------------------------------------------------------
+
+// InputT : type of activation (float or __nv_bfloat16)
+// Weight is always fp32; output is always fp32.
+// VPT = 16 / sizeof(InputT):  4 for fp32, 8 for bf16
+template <typename InputT, int kBlockSize, int kNumTokens, int kNumExperts,
+          int kHiddenDim>
+__global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
+    float* out, InputT const* mat_a, float const* mat_b) {
+  constexpr int VPT = 16 / sizeof(InputT);
+  constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
+  constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration;
+  constexpr int kWarpSize = 32;
+  constexpr int kNumWarps = kBlockSize / kWarpSize;
+
+  int const n_idx = blockIdx.x;
+  int const tid = threadIdx.x;
+  int const warpId = tid / kWarpSize;
+  int const laneId = tid % kWarpSize;
+
+  float acc[kNumTokens] = {};
+  __shared__ float sm_reduction[kNumTokens][kNumWarps];
+
+  float const* b_col = mat_b + n_idx * kHiddenDim;
+
+  int k_bases[k_iterations];
+#pragma unroll
+  for (int ki = 0; ki < k_iterations; ki++) {
+    k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
+  }
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  asm volatile("griddepcontrol.wait;");
+#endif
+
+  for (int ki = 0; ki < k_iterations; ki++) {
+    int const k_base = k_bases[ki];
+
+    float b_float[VPT];
+    load_weight<VPT>(b_col + k_base, b_float);
+
+#pragma unroll
+    for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
+      float a_float[VPT];
+      load_activation<InputT, VPT>(mat_a + m_idx * kHiddenDim + k_base,
+                                   a_float);
+#pragma unroll
+      for (int k = 0; k < VPT; k++) {
+        acc[m_idx] += a_float[k] * b_float[k];
+      }
+    }
+  }
+
+  // Warp-level butterfly reduction
+#pragma unroll
+  for (int m = 0; m < kNumTokens; m++) {
+    float sum = acc[m];
+    sum += __shfl_xor_sync(0xffffffff, sum, 16);
+    sum += __shfl_xor_sync(0xffffffff, sum, 8);
+    sum += __shfl_xor_sync(0xffffffff, sum, 4);
+    sum += __shfl_xor_sync(0xffffffff, sum, 2);
+    sum += __shfl_xor_sync(0xffffffff, sum, 1);
+    if (laneId == 0) sm_reduction[m][warpId] = sum;
+  }
+
+  __syncthreads();
+
+  if (tid == 0) {
+#pragma unroll
+    for (int m = 0; m < kNumTokens; m++) {
+      float final_sum = 0.0f;
+#pragma unroll
+      for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][w];
+      out[m * kNumExperts + n_idx] = final_sum;
+    }
+  }
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  asm volatile("griddepcontrol.launch_dependents;");
+#endif
+}
+
+// ---------------------------------------------------------------------------
+// Launcher
+// ---------------------------------------------------------------------------
+
+template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
+void invokeFp32RouterGemm(float* output, InputT const* mat_a,
+                          float const* mat_b, cudaStream_t stream) {
+  constexpr int kBlockSize = 128;
+  cudaLaunchConfig_t config;
+  config.gridDim = kNumExperts;
+  config.blockDim = kBlockSize;
+  config.dynamicSmemBytes = 0;
+  config.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = 1;
+  config.numAttrs = 1;
+  config.attrs = attrs;
+  cudaLaunchKernelEx(&config,
+                     fp32_router_gemm_kernel<InputT, kBlockSize, kNumTokens,
+                                             kNumExperts, kHiddenDim>,
+                     output, mat_a, mat_b);
+}
+
+// ---------------------------------------------------------------------------
+// Explicit instantiations: M=1..32, E=256, H=3072, for both input types
+// ---------------------------------------------------------------------------
+
+#define INSTANTIATE(T, M)                              \
+  template void invokeFp32RouterGemm<T, M, 256, 3072>( \
+      float*, T const*, float const*, cudaStream_t);
+
+#define INSTANTIATE_ALL(T) \
+  INSTANTIATE(T, 1)        \
+  INSTANTIATE(T, 2)        \
+  INSTANTIATE(T, 3)        \
+  INSTANTIATE(T, 4)        \
+  INSTANTIATE(T, 5)        \
+  INSTANTIATE(T, 6)        \
+  INSTANTIATE(T, 7)        \
+  INSTANTIATE(T, 8)        \
+  INSTANTIATE(T, 9)        \
+  INSTANTIATE(T, 10)       \
+  INSTANTIATE(T, 11)       \
+  INSTANTIATE(T, 12)       \
+  INSTANTIATE(T, 13)       \
+  INSTANTIATE(T, 14)       \
+  INSTANTIATE(T, 15)       \
+  INSTANTIATE(T, 16)       \
+  INSTANTIATE(T, 17)       \
+  INSTANTIATE(T, 18)       \
+  INSTANTIATE(T, 19)       \
+  INSTANTIATE(T, 20)       \
+  INSTANTIATE(T, 21)       \
+  INSTANTIATE(T, 22)       \
+  INSTANTIATE(T, 23)       \
+  INSTANTIATE(T, 24)       \
+  INSTANTIATE(T, 25)       \
+  INSTANTIATE(T, 26)       \
+  INSTANTIATE(T, 27)       \
+  INSTANTIATE(T, 28)       \
+  INSTANTIATE(T, 29)       \
+  INSTANTIATE(T, 30)       \
+  INSTANTIATE(T, 31)       \
+  INSTANTIATE(T, 32)
+
+INSTANTIATE_ALL(float)
+INSTANTIATE_ALL(__nv_bfloat16)
+
+#undef INSTANTIATE_ALL
+#undef INSTANTIATE

csrc/libtorch_stable/fp32_router_gemm_entry.cu

@@ -0,0 +1,127 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+#include "core/registration.h"
+#include "libtorch_stable/torch_utils.h"
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include <stdexcept>
+
+namespace {
+
+inline int getSMVersion() {
+  auto* props = get_device_prop();
+  return props->major * 10 + props->minor;
+}
+
+}  // namespace
+
+static constexpr int FP32_NUM_EXPERTS = 256;
+static constexpr int FP32_HIDDEN_DIM = 3072;
+static constexpr int FP32_MAX_TOKENS = 32;
+
+// Forward declarations — 4 template params must match fp32_router_gemm.cu
+template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
+void invokeFp32RouterGemm(float* output, InputT const* mat_a,
+                          float const* mat_b, cudaStream_t stream);
+
+// LoopUnroller templated on InputT
+template <typename InputT, int kBegin, int kEnd>
+struct Fp32LoopUnroller {
+  static void unroll(int num_tokens, float* output, InputT const* mat_a,
+                     float const* mat_b, cudaStream_t stream) {
+    if (num_tokens == kBegin) {
+      invokeFp32RouterGemm<InputT, kBegin, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
+          output, mat_a, mat_b, stream);
+    } else {
+      Fp32LoopUnroller<InputT, kBegin + 1, kEnd>::unroll(num_tokens, output,
+                                                         mat_a, mat_b, stream);
+    }
+  }
+};
+
+template <typename InputT, int kEnd>
+struct Fp32LoopUnroller<InputT, kEnd, kEnd> {
+  static void unroll(int num_tokens, float* output, InputT const* mat_a,
+                     float const* mat_b, cudaStream_t stream) {
+    if (num_tokens == kEnd) {
+      invokeFp32RouterGemm<InputT, kEnd, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
+          output, mat_a, mat_b, stream);
+    } else {
+      throw std::invalid_argument(
+          "fp32_router_gemm: num_tokens must be in [1, 32]");
+    }
+  }
+};
+
+void fp32_router_gemm(
+    torch::stable::Tensor& output,       // [num_tokens, num_experts]
+    torch::stable::Tensor const& mat_a,  // [num_tokens, hidden_dim]
+    torch::stable::Tensor const& mat_b   // [num_experts, hidden_dim]
+) {
+  STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
+  STD_TORCH_CHECK(output.is_cuda() && mat_a.is_cuda() && mat_b.is_cuda(),
+                  "fp32_router_gemm: all tensors must be CUDA tensors");
+  STD_TORCH_CHECK(output.get_device_index() == mat_a.get_device_index() &&
+                      output.get_device_index() == mat_b.get_device_index(),
+                  "fp32_router_gemm: all tensors must be on the same device");
+  STD_TORCH_CHECK(
+      output.is_contiguous() && mat_a.is_contiguous() && mat_b.is_contiguous(),
+      "fp32_router_gemm: all tensors must be contiguous");
+
+  const int num_tokens = mat_a.size(0);
+  const int num_experts = mat_b.size(0);
+  const int hidden_dim = mat_a.size(1);
+
+  STD_TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == num_experts,
+                  "fp32_router_gemm: output must have shape [num_tokens, "
+                  "num_experts]");
+  STD_TORCH_CHECK(
+      mat_a.size(1) == mat_b.size(1),
+      "fp32_router_gemm: mat_a and mat_b must have the same hidden_dim");
+  STD_TORCH_CHECK(hidden_dim == FP32_HIDDEN_DIM,
+                  "fp32_router_gemm: expected hidden_dim=3072");
+  STD_TORCH_CHECK(num_experts == FP32_NUM_EXPERTS,
+                  "fp32_router_gemm: expected num_experts=256");
+  STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
+                  "fp32_router_gemm: num_tokens must be in [0, 32]");
+  STD_TORCH_CHECK(
+      mat_a.scalar_type() == torch::headeronly::ScalarType::Float ||
+          mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "fp32_router_gemm: mat_a must be float32 or bfloat16");
+  STD_TORCH_CHECK(mat_b.scalar_type() == torch::headeronly::ScalarType::Float,
+                  "fp32_router_gemm: mat_b (weight) must be float32");
+  STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Float,
+                  "fp32_router_gemm: output must be float32");
+
+  if (num_tokens == 0) {
+    return;
+  }
+
+  STD_TORCH_CHECK(getSMVersion() >= 90, "fp32_router_gemm: requires SM90+");
+
+  auto stream = get_current_cuda_stream(mat_a.get_device_index());
+  float* out_ptr = reinterpret_cast<float*>(output.mutable_data_ptr());
+  float const* mat_b_ptr = reinterpret_cast<float const*>(mat_b.data_ptr());
+
+  if (mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
+    auto const* mat_a_ptr =
+        reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
+    Fp32LoopUnroller<__nv_bfloat16, 1, FP32_MAX_TOKENS>::unroll(
+        num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
+  } else {
+    auto const* mat_a_ptr = reinterpret_cast<float const*>(mat_a.data_ptr());
+    Fp32LoopUnroller<float, 1, FP32_MAX_TOKENS>::unroll(
+        num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
+  }
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("fp32_router_gemm", TORCH_BOX(&fp32_router_gemm));
+}

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/layernorm_kernels.cu

@@ -78,8 +78,7 @@ __global__ void rms_norm_kernel(
 #pragma unroll
     for (int j = 0; j < VEC_SIZE; j++) {
       float x = static_cast<float>(src1.val[j]);
-      float w = static_cast<float>(src2.val[j]);
-      dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
+      dst.val[j] = static_cast<scalar_t>(x * s_variance) * src2.val[j];
     }
     v_out[i] = dst;
   }
@@ -143,8 +142,7 @@ fused_add_rms_norm_kernel(
 #pragma unroll
     for (int j = 0; j < width; ++j) {
       float x = Converter::convert(res.data[j]);
-      float wf = Converter::convert(w.data[j]);
-      out.data[j] = Converter::convert(x * s_variance * wf);
+      out.data[j] = Converter::convert(x * s_variance) * w.data[j];
     }
     input_v[strided_id] = out;
   }
@@ -183,8 +181,8 @@ fused_add_rms_norm_kernel(
 
   for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
     float x = (float)residual[blockIdx.x * hidden_size + idx];
-    float w = (float)weight[idx];
-    input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
+    input[blockIdx.x * input_stride + idx] =
+        (scalar_t)(x * s_variance) * weight[idx];
   }
 }
 

csrc/libtorch_stable/layernorm_quant_kernels.cu

@@ -66,13 +66,8 @@ __global__ void rms_norm_static_fp8_quant_kernel(
 #pragma unroll
     for (int j = 0; j < VEC_SIZE; j++) {
       float x = static_cast<float>(src1.val[j]);
-      float w = static_cast<float>(src2.val[j]);
-      // Round normalized result through scalar_t to match the precision of the
-      // unfused composite (rms_norm writes scalar_t, then
-      // static_scaled_fp8_quant re-loads it as float before FP8 conversion).
-      // Without this round, the fused path is strictly more accurate and
-      // disagrees with the composite at exact E4M3 quantization tie boundaries.
-      scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
+      // Multiply in weight's native dtype to match rms_norm_kernel.
+      scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * src2.val[j];
       out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
           scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
                                                 scale_inv);
@@ -142,12 +137,8 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 #pragma unroll
     for (int i = 0; i < width; ++i) {
       float x = Converter::convert(res.data[i]);
-      float wf = Converter::convert(w.data[i]);
-      // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
-      // to match the unfused composite path at FP8 boundaries. We use the
-      // backend's hip_type for the intermediate since c10::Half/BFloat16 has
-      // ambiguous conversions on CUDA and no implicit conversion on ROCm.
-      HipT out_norm_h = Converter::convert(x * s_variance * wf);
+      // Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
+      HipT out_norm_h = Converter::convert(x * s_variance) * w.data[i];
       out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
           Converter::convert(out_norm_h), scale_inv);
     }
@@ -192,10 +183,8 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 
   for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
     float x = (float)residual[blockIdx.x * hidden_size + idx];
-    float w = (float)weight[idx];
-    // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
-    // to match the unfused composite path at FP8 boundaries.
-    scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
+    // Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
+    scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * weight[idx];
     out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
         static_cast<float>(out_norm), scale_inv);
   }

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

@@ -247,6 +247,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   ops.def(
       "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
 
+  // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+).
+  // conditionally compiled so impl registration is in source file
+  ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
+
   // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
   ops.def(
       "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
@@ -333,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, "
@@ -360,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) -> ()");
@@ -567,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,14 +170,14 @@ 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 |
 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
-| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
 | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any |
@@ -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.

docs/governance/committers.md

@@ -17,6 +17,7 @@ Sorted alphabetically by GitHub handle:
 - [@bbrowning](https://github.com/bbrowning): Tool use and reasoning parser
 - [@benchislett](https://github.com/benchislett): Engine core and spec decode
 - [@bigPYJ1151](https://github.com/bigPYJ1151): Intel CPU/XPU integration
+- [@BugenZhao](https://github.com/BugenZhao): Rust frontend
 - [@chaunceyjiang](https://github.com/chaunceyjiang): Tool use and reasoning parser
 - [@DarkLight1337](https://github.com/DarkLight1337): Multimodality, API server
 - [@esmeetu](https://github.com/esmeetu): developer marketing, community
@@ -130,6 +131,8 @@ If you have PRs touching the area, please feel free to ping the area owner for r
     - @DarkLight1337
 - API Server: The OpenAI-compatible API server
     - @DarkLight1337, @njhill, @aarnphm, @simon-mo, @heheda12345 (Responses API)
+- Rust Frontend: The experimental API server in Rust
+    - @BugenZhao, @njhill
 - Batch Runner: The OpenAI-compatible batch runner
     - @simon-mo
 

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(),
@@ -1165,9 +1235,7 @@ setup(
     install_requires=get_requirements(),
     extras_require={
         # AMD Zen CPU optimizations via zentorch
-        "zen": [
-            "zentorch-weekly==5.2.1.dev20260408"
-        ],  # Zentorch has weekly releases. This pulls the known-good version.
+        "zen": ["zentorch==2.11.0.0"],
         "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
         "tensorizer": ["tensorizer==2.10.1"],
         "fastsafetensors": ["fastsafetensors >= 0.2.2"],

tests/entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py

@@ -1,141 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import openai  # use the official client for correctness check
-import pytest
-import pytest_asyncio
-
-from tests.utils import RemoteOpenAIServer
-
-# a reasoning and tool calling model
-MODEL_NAME = "Qwen/QwQ-32B"
-
-
-@pytest.fixture(scope="module")
-def server():
-    args = [
-        "--max-model-len",
-        "8192",
-        "--enforce-eager",
-        "--reasoning-parser",
-        "deepseek_r1",
-        "--enable-auto-tool-choice",
-        "--tool-call-parser",
-        "hermes",
-    ]
-
-    with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
-        yield remote_server
-
-
-@pytest_asyncio.fixture
-async def client(server):
-    async with server.get_async_client() as async_client:
-        yield async_client
-
-
-TOOLS = [
-    {
-        "type": "function",
-        "function": {
-            "name": "get_current_weather",
-            "description": "Get the current weather in a given location",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "city": {
-                        "type": "string",
-                        "description": "The city to find the weather for, e.g. "
-                        "'San Francisco'",
-                    },
-                    "state": {
-                        "type": "string",
-                        "description": "the two-letter abbreviation for the state that "
-                        "the city is in, e.g. 'CA' which would mean 'California'",
-                    },
-                    "unit": {
-                        "type": "string",
-                        "description": "The unit to fetch the temperature in",
-                        "enum": ["celsius", "fahrenheit"],
-                    },
-                },
-                "required": ["city", "state", "unit"],
-            },
-        },
-    }
-]
-
-MESSAGES = [
-    {"role": "user", "content": "Hi! How are you doing today?"},
-    {"role": "assistant", "content": "I'm doing well! How can I help you?"},
-    {
-        "role": "user",
-        "content": "Can you tell me what the temperate will be in Dallas, "
-        "in fahrenheit?",
-    },
-]
-
-FUNC_NAME = "get_current_weather"
-FUNC_ARGS = """{"city": "Dallas", "state": "TX", "unit": "fahrenheit"}"""
-
-
-def extract_reasoning_and_calls(chunks: list):
-    reasoning = ""
-    tool_call_idx = -1
-    arguments = []
-    function_names = []
-    for chunk in chunks:
-        if chunk.choices[0].delta.tool_calls:
-            tool_call = chunk.choices[0].delta.tool_calls[0]
-            if tool_call.index != tool_call_idx:
-                tool_call_idx = chunk.choices[0].delta.tool_calls[0].index
-                arguments.append("")
-                function_names.append("")
-
-            if tool_call.function:
-                if tool_call.function.name:
-                    function_names[tool_call_idx] = tool_call.function.name
-
-                if tool_call.function.arguments:
-                    arguments[tool_call_idx] += tool_call.function.arguments
-        else:
-            if hasattr(chunk.choices[0].delta, "reasoning"):
-                reasoning += chunk.choices[0].delta.reasoning
-    return reasoning, arguments, function_names
-
-
-# test streaming
-@pytest.mark.asyncio
-async def test_chat_streaming_of_tool_and_reasoning(client: openai.AsyncOpenAI):
-    stream = await client.chat.completions.create(
-        model=MODEL_NAME,
-        messages=MESSAGES,
-        tools=TOOLS,
-        temperature=0.0,
-        stream=True,
-    )
-
-    chunks = []
-    async for chunk in stream:
-        chunks.append(chunk)
-
-    reasoning, arguments, function_names = extract_reasoning_and_calls(chunks)
-    assert len(reasoning) > 0
-    assert len(function_names) > 0 and function_names[0] == FUNC_NAME
-    assert len(arguments) > 0 and arguments[0] == FUNC_ARGS
-
-
-# test full generate
-@pytest.mark.asyncio
-async def test_chat_full_of_tool_and_reasoning(client: openai.AsyncOpenAI):
-    tool_calls = await client.chat.completions.create(
-        model=MODEL_NAME,
-        messages=MESSAGES,
-        tools=TOOLS,
-        temperature=0.0,
-        stream=False,
-    )
-
-    assert len(tool_calls.choices[0].message.reasoning) > 0
-    assert tool_calls.choices[0].message.tool_calls[0].function.name == FUNC_NAME
-    assert tool_calls.choices[0].message.tool_calls[0].function.arguments == FUNC_ARGS

tests/entrypoints/openai/chat_completion/test_serving_chat.py

@@ -1449,91 +1449,6 @@ class TestServingChatWithHarmony:
             ],
         )
 
-    @pytest.mark.asyncio
-    async def test_tools_and_reasoning(
-        self, serving_chat, stream, weather_tools, weather_messages_start
-    ):
-        tools = weather_tools
-        messages = list(weather_messages_start)
-
-        # Test the Harmony messages for the first turn's input
-        req = ChatCompletionRequest(model=MODEL_NAME, messages=messages, tools=tools)
-        input_messages, _ = (
-            serving_chat.openai_serving_render._make_request_with_harmony(req)
-        )
-        verify_harmony_messages(
-            input_messages,
-            [
-                {"role": "system"},
-                {"role": "developer", "tool_definitions": ["get_weather"]},
-                {"role": "user", "content": messages[0]["content"]},
-            ],
-        )
-
-        # Test the Chat Completion response for the first turn's output
-        reasoning_str = "I'll call get_weather."
-        tool_args_str = '{"location": "Paris"}'
-        response_str = (
-            f"<|channel|>analysis<|message|>{reasoning_str}<|end|>"
-            "<|start|>assistant to=functions.get_weather<|channel|>commentary"
-            f"<|constrain|>json<|message|>{tool_args_str}<|call|>"
-        )
-        response = await self.generate_response_from_harmony_str(
-            serving_chat, req, response_str, stream=stream
-        )
-        verify_chat_response(
-            response,
-            reasoning=reasoning_str,
-            tool_calls=[("get_weather", tool_args_str)],
-        )
-
-        tool_call = response.choices[0].message.tool_calls[0]
-
-        # Add the output messages from the first turn as input to the second turn
-        for choice in response.choices:
-            messages.append(choice.message.model_dump(exclude_none=True))
-
-        # Add our tool output message
-        messages.append(
-            {
-                "role": "tool",
-                "tool_call_id": tool_call.id,
-                "content": "20 degrees Celsius",
-            },
-        )
-
-        # Test the Harmony messages for the second turn's input
-        req_2 = ChatCompletionRequest(model=MODEL_NAME, messages=messages, tools=tools)
-        input_messages_2, _ = (
-            serving_chat.openai_serving_render._make_request_with_harmony(req_2)
-        )
-        verify_harmony_messages(
-            input_messages_2,
-            [
-                {"role": "system"},
-                {"role": "developer"},
-                {"role": "user"},
-                {
-                    "role": "assistant",
-                    "channel": "analysis",
-                    "content": reasoning_str,
-                },
-                {
-                    "role": "assistant",
-                    "channel": "commentary",
-                    "recipient": "functions.get_weather",
-                    "content": tool_args_str,
-                },
-                {
-                    "role": "tool",
-                    "author_name": "functions.get_weather",
-                    "channel": "commentary",
-                    "recipient": "assistant",
-                    "content": "20 degrees Celsius",
-                },
-            ],
-        )
-
     @pytest.mark.asyncio
     async def test_multi_turn_tools_and_reasoning(
         self, serving_chat, stream, weather_tools, weather_messages_start

tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py

@@ -121,7 +121,9 @@ class TestExtractHarmonyStreamingDelta:
 
         token_states = [
             TokenState(
-                channel=channel, recipient="functions.get_weather", text=args_text
+                channel=channel,
+                recipient="functions.get_weather",
+                text=args_text,
             )
         ]
 
@@ -168,7 +170,11 @@ class TestExtractHarmonyStreamingDelta:
         parser = MockStreamableParser(messages=messages)
 
         token_states = [
-            TokenState(channel="commentary", recipient="functions.tool2", text="args")
+            TokenState(
+                channel="commentary",
+                recipient="functions.tool2",
+                text="args",
+            )
         ]
 
         delta_message, _ = extract_harmony_streaming_delta(
@@ -199,75 +205,6 @@ class TestExtractHarmonyStreamingDelta:
         assert delta_message.content == delta_text
         assert tools_streamed is False
 
-    @pytest.mark.parametrize("channel", ["commentary", "analysis"])
-    @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id")
-    def test_new_tool_call_without_functions_prefix(
-        self, mock_make_tool_call_id, channel
-    ):
-        mock_make_tool_call_id.return_value = "call_bare123"
-        parser = MockStreamableParser()
-
-        token_states = [TokenState(channel=channel, recipient="get_weather", text="")]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient=None,
-            include_reasoning=False,
-        )
-
-        assert delta_message is not None
-        assert len(delta_message.tool_calls) == 1
-        tool_call = delta_message.tool_calls[0]
-        assert tool_call.id == "call_bare123"
-        assert tool_call.type == "function"
-        assert tool_call.function.name == "get_weather"
-        assert tool_call.function.arguments == ""
-        assert tool_call.index == 0
-        assert tools_streamed is True
-
-    @pytest.mark.parametrize("channel", ["commentary", "analysis"])
-    def test_tool_call_argument_streaming_without_functions_prefix(self, channel):
-        parser = MockStreamableParser()
-        args_text = '{"location": "Paris"}'
-
-        token_states = [
-            TokenState(channel=channel, recipient="get_weather", text=args_text)
-        ]
-
-        delta_message, tools_streamed = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient="get_weather",
-            include_reasoning=False,
-        )
-
-        assert delta_message is not None
-        tool_call = delta_message.tool_calls[0]
-        assert tool_call.id is None
-        assert tool_call.function.arguments == args_text
-        assert tool_call.index == 0
-        assert tools_streamed is True
-
-    def test_tool_call_index_from_previous_messages_without_functions_prefix(self):
-        messages = [
-            MockMessage(channel="commentary", recipient="tool1"),
-        ]
-        parser = MockStreamableParser(messages=messages)
-
-        token_states = [
-            TokenState(channel="commentary", recipient="tool2", text="args")
-        ]
-
-        delta_message, _ = extract_harmony_streaming_delta(
-            harmony_parser=parser,
-            token_states=token_states,
-            prev_recipient="tool2",
-            include_reasoning=False,
-        )
-
-        assert delta_message.tool_calls[0].index == 1
-
     @pytest.mark.parametrize("channel", ["commentary", "analysis"])
     @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id")
     def test_new_tool_call_dotted_function_name(self, mock_make_tool_call_id, channel):

tests/entrypoints/openai/test_dp_supervisor.py

@@ -21,7 +21,10 @@ import asyncio
 import contextlib
 import os
 import signal
+import subprocess
+import tempfile
 import time
+from pathlib import Path
 from types import SimpleNamespace
 
 import aiohttp
@@ -35,6 +38,7 @@ from vllm.entrypoints.openai.dp_supervisor import (
     DPSupervisor,
     _build_vllm_dp_server_args,
     infer_multi_port_external_lb_start_rank,
+    validate_multi_port_external_lb_args,
 )
 from vllm.logger import init_logger
 
@@ -75,6 +79,8 @@ def _make_unit_args(**overrides) -> argparse.Namespace:
         "ssl_keyfile": None,
         "ssl_certfile": None,
         "ssl_ca_certs": None,
+        "ssl_cert_reqs": 0,
+        "ssl_ciphers": None,
         "node_rank": 1,
         "tensor_parallel_size": 1,
         "pipeline_parallel_size": 1,
@@ -108,6 +114,8 @@ def _make_args(**overrides) -> argparse.Namespace:
         ssl_keyfile=None,
         ssl_certfile=None,
         ssl_ca_certs=None,
+        ssl_cert_reqs=0,
+        ssl_ciphers=None,
         node_rank=0,
         tensor_parallel_size=1,
         pipeline_parallel_size=1,
@@ -118,6 +126,33 @@ def _make_args(**overrides) -> argparse.Namespace:
     return argparse.Namespace(**base)
 
 
+def _generate_self_signed_cert(cert_dir: Path) -> tuple[Path, Path]:
+    """Generate a self-signed certificate for HTTPS lifecycle tests."""
+    cert_file = cert_dir / "cert.pem"
+    key_file = cert_dir / "key.pem"
+    subprocess.run(
+        [
+            "openssl",
+            "req",
+            "-x509",
+            "-newkey",
+            "rsa:2048",
+            "-keyout",
+            str(key_file),
+            "-out",
+            str(cert_file),
+            "-days",
+            "1",
+            "-nodes",
+            "-subj",
+            "/CN=localhost",
+        ],
+        check=True,
+        capture_output=True,
+    )
+    return cert_file, key_file
+
+
 # ---------------------------------------------------------------------------
 # Unit tests
 # ---------------------------------------------------------------------------
@@ -141,6 +176,15 @@ def test_build_multi_port_external_lb_child_args_sets_external_rank_server():
     assert child_args.api_server_count == 1
 
 
+def test_validate_multi_port_external_lb_args_allows_ssl():
+    args = _make_unit_args(
+        ssl_keyfile="/tmp/server.key",
+        ssl_certfile="/tmp/server.crt",
+        ssl_ca_certs="/tmp/ca.crt",
+    )
+    validate_multi_port_external_lb_args(args)
+
+
 def test_aggregates_health():
     supervisor = DPSupervisor(_make_unit_args())
     supervisor._is_ready = True
@@ -236,10 +280,18 @@ class MockVLLMServer:
     Health state is toggled by the test via set_healthy().
     """
 
-    def __init__(self, port: int, drain_seconds: float = 0.0) -> None:
+    def __init__(
+        self,
+        port: int,
+        drain_seconds: float = 0.0,
+        ssl_keyfile: str | None = None,
+        ssl_certfile: str | None = None,
+    ) -> None:
         self.port = port
         self._healthy = False
         self._drain_seconds = drain_seconds
+        self._ssl_keyfile = ssl_keyfile
+        self._ssl_certfile = ssl_certfile
         self._server: uvicorn.Server | None = None
         self._serve_task: asyncio.Task | None = None
 
@@ -274,6 +326,8 @@ class MockVLLMServer:
             port=self.port,
             log_level="warning",
             lifespan="off",
+            ssl_keyfile=self._ssl_keyfile,
+            ssl_certfile=self._ssl_certfile,
         )
         self._server = uvicorn.Server(config)
 
@@ -312,7 +366,11 @@ class MockVLLMServer:
 
 def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]):
     logger.info("Launching mock vLLM on port %s", child_args.port)
-    mock_vllm = MockVLLMServer(port=child_args.port)
+    mock_vllm = MockVLLMServer(
+        port=child_args.port,
+        ssl_keyfile=child_args.ssl_keyfile,
+        ssl_certfile=child_args.ssl_certfile,
+    )
     asyncio.run(mock_vllm.start())
 
 
@@ -320,7 +378,12 @@ def launch_mock_vllm_with_drain(
     child_args: argparse.Namespace, env_updates: dict[str, str]
 ):
     logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port)
-    mock_vllm = MockVLLMServer(port=child_args.port, drain_seconds=10.0)
+    mock_vllm = MockVLLMServer(
+        port=child_args.port,
+        drain_seconds=10.0,
+        ssl_keyfile=child_args.ssl_keyfile,
+        ssl_certfile=child_args.ssl_certfile,
+    )
     asyncio.run(mock_vllm.start())
 
 
@@ -329,15 +392,16 @@ def launch_mock_vllm_with_drain(
 # ---------------------------------------------------------------------------
 
 
-async def _poll_supervisor_health(expected_status: int) -> bool:
+async def _poll_supervisor_health(expected_status: int, use_ssl: bool = False) -> bool:
     """
     Poll GET /health on the supervisor until expected_status is seen.
     A connection error is treated as 503-equivalent when expected_status != 200.
     """
-    url = f"http://127.0.0.1:{_SUPERVISOR_PORT}/health"
+    scheme = "https" if use_ssl else "http"
+    url = f"{scheme}://127.0.0.1:{_SUPERVISOR_PORT}/health"
     async with aiohttp.ClientSession() as session:
         try:
-            async with session.get(url) as resp:
+            async with session.get(url, ssl=False if use_ssl else None) as resp:
                 if resp.status != expected_status:
                     print(f"expected: {expected_status=}, got: {resp.status=}")
                     return False
@@ -349,12 +413,15 @@ async def _poll_supervisor_health(expected_status: int) -> bool:
             return True
 
 
-async def _poll_until_api_server_running(port: int, retries: int = 10) -> None:
-    url = f"http://127.0.0.1:{port}/health"
+async def _poll_until_api_server_running(
+    port: int, retries: int = 10, use_ssl: bool = False
+) -> None:
+    scheme = "https" if use_ssl else "http"
+    url = f"{scheme}://127.0.0.1:{port}/health"
     async with aiohttp.ClientSession() as session:
         for _ in range(retries):
             try:
-                async with session.get(url) as resp:
+                async with session.get(url, ssl=False if use_ssl else None) as resp:
                     if resp.status != 200:
                         return
                 await asyncio.sleep(1.0)
@@ -363,22 +430,34 @@ async def _poll_until_api_server_running(port: int, retries: int = 10) -> None:
                 await asyncio.sleep(1.0)
 
 
-async def _set_healthy(port: int) -> None:
-    url = f"http://127.0.0.1:{port}/set_healthy"
-    async with aiohttp.ClientSession() as session, session.get(url) as resp:
+async def _set_healthy(port: int, use_ssl: bool = False) -> None:
+    scheme = "https" if use_ssl else "http"
+    url = f"{scheme}://127.0.0.1:{port}/set_healthy"
+    async with (
+        aiohttp.ClientSession() as session,
+        session.get(url, ssl=False if use_ssl else None) as resp,
+    ):
         assert resp.status == 200
 
 
-async def _set_unhealthy(port: int) -> None:
-    url = f"http://127.0.0.1:{port}/set_unhealthy"
-    async with aiohttp.ClientSession() as session, session.get(url) as resp:
+async def _set_unhealthy(port: int, use_ssl: bool = False) -> None:
+    scheme = "https" if use_ssl else "http"
+    url = f"{scheme}://127.0.0.1:{port}/set_unhealthy"
+    async with (
+        aiohttp.ClientSession() as session,
+        session.get(url, ssl=False if use_ssl else None) as resp,
+    ):
         assert resp.status == 200
 
 
-async def _kill_server(port: int) -> None:
-    url = f"http://127.0.0.1:{port}/kill"
+async def _kill_server(port: int, use_ssl: bool = False) -> None:
+    scheme = "https" if use_ssl else "http"
+    url = f"{scheme}://127.0.0.1:{port}/kill"
     try:
-        async with aiohttp.ClientSession() as session, session.get(url) as resp:
+        async with (
+            aiohttp.ClientSession() as session,
+            session.get(url, ssl=False if use_ssl else None) as resp,
+        ):
             assert resp.status != 200
     except Exception as e:
         assert isinstance(e, aiohttp.ClientConnectorError)
@@ -455,6 +534,34 @@ async def test_basic_lifecycle(monkeypatch):
         print("everything was cleaned up!")
 
 
+@pytest.mark.asyncio
+async def test_basic_lifecycle_with_ssl(monkeypatch):
+    with tempfile.TemporaryDirectory() as cert_dir:
+        cert_file, key_file = _generate_self_signed_cert(Path(cert_dir))
+        args = _make_args(
+            ssl_keyfile=str(key_file),
+            ssl_certfile=str(cert_file),
+        )
+
+        vllm_server_ports = [_CHILD_PORT_BASE + i for i in range(_N_CHILDREN)]
+
+        async with _run_supervisor(args, monkeypatch) as (supervisor, _task):
+            assert await _poll_supervisor_health(503, use_ssl=True)
+            assert not supervisor.is_ready
+
+            for port in vllm_server_ports:
+                assert await _poll_supervisor_health(503, use_ssl=True)
+                assert not supervisor.is_ready
+                await _poll_until_api_server_running(port, use_ssl=True)
+
+            for port in vllm_server_ports:
+                await _set_healthy(port, use_ssl=True)
+            await asyncio.sleep(1.0)
+
+            assert await _poll_supervisor_health(200, use_ssl=True)
+            assert supervisor.is_ready
+
+
 @pytest.mark.asyncio
 async def test_failed_startup(monkeypatch):
     """

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/moe/modular_kernel_tools/common.py

@@ -224,10 +224,6 @@ class Config:
         info = expert_info(self.fused_experts_type)
         return info.blocked_quantization_support
 
-    def supports_expert_map(self):
-        info = expert_info(self.fused_experts_type)
-        return info.supports_expert_map
-
     def supports_apply_weight_on_input(self):
         info = prepare_finalize_info(self.prepare_finalize_type)
         return info.supports_apply_weight_on_input
@@ -326,6 +322,15 @@ class Config:
         if self.needs_mori() and not has_mori():  # noqa: SIM103
             return False, "Needs MoRI, but MoRI not available."
 
+        try:
+            if not self.fused_experts_type._supports_current_device():
+                return (
+                    False,
+                    f"{self.fused_experts_type} not supported on the current device.",
+                )
+        except NotImplementedError:
+            pass
+
         return True, None
 
 
@@ -471,7 +476,7 @@ class RankTensors:
         topk_ids = topk_ids.to(device=device)
 
         expert_map = None
-        if config.world_size > 1 and config.supports_expert_map():
+        if config.world_size > 1:
             expert_map = torch.full(
                 (global_num_experts,), fill_value=-1, dtype=torch.int32
             )

tests/kernels/moe/modular_kernel_tools/mk_objects.py

@@ -67,7 +67,6 @@ class ExpertInfo:
     activation_format: mk.FusedMoEActivationFormat
     supported_dtypes: list[torch.dtype | str]
     blocked_quantization_support: bool
-    supports_expert_map: bool
     needs_matching_quant: bool = False
     needs_deep_gemm: bool = False
     needs_aiter: bool = False
@@ -129,7 +128,6 @@ def register_experts(
     activation_format: mk.FusedMoEActivationFormat,
     supported_dtypes: list[torch.dtype | str],
     blocked_quantization_support: bool,
-    supports_expert_map: bool,
     needs_matching_quant: bool = False,
     needs_deep_gemm: bool = False,
     needs_aiter: bool = False,
@@ -142,7 +140,6 @@ def register_experts(
         activation_format,
         supported_dtypes,
         blocked_quantization_support,
-        supports_expert_map,
         needs_matching_quant,
         needs_deep_gemm,
         needs_aiter,
@@ -176,7 +173,6 @@ register_experts(
     batched_format,
     common_float_types,
     blocked_quantization_support=True,
-    supports_expert_map=False,
     needs_matching_quant=True,
 )
 
@@ -185,7 +181,6 @@ register_experts(
     standard_format,
     common_float_and_int_types,
     blocked_quantization_support=True,
-    supports_expert_map=True,
     needs_matching_quant=True,
 )
 
@@ -194,7 +189,6 @@ register_experts(
     batched_format,
     common_float_and_int_types,
     blocked_quantization_support=True,
-    supports_expert_map=True,
 )
 
 # Disable on blackwell for now
@@ -260,7 +254,6 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
         nvfp4_types + fp8_types,
         blocked_quantization_support=True,
         # Note: this is a hack to get it to run for now
-        supports_expert_map=True,
     )
 else:
     FlashInferCutlassMoEPrepareAndFinalize = None
@@ -294,7 +287,6 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
         standard_format,
         nvfp4_types,
         blocked_quantization_support=False,
-        supports_expert_map=True,
     )
 
 if has_aiter():
@@ -307,7 +299,6 @@ if has_aiter():
         standard_format,
         fp8_types,
         blocked_quantization_support=True,
-        supports_expert_map=True,
         needs_aiter=True,
     )
 else:
@@ -319,7 +310,6 @@ if has_deep_gemm() and is_deep_gemm_supported():
         batched_format,
         fp8_types,
         blocked_quantization_support=True,
-        supports_expert_map=False,
         needs_matching_quant=False,
         needs_deep_gemm=True,
     )
@@ -328,7 +318,6 @@ if has_deep_gemm() and is_deep_gemm_supported():
         standard_format,
         fp8_types,
         blocked_quantization_support=True,
-        supports_expert_map=True,
         needs_matching_quant=False,
         needs_deep_gemm=True,
     )
@@ -337,7 +326,6 @@ if has_deep_gemm() and is_deep_gemm_supported():
         standard_format,
         common_float_and_int_types,
         blocked_quantization_support=True,
-        supports_expert_map=True,
         needs_matching_quant=True,
         needs_deep_gemm=True,
     )
@@ -353,14 +341,12 @@ if cutlass_fp8_supported():
         standard_format,
         fp8_types,
         blocked_quantization_support=False,
-        supports_expert_map=False,
     )
     register_experts(
         CutlassBatchedExpertsFp8,
         batched_format,
         fp8_types,
         blocked_quantization_support=False,
-        supports_expert_map=False,
     )
 else:
     CutlassBatchedExpertsFp8 = None
@@ -376,7 +362,6 @@ if cutlass_fp4_supported():
         standard_format,
         nvfp4_types,
         blocked_quantization_support=True,
-        supports_expert_map=False,
     )
 else:
     CutlassExpertsFp4 = None

tests/kernels/moe/test_modular_kernel_combinations.py

@@ -227,7 +227,7 @@ def is_nyi_config(config: Config) -> bool:
         ) == 1
         return unsupported_quant_config
 
-    return not info.supports_expert_map
+    return False
 
 
 def generate_valid_test_cases(

tests/kernels/test_fp32_router_gemm.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for fp32_router_gemm kernel: activation×weight→fp32, H=3072, E=256.
+
+Correctness baseline: torch.matmul in float64.
+"""
+
+import pytest
+import torch
+
+from vllm._custom_ops import fp32_router_gemm
+
+NUM_EXPERTS = 256
+HIDDEN_DIM = 3072
+# Absolute tolerance for fp32 kernel vs float64 reference
+ATOL_FP32 = 2e-4
+ATOL_BF16 = 2e-2  # bf16 activation has lower precision
+
+
+def _requires_sm90():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+    major, minor = torch.cuda.get_device_capability()
+    if major * 10 + minor < 90:
+        pytest.skip(f"fp32_router_gemm requires SM90+, got SM{major}{minor}")
+
+
+def _ref(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor:
+    """Reference: F.linear in float32 on GPU."""
+    return torch.nn.functional.linear(mat_a.float(), mat_b.float())
+
+
+@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
+def test_fp32_activation(num_tokens: int):
+    """fp32 activation → fp32 output should match reference closely."""
+    _requires_sm90()
+    torch.manual_seed(42)
+    device = torch.device("cuda")
+    mat_a = torch.randn(num_tokens, HIDDEN_DIM, dtype=torch.float32, device=device)
+    mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device)
+
+    out = fp32_router_gemm(mat_a, mat_b)
+    ref = _ref(mat_a, mat_b)
+
+    assert out.shape == (num_tokens, NUM_EXPERTS)
+    assert out.dtype == torch.float32
+    torch.testing.assert_close(out, ref, atol=ATOL_FP32, rtol=0)
+
+
+@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16, 32])
+def test_bf16_activation(num_tokens: int):
+    """bf16 activation → fp32 output should match reference within bf16 error."""
+    _requires_sm90()
+    torch.manual_seed(42)
+    device = torch.device("cuda")
+    mat_a_bf16 = torch.randn(
+        num_tokens, HIDDEN_DIM, dtype=torch.bfloat16, device=device
+    )
+    mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device)
+
+    out = fp32_router_gemm(mat_a_bf16, mat_b)
+    ref = _ref(mat_a_bf16, mat_b).to(device)
+
+    assert out.shape == (num_tokens, NUM_EXPERTS)
+    assert out.dtype == torch.float32
+    torch.testing.assert_close(out, ref, atol=ATOL_BF16, rtol=0)
+
+
+def test_output_shape_and_dtype():
+    """Basic shape and dtype checks."""
+    _requires_sm90()
+    device = torch.device("cuda")
+    mat_a = torch.randn(4, HIDDEN_DIM, dtype=torch.float32, device=device)
+    mat_b = torch.randn(NUM_EXPERTS, HIDDEN_DIM, dtype=torch.float32, device=device)
+    out = fp32_router_gemm(mat_a, mat_b)
+    assert out.shape == (4, NUM_EXPERTS)
+    assert out.dtype == torch.float32
+    assert out.device.type == "cuda"

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/models/multimodal/pooling/test_phi3v.py

@@ -8,6 +8,8 @@ from PIL import Image
 
 from vllm.assets.base import get_vllm_public_assets
 from vllm.assets.image import VLM_IMAGES_DIR
+from vllm.config import ModelConfig
+from vllm.multimodal import MULTIMODAL_REGISTRY
 
 from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
 from ....utils import large_gpu_test
@@ -37,6 +39,18 @@ HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
 
 MODELS = ["TIGER-Lab/VLM2Vec-Full"]
 
+SPECIAL_TOKEN_IMAGE_PROMPT = (
+    "\n<s><|user|>\n <|image_1|>\n\t <s>"
+    "Represent the given image for classification<|end|>"
+    "\n<|assistant|>\n"
+)
+
+
+def _get_cherry_blossom_image() -> Image.Image:
+    return Image.open(
+        get_vllm_public_assets(filename="cherry_blossom.jpg", s3_prefix=VLM_IMAGES_DIR)
+    )
+
 
 def _run_test(
     hf_runner: type[HfRunner],
@@ -123,19 +137,6 @@ def test_models_image(
     input_texts_images = [
         (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
     ]
-    # add cases for special_tokens
-    input_texts_images.append(
-        (
-            "\n<s><|user|>\n <|image_1|>\n\t <s>"
-            "Represent the given image for classification<|end|>"
-            "\n<|assistant|>\n",
-            Image.open(
-                get_vllm_public_assets(
-                    filename="cherry_blossom.jpg", s3_prefix=VLM_IMAGES_DIR
-                )
-            ),
-        )
-    )
     input_texts = [text for text, _ in input_texts_images]
     input_images = [image for _, image in input_texts_images]
 
@@ -147,3 +148,48 @@ def test_models_image(
         model,
         dtype=dtype,
     )
+
+
+@pytest.mark.core_model
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("dtype", ["half"])
+def test_models_image_special_tokens_processing(
+    model: str,
+    dtype: str,
+) -> None:
+    model_config = ModelConfig(
+        model,
+        runner="pooling",
+        trust_remote_code=True,
+        dtype=dtype,
+        max_model_len=1024,
+    )
+    processor = MULTIMODAL_REGISTRY.create_processor(model_config)
+    image = _get_cherry_blossom_image()
+
+    processed_inputs = processor(
+        SPECIAL_TOKEN_IMAGE_PROMPT,
+        mm_items=processor.info.parse_mm_data({"image": image}),
+        hf_processor_mm_kwargs={},
+    )
+
+    hf_processor = processor.info.get_hf_processor()
+    hf_inputs = hf_processor(
+        SPECIAL_TOKEN_IMAGE_PROMPT,
+        images=image,
+        return_tensors="pt",
+    )
+
+    image_token_id = hf_processor.get_special_image_token_id()
+    hf_prompt_token_ids = [
+        image_token_id if token_id < 0 else token_id
+        for token_id in hf_inputs["input_ids"][0].tolist()
+    ]
+
+    prompt_token_ids = processed_inputs["prompt_token_ids"]
+
+    assert prompt_token_ids == hf_prompt_token_ids
+    assert prompt_token_ids.count(image_token_id) == hf_prompt_token_ids.count(
+        image_token_id
+    )
+    assert prompt_token_ids.count(image_token_id) > 0

tests/models/multimodal/processing/test_tensor_schema.py

@@ -180,6 +180,7 @@ def test_model_tensor_schema(model_id: str):
         dummy_hf_overrides,
         model_arch=model_arch,
         exist_overrides=model_info.hf_overrides,
+        use_original_num_layers=getattr(model_info, "use_original_num_layers", False),
     )
 
     # ROCm: Detect if model uses AWQ quantization and set appropriate dtype

tests/models/quantization/test_gpt_oss.py

@@ -22,7 +22,14 @@ import pytest
 from packaging import version
 
 from vllm.platforms import current_platform
-from vllm.platforms.rocm import on_gfx950
+
+if current_platform.is_rocm():
+    from vllm.platforms.rocm import on_gfx950
+else:
+
+    def on_gfx950() -> bool:
+        return False
+
 
 MODEL_ACCURACIES = {
     # Full quantization: attention linears and MoE linears

tests/models/registry.py

@@ -1372,7 +1372,14 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         "stepfun-ai/Step3-VL-10B", trust_remote_code=True
     ),
     "Step3p7ForConditionalGeneration": _HfExamplesInfo(
-        "stepfun-ai/Step-3.7-Flash", is_available_online=False, trust_remote_code=True
+        "stepfun-ai/Step-3.7-Flash",
+        trust_remote_code=True,
+        use_original_num_layers=True,
+        # The MoE config lives in the nested ``text_config``, so the overrides
+        # must be nested too. Use 4 layers to initialize at least one MoE layer
+        # and shrink ``moe_num_experts`` (a non-standard key not handled by
+        # ``dummy_hf_overrides``) to avoid OOM during init.
+        hf_overrides={"text_config": {"num_hidden_layers": 4, "moe_num_experts": 8}},
     ),
     "UltravoxModel": _HfExamplesInfo(
         "fixie-ai/ultravox-v0_5-llama-3_2-1b",

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/_aiter_ops.py

@@ -2394,6 +2394,7 @@ class rocm_aiter_ops:
         alibi_slopes: torch.Tensor | None = None,
         return_lse: bool = False,
         out: torch.Tensor | None = None,
+        sink_ptr: torch.Tensor | None = None,
     ):
         """
         Flash attention with variable length sequences.
@@ -2422,6 +2423,7 @@ class rocm_aiter_ops:
             alibi_slopes=alibi_slopes,
             return_lse=return_lse,
             out=out,
+            sink_ptr=sink_ptr,
         )
 
     @staticmethod

vllm/_custom_ops.py

@@ -2412,6 +2412,31 @@ def dsv3_router_gemm(
     return output
 
 
+def fp32_router_gemm(
+    hidden_states: torch.Tensor,
+    router_weight: torch.Tensor,
+) -> torch.Tensor:
+    output = torch.empty(
+        hidden_states.shape[0],
+        router_weight.shape[0],
+        device=hidden_states.device,
+        dtype=torch.float32,
+    )
+    torch.ops._C.fp32_router_gemm(output, hidden_states, router_weight)
+    return output
+
+
+if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "fp32_router_gemm"):
+
+    @register_fake("_C::fp32_router_gemm")
+    def fp32_router_gemm_fake(
+        output: torch.Tensor,
+        mat_a: torch.Tensor,
+        mat_b: torch.Tensor,
+    ) -> None:
+        return
+
+
 def topk_softmax(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
@@ -2655,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/entrypoints/openai/chat_completion/serving.py

@@ -57,7 +57,6 @@ from vllm.entrypoints.openai.engine.serving import (
 )
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.openai.parser.harmony_utils import (
-    get_stop_tokens_for_assistant_actions,
     get_streamable_parser_for_assistant,
     parse_chat_output,
 )
@@ -158,13 +157,6 @@ class OpenAIServingChat(OpenAIServing):
             else getattr(mc, "override_generation_config", {}).get("max_new_tokens")
         )
         self.use_harmony = self.model_config.hf_config.model_type == "gpt_oss"
-        if self.use_harmony:
-            if "stop_token_ids" not in self.default_sampling_params:
-                self.default_sampling_params["stop_token_ids"] = []
-            self.default_sampling_params["stop_token_ids"].extend(
-                get_stop_tokens_for_assistant_actions()
-            )
-
         self.tool_call_id_type = get_tool_call_id_type(self.model_config)
 
         # NOTE(woosuk): While OpenAI's chat completion API supports browsing

vllm/entrypoints/openai/dp_supervisor.py

@@ -55,9 +55,9 @@ def validate_multi_port_external_lb_args(args: argparse.Namespace) -> None:
         raise ValueError(
             "Error: --data-parallel-multi-port-external-lb does not support --uds"
         )
-    if any((args.ssl_keyfile, args.ssl_certfile, args.ssl_ca_certs)):
+    if bool(args.ssl_keyfile) != bool(args.ssl_certfile):
         raise ValueError(
-            "Error: --data-parallel-multi-port-external-lb does not support HTTPS yet"
+            "Error: --ssl-keyfile and --ssl-certfile must be provided together"
         )
     if args.api_server_count not in (None, 1):
         raise ValueError(
@@ -151,7 +151,8 @@ def _child_base_url(args: argparse.Namespace, port: int) -> str:
         host = "127.0.0.1"
     elif host == "::":
         host = "::1"
-    return f"http://{host}:{port}"
+    scheme = "https" if args.ssl_keyfile and args.ssl_certfile else "http"
+    return f"{scheme}://{host}:{port}"
 
 
 def _join_processes_with_timeout(processes: list[BaseProcess], timeout: float) -> None:
@@ -178,7 +179,15 @@ async def _probe_endpoint(
     """
     for iteration in range(conn_err_failure_threshold):
         try:
-            async with session.get(_child_base_url(args, port) + path) as response:
+            probe_ssl = None
+            if args.ssl_keyfile and args.ssl_certfile:
+                # Probes target node-local child servers over loopback, so skip
+                # certificate verification to avoid SAN/hostname mismatches for
+                # localhost/127.0.0.1 deployments.
+                probe_ssl = False
+            async with session.get(
+                _child_base_url(args, port) + path, ssl=probe_ssl
+            ) as response:
                 # vLLM returns 503 on EngineDeadError, so we should return
                 # immediately if vLLM responds with a non-200 status code.
                 return response.status == HTTPStatus.OK
@@ -272,6 +281,11 @@ class DPSupervisor:
             host=host,
             port=self.supervisor_port,
             log_level=self.args.uvicorn_log_level,
+            ssl_keyfile=self.args.ssl_keyfile,
+            ssl_certfile=self.args.ssl_certfile,
+            ssl_ca_certs=self.args.ssl_ca_certs,
+            ssl_cert_reqs=self.args.ssl_cert_reqs,
+            ssl_ciphers=self.args.ssl_ciphers,
         )
         supervisor_server = uvicorn.Server(config)
         supervisor_server_task = asyncio.create_task(

vllm/entrypoints/openai/parser/harmony_utils.py

@@ -365,10 +365,6 @@ def render_for_completion(messages: list[Message]) -> list[int]:
     return token_ids
 
 
-def get_stop_tokens_for_assistant_actions() -> list[int]:
-    return get_encoding().stop_tokens_for_assistant_actions()
-
-
 def get_streamable_parser_for_assistant() -> StreamableParser:
     return StreamableParser(get_encoding(), role=Role.ASSISTANT)
 

vllm/entrypoints/openai/responses/protocol.py

@@ -372,8 +372,6 @@ class ResponsesRequest(OpenAIBaseModel):
         if (frequency_penalty := self.frequency_penalty) is None:
             frequency_penalty = default_sampling_params.get("frequency_penalty", 0.0)
 
-        stop_token_ids = default_sampling_params.get("stop_token_ids")
-
         # Structured output
         structured_outputs = self.structured_outputs
 
@@ -409,7 +407,6 @@ class ResponsesRequest(OpenAIBaseModel):
             top_k=top_k,
             max_tokens=max_tokens,
             logprobs=self.top_logprobs if self.is_include_output_logprobs() else None,
-            stop_token_ids=stop_token_ids,
             stop=stop,
             frequency_penalty=frequency_penalty,
             presence_penalty=presence_penalty,

vllm/entrypoints/openai/responses/serving.py

@@ -46,7 +46,6 @@ from vllm.entrypoints.openai.engine.serving import (
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.openai.parser.harmony_utils import (
     get_developer_message,
-    get_stop_tokens_for_assistant_actions,
     get_system_message,
     get_user_message,
     has_custom_tools,
@@ -222,13 +221,6 @@ class OpenAIServingResponses(OpenAIServing):
                 "For gpt-oss, we ignore --enable-auto-tool-choice "
                 "and always enable tool use."
             )
-            # OpenAI models have two EOS-like tokens: <|return|> and <|call|>.
-            # We need to add them to the stop token ids.
-            if "stop_token_ids" not in self.default_sampling_params:
-                self.default_sampling_params["stop_token_ids"] = []
-            self.default_sampling_params["stop_token_ids"].extend(
-                get_stop_tokens_for_assistant_actions()
-            )
 
         self.tool_call_id_type = get_tool_call_id_type(self.model_config)
 

vllm/model_executor/kernels/linear/__init__.py

@@ -61,6 +61,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
     XPUW4A8IntLinearKernel,
     XPUwNa16LinearKernel,
 )
+from vllm.model_executor.kernels.linear.mixed_precision.zentorch import (
+    ZentorchWNA16LinearKernel,
+)
 from vllm.model_executor.kernels.linear.mxfp4 import (
     MxFp4LinearKernel,
     MxFp4LinearLayerConfig,
@@ -87,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,
 )
@@ -160,6 +166,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.triton import (
 from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
     XPUFP8ScaledMMLinearKernel,
 )
+from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
+    ZentorchInt8ScaledMMLinearKernel,
+)
 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
 from vllm.platforms import PlatformEnum, current_platform
 
@@ -257,7 +266,7 @@ def _filter_kernels_by_backend(
 
 # in priority/performance order (when available)
 _POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]] = {
-    PlatformEnum.CPU: [CPUInt8ScaledMMLinearKernel],
+    PlatformEnum.CPU: [ZentorchInt8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel],
     PlatformEnum.CUDA: [
         CutlassInt8ScaledMMLinearKernel,
         TritonInt8ScaledMMLinearKernel,
@@ -353,6 +362,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
     ],
     PlatformEnum.CPU: [
         Dynamic4bitLinearKernel,
+        ZentorchWNA16LinearKernel,
         CPUWNA16LinearKernel,
     ],
 }
@@ -365,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: [
@@ -1023,6 +1036,8 @@ __all__ = [
     "RowWiseTorchFP8ScaledMMLinearKernel",
     "ROCmFP8ScaledMMLinearKernel",
     "TritonInt8ScaledMMLinearKernel",
+    "ZentorchInt8ScaledMMLinearKernel",
+    "ZentorchWNA16LinearKernel",
     "MPLinearKernel",
     "MPLinearLayerConfig",
     "AllSparkLinearKernel",

vllm/model_executor/kernels/linear/mixed_precision/__init__.py

@@ -39,6 +39,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
     XPUW4A8IntLinearKernel,
     XPUwNa16LinearKernel,
 )
+from vllm.model_executor.kernels.linear.mixed_precision.zentorch import (
+    ZentorchWNA16LinearKernel,
+)
 
 __all__ = [
     "MPLinearKernel",
@@ -55,4 +58,5 @@ __all__ = [
     "TritonW4A16LinearKernel",
     "XPUW4A8IntLinearKernel",
     "XPUwNa16LinearKernel",
+    "ZentorchWNA16LinearKernel",
 ]

vllm/model_executor/kernels/linear/mixed_precision/zentorch.py

@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Zentorch W4A16 GPTQ weight-only-quantized linear kernel for AMD Zen CPUs.
+
+Selected by ``choose_mp_linear_kernel`` ahead of the generic oneDNN-backed
+``CPUWNA16LinearKernel``. When ``can_implement`` rejects a layer, the selector
+falls through to the next kernel in ``_POSSIBLE_KERNELS[PlatformEnum.CPU]``.
+"""
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
+from vllm.platforms import current_platform
+from vllm.scalar_type import scalar_types
+
+from .cpu import CPUWNA16LinearKernel
+from .MPLinearKernel import MPLinearLayerConfig
+
+logger = init_logger(__name__)
+
+
+def _import_unpack_from_int32():
+    """Import compressed-tensors' ``unpack_from_int32`` across versions."""
+    try:
+        from compressed_tensors.compressors.pack_quantized.helpers import (
+            unpack_from_int32,
+        )
+    except ImportError:
+        from compressed_tensors.compressors.quantized_compressors.pack_quantized import (  # type: ignore[import-not-found]  # noqa: E501
+            unpack_from_int32,
+        )
+    return unpack_from_int32
+
+
+class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
+    """W4A16 GPTQ kernel backed by ``torch.ops.zentorch.zentorch_woq_linear``."""
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        ok, reason = super().can_implement(c)
+        if not ok:
+            return ok, reason
+
+        if not current_platform.is_zen_cpu():
+            return False, "ZentorchWNA16 requires an AMD Zen CPU."
+
+        if not has_zentorch_op(["zentorch_woq_repack_weight", "zentorch_woq_linear"]):
+            return (
+                False,
+                "torch.ops.zentorch.{zentorch_woq_repack_weight, "
+                "zentorch_woq_linear} are not registered.",
+            )
+
+        if c.has_g_idx:
+            return False, "ZentorchWNA16 does not support activation re-ordering."
+        return True, None
+
+    def _zentorch_woq_eligible(self, layer: torch.nn.Module) -> bool:
+        """Eligibility predicate for the zentorch W4A16 GPTQ fast path.
+
+        Constraints (any failure -> ``cpu_gemm_wna16`` path via ``super()``
+        with ``layer`` untouched).
+        """
+        if (
+            self.w_gidx_name is not None
+            and getattr(layer, self.w_gidx_name, None) is not None
+        ) or (getattr(self.config, "has_g_idx", False)):
+            return False
+
+        weight_packed = getattr(layer, self.w_q_name, None)
+        weight_scale = getattr(layer, self.w_s_name, None)
+        if weight_packed is None or weight_scale is None:
+            return False
+
+        bits = self.config.weight_type.mantissa
+        pack_factor = torch.iinfo(weight_packed.dtype).bits // bits
+        # 4-bit -> 8 values per int32;
+        if pack_factor != 8:
+            return False
+
+        # GPTQ-only. AWQ packs along the output dim instead.
+        in_dim = getattr(weight_packed, "input_dim", None)
+        pk_dim = getattr(weight_packed, "packed_dim", None)
+        if in_dim is None or pk_dim is None or in_dim != pk_dim:
+            return False
+
+        is_ct_format = in_dim == pk_dim == 1
+        if not is_ct_format:
+            return False
+
+        if weight_packed.dim() != 2 or weight_scale.dim() != 2:
+            return False
+
+        # 4-bit -> 8 values per int32; in_features must be divisible by num_groups.
+        in_features = weight_packed.shape[1] * 8
+        num_groups = weight_scale.shape[1]
+        return num_groups > 0 and in_features % num_groups == 0
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Repack CT GPTQ weights into the zentorch WOQ layout.
+
+        Falls back to ``CPUWNA16LinearKernel.process_weights_after_loading``
+        via ``super()`` when the layer doesn't satisfy
+        ``_zentorch_woq_eligible``.
+
+        On success, ``layer._zentorch_processed_weights`` is set to ``True``
+        """
+        if getattr(layer, "_zentorch_processed_weights", False):
+            return
+
+        if not self._zentorch_woq_eligible(layer):
+            logger.info_once(
+                "[zen_cpu] ZentorchWNA16 fast path not eligible for this "
+                "layer (AWQ pack layout, g_idx, or non-int32 storage); "
+                "falling back to CPUWNA16LinearKernel (cpu_gemm_wna16)."
+            )
+            super().process_weights_after_loading(layer)
+            return
+
+        if (not self.config.zero_points) and (self.w_zp_name is not None):
+            setattr(layer, self.w_zp_name, None)
+
+        if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
+            setattr(layer, self.w_gidx_name, None)
+
+        weight_q = getattr(layer, self.w_q_name)
+        weight_s = getattr(layer, self.w_s_name)
+        weight_packed = weight_q.data if hasattr(weight_q, "data") else weight_q
+        weight_scale = weight_s.data if hasattr(weight_s, "data") else weight_s
+
+        bits = self.config.weight_type.mantissa
+        pack_factor = torch.iinfo(weight_packed.dtype).bits // bits
+        out_features, num_groups = weight_scale.shape[0], weight_scale.shape[1]
+        in_features = weight_packed.shape[1] * pack_factor
+        original_shape = torch.Size([out_features, in_features])
+        unpack_from_int32 = _import_unpack_from_int32()
+        repack_op = torch.ops.zentorch.zentorch_woq_repack_weight.default
+
+        weight_unpacked = unpack_from_int32(
+            weight_packed,
+            bits,
+            original_shape,
+            packed_dim=weight_q.packed_dim,
+        )
+
+        zp_param = (
+            getattr(layer, self.w_zp_name, None) if self.w_zp_name is not None else None
+        )
+        needs_unsigned_offset = self.config.weight_type == scalar_types.uint4
+
+        if needs_unsigned_offset:
+            weight_unpacked = (weight_unpacked.to(torch.int32) + 8).clamp(0, 15)
+        repacked = repack_op(weight_unpacked.to(torch.int8).contiguous())
+
+        if zp_param is None:
+            zp_tc = None
+        else:
+            zp_tensor = zp_param.data if hasattr(zp_param, "data") else zp_param
+            zp = unpack_from_int32(
+                zp_tensor,
+                bits,
+                (out_features, num_groups),
+                packed_dim=zp_param.packed_dim,
+            )
+            if needs_unsigned_offset:
+                zp = (zp.to(torch.int32) + 8).clamp(0, 15)
+            zp_tc = zp.to(torch.int8).t().contiguous()
+
+        layer._zentorch_woq_packed = repacked.t()
+        layer._zentorch_woq_scale = weight_scale.t().contiguous()
+        layer._zentorch_woq_zero_point = zp_tc
+
+        for param_name in (self.w_q_name, self.w_s_name, self.w_zp_name):
+            if param_name is None:
+                continue
+            param = getattr(layer, param_name, None)
+            if param is None:
+                continue
+            if hasattr(param, "data"):
+                param.data = torch.empty(0)
+            else:
+                setattr(layer, param_name, torch.empty(0))
+
+        layer._zentorch_kind = "compressed_tensors_w4a16_gptq"
+        layer._zentorch_processed_weights = True
+        logger.info_once(
+            "[zen_cpu] Using zentorch_woq_linear for W4A16 GPTQ "
+            "(weight_type=%s, has_zp=%s)",
+            self.config.weight_type,
+            zp_tc is not None,
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if getattr(layer, "_zentorch_processed_weights", False):
+            return torch.ops.zentorch.zentorch_woq_linear.default(
+                x,
+                layer._zentorch_woq_packed,
+                layer._zentorch_woq_scale,
+                layer._zentorch_woq_zero_point,
+                bias,
+            )
+        return super().apply_weights(layer, x, bias)
+
+
+__all__ = ["ZentorchWNA16LinearKernel"]

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/kernels/linear/scaled_mm/__init__.py

@@ -39,6 +39,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
 from vllm.model_executor.kernels.linear.scaled_mm.triton import (
     TritonInt8ScaledMMLinearKernel,
 )
+from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
+    ZentorchInt8ScaledMMLinearKernel,
+)
 
 __all__ = [
     "FP8ScaledMMLinearKernel",
@@ -58,6 +61,7 @@ __all__ = [
     "RowWiseTorchFP8ScaledMMLinearKernel",
     "ROCmFP8ScaledMMLinearKernel",
     "TritonInt8ScaledMMLinearKernel",
+    "ZentorchInt8ScaledMMLinearKernel",
     "Fp8BlockScaledMMLinearKernel",
     "CPUFp8BlockScaledMMKernel",
 ]

vllm/model_executor/kernels/linear/scaled_mm/zentorch.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Zentorch dynamic-symmetric W8A8 int8 linear kernel for AMD Zen CPUs.
+
+Selected by ``choose_scaled_mm_linear_kernel`` ahead of the generic
+oneDNN-backed ``CPUInt8ScaledMMLinearKernel``. When ``is_supported`` or
+``can_implement`` rejects a layer, the selector falls through to the next
+kernel in ``_POSSIBLE_INT8_KERNELS[PlatformEnum.CPU]``.
+"""
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
+from vllm.model_executor.layers.quantization.utils import replace_parameter
+from vllm.platforms import current_platform
+
+from .ScaledMMLinearKernel import (
+    Int8ScaledMMLinearKernel,
+    Int8ScaledMMLinearLayerConfig,
+)
+
+logger = init_logger(__name__)
+
+
+class ZentorchInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cpu():
+            return False, "requires CPU."
+        if not current_platform.is_zen_cpu():
+            return False, "requires AMD Zen CPU."
+        if not has_zentorch_op(["zentorch_dynamic_qlinear"]):
+            return (
+                False,
+                "torch.ops.zentorch.zentorch_dynamic_qlinear is not registered.",
+            )
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        if c.is_static_input_scheme:
+            return False, "requires dynamic activation quantization."
+        if not c.input_symmetric:
+            return False, "requires symmetric activation quantization."
+        if not c.is_channelwise:
+            return False, "requires per-channel weight quantization."
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Prepare weights for ``zentorch_dynamic_qlinear``.
+
+        Keeps weight in [N, K] layout (int8, contiguous) and converts the
+        per-channel weight scale to bf16 with shape ``(N,)``.
+        """
+        w_q_name, w_s_name, _, _, _ = self.layer_param_names
+        weight = getattr(layer, w_q_name)
+        n = weight.shape[0]
+        replace_parameter(
+            layer,
+            w_q_name,
+            torch.nn.Parameter(weight.data.contiguous(), requires_grad=False),
+        )
+
+        weight_scale = getattr(layer, w_s_name)
+        ws = weight_scale.data
+        if ws.dim() == 2 and ws.shape[-1] == 1:
+            ws = ws.squeeze(-1)
+        ws = ws.to(torch.bfloat16).contiguous()
+        assert ws.shape == (n,), (
+            f"[zen_cpu] expected weight scale shape ({n},), got {tuple(ws.shape)}"
+        )
+
+        replace_parameter(
+            layer,
+            w_s_name,
+            torch.nn.Parameter(ws, requires_grad=False),
+        )
+        logger.info_once(
+            "[zen_cpu] Using zentorch_dynamic_qlinear for W8A8 (dynamic-symmetric)"
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        w_q_name, w_s_name, _, _, _ = self.layer_param_names
+        return torch.ops.zentorch.zentorch_dynamic_qlinear(
+            x,
+            getattr(layer, w_q_name),
+            getattr(layer, w_s_name),
+            bias,
+            zentorch_op_name="zentorch::zentorch_dynamic_qlinear",
+        )

vllm/model_executor/kernels/linear/zentorch_utils.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Gates zentorch CPU linear dispatch on platform/op availability."""
+
+from __future__ import annotations
+
+import torch
+
+from vllm.platforms import current_platform
+
+__all__ = ["has_zentorch_op"]
+
+
+def has_zentorch_op(op_names: list[str]) -> bool:
+    """Return ``True`` when running on Zen CPU with all named ops registered."""
+    if not op_names:
+        raise ValueError("has_zentorch_op requires at least one op name")
+    if not current_platform.is_zen_cpu():
+        return False
+    ns = getattr(torch.ops, "zentorch", None)
+    if ns is None:
+        return False
+    return all(hasattr(ns, op_name) for op_name in op_names)

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/aiter_mxfp4_w4a8_moe.py

@@ -248,9 +248,6 @@ class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
     ) -> bool:
         return True
 
-    def supports_expert_map(self) -> bool:
-        return False  # Expert parallelism not yet supported
-
     @property
     def expects_unquantized_inputs(self) -> bool:
         return True

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

@@ -316,9 +316,6 @@ class BatchedDeepGemmExperts(mk.FusedMoEExpertsModular):
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
         return True
 
-    def supports_expert_map(self) -> bool:
-        return False
-
     def supports_packed_ue8m0_act_scales(self) -> bool:
         """
         DeepGemm supports packed ue8m0 activation scales format in devices == sm100

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

@@ -100,9 +100,6 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic):
     ) -> bool:
         return True
 
-    def supports_expert_map(self) -> bool:
-        return False
-
     def apply(
         self,
         hidden_states: torch.Tensor,
@@ -256,9 +253,6 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic):
     ) -> bool:
         return True
 
-    def supports_expert_map(self) -> bool:
-        return False
-
     def apply(
         self,
         hidden_states: torch.Tensor,

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

@@ -378,7 +378,8 @@ class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular):
             topk_ids,
             activation,
             global_num_experts,
-            expert_map,
+            # the fp8 cutlass experts use their own expert map.
+            None,
             self.w1_scale,
             self.w2_scale,
             a1q_scale,
@@ -418,9 +419,6 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):
             or moe_parallel_config.use_fi_nvl_one_sided_kernels
         )
 
-    def supports_expert_map(self) -> bool:
-        return False
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         # topk weights and reduction are fused in moe_unpermute cuda kernel
         return TopKWeightAndReduceNoOP()
@@ -460,9 +458,6 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):
     def activation_format() -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.BatchedExperts
 
-    def supports_expert_map(self) -> bool:
-        return False
-
     def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
         return self.out_dtype if self.out_dtype is not None else act_dtype
 
@@ -741,9 +736,6 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
     def activation_format() -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
 
-    def supports_expert_map(self) -> bool:
-        return False
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 
@@ -1038,9 +1030,6 @@ class CutlassExpertsMxfp4(mk.FusedMoEExpertsModular):
     def activation_format() -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
 
-    def supports_expert_map(self) -> bool:
-        return False
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 
@@ -1340,9 +1329,6 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
         return True
 
-    def supports_expert_map(self) -> bool:
-        return True
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         # topk weights and reduction are fused in moe_unpermute cuda kernel
         return TopKWeightAndReduceNoOP()

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:
@@ -164,9 +192,6 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
             or moe_parallel_config.use_fi_nvl_one_sided_kernels
         )
 
-    def supports_expert_map(self) -> bool:
-        return True
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 
@@ -182,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
         )
@@ -204,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
@@ -224,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.
@@ -295,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)
@@ -313,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:
@@ -388,9 +439,6 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
             or moe_parallel_config.use_fi_nvl_one_sided_kernels
         )
 
-    def supports_expert_map(self) -> bool:
-        return True
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 

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

@@ -92,16 +92,6 @@ class FallbackExperts(mk.FusedMoEExpertsModular, ABC):
             moe_parallel_config
         ) and fallback_cls._supports_parallel_config(moe_parallel_config)
 
-    def supports_expert_map(self) -> bool:
-        assert (
-            self.experts.supports_expert_map()
-            == self.fallback_experts.supports_expert_map()
-        )
-        return (
-            self.experts.supports_expert_map()
-            and self.fallback_experts.supports_expert_map()
-        )
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         e_war = self.experts.finalize_weight_and_reduce_impl()
         fbe_war = self.fallback_experts.finalize_weight_and_reduce_impl()