[Bugfix] Guard mixed-dtype allreduce RMSNorm quant fusions (#48330)

Signed-off-by: hcenteno <hugo.centeno@estudiantat.upc.edu>
(cherry picked from commit 5f8e73cb8b)

cmake/external_projects/fmha_sm100.cmake

@@ -17,7 +17,7 @@ else()
   FetchContent_Declare(
     fmha_sm100
     GIT_REPOSITORY https://github.com/vllm-project/MSA.git
-    GIT_TAG fee783153f3efe57e3e933c5cb7e267a7cebcfb5
+    GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
     GIT_PROGRESS TRUE
     CONFIGURE_COMMAND ""
     BUILD_COMMAND ""

docs/models/supported_models.md

@@ -595,6 +595,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
 | `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B`, `allenai/MolmoWeb-4B`<sup>^</sup>, `allenai/MolmoWeb-8B`<sup>^</sup> | ✅︎ | ✅︎ |
 | `MossAudioModel` | MOSS-Audio | T + A<sup>+</sup> | `OpenMOSS-Team/MOSS-Audio-4B-Instruct`, `OpenMOSS-Team/MOSS-Audio-4B-Thinking`, `OpenMOSS-Team/MOSS-Audio-8B-Instruct`, `OpenMOSS-Team/MOSS-Audio-8B-Thinking` | ✅︎ | ✅︎ |
+| `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | T + A | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ |
 | `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |
 | `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
 | `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + I<sup>E+</sup> | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |
@@ -697,6 +698,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
 | `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ |
 | `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-4.0-1b-speech`, `ibm-granite/granite-speech-3.3-2b`, etc. | ✅︎ | ✅︎ |
 | `GraniteSpeechPlusForConditionalGeneration` | Granite Speech Plus | `ibm-granite/granite-speech-4.1-2b-plus` | ✅︎ | ✅︎ |
+| `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ |
 | `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | ✅︎ | ✅︎ |
 | `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, etc. | | ✅︎ |
 | `VoxtralForConditionalGeneration` | Voxtral (Mistral format) | `mistralai/Voxtral-Mini-3B-2507`, `mistralai/Voxtral-Small-24B-2507`, etc. | ✅︎ | ✅︎ |

requirements/cuda.txt

@@ -9,7 +9,7 @@ torchaudio==2.11.0
 # These must be updated alongside torch
 torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 torchcodec >= 0.14
-PyNvVideoCodec==2.1.0
+PyNvVideoCodec==2.0.4
 # FlashInfer should be updated together with the Dockerfile
 flashinfer-python==0.6.13
 flashinfer-cubin==0.6.13

tests/compile/passes/distributed/test_fusion_all_reduce.py

@@ -222,6 +222,25 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
         ]
 
 
+class TestAllReduceGemmaRMSNormStaticQuantFP8Model(
+    TestAllReduceRMSNormStaticQuantFP8Model
+):
+    def __init__(
+        self,
+        hidden_size=16,
+        token_num=16,
+        eps=1e-6,
+        dtype: torch.dtype = torch.float16,
+    ):
+        super().__init__(hidden_size, token_num, eps, dtype)
+        self.norm = [GemmaRMSNorm(hidden_size, eps) for _ in range(4)]
+        for norm in self.norm:
+            norm.weight.requires_grad_(False)
+
+    def ops_in_model_before(self):
+        return [torch.ops.vllm.all_reduce.default]
+
+
 class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module):
     """Exercises the new ROCm AITER AR+RMS+per-group-FP8-quant patterns.
 
@@ -416,6 +435,15 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
                 reason="Not supported on ROCm platform",
             ),
         ),
+        pytest.param(
+            TestAllReduceGemmaRMSNormStaticQuantFP8Model,
+            True,
+            False,
+            marks=pytest.mark.skipif(
+                current_platform.is_rocm(),
+                reason="Not supported on ROCm platform",
+            ),
+        ),
         pytest.param(
             TestAllReduceRMSNormStaticQuantFP8Model,
             False,
@@ -606,7 +634,10 @@ def all_reduce_fusion_pass_on_test_model(
         )
         backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)
         backend.check_after_ops(model.ops_in_model_after())
-        if test_model_cls is TestAllReduceGemmaRMSNormModel:
+        if test_model_cls in (
+            TestAllReduceGemmaRMSNormModel,
+            TestAllReduceGemmaRMSNormStaticQuantFP8Model,
+        ):
             fused_op = torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default
             fused_nodes = list(find_op_nodes(fused_op, backend.graph_post_pass))
             assert fused_nodes

tests/kernels/attention/test_minimax_m3.py

@@ -150,7 +150,7 @@ def _reference_index_topk(
         num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
         pages = block_table[req_id, :num_blocks]
         k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1)
-        score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale
+        score = sm_scale * torch.einsum("qhd,kd->hqk", q.float(), k.float())
 
         q_pos = prefix_len + torch.arange(q_len, device=idx_q.device)
         k_pos = torch.arange(k.shape[0], device=idx_q.device)
@@ -621,11 +621,12 @@ def test_decode_index_topk_fp8(num_idx_heads: int):
         init_blocks=init_blocks,
         local_blocks=local_blocks,
         num_kv_heads=num_idx_heads,
-        sm_scale=head_dim**-0.5,
         decode_query_len=decode_query_len,
+        max_decode_query_len=decode_query_len,
     )
     # Reference from the DEQUANTIZED fp8 values (the kernel computes the fp8 QK
-    # in fp32, so it must match an fp32 matmul of the same e4m3 values).
+    # in fp32 with no scaling, so it must match an unscaled fp32 matmul of the
+    # same e4m3 values).
     expected = _reference_index_topk(
         idx_q.float(),
         index_kv_cache.float(),
@@ -636,7 +637,6 @@ def test_decode_index_topk_fp8(num_idx_heads: int):
         topk,
         init_blocks,
         local_blocks,
-        head_dim**-0.5,
     )
     _assert_topk_indices_equal_unordered(actual, expected)
 

tests/kernels/mamba/test_cpu_short_conv.py

@@ -60,6 +60,12 @@ def test_short_conv_forward_native_prefill(vllm_config):
         layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix)
 
     layer.to("cpu")
+    # vLLM Linear layers allocate weights with torch.empty (uninitialized).
+    # On ARM these come back as zero-filled pages, so in_proj output is zero and
+    # the prefill state stays zero. Seed + init to make the test platform-safe.
+    torch.manual_seed(0)
+    for p in layer.parameters():
+        torch.nn.init.normal_(p)
     dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False)
     dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False)
 
@@ -117,6 +123,9 @@ def test_short_conv_forward_native_decode(vllm_config):
         layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix)
 
     layer.to("cpu")
+    torch.manual_seed(0)
+    for p in layer.parameters():
+        torch.nn.init.normal_(p)
     dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False)
     dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False)
 

tests/models/registry.py

@@ -1129,6 +1129,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         },
         trust_remote_code=True,
     ),
+    "MossTranscribeDiarizeForConditionalGeneration": _HfExamplesInfo(
+        "OpenMOSS-Team/MOSS-Transcribe-Diarize",
+        trust_remote_code=True,
+        is_available_online=False,
+    ),
     "HfMoondream": _HfExamplesInfo(
         "moondream/moondream3-preview",
         tokenizer="moondream/starmie-v1",

tests/models/transformers/fusers/test_linear.py

@@ -258,7 +258,8 @@ def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser):
         merged.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0))
         if q.bias is not None:
             merged.bias.copy_(torch.cat([q.bias, k.bias, v.bias], dim=0))
-    merged.split_sizes = [q.out_features, k.out_features, v.out_features]
+    merged.output_sizes = [q.out_features, k.out_features, v.out_features]
+    merged.tp_size = 1
     setattr(module, fuser.merged_name, merged)
     for name in (fuser.q_name, fuser.k_name, fuser.v_name):
         delattr(module, name)
@@ -332,7 +333,8 @@ def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
     # original semantics (branches, kwargs, attribute reads)
     code = fuser.fused_forward.__code__
     names = code.co_names
-    assert "qkv_proj" in names and "split_sizes" in names and "o_proj" in names
+    assert "qkv_proj" in names and "output_sizes" in names and "o_proj" in names
+    assert "tp_size" in names
     assert not {"q_proj", "k_proj", "v_proj"} & set(names)
     if attn_cls is FakeAttention:
         assert "update" in names  # the cache branch survives

tests/quantization/test_modelopt.py

@@ -133,6 +133,49 @@ def test_modelopt_mixed_precision_quantizes_parallel_lm_head():
     assert isinstance(method, ModelOptNvFp4LinearMethod)
 
 
+def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
+    from vllm.model_executor.layers.linear import LinearBase
+
+    config = _mixed_precision_config(
+        {
+            "model.layers.0.mlp.gate_proj": {"quant_algo": "NVFP4"},
+            "model.layers.0.mlp.up_proj": {"quant_algo": "NVFP4"},
+        }
+    )
+
+    fake_layer = MagicMock(spec=LinearBase)
+    with patch(
+        "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel"
+    ):
+        method = config.get_quant_method(fake_layer, "model.layers.0.mlp.gate_up_proj")
+
+    assert isinstance(method, ModelOptNvFp4LinearMethod)
+
+
+@pytest.mark.parametrize(
+    ("quantized_prefix", "missing_prefix"),
+    [
+        ("model.layers.0.mlp.gate_proj", "model.layers.0.mlp.down_proj"),
+        ("model.layers.0.self_attn.o_proj", "model.layers.0.self_attn.qkv_proj"),
+    ],
+)
+def test_modelopt_mixed_precision_does_not_infer_missing_sibling_linear(
+    quantized_prefix, missing_prefix
+):
+    from vllm.model_executor.layers.linear import LinearBase
+
+    config = _mixed_precision_config(
+        {
+            quantized_prefix: {"quant_algo": "NVFP4"},
+        }
+    )
+
+    fake_layer = MagicMock(spec=LinearBase)
+    method = config.get_quant_method(fake_layer, missing_prefix)
+
+    assert isinstance(method, UnquantizedLinearMethod)
+
+
 def test_vocab_parallel_embedding_weight_loader_accepts_scalar_scale():
     holder = Mock()
     scale = torch.nn.Parameter(torch.empty(1))

tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

@@ -28,6 +28,7 @@ hybrid_ssm_configs=(
   # GDN (Qwen3.5)
   "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B"
   "VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B"
+  "VLLM_SSM_CONV_STATE_LAYOUT=DS ENFORCE_EAGER=0 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B VLLM_SERVE_EXTRA_ARGS=--spec-method,mtp,--spec-tokens,1"
   # Mamba1 (Jamba)
   "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ai21labs/AI21-Jamba2-3B VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"
 )

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

@@ -76,6 +76,7 @@ DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}
 GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.2}
 PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-128}
 DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128}
+ENFORCE_EAGER=${ENFORCE_EAGER:-1}
 # Comma-separated extra args for vllm serve (e.g. --max-model-len,2048)
 VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
 
@@ -157,11 +158,13 @@ run_tests_for_model() {
     VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
     vllm serve $model_name \
     --port $PORT \
-    --enforce-eager \
     --block-size ${PREFILL_BLOCK_SIZE} \
     --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
     --tensor-parallel-size $PREFILLER_TP_SIZE \
     --kv-transfer-config '$KV_CONFIG_P'"
+    if [[ "$ENFORCE_EAGER" == "1" ]]; then
+      BASE_CMD="${BASE_CMD} --enforce-eager"
+    fi
     if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
       IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
       for arg in "${extra_args[@]}"; do
@@ -206,10 +209,12 @@ run_tests_for_model() {
     VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
     vllm serve $model_name \
     --port $PORT \
-    --enforce-eager \
     --block-size ${DECODE_BLOCK_SIZE} \
     --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
     --kv-transfer-config '$KV_CONFIG_D'"
+    if [[ "$ENFORCE_EAGER" == "1" ]]; then
+      BASE_CMD="${BASE_CMD} --enforce-eager"
+    fi
     if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
       IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS"
       for arg in "${extra_args[@]}"; do

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

@@ -752,7 +752,12 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
             return allreduce[4], allreduce[2]
 
         pm.register_replacement(
-            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
+            pattern,
+            replacement,
+            self.get_inputs(),
+            pm.fwd_only,
+            pm_pass,
+            extra_check=_norm_input_weight_dtype_match,
         )
 
 
@@ -941,7 +946,12 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
             return allreduce[4], allreduce[2], allreduce[5]
 
         pm.register_replacement(
-            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
+            pattern,
+            replacement,
+            self.get_inputs(),
+            pm.fwd_only,
+            pm_pass,
+            extra_check=_norm_input_weight_dtype_match,
         )
 
 

vllm/config/parallel.py

@@ -652,6 +652,7 @@ class ParallelConfig:
             )
             and self.enable_expert_parallel
             and self.tensor_parallel_size > 1
+            and self.data_parallel_size > 1
         )
 
     @property

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

@@ -210,7 +210,7 @@ def persistent_masked_m_silu_mul_quant(
         DeepGemmQuantScaleFMT.UE8M0,
     ]
 
-    device_capability = current_platform.get_device_capability(device_id=y.device.index)
+    device_capability = current_platform.get_device_capability()
     assert device_capability is not None
     cuda_arch = device_capability.to_int()
 

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

@@ -133,6 +133,8 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
             not moe_parallel_config.use_all2all_kernels
             or moe_parallel_config.use_ag_rs_all2all_kernels
             or moe_parallel_config.use_deepep_v2_kernels
+            or moe_parallel_config.use_fi_nvl_one_sided_kernels
+            or moe_parallel_config.use_fi_nvl_two_sided_kernels
         ) and not moe_parallel_config.enable_eplb
 
     @staticmethod

vllm/model_executor/layers/fused_moe/layer.py

@@ -127,6 +127,7 @@ def FusedMoE(
     num_redundant_experts: int = 0,
     has_bias: bool = False,
     is_sequence_parallel: bool = False,
+    reduce_results: bool = True,
     ckpt_names: tuple[str, str, str] = ("gate_proj", "down_proj", "up_proj"),
     n_shared_experts: int | None = None,
     router_logits_dtype: torch.dtype | None = None,
@@ -184,7 +185,11 @@ def FusedMoE(
         num_redundant_experts: Number of redundant experts for EPLB
         has_bias: Whether expert layers have bias terms
         is_sequence_parallel: Whether sequence parallelism is enabled
-        expert_mapping: Expert parameter mapping for weight loading
+        reduce_results: Whether to all-reduce the final output. Setting this
+            to False (to fuse the all-reduce downstream) is only honored on
+            the late-AR path.
+        ckpt_names: Checkpoint parameter name tuple (gate_proj, down_proj,
+            up_proj) used for weight loading
         n_shared_experts: Number of shared experts to fuse into the routed
             grouped GEMM (ROCm; requires aiter FSE or the router-append path)
         router_logits_dtype: Data type for router logits buffers
@@ -220,6 +225,14 @@ def FusedMoE(
         parallel_config=vllm_config.parallel_config,
     )
 
+    # Resolve the deferred all-reduce request against the parallel config.
+    skip_final_all_reduce = (
+        not reduce_results
+        and not moe_parallel_config.use_all2all_kernels
+        and not moe_parallel_config.is_sequence_parallel
+        and zero_expert_type is None
+    )
+
     global_num_experts, logical_num_experts, num_fused_shared_experts = (
         determine_expert_counts(
             num_experts,
@@ -340,6 +353,7 @@ def FusedMoE(
         swiglu_alpha=swiglu_alpha,
         swiglu_beta=swiglu_beta,
         max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size,
+        skip_final_all_reduce=skip_final_all_reduce,
     )
 
     logger.debug("FusedMoEConfig = %s", moe_config)

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -428,7 +428,6 @@ class MoERunner(MoERunnerInterface):
         if (
             shared_output is not None
             and not self.moe_config.is_sequence_parallel
-            and not self.moe_config.skip_final_all_reduce
             and self._fused_output_is_reduced
         ):
             shared_output = tensor_model_parallel_all_reduce(shared_output)
@@ -446,6 +445,13 @@ class MoERunner(MoERunnerInterface):
         here. Skipped when sequence-parallel is active (SP handles its
         own reduction) or when the early path already reduced both outputs.
         """
