From 106aa92f04c0d1c3c37947fd9d4921530562a961 Mon Sep 17 00:00:00 2001
From: bnellnm <49004751+bnellnm@users.noreply.github.com>
Date: Fri, 29 May 2026 17:19:31 -0400
Subject: [PATCH 01/35] [MoE Refactor] Migrate MoeWNA16Method quantization to
 MK oracle (#42647)

Signed-off-by: Bill Nell <bnell@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
---
 .../layers/fused_moe/experts/triton_moe.py    |  50 +++--
 .../layers/fused_moe/fused_moe.py             |   7 +-
 .../layers/fused_moe/oracle/int_wna16.py      |  97 +++++++++-
 .../layers/quantization/auto_gptq.py          |  64 +++----
 .../layers/quantization/awq_marlin.py         |   2 +
 .../compressed_tensors_moe_wna16_marlin.py    |  17 +-
 .../layers/quantization/moe_wna16.py          | 176 +++++++++++++++---
 .../layers/quantization/utils/gptq_utils.py   |  18 +-
 8 files changed, 333 insertions(+), 98 deletions(-)

diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
index 45b95b34102..ddcef519da0 100644
--- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
@@ -43,7 +43,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8Static128BlockSym,
     kFp8StaticChannelSym,
     kFp8StaticTensorSym,
+    kInt4Static,
+    kInt4Static32,
     kInt8DynamicTokenSym,
+    kInt8Static,
     kInt8StaticChannelSym,
 )
 from vllm.platforms import current_platform
@@ -459,40 +462,45 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
 class TritonWNA16Experts(TritonExperts):
     @staticmethod
     def _supports_current_device() -> bool:
-        raise NotImplementedError(
-            "TritonWNA16Experts is not yet used by an Oracle. "
-            "This method should not be called."
-        )
+        return current_platform.is_cuda_alike() or current_platform.is_xpu()
 
     @staticmethod
     def _supports_no_act_and_mul() -> bool:
-        raise NotImplementedError(
-            "TritonWNA16Experts is not yet used by an Oracle. "
-            "This method should not be called."
-        )
+        return True
 
     @staticmethod
     def _supports_quant_scheme(
         weight_key: QuantKey | None,
         activation_key: QuantKey | None,
     ) -> bool:
-        raise NotImplementedError(
-            "TritonWNA16Experts is not yet used by an Oracle. "
-            "This method should not be called."
-        )
+        SUPPORTED_W = [
+            kInt4Static,
+            kInt8Static,
+            kInt4Static32,
+            # other group sizes?
+        ]
+        return weight_key in SUPPORTED_W
 
     @staticmethod
     def _supports_activation(activation: MoEActivation) -> bool:
-        raise NotImplementedError(
-            "TritonWNA16Experts is not yet used by an Oracle. "
-            "This method should not be called."
-        )
+        return activation in [
+            MoEActivation.SILU,
+            MoEActivation.GELU,
+            MoEActivation.GELU_TANH,
+            MoEActivation.SWIGLUOAI,
+            MoEActivation.SWIGLUSTEP,
+            MoEActivation.SILU_NO_MUL,
+            MoEActivation.GELU_NO_MUL,
+            MoEActivation.GELU_TANH_NO_MUL,
+            MoEActivation.RELU2_NO_MUL,
+        ]
 
     @staticmethod
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
-        raise NotImplementedError(
-            "TritonWNA16Experts is not yet used by an Oracle. "
-            "This method should not be called."
+        # Why?
+        return not (
+            moe_parallel_config.use_fi_nvl_two_sided_kernels
+            or moe_parallel_config.use_fi_nvl_one_sided_kernels
         )
 
     def apply(
@@ -515,7 +523,9 @@ class TritonWNA16Experts(TritonExperts):
     ):
         # Check constraints.
         if self.quant_config.use_int4_w4a16:
-            assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"
+            assert hidden_states.size(-1) // 2 == w1.size(2), (
+                f"Hidden size mismatch {hidden_states.size(-1) // 2} == {w1.size(2)}"
+            )
         else:
             assert hidden_states.size(-1) == w1.size(2), (
                 f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}"
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
index 49957c8f5e3..d7b30ef7633 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -30,6 +30,7 @@ from vllm.model_executor.layers.fused_moe.utils import (
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import direct_register_custom_op
 
 logger = init_logger(__name__)
@@ -1243,7 +1244,11 @@ def get_default_config(
         bit = 4 if dtype == "int4_w4a16" else 8
         use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk, block_shape[1], E, bit)
         if use_moe_wna16_cuda:
-            config = {"BLOCK_SIZE_M": min(16, M), "SPLIT_K": 1}
+            config = {
+                "BLOCK_SIZE_M": min(16, next_power_of_2(M)),
+                "GROUP_SIZE_M": 1,
+                "SPLIT_K": 1,
+            }
         elif M <= 20:
             config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1, "SPLIT_K": 1}
         elif M <= 40:
diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
index 9de7a6ba119..dd96921d707 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
@@ -23,6 +23,9 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
     MarlinExperts,
     MarlinExpertsBase,
 )
+from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
+    TritonWNA16Experts,
+)
 from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import (
     TrtLlmMxint4ExpertsMonolithic,
 )
@@ -31,6 +34,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     marlin_act_int8_process_scales,
     marlin_moe_permute_scales,
     marlin_permute_bias,
+    marlin_zero_points,
     moe_awq_to_marlin_zero_points,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -45,6 +49,7 @@ class WNA16MoEBackend(Enum):
     MARLIN = "MARLIN"
     BATCHED_MARLIN = "BATCHED_MARLIN"
     FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
+    TRITON = "TRITON"
     XPU = "XPU"
 
 
@@ -58,6 +63,8 @@ def backend_to_kernel_cls(
         return [BatchedMarlinExperts]
     elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
         return [TrtLlmMxint4ExpertsMonolithic]
+    elif backend == WNA16MoEBackend.TRITON:
+        return [TritonWNA16Experts]
     elif backend == WNA16MoEBackend.XPU:
         from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
             XPUExpertsWNA16,
@@ -68,24 +75,38 @@ def backend_to_kernel_cls(
         raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")
 
 
-def _get_priority_backends() -> list[WNA16MoEBackend]:
+def _get_priority_backends(
+    may_have_zp: bool, may_have_bias: bool
+) -> list[WNA16MoEBackend]:
     """
     Get available backends in priority order based on platform and config.
     """
     if current_platform.is_xpu():
         return [WNA16MoEBackend.XPU]
 
-    _AVAILABLE_BACKENDS = [
-        WNA16MoEBackend.FLASHINFER_TRTLLM,
+    _AVAILABLE_BACKENDS = []
+
+    if not may_have_zp and not may_have_bias:
+        _AVAILABLE_BACKENDS.append(WNA16MoEBackend.FLASHINFER_TRTLLM)
+
+    # Marlin supports ZP and bias
+    _AVAILABLE_BACKENDS += [
         WNA16MoEBackend.MARLIN,
         WNA16MoEBackend.BATCHED_MARLIN,
     ]
     return _AVAILABLE_BACKENDS
 
+    if not may_have_bias:
+        _AVAILABLE_BACKENDS.append(WNA16MoEBackend.TRITON)
+
+    return _AVAILABLE_BACKENDS
+
 
 def select_wna16_moe_backend(
     config: FusedMoEConfig,
     weight_key: QuantKey,
+    may_have_zp: bool,
+    may_have_bias: bool,
 ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]:
     """Select the WNA16 MoE backend.
 
@@ -136,7 +157,7 @@ def select_wna16_moe_backend(
         raise ValueError(_make_log_unsupported(backend, reason))
 
     # Select kernels in order of backend.
-    AVAILABLE_BACKENDS = _get_priority_backends()
+    AVAILABLE_BACKENDS = _get_priority_backends(may_have_zp, may_have_bias)
 
     for backend in AVAILABLE_BACKENDS:
         activation_key = None  # always BF16 activation for WNA16 MoE
@@ -218,6 +239,7 @@ def make_wna16_moe_kernel(
     assert experts_cls in (
         MarlinExperts,
         BatchedMarlinExperts,
+        TritonWNA16Experts,
         TrtLlmMxint4ExpertsMonolithic,
         XPUExpertsWNA16,
     )
@@ -234,6 +256,7 @@ def make_wna16_moe_kernel(
     assert prepare_finalize is not None
 
     logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")
+    logger.info_once("Using %s", experts_cls.__name__, scope="local")
 
     extra_args: dict[str, Any] = {}
     if issubclass(experts_cls, MarlinExpertsBase):
@@ -403,6 +426,8 @@ def _process_weights_marlin(
     w2_input_global_scale: torch.Tensor | None = None
     w13_bias_out: torch.Tensor | None = None
     w2_bias_out: torch.Tensor | None = None
+    w13_qzeros_out: torch.Tensor | None = None
+    w2_qzeros_out: torch.Tensor | None = None
 
     # --- FP8 weight / scale adjustment ---
     if input_dtype == torch.float8_e4m3fn:
@@ -502,6 +527,24 @@ def _process_weights_marlin(
     if w2_bias is not None:
         w2_bias_out = marlin_permute_bias(w2_bias)
 
+    if w13_qzeros is not None:
+        w13_qzeros_out = marlin_zero_points(
+            w13_qzeros,
+            size_k=layer.intermediate_size_per_partition,
+            size_n=w13_qzeros.shape[2],
+            num_bits=num_bits,
+            is_a_8bit=is_a_8bit,
+        )
+
+    if w2_qzeros is not None:
+        w2_qzeros_out = marlin_zero_points(
+            w2_qzeros,
+            size_k=w2_qzeros.shape[1] * group_size_or_pack_factor,
+            size_n=w2_qzeros.shape[2],
+            num_bits=num_bits,
+            is_a_8bit=is_a_8bit,
+        )
+
     return (
         marlin_w13_qweight,
         marlin_w2_qweight,
@@ -511,8 +554,8 @@ def _process_weights_marlin(
         w2_g_idx,
         w13_g_idx_sort_indices,
         w2_g_idx_sort_indices,
-        w13_qzeros,
-        w2_qzeros,
+        w13_qzeros_out,
+        w2_qzeros_out,
         w13_input_global_scale,
         w2_input_global_scale,
         w13_bias_out,
@@ -780,6 +823,9 @@ def convert_to_wna16_moe_kernel_format(
         from vllm.model_executor.layers.quantization.awq_marlin import (
             AWQMarlinConfig,
         )
+        from vllm.model_executor.layers.quantization.moe_wna16 import (
+            MoeWNA16Config,
+        )
 
         if isinstance(quant_config, AWQMarlinConfig):
             if w13_qzeros is None or w2_qzeros is None:
@@ -814,11 +860,17 @@ def convert_to_wna16_moe_kernel_format(
             pack_factor = 32 // quant_config.num_bits
             group_size = quant_config.group_size
             actorder = quant_config.actorder
+        elif isinstance(quant_config, MoeWNA16Config):
+            num_bits = quant_config.weight_bits
+            pack_factor = quant_config.bit8_pack_factor
+            group_size = quant_config.group_size
+            actorder = None
         else:
             raise TypeError(
                 "Marlin WNA16 MoE backend requires AutoGPTQConfig, AWQMarlinConfig or "
                 f"QuantizationArgs, got {type(quant_config).__name__}."
             )
+
         if w13_g_idx is None or w2_g_idx is None:
             raise ValueError("GPTQ Marlin MoE requires g_idx tensors.")
         return _process_weights_marlin(
@@ -850,6 +902,39 @@ def convert_to_wna16_moe_kernel_format(
             w13_bias,
             w2_bias,
         )
+    elif backend == WNA16MoEBackend.TRITON:
+        # Convert from int32 to uint8 format for Triton kernel.
+        # This changes the shape from (E, N, K // 8) to (E, N, K // 2) for int4,
+        # which matches what the Triton kernel expects.
+        w13_uint8 = w13.view(torch.uint8)
+        w2_uint8 = w2.view(torch.uint8)
+        return (
+            w13_uint8,
+            w2_uint8,
+            w13_scale,
+            w2_scale,
+            None,
+            None,
+            None,
+            None,
+            w13_qzeros,
+            w2_qzeros,
+            None,
+            None,
+            w13_bias,
+            w2_bias,
+        )
+    elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
+        return _process_weights_flashinfer(
+            w13,
+            w2,
+            w13_scale,
+            w2_scale,
+            w13_g_idx,
+            w2_g_idx,
+            w13_bias,
+            w2_bias,
+        )
     elif backend == WNA16MoEBackend.XPU:
         assert quant_config is not None
         (
diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py
index 1821fd5c7f7..cb2dad3d524 100644
--- a/vllm/model_executor/layers/quantization/auto_gptq.py
+++ b/vllm/model_executor/layers/quantization/auto_gptq.py
@@ -485,6 +485,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
         self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
             moe,
             weight_key,
+            may_have_zp=True,
+            may_have_bias=True,
         )
 
     def create_weights(
@@ -644,12 +646,22 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
         layer.workspace = marlin_make_workspace_new(device, 4)
 
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
+        def replace_or_register(name: str, val: torch.Tensor | None):
+            if val is None:
+                return
+
+            if hasattr(layer, name):
+                replace_parameter(layer, name, val)
+            else:
+                layer.register_parameter(
+                    name, torch.nn.Parameter(val, requires_grad=False)
+                )
+
         is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
 
-        if is_a_8bit:
-            assert self.quant_config.quant_type.size_bits == 8, (
-                "W8A8-INT8 is not supported by marlin kernel."
-            )
+        assert not is_a_8bit or self.quant_config.quant_type.size_bits == 8, (
+            "W8A8-INT8 is not supported by marlin kernel."
+        )
 
         (
             w13,
@@ -660,8 +672,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
             w2_g_idx,
             w13_g_idx_sort_indices,
             w2_g_idx_sort_indices,
-            _w13_qzeros,
-            _w2_qzeros,
+            w13_qzeros,
+            w2_qzeros,
             w13_input_global_scale,
             w2_input_global_scale,
             w13_bias,
@@ -679,6 +691,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
             w2_g_idx=layer.w2_g_idx,
             w13_bias=getattr(layer, "w13_bias", None),
             w2_bias=getattr(layer, "w2_bias", None),
+            w13_qzeros=getattr(layer, "w13_qzeros", None),
+            w2_qzeros=getattr(layer, "w2_qzeros", None),
         )
 
         replace_parameter(layer, "w13_qweight", w13)
@@ -689,38 +703,12 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
         replace_parameter(layer, "w2_g_idx", w2_g_idx)
         replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
         replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
-        if w13_input_global_scale is not None:
-            if hasattr(layer, "w13_input_global_scale"):
-                replace_parameter(
-                    layer, "w13_input_global_scale", w13_input_global_scale
-                )
-            else:
-                layer.register_parameter(
-                    "w13_input_global_scale",
-                    torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
-                )
-        if w2_input_global_scale is not None:
-            if hasattr(layer, "w2_input_global_scale"):
-                replace_parameter(layer, "w2_input_global_scale", w2_input_global_scale)
-            else:
-                layer.register_parameter(
-                    "w2_input_global_scale",
-                    torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
-                )
-        if w13_bias is not None:
-            if hasattr(layer, "w13_bias"):
-                replace_parameter(layer, "w13_bias", w13_bias)
-            else:
-                layer.register_parameter(
-                    "w13_bias", torch.nn.Parameter(w13_bias, requires_grad=False)
-                )
-        if w2_bias is not None:
-            if hasattr(layer, "w2_bias"):
-                replace_parameter(layer, "w2_bias", w2_bias)
-            else:
-                layer.register_parameter(
-                    "w2_bias", torch.nn.Parameter(w2_bias, requires_grad=False)
-                )
+        replace_or_register("w13_input_global_scale", w13_input_global_scale)
+        replace_or_register("w2_input_global_scale", w2_input_global_scale)
+        replace_or_register("w13_bias", w13_bias)
+        replace_or_register("w2_bias", w2_bias)
+        replace_or_register("w13_qzeros", w13_qzeros)
+        replace_or_register("w2_qzeros", w2_qzeros)
 
         self._setup_kernel(layer)
 
diff --git a/vllm/model_executor/layers/quantization/awq_marlin.py b/vllm/model_executor/layers/quantization/awq_marlin.py
index 81c0fcb331e..48fbe6335a3 100644
--- a/vllm/model_executor/layers/quantization/awq_marlin.py
+++ b/vllm/model_executor/layers/quantization/awq_marlin.py
@@ -524,6 +524,8 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
         self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
             moe,
             kInt4Static,
+            may_have_zp=self.quant_config.zero_point,
+            may_have_bias=True,
         )
 
     def create_weights(
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
index 2d629d73edd..4b591fabcc9 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
@@ -88,7 +88,13 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
         self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
             config=self.moe,
             weight_key=weight_key,
+            may_have_zp=False,
+            may_have_bias=False,
         )
+        self.is_marlin = self.wna16_backend in [
+            WNA16MoEBackend.MARLIN,
+            WNA16MoEBackend.BATCHED_MARLIN,
+        ]
 
     def get_weight_shape(
         self,
@@ -114,7 +120,6 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
                 "num_groups_w2 must be provided for weight scales"
             )
         w13_num_shards = 2 if self.moe.is_act_and_mul else 1
-        is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM
         shape_map = {
             "w13_weight": {
                 "Flashinfer": (
@@ -157,7 +162,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
                 "Marlin": (num_experts, num_groups_w2, hidden_size),
             },
         }
-        backend_key = "Flashinfer" if is_flashinfer else "Marlin"
+        backend_key = "Marlin" if self.is_marlin else "Flashinfer"
         return shape_map[weight_name][backend_key]
 
     def create_weights(
@@ -174,9 +179,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
         # Will transpose the loaded weight along the
         # intermediate and hidden dim sizes. Will
         # shard for TP along the transposed dims
-        is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM
         extra_weight_attrs.update(
-            {"is_transposed": is_transposed, "quant_method": self.strategy}
+            {"is_transposed": self.is_marlin, "quant_method": self.strategy}
         )
 
         w13_weight = torch.nn.Parameter(
@@ -324,7 +328,6 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         # Process weights using the shared oracle infrastructure
-        is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM
         (
             w13_qweight,
             w2_qweight,
@@ -360,7 +363,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
         replace_parameter(layer, "w2_weight_scale", w2_scales)
 
         # Marlin-specific parameters (not needed for Flashinfer)
-        if not is_flashinfer:
+        if self.is_marlin:
             replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed)
             replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed)
             replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
@@ -392,7 +395,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
 
         # Add Marlin-specific arguments
         marlin_args: dict[str, Any] = {}
-        if not is_flashinfer:
+        if self.is_marlin:
             marlin_args = {
                 "w13_g_idx": layer.w13_weight_g_idx,
                 "w2_g_idx": layer.w2_weight_g_idx,
diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py
index 471febab044..fd99d520c64 100644
--- a/vllm/model_executor/layers/quantization/moe_wna16.py
+++ b/vllm/model_executor/layers/quantization/moe_wna16.py
@@ -13,11 +13,15 @@ from vllm.model_executor.layers.fused_moe import (
     RoutedExperts,
     SharedExperts,
 )
-from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
-    int4_w4a16_moe_quant_config,
-    int8_w8a16_moe_quant_config,
+)
+from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
+    WNA16MoEBackend,
+    convert_to_wna16_moe_kernel_format,
+    make_wna16_moe_kernel,
+    make_wna16_moe_quant_config,
+    select_wna16_moe_backend,
 )
 from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
     UnquantizedFusedMoEMethod,
@@ -31,7 +35,15 @@ from vllm.model_executor.layers.quantization.base_config import (
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     check_marlin_supports_layer,
 )
-from vllm.model_executor.utils import set_weight_attrs
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    INT4_DTYPE,
+    INT8_DTYPE,
+    QuantKey,
+    kInt4Static32GroupScale,
+    kInt4StaticGroupScale,
+    kInt8StaticGroupScale,
+)
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
 from vllm.platforms import current_platform
 
 
@@ -216,6 +228,33 @@ class MoeWNA16Method(FusedMoEMethodBase):
         super().__init__(moe)
         self.quant_config = quant_config
 
+        num_bits = self.quant_config.weight_bits
+        group_size = self.quant_config.group_size
+
+        if num_bits == 4:
+            quant_type = INT4_DTYPE
+            if group_size == 32:
+                scale = kInt4Static32GroupScale
+            else:
+                scale = kInt4StaticGroupScale
+        elif num_bits == 8:
+            assert group_size == -1
+            quant_type = INT8_DTYPE
+            scale = kInt8StaticGroupScale
+        else:
+            raise ValueError("MoeWNA16Method only supports int4 and int8 now.")
+
+        weight_key = QuantKey(quant_type, scale)
+
+        # Select WNA16 MoE backend via oracle.
+        # handle ZP?
+        self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
+            config=self.moe,
+            weight_key=weight_key,
+            may_have_zp=self.quant_config.has_zp,
+            may_have_bias=False,
+        )
+
     def create_weights(
         self,
         layer: RoutedExperts,
@@ -336,24 +375,89 @@ class MoeWNA16Method(FusedMoEMethodBase):
                 layer.register_parameter(key, param)
                 set_weight_attrs(param, extra_weight_attrs)
 
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        has_zp = self.quant_config.has_zp
+        (
+            w13_qweight,
+            w2_qweight,
+            w13_scales,
+            w2_scales,
+            w13_g_idx_processed,
+            w2_g_idx_processed,
+            w13_g_idx_sort_indices,
+            w2_g_idx_sort_indices,
+            w13_qzeros,
+            w2_qzeros,
+            w13_input_global_scale,
+            w2_input_global_scale,
+            _,  # w13_bias
+            _,  # w2_bias
+        ) = convert_to_wna16_moe_kernel_format(
+            backend=self.wna16_backend,
+            layer=layer,
+            quant_config=self.quant_config,
+            input_dtype=None,
+            w13=layer.w13_qweight,
+            w2=layer.w2_qweight,
+            w13_scale=layer.w13_scales,
+            w2_scale=layer.w2_scales,
+            w13_g_idx=getattr(layer, "w13_g_idx", None),
+            w2_g_idx=getattr(layer, "w2_g_idx", None),
+            w13_qzeros=layer.w13_qzeros if has_zp else None,
+            w2_qzeros=layer.w2_qzeros if has_zp else None,
+        )
+
+        # Replace common parameters
+        replace_parameter(layer, "w13_qweight", w13_qweight)
+        replace_parameter(layer, "w2_qweight", w2_qweight)
+        replace_parameter(layer, "w13_scales", w13_scales)
+        replace_parameter(layer, "w2_scales", w2_scales)
+
+        if has_zp:
+            assert w13_qzeros is not None and w2_qzeros is not None
+            replace_parameter(layer, "w13_qzeros", w13_qzeros)
+            replace_parameter(layer, "w2_qzeros", w2_qzeros)
+
+        # Marlin-specific parameters (not needed for Flashinfer)
+        if self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM:
+            replace_parameter(layer, "w13_g_idx", w13_g_idx_processed)
+            replace_parameter(layer, "w2_g_idx", w2_g_idx_processed)
+            replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
+            replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
+
+            # Register input global scales if present
+            if w13_input_global_scale is not None:
+                layer.register_parameter(
+                    "w13_input_global_scale",
+                    torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
+                )
+            if w2_input_global_scale is not None:
+                layer.register_parameter(
+                    "w2_input_global_scale",
+                    torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
+                )
+
+        assert self.experts_cls is not None
+        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
+        assert self.moe_quant_config is not None
+        self.moe_kernel = make_wna16_moe_kernel(
+            moe_quant_config=self.moe_quant_config,
+            moe_config=self.moe,
+            experts_cls=self.experts_cls,
+            routing_tables=layer._expert_routing_tables(),
+        )
+
     def get_fused_moe_quant_config(
         self, layer: RoutedExperts
     ) -> FusedMoEQuantConfig | None:
-        weight_bits = self.quant_config.weight_bits
         has_zp = self.quant_config.has_zp
-        assert weight_bits == 4 or weight_bits == 8
-        config_builder = (
-            int4_w4a16_moe_quant_config
-            if weight_bits == 4
-            else int8_w8a16_moe_quant_config
-        )
-
-        return config_builder(
+        return make_wna16_moe_quant_config(
             w1_scale=layer.w13_scales,
             w2_scale=layer.w2_scales,
             w1_zp=layer.w13_qzeros if has_zp else None,
             w2_zp=layer.w2_qzeros if has_zp else None,
-            block_shape=[0, layer.group_size],
+            group_size=layer.group_size,
+            num_bits=self.quant_config.weight_bits,
         )
 
     def apply(
@@ -365,22 +469,44 @@ class MoeWNA16Method(FusedMoEMethodBase):
         shared_experts: SharedExperts | None,
         shared_experts_input: torch.Tensor | None,
     ) -> torch.Tensor:
-        from vllm.model_executor.layers.fused_moe import fused_experts
-
-        assert layer.activation == MoEActivation.SILU, (
-            f"Only SiLU activation is supported, not {layer.activation}."
-        )
-
-        return fused_experts(
+        assert not self.is_monolithic
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply(
             x,
             layer.w13_qweight,
             layer.w2_qweight,
-            topk_weights=topk_weights,
-            topk_ids=topk_ids,
-            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            topk_weights,
+            topk_ids,
+            activation=layer.activation,
             global_num_experts=layer.global_num_experts,
             expert_map=layer.expert_map,
-            quant_config=self.moe_quant_config,
+            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            shared_experts=shared_experts,
+            shared_experts_input=shared_experts_input,
+        )
+
+    def apply_monolithic(
+        self,
+        layer: RoutedExperts,
+        x: torch.Tensor,
+        router_logits: torch.Tensor,
+        input_ids: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        assert self.is_monolithic
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply_monolithic(
+            x,
+            layer.w13_weight,
+            layer.w2_weight,
+            router_logits,
+            activation=layer.activation,
+            global_num_experts=layer.global_num_experts,
+            expert_map=layer.expert_map,
+            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            num_expert_group=layer.num_expert_group,
+            topk_group=layer.topk_group,
+            e_score_correction_bias=layer.e_score_correction_bias,
+            routed_scaling_factor=layer.routed_scaling_factor,
         )
 
     @staticmethod
diff --git a/vllm/model_executor/layers/quantization/utils/gptq_utils.py b/vllm/model_executor/layers/quantization/utils/gptq_utils.py
index 691d80b0b74..c73c42f4ff0 100644
--- a/vllm/model_executor/layers/quantization/utils/gptq_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/gptq_utils.py
@@ -3,7 +3,7 @@
 from collections.abc import Mapping
 from copy import deepcopy
 from types import MappingProxyType
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import regex as re
 import torch
@@ -69,6 +69,20 @@ def get_dynamic_override(
     return default_value
 
 
+def flatten_list(lst: list[Any]) -> list[Any]:
+    output = []
+
+    def _flatten(lst: list[Any]):
+        for i in lst:
+            if isinstance(i, list):
+                _flatten(i)
+            else:
+                output.append(i)
+
+    _flatten(lst)
+    return output
+
+
 def is_layer_gptq_quantized(
     prefix: str,
     quantized_layers: list[str],
@@ -83,6 +97,8 @@ def is_layer_gptq_quantized(
 
     proj_name = prefix.split(".")[-1]
 
+    quantized_layers = flatten_list(quantized_layers)
+
     # Fused layers like gate_up_proj or qkv_proj will not be fused
     # in the safetensors checkpoint. So, we convert the name
     # from the fused version to unfused + check to make sure that

From 7b98f498cdf0bf9cf0ecc37b5c0c994cb94513c3 Mon Sep 17 00:00:00 2001
From: bnellnm <49004751+bnellnm@users.noreply.github.com>
Date: Fri, 29 May 2026 17:26:56 -0400
Subject: [PATCH 02/35] [MoE Refactor] Remove supports_expert_map (#43108)

Signed-off-by: Bill Nell <bnell@redhat.com>
---
 .../moe/modular_kernel_tools/common.py        | 15 +++++++++-----
 .../moe/modular_kernel_tools/mk_objects.py    | 15 --------------
 .../moe/test_modular_kernel_combinations.py   |  2 +-
 .../fused_moe/experts/aiter_mxfp4_w4a8_moe.py |  3 ---
 .../experts/batched_deep_gemm_moe.py          |  3 ---
 .../layers/fused_moe/experts/cpu_moe.py       |  6 ------
 .../layers/fused_moe/experts/cutlass_moe.py   | 18 ++---------------
 .../layers/fused_moe/experts/deep_gemm_moe.py |  6 ------
 .../layers/fused_moe/experts/fallback.py      | 10 ----------
 .../experts/flashinfer_cutedsl_batched_moe.py |  3 ---
 .../experts/flashinfer_cutedsl_moe.py         |  3 ---
 .../experts/flashinfer_cutlass_moe.py         |  3 ---
 .../fused_moe/experts/fused_batched_moe.py    |  6 ------
 .../fused_moe/experts/fused_humming_moe.py    |  3 ---
 .../experts/gpt_oss_triton_kernels_moe.py     |  6 ------
 .../layers/fused_moe/experts/marlin_moe.py    |  6 ------
 .../fused_moe/experts/rocm_aiter_moe.py       |  3 ---
 .../layers/fused_moe/experts/triton_moe.py    |  3 ---
 .../fused_moe/experts/trtllm_bf16_moe.py      |  3 ---
 .../fused_moe/experts/trtllm_fp8_moe.py       |  3 ---
 .../fused_moe/experts/trtllm_mxfp4_moe.py     |  6 ------
 .../fused_moe/experts/trtllm_nvfp4_moe.py     |  3 ---
 .../layers/fused_moe/experts/xpu_moe.py       |  3 ---
 .../fused_moe/fused_moe_modular_method.py     |  7 +------
 .../layers/fused_moe/modular_kernel.py        | 13 ------------
 .../fused_moe/prepare_finalize/naive_dp_ep.py | 20 +++++++++++--------
 .../compressed_tensors_moe_w8a8_fp8.py        |  2 --
 27 files changed, 26 insertions(+), 148 deletions(-)

diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py
index ea52a2d3398..fdd00cfa27a 100644
--- a/tests/kernels/moe/modular_kernel_tools/common.py
+++ b/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
             )
diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py
index 7c3bde2eafa..78ee8084d90 100644
--- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py
+++ b/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
diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py
index c7295f3ed6e..0c0e1d61f90 100644
--- a/tests/kernels/moe/test_modular_kernel_combinations.py
+++ b/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(
diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
index 3906a7e057c..cc2adc31fcd 100644
--- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
+++ b/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
diff --git a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py
index 7bd383b9cda..c8611217a18 100644
--- a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py
+++ b/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
diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
index 54b264ef772..84740fc0570 100644
--- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+++ b/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,
diff --git a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
index 28a7d283b4b..feb49d260e1 100644
--- a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+++ b/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()
diff --git a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py
index e3e15e31618..3b354dd3ef1 100644
--- a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py
+++ b/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()
 
diff --git a/vllm/model_executor/layers/fused_moe/experts/fallback.py b/vllm/model_executor/layers/fused_moe/experts/fallback.py
index 40741d52af5..639b2bf2668 100644
--- a/vllm/model_executor/layers/fused_moe/experts/fallback.py
+++ b/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()
diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py
index 5eaaf46739f..253d1dae711 100644
--- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py
+++ b/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()
diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py
index 2310982792f..b512d51c135 100644
--- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py
+++ b/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()
 
diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
index b891583e3ef..fd9446c2a22 100644
--- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
+++ b/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()
 
diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py
index 0e31331e726..1f5724ac39c 100644
--- a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py
+++ b/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()
diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
index 8874228a142..53623f13254 100644
--- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+++ b/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
diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
index 98265abf7c8..03bf925fbd9 100644
--- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+++ b/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
diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
index 8bb9e5bdc06..1d0cf91d427 100644
--- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+++ b/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()
 
diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
index b272a458b17..8415ac02784 100644
--- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
+++ b/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()
 
diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
index ddcef519da0..cf2c2cff6a8 100644
--- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
@@ -128,9 +128,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()
 
diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
index 02b7450a5c9..592a1513d75 100644
--- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+++ b/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
diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
index b98b84cdc62..43126195205 100644
--- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+++ b/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):
     """
diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
index 1e2fff8eb66..43f800343c7 100644
--- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
@@ -113,9 +113,6 @@ class TrtLlmMxfp4ExpertsBase:
     def activation_format() -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
 
-    def supports_expert_map(self) -> bool:
-        return False
-
     @property
     def expects_unquantized_inputs(self) -> bool:
         return False
@@ -248,9 +245,6 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
         # routing is done externally, so accept any routing method.
         return True
 
-    def supports_expert_map(self) -> bool:
-        return True
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 
diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
index e4f292b7b1e..5ee023aa27c 100644
--- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
@@ -179,9 +179,6 @@ class TrtLlmNvFp4ExpertsBase:
             300000, _calc_max_supported_tokens(self.topk, self.moe_config.num_experts)
         )
 
-    def supports_expert_map(self) -> bool:
-        return False
-
 
 class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModular):
     """
diff --git a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
index 8cbf0a6ce02..82969dd8e25 100644
--- a/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
@@ -107,9 +107,6 @@ class XPUExperts(mk.FusedMoEExpertsModular):
         ]
         return (weight_key, activation_key) in SUPPORTED_W_A
 
-    def supports_expert_map(self) -> bool:
-        return True
-
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
index 3ebb63b0057..dd21ff58fc3 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
@@ -34,11 +34,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
         super().__init__(moe_kernel.moe_config)
         self.moe_quant_config = old_quant_method.moe_quant_config
         self.moe_kernel = moe_kernel
-        self.disable_expert_map = getattr(
-            old_quant_method,
-            "disable_expert_map",
-            not self.moe_kernel.supports_expert_map(),
-        )
         self.old_quant_method = old_quant_method
         logger.debug("Swapping out %s", self.old_quant_method.__class__.__name__)
 
@@ -103,7 +98,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
             activation=layer.activation,
             global_num_experts=layer.global_num_experts,
             apply_router_weight_on_input=layer.apply_router_weight_on_input,
-            expert_map=None if self.disable_expert_map else layer.expert_map,
+            expert_map=layer.expert_map,
             shared_experts=shared_experts,
             shared_experts_input=shared_experts_input,
         )
diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py
index 6fbc1bffaac..9c3ecee9f9b 100644
--- a/vllm/model_executor/layers/fused_moe/modular_kernel.py
+++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py
@@ -751,13 +751,6 @@ class FusedMoEExperts(ABC):
         """
         return False
 
-    @abstractmethod
-    def supports_expert_map(self) -> bool:
-        """
-        A flag indicating whether or not this class supports expert maps
-        """
-        raise NotImplementedError
-
     def supports_packed_ue8m0_act_scales(self) -> bool:
         """
         A flag indicating whether or not this class can process packed ue8m0
@@ -1567,12 +1560,6 @@ class FusedMoEKernel:
             == self.fused_experts.activation_format()
         )
 
-    def supports_expert_map(self) -> bool:
-        """
-        A flag indicating whether or not this class supports expert maps.
-        """
-        return self.fused_experts.supports_expert_map()
-
     def output_is_reduced(self) -> bool:
         """
         Indicates whether or not the output of fused MoE kernel
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py
index ffbb4c4a7d3..89f3843cc50 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py
@@ -17,7 +17,7 @@ def _quantize_and_setup_dispatch(
     a1: torch.Tensor,
     quant_config: FusedMoEQuantConfig,
     defer_input_quant: bool = False,
-) -> tuple[torch.Tensor, list[torch.Tensor] | None]:
+) -> tuple[torch.Tensor, list[torch.Tensor] | None, torch.Tensor | None]:
     # Defer input quantization to the MoE kernel.
     if defer_input_quant:
         a1q = a1
