[docs] Fix the docs build (#48008 )

Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> (cherry picked from commit 6cf7b26bd4)
Pin PyNvVideoCodec to tested 2.0.4 wheel (#48056 )
2026-07-11 11:25:04 +01:00 · 2026-07-09 12:43:21 -07:00 · 2026-07-09 03:24:39 -07:00 · 2026-07-09 03:24:39 -07:00 · 2026-07-08 21:12:19 -07:00 · 2026-07-08 21:12:19 -07:00
50 changed files with 1474 additions and 186 deletions
@@ -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 ""
@@ -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. | ✅︎ | ✅︎ |
@@ -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
@@ -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)

@@ -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)

@@ -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",
@@ -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
@@ -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))
@@ -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"
 )
@@ -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
@@ -652,6 +652,7 @@ class ParallelConfig:
            )
            and self.enable_expert_parallel
            and self.tensor_parallel_size > 1
+            and self.data_parallel_size > 1
        )

    @property
@@ -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()

@@ -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
@@ -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)
@@ -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.
@@ -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,
@@ -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

@@ -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,
@@ -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,
+        )
@@ -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",
@@ -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]
@@ -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):
@@ -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,
@@ -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)
@@ -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): ...
@@ -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
@@ -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): ...
@@ -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] = (
@@ -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):
@@ -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",
@@ -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, :]
@@ -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)
@@ -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
@@ -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
@@ -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)
@@ -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",
@@ -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",
@@ -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
@@ -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(
@@ -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),
        )
@@ -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.

@@ -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
@@ -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.

@@ -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)
@@ -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:
@@ -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)

@@ -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
@@ -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:
@@ -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
@@ -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
Author	SHA1	Message	Date
Martin HickeyandHarry Mellor	702f4814fe	[docs] Fix the docs build (#48008 ) Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> (cherry picked from commit `6cf7b26bd4`)	2026-07-11 11:25:04 +01:00
Brandon Pelfreyandkhluu	dd10e03f95	Pin PyNvVideoCodec to tested 2.0.4 wheel (#48056 ) (cherry picked from commit `753c5039f0`)	2026-07-09 12:43:21 -07:00
Chaunceyandkhluu	3d1c21a6fc	[P/D][Bugfix] Fix PD async KV load lookahead handling for MTP spec decode (#46694 ) Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> (cherry picked from commit `ae6170f874`)	2026-07-09 03:24:39 -07:00
liuzhenweiandkhluu	e89b9c26d9	[XPU] Fix Event init failure w/ blocking (#47868 ) Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com> (cherry picked from commit `bdaf27519f`)	2026-07-09 03:24:39 -07:00
Harry Mellorandkhluu	8bbc0062b1	Fix embed scaling + CUDA graphs in Transformers modelling backend (#48010 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> (cherry picked from commit `56da398dac`)	2026-07-08 21:12:19 -07:00
Thien Tranandkhluu	65168ef45c	Allow FlashInfer A2A backends for TRTLLM FP8 MoE Modular (#46661 ) Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> (cherry picked from commit `0d2f4e7c9c`)	2026-07-08 21:12:19 -07:00
Robert Shawandkhluu	51e3572d30	[Bug] Fix Batched DeepGEMM (#47884 ) Signed-off-by: Robert Shaw <robertgshaw2-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com> Co-authored-by: Robert Shaw <robertgshaw2-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com> Co-authored-by: Claude <noreply@anthropic.com> (cherry picked from commit `572b25b03e`)	2026-07-08 21:12:19 -07:00
Canlin Guoandkhluu	654e9e5b90	[Model] Support MOSS-Transcribe-Diarize (#47729 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com> (cherry picked from commit `285c08c036`)	2026-07-08 21:12:19 -07:00
Wei Zhaoandkhluu	15d7edbbcf	[Perf] Minimax M3 - Support cross-layer allreduce-norm fusion (#47631 ) Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `2afa3f7e95`)	2026-07-08 21:12:19 -07:00
Jeff (Junze) Maandkhluu	b53923c9ef	[Bugfix] DSV4 TP16 garbage output (#47493 ) Signed-off-by: Jeff Ma <jeffjma@umich.edu> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> (cherry picked from commit `80eb01e93d`)	2026-07-08 21:12:19 -07:00
Yongye Zhuandkhluu	cf3df33687	[Minimax-M3] Using tok_sparse_select from MSA instead of triton kernels (#47502 ) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> (cherry picked from commit `9021589498`)	2026-07-08 21:12:19 -07:00
Dakai Anandkhluu	ae3665a6ab	[Bugfix] Fix PD disagg + MTP correctness for Qwen3.5(GDN) (#47466 ) Signed-off-by: Dakai An <dakaian108@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> (cherry picked from commit `65dcde1695`)	2026-07-08 21:12:19 -07:00
Summer Yangandkhluu	f5ba5da988	[Perf] Use blocking CUDA events to avoid busy polling cuda driver lock (#47081 ) Signed-off-by: Jingyi Yang <girasoleyang@gmail.com> (cherry picked from commit `d3e69fd671`)	2026-07-08 21:12:19 -07:00
Wentao Yeandkhluu	7da76e4147	[CI Bug Fix] Temp fix for v3.2 accuracy (#47902 ) (cherry picked from commit `3f99883d97`)	2026-07-07 18:25:17 -07:00
Rahul Vishwakarmaandkhluu	b2dec4ac5a	[CPU][Bugfix] Fix flaky ShortConv prefill test on ARM (uninitialized weights) (#47848 ) Signed-off-by: Rahul Vishwakarma <Rahul.Vishwakarma2@ibm.com> (cherry picked from commit `f7efab58ec`)	2026-07-07 18:20:30 -07:00
Harry Mellorandkhluu	88550342ce	[CI] Fix Transformers modeling backend LoRA test (#47832 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> (cherry picked from commit `0ed05b6f82`)	2026-07-07 18:07:57 -07:00
danielafrimiandkhluu	ad49bb65d9	[BugFix] Fix ModelOpt mixed-precision quantization for sparse `quantized_layers` configs. (#47318 ) Signed-off-by: Daniel Afrimi <dafrimi@nvidia.com> Signed-off-by: <dafrimi@nvidia.com> (cherry picked from commit `0a2965b1b3`)	2026-07-07 18:05:03 -07:00