Merge branch 'main' into wentao-optimize-per-token-group-quant

.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

@@ -1275,10 +1275,12 @@ steps:
   - vllm/
   - tests/entrypoints/openai
   - tests/entrypoints/test_chat_utils
+  - tests/entrypoints/generate
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
   - pytest -v -s entrypoints/test_chat_utils.py
+  - pytest -v -s entrypoints/generate
 
 - label: Entrypoints Integration (API Server openai - Part 3) # TBD
   timeout_in_minutes: 180
@@ -1368,7 +1370,7 @@ steps:
   - vllm/platforms/rocm.py
   commands:
   - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling  --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
 
 - label: OpenAI API correctness # TBD
   timeout_in_minutes: 180
@@ -2782,10 +2784,12 @@ steps:
   - vllm/
   - tests/entrypoints/openai
   - tests/entrypoints/test_chat_utils
+  - tests/entrypoints/generate
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
   - pytest -v -s entrypoints/test_chat_utils.py
+  - pytest -v -s entrypoints/generate
 
 - label: Entrypoints Integration (API Server openai - Part 3) # TBD
   timeout_in_minutes: 180

.buildkite/test_areas/entrypoints.yaml

@@ -11,7 +11,7 @@ steps:
   - tests/entrypoints/
   commands:
   - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
 
 - label: Entrypoints Integration (LLM)
   key: entrypoints-integration-llm
@@ -60,9 +60,11 @@ steps:
   - vllm/
   - tests/entrypoints/openai
   - tests/entrypoints/test_chat_utils
+  - tests/entrypoints/generate
   commands:
   - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
   - pytest -v -s entrypoints/test_chat_utils.py
+  - pytest -v -s entrypoints/generate
   mirror:
     amd:
       device: mi325_1

.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
@@ -78,7 +84,7 @@
 /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
@@ -683,6 +683,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 +1256,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/cpu_extension.cmake

@@ -396,6 +396,13 @@ set(VLLM_EXT_SRC
     "csrc/cpu/cpu_attn.cpp"
     "csrc/cpu/torch_bindings.cpp")
 
+if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
+        VLLM_RVV_VLEN GREATER 0 AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
+    set(VLLM_EXT_SRC
+        "csrc/cpu/cpu_wna16.cpp"
+        ${VLLM_EXT_SRC})
+endif()
+
 if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
     set(VLLM_EXT_SRC
         "csrc/cpu/shm.cpp"

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/cpu/cpu_types_riscv_impl.hpp

@@ -94,6 +94,10 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
       : reg(RVVI(__riscv_vle16_v_f16, LMUL_256)(
             static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
 
+  explicit FP16Vec16(const c10::Half v)
+      : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _f16, LMUL_256)(
+            RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
+
   explicit FP16Vec16(const FP32Vec16& vec);
 
   void save(void* ptr) const {
@@ -165,6 +169,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
             reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
 
   explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {};
+  explicit BF16Vec16(const c10::BFloat16 v)
+      : reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16, LMUL_256)(
+            RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
   explicit BF16Vec16(const FP32Vec16&);
 
   void save(void* ptr) const {
@@ -290,6 +297,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
     }
     reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
   }
+  explicit BF16Vec16(const c10::BFloat16 v)
+      : reg_fp32(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(static_cast<float>(v),
+                                                      VEC_ELEM_NUM)) {}
   explicit BF16Vec16(const FP32Vec16&);
   void save(void* ptr) const {
     float tmp[16];
@@ -629,6 +639,19 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
       : reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)(
             data.reg, data.reg)) {};
   explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
+  explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
+    const uint64_t q_values = static_cast<uint64_t>(value);
+    auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM);
+    auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM);
+    auto shifts =
+        RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM);
+    auto shifted =
+        RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM);
+    auto idx64 =
+        RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM);
+    auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM);
+    reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM);
+  }
   explicit FP32Vec16(const FP16Vec16& v);
 
 #ifdef __riscv_zvfbfmin
@@ -641,6 +664,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
   explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
 #endif
 