@@ -33,7 +33,7 @@ def _quantize_and_setup_dispatch(
         # which makes the scales tensor different shape than
         # the hidden states, breaking the A2A kernel. So, we
         # delay the swizzling until after the A2A.
-        a1q, a1q_scale = a1q, a1q_scale = moe_kernel_quantize_input(
+        a1q, a1q_scale = moe_kernel_quantize_input(
             a1,
             input_sf,
             quant_dtype=quant_config.quant_dtype,
@@ -49,7 +49,7 @@ def _quantize_and_setup_dispatch(
     skip_gather_scales = a1q_scale is None or a1q_scale.ndim == 0
     scales = None if skip_gather_scales else [a1q_scale]
 
-    return a1q, scales
+    return a1q, scales, a1q_scale
 
 
 def _unwrap_scale_and_prepare_for_moe(
@@ -129,7 +129,9 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular
             )
             a1 = a1 * topk_weights.to(a1.dtype)
 
-        a1q, scales = _quantize_and_setup_dispatch(a1, quant_config, defer_input_quant)
+        a1q, scales, a1q_scale_orig = _quantize_and_setup_dispatch(
+            a1, quant_config, defer_input_quant
+        )
 
         # When LoRA is active, dispatch the per-token LoRA id along with
         # hidden_states so every rank receives the correct mapping for the
@@ -164,7 +166,7 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular
         if extra_tensors is None:
             assert len(res) == 3
             a1q, topk_weights, topk_ids = res
-            a1q_scale = None
+            a1q_scale = a1q_scale_orig
         else:
             assert len(res) == 4
             a1q, topk_weights, topk_ids, gathered_extras = res
@@ -178,7 +180,7 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular
                     gathered_extras, quant_config
                 )
             else:
-                a1q_scale = None
+                a1q_scale = a1q_scale_orig
 
         return a1q, a1q_scale, None, topk_ids, topk_weights
 
@@ -249,7 +251,9 @@ class MoEPrepareAndFinalizeNaiveDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMono
     ) -> mk.PrepareMonolithicResultType:
         """Quantize and Dispatch Router Logits."""
 
-        a1q, scales = _quantize_and_setup_dispatch(a1, quant_config, defer_input_quant)
+        a1q, scales, a1q_scale_orig = _quantize_and_setup_dispatch(
+            a1, quant_config, defer_input_quant
+        )
 
         res = get_ep_group().dispatch_router_logits(
             a1q,
@@ -261,7 +265,7 @@ class MoEPrepareAndFinalizeNaiveDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMono
         if scales is None:
             assert len(res) == 2
             a1q, router_logits = res
-            a1q_scale = None
+            a1q_scale = a1q_scale_orig
         else:
             assert len(res) == 3
             a1q, router_logits, scales = res
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py
index da5d85e4abc..14ef8bf614c 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py
@@ -405,8 +405,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
             topk_ids,
             activation=layer.activation,
             global_num_experts=layer.global_num_experts,
-            # TODO(rob): investigate the disable_expert_map introduced by:
-            # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501
             expert_map=layer.expert_map,
             apply_router_weight_on_input=layer.apply_router_weight_on_input,
             shared_experts=shared_experts,

From 8c6daf6e2fe8b8e731866700719741def43ca165 Mon Sep 17 00:00:00 2001
From: Flora Feng <4florafeng@gmail.com>
Date: Fri, 29 May 2026 18:52:46 -0400
Subject: [PATCH 03/35] [CI] Remove duplicate Harmony test coverage (#44023)

Signed-off-by: sfeng33 <4florafeng@gmail.com>
---
 .../chat_completion/test_serving_chat.py      | 85 -------------------
 .../test_serving_chat_stream_harmony.py       | 79 ++---------------
 2 files changed, 8 insertions(+), 156 deletions(-)

diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py
index 45fae821af3..11793ac8f49 100644
--- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py
+++ b/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
diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py b/tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py
index 0a0802a7939..1c058adaf0a 100644
--- a/tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py
+++ b/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):

From 8fad266507156d3666e9307a52a74252dc3e8bd7 Mon Sep 17 00:00:00 2001
From: "Kevin H. Luu" <khluu000@gmail.com>
Date: Fri, 29 May 2026 16:28:32 -0700
Subject: [PATCH 04/35] [CI] Fix smoke test step key to bypass block gate
 (#43974)

Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .buildkite/image_build/image_build.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.buildkite/image_build/image_build.yaml b/.buildkite/image_build/image_build.yaml
index f401f60ec96..35a074c98cc 100644
--- a/.buildkite/image_build/image_build.yaml
+++ b/.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:

From 187457a952cbaf21e28944920e0b93a28f6cb1bd Mon Sep 17 00:00:00 2001
From: bnellnm <49004751+bnellnm@users.noreply.github.com>
Date: Fri, 29 May 2026 19:45:29 -0400
Subject: [PATCH 05/35] =?UTF-8?q?Revert=20"[MoE=20Refactor]=20Migrate=20Mo?=
 =?UTF-8?q?eWNA16Method=20quantization=20to=20MK=20orac=E2=80=A6=20(#44033?=
 =?UTF-8?q?)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Bill Nell <bnell@redhat.com>
---
 .../layers/fused_moe/experts/triton_moe.py    |  50 ++---
 .../layers/fused_moe/fused_moe.py             |   7 +-
 .../layers/fused_moe/oracle/int_wna16.py      |  97 +---------
 .../layers/quantization/auto_gptq.py          |  64 ++++---
 .../layers/quantization/awq_marlin.py         |   2 -
 .../compressed_tensors_moe_wna16_marlin.py    |  17 +-
 .../layers/quantization/moe_wna16.py          | 176 +++---------------
 .../layers/quantization/utils/gptq_utils.py   |  18 +-
 8 files changed, 98 insertions(+), 333 deletions(-)

diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
index cf2c2cff6a8..25dd0584de0 100644
--- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
@@ -43,10 +43,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8Static128BlockSym,
     kFp8StaticChannelSym,
     kFp8StaticTensorSym,
-    kInt4Static,
-    kInt4Static32,
     kInt8DynamicTokenSym,
-    kInt8Static,
     kInt8StaticChannelSym,
 )
 from vllm.platforms import current_platform
@@ -459,45 +456,40 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
 class TritonWNA16Experts(TritonExperts):
     @staticmethod
     def _supports_current_device() -> bool:
-        return current_platform.is_cuda_alike() or current_platform.is_xpu()
+        raise NotImplementedError(
+            "TritonWNA16Experts is not yet used by an Oracle. "
+            "This method should not be called."
+        )
 
     @staticmethod
     def _supports_no_act_and_mul() -> bool:
-        return True
+        raise NotImplementedError(
+            "TritonWNA16Experts is not yet used by an Oracle. "
+            "This method should not be called."
+        )
 
     @staticmethod
     def _supports_quant_scheme(
         weight_key: QuantKey | None,
         activation_key: QuantKey | None,
     ) -> bool:
-        SUPPORTED_W = [
-            kInt4Static,
-            kInt8Static,
-            kInt4Static32,
-            # other group sizes?
-        ]
-        return weight_key in SUPPORTED_W
+        raise NotImplementedError(
+            "TritonWNA16Experts is not yet used by an Oracle. "
+            "This method should not be called."
+        )
 
     @staticmethod
     def _supports_activation(activation: MoEActivation) -> bool:
-        return activation in [
-            MoEActivation.SILU,
-            MoEActivation.GELU,
-            MoEActivation.GELU_TANH,
-            MoEActivation.SWIGLUOAI,
-            MoEActivation.SWIGLUSTEP,
-            MoEActivation.SILU_NO_MUL,
-            MoEActivation.GELU_NO_MUL,
-            MoEActivation.GELU_TANH_NO_MUL,
-            MoEActivation.RELU2_NO_MUL,
-        ]
+        raise NotImplementedError(
+            "TritonWNA16Experts is not yet used by an Oracle. "
+            "This method should not be called."
+        )
 
     @staticmethod
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
-        # Why?
-        return not (
-            moe_parallel_config.use_fi_nvl_two_sided_kernels
-            or moe_parallel_config.use_fi_nvl_one_sided_kernels
+        raise NotImplementedError(
+            "TritonWNA16Experts is not yet used by an Oracle. "
+            "This method should not be called."
         )
 
     def apply(
@@ -520,9 +512,7 @@ class TritonWNA16Experts(TritonExperts):
     ):
         # Check constraints.
         if self.quant_config.use_int4_w4a16:
-            assert hidden_states.size(-1) // 2 == w1.size(2), (
-                f"Hidden size mismatch {hidden_states.size(-1) // 2} == {w1.size(2)}"
-            )
+            assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"
         else:
             assert hidden_states.size(-1) == w1.size(2), (
                 f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}"
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
index d7b30ef7633..49957c8f5e3 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -30,7 +30,6 @@ from vllm.model_executor.layers.fused_moe.utils import (
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
-from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import direct_register_custom_op
 
 logger = init_logger(__name__)
@@ -1244,11 +1243,7 @@ def get_default_config(
         bit = 4 if dtype == "int4_w4a16" else 8
         use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk, block_shape[1], E, bit)
         if use_moe_wna16_cuda:
-            config = {
-                "BLOCK_SIZE_M": min(16, next_power_of_2(M)),
-                "GROUP_SIZE_M": 1,
-                "SPLIT_K": 1,
-            }
+            config = {"BLOCK_SIZE_M": min(16, M), "SPLIT_K": 1}
         elif M <= 20:
             config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1, "SPLIT_K": 1}
         elif M <= 40:
diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
index dd96921d707..9de7a6ba119 100644
--- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
@@ -23,9 +23,6 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
     MarlinExperts,
     MarlinExpertsBase,
 )
-from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
-    TritonWNA16Experts,
-)
 from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import (
     TrtLlmMxint4ExpertsMonolithic,
 )
@@ -34,7 +31,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     marlin_act_int8_process_scales,
     marlin_moe_permute_scales,
     marlin_permute_bias,
-    marlin_zero_points,
     moe_awq_to_marlin_zero_points,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -49,7 +45,6 @@ class WNA16MoEBackend(Enum):
     MARLIN = "MARLIN"
     BATCHED_MARLIN = "BATCHED_MARLIN"
     FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
-    TRITON = "TRITON"
     XPU = "XPU"
 
 
@@ -63,8 +58,6 @@ def backend_to_kernel_cls(
         return [BatchedMarlinExperts]
     elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
         return [TrtLlmMxint4ExpertsMonolithic]
-    elif backend == WNA16MoEBackend.TRITON:
-        return [TritonWNA16Experts]
     elif backend == WNA16MoEBackend.XPU:
         from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
             XPUExpertsWNA16,
@@ -75,38 +68,24 @@ def backend_to_kernel_cls(
         raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")
 
 
-def _get_priority_backends(
-    may_have_zp: bool, may_have_bias: bool
-) -> list[WNA16MoEBackend]:
+def _get_priority_backends() -> list[WNA16MoEBackend]:
     """
     Get available backends in priority order based on platform and config.
     """
     if current_platform.is_xpu():
         return [WNA16MoEBackend.XPU]
 
-    _AVAILABLE_BACKENDS = []
-
-    if not may_have_zp and not may_have_bias:
-        _AVAILABLE_BACKENDS.append(WNA16MoEBackend.FLASHINFER_TRTLLM)
-
-    # Marlin supports ZP and bias
-    _AVAILABLE_BACKENDS += [
+    _AVAILABLE_BACKENDS = [
+        WNA16MoEBackend.FLASHINFER_TRTLLM,
         WNA16MoEBackend.MARLIN,
         WNA16MoEBackend.BATCHED_MARLIN,
     ]
     return _AVAILABLE_BACKENDS
 
-    if not may_have_bias:
-        _AVAILABLE_BACKENDS.append(WNA16MoEBackend.TRITON)
-
-    return _AVAILABLE_BACKENDS
-
 
 def select_wna16_moe_backend(
     config: FusedMoEConfig,
     weight_key: QuantKey,
-    may_have_zp: bool,
-    may_have_bias: bool,
 ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]:
     """Select the WNA16 MoE backend.
 
@@ -157,7 +136,7 @@ def select_wna16_moe_backend(
         raise ValueError(_make_log_unsupported(backend, reason))
 
     # Select kernels in order of backend.
-    AVAILABLE_BACKENDS = _get_priority_backends(may_have_zp, may_have_bias)
+    AVAILABLE_BACKENDS = _get_priority_backends()
 
     for backend in AVAILABLE_BACKENDS:
         activation_key = None  # always BF16 activation for WNA16 MoE
@@ -239,7 +218,6 @@ def make_wna16_moe_kernel(
     assert experts_cls in (
         MarlinExperts,
         BatchedMarlinExperts,
-        TritonWNA16Experts,
         TrtLlmMxint4ExpertsMonolithic,
         XPUExpertsWNA16,
     )
@@ -256,7 +234,6 @@ def make_wna16_moe_kernel(
     assert prepare_finalize is not None
 
     logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")
-    logger.info_once("Using %s", experts_cls.__name__, scope="local")
 
     extra_args: dict[str, Any] = {}
     if issubclass(experts_cls, MarlinExpertsBase):
@@ -426,8 +403,6 @@ def _process_weights_marlin(
     w2_input_global_scale: torch.Tensor | None = None
     w13_bias_out: torch.Tensor | None = None
     w2_bias_out: torch.Tensor | None = None
-    w13_qzeros_out: torch.Tensor | None = None
-    w2_qzeros_out: torch.Tensor | None = None
 
     # --- FP8 weight / scale adjustment ---
     if input_dtype == torch.float8_e4m3fn:
@@ -527,24 +502,6 @@ def _process_weights_marlin(
     if w2_bias is not None:
         w2_bias_out = marlin_permute_bias(w2_bias)
 
-    if w13_qzeros is not None:
-        w13_qzeros_out = marlin_zero_points(
-            w13_qzeros,
-            size_k=layer.intermediate_size_per_partition,
-            size_n=w13_qzeros.shape[2],
-            num_bits=num_bits,
-            is_a_8bit=is_a_8bit,
-        )
-
-    if w2_qzeros is not None:
-        w2_qzeros_out = marlin_zero_points(
-            w2_qzeros,
-            size_k=w2_qzeros.shape[1] * group_size_or_pack_factor,
-            size_n=w2_qzeros.shape[2],
-            num_bits=num_bits,
-            is_a_8bit=is_a_8bit,
-        )
-
     return (
         marlin_w13_qweight,
         marlin_w2_qweight,
@@ -554,8 +511,8 @@ def _process_weights_marlin(
         w2_g_idx,
         w13_g_idx_sort_indices,
         w2_g_idx_sort_indices,
-        w13_qzeros_out,
-        w2_qzeros_out,
+        w13_qzeros,
+        w2_qzeros,
         w13_input_global_scale,
         w2_input_global_scale,
         w13_bias_out,
@@ -823,9 +780,6 @@ def convert_to_wna16_moe_kernel_format(
         from vllm.model_executor.layers.quantization.awq_marlin import (
             AWQMarlinConfig,
         )
-        from vllm.model_executor.layers.quantization.moe_wna16 import (
-            MoeWNA16Config,
-        )
 
         if isinstance(quant_config, AWQMarlinConfig):
             if w13_qzeros is None or w2_qzeros is None:
@@ -860,17 +814,11 @@ def convert_to_wna16_moe_kernel_format(
             pack_factor = 32 // quant_config.num_bits
             group_size = quant_config.group_size
             actorder = quant_config.actorder
-        elif isinstance(quant_config, MoeWNA16Config):
-            num_bits = quant_config.weight_bits
-            pack_factor = quant_config.bit8_pack_factor
-            group_size = quant_config.group_size
-            actorder = None
         else:
             raise TypeError(
                 "Marlin WNA16 MoE backend requires AutoGPTQConfig, AWQMarlinConfig or "
                 f"QuantizationArgs, got {type(quant_config).__name__}."
             )
-
         if w13_g_idx is None or w2_g_idx is None:
             raise ValueError("GPTQ Marlin MoE requires g_idx tensors.")
         return _process_weights_marlin(
@@ -902,39 +850,6 @@ def convert_to_wna16_moe_kernel_format(
             w13_bias,
             w2_bias,
         )
-    elif backend == WNA16MoEBackend.TRITON:
-        # Convert from int32 to uint8 format for Triton kernel.
-        # This changes the shape from (E, N, K // 8) to (E, N, K // 2) for int4,
-        # which matches what the Triton kernel expects.
-        w13_uint8 = w13.view(torch.uint8)
-        w2_uint8 = w2.view(torch.uint8)
-        return (
-            w13_uint8,
-            w2_uint8,
-            w13_scale,
-            w2_scale,
-            None,
-            None,
-            None,
-            None,
-            w13_qzeros,
-            w2_qzeros,
-            None,
-            None,
-            w13_bias,
-            w2_bias,
-        )
-    elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
-        return _process_weights_flashinfer(
-            w13,
-            w2,
-            w13_scale,
-            w2_scale,
-            w13_g_idx,
-            w2_g_idx,
-            w13_bias,
-            w2_bias,
-        )
     elif backend == WNA16MoEBackend.XPU:
         assert quant_config is not None
         (
diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py
index cb2dad3d524..1821fd5c7f7 100644
--- a/vllm/model_executor/layers/quantization/auto_gptq.py
+++ b/vllm/model_executor/layers/quantization/auto_gptq.py
@@ -485,8 +485,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
         self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
             moe,
             weight_key,
-            may_have_zp=True,
-            may_have_bias=True,
         )
 
     def create_weights(
@@ -646,22 +644,12 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
         layer.workspace = marlin_make_workspace_new(device, 4)
 
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
-        def replace_or_register(name: str, val: torch.Tensor | None):
-            if val is None:
-                return
-
-            if hasattr(layer, name):
-                replace_parameter(layer, name, val)
-            else:
-                layer.register_parameter(
-                    name, torch.nn.Parameter(val, requires_grad=False)
-                )
-
         is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
 
-        assert not is_a_8bit or self.quant_config.quant_type.size_bits == 8, (
-            "W8A8-INT8 is not supported by marlin kernel."
-        )
+        if is_a_8bit:
+            assert self.quant_config.quant_type.size_bits == 8, (
+                "W8A8-INT8 is not supported by marlin kernel."
+            )
 
         (
             w13,
@@ -672,8 +660,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
             w2_g_idx,
             w13_g_idx_sort_indices,
             w2_g_idx_sort_indices,
-            w13_qzeros,
-            w2_qzeros,
+            _w13_qzeros,
+            _w2_qzeros,
             w13_input_global_scale,
             w2_input_global_scale,
             w13_bias,
@@ -691,8 +679,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
             w2_g_idx=layer.w2_g_idx,
             w13_bias=getattr(layer, "w13_bias", None),
             w2_bias=getattr(layer, "w2_bias", None),
-            w13_qzeros=getattr(layer, "w13_qzeros", None),
-            w2_qzeros=getattr(layer, "w2_qzeros", None),
         )
 
         replace_parameter(layer, "w13_qweight", w13)
@@ -703,12 +689,38 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
         replace_parameter(layer, "w2_g_idx", w2_g_idx)
         replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
         replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
-        replace_or_register("w13_input_global_scale", w13_input_global_scale)
-        replace_or_register("w2_input_global_scale", w2_input_global_scale)
-        replace_or_register("w13_bias", w13_bias)
-        replace_or_register("w2_bias", w2_bias)
-        replace_or_register("w13_qzeros", w13_qzeros)
-        replace_or_register("w2_qzeros", w2_qzeros)
+        if w13_input_global_scale is not None:
+            if hasattr(layer, "w13_input_global_scale"):
+                replace_parameter(
+                    layer, "w13_input_global_scale", w13_input_global_scale
+                )
+            else:
+                layer.register_parameter(
+                    "w13_input_global_scale",
+                    torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
+                )
+        if w2_input_global_scale is not None:
+            if hasattr(layer, "w2_input_global_scale"):
+                replace_parameter(layer, "w2_input_global_scale", w2_input_global_scale)
+            else:
+                layer.register_parameter(
+                    "w2_input_global_scale",
+                    torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
+                )
+        if w13_bias is not None:
+            if hasattr(layer, "w13_bias"):
+                replace_parameter(layer, "w13_bias", w13_bias)
+            else:
+                layer.register_parameter(
+                    "w13_bias", torch.nn.Parameter(w13_bias, requires_grad=False)
+                )
+        if w2_bias is not None:
+            if hasattr(layer, "w2_bias"):
+                replace_parameter(layer, "w2_bias", w2_bias)
+            else:
+                layer.register_parameter(
+                    "w2_bias", torch.nn.Parameter(w2_bias, requires_grad=False)
+                )
 
         self._setup_kernel(layer)
 
diff --git a/vllm/model_executor/layers/quantization/awq_marlin.py b/vllm/model_executor/layers/quantization/awq_marlin.py
index 48fbe6335a3..81c0fcb331e 100644
--- a/vllm/model_executor/layers/quantization/awq_marlin.py
+++ b/vllm/model_executor/layers/quantization/awq_marlin.py
@@ -524,8 +524,6 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
         self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
             moe,
             kInt4Static,
-            may_have_zp=self.quant_config.zero_point,
-            may_have_bias=True,
         )
 
     def create_weights(
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
index 4b591fabcc9..2d629d73edd 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
@@ -88,13 +88,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
         self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
             config=self.moe,
             weight_key=weight_key,
-            may_have_zp=False,
-            may_have_bias=False,
         )
-        self.is_marlin = self.wna16_backend in [
-            WNA16MoEBackend.MARLIN,
-            WNA16MoEBackend.BATCHED_MARLIN,
-        ]
 
     def get_weight_shape(
         self,
@@ -120,6 +114,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
                 "num_groups_w2 must be provided for weight scales"
             )
         w13_num_shards = 2 if self.moe.is_act_and_mul else 1
+        is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM
         shape_map = {
             "w13_weight": {
                 "Flashinfer": (
@@ -162,7 +157,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
                 "Marlin": (num_experts, num_groups_w2, hidden_size),
             },
         }
-        backend_key = "Marlin" if self.is_marlin else "Flashinfer"
+        backend_key = "Flashinfer" if is_flashinfer else "Marlin"
         return shape_map[weight_name][backend_key]
 
     def create_weights(
@@ -179,8 +174,9 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
         # Will transpose the loaded weight along the
         # intermediate and hidden dim sizes. Will
         # shard for TP along the transposed dims
+        is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM
         extra_weight_attrs.update(
-            {"is_transposed": self.is_marlin, "quant_method": self.strategy}
+            {"is_transposed": is_transposed, "quant_method": self.strategy}
         )
 
         w13_weight = torch.nn.Parameter(
@@ -328,6 +324,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         # Process weights using the shared oracle infrastructure
+        is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM
         (
             w13_qweight,
             w2_qweight,
@@ -363,7 +360,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
         replace_parameter(layer, "w2_weight_scale", w2_scales)
 
         # Marlin-specific parameters (not needed for Flashinfer)
-        if self.is_marlin:
+        if not is_flashinfer:
             replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed)
             replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed)
             replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
@@ -395,7 +392,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
 
         # Add Marlin-specific arguments
         marlin_args: dict[str, Any] = {}
-        if self.is_marlin:
+        if not is_flashinfer:
             marlin_args = {
                 "w13_g_idx": layer.w13_weight_g_idx,
                 "w2_g_idx": layer.w2_weight_g_idx,
diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py
index fd99d520c64..471febab044 100644
--- a/vllm/model_executor/layers/quantization/moe_wna16.py
+++ b/vllm/model_executor/layers/quantization/moe_wna16.py
@@ -13,15 +13,11 @@ from vllm.model_executor.layers.fused_moe import (
     RoutedExperts,
     SharedExperts,
 )
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
-)
-from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
-    WNA16MoEBackend,
-    convert_to_wna16_moe_kernel_format,
-    make_wna16_moe_kernel,
-    make_wna16_moe_quant_config,
-    select_wna16_moe_backend,
+    int4_w4a16_moe_quant_config,
+    int8_w8a16_moe_quant_config,
 )
 from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
     UnquantizedFusedMoEMethod,
@@ -35,15 +31,7 @@ from vllm.model_executor.layers.quantization.base_config import (
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     check_marlin_supports_layer,
 )
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    INT4_DTYPE,
-    INT8_DTYPE,
-    QuantKey,
-    kInt4Static32GroupScale,
-    kInt4StaticGroupScale,
-    kInt8StaticGroupScale,
-)
-from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 
 
@@ -228,33 +216,6 @@ class MoeWNA16Method(FusedMoEMethodBase):
         super().__init__(moe)
         self.quant_config = quant_config
 
-        num_bits = self.quant_config.weight_bits
-        group_size = self.quant_config.group_size
-
-        if num_bits == 4:
-            quant_type = INT4_DTYPE
-            if group_size == 32:
-                scale = kInt4Static32GroupScale
-            else:
-                scale = kInt4StaticGroupScale
-        elif num_bits == 8:
-            assert group_size == -1
-            quant_type = INT8_DTYPE
-            scale = kInt8StaticGroupScale
-        else:
-            raise ValueError("MoeWNA16Method only supports int4 and int8 now.")
-
-        weight_key = QuantKey(quant_type, scale)
-
-        # Select WNA16 MoE backend via oracle.
-        # handle ZP?
-        self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
-            config=self.moe,
-            weight_key=weight_key,
-            may_have_zp=self.quant_config.has_zp,
-            may_have_bias=False,
-        )
-
     def create_weights(
         self,
         layer: RoutedExperts,
@@ -375,89 +336,24 @@ class MoeWNA16Method(FusedMoEMethodBase):
                 layer.register_parameter(key, param)
                 set_weight_attrs(param, extra_weight_attrs)
 
-    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        has_zp = self.quant_config.has_zp
-        (
-            w13_qweight,
-            w2_qweight,
-            w13_scales,
-            w2_scales,
-            w13_g_idx_processed,
-            w2_g_idx_processed,
-            w13_g_idx_sort_indices,
-            w2_g_idx_sort_indices,
-            w13_qzeros,
-            w2_qzeros,
-            w13_input_global_scale,
-            w2_input_global_scale,
-            _,  # w13_bias
-            _,  # w2_bias
-        ) = convert_to_wna16_moe_kernel_format(
-            backend=self.wna16_backend,
-            layer=layer,
-            quant_config=self.quant_config,
-            input_dtype=None,
-            w13=layer.w13_qweight,
-            w2=layer.w2_qweight,
-            w13_scale=layer.w13_scales,
-            w2_scale=layer.w2_scales,
-            w13_g_idx=getattr(layer, "w13_g_idx", None),
-            w2_g_idx=getattr(layer, "w2_g_idx", None),
-            w13_qzeros=layer.w13_qzeros if has_zp else None,
-            w2_qzeros=layer.w2_qzeros if has_zp else None,
-        )
-
-        # Replace common parameters
-        replace_parameter(layer, "w13_qweight", w13_qweight)
-        replace_parameter(layer, "w2_qweight", w2_qweight)
-        replace_parameter(layer, "w13_scales", w13_scales)
-        replace_parameter(layer, "w2_scales", w2_scales)
-
-        if has_zp:
-            assert w13_qzeros is not None and w2_qzeros is not None
-            replace_parameter(layer, "w13_qzeros", w13_qzeros)
-            replace_parameter(layer, "w2_qzeros", w2_qzeros)
-
-        # Marlin-specific parameters (not needed for Flashinfer)
-        if self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM:
-            replace_parameter(layer, "w13_g_idx", w13_g_idx_processed)
-            replace_parameter(layer, "w2_g_idx", w2_g_idx_processed)
-            replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
-            replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
-
-            # Register input global scales if present
-            if w13_input_global_scale is not None:
-                layer.register_parameter(
-                    "w13_input_global_scale",
-                    torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
-                )
-            if w2_input_global_scale is not None:
-                layer.register_parameter(
-                    "w2_input_global_scale",
-                    torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
-                )
-
-        assert self.experts_cls is not None
-        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
-        assert self.moe_quant_config is not None
-        self.moe_kernel = make_wna16_moe_kernel(
-            moe_quant_config=self.moe_quant_config,
-            moe_config=self.moe,
-            experts_cls=self.experts_cls,
-            routing_tables=layer._expert_routing_tables(),
-        )
-
     def get_fused_moe_quant_config(
         self, layer: RoutedExperts
     ) -> FusedMoEQuantConfig | None:
+        weight_bits = self.quant_config.weight_bits
         has_zp = self.quant_config.has_zp
-        return make_wna16_moe_quant_config(
+        assert weight_bits == 4 or weight_bits == 8
+        config_builder = (
+            int4_w4a16_moe_quant_config
+            if weight_bits == 4
+            else int8_w8a16_moe_quant_config
+        )
+
+        return config_builder(
             w1_scale=layer.w13_scales,
             w2_scale=layer.w2_scales,
             w1_zp=layer.w13_qzeros if has_zp else None,
             w2_zp=layer.w2_qzeros if has_zp else None,
-            group_size=layer.group_size,
-            num_bits=self.quant_config.weight_bits,
+            block_shape=[0, layer.group_size],
         )
 
     def apply(
@@ -469,44 +365,22 @@ class MoeWNA16Method(FusedMoEMethodBase):
         shared_experts: SharedExperts | None,
         shared_experts_input: torch.Tensor | None,
     ) -> torch.Tensor:
-        assert not self.is_monolithic
-        assert self.moe_kernel is not None
-        return self.moe_kernel.apply(
+        from vllm.model_executor.layers.fused_moe import fused_experts
+
+        assert layer.activation == MoEActivation.SILU, (
+            f"Only SiLU activation is supported, not {layer.activation}."
+        )
+
+        return fused_experts(
             x,
             layer.w13_qweight,
             layer.w2_qweight,
-            topk_weights,
-            topk_ids,
-            activation=layer.activation,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            apply_router_weight_on_input=layer.apply_router_weight_on_input,
             global_num_experts=layer.global_num_experts,
             expert_map=layer.expert_map,
-            apply_router_weight_on_input=layer.apply_router_weight_on_input,
-            shared_experts=shared_experts,
-            shared_experts_input=shared_experts_input,
-        )
-
-    def apply_monolithic(
-        self,
-        layer: RoutedExperts,
-        x: torch.Tensor,
-        router_logits: torch.Tensor,
-        input_ids: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        assert self.is_monolithic
-        assert self.moe_kernel is not None
-        return self.moe_kernel.apply_monolithic(
-            x,
-            layer.w13_weight,
-            layer.w2_weight,
-            router_logits,
-            activation=layer.activation,
-            global_num_experts=layer.global_num_experts,
-            expert_map=layer.expert_map,
-            apply_router_weight_on_input=layer.apply_router_weight_on_input,
-            num_expert_group=layer.num_expert_group,
-            topk_group=layer.topk_group,
-            e_score_correction_bias=layer.e_score_correction_bias,
-            routed_scaling_factor=layer.routed_scaling_factor,
+            quant_config=self.moe_quant_config,
         )
 
     @staticmethod
diff --git a/vllm/model_executor/layers/quantization/utils/gptq_utils.py b/vllm/model_executor/layers/quantization/utils/gptq_utils.py
index c73c42f4ff0..691d80b0b74 100644
--- a/vllm/model_executor/layers/quantization/utils/gptq_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/gptq_utils.py
@@ -3,7 +3,7 @@
 from collections.abc import Mapping
 from copy import deepcopy
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 import regex as re
 import torch
@@ -69,20 +69,6 @@ def get_dynamic_override(
     return default_value
 
 
-def flatten_list(lst: list[Any]) -> list[Any]:
-    output = []
-
-    def _flatten(lst: list[Any]):
-        for i in lst:
-            if isinstance(i, list):
-                _flatten(i)
-            else:
-                output.append(i)
-
-    _flatten(lst)
-    return output
-
-
 def is_layer_gptq_quantized(
     prefix: str,
     quantized_layers: list[str],
@@ -97,8 +83,6 @@ def is_layer_gptq_quantized(
 
     proj_name = prefix.split(".")[-1]
 
-    quantized_layers = flatten_list(quantized_layers)
-
     # Fused layers like gate_up_proj or qkv_proj will not be fused
     # in the safetensors checkpoint. So, we convert the name
     # from the fused version to unfused + check to make sure that

From 559d6710bf45e6fb3d48429855702b6ff24bdfe0 Mon Sep 17 00:00:00 2001
From: Jee Jee Li <pandaleefree@gmail.com>
Date: Sat, 30 May 2026 09:28:34 +0800
Subject: [PATCH 06/35] [PERF]MiniMax-M2 gate kernel (#38445)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Signed-off-by: qianlihuang <91178480+qianlihuang@users.noreply.github.com>
Co-authored-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com>
---
 CMakeLists.txt                                |  34 ++-
 benchmarks/kernels/benchmark_router_gemm.py   | 154 ++++++++++++
 cmake/utils.cmake                             |  10 +
 csrc/libtorch_stable/fp32_router_gemm.cu      | 223 ++++++++++++++++++
 .../libtorch_stable/fp32_router_gemm_entry.cu | 127 ++++++++++
 csrc/libtorch_stable/torch_bindings.cpp       |   4 +
 tests/kernels/test_fp32_router_gemm.py        |  78 ++++++
 vllm/_custom_ops.py                           |  25 ++
 .../layers/fused_moe/router/gate_linear.py    |  76 +++++-
 vllm/model_executor/models/minimax_m2.py      |   8 +-
 10 files changed, 716 insertions(+), 23 deletions(-)
 create mode 100644 benchmarks/kernels/benchmark_router_gemm.py
 create mode 100644 csrc/libtorch_stable/fp32_router_gemm.cu
 create mode 100644 csrc/libtorch_stable/fp32_router_gemm_entry.cu
 create mode 100644 tests/kernels/test_fp32_router_gemm.py

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 3f571170401..4d2d9785c89 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -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()
diff --git a/benchmarks/kernels/benchmark_router_gemm.py b/benchmarks/kernels/benchmark_router_gemm.py
new file mode 100644
index 00000000000..ba46a7fd78d
--- /dev/null
+++ b/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)
diff --git a/cmake/utils.cmake b/cmake/utils.cmake
index f81882ccbc2..f10ba93f7c6 100644
--- a/cmake/utils.cmake
+++ b/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
diff --git a/csrc/libtorch_stable/fp32_router_gemm.cu b/csrc/libtorch_stable/fp32_router_gemm.cu
new file mode 100644
index 00000000000..04397e0893c
--- /dev/null
+++ b/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
diff --git a/csrc/libtorch_stable/fp32_router_gemm_entry.cu b/csrc/libtorch_stable/fp32_router_gemm_entry.cu
new file mode 100644
index 00000000000..4baa740de93
--- /dev/null
+++ b/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));
+}
diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp
index 152391f9c20..13d75445009 100644
--- a/csrc/libtorch_stable/torch_bindings.cpp
+++ b/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, "
diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py
new file mode 100644
index 00000000000..f855eb7aa17
--- /dev/null
+++ b/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"
diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py
index 974828175dc..f12d128f083 100644
--- a/vllm/_custom_ops.py
+++ b/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,
diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py
index a868c9c8487..0a57a6f4dfe 100644
--- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py
+++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py
@@ -3,18 +3,22 @@
 import torch
 from torch.nn.parameter import Parameter
 
+import vllm._custom_ops as ops
 from vllm.model_executor.custom_op import PluggableLayer
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import direct_register_custom_op
 
 
 @PluggableLayer.register("gate_linear")
 class GateLinear(ReplicatedLinear):
-    """MoE gate linear layer with three-tier GEMM dispatch:
+    """MoE gate linear layer with multi-tier GEMM dispatch:
 
-    1. DSV3 specialized kernel (SM90+, batch<=16, supported dims)
-    2. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 + fp32 out_dtype)
-    3. F.linear via ReplicatedLinear (ultimate fallback)
+    1. DSV3 specialized kernel (SM90+, fp32 out, M<=16, H=7168, E=256/384)
+    2. fp32 specialized kernel  (SM90+, bf16/fp32 in, fp32 out,
+       M<=32, H=3072, E=256)
+    3. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
+    4. F.linear via ReplicatedLinear (ultimate fallback)
 
     The ``out_dtype`` attribute is mutable and can be set after init
     (e.g. when the required dtype depends on the expert quantization
@@ -25,6 +29,11 @@ class GateLinear(ReplicatedLinear):
     DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
     DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
 
+    # Dimensions supported by the fp32 specialized kernel
+    FP32_SUPPORTED_NUM_EXPERTS = [256]
+    FP32_SUPPORTED_HIDDEN_SIZES = [3072]
+    FP32_MAX_TOKENS = 32
+
     def __init__(
         self,
         input_size: int,
@@ -43,7 +52,7 @@ class GateLinear(ReplicatedLinear):
         )
 
         # If fp32 compute is required and no specialized kernel is available,
-        # store weights in fp32 so Tier 3 computes in fp32 natively.
+        # store weights in fp32 so the fallback linear path computes in fp32.
         if force_fp32_compute and not can_use_specialized_kernels:
             params_dtype = torch.float32
 
@@ -65,6 +74,16 @@ class GateLinear(ReplicatedLinear):
             and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES
         )
 
+        # fp32 specialized kernel eligibility (SM90+, exact dims, fp32 weight)
+        self.allow_fp32_router_gemm = (
+            not bias
+            and self.weight.dtype == torch.float32
+            and current_platform.is_cuda()
+            and is_hopper_or_blackwell
+            and output_size in self.FP32_SUPPORTED_NUM_EXPERTS
+            and input_size in self.FP32_SUPPORTED_HIDDEN_SIZES
+        )
+
         # cuBLAS bf16→fp32 eligibility
         self.allow_cublas_router_gemm = (
             self.allow_specialized_router_gemm
@@ -92,8 +111,6 @@ class GateLinear(ReplicatedLinear):
     def forward(
         self, x: torch.Tensor
     ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
-        import vllm._custom_ops as ops
-
         # Tier 1: DSV3 specialized kernel
         if self.allow_dsv3_router_gemm and x.shape[0] <= 16:
             output = ops.dsv3_router_gemm(
@@ -103,15 +120,56 @@ class GateLinear(ReplicatedLinear):
             )
             return output, None
 
-        # Tier 2: cuBLAS bf16→fp32
+        # Tier 2: fp32 specialized kernel (H=3072, E=256, M<=32)
+        # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph
+        # capture does not freeze the runtime num_tokens branch.
+        if self.allow_fp32_router_gemm and x.dtype in (
+            torch.float32,
+            torch.bfloat16,
+        ):
+            output = torch.ops.vllm.fp32_router_gemm_dispatch(x, self.weight)
+            return output, None
+
+        # Tier 3: cuBLAS bf16→fp32
         if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16:
             output = torch.mm(x, self.weight.T, out_dtype=torch.float32)
             return output, None
 
-        # Tier 3: F.linear (ReplicatedLinear)
+        # Tier 4: F.linear (ReplicatedLinear)
         if self.out_dtype is not None and x.dtype != self.weight.dtype:
             x = x.to(self.weight.dtype)
         output, output_bias = super().forward(x)
         if self.out_dtype is not None and output.dtype != self.out_dtype:
             output = output.to(self.out_dtype)
         return output, output_bias
+
+
+_FP32_ROUTER_GEMM_MAX_TOKENS = GateLinear.FP32_MAX_TOKENS
+
+
+def fp32_router_gemm_dispatch_impl(
+    x: torch.Tensor, weight: torch.Tensor
+) -> torch.Tensor:
+    """
+    Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS,
+    otherwise fall back to F.linear.
+    This must be wrapped in a custom op because our torch.compile integration
+    does not support runtime dispatching on num_tokens.
+    """
+    if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS:
+        return ops.fp32_router_gemm(x, weight)
+    else:
+        return torch.nn.functional.linear(x.float(), weight)
+
+
+def fp32_router_gemm_dispatch_fake(
+    x: torch.Tensor, weight: torch.Tensor
+) -> torch.Tensor:
+    return x.new_empty((x.shape[0], weight.shape[0]), dtype=torch.float32)
+
+
+direct_register_custom_op(
+    op_name="fp32_router_gemm_dispatch",
+    op_func=fp32_router_gemm_dispatch_impl,
+    fake_impl=fp32_router_gemm_dispatch_fake,
+)
diff --git a/vllm/model_executor/models/minimax_m2.py b/vllm/model_executor/models/minimax_m2.py
index 53cb29d1122..e5da9154150 100644
--- a/vllm/model_executor/models/minimax_m2.py
+++ b/vllm/model_executor/models/minimax_m2.py
@@ -43,10 +43,10 @@ from vllm.model_executor.layers.fused_moe import (
     FusedMoE,
     fused_moe_make_expert_params_mapping,
 )
+from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -113,12 +113,12 @@ class MiniMaxM2MoE(nn.Module):
             router_logits_dtype=torch.float32,
         )
 
-        self.gate = ReplicatedLinear(
+        self.gate = GateLinear(
             config.hidden_size,
             config.num_local_experts,
             bias=False,
             params_dtype=torch.float32,
-            quant_config=None,
+            out_dtype=torch.float32,
             prefix=f"{prefix}.gate",
         )
 
@@ -132,7 +132,7 @@ class MiniMaxM2MoE(nn.Module):
         hidden_states = hidden_states.view(-1, hidden_dim)
 
         # router_logits: (num_tokens, n_experts)
-        router_logits, _ = self.gate(hidden_states.to(torch.float32))
+        router_logits, _ = self.gate(hidden_states)
         final_hidden_states = self.experts(
             hidden_states=hidden_states, router_logits=router_logits
         )

From 1e2ce5d11a9136f03823663a299479cd1cbbacfc Mon Sep 17 00:00:00 2001
From: Gagan Dhakrey <59848316+gagandhakrey@users.noreply.github.com>
Date: Sat, 30 May 2026 07:06:34 +0530
Subject: [PATCH 07/35] offload prompt_embeds decode in render_prompts_async to
 avoid blocking (#43792)

Signed-off-by: Gagan Dhakrey <gagandhakrey@gmail.com>
---
 vllm/renderers/base.py | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py
index 41d8c0075fb..9fab3aff04e 100644
--- a/vllm/renderers/base.py
+++ b/vllm/renderers/base.py
@@ -104,6 +104,9 @@ class BaseRenderer(ABC, Generic[_T]):
         self._process_multimodal_async = make_async(
             self._process_multimodal, executor=self._mm_executor
         )
+        self._safe_load_prompt_embeds_async = make_async(
+            safe_load_prompt_embeds, executor=self._executor
+        )
         if mm_registry.supports_multimodal_inputs(config.model_config):
             mm_processor_cache = mm_registry.processor_cache_from_config(config)
 
@@ -376,11 +379,28 @@ class BaseRenderer(ABC, Generic[_T]):
 
         return [self.render_prompt(prompt) for prompt in prompts]
 
+    async def _render_prompt_async(
+        self,
+        prompt: DictPrompt | bytes,
+    ) -> DictPrompt:
+        if isinstance(prompt, bytes):
+            embeds = await self._safe_load_prompt_embeds_async(
+                self.model_config, prompt
+            )
+            return EmbedsPrompt(prompt_embeds=embeds)
+
+        return prompt
+
     async def render_prompts_async(
         self,
         prompts: Sequence[DictPrompt | bytes],
     ) -> list[DictPrompt]:
-        return self.render_prompts(prompts)
+        if len(prompts) == 0:
+            raise ValueError("You must pass at least one prompt")
+
+        return await asyncio.gather(
+            *(self._render_prompt_async(prompt) for prompt in prompts)
+        )
 
     @abstractmethod
     def render_messages(

From 1a096d82087bda7faaf4cad81d639419c4734869 Mon Sep 17 00:00:00 2001
From: Flora Feng <4florafeng@gmail.com>
Date: Fri, 29 May 2026 21:45:15 -0400
Subject: [PATCH 08/35] [Refactor] Remove dead current_tool_name_sent
 assignments from tool parsers (#43997)

Signed-off-by: sfeng33 <4florafeng@gmail.com>
---
 vllm/tool_parsers/ernie45_tool_parser.py      | 1 -
 vllm/tool_parsers/hunyuan_a13b_tool_parser.py | 3 ---
 vllm/tool_parsers/hy_v3_tool_parser.py        | 1 -
 vllm/tool_parsers/phi4mini_tool_parser.py     | 1 -
 4 files changed, 6 deletions(-)

diff --git a/vllm/tool_parsers/ernie45_tool_parser.py b/vllm/tool_parsers/ernie45_tool_parser.py
index 9722dddf734..f22eaca1f80 100644
--- a/vllm/tool_parsers/ernie45_tool_parser.py
+++ b/vllm/tool_parsers/ernie45_tool_parser.py
@@ -34,7 +34,6 @@ class Ernie45ToolParser(ToolParser):
         abc\n</think>\n\n\n<tool_call>\ndef\n</tool_call>\n
         """
         super().__init__(tokenizer, tools)
-        self.current_tool_name_sent = False
         self.prev_tool_call_arr: list[dict] = []
         self.current_tool_id = -1
         self.streamed_args_for_tool: list[str] = []
diff --git a/vllm/tool_parsers/hunyuan_a13b_tool_parser.py b/vllm/tool_parsers/hunyuan_a13b_tool_parser.py
index 29b2a5eae27..9723ef45d24 100644
--- a/vllm/tool_parsers/hunyuan_a13b_tool_parser.py
+++ b/vllm/tool_parsers/hunyuan_a13b_tool_parser.py
@@ -38,7 +38,6 @@ class HunyuanA13BToolParser(ToolParser):
         # Initialize state for streaming mode
         self.prev_tool_calls: list[dict] = []
         self.current_tool_id = -1
-        self.current_tool_name_sent = False
         self.streamed_args: list[str] = []  # Track arguments sent for each tool
 
         # For backward compatibility with tests
@@ -262,7 +261,6 @@ class HunyuanA13BToolParser(ToolParser):
                         )
                     else:
                         self.streaming_state["sent_tools"][0]["sent_name"] = True
-                    self.current_tool_name_sent = True
                     return delta
         return None
 
@@ -306,7 +304,6 @@ class HunyuanA13BToolParser(ToolParser):
                     ]
                 )
                 self.streaming_state["sent_tools"][current_idx]["sent_name"] = True
-                self.current_tool_name_sent = True
                 while len(self.streamed_args) <= current_idx:
                     self.streamed_args.append("")
                 return delta
diff --git a/vllm/tool_parsers/hy_v3_tool_parser.py b/vllm/tool_parsers/hy_v3_tool_parser.py
index 809a85ce417..496deb4f2d5 100644
--- a/vllm/tool_parsers/hy_v3_tool_parser.py
+++ b/vllm/tool_parsers/hy_v3_tool_parser.py
@@ -246,7 +246,6 @@ class HYV3ToolParser(ToolParser):
     def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
         super().__init__(tokenizer, tools)
 
-        self.current_tool_name_sent: bool = False
         self.prev_tool_call_arr: list[dict] = []
         self.current_tool_id: int = -1
         self.streamed_args_for_tool: list[
diff --git a/vllm/tool_parsers/phi4mini_tool_parser.py b/vllm/tool_parsers/phi4mini_tool_parser.py
index 2dc262bba2e..f2fa3ce9983 100644
--- a/vllm/tool_parsers/phi4mini_tool_parser.py
+++ b/vllm/tool_parsers/phi4mini_tool_parser.py
@@ -47,7 +47,6 @@ class Phi4MiniJsonToolParser(ToolParser):
         # streaming mode
         self.prev_tool_call_arr: list[dict[str, Any]] = []
         self.current_tool_id: int = -1
-        self.current_tool_name_sent: bool = False
         self.streamed_args_for_tool: list[
             str
         ] = []  # map what has been streamed for each tool so far to a list

From ef8840adc73bfbe3108811cebcd8af7252f9b6f0 Mon Sep 17 00:00:00 2001
From: Andreas Karatzas <akaratza@amd.com>
Date: Fri, 29 May 2026 23:14:37 -0500
Subject: [PATCH 09/35] [ROCm][CI] Fix failure in the Phi3V pooling test
 (#44028)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>
---
 tests/models/multimodal/pooling/test_phi3v.py | 72 +++++++++++++++----
 1 file changed, 59 insertions(+), 13 deletions(-)

diff --git a/tests/models/multimodal/pooling/test_phi3v.py b/tests/models/multimodal/pooling/test_phi3v.py
index 2794b0b2937..285ded375da 100644
--- a/tests/models/multimodal/pooling/test_phi3v.py
+++ b/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

From c0056b19bf4930ae7830e753b048b3daca0fbfee Mon Sep 17 00:00:00 2001
From: nemanjaudovic <152565955+nemanjaudovic@users.noreply.github.com>
Date: Sat, 30 May 2026 07:16:57 +0200
Subject: [PATCH 10/35] [ROCm] cmake: support PYTORCH_FOUND_HIP for torch 2.13
 native HIP language support (#43881)

Signed-off-by: nemanjaudovic <nudovic@amd.com>
Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
---
 CMakeLists.txt | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 4d2d9785c89..86c2214b249 100644
--- a/CMakeLists.txt
+++ b/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

From e9499996df8968f473db1f6bc7ec31207022aea0 Mon Sep 17 00:00:00 2001
From: Liangliang Ma <liangliang.ma@intel.com>
Date: Sat, 30 May 2026 14:16:49 +0800
Subject: [PATCH 11/35] [BugFix][Platform] Fix import vllm.platforms.rocm error
 on non-CUDA test_gpt_oss.py (#43571)

Signed-off-by: Ma, Liangliang <liangliang.ma@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
---
 tests/models/quantization/test_gpt_oss.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/tests/models/quantization/test_gpt_oss.py b/tests/models/quantization/test_gpt_oss.py
index fe9ddd2f6ba..783f1773d21 100644
--- a/tests/models/quantization/test_gpt_oss.py
+++ b/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

From 124fac10cb0ea83aee2ffeabac0b413d6b759b26 Mon Sep 17 00:00:00 2001
From: Lanze Liu <86434077+liulanze@users.noreply.github.com>
Date: Fri, 29 May 2026 23:16:53 -0700
Subject: [PATCH 12/35] [Bugfix] Fix RMSNorm kernels to multiply in weight's
 native dtype (#42379)

Signed-off-by: Lanze Liu <lanzetech@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
---
 csrc/libtorch_stable/layernorm_kernels.cu     | 10 ++++----
 .../layernorm_quant_kernels.cu                | 23 +++++--------------
 2 files changed, 10 insertions(+), 23 deletions(-)

diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu
index fb714b1b1e0..37df6be329f 100644
--- a/csrc/libtorch_stable/layernorm_kernels.cu
+++ b/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];
   }
 }
 
diff --git a/csrc/libtorch_stable/layernorm_quant_kernels.cu b/csrc/libtorch_stable/layernorm_quant_kernels.cu
index 26ffa76d6e1..32f3495f4e9 100644
--- a/csrc/libtorch_stable/layernorm_quant_kernels.cu
+++ b/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);
   }

From 3becc5db4034a65802c7d7b867fd236655c0ebcc Mon Sep 17 00:00:00 2001
From: Xiaoran <claire.rrchen@hotmail.com>
Date: Sat, 30 May 2026 03:13:18 -0700
Subject: [PATCH 13/35] [ROCm] Add attention sink support to AITer flash
 attention backend (#43817)

Signed-off-by: Xiaoran Chen <xiaoran@fb.com>
Co-authored-by: Xiaoran Chen <xiaoran@fb.com>
---
 docs/design/attention_backends.md           |  2 +-
 vllm/_aiter_ops.py                          |  2 ++
 vllm/v1/attention/backends/rocm_aiter_fa.py | 26 +++++++++++++++++----
 3 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md
index ed9f0d30162..329a4aacfb6 100644
--- a/docs/design/attention_backends.md
+++ b/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 |
diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py
index 2c9b939fa2f..5a8b690433c 100644
--- a/vllm/_aiter_ops.py
+++ b/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
diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py
index d0cc011fe4e..a9fa45debcf 100644
--- a/vllm/v1/attention/backends/rocm_aiter_fa.py
+++ b/vllm/v1/attention/backends/rocm_aiter_fa.py
@@ -709,6 +709,11 @@ class AiterFlashAttentionMetadataBuilder(
 
 class AiterFlashAttentionBackend(AttentionBackend):
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
+
+    @classmethod
+    def supports_sink(cls) -> bool:
+        return True
+
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
         "auto",
         "float16",
@@ -788,6 +793,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
         logits_soft_cap: float | None = None,
         attn_type: AttentionType = AttentionType.DECODER,
         kv_sharing_target_layer_name: int | None = None,
+        sinks: torch.Tensor | None = None,
     ) -> None:
         self.num_heads = num_heads
         self.head_size = head_size
@@ -806,6 +812,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
             logits_soft_cap = 0.0
         self.logits_soft_cap = logits_soft_cap
         self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
+        self.sinks = sinks
 
         assert self.num_heads % self.num_kv_heads == 0
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
@@ -878,6 +885,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
             alibi_slopes=self.alibi_slopes,
             return_lse=False,
             out=output,
+            sink_ptr=self.sinks,
         )
 
     def extend_forward(
@@ -927,6 +935,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
             window_size=self.sliding_window,
             alibi_slopes=self.alibi_slopes,
             return_lse=True,
+            sink_ptr=self.sinks,
         )
         assert attn_metadata.extend_metadata is not None
         chunk_context_metadata = attn_metadata.extend_metadata.chunk_context_metadata
@@ -974,6 +983,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
                 window_size=self.sliding_window,
                 alibi_slopes=self.alibi_slopes,
                 return_lse=True,
+                sink_ptr=self.sinks,
             )
             if chunked_output is None:
                 chunked_output = suf_out