+        # skip_final_all_reduce must not coexist with a pre-reduced fused
+        # output. This should be enforced by MoE config initialization.
+        if self.moe_config.skip_final_all_reduce:
+            assert not self._fused_output_is_reduced, (
+                "skip_final_all_reduce requires an un-reduced fused output"
+            )
+
         # We don't need to reduce the final output if:
         # - We are not running with TP or DP
         # - The MK already reduced the fused output itself.

vllm/model_executor/layers/linear.py

@@ -1000,16 +1000,12 @@ class QKVParallelLinear(ColumnParallelLinear):
             self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
             self.num_kv_head_replicas = 1
         input_size = self.hidden_size
-        output_size = (
-            self.num_heads * self.head_size
-            + self.num_kv_heads * self.head_size
-            + self.num_kv_heads * self.v_head_size
-        ) * tp_size
         self.output_sizes = [
             self.num_heads * self.head_size * tp_size,  # q_proj
             self.num_kv_heads * self.head_size * tp_size,  # k_proj
             self.num_kv_heads * self.v_head_size * tp_size,  # v_proj
         ]
+        output_size = sum(self.output_sizes)
 
         super().__init__(
             input_size=input_size,

vllm/model_executor/layers/quantization/modelopt.py

@@ -2457,16 +2457,29 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
                 if key.startswith(parent_dot):
                     return info["quant_algo"].upper()
 
-        # 4. Parent-prefix fallback for fused projections (qkv_proj, gate_up_proj).
-        for candidate in self._quantized_layer_prefix_candidates(prefix):
-            parent_dot = candidate.rsplit(".", 1)[0] + "."
-            algos = {
-                info["quant_algo"].upper()
-                for key, info in self.quantized_layers.items()
-                if key.startswith(parent_dot) and "." not in key[len(parent_dot) :]
-            }
-            if len(algos) == 1:
-                return algos.pop()
+        # 4. Parent-prefix fallback for fused projections whose config lists
+        # shard names instead of vLLM's packed module name.
+        fused_projection_shards = {
+            "qkv_proj": ("q_proj", "k_proj", "v_proj"),
+            "gate_up_proj": ("gate_proj", "up_proj"),
+        }
+        shard_names = fused_projection_shards.get(proj_name)
+        if shard_names is not None:
+            for candidate in self._quantized_layer_prefix_candidates(prefix):
+                parent_dot = candidate.rsplit(".", 1)[0] + "."
+                shard_algos: set[str] = set()
+                for shard_name in shard_names:
+                    shard_prefix = f"{parent_dot}{shard_name}"
+                    if shard_prefix in self.quantized_layers:
+                        algo = self.quantized_layers[shard_prefix]["quant_algo"].upper()
+                        shard_algos.add(algo)
+                if len(shard_algos) == 1:
+                    return shard_algos.pop()
+                if len(shard_algos) > 1:
+                    raise ValueError(
+                        f"Mixed quant_algo within fused layer {prefix}: "
+                        f"{shard_algos}. All shards must use the same quantization."
+                    )
 
         return None
 

vllm/model_executor/models/deepseek_v2.py

@@ -279,6 +279,7 @@ class DeepseekV2MoE(nn.Module):
         config: DeepseekV2Config | DeepseekV3Config,
         parallel_config: ParallelConfig,
         quant_config: QuantizationConfig | None = None,
+        reduce_results: bool = True,
         prefix: str = "",
     ):
         super().__init__()
@@ -379,6 +380,7 @@ class DeepseekV2MoE(nn.Module):
             enable_eplb=self.enable_eplb,
             num_redundant_experts=self.n_redundant_experts,
             is_sequence_parallel=self.is_sequence_parallel,
+            reduce_results=reduce_results,
             n_shared_experts=config.n_shared_experts
             if self.is_fusion_moe_shared_experts_enabled
             else None,

vllm/model_executor/models/moss_transcribe_diarize.py