+  // FP8 stub: dead code on RISC-V (fp8 KV cache is x86-only), needed for
+  // load_b_pair_vec template to compile on all platforms.
+  explicit FP32Vec16(const BF16Vec32&, int) : FP32Vec16() {}
+
   FP32Vec16 operator+(const FP32Vec16& b) const {
     return FP32Vec16(
         RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
@@ -891,6 +918,30 @@ inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
   acc = acc.fma(a, b);
 }
 
+template <typename VecT>
+static void interleave_save_16b(const VecT& vec0, const VecT& vec1, void* ptr) {
+  alignas(64) uint16_t values0[VecT::VEC_ELEM_NUM];
+  alignas(64) uint16_t values1[VecT::VEC_ELEM_NUM];
+  vec0.save(values0);
+  vec1.save(values1);
+
+  auto* packed = reinterpret_cast<uint32_t*>(ptr);
+  for (int32_t i = 0; i < VecT::VEC_ELEM_NUM; ++i) {
+    packed[i] = static_cast<uint32_t>(values0[i]) |
+                (static_cast<uint32_t>(values1[i]) << 16);
+  }
+}
+
+static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
+                            void* ptr) {
+  interleave_save_16b(vec0, vec1, ptr);
+}
+
+static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
+                            void* ptr) {
+  interleave_save_16b(vec0, vec1, ptr);
+}
+
 #ifdef __riscv_zvfbfmin
 template <>
 inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {

csrc/cpu/torch_bindings.cpp

@@ -518,7 +518,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
 
   // WNA16
-#if defined(__AVX512F__)
+#if defined(__AVX512F__) || defined(__riscv_v)
   ops.def(
       "cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
       "Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt "

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/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/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, "

docs/design/attention_backends.md

@@ -177,7 +177,7 @@ Priority is **1 = highest** (tried first).
 | `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 |

docs/features/multimodal_inputs.md

@@ -778,7 +778,7 @@ Then, you can use the OpenAI client as follows:
         base_url=openai_api_base,
     )
 
-    video_url = "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerFun.mp4"
+    video_url = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
 
     ## Use video url in the payload
     chat_completion_from_url = client.chat.completions.create(

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
 

docs/models/supported_models.md

@@ -437,6 +437,7 @@ th {
 | `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ |
 | `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ |
 | `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
+| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ |
 | `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
 | `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
 | `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |

docs/training/weight_transfer/nccl.md

@@ -84,7 +84,10 @@ Both the trainer (`NCCLTrainerSendWeightsArgs`) and inference side (`NCCLWeightT
 
 ## Receiving Weights (Inference Side)
 
-The inference side triggers weight reception using the four-phase protocol — `init_weight_transfer_engine`, `start_weight_update`, `update_weights`, `finish_weight_update`. The init phase is shown [above](#initialization). The remaining three steps are:
+The inference side triggers weight reception using the four-phase protocol:
+`init_weight_transfer_engine`, `start_weight_update`, `update_weights`,
+`finish_weight_update`. The init phase is shown [above](#initialization). The
+remaining three steps are:
 
 ```python
 from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
@@ -108,12 +111,24 @@ llm.update_weights(
 llm.finish_weight_update()
 ```
 
-The `names`, `dtype_names`, and `shapes` lists describe each parameter. These must match the order in which the trainer iterates over its parameters.
+The `names`, `dtype_names`, and `shapes` lists describe each parameter. These
+must match the order in which the trainer iterates over its parameters.
 
-`start_weight_update` must be called before `update_weights`, and `finish_weight_update` must be called after all weight chunks have been transferred. The `is_checkpoint_format` flag controls whether layerwise reload processing is applied (`True` for checkpoint-format weights, `False` for pre-processed kernel-format weights).
+`start_weight_update` must be called before `update_weights`, and
+`finish_weight_update` must be called after all weight chunks have been
+transferred. The `is_checkpoint_format` flag controls whether layerwise reload
+processing is applied (`True` for checkpoint-format weights, `False` for
+pre-processed kernel-format weights).
+
+Sparse NCCL patches still use `update_kind="sparse_flat"` inside
+`update_info`, but they should be wrapped in
+`start_weight_update(is_checkpoint_format=False)` because sparse patches apply
+directly to runtime/kernel-format parameters. The current sparse MVP requires
+`TP=1` and `PP=1`.
 
 ## Examples
 
 - [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
+- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `start_weight_update(is_checkpoint_format=False)` and currently require `TP=1` and `PP=1`
 - [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
 - [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane

examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py

@@ -203,7 +203,7 @@ def run_multi_image(model: str, max_completion_tokens: int) -> None:
 
 # Video input inference
 def run_video(model: str, max_completion_tokens: int) -> None:
-    video_url = "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerFun.mp4"
+    video_url = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
     video_base64 = encode_base64_content_from_url(video_url)
 
     ## Use video url in the payload

examples/rl/rlhf_sparse_nccl.py

@@ -0,0 +1,526 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Demonstrates dense-vs-sparse NCCL weight syncing with a real model.
+
+This example mirrors the validation story used for the sparse NCCL MVP:
+both the dense update path and the sparse patch path start from the same real
+checkpoint and apply the same deterministic trainer-side patch. The script then
+checks that greedy 1-token outputs match between the dense and sparse vLLM
+engines after the update.
+
+The example performs the following steps:
+* Load a training model on one GPU via a Ray actor.
+* Launch a vLLM engine with the same real model on a second GPU.
+* Verify trainer vs vLLM baseline agreement before any update.
+* Apply a deterministic patch to ``model.embed_tokens.weight`` on the trainer.
+* Run a dense NCCL update into a fresh vLLM engine and collect post-update
+  outputs.
+* Reset the trainer back to the baseline checkpoint.
+* Apply the same deterministic patch again.
+* Run a sparse NCCL update into another fresh vLLM engine and collect
+  post-update outputs.
+* Compare dense vs sparse baseline outputs, dense vs sparse post-update
+  outputs, estimated payload sizes, and trainer-side send times.
+
+Current sparse weight transfer MVP limitations:
+* ``TP=1`` and ``PP=1`` only
+* sparse updates use runtime/kernel-format parameter names
+* sparse updates are not composable with checkpoint-format or packed updates
+
+This example assumes a single-node cluster with two GPUs.
+"""
+
+import hashlib
+import os
+import time
+from collections.abc import Sequence
+
+import ray
+import torch
+from ray.util.placement_group import placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from vllm import LLM, SamplingParams
+from vllm.config import WeightTransferConfig
+from vllm.distributed.weight_transfer.base import SparseWeightPatch
+from vllm.distributed.weight_transfer.nccl_engine import (
+    NCCLTrainerSendWeightsArgs,
+    NCCLWeightTransferEngine,
+)
+from vllm.utils.network_utils import get_ip, get_open_port
+
+MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
+PATCHED_PARAM_NAME = "model.embed_tokens.weight"
+MAX_PATCH_ROWS = 32
+PROMPTS = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+SAMPLING_PARAMS = SamplingParams(temperature=0.0, max_tokens=1)
+
+
+class MyLLM(LLM):
+    """Configure the vLLM worker for Ray placement group execution."""
+
+    def __init__(self, *args, **kwargs):
+        os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
+        super().__init__(*args, **kwargs)
+
+
+@ray.remote(num_gpus=1)
+class TrainModel:
+    """Ray actor that owns the trainer-side model and deterministic patch state."""
+
+    def __init__(self, model_name: str):
+        self.model_name = model_name
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        self.model = None
+        self.patched_param = None
+        self.pending_sparse_patches: list[SparseWeightPatch] | None = None
+        self.model_update_group = None
+        self.master_address = get_ip()
+        self.port = get_open_port()
+        self.reset_model()
+
+    def reset_model(self) -> None:
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_name,
+            torch_dtype=torch.bfloat16,
+        ).to("cuda:0")
+        self.model.eval()
+
+        try:
+            self.patched_param = self.model.get_parameter(PATCHED_PARAM_NAME)
+        except AttributeError as exc:
+            raise RuntimeError(
+                f"Expected trainer model to expose `{PATCHED_PARAM_NAME}`"
+            ) from exc
+
+        self.pending_sparse_patches = None
+
+    def create_rendezvous(self) -> tuple[str, int]:
+        self.port = get_open_port()
+        return self.master_address, self.port
+
+    def init_weight_transfer_group(self, world_size: int) -> None:
+        self.model_update_group = NCCLWeightTransferEngine.trainer_init(
+            dict(
+                master_address=self.master_address,
+                master_port=self.port,
+                world_size=world_size,
+            )
+        )
+
+    def get_dense_update_info(self, packed: bool = False) -> tuple[dict, int]:
+        names = []
+        dtype_names = []
+        shapes = []
+        payload_bytes = 0
+        for name, param in self.model.named_parameters():
+            names.append(name)
+            dtype_names.append(str(param.dtype).split(".")[-1])
+            shapes.append(list(param.shape))
+            payload_bytes += param.numel() * param.element_size()
+
+        return (
+            dict(
+                names=names,
+                dtype_names=dtype_names,
+                shapes=shapes,
+                packed=packed,
+            ),
+            payload_bytes,
+        )
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        prompts: Sequence[str],
+        max_new_tokens: int = 1,
+    ) -> list[dict[str, object]]:
+        generations = []
+        for prompt in prompts:
+            model_inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda:0")
+            output = self.model.generate(
+                **model_inputs,
+                max_new_tokens=max_new_tokens,
+                do_sample=False,
+                pad_token_id=self.tokenizer.pad_token_id,
+            )
+            new_token_ids = output[0, model_inputs["input_ids"].shape[1] :].tolist()
+            generations.append(
+                {
+                    "token_ids": new_token_ids,
+                    "text": self.tokenizer.decode(
+                        new_token_ids,
+                        skip_special_tokens=False,
+                    ),
+                }
+            )
+        return generations
+
+    def prepare_sparse_patch(
+        self,
+        prompts: Sequence[str],
+        max_patch_rows: int = MAX_PATCH_ROWS,
+    ) -> tuple[dict[str, object], list[int], str, int]:
+        selected_token_ids: list[int] = []
+        special_ids = set(self.tokenizer.all_special_ids)
+        for prompt in prompts:
+            token_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"]
+            for token_id in token_ids:
+                if token_id in special_ids or token_id in selected_token_ids:
+                    continue
+                selected_token_ids.append(token_id)
+                if len(selected_token_ids) == max_patch_rows:
+                    break
+            if len(selected_token_ids) == max_patch_rows:
+                break
+
+        if not selected_token_ids:
+            raise ValueError("Could not derive any non-special token IDs to patch")
+
+        vocab_size = self.patched_param.shape[0]
+        next_token_id = selected_token_ids[-1]
+        while len(selected_token_ids) < max_patch_rows:
+            next_token_id = (next_token_id + 1) % vocab_size
+            if next_token_id in special_ids or next_token_id in selected_token_ids:
+                continue
+            selected_token_ids.append(next_token_id)
+
+        row_ids = torch.tensor(
+            selected_token_ids,
+            device=self.patched_param.device,
+            dtype=torch.long,
+        )
+        hidden_size = self.patched_param.shape[1]
+        column_offsets = torch.arange(
+            hidden_size,
+            device=self.patched_param.device,
+            dtype=torch.long,
+        )
+
+        with torch.no_grad():
+            # Rotate the selected embedding rows instead of zeroing them so the
+            # patch remains deterministic while avoiding a degenerate collapse
+            # to the same special token after the update.
+            replacement_rows = self.patched_param[row_ids].roll(shifts=1, dims=0)
+            self.patched_param[row_ids] = replacement_rows
+
+        flat_indices = (
+            row_ids.unsqueeze(1).mul(hidden_size).add(column_offsets).reshape(-1)
+        )
+        flat_values = self.patched_param[row_ids].reshape(-1).contiguous()
+        self.pending_sparse_patches = [
+            SparseWeightPatch(
+                name=PATCHED_PARAM_NAME,
+                indices=flat_indices.to(torch.int32),
+                values=flat_values,
+            )
+        ]
+        patch_digest = hashlib.sha256(
+            self.pending_sparse_patches[0].indices.cpu().numpy().tobytes()
+            + self.pending_sparse_patches[0]
+            .values.detach()
+            .float()
+            .cpu()
+            .numpy()
+            .tobytes()
+        ).hexdigest()
+
+        sparse_payload_bytes = (
+            flat_indices.numel() * torch.tensor([], dtype=torch.int32).element_size()
+            + flat_values.numel() * flat_values.element_size()
+        )
+        update_info = dict(
+            names=[PATCHED_PARAM_NAME],
+            dtype_names=[str(self.patched_param.dtype).split(".")[-1]],
+            shapes=[list(self.patched_param.shape)],
+            num_updates_list=[flat_indices.numel()],
+            update_kind="sparse_flat",
+        )
+        return update_info, selected_token_ids, patch_digest, sparse_payload_bytes
+
+    def broadcast_weights(self, packed: bool = False) -> float:
+        if self.model_update_group is None:
+            raise RuntimeError("Weight transfer group is not initialized")
+
+        trainer_args = NCCLTrainerSendWeightsArgs(
+            group=self.model_update_group,
+            packed=packed,
+        )
+        start = time.perf_counter()
+        NCCLWeightTransferEngine.trainer_send_weights(
+            iterator=self.model.named_parameters(),
+            trainer_args=trainer_args,
+        )
+        torch.accelerator.synchronize()
+        return (time.perf_counter() - start) * 1000.0
+
+    def broadcast_pending_sparse_patch(self) -> float:
+        if self.model_update_group is None:
+            raise RuntimeError("Weight transfer group is not initialized")
+        if self.pending_sparse_patches is None:
+            raise RuntimeError("Sparse patch has not been prepared")
+
+        start = time.perf_counter()
+        NCCLWeightTransferEngine.trainer_send_sparse_weights(
+            iter(self.pending_sparse_patches),
+            NCCLTrainerSendWeightsArgs(group=self.model_update_group),
+        )
+        torch.accelerator.synchronize()
+        self.pending_sparse_patches = None
+        return (time.perf_counter() - start) * 1000.0
+
+
+def launch_llm(
+    scheduling_inference: PlacementGroupSchedulingStrategy,
+):
+    return ray.remote(
+        num_cpus=0,
+        num_gpus=0,
+        scheduling_strategy=scheduling_inference,
+    )(MyLLM).remote(
+        model=MODEL_NAME,
+        enforce_eager=True,
+        tensor_parallel_size=1,
+        distributed_executor_backend="ray",
+        gpu_memory_utilization=0.7,
+        weight_transfer_config=WeightTransferConfig(backend="nccl"),
+    )
+
+
+def collect_vllm_generations(llm_handle) -> list[dict[str, object]]:
+    outputs = ray.get(llm_handle.generate.remote(PROMPTS, SAMPLING_PARAMS))
+    generations = []
+    for output in outputs:
+        generations.append(
+            {
+                "token_ids": output.outputs[0].token_ids,
+                "text": output.outputs[0].text,
+            }
+        )
+    return generations
+
+
+def token_sequences_match(
+    left: Sequence[dict[str, object]],
+    right: Sequence[dict[str, object]],
+) -> bool:
+    return [item["token_ids"] for item in left] == [item["token_ids"] for item in right]
+
+
+def print_generations(label: str, prompts: Sequence[str], generations) -> None:
+    print(f"\n{label}")
+    print("-" * 50)
+    for prompt, generation in zip(prompts, generations):
+        print(f"Prompt: {prompt!r}")
+        print(f"Token IDs: {generation['token_ids']}")
+        print(f"Text: {generation['text']!r}")
+        print("-" * 50)
+
+
+def run_dense_phase(
+    train_model,
+    scheduling_inference: PlacementGroupSchedulingStrategy,
+) -> dict[str, object]:
+    ray.get(train_model.reset_model.remote())
+    llm = launch_llm(scheduling_inference)
+    try:
+        dense_before = collect_vllm_generations(llm)
+
+        ray.get(llm.sleep.remote(level=0))
+        master_address, master_port = ray.get(train_model.create_rendezvous.remote())
+        world_size = ray.get(llm.get_world_size.remote()) + 1
+        inference_init = llm.init_weight_transfer_engine.remote(
+            dict(
+                init_info=dict(
+                    master_address=master_address,
+                    master_port=master_port,
+                    rank_offset=1,
+                    world_size=world_size,
+                )
+            )
+        )
+        trainer_init = train_model.init_weight_transfer_group.remote(world_size)
+        ray.get([trainer_init, inference_init])
+        ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
+
+        dense_update_info, dense_payload_bytes = ray.get(
+            train_model.get_dense_update_info.remote()
+        )
+        _, selected_token_ids, patch_digest, _ = ray.get(
+            train_model.prepare_sparse_patch.remote(PROMPTS)
+        )
+
+        inference_update = llm.update_weights.remote(
+            dict(update_info=dense_update_info)
+        )
+        dense_send_ms, _ = ray.get(
+            [
+                train_model.broadcast_weights.remote(packed=False),
+                inference_update,
+            ]
+        )
+        ray.get(llm.finish_weight_update.remote())
+        ray.get(llm.wake_up.remote(tags=["scheduling"]))
+
+        dense_after = collect_vllm_generations(llm)
+
+        return {
+            "dense_before": dense_before,
+            "dense_after": dense_after,
+            "selected_token_ids": selected_token_ids,
+            "patch_digest": patch_digest,
+            "dense_payload_bytes": dense_payload_bytes,
+            "dense_send_ms": dense_send_ms,
+        }
+    finally:
+        ray.kill(llm)
+
+
+def run_sparse_phase(
+    train_model,
+    scheduling_inference: PlacementGroupSchedulingStrategy,
+) -> dict[str, object]:
+    ray.get(train_model.reset_model.remote())
+    llm = launch_llm(scheduling_inference)
+    try:
+        sparse_before = collect_vllm_generations(llm)
+
+        ray.get(llm.sleep.remote(level=0))
+        master_address, master_port = ray.get(train_model.create_rendezvous.remote())
+        world_size = ray.get(llm.get_world_size.remote()) + 1
+        inference_init = llm.init_weight_transfer_engine.remote(
+            dict(
+                init_info=dict(
+                    master_address=master_address,
+                    master_port=master_port,
+                    rank_offset=1,
+                    world_size=world_size,
+                )
+            )
+        )
+        trainer_init = train_model.init_weight_transfer_group.remote(world_size)
+        ray.get([trainer_init, inference_init])
+        ray.get(llm.start_weight_update.remote(is_checkpoint_format=False))
+
+        sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = (
+            ray.get(train_model.prepare_sparse_patch.remote(PROMPTS))
+        )
+
+        inference_update = llm.update_weights.remote(
+            dict(update_info=sparse_update_info)
+        )
+        sparse_send_ms, _ = ray.get(
+            [
+                train_model.broadcast_pending_sparse_patch.remote(),
+                inference_update,
+            ]
+        )
+        ray.get(llm.finish_weight_update.remote())
+        ray.get(llm.wake_up.remote(tags=["scheduling"]))
+
+        sparse_after = collect_vllm_generations(llm)
+
+        return {
+            "sparse_before": sparse_before,
+            "sparse_after": sparse_after,
+            "selected_token_ids": selected_token_ids,
+            "patch_digest": patch_digest,
+            "sparse_payload_bytes": sparse_payload_bytes,
+            "sparse_send_ms": sparse_send_ms,
+        }
+    finally:
+        ray.kill(llm)
+
+
+ray.init()
+
+try:
+    train_model = TrainModel.remote(MODEL_NAME)
+
+    pg_inference = placement_group([{"GPU": 1, "CPU": 0}])
+    ray.get(pg_inference.ready())
+    scheduling_inference = PlacementGroupSchedulingStrategy(
+        placement_group=pg_inference,
+        placement_group_capture_child_tasks=True,
+        placement_group_bundle_index=0,
+    )
+
+    dense_results = run_dense_phase(train_model, scheduling_inference)
+    sparse_results = run_sparse_phase(train_model, scheduling_inference)
+
+    baseline_equal = token_sequences_match(
+        dense_results["dense_before"],
+        sparse_results["sparse_before"],
+    )
+    patch_selection_equal = (
+        dense_results["selected_token_ids"] == sparse_results["selected_token_ids"]
+    )
+    patch_digest_equal = dense_results["patch_digest"] == sparse_results["patch_digest"]
+    after_equal = token_sequences_match(
+        dense_results["dense_after"],
+        sparse_results["sparse_after"],
+    )
+    any_output_changed = any(
+        before["token_ids"] != after["token_ids"]
+        for before, after in zip(
+            dense_results["dense_before"],
+            dense_results["dense_after"],
+        )
+    )
+    dense_payload_mb = dense_results["dense_payload_bytes"] / (1024 * 1024)
+    sparse_payload_mb = sparse_results["sparse_payload_bytes"] / (1024 * 1024)
+
+    print_generations(
+        "Dense baseline outputs",
+        PROMPTS,
+        dense_results["dense_before"],
+    )
+    print_generations(
+        "Sparse baseline outputs", PROMPTS, sparse_results["sparse_before"]
+    )
+    print_generations(
+        "Dense outputs after update", PROMPTS, dense_results["dense_after"]
+    )
+    print_generations(
+        "Sparse outputs after update",
+        PROMPTS,
+        sparse_results["sparse_after"],
+    )
+
+    print(f"patched_token_ids = {dense_results['selected_token_ids']}")
+    print(f"patch_selection_equal = {patch_selection_equal}")
+    print(f"dense_patch_digest = {dense_results['patch_digest']}")
+    print(f"sparse_patch_digest = {sparse_results['patch_digest']}")
+    print(f"patch_digest_equal = {patch_digest_equal}")
+    print(f"baseline_equal = {baseline_equal}")
+    print(f"after_equal = {after_equal}")
+    print(f"any_output_changed = {any_output_changed}")
+    print(f"dense_payload_mb = {dense_payload_mb:.2f}")
+    print(f"sparse_payload_mb = {sparse_payload_mb:.2f}")
+    print(f"dense_send_ms = {dense_results['dense_send_ms']:.2f}")
+    print(f"sparse_send_ms = {sparse_results['sparse_send_ms']:.2f}")
+
+    if not baseline_equal:
+        raise RuntimeError(
+            "Dense and sparse phases did not start from the same baseline"
+        )
+    if not patch_selection_equal:
+        raise RuntimeError("Dense and sparse phases used different sparse patches")
+    if not patch_digest_equal:
+        raise RuntimeError("Dense and sparse phases produced different patch values")
+    if not after_equal:
+        raise RuntimeError("Dense and sparse updates produced different outputs")
+    if not any_output_changed:
+        raise RuntimeError("Patch did not change the observed outputs")
+finally:
+    ray.shutdown()

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, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
     }
 
     #[test]

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,
+    Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser,
+    MistralToolParser, Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser,
+    ToolParserError, ToolParserOutput,
 };
 
 use crate::parser::ParserFactory;
@@ -24,6 +24,9 @@ pub mod names {
     pub const GEMMA4: &str = "gemma4";
     pub const HERMES: &str = "hermes";
     pub const HY_V3: &str = "hy_v3";
+    // Matches the Python CLI name `--tool-call-parser internlm`, which Python
+    // also routes to `Internlm2ToolParser` despite the version-agnostic name.
+    pub const INTERNLM: &str = "internlm";
     pub const KIMI_K2: &str = "kimi_k2";
     pub const LLAMA3_JSON: &str = "llama3_json";
     pub const LLAMA4_JSON: &str = "llama4_json";
@@ -62,6 +65,7 @@ impl ToolParserFactory {
             .register_parser::<Gemma4ToolParser>(names::GEMMA4)
             .register_parser::<HermesToolParser>(names::HERMES)
             .register_parser::<HyV3ToolParser>(names::HY_V3)
+            .register_parser::<Internlm2ToolParser>(names::INTERNLM)
             .register_parser::<KimiK2ToolParser>(names::KIMI_K2)
             .register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
             .register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
@@ -80,6 +84,12 @@ impl ToolParserFactory {
             .register_pattern("hermes", names::HERMES)
             .register_pattern("hy3", names::HY_V3)
             .register_pattern("hy_v3", names::HY_V3)
+            // Narrow to `internlm2` substring so it matches `internlm2-chat-7b`
+            // and `internlm2_5-7b-chat` but NOT `internlm-chat-7b` (InternLM v1,
+            // routes to Llama), `internlm3-*` (also Llama-architecture per
+            // vllm/model_executor/models/registry.py:146), or `Intern-S1` /
+            // `Intern-S1-Pro` (separate intern-s1 parser, see PR #40115).
+            .register_pattern("internlm2", names::INTERNLM)
             .register_pattern("llama-4", names::LLAMA4_JSON)
             .register_pattern("llama-3.2", names::LLAMA3_JSON)
             .register_pattern("llama-3.1", names::LLAMA3_JSON)

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

@@ -161,4 +161,33 @@ fn factory_new_resolves_default_patterns() {
         factory.resolve_name_for_model("org/mm-m2-base"),
         Some(names::MINIMAX_M2)
     );
+
+    // InternLM2 positive: both dashed and underscored versioned names route.
+    assert_eq!(
+        factory.resolve_name_for_model("internlm/internlm2-chat-7b"),
+        Some(names::INTERNLM)
+    );
+    assert_eq!(
+        factory.resolve_name_for_model("internlm/internlm2_5-7b-chat"),
+        Some(names::INTERNLM)
+    );
+
+    // Negative: other internlm-org models do NOT route to the InternLM2 parser,
+    // since they use unrelated prompt formats.
+    //   - InternLM v1 (`internlm-chat-7b`) routes to Llama
+    //   - InternLM3 (`internlm3-8b-instruct`) routes to Llama
+    //   - Intern-S1 / Intern-S1-Pro have their own parser (Python PR #40115)
+    assert_eq!(
+        factory.resolve_name_for_model("internlm/internlm-chat-7b"),
+        None
+    );
+    assert_eq!(
+        factory.resolve_name_for_model("internlm/internlm3-8b-instruct"),
+        None
+    );
+    assert_eq!(factory.resolve_name_for_model("internlm/Intern-S1"), None);
+    assert_eq!(
+        factory.resolve_name_for_model("internlm/Intern-S1-Pro"),
+        None
+    );
 }

rust/src/server/src/routes/inference/generate.rs

@@ -3,23 +3,33 @@ mod types;
 mod validate;
 
 use std::collections::HashMap;
+use std::convert::Infallible;
+use std::result::Result;
 use std::sync::Arc;
 
+use asynk_strim_attr::{TryYielder, try_stream};
 use axum::Json;
 use axum::extract::State;
 use axum::http::HeaderMap;
+use axum::response::sse::{Event, Sse};
 use axum::response::{IntoResponse, Response};
+use futures::{Stream, StreamExt as _, pin_mut};
 use thiserror_ext::AsReport as _;
-use tracing::info;
+use tracing::{error, info, trace};
 use tracing_futures::Instrument as _;
 use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs};
-use vllm_llm::{CollectedGenerateOutput, GenerateOutputStreamExt as _};
+use vllm_llm::{
+    CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _,
+};
 
 use self::convert::prepare_generate_request;
-use self::types::{GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice};
-use crate::error::{ApiError, server_error};
+use self::types::{
+    GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice,
+    GenerateResponseStreamChoice, GenerateStreamResponse,
+};
+use crate::error::{ApiError, bail_server_error, server_error};
 use crate::routes::openai::utils::logprobs::clamp_logprob;
-use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb};
+use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb, Usage};
 use crate::routes::openai::utils::validated_json::ValidatedJson;
 use crate::state::AppState;
 use crate::utils::resolve_request_context;
@@ -46,6 +56,7 @@ pub async fn generate(
     let log_request = state.enable_log_requests;
     let include_logprobs = prepared.include_logprobs;
     let include_prompt_logprobs = prepared.include_prompt_logprobs;
+    let stream = prepared.stream;
 
     let raw_stream = match state
         .chat
@@ -64,6 +75,20 @@ pub async fn generate(
         }
     };
 
+    if stream {
+        let chunk_stream = generate_chunk_stream(
+            raw_stream,
+            prepared.request_id,
+            log_request,
+            prepared.include_usage,
+            prepared.include_continuous_usage,
+            include_logprobs,
+        );
+        let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span);
+
+        return Sse::new(sse_stream).into_response();
+    }
+
     let collected = match raw_stream.collect_output().instrument(request_span.clone()).await {
         Ok(collected) => collected,
         Err(error) => {
@@ -98,6 +123,102 @@ pub async fn generate(
     Json(response).into_response()
 }
 
+#[try_stream]
+async fn generate_chunk_stream(
+    stream: impl Stream<Item = vllm_llm::Result<GenerateOutput>>,
+    request_id: String,
+    log_request: bool,
+    include_usage: bool,
+    include_continuous_usage: bool,
+    include_logprobs: bool,
+    mut y: TryYielder<GenerateStreamResponse, ApiError>,
+) -> Result<(), ApiError> {
+    pin_mut!(stream);
+    let mut prompt_tokens: Option<u32> = None;
+    let mut output_tokens = 0_u32;
+
+    while let Some(next) = stream.next().await {
+        match next {
+            Ok(output) => {
+                if prompt_tokens.is_none() {
+                    prompt_tokens =
+                        output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len() as u32);
+                }
+                let usage_prompt_tokens = prompt_tokens.unwrap_or_default();
+
+                let token_ids = output.token_ids;
+                output_tokens = output_tokens.saturating_add(token_ids.len() as u32);
+                let finish_reason = output.finish_reason;
+
+                if matches!(finish_reason.as_ref(), Some(FinishReason::Error)) {
+                    bail_server_error!("Internal server error");
+                }
+
+                if let Some(finish_reason) = finish_reason.as_ref()
+                    && log_request
+                {
+                    info!(
+                        stream = true,
+                        prompt_tokens = usage_prompt_tokens,
+                        output_tokens,
+                        finish_reason = finish_reason.as_str(),
+                        "generate finished"
+                    );
+                }
+
+                if token_ids.is_empty() && finish_reason.is_none() {
+                    continue;
+                }
+
+                let logprobs = if include_logprobs && !token_ids.is_empty() {
+                    let logprobs = output.logprobs.as_ref().ok_or_else(|| {
+                        server_error!(
+                            "raw generate stream requested logprobs but generation returned none"
+                        )
+                    })?;
+                    Some(raw_logprobs_to_openai_chat(logprobs)?)
+                } else {
+                    None
+                };
+
+                y.yield_ok(GenerateStreamResponse {
+                    request_id: request_id.clone(),
+                    choices: vec![GenerateResponseStreamChoice {
+                        index: 0,
+                        logprobs,
+                        finish_reason: finish_reason.map(|reason| reason.as_str().to_string()),
+                        token_ids,
+                    }],
+                    usage: include_continuous_usage
+                        .then(|| Usage::from_counts(usage_prompt_tokens, output_tokens)),
+                })
+                .await;
+            }
+            Err(error) => {
+                error!(
+                    error = %error.as_report(),
+                    "raw generate stream failed"
+                );
+                bail_server_error!("{}", error.to_report_string());
+            }
+        }
+    }
+
+    if include_usage {
+        y.yield_ok(GenerateStreamResponse {
+            request_id,
+            choices: Vec::new(),
+            usage: Some(Usage::from_counts(
+                prompt_tokens.unwrap_or_default(),
+                output_tokens,
+            )),
+        })
+        .await;
+    }
+
+    Ok(())
+}
+
 fn collect_generate(
     collected: CollectedGenerateOutput,
     request_id: String,
@@ -213,3 +334,94 @@ fn position_to_logprob_map(position: &PositionLogprobs) -> HashMap<u32, Generate
 fn format_token_id(token_id: u32) -> String {
     format!("token_id:{token_id}")
 }
+
+/// Convert one raw-generate chunk stream into SSE events.
+#[try_stream]
+async fn generate_sse_stream(
+    stream: impl Stream<Item = Result<GenerateStreamResponse, ApiError>>,
+    mut y: TryYielder<Event, Infallible>,
+) -> Result<(), Infallible> {
+    pin_mut!(stream);
+
+    while let Some(next) = stream.next().await {
+        match next {
+            Ok(chunk) => y.yield_ok(to_sse_event(&chunk)).await,
+            Err(error) => {
+                y.yield_ok(to_error_sse_event(&error)).await;
+                break;
+            }
+        }
+    }
+
+    y.yield_ok(done_sse_event()).await;
+    Ok(())
+}
+
+fn to_sse_event(chunk: &GenerateStreamResponse) -> Event {
+    let payload = serde_json::to_string(chunk).expect("generate chunk must serialize to JSON");
+    trace!(payload, "generate emitting chunk");
+    Event::default().data(payload)
+}
+
+fn to_error_sse_event(error: &ApiError) -> Event {
+    let payload = serde_json::to_string(&error.to_error_response())
+        .expect("ErrorResponse must serialize to JSON");
+    trace!(payload, "generate emitting error");
+    Event::default().data(payload)
+}
+
+fn done_sse_event() -> Event {
+    trace!("generate emitting done");
+    Event::default().data("[DONE]")
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use futures::{TryStreamExt as _, stream};
+    use vllm_llm::GeneratePromptInfo;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn generate_chunk_stream_captures_late_prompt_info() {
+        let stream = stream::iter(vec![
+            Ok(GenerateOutput {
+                request_id: String::new(),
+                prompt_info: None,
+                token_ids: Vec::new(),
+                logprobs: None,
+                finish_reason: None,
+                kv_transfer_params: None,
+            }),
+            Ok(GenerateOutput {
+                request_id: String::new(),
+                prompt_info: Some(GeneratePromptInfo {
+                    prompt_token_ids: Arc::from([11_u32, 22_u32]),
+                    prompt_logprobs: None,
+                }),
+                token_ids: vec![33],
+                logprobs: None,
+                finish_reason: Some(FinishReason::stop_eos()),
+                kv_transfer_params: None,
+            }),
+        ]);
+
+        let chunks: Vec<_> =
+            generate_chunk_stream(stream, "raw-stream".to_string(), false, true, true, false)
+                .try_collect()
+                .await
+                .expect("collect chunks");
+
+        assert_eq!(chunks.len(), 2);
+        assert_eq!(
+            chunks[0].usage.as_ref().expect("chunk usage").prompt_tokens,
+            2
+        );
+        assert_eq!(
+            chunks[1].usage.as_ref().expect("final usage").prompt_tokens,
+            2
+        );
+    }
+}

rust/src/server/src/routes/inference/generate/convert.rs

@@ -10,6 +10,9 @@ use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params};
 pub struct PreparedRequest {
     pub request_id: String,
     pub text_request: TextRequest,
+    pub stream: bool,
+    pub include_usage: bool,
+    pub include_continuous_usage: bool,
     pub include_logprobs: bool,
     pub include_prompt_logprobs: bool,
 }
@@ -23,6 +26,18 @@ pub fn prepare_generate_request(
 ) -> Result<PreparedRequest, ApiError> {
     validate::validate_request_compat(&request, served_model_names)?;
 
+    let stream = request.stream;
+    let include_usage = request
+        .stream_options
+        .as_ref()
+        .and_then(|options| options.include_usage)
+        .unwrap_or(false);
+    let include_continuous_usage = include_usage
+        && request
+            .stream_options
+            .as_ref()
+            .and_then(|options| options.continuous_usage_stats)
+            .unwrap_or(false);
     let include_logprobs = request.sampling_params.logprobs.is_some();
     let include_prompt_logprobs = request.sampling_params.prompt_logprobs.is_some();
     let mut sampling_params = request.sampling_params;
@@ -47,6 +62,9 @@ pub fn prepare_generate_request(
     Ok(PreparedRequest {
         request_id: ctx.request_id,
         text_request,
+        stream,
+        include_usage,
+        include_continuous_usage,
         include_logprobs,
         include_prompt_logprobs,
     })
@@ -109,4 +127,28 @@ mod tests {
             Some(json!({"connector": "x"}))
         );
     }
+
+    #[test]
+    fn prepare_generate_request_gates_continuous_usage_on_include_usage() {
+        let request: GenerateRequest = serde_json::from_value(json!({
+            "model": "Qwen/Qwen1.5-0.5B-Chat",
+            "token_ids": [11, 22],
+            "stream": true,
+            "stream_options": {
+                "continuous_usage_stats": true
+            },
+            "sampling_params": {}
+        }))
+        .expect("parse request");
+
+        let prepared = prepare_generate_request(
+            request,
+            &["Qwen/Qwen1.5-0.5B-Chat".to_string()],
+            ResolvedRequestContext::default(),
+        )
+        .expect("prepare");
+
+        assert!(!prepared.include_usage);
+        assert!(!prepared.include_continuous_usage);
+    }
 }

rust/src/server/src/routes/inference/generate/types.rs

@@ -5,7 +5,7 @@ use serde_json::{Map, Value};
 use validator::Validate;
 use vllm_text::SamplingParams;
 
-use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable};
+use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage};
 
 /// vLLM-compatible request type for the token-in/token-out generate API.
 #[serde_with::skip_serializing_none]
@@ -17,6 +17,7 @@ pub struct GenerateRequest {
     pub sampling_params: SamplingParams,
     #[serde(default)]
     pub stream: bool,
+    pub stream_options: Option<StreamOptions>,
     pub cache_salt: Option<String>,
     #[serde(default)]
     pub priority: i32,
@@ -37,6 +38,25 @@ pub(super) struct GenerateResponseChoice {
     pub token_ids: Vec<u32>,
 }
 
+/// Mirrors the Python vLLM `GenerateResponseStreamChoice` class.
+#[serde_with::skip_serializing_none]
+#[derive(Debug, Clone, Serialize)]
+pub(super) struct GenerateResponseStreamChoice {
+    pub index: u32,
+    pub logprobs: Option<ChatLogProbs>,
+    pub finish_reason: Option<String>,
+    pub token_ids: Vec<u32>,
+}
+
+/// Mirrors the Python vLLM `GenerateStreamResponse` class.
+#[serde_with::skip_serializing_none]
+#[derive(Debug, Clone, Serialize)]
+pub(super) struct GenerateStreamResponse {
+    pub request_id: String,
+    pub choices: Vec<GenerateResponseStreamChoice>,
+    pub usage: Option<Usage>,
+}
+
 /// Mirrors the Python vLLM `GenerateResponse` class.
 #[serde_with::skip_serializing_none]
 #[derive(Debug, Clone, Serialize)]

rust/src/server/src/routes/inference/generate/validate.rs

@@ -13,8 +13,11 @@ pub(super) fn validate_request_compat(
         return Err(ApiError::model_not_found(model.clone()));
     }
 
-    if request.stream {
-        bail_invalid_request!(param = "stream", "stream=true is not supported.");
+    if request.stream_options.is_some() && !request.stream {
+        bail_invalid_request!(
+            param = "stream_options",
+            "stream_options are only supported when stream=true."
+        );
     }
 
     if request.token_ids.is_empty() {
@@ -65,11 +68,24 @@ mod tests {
     }
 
     #[test]
-    fn validate_request_compat_rejects_streaming() {
+    fn validate_request_compat_accepts_streaming() {
         let request = GenerateRequest {
             stream: true,
             ..base_request()
         };
+        assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_ok());
+    }
+
+    #[test]
+    fn validate_request_compat_rejects_stream_options_without_streaming() {
+        let request: GenerateRequest = serde_json::from_value(json!({
+            "model": "Qwen/Qwen1.5-0.5B-Chat",
+            "token_ids": [11, 22],
+            "stream": false,
+            "stream_options": {"include_usage": true},
+            "sampling_params": {}
+        }))
+        .expect("parse request");
         assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
     }
 

rust/src/server/src/routes/tests.rs

@@ -2425,8 +2425,72 @@ async fn non_stream_raw_generate_returns_token_output_envelope() {
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
-async fn raw_generate_rejects_streaming() {
-    let mut app = test_app().await;
+async fn stream_raw_generate_returns_sse_chunks_and_usage() {
+    let ipc = IpcNamespace::new().expect("create ipc namespace");
+    let handshake_address = ipc.handshake_endpoint();
+    let engine_id = b"engine-raw-generate-stream".to_vec();
+
+    let engine_task = MockEngineTask::new(spawn_mock_engine_task(
+        handshake_address.clone(),
+        engine_id.clone(),
+        |dealer, push| {
+            boxed_test_future(async move {
+                let add = recv_engine_message(dealer).await;
+                let request: EngineCoreRequest =
+                    rmp_serde::from_slice(&add[1]).expect("decode request");
+                assert_eq!(request.prompt_token_ids.as_deref(), Some(&[11, 22][..]));
+                assert_eq!(request.external_req_id.as_deref(), Some("raw-stream"));
+
+                send_outputs(
+                    push,
+                    EngineCoreOutputs {
+                        engine_index: 0,
+                        outputs: vec![
+                            request_output_with_logprobs(
+                                &request.request_id,
+                                vec![33],
+                                None,
+                                None,
+                                Some(sample_logprobs_for_token(33, 34)),
+                                None,
+                            ),
+                            request_output_with_logprobs(
+                                &request.request_id,
+                                vec![44],
+                                Some(EngineCoreFinishReason::Stop),
+                                None,
+                                Some(sample_logprobs_for_token(44, 45)),
+                                None,
+                            ),
+                        ],
+                        scheduler_stats: None,
+                        timestamp: 0.0,
+                        utility_output: None,
+                        finished_requests: None,
+                        wave_complete: None,
+                        start_wave: None,
+                    },
+                )
+                .await;
+            })
+        },
+    ));
+
+    let client = EngineCoreClient::connect(
+        EngineCoreClientConfig::new_single(handshake_address)
+            .with_model_name("test-model")
+            .with_local_input_output_addresses(
+                Some(ipc.input_endpoint()),
+                Some(ipc.output_endpoint()),
+            ),
+    )
+    .await
+    .expect("connect client");
+    let chat = ChatLlm::from_shared_backend(Llm::new(client), Arc::new(FakeChatBackend::new()));
+    let mut app = build_router(Arc::new(AppState::new(
+        vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
+        chat,
+    )));
 
     let response = app
         .call(
@@ -2437,9 +2501,17 @@ async fn raw_generate_rejects_streaming() {
                 .body(Body::from(
                     json!({
                         "model": "Qwen/Qwen1.5-0.5B-Chat",
+                        "request_id": "raw-stream",
                         "token_ids": [11, 22],
                         "stream": true,
-                        "sampling_params": {}
+                        "stream_options": {
+                            "include_usage": true,
+                            "continuous_usage_stats": true
+                        },
+                        "sampling_params": {
+                            "max_tokens": 2,
+                            "logprobs": 1
+                        }
                     })
                     .to_string(),
                 ))
@@ -2448,10 +2520,196 @@ async fn raw_generate_rejects_streaming() {
         .await
         .expect("call app");
 
-    assert_eq!(response.status(), StatusCode::BAD_REQUEST);
+    assert_eq!(response.status(), StatusCode::OK);
+    assert_eq!(
+        response.headers().get("content-type").and_then(|value| value.to_str().ok()),
+        Some("text/event-stream")
+    );
+
     let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
-    let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
-    assert_eq!(json["error"]["param"], "stream");
+    engine_task.await.expect("mock engine task");
+    let text = String::from_utf8(body.to_vec()).expect("utf8 body");
+    let payloads = sse_data_payloads(&text);
+    assert_eq!(payloads.len(), 4, "{text}");
+
+    let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
+    assert_eq!(first["request_id"], "raw-stream");
+    assert_eq!(first["choices"][0]["index"], 0);
+    assert_eq!(first["choices"][0]["token_ids"], json!([33]));
+    assert_eq!(
+        first["choices"][0]["logprobs"]["content"][0]["token"],
+        "token_id:33"
+    );
+    assert_eq!(first["usage"]["prompt_tokens"], 2);
+    assert_eq!(first["usage"]["completion_tokens"], 1);
+
+    let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
+    assert_eq!(second["choices"][0]["token_ids"], json!([44]));
+    assert_eq!(second["choices"][0]["finish_reason"], "stop");
+    assert_eq!(second["usage"]["completion_tokens"], 2);
+
+    let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json");
+    assert_eq!(usage["choices"], json!([]));
+    assert_eq!(usage["usage"]["prompt_tokens"], 2);
+    assert_eq!(usage["usage"]["completion_tokens"], 2);
+    assert_eq!(usage["usage"]["total_tokens"], 4);
+    assert_eq!(payloads[3], "[DONE]");
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn stream_raw_generate_emits_final_usage_without_continuous_usage() {
+    let (mut app, engine_task) = test_app_with_stream_output_specs(vec![
+        (vec![33], None),
+        (vec![44], Some(EngineCoreFinishReason::Stop)),
+    ])
+    .await;
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/inference/v1/generate")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "model": "Qwen/Qwen1.5-0.5B-Chat",
+                        "request_id": "raw-stream-final-usage",
+                        "token_ids": [11, 22],
+                        "stream": true,
+                        "stream_options": {
+                            "include_usage": true
+                        },
+                        "sampling_params": {
+                            "max_tokens": 2
+                        }
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    engine_task.await.expect("mock engine task");
+    let text = String::from_utf8(body.to_vec()).expect("utf8 body");
+    let payloads = sse_data_payloads(&text);
+    assert_eq!(payloads.len(), 4, "{text}");
+
+    let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
+    assert_eq!(first["choices"][0]["token_ids"], json!([33]));
+    assert!(first.get("usage").is_none());
+
+    let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
+    assert_eq!(second["choices"][0]["token_ids"], json!([44]));
+    assert_eq!(second["choices"][0]["finish_reason"], "stop");
+    assert!(second.get("usage").is_none());
+
+    let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json");
+    assert_eq!(usage["choices"], json!([]));
+    assert_eq!(usage["usage"]["prompt_tokens"], 2);
+    assert_eq!(usage["usage"]["completion_tokens"], 2);
+    assert_eq!(usage["usage"]["total_tokens"], 4);
+    assert_eq!(payloads[3], "[DONE]");
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn stream_raw_generate_emits_empty_finish_chunk() {
+    let (mut app, engine_task) = test_app_with_stream_output_specs(vec![
+        (vec![33], None),
+        (vec![], Some(EngineCoreFinishReason::Stop)),
+    ])
+    .await;
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/inference/v1/generate")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "model": "Qwen/Qwen1.5-0.5B-Chat",
+                        "request_id": "raw-stream-empty-finish",
+                        "token_ids": [11, 22],
+                        "stream": true,
+                        "sampling_params": {
+                            "max_tokens": 2
+                        }
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    engine_task.await.expect("mock engine task");
+    let text = String::from_utf8(body.to_vec()).expect("utf8 body");
+    let payloads = sse_data_payloads(&text);
+    assert_eq!(payloads.len(), 3, "{text}");
+
+    let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
+    assert_eq!(first["choices"][0]["token_ids"], json!([33]));
+    assert!(first["choices"][0].get("finish_reason").is_none());
+
+    let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
+    assert_eq!(second["choices"][0]["token_ids"], json!([]));
+    assert_eq!(second["choices"][0]["finish_reason"], "stop");
+    assert_eq!(payloads[2], "[DONE]");
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn stream_raw_generate_error_finish_returns_sse_error() {
+    let (mut app, engine_task) =
+        test_app_with_stream_output_specs(vec![(vec![], Some(EngineCoreFinishReason::Error))])
+            .await;
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/inference/v1/generate")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "model": "Qwen/Qwen1.5-0.5B-Chat",
+                        "request_id": "raw-stream-error",
+                        "token_ids": [11, 22],
+                        "stream": true,
+                        "stream_options": {
+                            "include_usage": true
+                        },
+                        "sampling_params": {
+                            "max_tokens": 2
+                        }
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    engine_task.await.expect("mock engine task");
+    let text = String::from_utf8(body.to_vec()).expect("utf8 body");
+
+    assert!(text.contains("\"type\":\"server_error\""), "{text}");
+    assert!(text.contains("Internal server error"), "{text}");
+    assert!(!text.contains("\"finish_reason\":\"error\""), "{text}");
+    assert!(!text.contains("\"usage\":"), "{text}");
+    assert!(text.trim_end().ends_with("data: [DONE]"), "{text}");
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]

rust/src/tool-parser/src/json/hermes.rs

@@ -8,7 +8,7 @@ const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
     marker_whitespace: JsonToolCallWhitespace::Optional,
     delimiter: None,
     name_key: "name",
-    arguments_key: "arguments",
+    arguments_key: &["arguments"],
 };
 
 /// Tool parser for Hermes XML-wrapped JSON tool calls.

rust/src/tool-parser/src/json/internlm2.rs

@@ -0,0 +1,351 @@
+use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
+use crate::{Result, Tool, ToolParser, ToolParserOutput};
+
+const INTERNLM2_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
+    parser_name: "InternLM2",
+    start_marker: "<|action_start|><|plugin|>",
+    end_marker: "<|action_end|>",
+    marker_whitespace: JsonToolCallWhitespace::Optional,
+    delimiter: None,
+    name_key: "name",
+    // The Python parser's `get_arguments()` accepts either `parameters` or
+    // `arguments` and prefers `parameters` when both are present. This Rust
+    // parser uses first-encountered semantics because the header parser only
+    // permits one args key per tool-call object; if a future model emits
+    // both keys in the same object, the Rust port will accept the first one
+    // and reject the trailing one as a syntax error rather than silently
+    // shadowing it.
+    arguments_key: &["parameters", "arguments"],
+};
+
+/// Tool parser for InternLM2 special-token wrapped JSON tool calls.
+///
+/// Example tool call content:
+///
+/// ```text
+/// <|action_start|><|plugin|>{"name": "get_weather", "parameters": {"location":"Tokyo"}}<|action_end|>
+/// ```
+///
+/// Arguments are already OpenAI-style JSON text, so they are streamed as raw
+/// argument deltas without schema conversion or JSON normalization.
+///
+/// # Divergences from the Python reference
+///
+/// This Rust port intentionally diverges from
+/// `vllm/tool_parsers/internlm2_tool_parser.py` in two user-visible ways:
+///
+/// - **Parallel tool calls are supported.** Python silently drops every
+///   `<|action_start|>` block after the first (`current_tool_id > 0` returns
+///   an empty delta); this parser emits every well-formed block with
+///   incrementing `tool_index`. Models that legitimately emit multiple action
+///   blocks therefore produce more tool calls under Rust than under Python.
+/// - **End-marker bytes inside JSON string values are preserved.** Python
+///   does `action.split("<|action_end|>")[0]` which truncates regardless of
+///   JSON context; this parser scans matched braces and quotes so a literal
+///   `<|action_end|>` inside an arguments string is forwarded intact.
+/// - **Only whitespace is allowed before the `{`.** Python's non-streaming
+///   `action[action.find("{"):]` drops any bytes before the first `{`, but
+///   its streaming path has no equivalent and the model format always emits
+///   `<|plugin|>{...`; this parser allows only whitespace there, matching the
+///   other JSON parsers in this crate.
+/// - **Truncated tool calls error rather than silently dropping.** Python's
+///   streaming wrapper swallows mid-stream errors with `except Exception:
+///   return None` (logging a traceback) while its non-streaming path raises
+///   `JSONDecodeError`; this parser returns an `incomplete InternLM2 tool
+///   call` error from `finish()`, matching the other JSON parsers and Python's
+///   non-streaming behavior.
+///
+/// # Known unaddressed divergences (TODO)
+///
+/// The following Python behaviors are NOT yet matched. They are deferred to
+/// follow-up work because they require non-local changes to the shared
+/// `JsonToolCallParser` core that would affect Hermes / Llama / Mistral /
+/// Qwen as well. If a real-world InternLM2 deployment hits one of these,
+/// prioritize the corresponding fix.
+///
+/// - **Arguments value type.** The shared core requires the arguments value
+///   to be a JSON object (`take_json_object` rejects anything not starting
+///   with `{`). Python's `json.dumps(action_dict.get("parameters", ...))`
+///   accepts `null`, arrays, strings, and numbers and round-trips them
+///   verbatim. Models that legitimately emit `"parameters":null` will hard-
+///   fail under Rust.
+/// - **Unknown arguments key.** Python falls back to `{}` via
+///   `action_dict.get("parameters", action_dict.get("arguments", {}))` when
+///   neither key is present; the Rust header parser raises
+///   `parsing failed: invalid InternLM2` for any unrecognized key. A model
+///   that emits a typo (e.g. `"params"`) breaks the whole response.
+/// - **Field order independence.** The header parser requires the JSON keys
+///   to appear in the order `name` then arguments key. Python's
+///   `json.loads` + `dict.get` is order-independent, so a model emitting
+///   `{"parameters":{...},"name":"foo"}` parses in Python but fails in Rust.
+pub struct Internlm2ToolParser {
+    inner: JsonToolCallParser,
+}
+
+impl Internlm2ToolParser {
+    /// Create an InternLM2 tool parser.
+    fn new(_tools: &[Tool]) -> Self {
+        Self {
+            inner: JsonToolCallParser::new(INTERNLM2_CONFIG),
+        }
+    }
+}
+
+impl ToolParser for Internlm2ToolParser {
+    /// Create a boxed InternLM2 tool parser.
+    fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
+    where
+        Self: Sized + 'static,
+    {
+        Ok(Box::new(Self::new(tools)))
+    }
+
+    /// Preserve special-token markers while decoding, since
+    /// `<|action_start|>`, `<|plugin|>`, and `<|action_end|>` are tokenizer
+    /// special tokens in InternLM2 models.
+    fn preserve_special_tokens(&self) -> bool {
+        true
+    }
+
+    /// Feed one decoded text chunk through the InternLM2 parser.
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+        self.inner.parse_into(chunk, output)
+    }
+
+    /// Flush any buffered partial state at end of stream.
+    fn finish(&mut self) -> Result<ToolParserOutput> {
+        self.inner.finish()
+    }
+
+    /// Clear parser state and return currently uncommitted buffered text.
+    fn reset(&mut self) -> String {
+        self.inner.reset()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use expect_test::expect;
+    use thiserror_ext::AsReport;
+
+    use super::Internlm2ToolParser;
+    use crate::test_utils::{collect_stream, split_by_chars, test_tools};
+    use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
+
+    const ACTION_START: &str = "<|action_start|><|plugin|>";
+    const ACTION_END: &str = "<|action_end|>";
+
+    fn build_tool_call(function_name: &str, args_key: &str, arguments: &str) -> String {
+        format!(
+            r#"{ACTION_START}{{"name":"{function_name}","{args_key}":{arguments}}}{ACTION_END}"#
+        )
+    }
+
+    #[test]
+    fn internlm2_parse_complete_without_tool_call_keeps_text() {
+        let mut parser = Internlm2ToolParser::new(&test_tools());
+        let result = parser.parse_complete("Hello, world!").unwrap();
+
+        assert_eq!(result.normal_text, "Hello, world!");
+        assert!(result.calls.is_empty());
+    }
+
+    #[test]
+    fn internlm2_parse_complete_extracts_parameters_key() {
+        let mut parser = Internlm2ToolParser::new(&test_tools());
+        let arguments = r#"{"location":"Tokyo","days":"3"}"#;
+        let result = parser
+            .parse_complete(&format!(
+                "Let me check.\n{}",
+                build_tool_call("get_weather", "parameters", arguments)
+            ))
+            .unwrap();
+
+        assert_eq!(result.normal_text, "Let me check.\n");
+        assert_eq!(result.calls.len(), 1);
+        assert_eq!(result.calls[0].tool_index, 0);
+        assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
+        assert_eq!(result.calls[0].arguments, arguments);
+    }
+
+    #[test]
+    fn internlm2_parse_complete_extracts_arguments_key_fallback() {
+        let mut parser = Internlm2ToolParser::new(&test_tools());
+        let arguments = r#"{"location":"Tokyo"}"#;
+        let result = parser
+            .parse_complete(&build_tool_call("get_weather", "arguments", arguments))
+            .unwrap();
+
+        assert_eq!(result.calls.len(), 1);
+        assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
+        assert_eq!(result.calls[0].arguments, arguments);
+    }
+
+    #[test]
+    fn internlm2_accepts_whitespace_after_plugin_marker() {
+        let mut parser = Internlm2ToolParser::new(&test_tools());
+        let result = parser
+            .parse_complete(&format!(
+                r#"{ACTION_START}
+{{"name":"get_weather","parameters":{{}}}}{ACTION_END}"#
+            ))
+            .unwrap();
+
+        assert_eq!(result.calls.len(), 1);
+        assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
+    }
+
+    #[test]
+    fn internlm2_does_not_validate_or_normalize_arguments() {
+        let mut parser = Internlm2ToolParser::new(&test_tools());
+        let arguments = r#"{"location":"Tokyo",}"#;
+        let result = parser
+            .parse_complete(&build_tool_call("get_weather", "parameters", arguments))
+            .unwrap();
+
+        assert_eq!(result.calls[0].arguments, arguments);
+    }
+
+    #[test]
+    fn internlm2_streaming_emits_argument_deltas() {
+        let mut parser = Internlm2ToolParser::new(&test_tools());
+        let chunks = [
+            "preface <|action",
+            "_start|><|plugin|>",
+            r#"{"name":"get_weather","parameters":"#,
+            r#"{"location":"#,
+            r#""Beijing""#,
+            r#"}"#,
+            r#"}<|action_end|> suffix"#,
+        ];
+
+        let mut result = ToolParserOutput::default();
+        let mut observed_arguments = Vec::new();
+        for chunk in chunks {
+            let next = parser.parse_chunk(chunk).unwrap();
+            observed_arguments.extend(
+                next.calls
+                    .iter()
+                    .filter(|call| call.name.is_none())
+                    .map(|call| call.arguments.clone()),
+            );
+            result.append(next);
+        }
+        result.append(parser.finish().unwrap());
+
+        assert_eq!(
+            observed_arguments,
+            [r#"{"location":"#, r#""Beijing""#, r#"}"#]
+        );
+        assert_eq!(result.normal_text, "preface  suffix");
+        assert_eq!(
+            result.coalesce_calls().calls[0].arguments,
+            r#"{"location":"Beijing"}"#
+        );
+    }
+
+    #[test]
+    fn internlm2_streaming_handles_split_markers() {
+        let input = format!(
+            "hello {}",
+            build_tool_call("get_weather", "parameters", r#"{"location":"Tokyo"}"#)
+        );
+        let chunks = split_by_chars(&input, 5);
+        let mut parser = Internlm2ToolParser::new(&test_tools());
+
+        let result = collect_stream(&mut parser, &chunks);
+
+        assert_eq!(result.normal_text, "hello ");
+        assert_eq!(result.calls.len(), 1);
+        assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#);
+    }
+
+    #[test]
+    fn internlm2_streaming_extracts_multiple_blocks() {
+        let input = format!(
+            "{}{}",
+            build_tool_call("get_weather", "parameters", r#"{"location":"Shanghai"}"#),
+            build_tool_call("add", "arguments", r#"{"x":1,"y":2}"#),
+        );
+        let chunks = split_by_chars(&input, 7);
+        let mut parser = Internlm2ToolParser::new(&test_tools());
+
+        let result = collect_stream(&mut parser, &chunks);
+
+        expect![[r#"
+            ToolParserOutput {
+                normal_text: "",
+                calls: [
+                    ToolCallDelta {
+                        tool_index: 0,
+                        name: Some(
+                            "get_weather",
+                        ),
+                        arguments: "{\"location\":\"Shanghai\"}",
+                    },
+                    ToolCallDelta {
+                        tool_index: 1,
+                        name: Some(
+                            "add",
+                        ),
+                        arguments: "{\"x\":1,\"y\":2}",
+                    },
+                ],
+            }
+        "#]]
+        .assert_debug_eq(&result);
+    }
+
+    #[test]
+    fn internlm2_keeps_end_marker_literal_inside_json_string() {
+        let mut parser = Internlm2ToolParser::new(&test_tools());
+        let arguments = format!(r#"{{"text":"literal {ACTION_END} inside"}}"#);
+        let input = build_tool_call("echo", "parameters", &arguments);
+
+        let result = parser.parse_complete(&input).unwrap();
+
+        assert_eq!(result.calls.len(), 1);
+        assert_eq!(result.calls[0].arguments, arguments);
+    }
+
+    #[test]
+    fn internlm2_finish_errors_on_truncated_tool_call() {
+        let mut parser = Internlm2ToolParser::new(&test_tools());
+        let pre_finish = parser
+            .parse_chunk(&format!(
+                r#"{ACTION_START}{{"name":"get_weather","parameters":{{"location""#
+            ))
+            .unwrap();
+        let error = parser.finish().unwrap_err();
+
+        assert_eq!(
+            pre_finish.calls[0].name.as_deref(),
+            Some("get_weather"),
+            "name delta is still emitted from parse_chunk() before truncation",
+        );
+        assert!(
+            error.to_report_string().contains("incomplete InternLM2 tool call"),
+            "finish() reports the truncated tool call as incomplete: {}",
+            error.to_report_string(),
+        );
+    }
+
+    #[test]
+    fn internlm2_unknown_arguments_key_fails() {
+        let mut parser = Internlm2ToolParser::new(&test_tools());
+        let input = build_tool_call("get_weather", "params", r#"{"location":"Tokyo"}"#);
+
+        let error = parser.parse_chunk(&input).unwrap_err();
+
+        expect![[r#"
+            tool parser parsing failed: invalid InternLM2
+            expected `parameters`, `arguments`"#]]
+        .assert_eq(&error.to_report_string());
+    }
+
+    #[test]
+    fn internlm2_preserve_special_tokens_is_true() {
+        let parser = Internlm2ToolParser::new(&test_tools());
+        assert!(parser.preserve_special_tokens());
+    }
+}

rust/src/tool-parser/src/json/llama.rs

@@ -203,7 +203,7 @@ fn llama_tool_call_header_event(input: &mut JsonToolInput<'_>) -> ModalResult<Ll
         marker_whitespace: JsonToolCallWhitespace::Optional,
         delimiter: Some(";"),
         name_key: "name",
-        arguments_key: "parameters",
+        arguments_key: &["parameters"],
     };
 
     match tool_call_header_event(input, CONFIG)? {

rust/src/tool-parser/src/json/mistral.rs

@@ -8,7 +8,7 @@ const MISTRAL_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
     marker_whitespace: JsonToolCallWhitespace::Optional,
     delimiter: Some(","),
     name_key: "name",
-    arguments_key: "arguments",
+    arguments_key: &["arguments"],
 };
 
 /// Tool parser for Mistral JSON-array tool calls.

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

@@ -1,20 +1,22 @@
 //! Shared parser core for JSON tool calls wrapped by text markers.
 
 pub use hermes::HermesToolParser;
+pub use internlm2::Internlm2ToolParser;
 pub use llama::Llama3JsonToolParser;
 pub use mistral::MistralToolParser;
 pub use qwen::Qwen3XmlToolParser;
 
 mod hermes;
+mod internlm2;
 mod llama;
 mod mistral;
 mod qwen;
 
 use winnow::ascii::multispace0 as ws0;
 use winnow::combinator::{alt, seq};
-use winnow::error::{ModalResult, StrContext, StrContextValue};
+use winnow::error::{AddContext, ModalResult, StrContext, StrContextValue};
 use winnow::prelude::*;
-use winnow::stream::Partial;
+use winnow::stream::{Partial, Stream};
 use winnow::token::literal;
 
 use super::utils::{
@@ -32,7 +34,10 @@ struct JsonToolCallConfig {
     marker_whitespace: JsonToolCallWhitespace,
     delimiter: Option<&'static str>,
     name_key: &'static str,
-    arguments_key: &'static str,
+    /// Candidate JSON keys naming the arguments payload, tried in order.
+    /// Most parsers use a single key like `["arguments"]`, but some accept
+    /// multiple (e.g. InternLM2 accepts `parameters` or `arguments`).
+    arguments_key: &'static [&'static str],
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -224,7 +229,7 @@ fn tool_call_header_event(
         _: ws0,
         _: literal(","),
         _: ws0,
-        _: |input: &mut JsonToolInput<'_>| json_key(input, config.arguments_key),
+        _: |input: &mut JsonToolInput<'_>| json_arguments_key(input, config.arguments_key),
         _: ws0,
         _: literal(":"),
         _: ws0,
@@ -246,6 +251,39 @@ fn json_key(input: &mut JsonToolInput<'_>, key: &'static str) -> ModalResult<()>
     .parse_next(input)
 }
 
+/// Parse a JSON object key accepting any of `candidates`.
+///
+/// The full quoted key is consumed and compared against the candidate list,
+/// so this works correctly under partial input regardless of key lengths.
+///
+/// On mismatch, each candidate is attached as its own `Expected` context so the
+/// error enumerates every valid key ("expected `a`, expected `b`"). Because
+/// `StrContextValue::StringLiteral` carries a single `&'static str`, the
+/// contexts are added in a loop over `candidates` rather than through chained
+/// `.context(...)` calls, which keeps the diagnostics complete for any number
+/// of candidates.
+fn json_arguments_key(
+    input: &mut JsonToolInput<'_>,
+    candidates: &'static [&'static str],
+) -> ModalResult<()> {
+    let start = input.checkpoint();
+    json_str
+        .verify(|key: &String| candidates.contains(&key.as_str()))
+        .void()
+        .parse_next(input)
+        .map_err(|err| {
+            err.map(|context_error| {
+                candidates.iter().fold(context_error, |context_error, candidate| {
+                    context_error.add_context(
+                        &*input,
+                        &start,
+                        StrContext::Expected(StrContextValue::StringLiteral(candidate)),
+                    )
+                })
+            })
+        })
+}
+
 /// Parse one event inside a marker-wrapped JSON tool-call arguments payload.
 fn parse_arguments_event(
     input: &mut JsonToolInput<'_>,
@@ -341,7 +379,7 @@ mod tests {
         marker_whitespace: JsonToolCallWhitespace::Optional,
         delimiter: Some("<"),
         name_key: "function",
-        arguments_key: "parameters",
+        arguments_key: &["parameters"],
     };
 
     fn build_tool_call(function_name: &str, arguments: &str) -> String {

rust/src/tool-parser/src/json/qwen.rs

@@ -8,7 +8,7 @@ const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
     marker_whitespace: JsonToolCallWhitespace::Exact("\n"),
     delimiter: None,
     name_key: "name",
-    arguments_key: "arguments",
+    arguments_key: &["arguments"],
 };
 
 /// Tool parser for Qwen XML-wrapped JSON tool calls.

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

@@ -24,7 +24,10 @@ pub use error::{Result, ToolParserError};
 pub use gemma4::Gemma4ToolParser;
 pub use glm_xml::{Glm45MoeToolParser, Glm47MoeToolParser};
 pub use hy_v3::HyV3ToolParser;
-pub use json::{HermesToolParser, Llama3JsonToolParser, MistralToolParser, Qwen3XmlToolParser};
+pub use json::{
+    HermesToolParser, Internlm2ToolParser, Llama3JsonToolParser, MistralToolParser,
+    Qwen3XmlToolParser,
+};
 pub use kimi_k2::KimiK2ToolParser;
 pub use minimax_m2::MinimaxM2ToolParser;
 pub use qwen_coder::Qwen3CoderToolParser;

setup.py

@@ -1165,9 +1165,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/distributed/test_weight_transfer.py

@@ -18,6 +18,7 @@ from torch.multiprocessing.reductions import reduce_tensor
 from vllm.config.parallel import ParallelConfig
 from vllm.config.weight_transfer import WeightTransferConfig
 from vllm.distributed.weight_transfer import WeightTransferEngineFactory
+from vllm.distributed.weight_transfer.base import SparseWeightPatch
 from vllm.distributed.weight_transfer.ipc_engine import (
     IPCWeightTransferEngine,
     IPCWeightTransferInitInfo,
@@ -89,6 +90,67 @@ class TestNCCLWeightTransferUpdateInfoValidation:
         )
         assert len(info.names) == 0
 
+    def test_valid_sparse_update_info(self):
+        """Test creating valid sparse NCCL update info."""
+        info = NCCLWeightTransferUpdateInfo(
+            names=["layer.weight", "layer.bias"],
+            dtype_names=["float32", "bfloat16"],
+            shapes=[[10, 10], [10]],
+            num_updates_list=[4, 2],
+            update_kind="sparse_flat",
+        )
+        assert info.update_kind == "sparse_flat"
+        assert info.num_updates_list == [4, 2]
+
+    def test_sparse_update_requires_num_updates_list(self):
+        with pytest.raises(ValueError, match="`num_updates_list` is required"):
+            NCCLWeightTransferUpdateInfo(
+                names=["layer.weight"],
+                dtype_names=["float32"],
+                shapes=[[10, 10]],
+                update_kind="sparse_flat",
+            )
+
+    def test_sparse_update_rejects_empty_num_updates_list(self):
+        with pytest.raises(ValueError, match="cannot be empty"):
+            NCCLWeightTransferUpdateInfo(
+                names=[],
+                dtype_names=[],
+                shapes=[],
+                num_updates_list=[],
+                update_kind="sparse_flat",
+            )
+
+    def test_sparse_update_rejects_packed(self):
+        with pytest.raises(ValueError, match="cannot be combined with `packed=True`"):
+            NCCLWeightTransferUpdateInfo(
+                names=["layer.weight"],
+                dtype_names=["float32"],
+                shapes=[[10, 10]],
+                num_updates_list=[3],
+                update_kind="sparse_flat",
+                packed=True,
+            )
+
+    def test_sparse_update_rejects_mismatched_num_updates(self):
+        with pytest.raises(ValueError, match="`num_updates_list`"):
+            NCCLWeightTransferUpdateInfo(
+                names=["layer.weight", "layer.bias"],
+                dtype_names=["float32", "float32"],
+                shapes=[[10, 10], [10]],
+                num_updates_list=[3],
+                update_kind="sparse_flat",
+            )
+
+    def test_dense_update_rejects_sparse_metadata(self):
+        with pytest.raises(ValueError, match="Sparse metadata"):
+            NCCLWeightTransferUpdateInfo(
+                names=["layer.weight"],
+                dtype_names=["float32"],
+                shapes=[[10, 10]],
+                num_updates_list=[3],
+            )
+
 
 # --- Unit Tests: Engine Parsing ---
 
@@ -222,6 +284,27 @@ def test_nccl_receive_weights_without_init_raises():
         engine.receive_weights(update_info, lambda x: None)
 
 
+def test_nccl_receive_sparse_weights_without_init_raises():
+    """Test that sparse receive raises if init_transfer_engine wasn't called."""
+    if torch.accelerator.device_count() < 1:
+        pytest.skip("Need at least 1 GPU for this test")
+
+    config = WeightTransferConfig(backend="nccl")
+    parallel_config = create_mock_parallel_config()
+    engine = NCCLWeightTransferEngine(config, parallel_config)
+
+    update_info = NCCLWeightTransferUpdateInfo(
+        names=["w"],
+        dtype_names=["float32"],
+        shapes=[[10]],
+        num_updates_list=[2],
+        update_kind="sparse_flat",
+    )
+
+    with pytest.raises(RuntimeError, match="not initialized"):
+        engine.receive_sparse_weights(update_info, lambda x: None)
+
+
 # --- Integration Test: NCCL Weight Transfer Between Ray Tasks ---
 
 
@@ -379,6 +462,136 @@ def test_nccl_weight_transfer_between_processes():
     )
 
 
+@ray.remote(num_gpus=1)
+def trainer_broadcast_sparse_tensor(
+    master_address: str,
+    master_port: int,
+    world_size: int,
+) -> bool:
+    """Trainer task that broadcasts sparse patches via NCCL."""
+    import torch
+
+    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+    from vllm.distributed.utils import StatelessProcessGroup
+    from vllm.distributed.weight_transfer.base import SparseWeightPatch
+    from vllm.distributed.weight_transfer.nccl_engine import (
+        NCCLTrainerSendWeightsArgs,
+        NCCLWeightTransferEngine,
+    )
+
+    pg = StatelessProcessGroup.create(
+        host=master_address,
+        port=master_port,
+        rank=0,
+        world_size=world_size,
+    )
+    comm = PyNcclCommunicator(pg, device=0)
+
+    patch = SparseWeightPatch(
+        name="test.weight",
+        indices=torch.tensor([1, 7, 25], dtype=torch.int32, device="cuda:0"),
+        values=torch.tensor([10.0, 20.0, 30.0], dtype=torch.float32, device="cuda:0"),
+    )
+    NCCLWeightTransferEngine.trainer_send_sparse_weights(
+        iter([patch]),
+        NCCLTrainerSendWeightsArgs(group=comm),
+    )
+    torch.accelerator.synchronize()
+    return True
+
+
+@ray.remote(num_gpus=1)
+def inference_receive_sparse_tensor(
+    master_address: str,
+    master_port: int,
+    world_size: int,
+) -> dict:
+    """Inference task that receives sparse patches via NCCLWeightTransferEngine."""
+    from unittest.mock import MagicMock
+
+    import torch
+
+    from vllm.config.parallel import ParallelConfig
+    from vllm.config.weight_transfer import WeightTransferConfig
+    from vllm.distributed.weight_transfer.nccl_engine import (
+        NCCLWeightTransferEngine,
+        NCCLWeightTransferInitInfo,
+        NCCLWeightTransferUpdateInfo,
+    )
+
+    config = WeightTransferConfig(backend="nccl")
+    parallel_config = MagicMock(spec=ParallelConfig)
+    parallel_config.rank = 0
+    parallel_config.world_size = 1
+    parallel_config.data_parallel_rank = 0
+    parallel_config.data_parallel_index = 0
+
+    engine = NCCLWeightTransferEngine(config, parallel_config)
+    engine.init_transfer_engine(
+        NCCLWeightTransferInitInfo(
+            master_address=master_address,
+            master_port=master_port,
+            rank_offset=1,
+            world_size=world_size,
+        )
+    )
+
+    target = torch.zeros(30, dtype=torch.float32, device="cuda")
+
+    def apply_sparse_patches(patches: list[SparseWeightPatch]):
+        for patch in patches:
+            target.index_copy_(0, patch.indices.to(torch.long), patch.values)
+
+    update_info = NCCLWeightTransferUpdateInfo(
+        names=["test.weight"],
+        dtype_names=["float32"],
+        shapes=[[30]],
+        num_updates_list=[3],
+        update_kind="sparse_flat",
+    )
+    engine.receive_sparse_weights(update_info, apply_sparse_patches)
+    torch.accelerator.synchronize()
+
+    expected = torch.zeros(30, dtype=torch.float32, device="cuda")
+    expected[[1, 7, 25]] = torch.tensor(
+        [10.0, 20.0, 30.0], dtype=torch.float32, device="cuda"
+    )
+    success = torch.equal(target, expected)
+    engine.shutdown()
+    return {
+        "success": success,
+        "selected_values": target[[1, 7, 25]].cpu().tolist(),
+    }
+
+
+@pytest.mark.skipif(
+    torch.accelerator.device_count() < 2,
+    reason="Need at least 2 GPUs to run NCCL sparse weight transfer test.",
+)
+def test_nccl_sparse_weight_transfer_between_processes():
+    """Test NCCL sparse weight transfer from trainer to inference process."""
+    ray.init(ignore_reinit_error=True)
+
+    master_address = "127.0.0.1"
+    master_port = get_open_port()
+    world_size = 2
+
+    inference_future = inference_receive_sparse_tensor.remote(
+        master_address, master_port, world_size
+    )
+    trainer_future = trainer_broadcast_sparse_tensor.remote(
+        master_address, master_port, world_size
+    )
+
+    trainer_result, result = ray.get([trainer_future, inference_future])
+
+    assert trainer_result, "Trainer should complete successfully"
+    assert result["success"], (
+        "Sparse weight transfer failed. "
+        f"Received selected values: {result['selected_values']}"
+    )
+
+
 # --- Unit Tests: IPCWeightTransferUpdateInfo Validation ---
 
 
@@ -461,9 +674,101 @@ class TestIPCWeightTransferUpdateInfoValidation:
                 ipc_handles=ipc_handles,
             )
 
-    def test_missing_ipc_handles_raises(self):
-        """Test that omitting ipc_handles raises TypeError."""
-        with pytest.raises(TypeError):
+    def test_sparse_update_kind_rejected(self):
+        """Test that IPC backend rejects sparse update metadata."""
+        if torch.accelerator.device_count() < 1:
+            pytest.skip("Need at least 1 GPU for this test")
+
+        dummy_tensor = torch.ones(10, 10, device="cuda:0")
+        ipc_handle = reduce_tensor(dummy_tensor)
+        gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
+        ipc_handles = [{gpu_uuid: ipc_handle}]
+
+        with pytest.raises(NotImplementedError, match="dense updates"):
+            IPCWeightTransferUpdateInfo(
+                names=["layer.weight"],
+                dtype_names=["float32"],
+                shapes=[[10, 10]],
+                num_updates_list=[1],
+                ipc_handles=ipc_handles,
+                update_kind="sparse_flat",
+            )
+
+    def test_sparse_methods_not_supported(self):
+        """Test that IPC engine inherits sparse rejection from the base class."""
+        config = WeightTransferConfig(backend="ipc")
+        parallel_config = create_mock_parallel_config()
+        engine = IPCWeightTransferEngine(
+            config, parallel_config, MagicMock(spec=torch.nn.Module)
+        )
+
+        with pytest.raises(NotImplementedError, match="(?i)sparse weight updates"):
+            engine.receive_sparse_weights(MagicMock(), lambda _: None)
+        with pytest.raises(NotImplementedError, match="(?i)sparse weight updates"):
+            engine.trainer_send_sparse_weights(
+                iter([]),
+                {"mode": "http", "url": "http://localhost:8000"},
+            )
+
+    def test_valid_update_info_from_pickled(self, monkeypatch):
+        """Test creating IPCWeightTransferUpdateInfo from pickled handles."""
+        if torch.accelerator.device_count() < 1:
+            pytest.skip("Need at least 1 GPU for this test")
+
+        monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+
+        dummy_tensor = torch.ones(10, 10, device="cuda:0")
+        ipc_handle = reduce_tensor(dummy_tensor)
+        gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
+        ipc_handles = [{gpu_uuid: ipc_handle}]
+
+        pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
+
+        info = IPCWeightTransferUpdateInfo(
+            names=["layer.weight"],
+            dtype_names=["float32"],
+            shapes=[[10, 10]],
+            ipc_handles_pickled=pickled,
+        )
+        assert info.ipc_handles == ipc_handles
+        assert info.ipc_handles_pickled is None
+
+    def test_pickled_requires_insecure_serialization_flag(self, monkeypatch):
+        """Test that pickled handles are rejected unless env flag is enabled."""
+        monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "0")
+
+        with pytest.raises(ValueError, match="VLLM_ALLOW_INSECURE_SERIALIZATION=1"):
+            IPCWeightTransferUpdateInfo(
+                names=[],
+                dtype_names=[],
+                shapes=[],
+                ipc_handles_pickled=base64.b64encode(pickle.dumps([])).decode("utf-8"),
+            )
+
+    def test_both_handles_and_pickled_raises(self):
+        """Test that providing both ipc_handles and ipc_handles_pickled raises."""
+        if torch.accelerator.device_count() < 1:
+            pytest.skip("Need at least 1 GPU for this test")
+
+        dummy_tensor = torch.ones(10, 10, device="cuda:0")
+        ipc_handle = reduce_tensor(dummy_tensor)
+        gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
+        ipc_handles = [{gpu_uuid: ipc_handle}]
+
+        pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
+
+        with pytest.raises(ValueError, match="Cannot specify both"):
+            IPCWeightTransferUpdateInfo(
+                names=["layer.weight"],
+                dtype_names=["float32"],
+                shapes=[[10, 10]],
+                ipc_handles=ipc_handles,
+                ipc_handles_pickled=pickled,
+            )
+
+    def test_neither_handles_nor_pickled_raises(self):
+        """Test that providing neither ipc_handles nor ipc_handles_pickled raises."""
+        with pytest.raises(ValueError, match="must be provided"):
             IPCWeightTransferUpdateInfo(
                 names=["layer.weight"],
                 dtype_names=["float32"],
@@ -558,6 +863,28 @@ class TestIPCEngineParsing:
         assert gpu_uuid in update_info.ipc_handles[0]
         assert gpu_uuid in update_info.ipc_handles[1]
 
+    def test_parse_update_info_ignores_none_pickled_handles(self):
+        """Test Ray/asdict payloads with a null pickled field use ipc_handles."""
+        config = WeightTransferConfig(backend="ipc")
+        parallel_config = create_mock_parallel_config()
+        engine = IPCWeightTransferEngine(
+            config, parallel_config, MagicMock(spec=torch.nn.Module)
+        )
+        ipc_handles = [{"gpu-uuid": ("ipc-args",)}]
+
+        update_info = engine.parse_update_info(
+            {
+                "names": ["w1"],
+                "dtype_names": ["float32"],
+                "shapes": [[1]],
+                "ipc_handles": ipc_handles,
+                "ipc_handles_pickled": None,
+            }
+        )
+
+        assert isinstance(update_info, IPCWeightTransferUpdateInfo)
+        assert update_info.ipc_handles == ipc_handles
+
     def test_parse_update_info_both_handles_and_pickled_raises(self):
         """Test that providing both ipc_handles and ipc_handles_pickled raises."""
         if torch.accelerator.device_count() < 1:

tests/entrypoints/generate/generative_scoring/test_generative_scoring.py

@@ -18,13 +18,13 @@ from unittest.mock import MagicMock
 import pytest
 
 from vllm.config.multimodal import MultiModalConfig
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse
-from vllm.entrypoints.openai.generative_scoring.serving import (
+from vllm.entrypoints.generate.generative_scoring.serving import (
     GenerativeScoringItemResult,
     GenerativeScoringRequest,
     GenerativeScoringResponse,
-    OpenAIServingGenerativeScoring,
+    ServingGenerativeScoring,
 )
+from vllm.entrypoints.openai.engine.protocol import ErrorResponse
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.logprobs import Logprob
@@ -86,13 +86,13 @@ def _create_mock_engine():
     return mock_engine
 
 
-def _create_serving(mock_engine) -> OpenAIServingGenerativeScoring:
-    """Create an OpenAIServingGenerativeScoring instance with mocks."""
+def _create_serving(mock_engine) -> ServingGenerativeScoring:
+    """Create an ServingGenerativeScoring instance with mocks."""
     models = OpenAIServingModels(
         engine_client=mock_engine,
         base_model_paths=BASE_MODEL_PATHS,
     )
-    return OpenAIServingGenerativeScoring(mock_engine, models, request_logger=None)
+    return ServingGenerativeScoring(mock_engine, models, request_logger=None)
 
 
 def _create_mock_request_output(logprobs_dict: dict[int, float]) -> RequestOutput:
@@ -186,7 +186,7 @@ class TestProbabilityComputation:
         self, label_logprobs, apply_softmax, should_sum_to_one
     ):
         """Test probability computation for softmax and true probability modes."""
-        serving = OpenAIServingGenerativeScoring.__new__(OpenAIServingGenerativeScoring)
+        serving = ServingGenerativeScoring.__new__(ServingGenerativeScoring)
         probs = serving._compute_probabilities(
             label_logprobs, apply_softmax=apply_softmax
         )
@@ -211,7 +211,7 @@ class TestProbabilityComputation:
 
     def test_score_formula(self):
         """Test the score formula: P(token[0]) / (P(token[0]) + P(token[1]))."""
-        serving = OpenAIServingGenerativeScoring.__new__(OpenAIServingGenerativeScoring)
+        serving = ServingGenerativeScoring.__new__(ServingGenerativeScoring)
 
         # With logprobs -0.5 and -2.0, softmax gives higher prob to first token
         logprobs = {9454: -0.5, 2753: -2.0}

tests/entrypoints/generate/generative_scoring/test_generative_scoring_e2e.py

@@ -8,7 +8,7 @@ Tests verify the full HTTP request/response flow using RemoteOpenAIServer.
 import pytest
 import requests
 
-from ....utils import RemoteOpenAIServer
+from tests.utils import RemoteOpenAIServer
 
 MODEL_NAME = "Qwen/Qwen3-0.6B"
 

tests/entrypoints/openai/chat_completion/test_audio_in_video.py

@@ -14,8 +14,12 @@ from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer
 MODEL_NAME = "Qwen/Qwen2.5-Omni-3B"
 
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 def server():
+    # Use module scope so the server is started once and shared across all
+    # tests in this file. Starting a new vLLM server per test on XPU can
+    # cause the second server startup to hang silently and exceed the
+    # wait-for-server timeout, resulting in RuntimeError.
     args = [
         "--max-model-len",
         "16384",

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/weight_transfer/test_weight_transfer_llm.py

@@ -48,6 +48,7 @@ class MockUpdateInfo(WeightTransferUpdateInfo):
     names: list[str] | None = None
     dtype_names: list[str] | None = None
     shapes: list[list[int]] | None = None
+    num_updates_list: list[int] | None = None
 
 
 class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo]):
@@ -87,6 +88,15 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo
         # (In real implementation, this would receive and load actual weights)
         load_weights([])
 
+    def receive_sparse_weights(
+        self,
+        update_info: MockUpdateInfo,
+        apply_patches: Callable[[list], None],
+    ) -> None:
+        MockWeightTransferEngine.receive_weights_called = True
+        MockWeightTransferEngine.last_update_info = update_info
+        apply_patches([])
+
     def shutdown(self) -> None:
         MockWeightTransferEngine.shutdown_called = True
 
@@ -198,8 +208,6 @@ def test_update_weights_calls_engine():
         llm.init_weight_transfer_engine(
             WeightTransferInitRequest(init_info={"test_param": "init"})
         )
-
-        # Start weight update (required before update_weights)
         llm.start_weight_update(is_checkpoint_format=True)
 
         # Call update_weights
@@ -232,14 +240,67 @@ def test_update_weights_calls_engine():
             assert dtypes == test_dtypes
             assert shapes == test_shapes
 
-        # Finish weight update
+        llm.finish_weight_update()
+
+
+@create_new_process_for_each_test()
+def test_update_weights_passes_sparse_metadata():
+    """Test sparse update metadata is forwarded unchanged to the engine."""
+    if torch.accelerator.device_count() < 1:
+        pytest.skip("Need at least 1 GPU for this test")
+
+    os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+    os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
+
+    with patch(
+        "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",
+        mock_create_engine,
+    ):
+        llm = LLM(
+            model=MODEL_NAME,
+            enforce_eager=True,
+            load_format="dummy",
+            tensor_parallel_size=1,
+            weight_transfer_config=WeightTransferConfig(backend="nccl"),
+        )
+
+        llm.init_weight_transfer_engine(
+            WeightTransferInitRequest(init_info={"test_param": "init"})
+        )
+        llm.start_weight_update(is_checkpoint_format=False)
+
+        llm.update_weights(
+            WeightTransferUpdateRequest(
+                update_info={
+                    "names": ["layer.weight"],
+                    "dtype_names": ["bfloat16"],
+                    "shapes": [[100]],
+                    "num_updates_list": [3],
+                    "update_kind": "sparse_flat",
+                }
+            )
+        )
+
+        def check_sparse_update_called(self):
+            engine = self.weight_transfer_engine
+            if not engine.receive_weights_called:
+                return None
+            info = engine.last_update_info
+            return (
+                info.update_kind,
+                info.num_updates_list,
+            )
+
+        results = llm.collective_rpc(check_sparse_update_called)
+        for result in results:
+            assert result == ("sparse_flat", [3])
+
         llm.finish_weight_update()
 
 
 @create_new_process_for_each_test()
 def test_full_weight_transfer_flow():
-    """Test the complete weight transfer flow:
-    init -> start -> update -> finish."""
+    """Test the complete weight transfer flow: init -> start -> update -> finish."""
     if torch.accelerator.device_count() < 1:
         pytest.skip("Need at least 1 GPU for this test")
 

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/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

@@ -522,6 +522,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
     "Qwen2MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen1.5-MoE-A2.7B-Chat"),
     "Qwen3ForCausalLM": _HfExamplesInfo("Qwen/Qwen3-8B"),
     "Qwen3MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"),
+    "MellumForCausalLM": _HfExamplesInfo("JetBrains/Mellum2-12B-A2.5B-Base"),
     "Qwen3NextForCausalLM": _HfExamplesInfo(
         "Qwen/Qwen3-Next-80B-A3B-Instruct",
         extras={"tiny-random": "tiny-random/qwen3-next-moe"},
@@ -1372,7 +1373,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/tool_use/test_chat_completion_request_validations.py

@@ -116,3 +116,68 @@ def test_no_reasoning_fields_unchanged():
     assistant_msg = request.messages[1]
     assert assistant_msg.get("reasoning") is None
     assert "reasoning_content" not in assistant_msg
+
+
+SAMPLE_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "get_weather",
+        "description": "Get the weather",
+        "parameters": {
+            "type": "object",
+            "properties": {"location": {"type": "string"}},
+        },
+    },
+}
+
+
+def test_structured_outputs_with_named_tool_choice_rejected():
+    """structured_outputs cannot be combined with a named tool_choice."""
+    with pytest.raises(
+        ValueError,
+        match="structured outputs or tools, not both",
+    ):
+        ChatCompletionRequest.model_validate(
+            {
+                "messages": [{"role": "user", "content": "Hello"}],
+                "model": "facebook/opt-125m",
+                "tools": [SAMPLE_TOOL],
+                "tool_choice": {
+                    "type": "function",
+                    "function": {"name": "get_weather"},
+                },
+                "structured_outputs": {"json": {"type": "object"}},
+            }
+        )
+
+
+def test_structured_outputs_with_auto_tool_choice_allowed():
+    """structured_outputs with tool_choice 'auto' should be allowed."""
+    request = ChatCompletionRequest.model_validate(
+        {
+            "messages": [{"role": "user", "content": "Hello"}],
+            "model": "facebook/opt-125m",
+            "tools": [SAMPLE_TOOL],
+            "tool_choice": "auto",
+            "structured_outputs": {"json": {"type": "object"}},
+        }
+    )
+    assert request.tool_choice == "auto"
+
+
+def test_multiple_structured_outputs_rejected():
+    """Only one kind of structured output constraint is allowed."""
+    with pytest.raises(
+        ValueError,
+        match="You can only use one kind of constraints",
+    ):
+        ChatCompletionRequest.model_validate(
+            {
+                "messages": [{"role": "user", "content": "Hello"}],
+                "model": "facebook/opt-125m",
+                "structured_outputs": {
+                    "json": {"type": "object"},
+                    "regex": ".*",
+                },
+            }
+        )

tests/utils_/test_import_utils.py

@@ -1,8 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from unittest.mock import MagicMock, patch
+
 import pytest
 
-from vllm.utils.import_utils import PlaceholderModule
+from vllm.utils.import_utils import PlaceholderModule, _has_module
 
 
 def _raises_module_not_found():
@@ -44,3 +46,94 @@ def test_placeholder_module_error_handling():
     with _raises_module_not_found():
         # Test conflict with internal __module attribute
         _ = placeholder_attr.module
+
+
+class TestHasModule:
+    """Tests for _has_module with trial import verification."""
+
+    def setup_method(self):
+        # Clear the @cache between tests so each test gets a fresh call
+        _has_module.cache_clear()
+
+    def test_returns_true_for_importable_stdlib_module(self):
+        assert _has_module("json") is True
+
+    def test_returns_false_for_nonexistent_module(self):
+        assert _has_module("nonexistent_module_xyz_12345") is False
+
+    def test_returns_false_when_find_spec_succeeds_but_import_fails(self):
+        """Simulate a native extension whose shared library is missing.
+
+        ``find_spec`` finds the package on disk, but the actual import
+        raises ``ImportError`` (e.g. missing ``libcudart.so``).
+        """
+        fake_spec = MagicMock()
+
+        with (
+            patch(
+                "vllm.utils.import_utils.importlib.util.find_spec",
+                return_value=fake_spec,
+            ),
+            patch(
+                "vllm.utils.import_utils.importlib.import_module",
+                side_effect=ImportError(
+                    "libcudart.so.12: cannot open shared object file"
+                ),
+            ),
+        ):
+            assert _has_module("fake_native_ext") is False
+
+    def test_returns_false_on_os_error_during_import(self):
+        """Some shared-library failures surface as ``OSError``."""
+        fake_spec = MagicMock()
+
+        with (
+            patch(
+                "vllm.utils.import_utils.importlib.util.find_spec",
+                return_value=fake_spec,
+            ),
+            patch(
+                "vllm.utils.import_utils.importlib.import_module",
+                side_effect=OSError("cannot load library"),
+            ),
+        ):
+            assert _has_module("fake_native_ext_os") is False
+
+    def test_returns_false_on_unexpected_error_during_import(self):
+        """A broken extension may raise a non-import error (e.g. ``RuntimeError``).
+
+        Such modules are not usable, so ``_has_module`` should still return
+        ``False`` rather than letting the exception propagate.
+        """
+        fake_spec = MagicMock()
+
+        with (
+            patch(
+                "vllm.utils.import_utils.importlib.util.find_spec",
+                return_value=fake_spec,
+            ),
+            patch(
+                "vllm.utils.import_utils.importlib.import_module",
+                side_effect=RuntimeError("CUDA driver version is insufficient"),
+            ),
+        ):
+            assert _has_module("fake_broken_ext") is False
+
+    def test_returns_false_when_find_spec_raises(self):
+        """``find_spec`` itself can raise for dotted names whose parent package
+        fails to import. This should be treated as the module being unavailable.
+        """
+        with patch(
+            "vllm.utils.import_utils.importlib.util.find_spec",
+            side_effect=ModuleNotFoundError("No module named 'fake_parent'"),
+        ):
+            assert _has_module("fake_parent.child") is False
+
+    def test_result_is_cached(self):
+        """Verify the @cache decorator prevents repeated imports."""
+        _has_module("json")  # prime the cache
+
+        with patch("vllm.utils.import_utils.importlib.util.find_spec") as mock_spec:
+            result = _has_module("json")  # should hit cache
+            mock_spec.assert_not_called()
+            assert result is True

tests/v1/spec_decode/test_acceptance_length.py

@@ -39,6 +39,8 @@ class Eagle3ModelConfig:
     marks: list = field(default_factory=list)
     # Custom relative tolerance (defaults to DEFAULT_RTOL if None)
     rtol: float | None = None
+    # ROCm-specific test configuration
+    rocm_expected_acceptance_lengths_per_pos: list[float] = field(default_factory=list)
 
 
 # Model configurations for EAGLE3 acceptance length tests.
@@ -69,6 +71,7 @@ EAGLE3_MODEL_CONFIGS = [
         # FLASHINFER incompatible: gpt-oss-20b uses sink attention which
         # FLASHINFER does not support ("sink setting not supported")
         excluded_backends={AttentionBackendEnum.FLASHINFER},
+        rocm_expected_acceptance_lengths_per_pos=[0.7040, 0.4820, 0.3350],
     ),
     Eagle3ModelConfig(
         verifier="Qwen/Qwen3-VL-30B-A3B-Instruct-FP8",
@@ -99,16 +102,14 @@ EXCLUDED_BACKENDS = {AttentionBackendEnum.FLEX_ATTENTION}
 
 
 def get_available_attention_backends() -> list[str]:
+    if current_platform.is_rocm():
+        return ["auto"]
+
     # Check if get_valid_backends is actually defined in the platform class
     # (not just returning None from __getattr__)
     get_valid_backends = getattr(current_platform.__class__, "get_valid_backends", None)
     if get_valid_backends is None:
-        if current_platform.is_rocm():
-            # ROCm uses Triton as its default attention backend since
-            # Flash Attention is not supported.
-            return ["TRITON_ATTN"]
-        else:
-            return ["FLASH_ATTN"]
+        return ["FLASH_ATTN"]
 
     device_capability = current_platform.get_device_capability()
     if device_capability is None:
@@ -167,6 +168,8 @@ def get_mt_bench_prompts(
         disable_shuffle=False,
         skip_chat_template=False,
         trust_remote_code=False,
+        enable_multimodal_chat=False,
+        request_id_prefix="",
     )
     samples = get_samples(args, tokenizer)
     prompt_ids = [
@@ -233,9 +236,12 @@ def test_eagle3_acceptance_length(
     monkeypatch: pytest.MonkeyPatch,
 ):
     # Skip if this backend is incompatible with the model
-    backend_enum = AttentionBackendEnum[attention_backend]
-    if backend_enum in model_config.excluded_backends:
-        pytest.skip(f"{attention_backend} is incompatible with {model_config.id}")
+    attention_config = None
+    if attention_backend != "auto":
+        backend_enum = AttentionBackendEnum[attention_backend]
+        if backend_enum in model_config.excluded_backends:
+            pytest.skip(f"{attention_backend} is incompatible with {model_config.id}")
+        attention_config = {"backend": attention_backend}
 
     with monkeypatch.context() as m:
         m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
@@ -247,11 +253,16 @@ def test_eagle3_acceptance_length(
                 "model": model_config.drafter,
                 "num_speculative_tokens": num_spec_tokens,
             },
-            attention_config={"backend": attention_backend},
+            attention_config=attention_config,
             tensor_parallel_size=tp_size,
             gpu_memory_utilization=0.7,
             disable_log_stats=False,
             max_model_len=DEFAULT_MAX_MODEL_LEN,
+            # Qwen/Qwen3-30B-A3B-FP8 with TP=4 needs EP
+            # https://github.com/vllm-project/vllm/issues/25292
+            enable_expert_parallel=(
+                tp_size == 4 and "Qwen3-VL" in model_config.verifier
+            ),
         ) as vllm_runner:
             tokenizer = vllm_runner.llm.get_tokenizer()
             prompt_ids = get_mt_bench_prompts(tokenizer, DEFAULT_NUM_PROMPTS)
@@ -272,6 +283,11 @@ def test_eagle3_acceptance_length(
             expected = model_config.expected_acceptance_length
             actual_per_pos = results["acceptance_lengths_per_pos"]
             expected_per_pos = model_config.expected_acceptance_lengths_per_pos
+            if (
+                current_platform.is_rocm()
+                and model_config.rocm_expected_acceptance_lengths_per_pos
+            ):
+                expected_per_pos = model_config.rocm_expected_acceptance_lengths_per_pos
 
             rel_error = abs(actual_acceptance_length - expected) / expected
 
@@ -294,14 +310,14 @@ def test_eagle3_acceptance_length(
                     zip(actual_per_pos, expected_per_pos)
                 ):
                     if exp > 0:
-                        pos_rel_error = abs(actual - exp) / exp
-                        assert pos_rel_error <= rtol, (
+                        min_expected = exp * (1 - rtol)
+                        assert actual >= min_expected, (
                             f"Per-position acceptance length regression at pos {pos} "
                             f"for {model_config.id}!\n"
                             f"  Expected: {exp:.3f}\n"
                             f"  Actual:   {actual:.3f}\n"
-                            f"  Relative error: {pos_rel_error:.2%} "
-                            f"(tolerance: {rtol:.2%})"
+                            f"  Minimum:  {min_expected:.3f}\n"
+                            f"  Tolerance: rtol={rtol:.2%}"
                         )
 
             print(

tests/v1/worker/test_gpu_model_runner.py

@@ -7,6 +7,7 @@ from unittest.mock import Mock
 import numpy as np
 import pytest
 import torch
+import torch.nn as nn
 
 import vllm.v1.worker.gpu_model_runner as gpu_model_runner_module
 from vllm.config import (
@@ -22,6 +23,7 @@ from vllm.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
+from vllm.distributed.weight_transfer.base import SparseWeightPatch
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
 from vllm.platforms import current_platform
@@ -784,6 +786,73 @@ def test_sample_passes_reordered_draft_probs_to_rejection_sampler():
     assert torch.equal(passed_draft_probs, expected_draft_probs)
 
 
+def test_apply_sparse_weight_patches_updates_only_selected_entries():
+    class DummyModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = nn.Parameter(torch.zeros(6, dtype=torch.float32))
+
+    runner = object.__new__(GPUModelRunner)
+    runner.model = DummyModel()
+
+    runner.apply_sparse_weight_patches(
+        [
+            SparseWeightPatch(
+                name="weight",
+                indices=torch.tensor([1, 4], dtype=torch.int32),
+                values=torch.tensor([3.5, -2.0], dtype=torch.float32),
+            )
+        ]
+    )
+
+    expected = torch.tensor([0.0, 3.5, 0.0, 0.0, -2.0, 0.0], dtype=torch.float32)
+    assert torch.equal(runner.get_model().weight.data, expected)
+
+
+def test_apply_sparse_weight_patches_rejects_mismatched_lengths():
+    class DummyModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
+
+    runner = object.__new__(GPUModelRunner)
+    runner.model = DummyModel()
+
+    with pytest.raises(ValueError, match="matching lengths"):
+        runner.apply_sparse_weight_patches(
+            [
+                SparseWeightPatch(
+                    name="weight",
+                    indices=torch.tensor([1, 2], dtype=torch.int32),
+                    values=torch.tensor([1.0], dtype=torch.float32),
+                )
+            ]
+        )
+
+
+def test_apply_sparse_weight_patches_rejects_non_contiguous_param():
+    class DummyModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = nn.Parameter(
+                torch.arange(12, dtype=torch.float32).view(3, 4).t()
+            )
+
+    runner = object.__new__(GPUModelRunner)
+    runner.model = DummyModel()
+
+    with pytest.raises(NotImplementedError, match="contiguous params"):
+        runner.apply_sparse_weight_patches(
+            [
+                SparseWeightPatch(
+                    name="weight",
+                    indices=torch.tensor([1], dtype=torch.int32),
+                    values=torch.tensor([1.0], dtype=torch.float32),
+                )
+            ]
+        )
+
+
 def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order(default_vllm_config):
     torch.set_default_dtype(torch.float16)
     layer_0 = "model.layers.0.self_attn.attn"

tests/v1/worker/test_gpu_worker_weight_transfer.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm.config.parallel import ParallelConfig
+from vllm.config.weight_transfer import WeightTransferConfig
+from vllm.distributed.weight_transfer.base import SparseWeightPatch
+from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferEngine
+from vllm.v1.worker.gpu_worker import Worker
+
+
+def _make_nccl_engine() -> NCCLWeightTransferEngine:
+    parallel_config = MagicMock(spec=ParallelConfig)
+    parallel_config.rank = 0
+    parallel_config.world_size = 1
+    parallel_config.data_parallel_rank = 0
+    parallel_config.data_parallel_index = 0
+    return NCCLWeightTransferEngine(
+        WeightTransferConfig(backend="nccl"),
+        parallel_config,
+        MagicMock(spec=torch.nn.Module),
+    )
+
+
+def test_update_weights_sparse_dispatches_to_sparse_receive(monkeypatch):
+    monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
+
+    worker = object.__new__(Worker)
+    worker.device = "cpu"
+    worker.parallel_config = SimpleNamespace(world_size=1)
+    worker.weight_transfer_engine = _make_nccl_engine()
+    worker._weight_update_active = True
+    worker._is_checkpoint_format = False
+
+    applied_patches = []
+
+    def apply_sparse_weight_patches(patches):
+        applied_patches.extend(patches)
+
+    worker.model_runner = SimpleNamespace(
+        apply_sparse_weight_patches=apply_sparse_weight_patches,
+    )
+
+    received_kinds = []
+
+    def receive_sparse_weights(update_info, apply_patches):
+        received_kinds.append(update_info.update_kind)
+        apply_patches(
+            [
+                SparseWeightPatch(
+                    name="layer.weight",
+                    indices=torch.tensor([1], dtype=torch.int32),
+                    values=torch.tensor([2.0], dtype=torch.float32),
+                )
+            ]
+        )
+
+    worker.weight_transfer_engine.receive_sparse_weights = receive_sparse_weights
+
+    Worker.update_weights(
+        worker,
+        {
+            "names": ["layer.weight"],
+            "dtype_names": ["float32"],
+            "shapes": [[4]],
+            "num_updates_list": [1],
+            "update_kind": "sparse_flat",
+        },
+    )
+
+    assert received_kinds == ["sparse_flat"]
+    assert len(applied_patches) == 1
+    assert torch.equal(applied_patches[0].indices, torch.tensor([1], dtype=torch.int32))
+
+
+def test_update_weights_sparse_rejects_tp_or_pp(monkeypatch):
+    monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
+
+    worker = object.__new__(Worker)
+    worker.device = "cpu"
+    worker.parallel_config = SimpleNamespace(world_size=2)
+    worker.weight_transfer_engine = _make_nccl_engine()
+    worker._weight_update_active = True
+    worker._is_checkpoint_format = False
+    worker.model_runner = SimpleNamespace(apply_sparse_weight_patches=lambda _: None)
+
+    with pytest.raises(NotImplementedError, match="TP=1 and PP=1"):
+        Worker.update_weights(
+            worker,
+            {
+                "names": ["layer.weight"],
+                "dtype_names": ["float32"],
+                "shapes": [[4]],
+                "num_updates_list": [1],
+                "update_kind": "sparse_flat",
+            },
+        )
+    assert worker._weight_update_active is False
+    assert worker._is_checkpoint_format is True
+
+
+def test_update_weights_sparse_rejects_checkpoint_format(monkeypatch):
+    monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
+
+    worker = object.__new__(Worker)
+    worker.device = "cpu"
+    worker.parallel_config = SimpleNamespace(world_size=1)
+    worker.weight_transfer_engine = _make_nccl_engine()
+    worker._weight_update_active = True
+    worker._is_checkpoint_format = True
+    worker.model_runner = SimpleNamespace(model=MagicMock())
+
+    with pytest.raises(ValueError, match="start_weight_update"):
+        Worker.update_weights(
+            worker,
+            {
+                "names": ["layer.weight"],
+                "dtype_names": ["float32"],
+                "shapes": [[4]],
+                "num_updates_list": [1],
+                "update_kind": "sparse_flat",
+            },
+        )
+    assert worker._weight_update_active is False
+    assert worker._is_checkpoint_format is True
+
+
+def test_update_weights_resets_state_when_update_info_is_invalid(monkeypatch):
+    monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None)
+
+    worker = object.__new__(Worker)
+    worker.device = "cpu"
+    worker.parallel_config = SimpleNamespace(world_size=1)
+    worker.weight_transfer_engine = _make_nccl_engine()
+    worker._weight_update_active = True
+    worker._is_checkpoint_format = False
+
+    with pytest.raises(ValueError, match="cannot be empty"):
+        Worker.update_weights(
+            worker,
+            {
+                "names": [],
+                "dtype_names": [],
+                "shapes": [],
+                "num_updates_list": [],
+                "update_kind": "sparse_flat",
+            },
+        )
+    assert worker._weight_update_active is False
+    assert worker._is_checkpoint_format is True

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,

vllm/distributed/weight_transfer/base.py

@@ -4,8 +4,8 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterator
-from dataclasses import dataclass, field
-from typing import Any, Generic, TypeVar
+from dataclasses import KW_ONLY, dataclass, field
+from typing import Any, Generic, Literal, TypeVar
 
 import torch
 
@@ -28,7 +28,44 @@ class WeightTransferInitInfo(ABC):  # noqa: B024
 class WeightTransferUpdateInfo(ABC):  # noqa: B024
     """Base class for backend-specific weight update info."""
 
-    pass
+    _: KW_ONLY
+    update_kind: Literal["dense", "sparse_flat"] = "dense"
+    """Weight update format."""
+    num_updates_list: list[int] | None = None
+    """Number of sparse entries to receive for each parameter in ``names``."""
+
+    def __post_init__(self) -> None:
+        if self.update_kind not in ("dense", "sparse_flat"):
+            raise ValueError(f"Unsupported update_kind: {self.update_kind}")
+        if self.update_kind == "dense":
+            if self.num_updates_list is not None:
+                raise ValueError(
+                    "Sparse metadata is only supported for `update_kind='sparse_flat'`"
+                )
+            return
+
+        if self.num_updates_list is None:
+            raise ValueError("`num_updates_list` is required for sparse updates")
+        if len(self.num_updates_list) == 0:
+            raise ValueError("`num_updates_list` cannot be empty for sparse updates")
+        if any(num_updates < 0 for num_updates in self.num_updates_list):
+            raise ValueError("Sparse `num_updates_list` entries must be non-negative")
+
+        names = getattr(self, "names", None)
+        if names is not None and len(self.num_updates_list) != len(names):
+            raise ValueError(
+                f"`num_updates_list` should be of the same size as `names`: "
+                f"got {len(self.num_updates_list)} and {len(names)}"
+            )
+
+
+@dataclass
+class SparseWeightPatch:
+    """A sparse in-place patch for one existing parameter."""
+
+    name: str
+    indices: torch.Tensor
+    values: torch.Tensor
 
 
 # API-level request classes (accept dicts for backend-agnostic serialization)
@@ -150,6 +187,16 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
         """
         raise NotImplementedError
 
+    def receive_sparse_weights(
+        self,
+        update_info: TUpdateInfo,
+        apply_patches: Callable[[list[SparseWeightPatch]], None],
+    ) -> None:
+        """Receive sparse weight patches from the trainer."""
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support sparse weight updates"
+        )
+
     @abstractmethod
     def shutdown(self) -> None:
         """
@@ -184,3 +231,11 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
             >>> engine.trainer_send_weights(param_iter, trainer_args)
         """
         raise NotImplementedError
+
+    @staticmethod
+    def trainer_send_sparse_weights(
+        _iterator: Iterator[SparseWeightPatch],
+        _trainer_args: dict[str, Any] | Any,
+    ) -> None:
+        """Send sparse weight patches from trainer to inference workers."""
+        raise NotImplementedError("Sparse weight updates are not supported")

vllm/distributed/weight_transfer/ipc_engine.py

@@ -74,10 +74,12 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
     names: list[str]
     dtype_names: list[str]
     shapes: list[list[int]]
-    ipc_handles: list[dict[str, tuple]] | dict[str, tuple]
+    ipc_handles: list[dict[str, tuple]] | dict[str, tuple] | None = None
     """IPC handles mapping physical GPU UUID to rebuild_cuda_tensor args.
     For non-packed mode: list of per-parameter handle dicts.
     For packed mode: single handle dict for the packed buffer."""
+    ipc_handles_pickled: str | None = None
+    """Base64-encoded pickled IPC handles, used for HTTP transport."""
     tensor_sizes: list[int] | None = None
     """Per-parameter sizes in bytes within the packed buffer.
     Required when packed=True, unused otherwise."""
@@ -85,6 +87,29 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
     """Whether this update uses packed tensor format."""
 
     def __post_init__(self):
+        super().__post_init__()
+        if self.update_kind != "dense":
+            raise NotImplementedError("IPC weight transfer only supports dense updates")
+
+        if self.ipc_handles_pickled is not None:
+            if self.ipc_handles is not None:
+                raise ValueError(
+                    "Cannot specify both `ipc_handles` and `ipc_handles_pickled`"
+                )
+
+            if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION:
+                raise ValueError(
+                    "Refusing to deserialize `ipc_handles_pickled` without "
+                    "VLLM_ALLOW_INSECURE_SERIALIZATION=1"
+                )
+
+            self.ipc_handles = pickle.loads(base64.b64decode(self.ipc_handles_pickled))
+            self.ipc_handles_pickled = None
+
+        if self.ipc_handles is None:
+            raise ValueError(
+                "Either `ipc_handles` or `ipc_handles_pickled` must be provided"
+            )
         num_params = len(self.names)
         if len(self.dtype_names) != num_params:
             raise ValueError(
@@ -153,8 +178,9 @@ class IPCWeightTransferEngine(
         Requires ``VLLM_ALLOW_INSECURE_SERIALIZATION=1`` because the
         payload is deserialized via ``pickle.loads``.
         """
-        if "ipc_handles_pickled" in update_dict:
-            if "ipc_handles" in update_dict:
+        pickled = update_dict.pop("ipc_handles_pickled", None)
+        if pickled is not None:
+            if update_dict.get("ipc_handles") is not None:
                 raise ValueError(
                     "Cannot specify both `ipc_handles` and `ipc_handles_pickled`"
                 )
@@ -165,7 +191,6 @@ class IPCWeightTransferEngine(
                     "VLLM_ALLOW_INSECURE_SERIALIZATION=1"
                 )
 
-            pickled = update_dict.pop("ipc_handles_pickled")
             update_dict["ipc_handles"] = pickle.loads(base64.b64decode(pickled))
 
         return super().parse_update_info(update_dict)

vllm/distributed/weight_transfer/nccl_engine.py

@@ -14,6 +14,7 @@ if TYPE_CHECKING:
 from vllm.config.parallel import ParallelConfig
 from vllm.config.weight_transfer import WeightTransferConfig
 from vllm.distributed.weight_transfer.base import (
+    SparseWeightPatch,
     WeightTransferEngine,
     WeightTransferInitInfo,
     WeightTransferUpdateInfo,
@@ -81,6 +82,7 @@ class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
 
     def __post_init__(self):
         """Validate that all lists have the same length."""
+        super().__post_init__()
         num_params = len(self.names)
         if len(self.dtype_names) != num_params:
             raise ValueError(
@@ -92,6 +94,13 @@ class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
                 f"`shapes` should be of the same size as `names`: "
                 f"got {len(self.shapes)} and {len(self.names)}"
             )
+        if self.update_kind == "dense":
+            return
+
+        if self.packed:
+            raise ValueError(
+                "`update_kind='sparse_flat'` cannot be combined with `packed=True`"
+            )
 
 
 class NCCLWeightTransferEngine(
@@ -178,6 +187,11 @@ class NCCLWeightTransferEngine(
                 "NCCL weight transfer not initialized. "
                 "Call init_transfer_engine() first."
             )
+        if update_info.update_kind != "dense":
+            raise ValueError(
+                "Sparse updates must use `receive_sparse_weights`, not "
+                "`receive_weights`"
+            )
 
         if update_info.packed:
             # Build iterator of (name, (shape, dtype)) from update_info
@@ -209,6 +223,42 @@ class NCCLWeightTransferEngine(
                 load_weights([(name, weight)])
                 del weight
 
+    def receive_sparse_weights(
+        self,
+        update_info: NCCLWeightTransferUpdateInfo,
+        apply_patches: Callable[[list[SparseWeightPatch]], None],
+    ) -> None:
+        """Receive sparse flat-index patches from trainer via NCCL."""
+        if self.model_update_group is None:
+            raise RuntimeError(
+                "NCCL weight transfer not initialized. "
+                "Call init_transfer_engine() first."
+            )
+        if update_info.update_kind != "sparse_flat":
+            raise ValueError("Sparse receive path requires `update_kind='sparse_flat'`")
+        assert update_info.num_updates_list is not None
+
+        for name, dtype_name, num_updates in zip(
+            update_info.names,
+            update_info.dtype_names,
+            update_info.num_updates_list,
+        ):
+            dtype = getattr(torch, dtype_name)
+            device = torch.accelerator.current_device_index()
+            indices = torch.empty(num_updates, dtype=torch.int32, device=device)
+            values = torch.empty(num_updates, dtype=dtype, device=device)
+            self.model_update_group.broadcast(
+                indices, src=0, stream=torch.cuda.current_stream()
+            )
+            self.model_update_group.broadcast(
+                values, src=0, stream=torch.cuda.current_stream()
+            )
+            apply_patches(
+                [SparseWeightPatch(name=name, indices=indices, values=values)]
+            )
+            del indices
+            del values
+
     def shutdown(self) -> None:
         if self.model_update_group is not None:
             # Clean up the communicator by removing the reference
@@ -272,6 +322,27 @@ class NCCLWeightTransferEngine(
                     stream=args.stream or torch.cuda.current_stream(),
                 )
 
+    @staticmethod
+    def trainer_send_sparse_weights(
+        iterator: Iterator[SparseWeightPatch],
+        trainer_args: dict[str, Any] | NCCLTrainerSendWeightsArgs,
+    ) -> None:
+        """Broadcast sparse flat-index patches from trainer to vLLM workers."""
+        if isinstance(trainer_args, dict):
+            args = NCCLTrainerSendWeightsArgs(**trainer_args)
+        else:
+            args = trainer_args
+
+        if args.packed:
+            raise ValueError(
+                "Sparse NCCL updates cannot be combined with `packed=True`"
+            )
+
+        stream = args.stream or torch.cuda.current_stream()
+        for patch in iterator:
+            args.group.broadcast(patch.indices, src=args.src, stream=stream)
+            args.group.broadcast(patch.values, src=args.src, stream=stream)
+
     @staticmethod
     def trainer_init(
         init_info: NCCLWeightTransferInitInfo | dict,

vllm/entrypoints/generate/api_router.py

@@ -41,6 +41,10 @@ def register_generate_api_routers(app: FastAPI):
 
     register_anthropic_api_router(app)
 
+    from .generative_scoring.api_router import register_generative_scoring_api_router
+
+    register_generative_scoring_api_router(app)
+
 
 async def init_generate_state(
     engine_client: "EngineClient",
@@ -185,3 +189,11 @@ async def init_generate_state(
         if "generate" in supported_tasks
         else None
     )
+
+    from .generative_scoring.serving import ServingGenerativeScoring
+
+    state.serving_generative_scoring = ServingGenerativeScoring(
+        engine_client,
+        state.openai_serving_models,
+        request_logger=request_logger,
+    )

vllm/entrypoints/generate/generative_scoring/api_router.py

@@ -1,34 +1,25 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from http import HTTPStatus
-from typing import TYPE_CHECKING
 
 from fastapi import APIRouter, Depends, FastAPI, Request
 from fastapi.responses import JSONResponse
 
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse
-from vllm.entrypoints.openai.generative_scoring.serving import (
+from vllm.entrypoints.generate.generative_scoring.serving import (
     GenerativeScoringResponse,
-    OpenAIServingGenerativeScoring,
+    ServingGenerativeScoring,
 )
+from vllm.entrypoints.openai.engine.protocol import ErrorResponse
 from vllm.entrypoints.openai.utils import validate_json_request
 from vllm.entrypoints.utils import load_aware_call, with_cancellation
 from vllm.logger import init_logger
 
-if TYPE_CHECKING:
-    from argparse import Namespace
-
-    from starlette.datastructures import State
-
-    from vllm.engine.protocol import EngineClient
-    from vllm.entrypoints.logger import RequestLogger
-
 router = APIRouter()
 
 logger = init_logger(__name__)
 
 
-def generative_scoring(request: Request) -> OpenAIServingGenerativeScoring | None:
+def generative_scoring(request: Request) -> ServingGenerativeScoring | None:
     return request.app.state.serving_generative_scoring
 
 
@@ -51,7 +42,7 @@ async def create_generative_scoring(raw_request: Request):
 
     raw_body = await raw_request.json()
 
-    from vllm.entrypoints.openai.generative_scoring.serving import (
+    from vllm.entrypoints.generate.generative_scoring.serving import (
         GenerativeScoringRequest,
     )
 
@@ -68,20 +59,3 @@ async def create_generative_scoring(raw_request: Request):
 
 def register_generative_scoring_api_router(app: FastAPI):
     app.include_router(router)
-
-
-async def init_generative_scoring_state(
-    engine_client: "EngineClient",
-    state: "State",
-    args: "Namespace",
-    request_logger: "RequestLogger | None",
-):
-    from vllm.entrypoints.openai.generative_scoring.serving import (
-        OpenAIServingGenerativeScoring,
-    )
-
-    state.serving_generative_scoring = OpenAIServingGenerativeScoring(
-        engine_client,
-        state.openai_serving_models,
-        request_logger=request_logger,
-    )

vllm/entrypoints/generate/generative_scoring/serving.py

@@ -142,7 +142,7 @@ class GenerativeScoringResponse(OpenAIBaseModel):
 # ============================================================================
 
 
-class OpenAIServingGenerativeScoring(OpenAIServing):
+class ServingGenerativeScoring(OpenAIServing):
     """Serving class for generative scoring computation.
 
     This class handles computing the probability of specified token IDs

vllm/entrypoints/llm.py

@@ -873,14 +873,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
         )
 
     def start_weight_update(self, is_checkpoint_format: bool = True) -> None:
-        """
-        Start a new weight update.
-
-        Args:
-            is_checkpoint_format: Whether incoming weights are in checkpoint
-                format (need layerwise processing) or kernel format (direct
-                copy).
-        """
+        """Start a new weight update."""
         self.llm_engine.collective_rpc(
             "start_weight_update",
             kwargs={"is_checkpoint_format": is_checkpoint_format},
@@ -902,9 +895,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
         )
 
     def finish_weight_update(self) -> None:
-        """
-        Finish the current weight update.
-        """
+        """Finish the current weight update."""
         self.llm_engine.collective_rpc("finish_weight_update")
 
     def __repr__(self) -> str:

vllm/entrypoints/openai/api_server.py

@@ -196,7 +196,7 @@ def build_app(
     register_sagemaker_api_router(app, supported_tasks, model_config)
 
     if "generate" in supported_tasks:
-        from vllm.entrypoints.openai.generate.api_router import (
+        from vllm.entrypoints.generate.api_router import (
             register_generate_api_routers,
         )
 
@@ -220,12 +220,6 @@ def build_app(
 
         elastic_ep_attach_router(app)
 
-        from vllm.entrypoints.openai.generative_scoring.api_router import (
-            register_generative_scoring_api_router,
-        )
-
-        register_generative_scoring_api_router(app)
-
     if "generate" in supported_tasks or "render" in supported_tasks:
         from vllm.entrypoints.serve.render.api_router import (
             attach_router as attach_render_router,
@@ -402,18 +396,12 @@ async def init_app_state(
     )
 
     if "generate" in supported_tasks:
-        from vllm.entrypoints.openai.generate.api_router import init_generate_state
+        from vllm.entrypoints.generate.api_router import init_generate_state
 
         await init_generate_state(
             engine_client, state, args, request_logger, supported_tasks
         )
 
-        from vllm.entrypoints.openai.generative_scoring.api_router import (
-            init_generative_scoring_state,
-        )
-
-        await init_generative_scoring_state(engine_client, state, args, request_logger)
-
     if "transcription" in supported_tasks or "realtime" in supported_tasks:
         from vllm.entrypoints.speech_to_text.factories import init_speech_to_text_state
 

vllm/entrypoints/openai/chat_completion/protocol.py

@@ -740,19 +740,19 @@ class ChatCompletionRequest(OpenAIBaseModel):
         )
         # you can only use one kind of constraints for structured outputs
         if count > 1:
-            raise ValueError(
+            raise VLLMValidationError(
                 "You can only use one kind of constraints for structured "
-                "outputs ('json', 'regex' or 'choice')."
+                "outputs ('json', 'regex' or 'choice').",
             )
         # you can only either use structured outputs or tools, not both
-        if count > 1 and data.get("tool_choice", "none") not in (
+        if count > 0 and data.get("tool_choice", "none") not in (
             "none",
             "auto",
             "required",
         ):
-            raise ValueError(
+            raise VLLMValidationError(
                 "You can only either use constraints for structured outputs "
-                "or tools, not both."
+                "or tools, not both.",
             )
         return data
 
@@ -784,17 +784,21 @@ class ChatCompletionRequest(OpenAIBaseModel):
         if "tool_choice" in data and data["tool_choice"] is not None:
             # ensure that if "tool choice" is specified, tools are present
             if "tools" not in data or data["tools"] is None:
-                raise ValueError("When using `tool_choice`, `tools` must be set.")
+                raise VLLMValidationError(
+                    "When using `tool_choice`, `tools` must be set.",
+                    parameter="tool_choice",
+                )
 
             # make sure that tool choice is either a named tool
             # OR that it's set to "auto" or "required"
             if data["tool_choice"] not in ["auto", "required"] and not isinstance(
                 data["tool_choice"], dict
             ):
-                raise ValueError(
+                raise VLLMValidationError(
                     f"Invalid value for `tool_choice`: {data['tool_choice']}! "
                     'Only named tools, "none", "auto" or "required" '
-                    "are supported."
+                    "are supported.",
+                    parameter="tool_choice",
                 )
 
             # ensure that if "tool_choice" is specified as an object,
@@ -807,29 +811,33 @@ class ChatCompletionRequest(OpenAIBaseModel):
                 valid_tool = False
                 function = data["tool_choice"].get("function")
                 if not isinstance(function, dict):
-                    raise ValueError(
+                    raise VLLMValidationError(
                         f"Invalid value for `function`: `{function}` in "
-                        f"`tool_choice`! {correct_usage_message}"
+                        f"`tool_choice`! {correct_usage_message}",
+                        parameter="tool_choice.function",
                     )
                 if "name" not in function:
-                    raise ValueError(
+                    raise VLLMValidationError(
                         f"Expected field `name` in `function` in "
-                        f"`tool_choice`! {correct_usage_message}"
+                        f"`tool_choice`! {correct_usage_message}",
+                        parameter="tool_choice.function.name",
                     )
                 function_name = function["name"]
                 if not isinstance(function_name, str) or len(function_name) == 0:
-                    raise ValueError(
+                    raise VLLMValidationError(
                         f"Invalid `name` in `function`: `{function_name}`"
-                        f" in `tool_choice`! {correct_usage_message}"
+                        f" in `tool_choice`! {correct_usage_message}",
+                        parameter="tool_choice.function.name",
                     )
                 for tool in data["tools"]:
                     if tool["function"]["name"] == function_name:
                         valid_tool = True
                         break
                 if not valid_tool:
-                    raise ValueError(
+                    raise VLLMValidationError(
                         "The tool specified in `tool_choice` does not match any"
-                        " of the specified `tools`"
+                        " of the specified `tools`",
+                        parameter="tool_choice",
                     )
         return data
 
@@ -837,9 +845,9 @@ class ChatCompletionRequest(OpenAIBaseModel):
     @classmethod
     def check_generation_prompt(cls, data):
         if data.get("continue_final_message") and data.get("add_generation_prompt"):
-            raise ValueError(
+            raise VLLMValidationError(
                 "Cannot set both `continue_final_message` and "
-                "`add_generation_prompt` to True."
+                "`add_generation_prompt` to True.",
             )
         return data
 
@@ -849,8 +857,9 @@ class ChatCompletionRequest(OpenAIBaseModel):
         if data.get("cache_salt") is not None and (
             not isinstance(data["cache_salt"], str) or not data["cache_salt"]
         ):
-            raise ValueError(
-                "Parameter 'cache_salt' must be a non-empty string if provided."
+            raise VLLMValidationError(
+                "Parameter 'cache_salt' must be a non-empty string if provided.",
+                parameter="cache_salt",
             )
         return data
 

vllm/entrypoints/openai/completion/protocol.py

@@ -447,8 +447,9 @@ class CompletionRequest(OpenAIBaseModel):
         )
 
         if prompt_is_empty and embeds_is_empty:
-            raise ValueError(
-                "Either prompt or prompt_embeds must be provided and non-empty."
+            raise VLLMValidationError(
+                "Either prompt or prompt_embeds must be provided and non-empty.",
+                parameter="prompt",
             )
 
         return data
@@ -459,8 +460,9 @@ class CompletionRequest(OpenAIBaseModel):
         if data.get("cache_salt") is not None and (
             not isinstance(data["cache_salt"], str) or not data["cache_salt"]
         ):
-            raise ValueError(
-                "Parameter 'cache_salt' must be a non-empty string if provided."
+            raise VLLMValidationError(
+                "Parameter 'cache_salt' must be a non-empty string if provided.",
+                parameter="cache_salt",
             )
         return data
 

vllm/entrypoints/sagemaker/api_router.py

@@ -11,9 +11,9 @@ from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request
 from fastapi.responses import JSONResponse, Response
 
 from vllm.config import ModelConfig
+from vllm.entrypoints.generate.factories import get_generate_invocation_types
 from vllm.entrypoints.openai.engine.protocol import ErrorResponse
 from vllm.entrypoints.openai.engine.serving import OpenAIServing
-from vllm.entrypoints.openai.generate.factories import get_generate_invocation_types
 from vllm.entrypoints.openai.utils import validate_json_request
 from vllm.entrypoints.pooling.base.serving import PoolingServingBase
 from vllm.entrypoints.pooling.factories import get_pooling_invocation_types

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,
@@ -160,6 +163,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 +263,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 +359,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
     ],
     PlatformEnum.CPU: [
         Dynamic4bitLinearKernel,
+        ZentorchWNA16LinearKernel,
         CPUWNA16LinearKernel,
     ],
 }
@@ -1023,6 +1030,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/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/cutlass.py

@@ -312,7 +312,7 @@ class CutlassFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
     ) -> torch.Tensor:
         out_dtype = self.config.out_dtype
         if self.is_hopper:
-            return torch.ops.vllm.padded_cutlass(
+            return torch.ops.vllm.dynamic_padded_cutlass(
                 A,
                 B,
                 As,
@@ -320,14 +320,14 @@ class CutlassFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
                 list(self.weight_group_shape),
                 out_dtype,
             )
-        else:
-            return ops.cutlass_scaled_mm(
-                A,
-                B.T,
-                out_dtype=out_dtype,
-                scale_a=As,
-                scale_b=Bs.T,
-            )
+
+        return ops.cutlass_scaled_mm(
+            A,
+            B.T,
+            out_dtype=out_dtype,
+            scale_a=As,
+            scale_b=Bs.T,
+        )
 
 
 def cutlass_scaled_mm(
@@ -397,8 +397,56 @@ def _padded_cutlass_fake(
     )
 
 
+def _dynamic_padded_cutlass(
+    qx: torch.Tensor,
+    weight: torch.Tensor,
+    x_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    block_size: list[int],
+    output_dtype: torch.dtype,
+) -> torch.Tensor:
+    def run_padded(
+        qx: torch.Tensor,
+        weight: torch.Tensor,
+        x_scale: torch.Tensor,
+        weight_scale: torch.Tensor,
+    ) -> torch.Tensor:
+        return _padded_cutlass(
+            qx, weight, x_scale, weight_scale, block_size, output_dtype
+        )
+
+    def run_direct(
+        qx: torch.Tensor,
+        weight: torch.Tensor,
+        x_scale: torch.Tensor,
+        weight_scale: torch.Tensor,
+    ) -> torch.Tensor:
+        return cutlass_scaled_mm(
+            qx, weight, x_scale, weight_scale, block_size, output_dtype
+        )
+
+    if torch.compiler.is_compiling():
+        return torch.cond(
+            qx.shape[0] % 4 != 0,
+            run_padded,
+            run_direct,
+            (qx, weight, x_scale, weight_scale),
+        )
+
+    if qx.shape[0] % 4 != 0:
+        return run_padded(qx, weight, x_scale, weight_scale)
+
+    return run_direct(qx, weight, x_scale, weight_scale)
+
+
 direct_register_custom_op(
     "padded_cutlass",
     _padded_cutlass,
     fake_impl=_padded_cutlass_fake,
 )
+
+direct_register_custom_op(
+    "dynamic_padded_cutlass",
+    _dynamic_padded_cutlass,
+    fake_impl=_padded_cutlass_fake,
+)

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/attention/mla_attention.py

@@ -11,12 +11,12 @@ Sq      as Q sequence length
 Skv     as KV sequence length
 
 MLA has two possible ways of computing, a data-movement friendly approach and a
-compute friendly approach, we generally want to use the compute friendly
-approach for "prefill" (i.e. the ratio Sq / Skv is "small", is near 1)
-and the data-movement friendly approach for "decode" (i.e. the ratio
-Sq / Skv is "large").
+compute friendly approach. We generally want to use the compute friendly
+approach for "prefill" (i.e. the ratio Sq / Skv is relatively large, often near
+1) and the data-movement friendly approach for "decode" (i.e. the ratio
+Sq / Skv is small).
 
-NOTE what we deem small and large is currently determined by if its labelled
+NOTE what we deem small and large is currently determined by if it is labelled
 prefill or decode by the scheduler, but this is something we should probably
 tune.
 
@@ -96,7 +96,7 @@ NOTE: in the actual code,
 Runtime
 q_c      = h_t @ W_DQ
 q_nope   = (q_c @ W_UQ).view(-1, N, P)
-ql_nope  = einsum("snh,lnh->snl", q, W_UK)
+ql_nope  = einsum("snh,lnh->snl", q_nope, W_UK)
 q_pe     = RoPE(q_c @ W_QR).view(Sq, N, R)
 new_kv_c = h_t @ W_DKV
 new_k_pe = RoPE(h_t @ W_KR)
@@ -115,7 +115,7 @@ spda_o = scaled_dot_product_attention(
 )
 
 o = einsum("snl,lnv->snv", spda_o.reshape(-1, N, Lkv), W_UV)
-return o.view(-1, N * V) @ self.num_heads @ W_O
+return o.view(-1, N * V) @ W_O
 
 
 ## Chunked Prefill

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

@@ -164,9 +164,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()
 
@@ -388,9 +385,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()

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

@@ -89,9 +89,6 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular):
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
         return True
 
-    def supports_expert_map(self) -> bool:
-        return False
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         # Let PrepareAndFinalize::finalize() decide the impl.
         return TopKWeightAndReduceDelegate()

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

@@ -98,9 +98,6 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
     ) -> bool:
         return True
 
-    def supports_expert_map(self) -> bool:
-        return False
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 

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

@@ -207,9 +207,6 @@ class FlashInferExperts(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()
 

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

@@ -555,9 +555,6 @@ class NaiveBatchedExperts(mk.FusedMoEExpertsModular):
             "This method should not be called."
         )
 
-    def supports_expert_map(self) -> bool:
-        return False
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         # Let PrepareAndFinalize::finalize() decide the impl.
         return TopKWeightAndReduceDelegate()
@@ -799,9 +796,6 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
         return True
 
-    def supports_expert_map(self) -> bool:
-        return False
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         # Let PrepareAndFinalize::finalize() decide the impl.
         return TopKWeightAndReduceDelegate()

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

@@ -156,9 +156,6 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
     ) -> bool:
         return True
 
-    def supports_expert_map(self) -> bool:
-        return True
-
     @staticmethod
     def _supports_current_device() -> bool:
         platform = current_platform

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

@@ -608,9 +608,6 @@ class BaseOAITritonExperts(mk.FusedMoEExpertsModular):
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
         return True
 
-    def supports_expert_map(self) -> bool:
-        return True
-
     def moe_problem_size(
         self,
         a1: torch.Tensor,
@@ -1036,9 +1033,6 @@ class OAITritonMxfp4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
     ) -> bool:
         return True
 
-    def supports_expert_map(self) -> bool:
-        return True
-
     @property
     def expects_unquantized_inputs(self) -> bool:
         return True

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

@@ -686,9 +686,6 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
 class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
     """Marlin-based fused MoE expert implementation."""
 
-    def supports_expert_map(self) -> bool:
-        return True
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 
@@ -920,9 +917,6 @@ class BatchedMarlinExperts(MarlinExpertsBase):
             is_k_full=is_k_full,
         )
 
-    def supports_expert_map(self) -> bool:
-        return True
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceDelegate()
 

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

@@ -441,9 +441,6 @@ class AiterExperts(mk.FusedMoEExpertsModular):
             or moe_parallel_config.use_fi_nvl_one_sided_kernels
         )
 
-    def supports_expert_map(self):
-        return True
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 

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

@@ -125,9 +125,6 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
     def _supports_batch_invariance():
         return True
 
-    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/trtllm_bf16_moe.py

@@ -99,9 +99,6 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
     ) -> bool:
         return True
 
-    def supports_expert_map(self) -> bool:
-        return False
-
     @property
     def expects_unquantized_inputs(self) -> bool:
         return True

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

@@ -88,9 +88,6 @@ class TrtLlmFp8ExpertsBase:
             or moe_parallel_config.use_ag_rs_all2all_kernels
         ) and not moe_parallel_config.enable_eplb
 
-    def supports_expert_map(self) -> bool:
-        return False
-
 
 class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
     """