@@ -1092,6 +1102,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
                     window_size=self.sliding_window,
                     alibi_slopes=self.alibi_slopes,
                     out=output_actual_tokens[num_decode_tokens + num_extend_tokens :],
+                    sink_ptr=self.sinks,
                 )
 
             # calculate for extends
@@ -1136,11 +1147,17 @@ class AiterFlashAttentionImpl(AttentionImpl):
                 assert attn_metadata.decode_metadata is not None
                 decode_max_query_len = attn_metadata.decode_metadata.max_query_len
 
-                # Multi-token speculative decode path.
-                if decode_max_query_len > 1:
+                # Use unified_attention for speculative decoding (multi-token),
+                # sliding window, or sinks
+                # (pa_fwd_asm and paged_attention_v1 don't support sinks)
+                if (
+                    self.sliding_window[0] != -1
+                    or decode_max_query_len > 1
+                    or self.sinks is not None
+                ):
                     assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
-                        "Shuffle KV cache layout is not supported with "
-                        "speculative decoding (multi-token decode)."
+                        "Shuffle KV cache layout is not supported with sliding "
+                        "window, sinks, or speculative decoding (multi-token decode)."
                     )
                     if not attn_metadata.causal:
                         from aiter.ops.triton.attention.mha_v3 import (
@@ -1207,6 +1224,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
                             q_descale=None,
                             k_descale=layer._k_scale.expand(descale_shape),
                             v_descale=layer._v_scale.expand(descale_shape),
+                            sinks=self.sinks,
                         )
                     return
 