@@ -0,0 +1,801 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Inference-only MOSS-Transcribe-Diarize ASR model.
+
+The checkpoint layout is:
+
+* ``model.whisper_encoder.*``: Whisper-medium encoder weights.
+* ``model.vq_adaptor.*``: 4x time-merge projector.
+* ``model.language_model.*``: Qwen3-0.6B decoder weights.
+"""
+
+import math
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Annotated, Any, Literal, TypeAlias
+
+import torch
+from torch import nn
+from transformers import BatchFeature
+
+from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
+from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.speech_to_text import SpeechToTextParams
+from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TextPrompt
+from vllm.model_executor.models.interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    SupportsPP,
+    SupportsTranscription,
+    _require_is_multimodal,
+)
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    _merge_multimodal_embeddings,
+    init_vllm_registered_model,
+    maybe_prefix,
+)
+from vllm.model_executor.models.whisper import (
+    WhisperEncoder,
+    _create_fake_bias_for_k_proj,
+)
+from vllm.model_executor.models.whisper_utils import ISO639_1_SUPPORTED_LANGS
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import (
+    AudioItem,
+    MultiModalFieldConfig,
+    MultiModalKwargsItems,
+)
+from vllm.multimodal.parse import (
+    DictEmbeddingItems,
+    ModalityDataItems,
+    MultiModalDataItems,
+    MultiModalDataParser,
+)
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    PromptReplacement,
+    PromptUpdate,
+    PromptUpdateDetails,
+)
+from vllm.multimodal.processing.processor import ProcessorInputs
+from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.processor import cached_processor_from_config
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
+
+WHISPER_ENCODER_STRIDE = 2
+
+AUDIO_PLACEHOLDER = "<|audio_start|><|audio_pad|><|audio_end|>"
+
+DEFAULT_MOSS_TRANSCRIBE_DIARIZE_PROMPT = (
+    "请将音频转写为文本，每一段需以起始时间戳和说话人编号"
+    "（[S01]、[S02]、[S03]…）开头，正文为对应的语音内容，"
+    "并在段末标注结束时间戳，以清晰标明该段语音范围。"
+)
+
+
+class MossTranscribeDiarizeAudioInputs(TensorSchema):
+    """
+    Dimensions:
+        - c: Audio chunks
+        - m: Mel bins
+        - f: Mel frames
+        - n: Number of audio items
+    """
+
+    type: Literal["audio_features"] = "audio_features"
+
+    input_features: Annotated[
+        torch.Tensor | None,
+        TensorShape("c", "m", "f"),
+    ]
+    audio_feature_lengths: Annotated[
+        torch.Tensor | None,
+        TensorShape("c"),
+    ]
+    audio_chunk_counts: Annotated[
+        torch.Tensor | None,
+        TensorShape("n"),
+    ]
+
+
+class MossTranscribeDiarizeEmbeddingInputs(TensorSchema):
+    """
+    Dimensions:
+        - n: Number of audio items
+        - t: Number of audio tokens
+        - h: Hidden size
+    """
+
+    type: Literal["audio_embeds"] = "audio_embeds"
+
+    audio_embeds: Annotated[
+        list[torch.Tensor],
+        TensorShape("n", "t", "h", dynamic_dims={"t"}),
+    ]
+
+
+MossTranscribeDiarizeInputs: TypeAlias = (
+    MossTranscribeDiarizeAudioInputs | MossTranscribeDiarizeEmbeddingInputs
+)
+
+
+def _compute_total_audio_tokens(
+    num_samples: int,
+    feature_extractor: Any,
+    audio_merge_size: int,
+) -> int:
+    if num_samples <= 0:
+        return 0
+
+    n_samples = int(feature_extractor.n_samples)
+    stride = (
+        int(feature_extractor.hop_length)
+        * WHISPER_ENCODER_STRIDE
+        * int(audio_merge_size)
+    )
+    total = 0
+    for start in range(0, num_samples, n_samples):
+        chunk_samples = min(n_samples, num_samples - start)
+        total += (chunk_samples - 1) // stride + 1
+    return total
+
+
+def _get_max_audio_samples(feature_extractor: Any) -> int:
+    if hasattr(feature_extractor, "chunk_length"):
+        return int(feature_extractor.chunk_length * feature_extractor.sampling_rate)
+    return int(feature_extractor.n_samples)
+
+
+def _as_audio_embedding_list(audio_embeds: object) -> list[torch.Tensor]:
+    if isinstance(audio_embeds, torch.Tensor):
+        if audio_embeds.ndim == 2:
+            return [audio_embeds]
+        if audio_embeds.ndim == 3:
+            return list(audio_embeds.unbind(dim=0))
+        raise ValueError(
+            f"`audio_embeds` must be a 2D or 3D tensor, got {audio_embeds.ndim}D."
+        )
+
+    if isinstance(audio_embeds, (list, tuple)) and all(
+        isinstance(audio_embed, torch.Tensor) for audio_embed in audio_embeds
+    ):
+        return list(audio_embeds)
+
+    raise TypeError(
+        "`audio_embeds` must be a torch.Tensor or a list of torch.Tensor "
+        f"objects, got {type(audio_embeds)!r}."
+    )
+
+
+def _get_required_token_id(tokenizer: Any, token: str) -> int:
+    token_id = tokenizer.convert_tokens_to_ids(token)
+    if token_id is None or (isinstance(token_id, int) and token_id < 0):
+        raise ValueError(f"Tokenizer is missing required token {token!r}.")
+    return int(token_id)
+
+
+def _get_audios_from_mm_data(mm_data: Mapping[str, object]) -> list[Any]:
+    mm_data_dict = dict(mm_data)
+    audio_data = mm_data_dict.pop("audios", mm_data_dict.pop("audio", []))
+    if isinstance(audio_data, list):
+        audios = audio_data
+    elif isinstance(audio_data, tuple) and not (
+        len(audio_data) == 2 and isinstance(audio_data[1], (int, float))
+    ):
+        audios = list(audio_data)
+    elif audio_data is None:
+        audios = []
+    else:
+        audios = [audio_data]
+
+    audio_arrays: list[Any] = []
+    for audio in audios:
+        if isinstance(audio, (tuple, list)) and len(audio) == 2:
+            audio = audio[0]
+        audio_arrays.append(audio)
+    return audio_arrays
+
+
+def _add_vllm_audio_metadata(
+    processed: BatchFeature,
+    num_audios: int,
+) -> BatchFeature:
+    audio_feature_lengths = processed["audio_feature_lengths"]
+    if not isinstance(audio_feature_lengths, torch.Tensor):
+        audio_feature_lengths = torch.tensor(audio_feature_lengths, dtype=torch.long)
+    audio_feature_lengths = audio_feature_lengths.to(dtype=torch.long)
+
+    audio_chunk_mapping = processed.get("audio_chunk_mapping")
+    if audio_chunk_mapping is None:
+        if num_audios == 1:
+            audio_chunk_mapping = torch.zeros_like(audio_feature_lengths)
+        elif audio_feature_lengths.numel() == num_audios:
+            audio_chunk_mapping = torch.arange(num_audios, dtype=torch.long)
+        else:
+            raise ValueError(
+                "The MOSS processor did not return `audio_chunk_mapping`, and "
+                "the chunk-to-audio mapping cannot be inferred."
+            )
+    elif not isinstance(audio_chunk_mapping, torch.Tensor):
+        audio_chunk_mapping = torch.tensor(audio_chunk_mapping, dtype=torch.long)
+    audio_chunk_mapping = audio_chunk_mapping.to(dtype=torch.long)
+    if audio_chunk_mapping.numel() != audio_feature_lengths.numel():
+        raise ValueError(
+            "`audio_chunk_mapping` must contain one item per audio chunk: got "
+            f"{audio_chunk_mapping.numel()} mappings for "
+            f"{audio_feature_lengths.numel()} chunks."
+        )
+    if audio_chunk_mapping.numel() > 0 and (
+        audio_chunk_mapping.min().item() < 0
+        or audio_chunk_mapping.max().item() >= num_audios
+    ):
+        raise ValueError("`audio_chunk_mapping` contains an out-of-range audio index.")
+
+    audio_chunk_counts = torch.bincount(
+        audio_chunk_mapping.cpu(),
+        minlength=num_audios,
+    ).to(dtype=torch.long)
+    audio_token_lengths = torch.zeros(
+        num_audios,
+        dtype=torch.long,
+        device=audio_feature_lengths.device,
+    )
+    audio_token_lengths.scatter_add_(
+        0,
+        audio_chunk_mapping.to(device=audio_feature_lengths.device),
+        audio_feature_lengths,
+    )
+
+    processed["audio_feature_lengths"] = audio_feature_lengths
+    processed["audio_chunk_counts"] = audio_chunk_counts
+    processed["audio_token_lengths"] = audio_token_lengths
+    return processed
+
+
+class MossTranscribeDiarizeWhisperEncoder(WhisperEncoder):
+    packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."},
+        orig_to_new_stacked={
+            ".self_attn.q_proj": (".self_attn.qkv_proj", "q"),
+            ".self_attn.k_proj": (".self_attn.qkv_proj", "k"),
+            ".self_attn.v_proj": (".self_attn.qkv_proj", "v"),
+        },
+    )
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ) -> None:
+        super().__init__(
+            vllm_config=vllm_config.with_hf_config(
+                vllm_config.model_config.hf_config.audio_config
+            ),
+            prefix=prefix,
+        )
+        self.audio_merge_size = int(vllm_config.model_config.hf_config.audio_merge_size)
+        self.max_encoder_batch: int | None = None
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight")
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+
+    def forward(
+        self,
+        input_features: torch.Tensor,
+        audio_feature_lengths: torch.Tensor,
+    ) -> torch.Tensor:
+        if input_features.numel() == 0:
+            return input_features.new_empty((1, 0, self.conv1.out_channels))
+        device = self.conv1.weight.device
+        dtype = self.conv1.weight.dtype
+        input_features = input_features.to(device=device, dtype=dtype)
+        audio_feature_lengths = audio_feature_lengths.to(device=device)
+        batch_size = self.max_encoder_batch or input_features.shape[0]
+        encoded_parts: list[torch.Tensor] = []
+
+        for start in range(0, input_features.shape[0], batch_size):
+            encoded_parts.append(
+                super().forward([input_features[start : start + batch_size]])
+            )
+
+        hidden = torch.cat(encoded_parts, dim=0)
+        chunks = [
+            hidden[idx : idx + 1, : int(token_len.item()) * self.audio_merge_size]
+            for idx, token_len in enumerate(audio_feature_lengths)
+        ]
+        return torch.cat(chunks, dim=1)
+
+
+class MossTranscribeDiarizeVQAdaptor(nn.Module):
+    def __init__(self, input_dim: int, hidden_size: int, eps: float) -> None:
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Linear(input_dim, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+            nn.LayerNorm(hidden_size, eps=eps, bias=True),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.layers(x)
+
+
+def _mtd_field_config(
+    hf_inputs: Mapping[str, torch.Tensor],
+) -> Mapping[str, MultiModalFieldConfig]:
+    fields: dict[str, MultiModalFieldConfig] = {}
+    if "audio_embeds" in hf_inputs:
+        fields["audio_embeds"] = MultiModalFieldConfig.batched("audio")
+    if "audio_chunk_counts" in hf_inputs:
+        audio_chunk_counts = hf_inputs["audio_chunk_counts"]
+        fields.update(
+            input_features=MultiModalFieldConfig.flat_from_sizes(
+                "audio",
+                audio_chunk_counts,
+            ),
+            audio_feature_lengths=MultiModalFieldConfig.flat_from_sizes(
+                "audio",
+                audio_chunk_counts,
+            ),
+            audio_chunk_counts=MultiModalFieldConfig.batched("audio"),
+            audio_token_lengths=MultiModalFieldConfig.batched("audio"),
+        )
+    return fields
+
+
+class MossTranscribeDiarizeMultiModalDataParser(MultiModalDataParser):
+    def _parse_audio_data(
+        self,
+        data: dict[str, torch.Tensor] | ModalityData[AudioItem],
+    ) -> ModalityDataItems[Any, Any] | None:
+        if isinstance(data, dict):
+            return DictEmbeddingItems(
+                data,
+                modality="audio",
+                required_fields={"audio_embeds"},
+                fields_factory=_mtd_field_config,
+            )
+
+        return super()._parse_audio_data(data)
+
+
+class MossTranscribeDiarizeProcessingInfo(BaseProcessingInfo):
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"audio": 1}
+
+    def get_hf_processor(self, **kwargs: object) -> Any:
+        return self.ctx.get_hf_processor(**kwargs)
+
+    def get_feature_extractor(self, **kwargs: object) -> Any:
+        return self.get_hf_processor(**kwargs).feature_extractor
+
+    def get_data_parser(self) -> MultiModalDataParser:
+        feature_extractor = self.get_feature_extractor()
+        return MossTranscribeDiarizeMultiModalDataParser(
+            target_sr=feature_extractor.sampling_rate,
+            target_channels=1,
+            expected_hidden_size=self._get_expected_hidden_size(),
+        )
+
+    def get_mm_max_tokens_per_item(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> Mapping[str, int]:
+        if mm_counts.get("audio", 0) <= 0:
+            return {}
+
+        feature_extractor = self.get_feature_extractor()
+        max_audio_samples = _get_max_audio_samples(feature_extractor)
+        max_audio_tokens = _compute_total_audio_tokens(
+            max_audio_samples,
+            feature_extractor,
+            self.get_hf_processor().audio_merge_size,
+        )
+        return {"audio": min(seq_len, max_audio_tokens)}
+
+
+class MossTranscribeDiarizeDummyInputsBuilder(
+    BaseDummyInputsBuilder[MossTranscribeDiarizeProcessingInfo]
+):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        return AUDIO_PLACEHOLDER * mm_counts.get("audio", 0)
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions],
+    ) -> MultiModalDataDict:
+        num_audios = mm_counts.get("audio", 0)
+        if num_audios == 0:
+            return {}
+
+        feature_extractor = self.info.get_feature_extractor()
+        return {
+            "audio": self._get_dummy_audios(
+                length=_get_max_audio_samples(feature_extractor),
+                num_audios=num_audios,
+                overrides=mm_options.get("audio"),
+            )
+        }
+
+    def get_dummy_processor_inputs(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions],
+    ) -> ProcessorInputs:
+        dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
+        dummy_mm_items = self.info.parse_mm_data(dummy_mm_data)
+        num_audios = mm_counts.get("audio", 0)
+        tokenizer = self.info.get_tokenizer()
+        prompt = tokenizer.encode(
+            AUDIO_PLACEHOLDER * num_audios,
+            add_special_tokens=False,
+        ) or tokenizer.encode(
+            "\n",
+            add_special_tokens=False,
+        )
+        return ProcessorInputs(prompt=prompt, mm_data_items=dummy_mm_items)
+
+
+class MossTranscribeDiarizeMultiModalProcessor(
+    BaseMultiModalProcessor[MossTranscribeDiarizeProcessingInfo]
+):
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object],
+        tok_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        tokenizer = self.info.get_tokenizer()
+        audios = _get_audios_from_mm_data(mm_data)
+        if not audios:
+            input_ids = tokenizer.encode(
+                prompt,
+                add_special_tokens=tok_kwargs.get("add_special_tokens", False),
+            )
+            return BatchFeature({"input_ids": [input_ids]}, tensor_type="pt")
+
+        processed = self.info.ctx.call_hf_processor(
+            self.info.get_hf_processor(**mm_kwargs),
+            dict(text=prompt, audio=audios),
+            dict(**mm_kwargs, **tok_kwargs),
+        )
+        return _add_vllm_audio_metadata(processed, len(audios))
+
+    def _hf_processor_applies_updates(
+        self,
+        prompt_text: str,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        tokenization_kwargs: Mapping[str, object],
+    ) -> bool:
+        return mm_items.get_count("audio", strict=False) > 0
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        return _mtd_field_config(hf_inputs)
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        out_mm_data = out_mm_kwargs.get_data()
+        audio_token_lengths_tensor = out_mm_data.get("audio_token_lengths")
+        if audio_token_lengths_tensor is None:
+            audio_embeds = out_mm_data.get("audio_embeds")
+            if audio_embeds is None:
+                audio_token_lengths: list[int] = []
+            else:
+                audio_token_lengths = [
+                    int(audio_embed.shape[0])
+                    for audio_embed in _as_audio_embedding_list(audio_embeds)
+                ]
+        else:
+            if not isinstance(audio_token_lengths_tensor, torch.Tensor):
+                raise TypeError(
+                    "`audio_token_lengths` must be a torch.Tensor, got "
+                    f"{type(audio_token_lengths_tensor)!r}."
+                )
+            audio_token_lengths = [
+                int(length) for length in audio_token_lengths_tensor.tolist()
+            ]
+        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+        tokenizer = self.info.get_tokenizer()
+        audio_start_id = _get_required_token_id(tokenizer, processor.audio_start_token)
+        audio_token_id = int(processor.audio_token_id)
+        audio_end_id = _get_required_token_id(tokenizer, processor.audio_end_token)
+
+        def get_num_tokens(item_idx: int) -> int:
+            if item_idx >= len(audio_token_lengths):
+                raise ValueError(
+                    "Cannot determine the number of audio tokens for audio item "
+                    f"{item_idx}."
+                )
+            num_tokens = audio_token_lengths[item_idx]
+            if num_tokens <= 0:
+                raise ValueError("Audio input is too short to produce any tokens.")
+            return num_tokens
+
+        def get_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]:
+            num_tokens = get_num_tokens(item_idx)
+            audio_tokens = processor._audio_span_ids(num_tokens)
+            return PromptUpdateDetails.select_token_id(
+                [audio_start_id] + audio_tokens + [audio_end_id],
+                embed_token_id=audio_token_id,
+            )
+
+        return [
+            PromptReplacement(
+                modality="audio",
+                target=AUDIO_PLACEHOLDER,
+                replacement=get_replacement,
+            ),
+        ]
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    MossTranscribeDiarizeMultiModalProcessor,
+    info=MossTranscribeDiarizeProcessingInfo,
+    dummy_inputs=MossTranscribeDiarizeDummyInputsBuilder,
+)
+class MossTranscribeDiarizeForConditionalGeneration(
+    nn.Module,
+    SupportsMultiModal,
+    SupportsPP,
+    SupportsTranscription,
+):
+    supports_transcription = True
+    supports_transcription_only = True
+    supports_segment_timestamp = False
+    supported_languages = ISO639_1_SUPPORTED_LANGS
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "language_model.layers.": "language_model.model.layers.",
+            "language_model.embed_tokens.": "language_model.model.embed_tokens.",
+            "language_model.norm.": "language_model.model.norm.",
+            "model.language_model.model.": "language_model.model.",
+            "model.language_model.lm_head.": "language_model.lm_head.",
+            "model.language_model.": "language_model.model.",
+            "model.whisper_encoder.": "whisper_encoder.",
+            "model.vq_adaptor.": "vq_adaptor.",
+            "model.": None,
+        },
+    )
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        return AUDIO_PLACEHOLDER if modality.startswith("audio") else None
+
+    @classmethod
+    def get_speech_to_text_config(
+        cls,
+        model_config: ModelConfig,
+        task_type: str,
+    ) -> SpeechToTextConfig:
+        processor = cached_processor_from_config(model_config)
+        return SpeechToTextConfig(
+            max_audio_clip_s=None,
+            sample_rate=processor.feature_extractor.sampling_rate,
+            min_energy_split_window_size=None,
+        )
+
+    @classmethod
+    def get_num_audio_tokens(
+        cls,
+        audio_duration_s: float,
+        stt_config: SpeechToTextConfig,
+        model_config: ModelConfig,
+    ) -> int | None:
+        processor = cached_processor_from_config(model_config)
+        num_samples = math.ceil(audio_duration_s * stt_config.sample_rate)
+        return _compute_total_audio_tokens(
+            num_samples,
+            processor.feature_extractor,
+            processor.audio_merge_size,
+        )
+
+    @classmethod
+    def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType:
+        stt_config = stt_params.stt_config
+        question = stt_params.request_prompt or DEFAULT_MOSS_TRANSCRIBE_DIARIZE_PROMPT
+        question = question.strip() or DEFAULT_MOSS_TRANSCRIBE_DIARIZE_PROMPT
+        prompt = (
+            "<|im_start|>system\n"
+            "You are a helpful assistant.<|im_end|>\n"
+            f"<|im_start|>user\n{AUDIO_PLACEHOLDER}\n"
+            f"{question}<|im_end|>\n"
+            "<|im_start|>assistant\n"
+        )
+        return TextPrompt(
+            prompt=prompt,
+            multi_modal_data={"audio": (stt_params.audio, stt_config.sample_rate)},
+        )
+
+    @classmethod
+    def post_process_output(cls, text: str) -> str:
+        return text.strip()
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        self.config = vllm_config.model_config.hf_config
+        self.dtype = vllm_config.model_config.dtype
+
+        with self._mark_tower_model(vllm_config, "audio"):
+            self.whisper_encoder = MossTranscribeDiarizeWhisperEncoder(
+                vllm_config=vllm_config,
+                prefix=maybe_prefix(prefix, "whisper_encoder"),
+            )
+            self.vq_adaptor = MossTranscribeDiarizeVQAdaptor(
+                input_dim=int(self.config.adaptor_input_dim),
+                hidden_size=int(self.config.text_config.hidden_size),
+                eps=float(self.config.text_config.rms_norm_eps),
+            )
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                hf_config=self.config.text_config,
+                prefix=maybe_prefix(prefix, "language_model"),
+                architectures=["Qwen3ForCausalLM"],
+            )
+
+        self.make_empty_intermediate_tensors = (
+            self.language_model.make_empty_intermediate_tensors
+        )
+
+    def _time_merge(self, features: torch.Tensor) -> torch.Tensor:
+        batch, seq_len, dim = features.shape
+        merge_size = int(self.config.audio_merge_size)
+        seq_len_trim = (seq_len // merge_size) * merge_size
+        return features[:, :seq_len_trim, :].reshape(
+            batch,
+            seq_len_trim // merge_size,
+            dim * merge_size,
+        )
+
+    def _parse_and_validate_audio_input(
+        self,
+        **kwargs: object,
+    ) -> MossTranscribeDiarizeInputs | None:
+        input_features = kwargs.pop("input_features", None)
+        audio_embeds = kwargs.pop("audio_embeds", None)
+        audio_feature_lengths = kwargs.pop("audio_feature_lengths", None)
+        audio_chunk_counts = kwargs.pop("audio_chunk_counts", None)
+        if input_features is None and audio_embeds is None:
+            return None
+        if audio_embeds is not None:
+            return MossTranscribeDiarizeEmbeddingInputs(
+                type="audio_embeds",
+                audio_embeds=_as_audio_embedding_list(audio_embeds),
+            )
+        return MossTranscribeDiarizeAudioInputs(
+            type="audio_features",
+            input_features=input_features,
+            audio_feature_lengths=audio_feature_lengths,
+            audio_chunk_counts=audio_chunk_counts,
+        )
+
+    def _process_audio_input(
+        self, audio_input: MossTranscribeDiarizeInputs
+    ) -> list[torch.Tensor]:
+        if audio_input["type"] == "audio_embeds":
+            return list(audio_input["audio_embeds"])
+
+        input_features = audio_input["input_features"]
+        audio_feature_lengths = audio_input["audio_feature_lengths"]
+        audio_chunk_counts = audio_input["audio_chunk_counts"]
+        if input_features is None or audio_feature_lengths is None:
+            raise ValueError(
+                "MOSS-Transcribe-Diarize audio inputs require both "
+                "`input_features` and `audio_feature_lengths`."
+            )
+        if audio_feature_lengths.numel() != input_features.shape[0]:
+            raise ValueError(
+                "`audio_feature_lengths` must contain one length per "
+                "`input_features` chunk: got "
+                f"{audio_feature_lengths.numel()} lengths for "
+                f"{input_features.shape[0]} chunks."
+            )
+        if audio_chunk_counts is None:
+            audio_chunk_counts = audio_feature_lengths.new_tensor(
+                [input_features.shape[0]],
+                dtype=torch.long,
+            )
+        else:
+            audio_chunk_counts = audio_chunk_counts.to(dtype=torch.long)
+        if audio_chunk_counts.numel() == 0:
+            raise ValueError("`audio_chunk_counts` must contain at least one item.")
+        if torch.any(audio_chunk_counts <= 0):
+            raise ValueError("`audio_chunk_counts` must contain positive counts.")
+        num_audio_chunks = int(audio_chunk_counts.sum().item())
+        if num_audio_chunks != input_features.shape[0]:
+            raise ValueError(
+                "`audio_chunk_counts` must sum to the number of input feature chunks: "
+                f"got {num_audio_chunks} chunks for "
+                f"{input_features.shape[0]} input chunks."
+            )
+
+        features = self.whisper_encoder(input_features, audio_feature_lengths)
+        merged = self._time_merge(features.to(dtype=self.dtype))
+        projected = self.vq_adaptor(merged).squeeze(0)
+
+        audio_chunk_offsets = torch.cumsum(audio_chunk_counts, dim=0)
+        audio_chunk_offsets = torch.cat(
+            [audio_chunk_offsets.new_zeros(1), audio_chunk_offsets]
+        )
+        tokens_per_item = [
+            int(audio_feature_lengths[start:end].sum().item())
+            for start, end in zip(
+                audio_chunk_offsets[:-1].tolist(),
+                audio_chunk_offsets[1:].tolist(),
+            )
+        ]
+        if any(num_tokens <= 0 for num_tokens in tokens_per_item):
+            raise ValueError("Audio input is too short to produce any tokens.")
+        return list(projected.split(tokens_per_item, dim=0))
+
+    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
+        audio_input = self._parse_and_validate_audio_input(**kwargs)
+        if audio_input is None:
+            return []
+        return self._process_audio_input(audio_input)
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        inputs_embeds = self.language_model.embed_input_ids(input_ids)
+        if not multimodal_embeddings:
+            return inputs_embeds
+        return _merge_multimodal_embeddings(
+            inputs_embeds=inputs_embeds,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=_require_is_multimodal(is_multimodal),
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+        return self.language_model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        return self.language_model.compute_logits(hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(
+            weights,
+            mapper=self.hf_to_vllm_mapper,
+        )

vllm/model_executor/models/registry.py

@@ -447,6 +447,10 @@ _MULTIMODAL_MODELS = {
     "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"),
     "KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"),
     "MoonshotKimiaForCausalLM": ("kimi_audio", "KimiAudioForConditionalGeneration"),
+    "MossTranscribeDiarizeForConditionalGeneration": (
+        "moss_transcribe_diarize",
+        "MossTranscribeDiarizeForConditionalGeneration",
+    ),
     "LightOnOCRForConditionalGeneration": (
         "lightonocr",
         "LightOnOCRForConditionalGeneration",

vllm/model_executor/models/transformers/base.py

@@ -78,6 +78,17 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+class ScaledVocabParallelEmbedding(VocabParallelEmbedding):
+    """`VocabParallelEmbedding` that scales its output."""
+
+    def __init__(self, *args, embed_scale: float, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.embed_scale = embed_scale
+
+    def forward(self, input_: torch.Tensor) -> torch.Tensor:
+        return super().forward(input_) * self.embed_scale
+
+
 class Base(
     nn.Module,
     VllmModel,
@@ -156,22 +167,26 @@ class Base(
         self.attention_instances = self.create_attention_instances()
 
         # Input embeddings
-        self.embed_scale = None
         input_embeddings = self.model.get_input_embeddings()
         if not isinstance(input_embeddings, PPMissingLayer):
-            # Some models scale embeddings inside the input embedding layer
-            self.embed_scale = getattr(input_embeddings, "embed_scale", None)
             names = ("embedding_size", "hidden_size")
             embedding_dim = getattr_iter(self.text_config, names, None)
             assert embedding_dim is not None
-            self.model.set_input_embeddings(
-                VocabParallelEmbedding(
-                    self.text_config.vocab_size,
-                    embedding_dim=embedding_dim,
-                    org_num_embeddings=self.text_config.vocab_size,
-                    quant_config=self.quant_config,
-                )
+            embedding_kwargs = dict(
+                num_embeddings=self.text_config.vocab_size,
+                embedding_dim=embedding_dim,
+                org_num_embeddings=self.text_config.vocab_size,
+                quant_config=self.quant_config,
             )
+            embed_scale = getattr(input_embeddings, "embed_scale", None)
+            if embed_scale is not None:
+                # Some models scale embeddings inside the input embedding layer
+                new_input_embeddings = ScaledVocabParallelEmbedding(
+                    **embedding_kwargs, embed_scale=float(embed_scale)
+                )
+            else:
+                new_input_embeddings = VocabParallelEmbedding(**embedding_kwargs)
+            self.model.set_input_embeddings(new_input_embeddings)
 
         # Initialize any parameters that have not had their modules replaced
         self.init_parameters(self.model)
@@ -590,10 +605,7 @@ class Base(
         _init_parameters(module, dtype)
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        inputs_embeds = self.model.get_input_embeddings()(input_ids)
-        if self.embed_scale is not None:
-            inputs_embeds *= self.embed_scale
-        return inputs_embeds
+        return self.model.get_input_embeddings()(input_ids)
 
     def forward(
         self,
@@ -608,16 +620,6 @@ class Base(
             input_ids = None
             inputs_embeds = intermediate_tensors["hidden_states"]
 
-        # If the model scales embeddings inside the input embedding layer we must
-        # ensure they are scaled here since VocabParallelEmbedding will not do it
-        if (
-            self.embed_scale is not None
-            and input_ids is not None
-            and inputs_embeds is None
-        ):
-            inputs_embeds = self.embed_input_ids(input_ids)
-            input_ids = None
-
         # Add batch dimension before entering Transformers model
         if input_ids is not None and input_ids.ndim == 1:
             # [seq_len] -> [1, seq_len]

vllm/model_executor/models/transformers/fusers/qkv.py

@@ -127,15 +127,14 @@ class QKVFuser(StackedFuser):
         if len({id(block) for block, _ in blocks}) != 1:
             raise ValueError("projection calls are in different blocks")
 
-        # q(x), k(x), v(x) -> q, k, v = qkv(x).split(self.qkv.split_sizes, -1)
+        # q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1)
         names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)}
         temps = [f"{name}_fused" for name in (self.q_name, self.k_name, self.v_name)]
         if names & set(temps):
             raise ValueError("fused temporaries would shadow existing names")
         merged = f"self.{self.merged_name}"
-        template = (
-            f"{', '.join(temps)} = {merged}(__arg__).split({merged}.split_sizes, -1)"
-        )
+        sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]"
+        template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)"
         assign = ast.parse(template).body[0]
         arg = next(
             node
@@ -198,13 +197,6 @@ class QKVFuser(StackedFuser):
             self.merged_name,
             merged,
         )
-        # The rewritten forward splits the merged projection into this rank's
-        # shard sizes (see `update_forward`)
-        merged.split_sizes = [
-            merged.num_heads * merged.head_size,
-            merged.num_kv_heads * merged.head_size,
-            merged.num_kv_heads * merged.v_head_size,
-        ]
         setattr(module, self.merged_name, merged)
         # Drop the consumed submodules so their (meta) params are not expected.
         for name in (self.q_name, self.k_name, self.v_name):

vllm/models/deepseek_v32/nvidia/model.py

@@ -72,16 +72,16 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
             and layer_idx >= config.first_k_dense_replace
             and layer_idx % moe_layer_freq == 0
         ):
+            # Defer the MoE cross-rank all-reduce; it is fused into the next
+            # layer's input_layernorm (or the final norm) via
+            # fused_allreduce_rms_norm.
             self.mlp = DeepseekV2MoE(
                 config=config,
                 parallel_config=parallel_config,
                 quant_config=quant_config,
+                reduce_results=False,
                 prefix=f"{prefix}.mlp",
             )
-            # Defer the MoE cross-rank all-reduce; it is fused into the next
-            # layer's input_layernorm (or the final norm) via
-            # fused_allreduce_rms_norm. self.mlp.experts is the MoERunner.
-            self.mlp.experts.moe_config.skip_final_all_reduce = True
         else:
             self.mlp = DeepseekV2MLP(
                 hidden_size=config.hidden_size,

vllm/models/deepseek_v32/nvidia/mtp.py

@@ -89,7 +89,7 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
         hidden_states, residual = self.mtp_block(
             positions=positions, hidden_states=hidden_states, residual=None
         )
-        # mtp_block's MoE output is left un-reduced (skip_final_all_reduce); the
+        # mtp_block's MoE output is left un-reduced (reduce_results=False); the
         # main model fuses that all-reduce into the next norm, but here the
         # recycle hidden is consumed directly, so reduce it now.
         hidden_states = tensor_model_parallel_all_reduce(hidden_states)

vllm/models/deepseek_v4/attention.py

@@ -618,7 +618,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
             dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
             compress_ratio=self.compress_ratio,
             cache_dtype_str=self.kv_cache_dtype,
-            alignment=576 if uses_fp8_ds_mla_layout else None,
+            alignment=576 if uses_fp8_ds_mla_layout else 512,
             model_version="deepseek_v4",
             kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
         )
@@ -648,15 +648,15 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
         # head_dim already carries the fp8 scale padding
         # compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout.
+        uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
         return MLAAttentionSpec(
             block_size=self.cache_config.block_size,
             num_kv_heads=1,
             head_size=self.head_dim,
             dtype=self.dtype,
             compress_ratio=self.compress_ratio,
-            # DeepseekV4 aligns indexer pages to FlashMLA's 576B so they can pack with
-            # the indexer's compressor state cache. V3.2 keeps the legacy layout.
-            alignment=576,
+            # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
+            alignment=576 if uses_fp8_ds_mla_layout else 512,
         )
 
     def forward(self): ...

vllm/models/deepseek_v4/common/ops/cache_utils.py

@@ -654,6 +654,8 @@ def build_flashinfer_mixed_sparse_indices(
     topk: int,
     decode_compressed_indices_are_local: bool = False,
     decode_is_valid_token: torch.Tensor | None = None,
+    swa_block_span: int | None = None,
+    compressed_block_span: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches.
 
@@ -730,6 +732,13 @@ def build_flashinfer_mixed_sparse_indices(
     max_block_size = max(window_block_size, topk_block_size)
     num_warps = 4 if max_block_size >= 256 else 1
 
+    # block_span = page_stride / token_stride; == block_size (no-op) for unpacked KV.
+    swa_span = swa_block_size if swa_block_span is None else swa_block_span
+    compressed_span = (
+        compressed_block_size
+        if compressed_block_span is None
+        else compressed_block_span
+    )
     _build_flashinfer_mixed_sparse_indices_kernel[(num_tokens,)](
         sparse_indices,
         sparse_indices.stride(0),
@@ -748,9 +757,11 @@ def build_flashinfer_mixed_sparse_indices(
         swa_block_table,
         swa_block_table.stride(0),
         swa_block_size,
+        swa_span,
         compressed_block_table,
         compressed_block_table.stride(0),
         compressed_block_size,
+        compressed_span,
         NUM_DECODE_TOKENS=num_decode_tokens,
         WINDOW_SIZE=window_size,
         COMPRESS_RATIO=compress_ratio,
@@ -767,6 +778,18 @@ def build_flashinfer_mixed_sparse_indices(
     return sparse_indices, sparse_topk_lens
 
 
+@triton.jit
+def _remap_flashinfer_index(values, block_size, block_span):
+    # FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so
+    # packed pages (#44577) need block*block_size+off -> block*block_span+off.
+    # TODO: remove once flashinfer-ai/flashinfer#3856 is fixed.
+    is_valid = values >= 0
+    safe_values = tl.where(is_valid, values, 0)
+    values = (safe_values // block_size) * block_span
+    values += safe_values % block_size
+    return tl.where(is_valid, values, -1)
+
+
 @triton.jit(
     do_not_specialize=[
         "sparse_indices_stride",
@@ -775,8 +798,10 @@ def build_flashinfer_mixed_sparse_indices(
         "prefill_topk_stride",
         "swa_block_table_stride",
         "swa_block_size",
+        "swa_block_span",
         "compressed_block_table_stride",
         "compressed_block_size",
+        "compressed_block_span",
         "NUM_DECODE_TOKENS",
         "PREFILL_TOPK_STRIDE",
     ]
@@ -799,9 +824,11 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
     swa_block_table_ptr,
     swa_block_table_stride,
     swa_block_size,
+    swa_block_span,
     compressed_block_table_ptr,
     compressed_block_table_stride,
     compressed_block_size,
+    compressed_block_span,
     NUM_DECODE_TOKENS,
     WINDOW_SIZE: tl.constexpr,
     COMPRESS_RATIO: tl.constexpr,
@@ -825,6 +852,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
                 mask=mask,
                 other=-1,
             )
+            values = _remap_flashinfer_index(values, swa_block_size, swa_block_span)
             tl.store(
                 sparse_indices_ptr + token_idx * sparse_indices_stride + offset,
                 values,
@@ -858,6 +886,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
                 values = block_numbers * compressed_block_size + block_offsets
                 values = tl.where(is_valid, values, -1)
                 compressed_len += tl.sum((is_valid & token_valid).to(tl.int32), axis=0)
+            values = _remap_flashinfer_index(
+                values, compressed_block_size, compressed_block_span
+            )
             tl.store(
                 sparse_indices_ptr
                 + token_idx * sparse_indices_stride
@@ -904,6 +935,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
         block_offsets = pos_offset % swa_block_size
         slot_ids = block_numbers * swa_block_size + block_offsets
         slot_ids = tl.where(offset < swa_len, slot_ids, -1)
+        slot_ids = _remap_flashinfer_index(slot_ids, swa_block_size, swa_block_span)
         tl.store(
             sparse_indices_ptr + token_idx * sparse_indices_stride + offset,
             slot_ids,
@@ -930,6 +962,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
         block_offsets = local_idx % compressed_block_size
         slot_ids = block_numbers * compressed_block_size + block_offsets
         slot_ids = tl.where((offset < topk_len) & is_valid, slot_ids, -1)
+        slot_ids = _remap_flashinfer_index(
+            slot_ids, compressed_block_size, compressed_block_span
+        )
         tl.store(
             sparse_indices_ptr
             + token_idx * sparse_indices_stride

vllm/models/deepseek_v4/compressor.py

@@ -162,7 +162,7 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
             head_size=self.state_dim,
             dtype=self.dtype,
             sliding_window=self.sliding_window,
-            alignment=576 if uses_fp8_ds_mla_layout else None,
+            alignment=576 if uses_fp8_ds_mla_layout else 512,
         )
 
     def forward(self): ...

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

@@ -45,6 +45,21 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor:
     return workspace
 
 
+def _packed_block_span(pool: torch.Tensor) -> int:
+    """Per-block stride of ``pool`` in tokens (``stride(0)//stride(-2)``): ==
+    block_size for unpacked KV, larger when packed (#44577). Raises if not
+    token-aligned."""
+    block_stride = pool.stride(0)
+    token_stride = pool.stride(-2)
+    if block_stride % token_stride != 0:
+        raise NotImplementedError(
+            "FLASHINFER_MLA_SPARSE_DSV4 packed KV requires the per-block stride "
+            f"({block_stride}) to be a multiple of the per-token stride "
+            f"({token_stride}); this layout is not supported yet."
+        )
+    return block_stride // token_stride
+
+
 class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend):
     """FlashInfer backend using the DSv4 sparse metadata/cache layout.
 
@@ -368,6 +383,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
         )
         cached_sparse = swa_metadata.flashinfer_sparse_index_cache.get(cache_key, None)
         if cached_sparse is None:
+            swa_block_span = _packed_block_span(swa_k_cache)
+            compressed_block_span = _packed_block_span(compressed_kv_cache)
             sparse_indices, sparse_topk_lens = build_flashinfer_mixed_sparse_indices(
                 decode_swa_indices,
                 decode_compressed_indices,
@@ -385,6 +402,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
                 top_k,
                 decode_compressed_indices_are_local=decode_compressed_indices_are_local,
                 decode_is_valid_token=decode_is_valid_token,
+                swa_block_span=swa_block_span,
+                compressed_block_span=compressed_block_span,
             )
             if cache_key != "c4a":
                 swa_metadata.flashinfer_sparse_index_cache[cache_key] = (

vllm/models/minimax_m3/common/indexer.py

@@ -256,6 +256,13 @@ class MiniMaxM3IndexerMetadataBuilder(
             dtype=torch.int32,
             device=device,
         )
+        # Stable per-token causal page-count buffer for decode cudagraph replays
+        # (consumed by the MSA top-k path's sparse_topk_select num_valid_pages).
+        self.num_valid_pages_buffer = torch.empty(
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            dtype=torch.int32,
+            device=device,
+        )
 
 
 class MiniMaxM3IndexerTritonMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):

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

@@ -4,6 +4,7 @@
 
 from .index_topk import (
     minimax_m3_index_decode,
+    minimax_m3_index_decode_score,
     minimax_m3_index_score,
     minimax_m3_index_topk,
 )
@@ -11,6 +12,7 @@ from .sparse_attn import minimax_m3_sparse_attn, minimax_m3_sparse_attn_decode
 
 __all__ = [
     "minimax_m3_index_decode",
+    "minimax_m3_index_decode_score",
     "minimax_m3_index_score",
     "minimax_m3_index_topk",
     "minimax_m3_sparse_attn",

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

@@ -757,25 +757,26 @@ def minimax_m3_index_topk(
 
 
 @torch.no_grad()
-def minimax_m3_index_decode(
+def minimax_m3_index_decode_score(
     idx_q: torch.Tensor,  # [total_q, num_idx_heads, head_dim]
     index_kv_cache: torch.Tensor,  # [num_blocks, 128, head_dim]
     block_table: torch.Tensor,  # [num_reqs, max_blocks]
     seq_lens: torch.Tensor,  # [num_reqs] int32
     max_seq_len: int,
-    topk: int,
     init_blocks: int,
     local_blocks: int,
     num_kv_heads: int,
     decode_query_len: int,
     max_decode_query_len: int,
-    out: torch.Tensor | None = None,
+    score_out: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """Decode index block-score + top-k, both split-K (cudagraph-safe).
+    """Decode index block-score (split-K, cudagraph-safe); no top-k.
 
-    Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad).
-    When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into
-    ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating.
+    Returns score [num_kv_heads, total_q, >=max_block] (fp32; init/local blocks
+    forced to 1e30/1e29). When ``score_out`` is given the scores are written into
+    it (read/written by strides, so a transposed view of a unified buffer is
+    accepted) instead of a fresh tensor -- used to share a unified score buffer
+    with the prefill side and run a single top-k over both.
     """
     total_q, num_idx_heads, head_dim = idx_q.shape
     assert num_idx_heads == num_kv_heads, (
@@ -783,7 +784,6 @@ def minimax_m3_index_decode(
     )
     assert decode_query_len <= max_decode_query_len
     assert total_q == seq_lens.shape[0] * decode_query_len
-    batch = total_q
     max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
     use_pdl = current_platform.is_arch_support_pdl()
     # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA
@@ -800,13 +800,16 @@ def minimax_m3_index_decode(
     if num_idx_heads > 1 and max_decode_query_len > 1:
         score_kwargs.update({"num_warps": 4, "num_stages": 2})
 
-    # Keep score strides 16-divisible to avoid Triton recompiles.
-    score_block_stride = round_up(max_block, 16)
-    score = torch.empty(
-        (num_idx_heads, total_q, score_block_stride),
-        dtype=torch.float32,
-        device=idx_q.device,
-    )
+    if score_out is not None:
+        score = score_out
+    else:
+        # Keep score strides 16-divisible to avoid Triton recompiles.
+        score_block_stride = round_up(max_block, 16)
+        score = torch.empty(
+            (num_idx_heads, total_q, score_block_stride),
+            dtype=torch.float32,
+            device=idx_q.device,
+        )
     # split-K over seq blocks; chunk count depends only on shape constants so
     # the grid is fixed within a cuda graph.
     TARGET_GRID = 512
@@ -848,6 +851,55 @@ def minimax_m3_index_decode(
         USE_PDL=use_pdl,
         **score_kwargs,
     )
+    return score
+
+
+@torch.no_grad()
+def minimax_m3_index_decode(
+    idx_q: torch.Tensor,  # [total_q, num_idx_heads, head_dim]
+    index_kv_cache: torch.Tensor,  # [num_blocks, 128, head_dim]
+    block_table: torch.Tensor,  # [num_reqs, max_blocks]
+    seq_lens: torch.Tensor,  # [num_reqs] int32
+    max_seq_len: int,
+    topk: int,
+    init_blocks: int,
+    local_blocks: int,
+    num_kv_heads: int,
+    decode_query_len: int,
+    max_decode_query_len: int,
+    out: torch.Tensor | None = None,
+    score_out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Decode index block-score + top-k, both split-K (cudagraph-safe).
+
+    Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad).
+    When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into
+    ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating.
+    When ``score_out`` ([num_kv_heads, total_q, >=max_block]) is given, the block
+    scores are written into it (read back by the top-k) instead of a fresh
+    tensor -- used to share a unified score buffer with the prefill side. Reads
+    via strides, so a transposed view of a block-major buffer is accepted.
+    """
+    total_q, num_idx_heads, _ = idx_q.shape
+    batch = total_q
+    max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)
+    use_pdl = current_platform.is_arch_support_pdl()
+    pdl_kwargs: dict[str, bool | int] = {}
+    if use_pdl:
+        pdl_kwargs.update({"launch_pdl": True})
+    score = minimax_m3_index_decode_score(
+        idx_q,
+        index_kv_cache,
+        block_table,
+        seq_lens,
+        max_seq_len,
+        init_blocks,
+        local_blocks,
+        num_kv_heads,
+        decode_query_len,
+        max_decode_query_len,
+        score_out=score_out,
+    )
 
     if out is not None:
         topk_idx = out[:, :total_q, :]

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

@@ -43,7 +43,6 @@ _FP8_DTYPES = (
     {
         "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
         "BLOCK_SIZE_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]),
-        "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
         "BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"]
         * triton.next_power_of_2(args["gqa_group_size"]),
     }
@@ -84,7 +83,6 @@ def _gqa_sparse_fwd_kernel(
     BLOCK_SIZE_K: tl.constexpr,  # == SPARSE_BLOCK_SIZE (128)
     BLOCK_SIZE_D: tl.constexpr,
     BLOCK_SIZE_H: tl.constexpr,
-    BLOCK_SIZE_T: tl.constexpr,
     BLOCK_SIZE_QH: tl.constexpr,
     USE_FP8: tl.constexpr,  # fp8 KV cache: dequantize K/V to q.dtype on load
 ):
@@ -109,9 +107,10 @@ def _gqa_sparse_fwd_kernel(
     for j in range(real_q_loop):
         pid_q_j = pid_q * num_q_loop + j
         t_ptr_j = t_ptr + (q_block_start + pid_q_j) * stride_tn + pid_kh * stride_th
-        off_t = tl.arange(0, BLOCK_SIZE_T)
-        topk_idx = tl.load(t_ptr_j + off_t * stride_tk, mask=off_t < max_topk, other=-1)
-        real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0)
+        # Valid block count from seq position (no sentinel): block_size_q == 1.
+        q_abs = prefix_len + pid_q_j * BLOCK_SIZE_Q
+        valid_blocks = (q_abs + BLOCK_SIZE_K) // BLOCK_SIZE_K
+        real_topk = tl.minimum(max_topk, valid_blocks)
         q_ptrs = tl.make_block_ptr(
             base=q_ptr + q_start * stride_qn + pid_h * stride_qh,
             shape=(q_len, gqa_group_size, head_dim),
@@ -202,7 +201,6 @@ def _gqa_sparse_fwd_kernel(
             16, triton.next_power_of_2(args["gqa_group_size"])
         ),
         "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
-        "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
     }
 )
 @triton.jit(do_not_specialize=["decode_query_len"])
@@ -243,7 +241,6 @@ def _gqa_sparse_decode_kernel(
     NUM_TOPK_CHUNKS: tl.constexpr,
     BLOCK_SIZE_H: tl.constexpr,
     BLOCK_SIZE_D: tl.constexpr,
-    BLOCK_SIZE_T: tl.constexpr,
     USE_FP8: tl.constexpr,  # fp8 KV cache: dequantize K/V to q.dtype on load
     USE_PDL: tl.constexpr,
 ):
@@ -268,11 +265,10 @@ def _gqa_sparse_decode_kernel(
     # attention range instead of letting padded rows produce negative lengths.
     kv_len = tl.maximum(query_pos + 1, 0)
 
-    # number of valid (non-padded) selected blocks for this query token
-    off_t = tl.arange(0, BLOCK_SIZE_T)
+    # Valid block count from seq_len (no sentinel): min(topk, cdiv(kv_len, blk)).
     idx_base = t_ptr + pid_kh * stride_th + pid_b * stride_tn
-    topk_idx = tl.load(idx_base + off_t * stride_tk, mask=off_t < max_topk, other=-1)
-    real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0)
+    num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
+    real_topk = tl.minimum(max_topk, num_blocks)
     chunk_end_topk = tl.minimum(chunk_end_compiletime, real_topk)
 
     off_n = tl.arange(0, BLOCK_SIZE_K)

vllm/models/minimax_m3/nvidia/indexer_msa.py

@@ -2,16 +2,21 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """MSA (SM100/Blackwell) indexer impl for MiniMax M3.
 
-Prefill scores with ``fmha_sm100``'s score-only (``OnlyScore``) path then selects
-top-k blocks with the Triton ``minimax_m3_index_topk`` kernel -- fmha is much
-faster than Triton for the wide prefill score (benchmarked ~3-5x).
+Both sides write block scores into one unified token-major buffer
+``[total_q, H, max_k_tiles]``, then a single ``fmha_sm100.sparse_topk_select``
+selects the top-k blocks for the whole batch (decode ``[:nd]`` + prefill
+``[nd:]``) into the shared ``topk_indices_buffer``. It bounds each row by its
+causal page count and force-includes the init/local blocks, so the unwritten
+tail of the buffer is pre-filled with ``-inf``.
 
-Decode uses the Triton fused ``minimax_m3_index_decode`` (the same kernel the
-Triton indexer impl uses): for q_len==1 it is a purpose-built vector x matrix
-score (no wasted tensor-core tiles) with a 256-way split-K and a fused split-K
-top-k, which beats fmha's OnlyScore (wasted MMA on a single query, 64-split cap)
-by ~1.1-3.7x. It is cudagraph-safe by construction (shape-constant split grids)
-and writes the shared ``topk_indices_buffer`` via ``out=``.
+Prefill scores with ``fmha_sm100``'s score-only (``OnlyScore``) path (much faster
+than Triton for the wide prefill score, benchmarked ~3-5x), writing its
+``max_score`` straight into the buffer's prefill region (stride-aware, no copy).
+
+Decode scores with the Triton split-K ``minimax_m3_index_decode_score`` (a
+purpose-built vector x matrix score, no wasted tensor-core tiles, 256-way
+split-K, cudagraph-safe by shape-constant grids), writing into the decode
+region. Its tuning heuristics are kept; only the top-k is shared with prefill.
 
 ``fmha_sm100`` imports are function-local so this module is import-safe on
 AMD / non-SM100.
@@ -22,6 +27,7 @@ from typing import ClassVar
 
 import torch
 
+from vllm.config import VllmConfig
 from vllm.forward_context import get_forward_context
 from vllm.models.minimax_m3.common.indexer import (
     MiniMaxM3IndexerBackend,
@@ -31,8 +37,7 @@ from vllm.models.minimax_m3.common.indexer import (
     MiniMaxM3IndexerMetadataBuilder,
 )
 from vllm.models.minimax_m3.common.ops.index_topk import (
-    minimax_m3_index_decode,
-    minimax_m3_index_topk,
+    minimax_m3_index_decode_score,
 )
 from vllm.v1.attention.backend import (
     AttentionBackend,
@@ -40,10 +45,22 @@ from vllm.v1.attention.backend import (
     CommonAttentionMetadata,
 )
 from vllm.v1.attention.backends.utils import split_decodes_and_prefills
+from vllm.v1.kv_cache_interface import AttentionSpec
 
 # Page size == sparse block size == index-K block; fmha tile id == M3 block id.
 PAGE_SIZE = 128
 
+# Fill for unwritten score tiles: -inf so they never win the top-k (score kernels
+# only write causally-valid blocks).
+_SCORE_SENTINEL = float("-inf")
+
+# Tile (KV-block) dim of the unified score buffer, hardcoded as a cudagraph
+# capture-time constant so the decode score kernel's buffer shape is frozen
+# across replays. 8192 tiles == 1M tokens of context; -inf padding +
+# num_valid_pages bound each row to its causal range, so shorter replays reuse
+# the same buffer safely.
+MAX_K_TILES = 8192
+
 
 class MiniMaxM3IndexerMSABackend(MiniMaxM3IndexerBackend):
     """Indexer side-cache backend selecting the MSA builder."""
@@ -71,6 +88,25 @@ class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata):
     (the base ``prefill`` field is unused on this path)."""
 
     prefill_msa: MiniMaxM3IndexerMSAPrefillMetadata | None = None
+    # Per-forward view (``[:num_tokens]``) of the builder's persistent unified
+    # score buffer ``[total_q, H, MAX_K_TILES]``, shared by decode and prefill
+    # and reused across all layers. Pre-filled with the -inf sentinel in
+    # ``build()``; each layer overwrites its valid tiles before its own top-k.
+    unified_scores: torch.Tensor | None = None
+    # Tile (KV-block) dim of the unified score buffer (== ``MAX_K_TILES``).
+    # Forced as the fmha plan's ``max_k_tiles`` so prefill writes its max_score
+    # straight into the shared buffer.
+    max_k_tiles: int = 0
+    # Batch-wide inputs for the single top-k over the unified buffer (decode +
+    # prefill in one call). ``cu_seqlens_q`` is the per-request query-start
+    # offsets (== query_start_loc, a stable view for cudagraph); ``prefix_lens``
+    # is the per-request context length (== context_lens).
+    topk_cu_seqlens_q: torch.Tensor | None = None
+    topk_prefix_lens: torch.Tensor | None = None
+    topk_max_query_len: int = 0
+    # Per-token causal page count cdiv(seq_pos+1, PAGE_SIZE), [total_q] int32:
+    # drives sparse_topk_select force_end_blocks + -1 out-of-range clamp.
+    topk_num_valid_pages: torch.Tensor | None = None
 
 
 class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
@@ -79,6 +115,28 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
 
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
 
+    def __init__(
+        self,
+        kv_cache_spec: AttentionSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+    ) -> None:
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        # Persistent unified score buffer [T, H, MAX_K_TILES] shared by all
+        # indexer layers and reused across forwards. Stable address (required
+        # for the captured decode path) + fixed tile dim so the decode score
+        # kernel's shape is frozen at capture. Filled with -inf per forward in
+        # build(); the valid/padding partition is a per-forward constant, so
+        # every layer overwrites the same valid tiles before its own top-k.
+        self.unified_scores_buffer = torch.empty(
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            self.num_index_heads,
+            MAX_K_TILES,
+            dtype=torch.float32,
+            device=device,
+        )
+
     def build(
         self,
         common_prefix_len: int,
@@ -107,6 +165,19 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
             common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
         )
 
+        # Per-token causal page count for the top-k, into the stable cg buffer.
+        positions = common_attn_metadata.positions
+        assert positions is not None
+        num_valid_pages = self.num_valid_pages_buffer[:num_tokens]
+        num_valid_pages.copy_(positions[:num_tokens] // PAGE_SIZE + 1)
+
+        # Unified score buffer: a per-forward view of the persistent buffer,
+        # reset to the -inf sentinel once here and shared by every layer (the
+        # tile dim is the capture-time constant MAX_K_TILES).
+        max_k_tiles = MAX_K_TILES
+        unified_scores = self.unified_scores_buffer[:num_tokens]
+        unified_scores.fill_(_SCORE_SENTINEL)
+
         decode: MiniMaxM3IndexerDecodeMetadata | None = None
         if num_decodes > 0:
             qsl_cpu = common_attn_metadata.query_start_loc_cpu
@@ -148,6 +219,11 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
                 causal=True,
                 num_kv_splits=1,
             )
+            # Force the plan's tile dim to the unified buffer's so prefill writes
+            # its max_score straight into unified[:, :, nd:] (the stride-aware
+            # binding shape-matches the tile dim exactly). max_k_tiles >= the
+            # plan's natural value, so the extra tiles are simply never written.
+            plan["max_k_tiles"] = max_k_tiles
             cols = torch.arange(block_table.shape[1], device=block_table.device)
             valid = cols[None, :] < nvp[lo:hi].to(block_table.device)[:, None]
             prefill = MiniMaxM3IndexerMSAPrefillMetadata(
@@ -171,6 +247,12 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
             num_prefill_tokens=num_prefill_tokens,
             decode=decode,
             prefill_msa=prefill,
+            unified_scores=unified_scores,
+            max_k_tiles=max_k_tiles,
+            topk_cu_seqlens_q=query_start_loc[: num_reqs + 1],
+            topk_prefix_lens=context_lens,
+            topk_max_query_len=common_attn_metadata.max_query_len,
+            topk_num_valid_pages=num_valid_pages,
         )
 
 
@@ -195,36 +277,49 @@ class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl):
             -1, self.num_index_heads, self.index_head_dim
         )
         kv = self.index_cache.kv_cache