From 50c80d792307076bdb811a12f5a80e9e1ea8b27d Mon Sep 17 00:00:00 2001
From: Bugen Zhao <i@bugenzhao.com>
Date: Sat, 30 May 2026 22:23:54 +0800
Subject: [PATCH 14/35] [Governance] Add @BugenZhao as Rust frontend code owner
 (#44047)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>
---
 .github/CODEOWNERS            | 8 +++++++-
 docs/governance/committers.md | 3 +++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index bd5deff1b82..beaaa5d8642 100644
--- a/.github/CODEOWNERS
+++ b/.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
diff --git a/docs/governance/committers.md b/docs/governance/committers.md
index 386e4f2a4bb..738c59df445 100644
--- a/docs/governance/committers.md
+++ b/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
 

From e1105064b282bb807ba9c309741b40a3b64e2261 Mon Sep 17 00:00:00 2001
From: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Date: Sat, 30 May 2026 10:34:33 -0400
Subject: [PATCH 15/35] [Bug] Fix gemma4 MTP IMA issue when TP>1, `CUDA error:
 an illegal memory access was encountered` (#43909)

Signed-off-by: yewentao256 <zhyanwentao@126.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
---
 vllm/model_executor/models/gemma4_mtp.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/vllm/model_executor/models/gemma4_mtp.py b/vllm/model_executor/models/gemma4_mtp.py
index c294ffc6f9a..122855400d9 100644
--- a/vllm/model_executor/models/gemma4_mtp.py
+++ b/vllm/model_executor/models/gemma4_mtp.py
@@ -501,6 +501,7 @@ class Gemma4MTP(nn.Module):
         config = vllm_config.speculative_config.draft_model_config.hf_config
         text_config = _get_text_config(config)
         self.config = config
+        self._stable_full_lm_head_weight: torch.Tensor | None = None
 
         self.model = Gemma4MultiTokenPredictor(
             vllm_config=vllm_config,
@@ -567,6 +568,8 @@ class Gemma4MTP(nn.Module):
         )
 
     def _get_full_lm_head_weight(self) -> torch.Tensor:
+        if self._stable_full_lm_head_weight is not None:
+            return self._stable_full_lm_head_weight
         lm_head_weight = self.lm_head.weight
         tp_size = get_tensor_model_parallel_world_size()
         if tp_size > 1:
@@ -574,7 +577,11 @@ class Gemma4MTP(nn.Module):
                 lm_head_weight,
                 dim=0,
             )
-        return lm_head_weight[: self.masked_embedding.vocab_size]
+        lm_head_weight = lm_head_weight[: self.masked_embedding.vocab_size]
+        if tp_size > 1:
+            lm_head_weight = lm_head_weight.contiguous()
+            self._stable_full_lm_head_weight = lm_head_weight
+        return lm_head_weight
 
     def compute_logits(
         self,
@@ -599,5 +606,6 @@ class Gemma4MTP(nn.Module):
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        self._stable_full_lm_head_weight = None
         loader = AutoWeightsLoader(self)
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

From 27fa5aa3b952a6108de127423397e50364a95fcb Mon Sep 17 00:00:00 2001
From: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Date: Sat, 30 May 2026 09:40:52 -0700
Subject: [PATCH 16/35] [MRV2] Support breakable CUDA graph (#44050)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
---
 vllm/v1/worker/gpu/cudagraph_utils.py         | 38 +++++++++++++++++--
 vllm/v1/worker/gpu/model_runner.py            | 12 +++++-
 .../gpu/spec_decode/eagle/speculator.py       | 14 ++++++-
 3 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py
index c7a7ffe442d..384f1192435 100644
--- a/vllm/v1/worker/gpu/cudagraph_utils.py
+++ b/vllm/v1/worker/gpu/cudagraph_utils.py
@@ -9,6 +9,10 @@ import torch
 import torch.nn as nn
 from tqdm import tqdm
 
+from vllm.compilation.breakable_cudagraph import (
+    BreakableCUDAGraphWrapper,
+    is_breakable_cudagraph_enabled,
+)
 from vllm.compilation.counter import compilation_counter
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
@@ -116,6 +120,13 @@ class CudaGraphManager:
         )
         self._init_candidates()
 
+        # Breakable CUDA graph (PW CUDA graph without torch.compile)
+        self.use_breakable_cg = (
+            is_breakable_cudagraph_enabled()
+            and self.cudagraph_mode.has_piecewise_cudagraphs()
+        )
+        self.breakable_cg_runner: BreakableCUDAGraphWrapper | None = None
+
     def _init_candidates(self) -> None:
         """Build priority-ordered candidate lists for each token count."""
         capture_sizes = self.compilation_config.cudagraph_capture_sizes
@@ -283,6 +294,20 @@ class CudaGraphManager:
         get_offloader().sync_prev_onload()
         self.graphs[desc].replay()
 
+    def init_breakable_cg_runner(self, model: nn.Module) -> None:
+        if self.breakable_cg_runner is None:
+            self.breakable_cg_runner = BreakableCUDAGraphWrapper(
+                model, self.vllm_config
+            )
+            self.breakable_cg_runner.graph_pool = self.pool
+
+    def run_pw_graph(self, model: nn.Module, model_inputs: dict[str, Any]) -> Any:
+        if not self.use_breakable_cg:
+            # Default: Use torch-compiled piecewise cudagraph.
+            return model(**model_inputs)
+        assert self.breakable_cg_runner is not None
+        return self.breakable_cg_runner(**model_inputs)
+
 
 class ModelCudaGraphManager(CudaGraphManager):
     """CudaGraphManager with model-specific capture and hidden state management."""
@@ -316,6 +341,8 @@ class ModelCudaGraphManager(CudaGraphManager):
     ) -> dict[BatchExecutionDescriptor, CapturedAttentionState]:
         """Capture CUDA graphs for model forward pass."""
         self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs
+        if self.use_breakable_cg:
+            self.init_breakable_cg_runner(model)
 
         def create_forward_fn(
             desc: BatchExecutionDescriptor,
@@ -370,11 +397,16 @@ class ModelCudaGraphManager(CudaGraphManager):
                     slot_mapping=slot_mappings,
                     batch_descriptor=batch_descriptor,
                 ):
-                    model_output = model(**model_inputs)
+                    if cg_mode == CUDAGraphMode.PIECEWISE:
+                        # PIECEWISE graph (compiled PW or breakable, chosen inside
+                        # run_pw_graph).
+                        model_output = self.run_pw_graph(model, model_inputs)
+                    else:
+                        model_output = model(**model_inputs)
 
                 if cg_mode == CUDAGraphMode.PIECEWISE:
-                    # PW CUDA graph internally handles the model outputs.
-                    # No need to keep track of the hidden states.
+                    # PW CUDA graph (compiled or breakable) internally handles the
+                    # model outputs. No need to keep track of the hidden states.
                     return None
 
                 if self.is_last_pp_rank:
diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
index 5cba66e5c9f..519395e90b2 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -1189,7 +1189,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 skip_compiled=skip_compiled,
             ):
                 self.kv_connector.pre_forward(scheduler_output)
-                model_output = self.model(**model_inputs)
+                if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE:
+                    # Run the PIECEWISE graph (compiled PW cudagraph or breakable
+                    # cudagraph, chosen inside run_pw_graph). cg_mode is only
+                    # PIECEWISE after the cudagraph manager exists.
+                    assert self.cudagraph_manager is not None
+                    model_output = self.cudagraph_manager.run_pw_graph(
+                        self.model, model_inputs
+                    )
+                else:
+                    # Eager (NONE): call the raw model directly.
+                    model_output = self.model(**model_inputs)
 
         if self.is_last_pp_rank:
             if self.use_aux_hidden_state_outputs:
diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
index 6ae3fe793bd..d5095add1e0 100644
--- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
@@ -215,12 +215,22 @@ class EagleSpeculator:
                 )
                 inputs_embeds = self.inputs_embeds[:num_tokens]
 
-            ret_hidden_states = self.model(
+            model_inputs = dict(
                 input_ids=self.input_buffers.input_ids[:num_tokens],
                 positions=self.input_buffers.positions[:num_tokens],
                 hidden_states=self.hidden_states[:num_tokens],
                 inputs_embeds=inputs_embeds,
             )
+            if cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE:
+                # Draft prefill with PIECEWISE cudagraph (compiled PW or breakable),
+                # chosen inside run_pw_graph.
+                assert self.prefill_cudagraph_manager is not None
+                ret_hidden_states = self.prefill_cudagraph_manager.run_pw_graph(
+                    self.model, model_inputs
+                )
+            else:
+                # Eager (NONE): call the raw model directly.
+                ret_hidden_states = self.model(**model_inputs)
         if self.method == "mtp":
             last_hidden_states = ret_hidden_states
             hidden_states = ret_hidden_states
@@ -431,6 +441,8 @@ class EagleSpeculator:
         # For PIECEWISE, only the model's compiled regions are captured
         # and the rest (compute_logits, gumbel_sample) runs eagerly.
         assert self.prefill_cudagraph_manager is not None
+        if self.prefill_cudagraph_manager.use_breakable_cg:
+            self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model)
         self.prefill_cudagraph_manager.capture(
             self.prefill,
             attn_states,

From 3fd9d2d35714e80b4cb3fcd3c408a0398fa2525f Mon Sep 17 00:00:00 2001
From: Aakar Dwivedi <82587125+aadwived@users.noreply.github.com>
Date: Sun, 31 May 2026 00:47:21 +0530
Subject: [PATCH 17/35] [CPU][Zen] Route W8A8 and W4A16 linear inference
 through zentorch on AMD Zen CPUs (#41813)

Signed-off-by: R <Ganesh.R@amd.com>
Signed-off-by: Harshal Adhav <harshal.adhav@amd.com>
Signed-off-by: Aakar Dwivedi <aadwived@amd.com>
Co-authored-by: R <Ganesh.R@amd.com>
Co-authored-by: Harshal Adhav <harshal.adhav@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>
---
 setup.py                                      |   4 +-
 .../model_executor/kernels/linear/__init__.py |  11 +-
 .../linear/mixed_precision/__init__.py        |   4 +
 .../linear/mixed_precision/zentorch.py        | 211 ++++++++++++++++++
 .../kernels/linear/scaled_mm/__init__.py      |   4 +
 .../kernels/linear/scaled_mm/zentorch.py      |  98 ++++++++
 .../kernels/linear/zentorch_utils.py          |  23 ++
 7 files changed, 351 insertions(+), 4 deletions(-)
 create mode 100644 vllm/model_executor/kernels/linear/mixed_precision/zentorch.py
 create mode 100644 vllm/model_executor/kernels/linear/scaled_mm/zentorch.py
 create mode 100644 vllm/model_executor/kernels/linear/zentorch_utils.py

diff --git a/setup.py b/setup.py
index 3221c1d1bc0..07374807bee 100644
--- a/setup.py
+++ b/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"],
diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py
index 3a764e15657..37e5b8e1d54 100644
--- a/vllm/model_executor/kernels/linear/__init__.py
+++ b/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",
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py b/vllm/model_executor/kernels/linear/mixed_precision/__init__.py
index b95197db426..c0b8c35bbd5 100644
--- a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py
+++ b/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",
 ]
diff --git a/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py b/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py
new file mode 100644
index 00000000000..c3e8b17cc9a
--- /dev/null
+++ b/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"]
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py
index f8f12f7b0cb..9bd644b6299 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py
+++ b/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",
 ]
diff --git a/vllm/model_executor/kernels/linear/scaled_mm/zentorch.py b/vllm/model_executor/kernels/linear/scaled_mm/zentorch.py
new file mode 100644
index 00000000000..c434c9d465f
--- /dev/null
+++ b/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",
+        )
diff --git a/vllm/model_executor/kernels/linear/zentorch_utils.py b/vllm/model_executor/kernels/linear/zentorch_utils.py
new file mode 100644
index 00000000000..310aed579ef
--- /dev/null
+++ b/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)