-        # Both sides write into the single shared persistent topk_indices_buffer:
-        # decode at [:, :nd], prefill at [:, nd:] (each kernel writes [:, :total_q]).
+        # Shared persistent top-k output buffer; the unified top-k below writes
+        # the selected block ids into buf[:num_tokens].
         buf = self.topk_indices_buffer
+        assert buf is not None
 
-        decode_topk: torch.Tensor | None = None
+        # Unified token-major score buffer [total_q, H, MAX_K_TILES]: the tile
+        # dim is innermost/contiguous, so both fmha writes (native [T,H,K]) and
+        # the block-iterating top-k reads hit contiguous tiles. Each side gets a
+        # contiguous slice on dim 0: decode [:nd], prefill [nd:]; the kernels
+        # read/write by strides. The builder allocates it once (persistent,
+        # shared by all layers) and resets it to the sentinel each forward, so
+        # the top-k never picks an unwritten tile.
+        unified_scores = md.unified_scores
+        assert unified_scores is not None
+
+        # Decode scores -> unified[:nd] (transposed [H, nd, MK] view; the kernel
+        # writes by strides). Top-k is deferred to the single unified call below.
         if md.decode is not None:
             d = md.decode
-            decode_topk = minimax_m3_index_decode(
+            minimax_m3_index_decode_score(
                 index_q[:nd],
                 kv,
                 d.block_table,
                 d.seq_lens,
                 d.max_seq_len,
-                self.topk_blocks,
                 self.init_blocks,
                 self.local_blocks,
                 self.num_kv_heads,
                 d.decode_query_len,
                 d.max_decode_query_len,
-                out=buf,
+                score_out=unified_scores[:nd].transpose(0, 1),
             )
 
-        prefill_topk: torch.Tensor | None = None
         if md.prefill_msa is not None:
             from vllm.third_party.fmha_sm100.api import _fmha_sm100
 
             p = md.prefill_msa
             # Index-K cache (num_blocks, 128, D) -> paged MQA (num_blocks,1,128,D).
             k_pages = kv.view(kv.shape[0], 1, PAGE_SIZE, self.index_head_dim)
-            _, max_score = _fmha_sm100(
+            # fmha writes its max_score natively as [nnz_p, H, max_k_tiles] into
+            # the prefill region (stride-aware; plan max_k_tiles forced to
+            # md.max_k_tiles so the shape matches exactly -> no copy).
+            _fmha_sm100(
                 index_q[nd:],
                 k_pages,
                 k_pages,  # V placeholder; not read in OnlyScore
@@ -233,19 +328,24 @@ class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl):
                 output_o=False,
                 output_maxscore=True,
                 sm_scale=self.scale,
-            )
-            # Triton top-k wants [num_index_heads, num_tokens, max_block]; the
-            # transpose is a strided view (the kernel reads via strides).
-            out = buf[:, nd:, :] if buf is not None else None
-            prefill_topk = minimax_m3_index_topk(
-                max_score.transpose(1, 2),
-                p.cu_seqlens_q,
-                p.prefix_lens,
-                p.max_query_len,
-                self.topk_blocks,
-                self.init_blocks,
-                self.local_blocks,
-                out=out,
+                max_score=unified_scores[nd:],
             )
 
-        return decode_topk, prefill_topk
+        # Single top-k over the unified buffer via fmha_sm100 sparse_topk_select
+        # (THK, no transpose) into ``buf``. num_valid_pages drives force_end_blocks
+        # (always keep each token's local block; fmha OnlyScore won't) + the -1
+        # out-of-range clamp.
+        from vllm.third_party.fmha_sm100.api import sparse_topk_select
+
+        sparse_topk_select(
+            unified_scores,
+            self.topk_blocks,
+            num_valid_pages=md.topk_num_valid_pages,
+            force_begin_blocks=self.init_blocks,
+            force_end_blocks=self.local_blocks,
+            output=buf[:num_tokens],
+            max_score_layout="THK",
+        )
+
+        # The attend reads ``buf`` directly; this return is vestigial.
+        return None, None

vllm/models/minimax_m3/nvidia/model.py

@@ -196,6 +196,7 @@ class MiniMaxM3MoE(nn.Module):
         config: PretrainedConfig,
         layer_id: int,
         quant_config: QuantizationConfig | None = None,
+        reduce_results: bool = True,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -262,6 +263,7 @@ class MiniMaxM3MoE(nn.Module):
             router_logits_dtype=self.gate.out_dtype,
             shared_experts=self.shared_experts,
             quant_config=quant_config,
+            reduce_results=reduce_results,
             prefix=f"{prefix}.experts",
         )
 