From 6bdabbad5bce747865fd3a249658518a4269cc22 Mon Sep 17 00:00:00 2001
From: Jee Jee Li <pandaleefree@gmail.com>
Date: Sun, 31 May 2026 13:16:12 +0800
Subject: [PATCH 18/35] [CI/Build] Enable Step3p7ForConditionalGeneration
 testing (#43956)

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
---
 tests/models/multimodal/processing/test_tensor_schema.py | 1 +
 tests/models/registry.py                                 | 9 ++++++++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py
index 5afcab9f324..12c5071978f 100644
--- a/tests/models/multimodal/processing/test_tensor_schema.py
+++ b/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
diff --git a/tests/models/registry.py b/tests/models/registry.py
index 3ef6997621d..2d119e78d57 100644
--- a/tests/models/registry.py
+++ b/tests/models/registry.py
@@ -1372,7 +1372,14 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         "stepfun-ai/Step3-VL-10B", trust_remote_code=True
     ),
     "Step3p7ForConditionalGeneration": _HfExamplesInfo(
-        "stepfun-ai/Step-3.7-Flash", is_available_online=False, trust_remote_code=True
+        "stepfun-ai/Step-3.7-Flash",
+        trust_remote_code=True,
+        use_original_num_layers=True,
+        # The MoE config lives in the nested ``text_config``, so the overrides
+        # must be nested too. Use 4 layers to initialize at least one MoE layer
+        # and shrink ``moe_num_experts`` (a non-standard key not handled by
+        # ``dummy_hf_overrides``) to avoid OOM during init.
+        hf_overrides={"text_config": {"num_hidden_layers": 4, "moe_num_experts": 8}},
     ),
     "UltravoxModel": _HfExamplesInfo(
         "fixie-ai/ultravox-v0_5-llama-3_2-1b",

From 8b8546da1c3ba65097357523bc24199e36eddf65 Mon Sep 17 00:00:00 2001
From: nightcityblade <jackchen@haloailabs.com>
Date: Mon, 1 Jun 2026 03:28:38 +0800
Subject: [PATCH 19/35] docs: fix MLA attention docstring examples (#44118)

Co-authored-by: nightcityblade <nightcityblade@gmail.com>
---
 .../layers/attention/mla_attention.py              | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py
index 71fd297a7ed..140e071c746 100644
--- a/vllm/model_executor/layers/attention/mla_attention.py
+++ b/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

From f46e6be169909d4bf2c383b1852123a121fe90e2 Mon Sep 17 00:00:00 2001
From: Umut Polat <52835619+umut-polat@users.noreply.github.com>
Date: Mon, 1 Jun 2026 07:04:11 +0300
Subject: [PATCH 20/35] [Misc] Use VLLMValidationError consistently in chat
 completion and completion protocol validators (#36254)

Signed-off-by: umut-polat <52835619+umut-polat@users.noreply.github.com>
---
 ...est_chat_completion_request_validations.py | 65 +++++++++++++++++++
 .../openai/chat_completion/protocol.py        | 49 ++++++++------
 .../entrypoints/openai/completion/protocol.py | 10 +--
 3 files changed, 100 insertions(+), 24 deletions(-)

diff --git a/tests/tool_use/test_chat_completion_request_validations.py b/tests/tool_use/test_chat_completion_request_validations.py
index d832feda7f5..7adf4beb9d8 100644
--- a/tests/tool_use/test_chat_completion_request_validations.py
+++ b/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": ".*",
+                },
+            }
+        )
diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py
index 73ecb3f35a1..0be220fff77 100644
--- a/vllm/entrypoints/openai/chat_completion/protocol.py
+++ b/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
 
diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py
index cb793a41563..a6c3f9c93dc 100644
--- a/vllm/entrypoints/openai/completion/protocol.py
+++ b/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
 

From 4721bb3aa43078167eb893a9ebf9e50565030c1c Mon Sep 17 00:00:00 2001
From: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Date: Mon, 1 Jun 2026 01:00:33 -0400
Subject: [PATCH 21/35] [MRV2] Remove Eagle's dedicated CUDA graph pool
 (#44078)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
---
 .../worker/gpu/spec_decode/eagle/cudagraph.py | 25 ++-----------------
 .../gpu/spec_decode/eagle/speculator.py       |  3 ---
 2 files changed, 2 insertions(+), 26 deletions(-)

diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py
index 43bece01d0e..300a57ec705 100644
--- a/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py
+++ b/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py
@@ -4,7 +4,6 @@ from collections.abc import Callable
 
 import torch
 
-from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu.block_table import BlockTables
@@ -19,27 +18,7 @@ from vllm.v1.worker.gpu.model_states.interface import ModelState
 from vllm.v1.worker.utils import AttentionGroup
 
 
-class EagleCudaGraphManagerBase(CudaGraphManager):
-    """Base CudaGraphManager for Eagle with a dedicated graph pool."""
-
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        device: torch.device,
-        cudagraph_mode: CUDAGraphMode,
-        decode_query_len: int,
-    ):
-        super().__init__(vllm_config, device, cudagraph_mode, decode_query_len)
-
-        # Use a dedicated pool for Eagle to avoid memory overlap with the main
-        # model's cudagraph. The base class uses a shared global pool, but Eagle's
-        # internal allocations (e.g., gumbel_sample temporaries) can conflict with
-        # the main model's allocations when sharing the same pool.
-        if cudagraph_mode:
-            self.pool = torch.cuda.graph_pool_handle()
-
-
-class PrefillEagleCudaGraphManager(EagleCudaGraphManagerBase):
+class PrefillEagleCudaGraphManager(CudaGraphManager):
     """Eagle CudaGraphManager for prefill, using pre-built attention states
     from the target model's capture."""
 
@@ -74,7 +53,7 @@ class PrefillEagleCudaGraphManager(EagleCudaGraphManagerBase):
         super().capture(create_forward_fn, progress_bar_desc)
 
 
-class DecodeEagleCudaGraphManager(EagleCudaGraphManagerBase):
+class DecodeEagleCudaGraphManager(CudaGraphManager):
     """Eagle CudaGraphManager for decode draft generation, building its own
     attention metadata from scratch."""
 
diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
index d5095add1e0..8ca2882e3a9 100644
--- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
@@ -145,9 +145,6 @@ class EagleSpeculator:
             cudagraph_mode,
             decode_query_len=1,
         )
-        # Share a single pool between prefill and decode since they never
-        # execute concurrently.
-        self.decode_cudagraph_manager.pool = self.prefill_cudagraph_manager.pool
 
     def load_model(self, target_model: nn.Module) -> None:
         target_attn_layer_names = get_layers_from_vllm_config(

From 29d69332aa658a96698a456e044763321f4bbd82 Mon Sep 17 00:00:00 2001
From: Jeffrey Wang <jeffreywang@anyscale.com>
Date: Sun, 31 May 2026 22:06:33 -0700
Subject: [PATCH 22/35] [BugFix] Fix `_has_module` to verify native deps via
 trial import (#44035)

Signed-off-by: esmeetu <jasonailu87@gmail.com>
Signed-off-by: Jeffrey Wang <jeffreywang@anyscale.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: esmeetu <jasonailu87@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
---
 tests/utils_/test_import_utils.py | 95 ++++++++++++++++++++++++++++++-
 vllm/utils/import_utils.py        | 20 +++++--
 2 files changed, 110 insertions(+), 5 deletions(-)

diff --git a/tests/utils_/test_import_utils.py b/tests/utils_/test_import_utils.py
index d42685b3fc9..d1f822037ac 100644
--- a/tests/utils_/test_import_utils.py
+++ b/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
diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py
index e97228bfa60..e008e17d806 100644
--- a/vllm/utils/import_utils.py
+++ b/vllm/utils/import_utils.py
@@ -392,12 +392,24 @@ class LazyLoader(ModuleType):
 # Optional dependency detection utilities
 @cache
 def _has_module(module_name: str) -> bool:
-    """Return True if *module_name* can be found in the current environment.
+    """Return True if *module_name* can be imported in the current environment.
 
-    The result is cached so that subsequent queries for the same module incur
-    no additional overhead.
+    Uses ``importlib.util.find_spec`` as a fast pre-check, then performs a
+    trial import to verify that native dependencies (shared libraries, etc.)
+    are also satisfied. Any failure during the trial import is treated as the
+    module being unavailable. The result is cached so that subsequent queries
+    for the same module incur no additional overhead.
     """
-    return importlib.util.find_spec(module_name) is not None
+    try:
+        if importlib.util.find_spec(module_name) is None:
+            return False
+        importlib.import_module(module_name)
+    except ImportError:
+        logger.warning(
+            "Module %s was found but failed to import", module_name, exc_info=True
+        )
+        return False
+    return True
 
 
 def has_deep_ep() -> bool:

From 1fd8bd02a4b45006b42c0920e19e5fefb54be992 Mon Sep 17 00:00:00 2001
From: Isotr0py <mozf@mail2.sysu.edu.cn>
Date: Mon, 1 Jun 2026 14:01:10 +0800
Subject: [PATCH 23/35] [Docs] Replace broken video url in examples (#44159)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
---
 docs/features/multimodal_inputs.md                              | 2 +-
 .../multimodal/openai_chat_completion_client_for_multimodal.py  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md
index f6d4f3f86d8..847743dfff1 100644
--- a/docs/features/multimodal_inputs.md
+++ b/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(
diff --git a/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py b/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py
index 3a007731c74..e1c4fd76d7d 100644
--- a/examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py
+++ b/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

From 98f1279815db9201f3ce437843c002f0b5ec5733 Mon Sep 17 00:00:00 2001
From: wcy <86111164+wcynb1023@users.noreply.github.com>
Date: Mon, 1 Jun 2026 14:56:41 +0800
Subject: [PATCH 24/35] [CPU][RISC-V] Add missing RVV cpu_types helpers for
 WNA16 (#42730)

Signed-off-by: wcy <233313160abc@gmail.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
---
 cmake/cpu_extension.cmake         |  7 +++++
 csrc/cpu/cpu_types_riscv_impl.hpp | 51 +++++++++++++++++++++++++++++++
 csrc/cpu/torch_bindings.cpp       |  2 +-
 3 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake
index ffab4015f49..c51384e3196 100644
--- a/cmake/cpu_extension.cmake
+++ b/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"
diff --git a/csrc/cpu/cpu_types_riscv_impl.hpp b/csrc/cpu/cpu_types_riscv_impl.hpp
index d6cae76c45c..06a38c780a2 100644
--- a/csrc/cpu/cpu_types_riscv_impl.hpp
+++ b/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) {
diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp
index 35350cf247e..29f86ba0eb9 100644
--- a/csrc/cpu/torch_bindings.cpp
+++ b/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 "

From 1f6048abe57511ed789deb8f9db4760546bcf4f5 Mon Sep 17 00:00:00 2001
From: Uranus <109661872+UranusSeven@users.noreply.github.com>
Date: Mon, 1 Jun 2026 15:14:47 +0800
Subject: [PATCH 25/35] fix: glm5.1 pp model loading (#42944)

Signed-off-by: UranusSeven <109661872+UranusSeven@users.noreply.github.com>
---
 vllm/model_executor/models/deepseek_mtp.py | 10 ++++++++--
 vllm/model_executor/models/deepseek_v2.py  | 20 +++++++++++++++++---
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py
index 37f94c687a2..b8987a99872 100644
--- a/vllm/model_executor/models/deepseek_mtp.py
+++ b/vllm/model_executor/models/deepseek_mtp.py
@@ -35,7 +35,7 @@ from .deepseek_v2 import (
     _try_load_fp8_indexer_wk,
     get_spec_layer_idx_from_weight_name,
 )
-from .utils import maybe_prefix
+from .utils import get_pp_missing_layer_names, maybe_prefix
 
 logger = init_logger(__name__)
 
@@ -267,6 +267,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
             ),
         )
 
+        pp_missing_layer_names = get_pp_missing_layer_names(self)
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
         _pending_wk_fp8: dict = {}  # FP8 indexer wk dequant buffer
@@ -282,7 +283,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
             name = self._rewrite_spec_layer_name(spec_layer, name)
 
             if _try_load_fp8_indexer_wk(
-                name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params
+                name,
+                loaded_weight,
+                _pending_wk_fp8,
+                params_dict,
+                loaded_params,
+                pp_missing_layer_names,
             ):
                 continue
 
diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py
index ee78f681c37..e80d00437c7 100644
--- a/vllm/model_executor/models/deepseek_v2.py
+++ b/vllm/model_executor/models/deepseek_v2.py
@@ -105,6 +105,7 @@ from .interfaces import (
 )
 from .utils import (
     PPMissingLayer,
+    get_pp_missing_layer_names,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
     make_layers,
@@ -742,7 +743,9 @@ class Indexer(nn.Module):
         return self.indexer_op(hidden_states, q_fp8, k, weights)
 
 
-def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
+def _try_load_fp8_indexer_wk(
+    name, tensor, buf, params_dict, loaded_params, pp_missing_layer_names
+):
     """
     We fuse the WK and weights_proj projections, but in some checkpoints WK is stored
     in FP8 with a separate weight_scale_inv, while weights_proj is stored in BF16.
@@ -758,6 +761,12 @@ def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
         return False  # WK is not in FP8 format, ignore.
     # Buffer this tensor (weight or scale) until both have arrived.
     layer_prefix = name.rsplit(".wk.", 1)[0]  # e.g. "model.layers.0.self_attn.indexer"
+    fused_name = f"{layer_prefix}.wk_weights_proj.weight"
+    if any(
+        name.startswith(missing_layer_name)
+        for missing_layer_name in pp_missing_layer_names
+    ):
+        return True
     entry = buf.setdefault(layer_prefix, {})
     entry["weight" if is_weight else "scale"] = tensor
     if "weight" not in entry or "scale" not in entry:
@@ -775,7 +784,6 @@ def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
     )
 
     # Load the dequantized weight into shard 0 of the fused buffer.
-    fused_name = f"{layer_prefix}.wk_weights_proj.weight"
     param = params_dict[fused_name]
     param.weight_loader(param, weight_bf16, 0)
     loaded_params.add(fused_name)
@@ -1379,6 +1387,7 @@ class DeepseekV2Model(nn.Module):
             num_redundant_experts=self.num_redundant_experts,
         )
 
+        pp_missing_layer_names = get_pp_missing_layer_names(self)
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
         for name, loaded_weight in weights:
@@ -1394,7 +1403,12 @@ class DeepseekV2Model(nn.Module):
             )
 
             if _try_load_fp8_indexer_wk(
-                name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params
+                name,
+                loaded_weight,
+                _pending_wk_fp8,
+                params_dict,
+                loaded_params,
+                pp_missing_layer_names,
             ):
                 continue
 

From 0910f7e0e1f5ee9f2e5a6f76d0d09e68fddadc01 Mon Sep 17 00:00:00 2001
From: "wang.yuqi" <yuqi.wang@daocloud.io>
Date: Mon, 1 Jun 2026 15:54:59 +0800
Subject: [PATCH 26/35] [Frontend] Resettle generative scoring entrypoint.
 (#44153)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
---
 .buildkite/test-amd.yaml                      |  6 +++-
 .buildkite/test_areas/entrypoints.yaml        |  4 ++-
 .../entrypoints}/generate/__init__.py         |  0
 .../generative_scoring/__init__.py            |  0
 .../test_generative_scoring.py                | 16 ++++-----
 .../test_generative_scoring_e2e.py            |  2 +-
 .../{openai => }/generate/api_router.py       | 12 +++++++
 .../{openai => }/generate/factories.py        |  0
 .../generative_scoring/__init__.py            |  0
 .../generative_scoring/api_router.py          | 36 +++----------------
 .../generative_scoring/serving.py             |  2 +-
 vllm/entrypoints/openai/api_server.py         | 16 ++-------
 vllm/entrypoints/sagemaker/api_router.py      |  2 +-
 13 files changed, 38 insertions(+), 58 deletions(-)
 rename {vllm/entrypoints/openai => tests/entrypoints}/generate/__init__.py (100%)
 rename tests/entrypoints/{openai => generate}/generative_scoring/__init__.py (100%)
 rename tests/entrypoints/{openai => generate}/generative_scoring/test_generative_scoring.py (95%)
 rename tests/entrypoints/{openai => generate}/generative_scoring/test_generative_scoring_e2e.py (99%)
 rename vllm/entrypoints/{openai => }/generate/api_router.py (95%)
 rename vllm/entrypoints/{openai => }/generate/factories.py (100%)
 rename vllm/entrypoints/{openai => generate}/generative_scoring/__init__.py (100%)
 rename vllm/entrypoints/{openai => generate}/generative_scoring/api_router.py (65%)
 rename vllm/entrypoints/{openai => generate}/generative_scoring/serving.py (99%)

diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml
index 50c13e4a89a..a04e88b3d7e 100644
--- a/.buildkite/test-amd.yaml
+++ b/.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
diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml
index 1ae8c79fab7..ebaec9954a3 100644
--- a/.buildkite/test_areas/entrypoints.yaml
+++ b/.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
diff --git a/vllm/entrypoints/openai/generate/__init__.py b/tests/entrypoints/generate/__init__.py
similarity index 100%
rename from vllm/entrypoints/openai/generate/__init__.py
rename to tests/entrypoints/generate/__init__.py
diff --git a/tests/entrypoints/openai/generative_scoring/__init__.py b/tests/entrypoints/generate/generative_scoring/__init__.py
similarity index 100%
rename from tests/entrypoints/openai/generative_scoring/__init__.py
rename to tests/entrypoints/generate/generative_scoring/__init__.py
diff --git a/tests/entrypoints/openai/generative_scoring/test_generative_scoring.py b/tests/entrypoints/generate/generative_scoring/test_generative_scoring.py
similarity index 95%
rename from tests/entrypoints/openai/generative_scoring/test_generative_scoring.py
rename to tests/entrypoints/generate/generative_scoring/test_generative_scoring.py
index 632c4bcc90a..d8008299229 100644
--- a/tests/entrypoints/openai/generative_scoring/test_generative_scoring.py
+++ b/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}
diff --git a/tests/entrypoints/openai/generative_scoring/test_generative_scoring_e2e.py b/tests/entrypoints/generate/generative_scoring/test_generative_scoring_e2e.py
similarity index 99%
rename from tests/entrypoints/openai/generative_scoring/test_generative_scoring_e2e.py
rename to tests/entrypoints/generate/generative_scoring/test_generative_scoring_e2e.py
index 64a59b270f1..4fe8dbe791b 100644
--- a/tests/entrypoints/openai/generative_scoring/test_generative_scoring_e2e.py
+++ b/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"
 
diff --git a/vllm/entrypoints/openai/generate/api_router.py b/vllm/entrypoints/generate/api_router.py
similarity index 95%
rename from vllm/entrypoints/openai/generate/api_router.py
rename to vllm/entrypoints/generate/api_router.py
index 84a7fddeabe..713e2566bc5 100644
--- a/vllm/entrypoints/openai/generate/api_router.py
+++ b/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,
+    )
diff --git a/vllm/entrypoints/openai/generate/factories.py b/vllm/entrypoints/generate/factories.py
similarity index 100%
rename from vllm/entrypoints/openai/generate/factories.py
rename to vllm/entrypoints/generate/factories.py
diff --git a/vllm/entrypoints/openai/generative_scoring/__init__.py b/vllm/entrypoints/generate/generative_scoring/__init__.py
similarity index 100%
rename from vllm/entrypoints/openai/generative_scoring/__init__.py
rename to vllm/entrypoints/generate/generative_scoring/__init__.py
diff --git a/vllm/entrypoints/openai/generative_scoring/api_router.py b/vllm/entrypoints/generate/generative_scoring/api_router.py
similarity index 65%
rename from vllm/entrypoints/openai/generative_scoring/api_router.py
rename to vllm/entrypoints/generate/generative_scoring/api_router.py
index ed0a81d149c..e6918b7f03b 100644
--- a/vllm/entrypoints/openai/generative_scoring/api_router.py
+++ b/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,
-    )
diff --git a/vllm/entrypoints/openai/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py
similarity index 99%
rename from vllm/entrypoints/openai/generative_scoring/serving.py
rename to vllm/entrypoints/generate/generative_scoring/serving.py
index fd8f89cadad..0592d0b29af 100644
--- a/vllm/entrypoints/openai/generative_scoring/serving.py
+++ b/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
diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py
index 461128ed905..5455f1ca427 100644
--- a/vllm/entrypoints/openai/api_server.py
+++ b/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
 
diff --git a/vllm/entrypoints/sagemaker/api_router.py b/vllm/entrypoints/sagemaker/api_router.py
index b3b11cd07b4..00dd7db2818 100644
--- a/vllm/entrypoints/sagemaker/api_router.py
+++ b/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