@@ -659,14 +661,10 @@ class MiniMaxM3DecoderLayer(nn.Module):
         layer_id = int(prefix.split(sep=".")[-1])
         self.layer_id = layer_id
 
-        # Complete the preceding dense MLP's deferred all-reduce
-        # (reduce_results=False), fused into this layer's input_layernorm.
-        # Disable this fusion when PP is set
-        self.fuse_input_allreduce = (
-            layer_id > 0
-            and not _is_moe_layer(config, layer_id - 1)
-            and vllm_config.parallel_config.pipeline_parallel_size == 1
-        )
+        # When set, complete the preceding layer's deferred FFN all-reduce
+        # fused into this layer's input_layernorm.
+        # Configured by MiniMaxM3Model.__init__
+        self.fuse_input_allreduce = False
 
         is_sparse_attention_layer = (
             force_sparse_attn or layer_id in _sparse_attention_layer_ids(config)
@@ -692,12 +690,19 @@ class MiniMaxM3DecoderLayer(nn.Module):
 
         # Dense layers store the FFN under `mlp`; MoE layers under
         # `block_sparse_moe` -- matching the checkpoint's naming.
+        # Leave the FFN output un-reduced so its all-reduce fuses into the
+        # next RMSNorm. MTP blocks add the residual directly and PP sends
+        # hidden states across stages, so both must reduce.
+        reduce_results = (
+            is_mtp_block or vllm_config.parallel_config.pipeline_parallel_size > 1
+        )
         self.is_moe_layer = force_moe or _is_moe_layer(config, layer_id)
         if self.is_moe_layer:
             self.block_sparse_moe = MiniMaxM3MoE(
                 config=config,
                 layer_id=layer_id,
                 quant_config=quant_config,
+                reduce_results=reduce_results,
                 prefix=f"{prefix}.block_sparse_moe",
             )
         else:
@@ -706,7 +711,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
                 intermediate_size=config.dense_intermediate_size,
                 quant_config=quant_config,
                 prefix=f"{prefix}.mlp",
-                reduce_results=vllm_config.parallel_config.pipeline_parallel_size > 1,
+                reduce_results=reduce_results,
             )
 
         # config.use_gemma_norm is True for M3 -> Gemma-style RMSNorm.
@@ -745,6 +750,14 @@ class MiniMaxM3DecoderLayer(nn.Module):
         hidden_states = ffn(hidden_states)
         return hidden_states, residual
 
+    @property
+    def ffn_all_reduce_deferred(self) -> bool:
+        """This layer's FFN output is left un-reduced; the caller fuses the
+        all-reduce into the next RMSNorm."""
+        if self.is_moe_layer:
+            return self.block_sparse_moe.experts.moe_config.skip_final_all_reduce
+        return not self.mlp.down_proj.reduce_results
+
 
 class MiniMaxM3Model(nn.Module, EagleModelMixin):
     fall_back_to_pt_during_load = False
@@ -770,8 +783,9 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
 
         # Reserved top-k indices buffer shared by all sparse-attention indexer
         # layers (mirrors DeepseekV4); kept at a stable address so the indexer's
-        # top-k output survives cudagraph capture/replay. Shape matches the
-        # per-head index top-k output [num_index_heads, total_q, topk].
+        # top-k output survives cudagraph capture/replay. Token-major
+        # [total_q, num_index_heads, topk] so the indexer writes its native
+        # [token, head, topk] top-k; the attend transposes to [H, tokens, topk].
         sparse_cfg = getattr(config, "sparse_attention_config", None)
         if sparse_cfg is not None:
             tp_size = get_tensor_model_parallel_world_size()