From de218634194cd5ca4335eb478fbba5246cb54dbf Mon Sep 17 00:00:00 2001
From: "Will.hou" <1205157517@qq.com>
Date: Mon, 1 Jun 2026 16:58:46 +0800
Subject: [PATCH 27/35] [Rust Frontend] Add InternLM2 tool parser (#43481)

Signed-off-by: Will.hou <1205157517@qq.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
---
 rust/src/chat/src/lib.rs                   |   2 +-
 rust/src/chat/src/parser/tool/mod.rs       |  16 +-
 rust/src/chat/src/parser/tool/tests.rs     |  29 ++
 rust/src/tool-parser/src/json/hermes.rs    |   2 +-
 rust/src/tool-parser/src/json/internlm2.rs | 351 +++++++++++++++++++++
 rust/src/tool-parser/src/json/llama.rs     |   2 +-
 rust/src/tool-parser/src/json/mistral.rs   |   2 +-
 rust/src/tool-parser/src/json/mod.rs       |  48 ++-
 rust/src/tool-parser/src/json/qwen.rs      |   2 +-
 rust/src/tool-parser/src/lib.rs            |   5 +-
 10 files changed, 445 insertions(+), 14 deletions(-)
 create mode 100644 rust/src/tool-parser/src/json/internlm2.rs

diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs
index 6c8d85b54dd..0669af8daff 100644
--- a/rust/src/chat/src/lib.rs
+++ b/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]
diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs
index 7dc72672299..ad220b5a787 100644
--- a/rust/src/chat/src/parser/tool/mod.rs
+++ b/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)
diff --git a/rust/src/chat/src/parser/tool/tests.rs b/rust/src/chat/src/parser/tool/tests.rs
index f18aa69f950..65e9f4e075b 100644
--- a/rust/src/chat/src/parser/tool/tests.rs
+++ b/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
+    );
 }
diff --git a/rust/src/tool-parser/src/json/hermes.rs b/rust/src/tool-parser/src/json/hermes.rs
index c8e6a49a7c5..f6b130ec472 100644
--- a/rust/src/tool-parser/src/json/hermes.rs
+++ b/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.
diff --git a/rust/src/tool-parser/src/json/internlm2.rs b/rust/src/tool-parser/src/json/internlm2.rs
new file mode 100644
index 00000000000..5c3b024a936
--- /dev/null
+++ b/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());
+    }
+}
diff --git a/rust/src/tool-parser/src/json/llama.rs b/rust/src/tool-parser/src/json/llama.rs
index 5ad6918b617..36bc8a8347d 100644
--- a/rust/src/tool-parser/src/json/llama.rs
+++ b/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)? {
diff --git a/rust/src/tool-parser/src/json/mistral.rs b/rust/src/tool-parser/src/json/mistral.rs
index a8756bb2cf6..9ca40fcaf97 100644
--- a/rust/src/tool-parser/src/json/mistral.rs
+++ b/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.
diff --git a/rust/src/tool-parser/src/json/mod.rs b/rust/src/tool-parser/src/json/mod.rs
index 4b589b891d4..5102c025c2c 100644
--- a/rust/src/tool-parser/src/json/mod.rs
+++ b/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 {
diff --git a/rust/src/tool-parser/src/json/qwen.rs b/rust/src/tool-parser/src/json/qwen.rs
index f58ca6e0fa6..b8caff0fefd 100644
--- a/rust/src/tool-parser/src/json/qwen.rs
+++ b/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.
diff --git a/rust/src/tool-parser/src/lib.rs b/rust/src/tool-parser/src/lib.rs
index d8f648b81d5..f1dc0455843 100644
--- a/rust/src/tool-parser/src/lib.rs
+++ b/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;

From 8796838910a0c12149084c151a221ed384dea2dd Mon Sep 17 00:00:00 2001
From: zzt <mf1732009@smail.nju.edu.cn>
Date: Mon, 1 Jun 2026 17:42:49 +0800
Subject: [PATCH 28/35] [Bugfix] fix wrong partial_rotary_factor calculation
 for bailing_moe model. (#43770)

Signed-off-by: zzt <zengzetang.zzt@antgroup.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
---
 vllm/model_executor/models/bailing_moe.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py
index 56e119207da..a45d0ca81de 100644
--- a/vllm/model_executor/models/bailing_moe.py
+++ b/vllm/model_executor/models/bailing_moe.py
@@ -130,7 +130,12 @@ class BailingAttention(nn.Module):
             prefix=f"{prefix}.dense",
         )
 
-        rotary_dim = getattr(config, "rotary_dim", self.head_dim)
+        rotary_dim = getattr(config, "rotary_dim", None)
+        if rotary_dim is None:
+            partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
+            rotary_dim = int(self.head_dim * partial_rotary_factor)
+        if rotary_dim is None:
+            rotary_dim = self.head_dim
         config.rope_parameters["partial_rotary_factor"] = rotary_dim / self.head_dim
 
         self.rotary_emb = get_rope(

From bd0aecdc087382c2e3411c24e5d252d8b83cfe25 Mon Sep 17 00:00:00 2001
From: Chaojun Zhang <chaojun.zhang@intel.com>
Date: Mon, 1 Jun 2026 19:21:36 +0800
Subject: [PATCH 29/35] [XPU][CI] Fix test_audio_in_video flake by using
 module-scoped server fixture (#44146)

Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
---
 .../openai/chat_completion/test_audio_in_video.py           | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/entrypoints/openai/chat_completion/test_audio_in_video.py b/tests/entrypoints/openai/chat_completion/test_audio_in_video.py
index 61ee91eab4d..14ac7ce0439 100644
--- a/tests/entrypoints/openai/chat_completion/test_audio_in_video.py
+++ b/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",

From 985c97a6a884f1a29e7584c05e11214f7fb96dbf Mon Sep 17 00:00:00 2001
From: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Date: Mon, 1 Jun 2026 09:05:21 -0400
Subject: [PATCH 30/35] [Perf] Optimize cutlass fp8 scaled mm bypassing
 padding, 20% kernel performance improvement (#43706)

Signed-off-by: yewentao256 <zhyanwentao@126.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
---
 .../kernels/linear/scaled_mm/cutlass.py       | 66 ++++++++++++++++---
 1 file changed, 57 insertions(+), 9 deletions(-)

diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
index 9e65edb851e..b52d2c5b101 100644
--- a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
+++ b/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,
+)

From 023808c23d234387298732ebc942fff5939dbd8b Mon Sep 17 00:00:00 2001
From: Madeesh Kannan <shadeMe@users.noreply.github.com>
Date: Mon, 1 Jun 2026 14:11:35 +0000
Subject: [PATCH 31/35] [Feature] Add support for JetBrains' Mellum v2 code
 generation model (#43992)

Signed-off-by: Madeesh Kannan <madeeswaran.kannan@jetbrains.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
---
 docs/models/supported_models.md             |   1 +
 tests/models/registry.py                    |   1 +
 vllm/model_executor/models/mellum.py        | 253 ++++++++++++++++++++
 vllm/model_executor/models/registry.py      |   1 +
 vllm/transformers_utils/config.py           |   1 +
 vllm/transformers_utils/configs/__init__.py |   2 +
 vllm/transformers_utils/configs/mellum.py   |   7 +
 7 files changed, 266 insertions(+)
 create mode 100644 vllm/model_executor/models/mellum.py
 create mode 100644 vllm/transformers_utils/configs/mellum.py

diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index 0654a59caf1..4612b4c423f 100644
--- a/docs/models/supported_models.md
+++ b/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. | | ✅︎ |
diff --git a/tests/models/registry.py b/tests/models/registry.py
index 2d119e78d57..36e201eac8c 100644
--- a/tests/models/registry.py
+++ b/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"},
diff --git a/vllm/model_executor/models/mellum.py b/vllm/model_executor/models/mellum.py
new file mode 100644
index 00000000000..bdbf0df7fd1
--- /dev/null
+++ b/vllm/model_executor/models/mellum.py
@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+from torch import nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+
+from .qwen3_moe import (
+    Qwen3MoeAttention,
+    Qwen3MoeDecoderLayer,
+    Qwen3MoeForCausalLM,
+    Qwen3MoeMLP,
+    Qwen3MoeModel,
+    Qwen3MoeSparseMoeBlock,
+)
+from .utils import PPMissingLayer, extract_layer_index, maybe_prefix
+
+
+class MellumAttention(Qwen3MoeAttention):
+    """
+    Differences from `Qwen3MoeAttention`:
+    - Supports `per_layer_sliding_window` for `Attention`.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        rope_parameters: dict[str, Any],
+        max_position_embeddings: int = 8192,
+        head_dim: int | None = None,
+        rms_norm_eps: float = 1e-06,
+        qkv_bias: bool = False,
+        cache_config: Any | None = None,
+        quant_config: Any | None = None,
+        prefix: str = "",
+        dual_chunk_attention_config: dict[str, Any] | None = None,
+        per_layer_sliding_window: int | None = None,
+    ) -> None:
+        nn.Module.__init__(self)
+
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = head_dim or (hidden_size // self.total_num_heads)
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.max_position_embeddings = max_position_embeddings
+        self.dual_chunk_attention_config = dual_chunk_attention_config
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size,
+            self.head_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=qkv_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.head_dim,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=max_position_embeddings,
+            rope_parameters=rope_parameters,
+            dual_chunk_attention_config=dual_chunk_attention_config,
+        )
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            per_layer_sliding_window=per_layer_sliding_window,
+            prefix=f"{prefix}.attn",
+            **(
+                {
+                    "layer_idx": extract_layer_index(prefix),
+                    "dual_chunk_attention_config": dual_chunk_attention_config,
+                }
+                if dual_chunk_attention_config
+                else {}
+            ),
+        )
+
+        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+
+
+class MellumDecoderLayer(Qwen3MoeDecoderLayer):
+    """
+    Differences from `Qwen3MoeDecoderLayer`:
+    - Supports interleaved SWA and per-layer RoPE scaling.
+    """
+
+    def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
+        nn.Module.__init__(self)
+
+        config = vllm_config.model_config.hf_text_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+
+        self.hidden_size = config.hidden_size
+        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
+        dual_chunk_attention_config = getattr(
+            config, "dual_chunk_attention_config", None
+        )
+
+        layer_idx = extract_layer_index(prefix)
+        layer_type = config.layer_types[layer_idx]
+        if layer_type == "sliding_attention":
+            sliding_window = getattr(config, "sliding_window", None)
+        else:
+            sliding_window = None
+        rope_parameters = config.rope_parameters[layer_type]
+
+        self.self_attn = MellumAttention(
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=config.num_key_value_heads,
+            rope_parameters=rope_parameters,
+            max_position_embeddings=max_position_embeddings,
+            rms_norm_eps=config.rms_norm_eps,
+            qkv_bias=getattr(config, "attention_bias", False),
+            head_dim=getattr(config, "head_dim", None),
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+            dual_chunk_attention_config=dual_chunk_attention_config,
+            per_layer_sliding_window=sliding_window,
+        )
+
+        if config.mlp_layer_types[layer_idx] == "sparse":
+            self.mlp = Qwen3MoeSparseMoeBlock(
+                vllm_config=vllm_config, prefix=f"{prefix}.mlp"
+            )
+        else:
+            self.mlp = Qwen3MoeMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+
+@support_torch_compile
+class MellumModel(Qwen3MoeModel):
+    """
+    Differences from `Qwen3MoeModel`:
+    - Uses `MellumDecoderLayer`.
+    """
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ):
+        super().__init__(
+            vllm_config=vllm_config,
+            prefix=prefix,
+            decoder_layer_type=MellumDecoderLayer,
+        )
+
+
+class MellumForCausalLM(Qwen3MoeForCausalLM):
+    """
+    Differences from `Qwen3MoeForCausalLM`:
+    - Uses `MellumModel`.
+    """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        nn.Module.__init__(self)
+        config = vllm_config.model_config.hf_text_config
+        quant_config = vllm_config.quant_config
+        self.config = config
+        self.quant_config = quant_config
+        if "dense" in getattr(config, "mlp_layer_types", []):
+            self.packed_modules_mapping["gate_up_proj"] = ["gate_proj", "up_proj"]
+        self.model = MellumModel(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+        self.lm_head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "lm_head"),
+        )
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+
+        self.expert_weights = []
+
+        self.moe_layers = []
+        example_layer = None
+        for layer in self.model.layers:
+            if isinstance(layer, PPMissingLayer):
+                continue
+
+            assert isinstance(layer, Qwen3MoeDecoderLayer)
+            if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock):
+                example_layer = layer.mlp
+                self.moe_layers.append(layer.mlp.experts)
+
+        if example_layer is None:
+            raise RuntimeError("No MoE layer found in the model.layers.")
+
+        self.num_moe_layers = len(self.moe_layers)
+        self.num_expert_groups = 1
+        self.num_shared_experts = 0
+        self.num_logical_experts = example_layer.n_logical_experts
+        self.num_physical_experts = example_layer.n_physical_experts
+        self.num_local_physical_experts = example_layer.n_local_physical_experts
+        self.num_routed_experts = example_layer.n_routed_experts
+        self.num_redundant_experts = example_layer.n_redundant_experts
diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py
index 94472d27e1c..d96ceeb4b50 100644
--- a/vllm/model_executor/models/registry.py
+++ b/vllm/model_executor/models/registry.py
@@ -159,6 +159,7 @@ _TEXT_GENERATION_MODELS = {
     "LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"),
     "MambaForCausalLM": ("mamba", "MambaForCausalLM"),
     "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
+    "MellumForCausalLM": ("mellum", "MellumForCausalLM"),
     "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
     "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
     "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py
index f940739a96c..8339c183c0f 100644
--- a/vllm/transformers_utils/config.py
+++ b/vllm/transformers_utils/config.py
@@ -116,6 +116,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
     RefinedWebModel="RWConfig",  # For tiiuae/falcon-7b(-instruct)
     mlp_speculator="MLPSpeculatorConfig",
     medusa="MedusaConfig",
+    mellum="MellumConfig",
     midashenglm="MiDashengLMConfig",
     moondream3="Moondream3Config",
     eagle="EAGLEConfig",
diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py
index 5998e61dfd8..71f7723e4c8 100644
--- a/vllm/transformers_utils/configs/__init__.py
+++ b/vllm/transformers_utils/configs/__init__.py
@@ -49,6 +49,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
     "LagunaConfig": "vllm.transformers_utils.configs.laguna",
     "Lfm2MoeConfig": "vllm.transformers_utils.configs.lfm2_moe",
     "MedusaConfig": "vllm.transformers_utils.configs.medusa",
+    "MellumConfig": "vllm.transformers_utils.configs.mellum",
     "MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm",
     "MLPSpeculatorConfig": "vllm.transformers_utils.configs.mlp_speculator",
     "Moondream3Config": "vllm.transformers_utils.configs.moondream3",
@@ -117,6 +118,7 @@ __all__ = [
     "LagunaConfig",
     "Lfm2MoeConfig",
     "MedusaConfig",
+    "MellumConfig",
     "MiDashengLMConfig",
     "MLPSpeculatorConfig",
     "Moondream3Config",
diff --git a/vllm/transformers_utils/configs/mellum.py b/vllm/transformers_utils/configs/mellum.py
new file mode 100644
index 00000000000..2bed53394b2
--- /dev/null
+++ b/vllm/transformers_utils/configs/mellum.py
@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from transformers import Qwen3MoeConfig
+
+
+class MellumConfig(Qwen3MoeConfig):
+    model_type = "mellum"

From 035733515f25764cfa828b269cd762d38e4959b9 Mon Sep 17 00:00:00 2001
From: Yongye Zhu <zyy1102000@gmail.com>
Date: Mon, 1 Jun 2026 12:18:32 -0400
Subject: [PATCH 32/35] [Kernel][DSv4] Optimize sparse FP8 compressor kernels
 (#44161)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
---
 .../ops/sparse_attn_compress_cutedsl.py       | 228 +++++++++++-------
 1 file changed, 138 insertions(+), 90 deletions(-)

diff --git a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
index 0eba82126af..1e6c2ed829d 100644
--- a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
+++ b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
@@ -96,9 +96,16 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
         self.quant_block = quant_block
         self.token_stride = token_stride
         self.scale_dim = scale_dim
-        self.num_warps = head_size // quant_block
+        self.elems_per_lane = 8
+        self.copy_elems = 4
+        self.copy_chunks = self.elems_per_lane // self.copy_elems
+        self.lanes_per_group = quant_block // self.elems_per_lane
+        self.groups_per_warp = 32 // self.lanes_per_group
+        self.scale_reduce_steps = self.lanes_per_group.bit_length() - 1
+        self.scale_reduce_offset = self.lanes_per_group // 2
+        self.num_warps = (head_size // quant_block) // self.groups_per_warp
         self.nope_blocks = self.nope_dim // quant_block
-        self.tb_size = head_size // 2
+        self.tb_size = self.num_warps * 32
         self.compress_ratio = compress_ratio
         self.overlap = overlap
         self.window = (1 + int(overlap)) * compress_ratio
@@ -156,8 +163,9 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
         tid, _, _ = cute.arch.thread_idx()
         warp_id = cute.arch.make_warp_uniform(tid // 32)
         lane_id = tid % 32
-        elem0 = tid * 2
-        elem1 = elem0 + 1
+        group_lane = lane_id % self.lanes_per_group
+        group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group
+        elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane
 
         slot_id = slot_mapping[token_idx]
         has_position = token_idx < positions.shape[0]
@@ -201,12 +209,24 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
                     s_block_numbers[row] = block_number_i32
             cute.arch.sync_threads()
 
-            max0 = -Float32.inf
-            max1 = -Float32.inf
-            sum0 = Float32(0.0)
-            sum1 = Float32(0.0)
-            product0 = Float32(0.0)
-            product1 = Float32(0.0)
+            local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
+            local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
+            local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
+
+            for e in cutlass.range_constexpr(self.elems_per_lane):
+                local_max[e] = -Float32.inf
+                local_sum[e] = Float32(0.0)
+                local_product[e] = Float32(0.0)
+
+            cp_f32x4 = cute.make_copy_atom(
+                cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128
+            )
+            copy_layout = cute.make_layout(
+                (self.copy_chunks, self.copy_elems),
+                stride=(self.copy_elems, 1),
+            )
+            kv_vals = cute.make_rmem_tensor(copy_layout, Float32)
+            score_vals = cute.make_rmem_tensor(copy_layout, Float32)
 
             for row in cutlass.range_constexpr(self.window):
                 pos = start + Int64(row)
@@ -215,46 +235,51 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
                     block_offset = pos - block_index * block_size
                     block_number = s_block_numbers[row].to(Int64)
                     head_offset = Int64((row // self.compress_ratio) * self.head_dim)
-                    row_base = (
-                        block_number * state_cache.stride[0]
-                        + block_offset * state_cache.stride[1]
-                        + head_offset
-                    )
+                    row_tensor = state_cache[block_number, block_offset, None]
+                    for chunk in cutlass.range_constexpr(self.copy_chunks):
+                        copy_elem = const_expr(chunk * self.copy_elems)
+                        col_tile = (
+                            head_offset + (elem_base + Int32(copy_elem)).to(Int64)
+                        ) // Int64(self.copy_elems)
+                        kv_src = cute.local_tile(
+                            row_tensor,
+                            tiler=(self.copy_elems,),
+                            coord=(col_tile,),
+                        )
+                        score_src = cute.local_tile(
+                            row_tensor,
+                            tiler=(self.copy_elems,),
+                            coord=(
+                                col_tile + Int64(self.state_width // self.copy_elems),
+                            ),
+                        )
+                        cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None])
+                        cute.copy(cp_f32x4, score_src, score_vals[chunk, None])
 
-                    score0 = state_cache.iterator[
-                        row_base + Int64(self.state_width) + elem0.to(Int64)
-                    ]
-                    kv0 = state_cache.iterator[row_base + elem0.to(Int64)]
-                    new_max0 = cute.arch.fmax(max0, score0)
-                    old_scale0 = cute.math.exp2(
-                        (max0 - new_max0) * Float32(self.rcp_ln2), fastmath=True
-                    )
-                    new_scale0 = cute.math.exp2(
-                        (score0 - new_max0) * Float32(self.rcp_ln2), fastmath=True
-                    )
-                    sum0 = sum0 * old_scale0 + new_scale0
-                    product0 = product0 * old_scale0 + kv0 * new_scale0
-                    max0 = new_max0
+                    for e in cutlass.range_constexpr(self.elems_per_lane):
+                        chunk = const_expr(e // self.copy_elems)
+                        copy_elem = const_expr(e % self.copy_elems)
+                        score = score_vals[chunk, copy_elem]
+                        kv = kv_vals[chunk, copy_elem]
+                        new_max = cute.arch.fmax(local_max[e], score)
+                        old_scale = cute.math.exp2(
+                            (local_max[e] - new_max) * Float32(self.rcp_ln2),
+                            fastmath=True,
+                        )
+                        new_scale = cute.math.exp2(
+                            (score - new_max) * Float32(self.rcp_ln2),
+                            fastmath=True,
+                        )
+                        local_sum[e] = local_sum[e] * old_scale + new_scale
+                        local_product[e] = local_product[e] * old_scale + kv * new_scale
+                        local_max[e] = new_max
 
-                    score1 = state_cache.iterator[
-                        row_base + Int64(self.state_width) + elem1.to(Int64)
-                    ]
-                    kv1 = state_cache.iterator[row_base + elem1.to(Int64)]
-                    new_max1 = cute.arch.fmax(max1, score1)
-                    old_scale1 = cute.math.exp2(
-                        (max1 - new_max1) * Float32(self.rcp_ln2), fastmath=True
-                    )
-                    new_scale1 = cute.math.exp2(
-                        (score1 - new_max1) * Float32(self.rcp_ln2), fastmath=True
-                    )
-                    sum1 = sum1 * old_scale1 + new_scale1
-                    product1 = product1 * old_scale1 + kv1 * new_scale1
-                    max1 = new_max1
+            x = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
+            local_sumsq = Float32(0.0)
+            for e in cutlass.range_constexpr(self.elems_per_lane):
+                x[e] = local_product[e] / local_sum[e]
+                local_sumsq += x[e] * x[e]
 
-            x0 = product0 / sum0
-            x1 = product1 / sum1
-
-            local_sumsq = x0 * x0 + x1 * x1
             warp_sum = local_sumsq
             for step in cutlass.range_constexpr(5):
                 offset = const_expr(16 >> step)
@@ -273,8 +298,9 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
             cute.arch.sync_threads()
 
             rrms = rrms_shared[0]
-            x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32)
-            x1 = x1 * rrms * rms_norm_weight[elem1].to(Float32)
+            for e in cutlass.range_constexpr(self.elems_per_lane):
+                elem = elem_base + e
+                x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32)
 
             k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
             k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
@@ -287,31 +313,53 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
                 + kv_offset * Int64(self.scale_dim)
             )
 
-            if warp_id == self.nope_blocks:
-                pair_idx = lane_id
+            if group_idx == self.nope_blocks:
                 compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
                     self.compress_ratio
                 )
-                cos_v = cos_sin_cache[compressed_pos, pair_idx]
-                sin_v = cos_sin_cache[
-                    compressed_pos, pair_idx + Int32(self.rope_dim // 2)
-                ]
-                real = x0 * cos_v - x1 * sin_v
-                imag = x0 * sin_v + x1 * cos_v
-                packed = _fp32x2_to_bf16x2(real, imag)
-                out_base = value_base + Int64(self.nope_dim) + (lane_id * 4).to(Int64)
-                k_cache_u32.iterator[out_base // Int64(4)] = packed
+                for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
+                    elem = const_expr(pair * 2)
+                    pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair)
+                    cos_v = cos_sin_cache[compressed_pos, pair_idx]
+                    sin_v = cos_sin_cache[
+                        compressed_pos, pair_idx + Int32(self.rope_dim // 2)
+                    ]
+                    real = x[elem] * cos_v - x[elem + 1] * sin_v
+                    imag = x[elem] * sin_v + x[elem + 1] * cos_v
+                    packed = _fp32x2_to_bf16x2(real, imag)
+                    out_base = (
+                        value_base
+                        + Int64(self.nope_dim)
+                        + ((elem_base - self.nope_dim + Int32(elem)) * 2).to(Int64)
+                    )
+                    k_cache_u32.iterator[out_base // Int64(4)] = packed
             else:
-                q_packed = _fp32x2_to_bf16x2(x0, x1)
-                q0, q1 = _bf16x2_to_fp32(q_packed)
-                abs0 = cute.math.absf(q0)
-                abs1 = cute.math.absf(q1)
-                local_absmax = cute.arch.fmax(abs0, abs1)
+                q = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
+                local_absmax = Float32(0.0)
+                for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
+                    elem = const_expr(pair * 2)
+                    q_packed = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
+                    q0, q1 = _bf16x2_to_fp32(q_packed)
+                    q[elem] = q0
+                    q[elem + 1] = q1
+                    local_absmax = cute.arch.fmax(
+                        local_absmax,
+                        cute.arch.fmax(cute.math.absf(q0), cute.math.absf(q1)),
+                    )
                 absmax = local_absmax
-                for step in cutlass.range_constexpr(5):
-                    offset = const_expr(16 >> step)
+                group_mask_and_clamp = const_expr(
+                    (cute.arch.WARP_SIZE - self.lanes_per_group) << 8
+                    | (cute.arch.WARP_SIZE - 1)
+                )
+                for step in cutlass.range_constexpr(self.scale_reduce_steps):
+                    offset = const_expr(self.scale_reduce_offset >> step)
                     absmax = cute.arch.fmax(
-                        absmax, cute.arch.shuffle_sync_bfly(absmax, offset)
+                        absmax,
+                        cute.arch.shuffle_sync_bfly(
+                            absmax,
+                            offset=offset,
+                            mask_and_clamp=group_mask_and_clamp,
+                        ),
                     )
                 scale_raw = cute.arch.fmax(
                     Float32(self.min_scale),
@@ -320,22 +368,22 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
                 bits = _recast_val(scale_raw, Uint32)
                 ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
                 inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32)
-                y0 = cute.arch.fmin(
-                    cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)),
-                    Float32(self.fp8_max),
-                )
-                y1 = cute.arch.fmin(
-                    cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)),
-                    Float32(self.fp8_max),
-                )
-                packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
-                out_base = value_base + (warp_id * self.quant_block + lane_id * 2).to(
-                    Int64
-                )
-                k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
-                if lane_id == 0:
-                    k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(Uint8)
-                    if warp_id == 0:
+                for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
+                    elem = const_expr(pair * 2)
+                    y0 = cute.arch.fmin(
+                        cute.arch.fmax(q[elem] * inv_scale, Float32(-self.fp8_max)),
+                        Float32(self.fp8_max),
+                    )
+                    y1 = cute.arch.fmin(
+                        cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-self.fp8_max)),
+                        Float32(self.fp8_max),
+                    )
+                    packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
+                    out_base = value_base + (elem_base + Int32(elem)).to(Int64)
+                    k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
+                if group_lane == 0:
+                    k_cache.iterator[scale_base + group_idx.to(Int64)] = ue8m0.to(Uint8)
+                    if group_idx == 0:
                         k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8(
                             0
                         )
@@ -462,11 +510,11 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
 
 class SparseAttnCompressKernel:
     head_tile = 64
-    rows_per_warp = 8
+    rows_per_warp = 16
     row_pairs_per_warp = rows_per_warp // 2
     elems_per_lane = 4
     lanes_per_row = head_tile // elems_per_lane
-    num_warps = 16
+    num_warps = 8
     stats_warp_stride = num_warps + 1
     tb_size = num_warps * 32
     rcp_ln2 = 1.4426950408889634
@@ -715,8 +763,8 @@ class SparseAttnCompressKernel:
 
                 local_warp_max = s_max[out_lane, out_elem, final_lane]
                 global_max = local_warp_max
-                for step in cutlass.range_constexpr(4):
-                    offset = const_expr(8 >> step)
+                for step in cutlass.range_constexpr(3):
+                    offset = const_expr(4 >> step)
                     global_max = cute.arch.fmax(
                         global_max,
                         cute.arch.shuffle_sync_bfly(
@@ -732,8 +780,8 @@ class SparseAttnCompressKernel:
                 )
                 global_sum = s_sum[out_lane, out_elem, final_lane] * scale
                 global_product = s_product[out_lane, out_elem, final_lane] * scale
-                for step in cutlass.range_constexpr(4):
-                    offset = const_expr(8 >> step)
+                for step in cutlass.range_constexpr(3):
+                    offset = const_expr(4 >> step)
                     global_sum += cute.arch.shuffle_sync_bfly(
                         global_sum,
                         offset=offset,

From fd9e91d7e4116c9f3d1a3fc237677c925bf9d6d9 Mon Sep 17 00:00:00 2001
From: Andreas Karatzas <akaratza@amd.com>
Date: Mon, 1 Jun 2026 12:40:01 -0500
Subject: [PATCH 33/35] [ROCm][CI] Fix and stabilize EAGLE3 acceptance tests
 (#41294)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Co-authored-by: Micah Williamson <micah.williamson@amd.com>
---
 .../v1/spec_decode/test_acceptance_length.py  | 44 +++++++++++++------
 1 file changed, 30 insertions(+), 14 deletions(-)

diff --git a/tests/v1/spec_decode/test_acceptance_length.py b/tests/v1/spec_decode/test_acceptance_length.py
index 62ff100fdbf..90e3821e2f1 100644
--- a/tests/v1/spec_decode/test_acceptance_length.py
+++ b/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(

From 182c67daf195bf787a32508bafdf3ae56561cc00 Mon Sep 17 00:00:00 2001
From: Xunzhuo <xunzhuo@vllm-semantic-router.ai>
Date: Tue, 2 Jun 2026 03:30:55 +0800
Subject: [PATCH 34/35] [Rust Frontend] Support streaming `generate` endpoint
 (#43779)

Signed-off-by: xunzhuo <xunzhuo@vllm-semantic-router.ai>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
---
 .../server/src/routes/inference/generate.rs   | 222 +++++++++++++-
 .../src/routes/inference/generate/convert.rs  |  42 +++
 .../src/routes/inference/generate/types.rs    |  22 +-
 .../src/routes/inference/generate/validate.rs |  22 +-
 rust/src/server/src/routes/tests.rs           | 270 +++++++++++++++++-
 5 files changed, 563 insertions(+), 15 deletions(-)

diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs
index ff7ea3c6302..b256b7721f9 100644
--- a/rust/src/server/src/routes/inference/generate.rs
+++ b/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
+        );
+    }
+}
diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs
index 70374c2f1eb..844a606bcfd 100644
--- a/rust/src/server/src/routes/inference/generate/convert.rs
+++ b/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);
+    }
 }
diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs
index de7a196c3c6..d4567c44aa6 100644
--- a/rust/src/server/src/routes/inference/generate/types.rs
+++ b/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)]
diff --git a/rust/src/server/src/routes/inference/generate/validate.rs b/rust/src/server/src/routes/inference/generate/validate.rs
index 74a5bbb690a..43347c60b57 100644
--- a/rust/src/server/src/routes/inference/generate/validate.rs
+++ b/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());
     }
 
diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs
index d166b800447..b1a4f0705fd 100644
--- a/rust/src/server/src/routes/tests.rs
+++ b/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)]

From 266b9d9c64ddb64c719d422280d4522382a229d3 Mon Sep 17 00:00:00 2001
From: Siddharth Bedekar <104613085+bedeks@users.noreply.github.com>
Date: Mon, 1 Jun 2026 12:37:30 -0700
Subject: [PATCH 35/35] [Frontend][Core] Add sparse NCCL weight transfer
 support for in-place updates (#40096)

Signed-off-by: Siddharth Bedekar <bedeksid@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
---
 docs/training/weight_transfer/nccl.md         |  21 +-
 examples/rl/rlhf_sparse_nccl.py               | 526 ++++++++++++++++++
 tests/distributed/test_weight_transfer.py     | 333 ++++++++++-
 .../test_weight_transfer_llm.py               |  71 ++-
 tests/v1/worker/test_gpu_model_runner.py      |  69 +++
 .../worker/test_gpu_worker_weight_transfer.py | 155 ++++++
 vllm/distributed/weight_transfer/base.py      |  61 +-
 .../distributed/weight_transfer/ipc_engine.py |  33 +-
 .../weight_transfer/nccl_engine.py            |  71 +++
 vllm/entrypoints/llm.py                       |  13 +-
 vllm/v1/worker/gpu_model_runner.py            |  39 ++
 vllm/v1/worker/gpu_worker.py                  | 106 ++--
 12 files changed, 1423 insertions(+), 75 deletions(-)
 create mode 100644 examples/rl/rlhf_sparse_nccl.py
 create mode 100644 tests/v1/worker/test_gpu_worker_weight_transfer.py

diff --git a/docs/training/weight_transfer/nccl.md b/docs/training/weight_transfer/nccl.md
index bfde1ee2ae3..7b531218568 100644
--- a/docs/training/weight_transfer/nccl.md
+++ b/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
diff --git a/examples/rl/rlhf_sparse_nccl.py b/examples/rl/rlhf_sparse_nccl.py
new file mode 100644
index 00000000000..bddd28b6485
--- /dev/null
+++ b/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()
diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py
index 295e812a124..467e3934a05 100644
--- a/tests/distributed/test_weight_transfer.py
+++ b/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:
diff --git a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py
index 6c626986507..1dd89afcf80 100644
--- a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py
+++ b/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")
 
diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py
index 56811982b91..1a1352249c3 100644
--- a/tests/v1/worker/test_gpu_model_runner.py
+++ b/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"
diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py
new file mode 100644
index 00000000000..dba0f658542
--- /dev/null
+++ b/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
diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py
index 6e99adde1ca..eda209c3f6b 100644
--- a/vllm/distributed/weight_transfer/base.py
+++ b/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")
diff --git a/vllm/distributed/weight_transfer/ipc_engine.py b/vllm/distributed/weight_transfer/ipc_engine.py
index b138c7dd937..a77aab751ff 100644
--- a/vllm/distributed/weight_transfer/ipc_engine.py
+++ b/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)
diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py
index 3b04a5f65ba..674f5b524da 100644
--- a/vllm/distributed/weight_transfer/nccl_engine.py
+++ b/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,
diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py
index c8e9b8c08f0..802d7a6d796 100644
--- a/vllm/entrypoints/llm.py
+++ b/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:
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
index f82d2224a41..9a05c765894 100644
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -49,6 +49,7 @@ from vllm.distributed.parallel_state import (
     is_global_first_rank,
     prepare_communication_buffer_for_model,
 )
+from vllm.distributed.weight_transfer.base import SparseWeightPatch
 from vllm.forward_context import (
     BatchDescriptor,
     set_forward_context,
@@ -3189,6 +3190,44 @@ class GPUModelRunner(
             return self.model.unwrap()
         return self.model
 
+    def apply_sparse_weight_patches(self, patches: Iterable[SparseWeightPatch]) -> None:
+        """Apply sparse flat-index patches directly to existing model params."""
+        model = self.get_model()
+        for patch in patches:
+            param = model.get_parameter(patch.name)
+            if not param.data.is_contiguous():
+                raise NotImplementedError(
+                    "Sparse weight updates currently require contiguous params: "
+                    f"{patch.name}"
+                )
+
+            if patch.indices.dtype != torch.int32:
+                raise ValueError(
+                    "Sparse weight updates currently require int32 indices: "
+                    f"{patch.name}"
+                )
+            if patch.indices.ndim != 1 or patch.values.ndim != 1:
+                raise ValueError(
+                    f"Sparse weight patches must be 1D flattened updates: {patch.name}"
+                )
+            if patch.indices.numel() != patch.values.numel():
+                raise ValueError(
+                    "`indices` and `values` must have matching lengths for "
+                    f"{patch.name}"
+                )
+            if patch.values.dtype != param.dtype:
+                raise ValueError(
+                    f"Sparse values dtype {patch.values.dtype} does not match "
+                    f"parameter dtype {param.dtype} for {patch.name}"
+                )
+
+            flat_param = param.data.view(-1)
+            flat_param.index_copy_(
+                0,
+                patch.indices.to(device=flat_param.device, dtype=torch.long),
+                patch.values.to(device=flat_param.device),
+            )
+
     def get_supported_generation_tasks(self) -> list[GenerationTask]:
         model = self.get_model()
         supported_tasks = list[GenerationTask]()
diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py
index e63f50bc8dc..121fc69f532 100644
--- a/vllm/v1/worker/gpu_worker.py
+++ b/vllm/v1/worker/gpu_worker.py
@@ -992,23 +992,19 @@ class Worker(WorkerBase):
 
     def start_weight_update(self, is_checkpoint_format: bool = True) -> None:
         """
-        Start a new weight update.
-
-        Prepares the model for receiving weights. For checkpoint format,
-        this initializes state for layerwise processing. For kernel format, this is
-        a no-op but must still be called for consistency.
+        Start a new weight update session.
 
         Args:
             is_checkpoint_format: Whether incoming weights are in checkpoint
                 format (need layerwise processing) or kernel format (direct
-                copy). Stored as state for finish_weight_update.
+                copy / sparse patch application).
         """
         self._check_weight_transfer_engine()
 
         if self._weight_update_active:
             raise RuntimeError(
-                "start_weight_update called while a weight update is "
-                "already active. Call finish_weight_update first."
+                "start_weight_update called while a weight update is already "
+                "active. Call finish_weight_update first."
             )
 
         if is_checkpoint_format:
@@ -1020,16 +1016,15 @@ class Worker(WorkerBase):
             with torch.device(self.device):
                 initialize_layerwise_reload(model)
 
-        # Store state so update_weights/finish_weight_update can check
         self._is_checkpoint_format = is_checkpoint_format
         self._weight_update_active = True
 
     def update_weights(self, update_info: dict) -> None:
         """
-        Receive weights from the trainer (one or more chunks).
+        Receive one weight update chunk from the trainer.
 
         start_weight_update must be called before update_weights and
-        finish_weight_update must be called after.
+        finish_weight_update must be called after all chunks have been sent.
 
         Args:
             update_info: Dictionary containing backend-specific update info
@@ -1042,52 +1037,72 @@ class Worker(WorkerBase):
                 "start_weight_update must be called before update_weights."
             )
 
-        # Parse dict into backend-specific typed dataclass
-        typed_update_info = self.weight_transfer_engine.parse_update_info(update_info)
+        update_succeeded = False
+        try:
+            # Parse dict into backend-specific typed dataclass
+            typed_update_info = self.weight_transfer_engine.parse_update_info(
+                update_info
+            )
 
-        model = self.model_runner.model
+            with torch.device(self.device):
+                if self._is_checkpoint_format:
+                    if typed_update_info.update_kind != "dense":
+                        raise ValueError(
+                            "Sparse weight updates require "
+                            "`start_weight_update(is_checkpoint_format=False)`."
+                        )
 
-        with torch.device(self.device):
-            if self._is_checkpoint_format:
-                self.weight_transfer_engine.receive_weights(
-                    typed_update_info,
-                    load_weights=model.load_weights,
-                )
-            else:
-                # Weights are already in kernel format, copy directly
-                def load_weights_direct(
-                    weights: list[tuple[str, torch.Tensor]],
-                ) -> None:
-                    for name, weight in weights:
-                        param = model.get_parameter(name)
-                        param.copy_(weight)
+                    model = self.model_runner.model
 
-                self.weight_transfer_engine.receive_weights(
-                    typed_update_info,
-                    load_weights=load_weights_direct,
-                )
+                    # Use layerwise reload pattern for checkpoint format weights
+                    self.weight_transfer_engine.receive_weights(
+                        typed_update_info,
+                        load_weights=model.load_weights,
+                    )
+                elif typed_update_info.update_kind == "sparse_flat":
+                    if self.parallel_config.world_size != 1:
+                        raise NotImplementedError(
+                            "Sparse weight updates currently require TP=1 and PP=1"
+                        )
+                    self.weight_transfer_engine.receive_sparse_weights(
+                        typed_update_info,
+                        apply_patches=self.model_runner.apply_sparse_weight_patches,
+                    )
+                else:
+                    model = self.model_runner.model
 
-        # NCCL broadcast/packed path are asynchronous.
-        # Sync here so the next step uses the new weights.
-        torch.accelerator.synchronize()
+                    # Weights are already in kernel format, copy directly.
+                    def load_weights_direct(
+                        weights: list[tuple[str, torch.Tensor]],
+                    ) -> None:
+                        for name, weight in weights:
+                            param = model.get_parameter(name)
+                            param.copy_(weight)
+
+                    self.weight_transfer_engine.receive_weights(
+                        typed_update_info,
+                        load_weights=load_weights_direct,
+                    )
+
+            # NCCL broadcast/packed path are asynchronous.
+            # Sync here so the next step uses the new weights.
+            torch.accelerator.synchronize()
+            update_succeeded = True
+        finally:
+            if not update_succeeded:
+                self._weight_update_active = False
+                self._is_checkpoint_format = True
 
     def finish_weight_update(self) -> None:
-        """
-        Finish the current weight update.
-
-        For checkpoint format, this runs layerwise postprocessing.
-        Uses the is_checkpoint_format state stored by start_weight_update.
-        """
+        """Finish the current weight update session."""
         self._check_weight_transfer_engine()
 
         if not self._weight_update_active:
             raise RuntimeError(
-                "start_weight_update must be called before finish_weight_update."
+                "finish_weight_update called without a matching start_weight_update."
             )
 
-        is_checkpoint_format = self._is_checkpoint_format
-
-        if is_checkpoint_format:
+        if self._is_checkpoint_format:
             from vllm.model_executor.model_loader.reload import (
                 finalize_layerwise_reload,
             )
@@ -1096,7 +1111,6 @@ class Worker(WorkerBase):
             with torch.device(self.device):
                 finalize_layerwise_reload(model, self.model_config)
 
-        # Reset state
         self._weight_update_active = False
         self._is_checkpoint_format = True