@@ -781,8 +795,8 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
             max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
             padded_num_tokens = (max_num_batched_tokens + 3) // 4 * 4
             self.topk_indices_buffer = torch.empty(
-                num_index_heads,
                 padded_num_tokens,
+                num_index_heads,
                 sparse_cfg["sparse_topk_blocks"],
                 dtype=torch.int32,
             )
@@ -807,6 +821,15 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
             ["hidden_states", "residual"], config.hidden_size
         )
 
+        # Configure cross-layer all-reduce/RMSNorm fusion: a layer whose FFN output
+        # is left un-reduced has that all-reduce fused into the next layer's
+        # input_layernorm (or the final norm).
+        prev_defers = False
+        for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]):
+            layer.fuse_input_allreduce = idx > 0 and prev_defers
+            prev_defers = layer.ffn_all_reduce_deferred
+        self.fuse_final_norm_allreduce = prev_defers
+
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids)
 
@@ -841,7 +864,12 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
                 {"hidden_states": hidden_states, "residual": residual}
             )
 
-        hidden_states, _ = self.norm(hidden_states, residual)
+        if self.fuse_final_norm_allreduce:
+            hidden_states, _ = fused_allreduce_gemma_rms_norm(
+                hidden_states, residual, self.norm
+            )
+        else:
+            hidden_states, _ = self.norm(hidden_states, residual)
 
         if len(aux_hidden_states) > 0:
             return hidden_states, aux_hidden_states

vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

@@ -39,7 +39,8 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
 
         nd = main_md.num_decode_tokens
         num_tokens = main_md.num_actual_tokens
-        # Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:].
+        # Indexer top-k from the shared token-major buffer [total_q, H, MK]; the
+        # kernels want [H, tokens, MK], so slice tokens on dim 0 then transpose.
         topk = layer.topk_indices_buffer  # type: ignore[attr-defined]
         assert topk is not None
         hd = self.head_size
@@ -56,7 +57,7 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
             minimax_m3_sparse_attn_decode(
                 q[:nd],
                 kv_cache,
-                topk[:, :nd, :],
+                topk[:nd].transpose(0, 1),
                 d.block_table,
                 d.seq_lens,
                 self.num_kv_heads,
@@ -74,7 +75,9 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
 
             p = main_md.prefill
             assert p is not None
-            prefill_topk = topk[:, nd:num_tokens, :]
+            # [H, prefill, MK] transposed view; build_k2q_csr consumes the
+            # strided view directly (topK stays innermost-contiguous).
+            prefill_topk = topk[nd:num_tokens].transpose(0, 1)
             qp = q[nd:]
             k_cache = kv_cache[:, 0].transpose(1, 2)
             v_cache = kv_cache[:, 1].transpose(1, 2)

vllm/multimodal/video.py

@@ -15,7 +15,7 @@ import torch
 
 from vllm import envs
 from vllm.logger import init_logger
-from vllm.utils.import_utils import PlaceholderModule
+from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available
 from vllm.utils.mem_constants import MiB_bytes
 from vllm.utils.registry import ExtensionManager
 
@@ -33,7 +33,7 @@ except ImportError:
 
 try:
     from torchcodec.decoders import VideoDecoder
-except ImportError:
+except (ImportError, RuntimeError):
     VideoDecoder = PlaceholderModule("torchcodec").placeholder_attr(  # type: ignore[assignment]
         "decoders.VideoDecoder"
     )
@@ -956,6 +956,7 @@ class VideoBackend(
             assert not frame_recovery, (
                 "frame_recovery is only available for `opencv` backend"
             )
+            check_torchcodec_available()
             decoder = cls.make_torchcodec_decoder(
                 data,
                 num_ffmpeg_threads=num_ffmpeg_threads,

vllm/transformers_utils/config.py

@@ -106,6 +106,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
     minimax_m3_vl="MiniMaxM3Config",
     minimax_m3_mtp="MiniMaxM3MTPConfig",
     moondream3="Moondream3Config",
+    moss_transcribe_diarize="MossTranscribeDiarizeConfig",
     eagle="EAGLEConfig",
     speculators="SpeculatorsConfig",
     nemotron="NemotronConfig",

vllm/transformers_utils/configs/__init__.py

@@ -60,6 +60,9 @@ _CLASS_TO_MODULE: dict[str, str] = {
     "Moondream3Config": "vllm.transformers_utils.configs.moondream3",
     "Moondream3TextConfig": "vllm.transformers_utils.configs.moondream3",
     "Moondream3VisionConfig": "vllm.transformers_utils.configs.moondream3",
+    "MossTranscribeDiarizeConfig": (
+        "vllm.transformers_utils.configs.moss_transcribe_diarize"
+    ),
     "MoonViTConfig": "vllm.transformers_utils.configs.moonvit",
     "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear",
     "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl",
@@ -134,6 +137,7 @@ __all__ = [
     "Moondream3Config",
     "Moondream3TextConfig",
     "Moondream3VisionConfig",
+    "MossTranscribeDiarizeConfig",
     "MoonViTConfig",
     "KimiLinearConfig",
     "KimiVLConfig",

vllm/transformers_utils/configs/moss_transcribe_diarize.py

@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+from transformers import PretrainedConfig, Qwen3Config
+from transformers.models.whisper.configuration_whisper import WhisperConfig
+
+
+class MossTranscribeDiarizeConfig(PretrainedConfig):
+    """Configuration for MOSS-Transcribe-Diarize."""
+
+    model_type = "moss_transcribe_diarize"
+    sub_configs = {"text_config": Qwen3Config, "audio_config": WhisperConfig}
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        text_config: dict[str, Any] | Qwen3Config | None = None,
+        audio_config: dict[str, Any] | WhisperConfig | None = None,
+        audio_token_id: int = 151671,
+        audio_merge_size: int = 4,
+        adaptor_input_dim: int | None = None,
+        tie_word_embeddings: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        text_config_obj: Qwen3Config
+        if text_config is None:
+            text_config_obj = Qwen3Config(
+                vocab_size=151936,
+                hidden_size=1024,
+                intermediate_size=3072,
+                num_hidden_layers=28,
+                num_attention_heads=16,
+                num_key_value_heads=8,
+                head_dim=128,
+                max_position_embeddings=40960,
+                tie_word_embeddings=tie_word_embeddings,
+                rope_theta=1_000_000.0,
+                layer_types=["full_attention"] * 28,
+            )
+        elif isinstance(text_config, dict):
+            text_config_obj = Qwen3Config(**text_config)
+        else:
+            text_config_obj = text_config
+
+        audio_config_obj: WhisperConfig
+        if audio_config is None:
+            audio_config_obj = WhisperConfig(
+                num_mel_bins=80,
+                d_model=1024,
+                encoder_layers=24,
+                encoder_attention_heads=16,
+                encoder_ffn_dim=4096,
+                max_source_positions=1500,
+                dropout=0.0,
+                attention_dropout=0.0,
+                activation_dropout=0.0,
+                activation_function="gelu",
+                encoder_layerdrop=0.0,
+                scale_embedding=False,
+            )
+        elif isinstance(audio_config, dict):
+            audio_config_obj = WhisperConfig(**audio_config)
+        else:
+            audio_config_obj = audio_config
+
+        text_config_obj.tie_word_embeddings = tie_word_embeddings
+        if not getattr(text_config_obj, "layer_types", None):
+            text_config_obj.layer_types = [
+                "full_attention"
+            ] * text_config_obj.num_hidden_layers
+
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+        self.text_config = text_config_obj
+        self.audio_config = audio_config_obj
+        self.audio_token_id = int(audio_token_id)
+        self.audio_merge_size = int(audio_merge_size)
+        self.adaptor_input_dim = (
+            int(adaptor_input_dim)
+            if adaptor_input_dim is not None
+            else int(audio_config_obj.d_model) * int(audio_merge_size)
+        )
+
+        self.vocab_size = int(text_config_obj.vocab_size)
+        self.hidden_size = int(text_config_obj.hidden_size)
+        self.intermediate_size = int(text_config_obj.intermediate_size)
+        self.num_hidden_layers = int(text_config_obj.num_hidden_layers)
+        self.num_attention_heads = int(text_config_obj.num_attention_heads)
+        self.num_key_value_heads = int(text_config_obj.num_key_value_heads)
+        self.head_dim = int(text_config_obj.head_dim)
+        self.hidden_act = text_config_obj.hidden_act
+        self.max_position_embeddings = int(text_config_obj.max_position_embeddings)
+        self.rms_norm_eps = float(text_config_obj.rms_norm_eps)
+        rope_parameters = getattr(text_config_obj, "rope_parameters", None)
+        rope_theta = float(getattr(text_config_obj, "rope_theta", 1_000_000.0))
+        if rope_parameters is None:
+            rope_parameters = {
+                "rope_type": "default",
+                "rope_theta": rope_theta,
+            }
+            text_config_obj.rope_parameters = rope_parameters
+        self.rope_parameters = rope_parameters
+        self.rope_theta = float(rope_parameters.get("rope_theta", rope_theta))
+        self.attention_bias = bool(text_config_obj.attention_bias)
+        self.attention_dropout = float(text_config_obj.attention_dropout)
+        self.is_causal = True

vllm/utils/import_utils.py

@@ -552,3 +552,21 @@ def has_cutedsl() -> bool:
 def has_humming() -> bool:
     """Whether the optional `humming` package is available."""
     return _has_module("humming")
+
+
+def check_torchcodec_available():
+    """Whether the optional `torchcodec` package is available."""
+    try:
+        import torchcodec  # noqa: F401
+    except RuntimeError as e:
+        # torchcodec will raise RuntimeError during import instead
+        # of ImportError when system ffmpeg unavailable, with a
+        # message that can leak sensitive system information.
+        # Trim it down to avoid it.
+        marker = (
+            "The following exceptions were raised as we tried to load libtorchcodec:"
+        )
+        message = str(e)
+        if marker in message:
+            raise RuntimeError(message.split(marker, 1)[0].rstrip()) from None
+        raise e

vllm/v1/attention/backends/mla/prefill/selector.py

@@ -134,7 +134,7 @@ def get_mla_prefill_backend(
                 f"Reason: {invalid_reasons}"
             )
         assert backend_cls is not None
-        logger.info("Using %s MLA prefill backend.", selected_backend.name)
+        logger.info_once("Using %s MLA prefill backend.", selected_backend.name)
         return backend_cls
 
     return _auto_select_mla_prefill_backend(

vllm/v1/attention/backends/mla/sparse_swa.py

@@ -90,7 +90,8 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
             dtype=self.dtype,
             sliding_window=self.window_size,
             cache_dtype_str=self.cache_config.cache_dtype,
-            alignment=576 if uses_fp8_ds_mla_layout else None,
+            # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
+            alignment=576 if uses_fp8_ds_mla_layout else 512,
             model_version="deepseek_v4",
             kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype),
         )

vllm/v1/core/kv_cache_manager.py

@@ -8,10 +8,14 @@ from typing import Literal, overload
 
 from vllm.distributed.kv_events import BlockStored, KVCacheEvent
 from vllm.logger import init_logger
+from vllm.utils.math_utils import cdiv
 from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator
 from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
 from vllm.v1.core.kv_cache_utils import KVCacheBlock
 from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    CrossAttentionSpec,
+    EncoderOnlyAttentionSpec,
     KVCacheConfig,
     get_kv_cache_spec_kind,
     get_kv_cache_spec_sliding_window,
@@ -593,6 +597,26 @@ class KVCacheManager:
         """Get the block ids of a request."""
         return self.get_blocks(request_id).get_block_ids()
 
+    def get_block_ids_for_computed_tokens(
+        self,
+        request_id: str,
+        num_computed_tokens: int,
+    ) -> tuple[list[int], ...]:
+        """Get block ids covering the request's computed tokens."""
+        block_ids = self.get_block_ids(request_id)
+        clipped_block_ids: list[list[int]] = []
+        for group, ids in zip(self.kv_cache_config.kv_cache_groups, block_ids):
+            spec = group.kv_cache_spec
+            if not isinstance(spec, AttentionSpec) or isinstance(
+                spec, (CrossAttentionSpec, EncoderOnlyAttentionSpec)
+            ):
+                clipped_block_ids.append(ids)
+                continue
+
+            num_valid_blocks = cdiv(num_computed_tokens, spec.block_size)
+            clipped_block_ids.append(ids[:num_valid_blocks])
+        return tuple(clipped_block_ids)
+
     def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
         """Cache the blocks for the request, if enabled.
 

vllm/v1/core/sched/scheduler.py

@@ -872,12 +872,10 @@ class Scheduler(SchedulerInterface):
                     if num_new_tokens == 0:
                         break
 
-                # Handles an edge case when P/D Disaggregation
-                # is used with Spec Decoding where an
-                # extra block gets allocated which
-                # creates a mismatch between the number
-                # of local and remote blocks.
-                limit_lookahead_tokens = load_kv_async and self.use_eagle
+                # During async KV load, no forward pass is run yet.
+                # Allocate speculative lookahead slots later to avoid
+                # mismatching local and remote block counts.
+                limit_lookahead_tokens = load_kv_async and self.num_lookahead_tokens > 0
                 effective_lookahead_tokens = (
                     0 if limit_lookahead_tokens else self.num_lookahead_tokens
                 )
@@ -2371,7 +2369,10 @@ class Scheduler(SchedulerInterface):
             num_prompt_tokens=request.num_prompt_tokens,
         )
 
-        block_ids = self.kv_cache_manager.get_block_ids(request.request_id)
+        block_ids = self.kv_cache_manager.get_block_ids_for_computed_tokens(
+            request_id=request.request_id,
+            num_computed_tokens=request.num_computed_tokens,
+        )
 
         if not isinstance(self.connector, SupportsHMA):
             # NOTE(Kuntai): We should deprecate this code path after we enforce

vllm/v1/kv_offload/tiering/p2p/manager.py

@@ -11,7 +11,7 @@ from __future__ import annotations
 import time
 from collections.abc import Iterable, Sequence
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from typing_extensions import override
 
@@ -115,7 +115,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
         port: int = 7777,
         backends: list[str] | None = None,
         num_threads: int = 4,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         """Initialize the P2P secondary tier manager.
 

vllm/v1/worker/gpu/async_utils.py

@@ -24,7 +24,8 @@ class AsyncOutput(AsyncModelRunnerOutput):
         self.model_runner_output = model_runner_output
         self.sampler_output = sampler_output
         self.num_sampled_tokens = num_sampled_tokens
-        self.copy_event = torch.cuda.Event()
+        # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
+        self.copy_event = torch.cuda.Event(blocking=True)
 
         with stream(copy_stream, main_stream):
             copy_stream.wait_stream(main_stream)
@@ -81,7 +82,8 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput):
         self.model_runner_output = model_runner_output
         self.pooler_output = pooler_output
         self.is_valid = is_valid
-        self.copy_event = torch.cuda.Event()
+        # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
+        self.copy_event = torch.cuda.Event(blocking=True)
 
         with stream(copy_stream, main_stream):
             copy_stream.wait_stream(main_stream)

vllm/v1/worker/gpu/input_batch.py

@@ -340,10 +340,12 @@ def _combine_sampled_and_draft_tokens_kernel(
         # Handling prefill tokens. No sampled or draft tokens.
         return
 
-    if NUM_NEW_SAMPLED_TOKENS > 0:
+    # Keep prompt-tail slots intact; only rewrite generated-token slots.
+    first_logit_seq_pos = seq_len - num_logits
+    if NUM_NEW_SAMPLED_TOKENS > 0 and first_logit_seq_pos >= prefill_len:
         # Write the last sampled token ID to input_ids.
         last_token_id = tl.load(last_sampled_tokens_ptr + req_state_idx)
-        tl.store(input_ids_ptr + query_end - num_logits, last_token_id)
+        tl.store(input_ids_ptr + logits_start, last_token_id)
 
     # Write the draft tokens (if any) to input_ids.
     if num_draft_tokens > 0:

vllm/v1/worker/gpu/model_states/mamba_hybrid.py

@@ -255,11 +255,16 @@ class MambaHybridModelState(DefaultModelState):
             # GDN uses >= 0 to select spec-decode rows, so non-decode rows
             # need the -1 sentinel rather than a raw zero draft count.
             num_decode_draft_tokens_np = np.full(num_reqs, -1, dtype=np.int32)
-            if input_batch.num_draft_tokens_per_req is not None:
-                has_draft_tokens = input_batch.num_draft_tokens_per_req > 0
-                spec_decode_mask = has_draft_tokens & ~input_batch.is_prefilling_np
+            num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req
+            if num_draft_tokens_per_req is not None:
+                # A row is a spec-decode row only when its whole prompt is already
+                # computed, i.e. exactly one non-draft (decode) token is scheduled.
+                is_decode = (
+                    input_batch.num_scheduled_tokens == num_draft_tokens_per_req + 1
+                )
+                spec_decode_mask = (num_draft_tokens_per_req > 0) & is_decode
                 num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where(
-                    spec_decode_mask, input_batch.num_draft_tokens_per_req, -1
+                    spec_decode_mask, num_draft_tokens_per_req, -1
                 )
             num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np)
 

vllm/v1/worker/gpu/spec_decode/utils.py

@@ -12,7 +12,8 @@ class DraftTokensHandler:
     def __init__(self, device: torch.device | None = None):
         self.device = device
         self.copy_stream = torch.cuda.Stream(device)
-        self.copy_event = torch.cuda.Event()
+        # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
+        self.copy_event = torch.cuda.Event(blocking=True)
 
         self.req_ids: list[str] = []
         self.draft_tokens_np: np.ndarray | None = None

vllm/v1/worker/gpu/states.py

@@ -110,16 +110,6 @@ class RequestState:
         self.num_computed_tokens_np[req_idx] = num_computed_tokens
         self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens)
 
-        if 0 < num_computed_tokens <= prefill_len:
-            # For PD disagg or resumed requests: set last_sampled to the last
-            # computed token so the first decode step gets the right input_id.
-            # For fresh prefill requests (num_computed_tokens == 0) the tensor
-            # is not read by combine_sampled_and_draft_tokens so we skip the
-            # write. Use a slice assignment rather than scalar indexing so the
-            # write is dispatched through fill_ without a host/device sync.
-            self.last_sampled_tokens[req_idx : req_idx + 1] = all_token_ids[
-                num_computed_tokens - 1
-            ]
         self.draft_tokens[req_idx].zero_()
 
     def apply_staged_writes(self) -> None:

vllm/v1/worker/gpu_model_runner.py

@@ -259,7 +259,8 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
         self._invalid_req_indices = invalid_req_indices
 
         # Event on the copy stream so we can synchronize the non-blocking copy.
-        self.async_copy_ready_event = torch.Event()
+        # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
+        self.async_copy_ready_event = torch.cuda.Event(blocking=True)
 
         # Keep a reference to the device tensor to avoid it being
         # deallocated until we finish copying it to the host.
@@ -392,7 +393,8 @@ class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput):
         self._model_runner_output = model_runner_output
 
         # Event on the copy stream so we can synchronize the non-blocking copy.
-        self.async_copy_ready_event = torch.Event()
+        # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
+        self.async_copy_ready_event = torch.cuda.Event(blocking=True)
 
         # Keep a reference to the device tensors to avoid them being
         # deallocated until we finish copying it to the host.
@@ -718,7 +720,9 @@ class GPUModelRunner(
         self.prepare_inputs_event: torch.Event | None = None
         if self.use_async_scheduling:
             self.async_output_copy_stream = torch.cuda.Stream()
-            self.prepare_inputs_event = torch.Event()
+            # Blocking (sleep) event to avoid busy-polling the CUDA driver lock;
+            # under TP contention that spin can balloon and make the rank a straggler.
+            self.prepare_inputs_event = torch.cuda.Event(blocking=True)
 
         # self.cudagraph_batch_sizes sorts in ascending order.
         if (
@@ -2204,10 +2208,7 @@ class GPUModelRunner(
                 req_idx = self.input_batch.req_id_to_index[req_id]
                 draft_len = len(draft_token_ids)
                 num_draft_tokens[req_idx] = draft_len
-                if (
-                    self.input_batch.num_computed_tokens_cpu[req_idx]
-                    >= self.input_batch.num_prompt_tokens[req_idx]
-                ):
+                if num_scheduled_tokens[req_idx] == draft_len + 1:
                     num_decode_draft_tokens[req_idx] = draft_len
             spec_decode_metadata = self._calc_spec_decode_metadata(
                 num_draft_tokens, cu_num_tokens

vllm/v1/worker/xpu_model_runner.py

@@ -49,7 +49,13 @@ def _torch_cuda_wrapper():
     torch.cuda.current_stream = partial(torch.xpu.current_stream)
     torch.cuda.stream = partial(torch.xpu.stream)
     torch.cuda.set_stream = partial(torch.xpu.set_stream)
-    torch.cuda.Event = partial(torch.xpu.Event)
+
+    # torch.xpu.Event does not accept the ``blocking`` kwarg that
+    # torch.cuda.Event supports, so drop it here.
+    def _xpu_event(*args, blocking=None, **kwargs):
+        return torch.xpu.Event(*args, **kwargs)
+
+    torch.cuda.Event = _xpu_event
     if supports_xpu_graph():
         torch.cuda.graph = partial(torch.xpu.graph)
         torch.cuda.CUDAGraph = torch.xpu.XPUGraph