[ROCm][CI] Fix Kimi K3 KDA on ROCm (#50262 )

Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
[Spec Decode][Perf] Replicate DSpark Markov head across TP ranks (#49731 )
2026-07-29 15:47:50 +00:00 · 2026-07-29 11:43:51 -04:00 · 2026-07-29 06:41:54 -07:00 · 2026-07-29 06:30:45 -07:00 · 2026-07-29 13:16:13 +00:00 · 2026-07-29 12:21:09 +00:00
158 changed files with 4289 additions and 3580 deletions
@@ -28,11 +28,6 @@
        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
      }
    },
-    {
-        "dataset_name": "sharegpt",
-        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
-      }
-    },
    {
      "test_name": "serving_llama8B_tp1_random_128_128",
      "server_parameters": {
@@ -369,7 +369,7 @@ export HF_TOKEN ZE_AFFINITY_MASK
    -e CMDS \
    --name "${container_name}" \
    "${IMAGE}" \
-    bash -c 'set -e; source /opt/intel/oneapi/setvars.sh --force; source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
+    bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
    >/dev/null
 } 9>/tmp/docker-pull.lock

@@ -1576,10 +1576,14 @@ steps:
  - vllm/third_party/flash_linear_attention/ops/kda.py
  - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
  - vllm/third_party/flash_linear_attention/ops/l2norm.py
-  - tests/kernels/test_kda.py
+  - vllm/models/kimi_k3/nvidia/kda.py
+  - vllm/models/kimi_k3/nvidia/kda_metadata.py
+  - vllm/models/kimi_k3/nvidia/ops/third_party/kda/
+  - tests/models/kimi_k3/test_kda.py
+  - tests/models/kimi_k3/test_kda_metadata.py
  - vllm/platforms/rocm.py
  commands:
-  - pytest -v -s kernels/test_kda.py
+  - pytest -v -s models/kimi_k3/test_kda.py models/kimi_k3/test_kda_metadata.py

 - label: Kernels Mamba Test # TBD
  timeout_in_minutes: 180
@@ -28,9 +28,9 @@ if(DEEPGEMM_SRC_DIR)
  message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
 else()
  # Keep in sync with tools/install_deepgemm.sh
-  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git")
+  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
  # NOTE: This is currently targeting nv-dev branch due to sm120 support
-  set(_DEEPGEMM_UPSTREAM_TAG "f5a76426fa084087169693fd0cd815223576d6e9")
+  set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a")

  set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
  if(NOT _deepgemm_fc_root)
@@ -63,21 +63,6 @@ RUN apt-get update -y && \
    python3-pip && \
    rm -rf /var/lib/apt/lists/*

-# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
-RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
-    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
-    printf '%s\n' \
-    'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \
-    'Pin: version 2025.3*' \
-    'Pin-Priority: 1001' \
-    > /etc/apt/preferences.d/oneapi-2025.3.pref && \
-    apt-get update -y && \
-    apt-get install -y --no-install-recommends \
-    intel-oneapi-compiler-dpcpp-cpp-2025.3 \
-    intel-oneapi-mkl-devel-2025.3 \
-    intel-oneapi-dnnl-devel-2025.3 && \
-    rm -rf /var/lib/apt/lists/*
-
 # Install UMD
 RUN mkdir neo && \
    cd neo && \
@@ -100,22 +85,6 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh \
    && uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"

-# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
-ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
-RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
-    bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
-    rm "${ONECCL_INSTALLER}" && \
-    echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
-    echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc && \
-    rm -f /opt/intel/oneapi/ccl/latest && \
-    ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest && \
-    printf '%s\n' \
-    '/opt/intel/oneapi/ccl/2021.15/lib' \
-    '/opt/intel/oneapi/mpi/2021.15/lib' \
-    '/opt/intel/oneapi/compiler/2025.3/lib' \
-    > /etc/ld.so.conf.d/oneapi-ccl.conf && \
-    ldconfig
-
 SHELL ["bash", "-c"]
 CMD ["bash", "-c", "source /root/.bashrc && exec bash"]

@@ -135,8 +104,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --upgrade pip


-
-ENV LD_LIBRARY_PATH=/opt/intel/oneapi/ccl/2021.15/lib:/opt/intel/oneapi/mpi/2021.15/lib:/opt/intel/oneapi/compiler/2025.3/lib:/usr/local/lib
+ENV LD_LIBRARY_PATH=/opt/venv/lib:/usr/local/lib
 CMD ["/bin/bash"]

 ######################### UCX + NIXL BUILD STAGE #########################
@@ -216,8 +184,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install -r /workspace/vllm/requirements/xpu.txt && \
    uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt && \
    uv pip uninstall triton triton-xpu && \
-    uv pip install triton-xpu==3.7.1 && \
-    uv pip uninstall oneccl oneccl-devel
+    uv pip install triton-xpu==3.7.2

 # Keep source-dependent layers near the end so frequent code-only changes
 # don't invalidate heavy dependency and UCX/NIXL layers.
@@ -42,12 +42,12 @@ pip install -v -r requirements/xpu.txt

    ```bash
    pip uninstall -y triton triton-xpu
-    pip install triton-xpu==3.7.1 --extra-index-url https://download.pytorch.org/whl/xpu
+    pip install triton-xpu==3.7.2 --extra-index-url https://download.pytorch.org/whl/xpu
    ```

    !!! note
        - `triton` (without suffix) is for NVIDIA GPUs only. On XPU, using it instead of `triton-xpu` can cause correctness or runtime issues.
-        - For torch 2.12 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.1`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
+        - For torch 2.13 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.2`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).

 - Finally, build and install vLLM XPU backend:

@@ -547,7 +547,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ |
 | `KimiAudioForConditionalGeneration` | Kimi-Audio | T + A<sup>+</sup> | `moonshotai/Kimi-Audio-7B-Instruct` | | ✅︎ |
 | `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I<sup>+</sup> | `moonshotai/Kimi-K2.5` | | ✅︎ |
-| `KimiK3ForConditionalGeneration` | Kimi-K3 | T + I<sup>+</sup> | `moonshotai/Kimi-K3` | | ✅︎ |
 | `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
 | `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I<sup>+</sup> | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ |
 | `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I<sup>+</sup> | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |
@@ -155,8 +155,6 @@ view_seperator = "view_seperator"
 inverse_std_variences = "inverse_std_variences"

 [tool.typos.default.extend-words]
-Hel = "Hel"
-wether = "wether"
 iy = "iy"
 indx = "indx"
 # intel cpu features
@@ -16,8 +16,8 @@ PyNvVideoCodec==2.0.4
 --extra-index-url https://flashinfer.ai/whl/
 flashinfer-python==0.6.15.post1
 flashinfer-cubin==0.6.15.post1
-apache-tvm-ffi==0.1.11
-tilelang==0.1.12
+apache-tvm-ffi==0.1.10
+tilelang==0.1.9
 nvidia-cudnn-frontend>=1.19.1
 # Required for LLM_NVTX_SCOPES_FOR_PROFILING=1
 nvtx==0.2.15
@@ -45,7 +45,7 @@ anyio==4.14.1
    #   sse-starlette
    #   starlette
    #   watchfiles
-apache-tvm-ffi==0.1.11
+apache-tvm-ffi==0.1.10
    # via
    #   -c requirements/cuda.txt
    #   xgrammar
@@ -140,7 +140,7 @@ docopt==0.6.2
    # via num2words
 docstring-parser==0.18.0
    # via anthropic
-dpcpp-cpp-rt==2025.3.2
+dpcpp-cpp-rt==2026.0.0
    # via
    #   onemkl-sycl-blas
    #   onemkl-sycl-dft
@@ -253,27 +253,27 @@ ijson==3.5.0
    # via -r requirements/test/../common.txt
 imageio==2.37.3
    # via scikit-image
-impi-rt==2021.17.2
+impi-rt==2021.18.0
    # via
    #   oneccl
    #   torch
 iniconfig==2.3.0
    # via pytest
-intel-cmplr-lib-rt==2025.3.2
+intel-cmplr-lib-rt==2026.0.0
    # via
    #   intel-sycl-rt
    #   torch
-intel-cmplr-lib-ur==2025.3.2
+intel-cmplr-lib-ur==2026.0.0
    # via
    #   intel-openmp
    #   intel-sycl-rt
    #   torch
-intel-cmplr-lic-rt==2025.3.2
+intel-cmplr-lic-rt==2026.0.0
    # via
    #   intel-opencl-rt
    #   intel-sycl-rt
    #   torch
-intel-opencl-rt==2025.3.2
+intel-opencl-rt==2026.0.0
    # via
    #   dpcpp-cpp-rt
    #   onemkl-sycl-blas
@@ -282,14 +282,14 @@ intel-opencl-rt==2025.3.2
    #   onemkl-sycl-rng
    #   onemkl-sycl-sparse
    #   torch
-intel-openmp==2025.3.2
+intel-openmp==2026.0.0
    # via
    #   dpcpp-cpp-rt
    #   mkl
    #   torch
-intel-pti==0.16.0
+intel-pti==0.17.0
    # via torch
-intel-sycl-rt==2025.3.2
+intel-sycl-rt==2026.0.0
    # via
    #   dpcpp-cpp-rt
    #   oneccl
@@ -378,7 +378,7 @@ mistral-common==1.11.5
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
    #   -r requirements/test/xpu.in
-mkl==2025.3.1
+mkl==2026.0.0
    # via
    #   onemkl-sycl-blas
    #   onemkl-sycl-dft
@@ -453,28 +453,28 @@ numpy==2.2.6
    #   torchvision
    #   transformers
    #   xgrammar
-oneccl==2021.17.2
+oneccl==2022.0.0
    # via
    #   oneccl-devel
    #   torch
-oneccl-devel==2021.17.2
+oneccl-devel==2022.0.0
    # via torch
-onemkl-license==2025.3.1
+onemkl-license==2026.0.0
    # via
    #   mkl
    #   torch
-onemkl-sycl-blas==2025.3.1
+onemkl-sycl-blas==2026.0.0
    # via
    #   onemkl-sycl-lapack
    #   onemkl-sycl-sparse
    #   torch
-onemkl-sycl-dft==2025.3.1
+onemkl-sycl-dft==2026.0.0
    # via torch
-onemkl-sycl-lapack==2025.3.1
+onemkl-sycl-lapack==2026.0.0
    # via torch
-onemkl-sycl-rng==2025.3.1
+onemkl-sycl-rng==2026.0.0
    # via torch
-onemkl-sycl-sparse==2025.3.1
+onemkl-sycl-sparse==2026.0.0
    # via torch
 openai==2.44.0
    # via
@@ -719,6 +719,8 @@ pyyaml==6.0.3
    #   timm
    #   transformers
    #   uvicorn
+pyzes==0.1.1
+    # via torch
 pyzmq==27.1.0
    # via
    #   -c requirements/common.txt
@@ -871,14 +873,14 @@ tabledata==1.3.4
    # via pytablewriter
 tabulate==0.10.0
    # via sacrebleu
-tbb==2022.3.1
+tbb==2023.0.0
    # via
    #   intel-opencl-rt
    #   mkl
    #   torch
 tblib==3.1.0
    # via -r requirements/test/xpu.in
-tcmlib==1.4.1
+tcmlib==1.5.0
    # via
    #   tbb
    #   torch
@@ -910,7 +912,7 @@ tokenizers==0.22.2
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
    #   transformers
-torch==2.12.0+xpu
+torch==2.13.0+xpu
    # via
    #   -c requirements/xpu.txt
    #   accelerate
@@ -920,7 +922,7 @@ torch==2.12.0+xpu
    #   timm
    #   torchvision
    #   xgrammar
-torchvision==0.27.0+xpu
+torchvision==0.28.0+xpu
    # via timm
 tqdm==4.67.3
    # via
@@ -946,7 +948,7 @@ transformers==5.14.1
    #   xgrammar
 triton==3.7.1
    # via xgrammar
-triton-xpu==3.7.1
+triton-xpu==3.7.2
    # via torch
 typepy==1.3.4
    # via
@@ -1001,7 +1003,7 @@ typing-inspection==0.4.2
    #   mcp
    #   pydantic
    #   pydantic-settings
-umf==1.0.3
+umf==1.1.0
    # via
    #   intel-cmplr-lib-ur
    #   torch
@@ -12,10 +12,10 @@ jinja2>=3.1.6
 datasets # for benchmark scripts
 numba == 0.65.0 # Required for N-gram speculative decoding
 --extra-index-url=https://download.pytorch.org/whl/xpu
-torch==2.12.0
+torch==2.13.0
 torchaudio
 torchvision
 torchcodec >= 0.14 # Required for the torchcodec video decoding backend

-auto_round_lib==0.14.1
-vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.11.1/vllm_xpu_kernels-0.1.11.1-cp38-abi3-manylinux_2_28_x86_64.whl
+auto_round_lib==0.14.2
+vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.12/vllm_xpu_kernels-0.1.12-cp38-abi3-manylinux_2_28_x86_64.whl
@@ -57,6 +57,7 @@ impl HfChatBackend {
                    processor_config: files.processor_config_path.as_deref(),
                },
                tokenizer.clone(),
+                options.limit_mm_per_prompt.clone(),
            )?
        };
        let multimodal_render_info = resolve_multimodal_render_info(multimodal_model_info.as_ref());
@@ -231,6 +232,7 @@ mod tests {
                chat_template_content_format: Default::default(),
                chat_template: None,
                default_chat_template_kwargs: HashMap::new(),
+                limit_mm_per_prompt: HashMap::new(),
            },
            test_tokenizer(),
        )
@@ -8,7 +8,7 @@ use serde_json::Value;
 use vllm_text::{DynTextBackend, TextBackend};

 use crate::error::Result;
-use crate::multimodal::MultimodalModelInfo;
+use crate::multimodal::{MmLimitPerPrompt, MultimodalModelInfo};
 use crate::output::DynChatOutputProcessor;
 use crate::renderer::DynChatRenderer;
 use crate::request::ChatRequest;
@@ -74,6 +74,9 @@ pub struct LoadModelBackendsOptions {
    /// Optional server-default keyword arguments merged into every
    /// chat-template render before request-level `chat_template_kwargs`.
    pub default_chat_template_kwargs: HashMap<String, Value>,
+    /// Maximum number of input items allowed per prompt for each modality.
+    /// Unspecified modalities are unlimited.
+    pub limit_mm_per_prompt: MmLimitPerPrompt,
 }

 /// Shared backends loaded from a model id.
@@ -23,6 +23,8 @@ pub enum Error {
    UnsupportedMultimodalContent(&'static str),
    #[error("`{modality}` input is not supported by this model")]
    UnsupportedModality { modality: String },
+    #[error("At most {limit} {modality}(s) may be provided in one prompt.")]
+    MmLimitExceeded { modality: String, limit: usize },
    #[error("multimodal preprocessing error: {0}")]
    Multimodal(#[message] String),
    #[error("{kind} parsing is not available for model `{model_id}`")]
@@ -87,7 +89,8 @@ impl Error {
            Self::Text(error) => error.is_request_validation_error(),
            Self::UnsupportedMultimodalRenderer
            | Self::UnsupportedMultimodalContent(_)
-            | Self::UnsupportedModality { .. } => true,
+            | Self::UnsupportedModality { .. }
+            | Self::MmLimitExceeded { .. } => true,

            _ => false,
        }
@@ -11,7 +11,7 @@
 //! Raw media stays above `vllm-text`; this module lowers it into token IDs and
 //! opaque tensor payloads before the request is handed to text generation.

-use std::collections::HashSet;
+use std::collections::{BTreeMap, HashMap, HashSet};
 use std::fs;
 use std::path::Path;
 use std::sync::{Arc, LazyLock};
@@ -24,6 +24,7 @@ use llm_multimodal::{
    PromptReplacement, Tokenizer as TokenResolver, TrackedMedia, VideoClip, VisionPreProcessor,
    VisionProcessorRegistry,
 };
+use serde::{Deserialize, Serialize};
 use thiserror_ext::AsReport as _;
 use tracing::warn;
 use vllm_engine_core_client::protocol::dtype::ModelDtype;
@@ -52,6 +53,71 @@ pub struct MultimodalModelInfo {
    video: Option<ModalitySupport>,
    audio: Option<AudioModalitySupport>,
    media_connector: Arc<MediaConnector>,
+    /// Maximum number of input items allowed per prompt for each modality.
+    limit_mm_per_prompt: MmLimitPerPrompt,
+}
+
+/// Per-modality item-count limits configured by `--limit-mm-per-prompt`.
+///
+/// Modalities absent from the map are unlimited.
+pub type MmLimitPerPrompt = HashMap<MmLimitModality, MmLimitSpec>;
+
+/// Modalities that `--limit-mm-per-prompt` can be keyed by.
+///
+/// Closed on purpose: these are exactly the keys Python accepts, per
+/// `MultiModalDummyOptionsBuiltins` in `vllm/config/multimodal.py`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum MmLimitModality {
+    Image,
+    Audio,
+    Video,
+}
+
+impl MmLimitModality {
+    /// The wire name, matching Python's modality strings.
+    pub fn as_str(self) -> &'static str {
+        match self {
+            Self::Image => "image",
+            Self::Audio => "audio",
+            Self::Video => "video",
+        }
+    }
+}
+
+/// One modality's limit, in either of the two shapes Python accepts.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(
+    untagged,
+    expecting = "an item count, or an object with an optional `count` field"
+)]
+pub enum MmLimitSpec {
+    /// Legacy form: `"image": 16`
+    Count(usize),
+
+    /// Configurable form:
+    /// `"video": {"count": 1, "num_frames": 32}`
+    Options {
+        /// Absent means unlimited, matching an absent modality.
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        count: Option<usize>,
+
+        /// Preserve Python-owned options for forwarding. Never interpreted
+        /// here: they size the engine's dummy-profiling encoder cache, which
+        /// has no Rust counterpart.
+        #[serde(flatten)]
+        extra: BTreeMap<String, serde_json::Value>,
+    },
+}
+
+impl MmLimitSpec {
+    /// The configured item count, or `None` when this modality is unlimited.
+    pub fn count(&self) -> Option<usize> {
+        match self {
+            Self::Count(count) => Some(*count),
+            Self::Options { count, .. } => *count,
+        }
+    }
 }

 /// Model metadata and tokenizer access shared by all multimodal specs.
@@ -287,6 +353,7 @@ impl MultimodalModelInfo {
        model_type: Option<String>,
        files: MultimodalConfigFiles<'_>,
        tokenizer: DynTokenizer,
+        limit_mm_per_prompt: MmLimitPerPrompt,
    ) -> Result<Option<Self>> {
        let config = match files.config {
            Some(path) => {
@@ -319,7 +386,12 @@ impl MultimodalModelInfo {
            tokenizer: TokenizerResolver(tokenizer),
        };

-        Self::from_loaded(context, preprocessor_config, video_preprocessor_config)
+        Self::from_loaded(
+            context,
+            preprocessor_config,
+            video_preprocessor_config,
+            limit_mm_per_prompt,
+        )
    }

    /// Resolve multimodal support from an assembled context and parsed
@@ -328,6 +400,7 @@ impl MultimodalModelInfo {
        context: MultimodalModelContext,
        preprocessor_config: PreProcessorConfig,
        video_preprocessor_config: PreProcessorConfig,
+        limit_mm_per_prompt: MmLimitPerPrompt,
    ) -> Result<Option<Self>> {
        let (image, video) = Self::resolve_vision_lanes(
            &context,
@@ -356,6 +429,7 @@ impl MultimodalModelInfo {
            video,
            audio,
            media_connector,
+            limit_mm_per_prompt,
        }))
    }

@@ -567,7 +641,55 @@ fn input_audio_data_url(data: &str, format: Option<&str>) -> Result<String> {
    Ok(format!("data:{mime_type};base64,{data}"))
 }

+/// The modality a content part counts against, or `None` for plain text.
+///
+/// Embedding inputs share their base modality's budget rather than getting one
+/// of their own, matching Python's `modality.replace("_embeds", "")` in
+/// `vllm/entrypoints/chat_utils.py`.
+fn media_part_limit_modality(part: &MediaContentPart) -> Option<MmLimitModality> {
+    match part {
+        MediaContentPart::Text { .. } => None,
+        MediaContentPart::ImageUrl { .. }
+        | MediaContentPart::ImageData { .. }
+        | MediaContentPart::ImageEmbeds { .. } => Some(MmLimitModality::Image),
+        MediaContentPart::AudioUrl { .. } | MediaContentPart::AudioData { .. } => {
+            Some(MmLimitModality::Audio)
+        }
+        MediaContentPart::VideoUrl { .. } | MediaContentPart::VideoData { .. } => {
+            Some(MmLimitModality::Video)
+        }
+    }
+}
+
 impl MultimodalModelInfo {
+    /// Reject requests exceeding `--limit-mm-per-prompt`'s configured
+    /// per-modality item count, before any fetch/decode work is spent on them.
+    ///
+    /// Modalities without a configured count are unlimited.
+    fn validate_mm_limits(&self, media_parts: &[MediaContentPart]) -> Result<()> {
+        let mut counts: HashMap<MmLimitModality, usize> = HashMap::new();
+        for part in media_parts {
+            if let Some(modality) = media_part_limit_modality(part) {
+                *counts.entry(modality).or_default() += 1;
+            }
+        }
+
+        for (modality, count) in counts {
+            let Some(limit) = self.limit_mm_per_prompt.get(&modality).and_then(MmLimitSpec::count)
+            else {
+                continue;
+            };
+            if count > limit {
+                return Err(Error::MmLimitExceeded {
+                    modality: modality.as_str().to_string(),
+                    limit,
+                });
+            }
+        }
+
+        Ok(())
+    }
+
    /// Run media fetch, per-modality preprocessing, prompt expansion, and
    /// feature build.
    ///
@@ -584,8 +706,7 @@ impl MultimodalModelInfo {
        if media_parts_len == 0 {
            return Ok(Vec::new());
        }
-        // TODO: enforce per-modality item-count limits, aligned with the
-        // engine's `--limit-mm-per-prompt` semantics.
+        self.validate_mm_limits(&media_parts)?;
        let fetched = self.fetch_media(media_parts).await?;

        let mut prepared = Vec::new();
@@ -753,10 +874,11 @@ mod tests {
            .with_regular_token("<|video_pad|>", QWEN3_VIDEO_PAD_ID)
    }

-    fn test_info(
+    fn test_info_with_limits(
        model_type: &str,
        config: serde_json::Value,
        tokenizer: TestTokenizer,
+        limit_mm_per_prompt: MmLimitPerPrompt,
    ) -> MultimodalModelInfo {
        let context = MultimodalModelContext {
            model_id: format!("{model_type}-test"),
@@ -769,11 +891,20 @@ mod tests {
            context,
            PreProcessorConfig::default(),
            PreProcessorConfig::default(),
+            limit_mm_per_prompt,
        )
        .unwrap()
        .unwrap_or_else(|| panic!("{model_type} multimodal support should resolve"))
    }

+    fn test_info(
+        model_type: &str,
+        config: serde_json::Value,
+        tokenizer: TestTokenizer,
+    ) -> MultimodalModelInfo {
+        test_info_with_limits(model_type, config, tokenizer, HashMap::new())
+    }
+
    fn llama4_info() -> MultimodalModelInfo {
        let config = serde_json::json!({
            "model_type": "llama4",
@@ -783,16 +914,19 @@ mod tests {
        test_info("llama4", config, llama4_tokenizer())
    }

-    pub(super) fn qwen3_vl_info() -> MultimodalModelInfo {
-        let config = serde_json::json!({
+    fn qwen3_vl_config() -> serde_json::Value {
+        serde_json::json!({
            "model_type": "qwen3_vl",
            "image_token_id": QWEN3_IMAGE_PAD_ID,
            "video_token_id": QWEN3_VIDEO_PAD_ID,
            "vision_start_token_id": 151652,
            "vision_end_token_id": 151653,
            "vision_config": {"patch_size": 16}
-        });
-        test_info("qwen3_vl", config, qwen3_vl_tokenizer())
+        })
+    }
+
+    pub(super) fn qwen3_vl_info() -> MultimodalModelInfo {
+        test_info("qwen3_vl", qwen3_vl_config(), qwen3_vl_tokenizer())
    }

    #[test]
@@ -853,4 +987,123 @@ mod tests {
        );
        assert!(input_audio_data_url("AAAA", Some("flac")).is_err());
    }
+
+    fn image_url_part() -> MediaContentPart {
+        MediaContentPart::ImageUrl {
+            url: "https://example.com/image.png".to_string(),
+            detail: None,
+            uuid: None,
+        }
+    }
+
+    fn qwen3_vl_info_with_limits(limit_mm_per_prompt: MmLimitPerPrompt) -> MultimodalModelInfo {
+        test_info_with_limits(
+            "qwen3_vl",
+            qwen3_vl_config(),
+            qwen3_vl_tokenizer(),
+            limit_mm_per_prompt,
+        )
+    }
+
+    #[test]
+    fn validate_mm_limits_ignores_text_parts() {
+        let info = qwen3_vl_info();
+        let parts = vec![
+            MediaContentPart::Text {
+                text: "hello".to_string(),
+            },
+            MediaContentPart::Text {
+                text: "world".to_string(),
+            },
+        ];
+        assert!(info.validate_mm_limits(&parts).is_ok());
+    }
+
+    #[test]
+    fn validate_mm_limits_leaves_unconfigured_modalities_unlimited() {
+        let info = qwen3_vl_info();
+        let parts: Vec<_> = std::iter::repeat_with(image_url_part).take(1_000).collect();
+
+        assert!(info.validate_mm_limits(&parts).is_ok());
+    }
+
+    #[test]
+    fn validate_mm_limits_enforces_configured_limit_at_the_boundary() {
+        let info = qwen3_vl_info_with_limits(HashMap::from([(
+            MmLimitModality::Image,
+            MmLimitSpec::Count(1),
+        )]));
+        assert!(info.validate_mm_limits(&[image_url_part()]).is_ok());
+
+        let error = info.validate_mm_limits(&[image_url_part(), image_url_part()]).unwrap_err();
+        assert_eq!(
+            error.to_report_string(),
+            "At most 1 image(s) may be provided in one prompt."
+        );
+        // Confirms the HTTP-mapping bug found during implementation stays fixed:
+        // this must map to 400, not 500.
+        assert!(error.is_request_validation_error());
+    }
+
+    #[test]
+    fn validate_mm_limits_counts_image_embeds_against_the_image_limit() {
+        let info = qwen3_vl_info_with_limits(HashMap::from([(
+            MmLimitModality::Image,
+            MmLimitSpec::Count(1),
+        )]));
+        let image_embeds_part = MediaContentPart::ImageEmbeds {
+            payload: serde_json::Value::String("AAAA".to_string()),
+            uuid: None,
+        };
+
+        let error = info.validate_mm_limits(&[image_url_part(), image_embeds_part]).unwrap_err();
+        assert_eq!(
+            error.to_report_string(),
+            "At most 1 image(s) may be provided in one prompt."
+        );
+    }
+
+    /// An options object without a `count` carries only profiling keys, which
+    /// say nothing about how many items are allowed.
+    #[test]
+    fn validate_mm_limits_treats_a_count_less_options_object_as_unlimited() {
+        let info = qwen3_vl_info_with_limits(HashMap::from([(
+            MmLimitModality::Image,
+            MmLimitSpec::Options {
+                count: None,
+                extra: BTreeMap::from([("width".to_string(), serde_json::json!(512))]),
+            },
+        )]));
+
+        assert!(info.validate_mm_limits(&[image_url_part(), image_url_part()]).is_ok());
+    }
+
+    fn parse_limits(json: &str) -> MmLimitPerPrompt {
+        serde_json::from_str(json).expect("limit map should parse")
+    }
+
+    #[test]
+    fn limit_map_parses_both_shapes_python_accepts() {
+        let limits = parse_limits(r#"{"image": 16, "video": {"count": 1, "num_frames": 32}}"#);
+
+        assert_eq!(limits[&MmLimitModality::Image].count(), Some(16));
+        assert_eq!(limits[&MmLimitModality::Video].count(), Some(1));
+        assert_eq!(limits.get(&MmLimitModality::Audio), None);
+    }
+
+    #[test]
+    fn limit_map_rejects_keys_python_does_not_accept() {
+        assert!(serde_json::from_str::<MmLimitPerPrompt>(r#"{"image_embeds": 1}"#).is_err());
+    }
+
+    /// Managed mode forwards this map back to Python as JSON, where
+    /// `BaseDummyOptions.count` is a non-optional `int` under `extra="forbid"`.
+    /// Emitting `"count": null` would make the engine subprocess fail to start.
+    #[test]
+    fn limit_map_round_trips_without_emitting_a_null_count() {
+        let source = r#"{"video":{"num_frames":32}}"#;
+        let encoded = serde_json::to_string(&parse_limits(source)).expect("map should serialize");
+
+        assert_eq!(encoded, source);
+    }
 }
@@ -107,6 +107,7 @@ mod tests {
            context,
            PreProcessorConfig::default(),
            PreProcessorConfig::default(),
+            HashMap::new(),
        )
        .unwrap()
        .expect("Inkling multimodal support")
@@ -126,6 +127,7 @@ mod tests {
            context,
            PreProcessorConfig::default(),
            PreProcessorConfig::default(),
+            HashMap::new(),
        )
        .unwrap()
        .expect("Qwen3-ASR multimodal support")
@@ -222,7 +222,9 @@ mod tests {
        assert_eq!(tensor.shape, vec![expected.len()]);
        assert_eq!(
            tensor.data,
-            WireArrayData::RawView(expected.iter().map(|value| u8::from(*value)).collect())
+            WireArrayData::RawView(
+                expected.iter().map(|value| u8::from(*value)).collect::<Vec<_>>().into(),
+            )
        );
    }

@@ -2,28 +2,58 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 use std::collections::HashMap;
+use std::mem::size_of;

 use half::{bf16, f16};
 use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs};
 use vllm_engine_core_client::protocol::dtype::ModelDtype;
 use vllm_engine_core_client::protocol::multimodal::MmKwargValue as ProtocolKwargValue;
-use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
+use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireArrayData, WireTensor};

 use crate::error::{Error, Result, bail_multimodal, multimodal};

+/// Element type retained alongside an encoded tensor during multimodal lowering.
+#[derive(Debug, Clone, Copy)]
+pub(super) enum TensorKind {
+    /// 32-bit floating point.
+    F32,
+    /// IEEE 16-bit floating point.
+    F16,
+    /// Brain floating point.
+    Bf16,
+    /// Signed 64-bit integer.
+    I64,
+    /// Unsigned 32-bit integer.
+    U32,
+}
+
+impl TensorKind {
+    const fn element_size(self) -> usize {
+        match self {
+            Self::F32 => size_of::<f32>(),
+            Self::F16 => size_of::<f16>(),
+            Self::Bf16 => size_of::<bf16>(),
+            Self::I64 => size_of::<i64>(),
+            Self::U32 => size_of::<u32>(),
+        }
+    }
+
+    const fn wire_dtype(self) -> &'static str {
+        match self {
+            Self::F32 => "float32",
+            Self::F16 => "float16",
+            Self::Bf16 => "bfloat16",
+            Self::I64 => "int64",
+            Self::U32 => "uint32",
+        }
+    }
+}
+
 /// Representation for multimodal kwarg values for transformation.
 #[derive(Debug)]
 pub(super) enum KwargValue {
-    /// Float tensor with row-major flat data and shape.
-    F32Tensor { data: Vec<f32>, shape: Vec<usize> },
-    /// Float16 tensor with row-major flat data and shape.
-    F16Tensor { data: Vec<f16>, shape: Vec<usize> },
-    /// BFloat16 tensor with row-major flat data and shape.
-    Bf16Tensor { data: Vec<bf16>, shape: Vec<usize> },
-    /// Signed integer tensor with row-major flat data and shape.
-    I64Tensor { data: Vec<i64>, shape: Vec<usize> },
-    /// Unsigned integer tensor with row-major flat data and shape.
-    U32Tensor { data: Vec<u32>, shape: Vec<usize> },
+    /// Tensor with row-major flat data and shape.
+    Tensor { kind: TensorKind, wire: WireTensor },
    /// Non-tensor kwarg value that is shared or copied as-is.
    Passthrough(ProtocolKwargValue),
 }
@@ -67,8 +97,14 @@ impl KwargValue {
            ModelSpecificValue::Tensor { data, shape } => {
                Self::from_f32_tensor(data, shape, float_dtype)?
            }
-            ModelSpecificValue::IntTensor { data, shape } => Self::I64Tensor { data, shape },
-            ModelSpecificValue::UintTensor { data, shape } => Self::U32Tensor { data, shape },
+            ModelSpecificValue::IntTensor { data, shape } => {
+                let wire = WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?;
+                Self::tensor(TensorKind::I64, wire)
+            }
+            ModelSpecificValue::UintTensor { data, shape } => {
+                let wire = WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?;
+                Self::tensor(TensorKind::U32, wire)
+            }
            ModelSpecificValue::Int(value) => Self::Passthrough(Int(value)),
            ModelSpecificValue::Float(value) => Self::Passthrough(Float(value)),
            ModelSpecificValue::IntVec(values) => {
@@ -93,17 +129,23 @@ impl KwargValue {
    /// Convert a float tensor to the target float dtype if needed, keeping the
    /// same shape.
    fn from_f32_tensor(data: Vec<f32>, shape: Vec<usize>, float_dtype: ModelDtype) -> Result<Self> {
-        match float_dtype {
-            ModelDtype::Float16 => Ok(Self::F16Tensor {
-                data: data.into_iter().map(f16::from_f32).collect(),
-                shape,
-            }),
-            ModelDtype::BFloat16 => Ok(Self::Bf16Tensor {
-                data: data.into_iter().map(bf16::from_f32).collect(),
-                shape,
-            }),
-            ModelDtype::Float32 => Ok(Self::F32Tensor { data, shape }),
-        }
+        let (kind, wire) = match float_dtype {
+            ModelDtype::Float16 => (
+                TensorKind::F16,
+                WireTensor::from_f16(shape, data.into_iter().map(f16::from_f32).collect()),
+            ),
+            ModelDtype::BFloat16 => (
+                TensorKind::Bf16,
+                WireTensor::from_bf16(shape, data.into_iter().map(bf16::from_f32).collect()),
+            ),
+            ModelDtype::Float32 => (TensorKind::F32, WireTensor::from_f32(shape, data)),
+        };
+        wire.map(|wire| Self::tensor(kind, wire)).map_err(Error::Multimodal)
+    }
+
+    fn tensor(kind: TensorKind, wire: WireTensor) -> Self {
+        debug_assert_eq!(wire.dtype, kind.wire_dtype());
+        Self::Tensor { kind, wire }
    }
 }

@@ -111,15 +153,11 @@ impl TryFrom<&KwargValue> for ProtocolKwargValue {
    type Error = Error;

    fn try_from(value: &KwargValue) -> Result<Self> {
-        let tensor = match value {
-            KwargValue::F32Tensor { data, shape } => WireTensor::from_f32(shape.clone(), data),
-            KwargValue::F16Tensor { data, shape } => WireTensor::from_f16(shape.clone(), data),
-            KwargValue::Bf16Tensor { data, shape } => WireTensor::from_bf16(shape.clone(), data),
-            KwargValue::I64Tensor { data, shape } => WireTensor::from_i64(shape.clone(), data),
-            KwargValue::U32Tensor { data, shape } => WireTensor::from_u32(shape.clone(), data),
+        let wire = match value {
+            KwargValue::Tensor { wire, .. } => wire.clone(),
            KwargValue::Passthrough(value) => return Ok(value.clone()),
        };
-        tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
+        Ok(ProtocolKwargValue::Tensor(wire))
    }
 }

@@ -127,11 +165,7 @@ impl KwargValue {
    /// First-axis length for tensor values; `None` for passthrough kwargs.
    pub(super) fn first_dim(&self) -> Option<usize> {
        match self {
-            Self::F32Tensor { shape, .. }
-            | Self::F16Tensor { shape, .. }
-            | Self::Bf16Tensor { shape, .. }
-            | Self::I64Tensor { shape, .. }
-            | Self::U32Tensor { shape, .. } => shape.first().copied(),
+            Self::Tensor { wire, .. } => wire.shape.first().copied(),
            Self::Passthrough(_) => None,
        }
    }
@@ -161,30 +195,13 @@ impl KwargValue {
        end: usize,
        drop_axis: bool,
    ) -> Result<ProtocolKwargValue> {
-        let tensor = match self {
-            Self::F32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_f32(shape, data)
-            }
-            Self::F16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_f16(shape, data)
-            }
-            Self::Bf16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_bf16(shape, data)
-            }
-            Self::I64Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_i64(shape, data)
-            }
-            Self::U32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
-                WireTensor::from_u32(shape, data)
+        let wire = match self {
+            Self::Tensor { kind, wire } => {
+                slice_first_axis_range(wire, kind.element_size(), start, end, drop_axis)
            }
            Self::Passthrough(value) => return Ok(value.clone()),
        };
-        tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
+        wire.map(ProtocolKwargValue::Tensor)
    }
 }

@@ -208,44 +225,53 @@ pub(super) fn flat_range_for_index(
 /// Read a tensor value as per-image sizes for flat slicing.
 fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
    match tensor {
-        KwargValue::I64Tensor { data, .. } => data
-            .iter()
+        KwargValue::Tensor {
+            kind: TensorKind::I64,
+            wire,
+        } => raw_tensor_bytes(wire, size_of::<i64>())?
+            .chunks_exact(size_of::<i64>())
+            .map(|bytes| i64::from_ne_bytes(bytes.try_into().expect("exact int64 chunk")))
            .map(|value| {
-                usize::try_from(*value)
+                usize::try_from(value)
                    .map_err(|_| multimodal!("negative flat tensor size `{value}`"))
            })
            .collect(),
-        KwargValue::U32Tensor { data, .. } => {
-            Ok(data.iter().map(|value| *value as usize).collect())
-        }
+        KwargValue::Tensor {
+            kind: TensorKind::U32,
+            wire,
+        } => Ok(raw_tensor_bytes(wire, size_of::<u32>())?
+            .chunks_exact(size_of::<u32>())
+            .map(|bytes| u32::from_ne_bytes(bytes.try_into().expect("exact uint32 chunk")) as usize)
+            .collect()),
        _ => Err(multimodal!("flat tensor sizes must be int64 or uint32")),
    }
 }

 /// Slice a flat row-major tensor along its first axis.
-fn slice_first_axis_range<'a, T>(
-    shape: &[usize],
-    data: &'a [T],
+fn slice_first_axis_range(
+    tensor: &WireTensor,
+    element_size: usize,
    start: usize,
    end: usize,
    drop_axis: bool,
-) -> Result<(Vec<usize>, &'a [T])> {
+) -> Result<WireTensor> {
+    let shape = tensor.shape.as_slice();
+    raw_tensor_bytes(tensor, element_size)?;
    let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
    if start > end || end > first_dim {
        bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
    }
-    let expected_len = shape
-        .checked_numel()
-        .ok_or_else(|| multimodal!("tensor shape {shape:?} has too many elements"))?;
-    if expected_len != data.len() {
-        bail_multimodal!(
-            "tensor shape {shape:?} expects {expected_len} elements, got {}",
-            data.len()
-        );
-    }
-    let stride = shape[1..].iter().product::<usize>();
-    let data_start = start * stride;
-    let data_end = end * stride;
+    let stride = shape[1..]
+        .iter()
+        .try_fold(1usize, |acc, dim| acc.checked_mul(*dim))
+        .and_then(|stride| stride.checked_mul(element_size))
+        .ok_or_else(|| multimodal!("tensor shape {shape:?} byte stride overflowed usize"))?;
+    let data_start = start
+        .checked_mul(stride)
+        .ok_or_else(|| multimodal!("tensor slice start byte offset overflowed usize"))?;
+    let data_end = end
+        .checked_mul(stride)
+        .ok_or_else(|| multimodal!("tensor slice end byte offset overflowed usize"))?;
    let out_shape = if drop_axis {
        shape[1..].to_vec()
    } else {
@@ -253,7 +279,38 @@ fn slice_first_axis_range<'a, T>(
        shape[0] = end - start;
        shape
    };
-    Ok((out_shape, &data[data_start..data_end]))
+    let WireArrayData::RawView(data) = &tensor.data else {
+        return Err(multimodal!("cannot slice an aux tensor buffer"));
+    };
+    Ok(WireTensor::from_raw_bytes(
+        tensor.dtype.clone(),
+        out_shape,
+        data.slice(data_start..data_end),
+    ))
+}
+
+fn raw_tensor_bytes(tensor: &WireTensor, element_size: usize) -> Result<&[u8]> {
+    let WireArrayData::RawView(data) = &tensor.data else {
+        return Err(multimodal!("expected an inline tensor buffer"));
+    };
+    let expected_bytes = tensor
+        .shape
+        .checked_numel()
+        .and_then(|numel| numel.checked_mul(element_size))
+        .ok_or_else(|| {
+            multimodal!(
+                "tensor shape {:?} byte length overflowed usize",
+                tensor.shape
+            )
+        })?;
+    if expected_bytes != data.len() {
+        bail_multimodal!(
+            "tensor shape {:?} expects {expected_bytes} bytes, got {}",
+            tensor.shape,
+            data.len()
+        );
+    }
+    Ok(data)
 }

 #[cfg(test)]
@@ -262,28 +319,32 @@ mod tests {

    #[test]
    fn batched_wire_value_at_drops_first_axis() {
-        let value = KwargValue::F32Tensor {
-            data: vec![1.0, 2.0, 3.0, 4.0],
-            shape: vec![2, 2],
-        };
+        let data = vec![1.0_f32, 2.0, 3.0, 4.0];
+        let expected_ptr = data.as_ptr().cast::<u8>().wrapping_add(2 * size_of::<f32>());
+        let value = KwargValue::tensor(
+            TensorKind::F32,
+            WireTensor::from_f32(vec![2, 2], data).unwrap(),
+        );

        let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
            panic!("expected tensor");
        };

        assert_eq!(tensor.shape, vec![2]);
+        let raw_view = tensor.data.into_raw_view().unwrap();
+        assert_eq!(raw_view.as_ptr(), expected_ptr);
        assert_eq!(
-            tensor.data.into_raw_view().unwrap(),
+            raw_view,
            [3.0_f32, 4.0].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
        );
    }

    #[test]
    fn flat_wire_value_range_keeps_first_axis() {
-        let value = KwargValue::U32Tensor {
-            data: (0..10).collect(),
-            shape: vec![5, 2],
-        };
+        let value = KwargValue::tensor(
+            TensorKind::U32,
+            WireTensor::from_u32(vec![5, 2], (0..10_u32).collect()).unwrap(),
+        );

        let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
            panic!("expected tensor");
@@ -298,10 +359,10 @@ mod tests {

    #[test]
    fn flat_range_for_index_uses_size_tensor() {
-        let sizes = KwargValue::I64Tensor {
-            data: vec![2, 3, 4],
-            shape: vec![3],
-        };
+        let sizes = KwargValue::tensor(
+            TensorKind::I64,
+            WireTensor::from_i64(vec![3], vec![2_i64, 3, 4]).unwrap(),
+        );

        assert_eq!(
            flat_range_for_index(&sizes, "image_grid_thw", 1).unwrap(),
@@ -311,10 +372,11 @@ mod tests {

    #[test]
    fn slice_first_axis_range_errors_on_shape_data_mismatch() {
-        let error = slice_first_axis_range(&[2, 2], &[1.0_f32, 2.0, 3.0], 0, 1, true).unwrap_err();
+        let tensor = WireTensor::from_raw("float32", vec![2, 2], vec![0; 3 * size_of::<f32>()]);
+        let error = slice_first_axis_range(&tensor, size_of::<f32>(), 0, 1, true).unwrap_err();

        assert!(
-            matches!(error, Error::Multimodal(message) if message.contains("expects 4 elements"))
+            matches!(error, Error::Multimodal(message) if message.contains("expects 16 bytes"))
        );
    }

@@ -207,6 +207,7 @@ mod tests {
                Some("qwen3_vl".to_string()),
                files,
                Arc::new(qwen3_vl_tokenizer()),
+                std::collections::HashMap::new(),
            )
        };

@@ -23,6 +23,7 @@ use serde_with::{DefaultOnNull, OneOrMany, serde_as};
 use thiserror_ext::AsReport as _;
 use uuid::Uuid;
 use vllm_chat::ReasoningParserFactory;
+use vllm_chat::multimodal::MmLimitPerPrompt;
 use vllm_engine_core_client::TransportMode;
 use vllm_managed_engine::ManagedEngineConfig;
 use vllm_managed_engine::cli::{ManagedEngineArgs, repartition_managed_engine_args};
@@ -186,6 +187,17 @@ pub struct SharedRuntimeArgs {
    #[serde(default)]
    pub default_chat_template_kwargs: Option<HashMap<String, Value>>,

+    /// The maximum number of input items allowed per prompt for each
+    /// modality, as a JSON object (e.g. `{"image": 16, "video": 2}`).
+    ///
+    /// Also accepts the engine's configurable form
+    /// (e.g. `{"video": {"count": 1, "num_frames": 32}}`); the extra
+    /// profiling options are forwarded to the engine untouched.
+    /// Unspecified modalities are unlimited.
+    #[arg(long, value_parser = parse_json::<MmLimitPerPrompt>, value_name = "JSON", default_value = "{}")]
+    #[serde(default)]
+    pub limit_mm_per_prompt: MmLimitPerPrompt,
+
    /// The format to render message content within a chat template.
    ///
    /// * "auto" detects the format from the template
@@ -348,6 +360,18 @@ impl SharedRuntimeArgs {
            .expect("profiler config serialization should not fail")
    }

+    /// Return the per-modality limits as JSON for managed Python engine
+    /// forwarding, or `None` when nothing is configured.
+    ///
+    /// Round-tripping the parsed map rather than the raw argument keeps the
+    /// engine's own profiling options (`num_frames`, `width`, ...) intact.
+    pub fn limit_mm_per_prompt_json(&self) -> Option<String> {
+        (!self.limit_mm_per_prompt.is_empty()).then(|| {
+            serde_json::to_string(&self.limit_mm_per_prompt)
+                .expect("limit-mm-per-prompt serialization should not fail")
+        })
+    }
+
    /// Apply fallback logic for API key configuration from env variables.
    fn apply_env_api_key_fallback(&mut self) {
        if self.api_key.is_empty()
@@ -399,6 +423,7 @@ impl SharedRuntimeArgs {
            language_model_only: self.language_model_only,
            chat_template: self.chat_template,
            default_chat_template_kwargs: self.default_chat_template_kwargs,
+            limit_mm_per_prompt: self.limit_mm_per_prompt,
            chat_template_content_format: self.chat_template_content_format,
            max_logprobs: self.max_logprobs,
            api_server_options,
@@ -451,6 +476,7 @@ impl SharedRuntimeArgs {
            language_model_only: self.language_model_only,
            chat_template: self.chat_template,
            default_chat_template_kwargs: self.default_chat_template_kwargs,
+            limit_mm_per_prompt: self.limit_mm_per_prompt,
            chat_template_content_format: self.chat_template_content_format,
            max_logprobs: self.max_logprobs,
            api_server_options,
@@ -666,6 +692,7 @@ impl ServeArgs {
            self.runtime.disable_log_stats,
            self.runtime.shutdown_timeout,
            handshake_port,
+            self.runtime.limit_mm_per_prompt_json(),
        )
    }
 }
@@ -64,6 +64,7 @@ fn serve_args_forward_python_flags_with_separator() {
                        http_timeout_keep_alive: None,
                        chat_template: None,
                        default_chat_template_kwargs: None,
+                        limit_mm_per_prompt: {},
                        chat_template_content_format: Auto,
                        enable_log_requests: false,
                        enable_prompt_tokens_details: false,
@@ -762,6 +763,7 @@ fn frontend_args_accept_json() {
                        http_timeout_keep_alive: None,
                        chat_template: None,
                        default_chat_template_kwargs: None,
+                        limit_mm_per_prompt: {},
                        chat_template_content_format: Auto,
                        enable_log_requests: false,
                        enable_prompt_tokens_details: false,
@@ -1117,6 +1119,24 @@ fn frontend_args_json_rejects_malformed_json() {
    "#]].assert_eq(&error.to_string());
 }

+#[test]
+fn serve_args_reject_unsupported_modality_in_limit_mm_per_prompt() {
+    let error = Cli::try_parse_from([
+        "vllm-rs",
+        "serve",
+        "Qwen/Qwen3-0.6B",
+        "--limit-mm-per-prompt",
+        r#"{"unsupported_modality": 1}"#,
+    ])
+    .unwrap_err();
+
+    expect![[r#"
+        error: invalid value '{"unsupported_modality": 1}' for '--limit-mm-per-prompt <JSON>': invalid JSON object: unknown variant `unsupported_modality`, expected one of `image`, `audio`, `video` at line 1 column 23
+
+        For more information, try '--help'.
+    "#]].assert_eq(&error.to_string());
+}
+
 #[test]
 fn serve_args_reject_flags_before_model() {
    let error = Cli::try_parse_from(["vllm-rs", "serve", "--python", "python3", "Qwen/Qwen3-0.6B"])
@@ -1331,6 +1351,7 @@ fn serve_args_accept_handshake_aliases() {
                        http_timeout_keep_alive: None,
                        chat_template: None,
                        default_chat_template_kwargs: None,
+                        limit_mm_per_prompt: {},
                        chat_template_content_format: Auto,
                        enable_log_requests: false,
                        enable_prompt_tokens_details: false,
@@ -1474,6 +1495,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
            language_model_only: false,
            chat_template: None,
            default_chat_template_kwargs: None,
+            limit_mm_per_prompt: {},
            chat_template_content_format: Auto,
            max_logprobs: None,
            api_server_options: ApiServerOptions {
@@ -1558,6 +1580,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
            language_model_only: false,
            chat_template: None,
            default_chat_template_kwargs: None,
+            limit_mm_per_prompt: {},
            chat_template_content_format: Auto,
            max_logprobs: None,
            api_server_options: ApiServerOptions {
@@ -1660,6 +1683,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
            language_model_only: false,
            chat_template: None,
            default_chat_template_kwargs: None,
+            limit_mm_per_prompt: {},
            chat_template_content_format: Auto,
            max_logprobs: None,
            api_server_options: ApiServerOptions {
@@ -299,11 +299,6 @@ pub struct EngineUnsupportedArgs {
    )]
    pub kv_sharing_fast_prefill: Option<Unsupported>,

-    /// The maximum number of input items and options allowed per
-    /// prompt for each modality.
-    #[arg(long)]
-    pub limit_mm_per_prompt: Option<Unsupported>,
-
    /// Additional args passed to process media inputs, keyed by modalities.
    #[arg(long)]
    pub media_io_kwargs: Option<Unsupported>,
@@ -519,7 +519,7 @@ impl EngineCoreClient {
                "registered request to engine"
            );

-            self.inner.send_to_engine(&engine_id, EngineCoreRequestType::Add, &req).await?;
+            self.inner.send_request_to_engine(&engine_id, req).await?;
            Ok(())
        }
        .await;
@@ -6,6 +6,7 @@ use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};

 use arc_swap::ArcSwapOption;
+use bytes::Bytes;
 use parking_lot::Mutex;
 use thiserror_ext::AsReport as _;
 use tokio::runtime::Handle;
@@ -21,17 +22,30 @@ use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_outpu
 use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder};
 use crate::protocol::encode_msgpack;
 use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs};
-use crate::protocol::request::EngineCoreRequestType;
+use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
 use crate::protocol::stats::SchedulerStats;
 use crate::protocol::utility::UtilityOutput;
 use crate::transport::{ConnectedEngine, EngineId};
 use crate::{Error, Result, transport};

+const MSGPACK_ZERO_COPY_THRESHOLD_ENV: &str = "VLLM_MSGPACK_ZERO_COPY_THRESHOLD";
+const DEFAULT_MSGPACK_ZERO_COPY_THRESHOLD: usize = 256;
+
+fn msgpack_zero_copy_threshold() -> usize {
+    std::env::var(MSGPACK_ZERO_COPY_THRESHOLD_ENV)
+        .ok()
+        .and_then(|value| value.parse().ok())
+        .unwrap_or(DEFAULT_MSGPACK_ZERO_COPY_THRESHOLD)
+}
+
 pub(crate) struct ClientInner {
    input_send: RouterSendHalf,
    /// The runtime handle used for sending messages to the engine.
    handle: Handle,
    model_name: String,
+    /// Per-tensor byte threshold loaded from env variable
+    /// `VLLM_MSGPACK_ZERO_COPY_THRESHOLD` when this inner client is created.
+    msgpack_zero_copy_threshold: usize,
    scheduler_stats_recorder: SchedulerStatsRecorder,
    request_reg: Mutex<RequestRegistry>,
    utility_reg: Mutex<UtilityRegistry>,
@@ -54,6 +68,7 @@ impl ClientInner {
            input_send,
            handle,
            model_name,
+            msgpack_zero_copy_threshold: msgpack_zero_copy_threshold(),
            scheduler_stats_recorder,
            request_reg: Mutex::new(RequestRegistry::new(engines)),
            utility_reg: Mutex::new(UtilityRegistry::default()),
@@ -229,9 +244,29 @@ impl ClientInner {
    where
        T: serde::Serialize + std::fmt::Debug,
    {
-        // TODO: for `EngineCoreRequest`, split outbound tensor raw views into aux
-        // frames instead of always producing a single msgpack frame.
-        let payload = encode_msgpack(payload)?;
+        let payload = Bytes::from(encode_msgpack(payload)?);
+        self.send_encoded_to_engine(engine_id, request_type, payload, Vec::new()).await
+    }
+
+    /// Send an add request, moving large tensor buffers into auxiliary frames.
+    pub async fn send_request_to_engine(
+        &self,
+        engine_id: &EngineId,
+        mut payload: EngineCoreRequest,
+    ) -> Result<()> {
+        let aux_frames = payload.extract_aux_frames(self.msgpack_zero_copy_threshold);
+        let payload = Bytes::from(encode_msgpack(&payload)?);
+        self.send_encoded_to_engine(engine_id, EngineCoreRequestType::Add, payload, aux_frames)
+            .await
+    }
+
+    async fn send_encoded_to_engine(
+        &self,
+        engine_id: &EngineId,
+        request_type: EngineCoreRequestType,
+        payload: Bytes,
+        aux_frames: Vec<Bytes>,
+    ) -> Result<()> {
        let mut input_send = self.input_send.clone();
        let engine_id = engine_id.clone();

@@ -242,6 +277,7 @@ impl ClientInner {
                    &engine_id,
                    request_type.to_frame(),
                    payload,
+                    aux_frames,
                )
                .await
            })
@@ -209,17 +209,17 @@ impl WireLogprobs {
            logprob_token_ids: WireNdArray {
                dtype: "<i8".to_string(),
                shape: vec![rows, cols],
-                data: WireArrayData::RawView(token_ids),
+                data: WireArrayData::RawView(token_ids.into()),
            },
            logprobs: WireNdArray {
                dtype: "<f4".to_string(),
                shape: vec![rows, cols],
-                data: WireArrayData::RawView(logprobs),
+                data: WireArrayData::RawView(logprobs.into()),
            },
            token_ranks: WireNdArray {
                dtype: "<i8".to_string(),
                shape: vec![rows],
-                data: WireArrayData::RawView(token_ranks),
+                data: WireArrayData::RawView(token_ranks.into()),
            },
            cu_num_generated_tokens: None,
        })
@@ -4,6 +4,7 @@
 use std::io::Cursor;

 use byteorder::{BigEndian, LittleEndian, NativeEndian, ReadBytesExt};
+use bytes::Bytes;
 use itertools::Itertools as _;

 use crate::error::{Error, Result, ext_value_decode};
@@ -126,7 +127,7 @@ pub(super) fn decode_array_metadata<Frame>(
    field: &str,
    frames: &[Frame],
    expected_scalars: &[ScalarType],
-) -> Result<(Vec<usize>, Vec<u8>, ScalarType, Endianness)>
+) -> Result<(Vec<usize>, Bytes, ScalarType, Endianness)>
 where
    Frame: AsRef<[u8]>,
 {
@@ -171,7 +172,7 @@ pub(super) fn resolve_array_bytes<Frame>(
    value: WireArrayData,
    field: &str,
    frames: &[Frame],
-) -> Result<Vec<u8>>
+) -> Result<Bytes>
 where
    Frame: AsRef<[u8]>,
 {
@@ -187,7 +188,7 @@ where
                    ),
                )
            })?;
-            Ok(frame.as_ref().to_vec())
+            Ok(Bytes::copy_from_slice(frame.as_ref()))
        }
    }
 }
@@ -3,6 +3,7 @@

 use std::collections::BTreeMap;

+use bytes::Bytes;
 use serde::{Deserialize, Serialize};
 use serde_tuple::{Deserialize_tuple, Serialize_tuple};

@@ -105,6 +106,44 @@ pub enum MmKwargValue {
    List(Vec<MmKwargValue>),
 }

+impl MmFeatureSpec {
+    /// Extract large tensor buffers from this feature in serialized field order.
+    pub(crate) fn extract_aux_frames(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        if let Some(data) = &mut self.data {
+            for elem in data.values_mut() {
+                elem.extract_aux_frames(aux_frames, threshold);
+            }
+        }
+        if let Some(is_embed) = &mut self.mm_position.is_embed {
+            is_embed.extract_aux_frame(aux_frames, threshold);
+        }
+    }
+}
+
+impl MmFieldElem {
+    /// Extract large tensor buffers from this field element.
+    fn extract_aux_frames(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        if let Some(data) = &mut self.data {
+            data.extract_aux_frames(aux_frames, threshold);
+        }
+    }
+}
+
+impl MmKwargValue {
+    /// Recursively extract large tensor buffers from this nested value.
+    fn extract_aux_frames(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        match self {
+            Self::Tensor(tensor) => tensor.extract_aux_frame(aux_frames, threshold),
+            Self::List(values) => {
+                for value in values {
+                    value.extract_aux_frames(aux_frames, threshold);
+                }
+            }
+            Self::Int(_) | Self::Float(_) => {}
+        }
+    }
+}
+
 /// Defines how to interpret tensor data belonging to a keyword argument for
 /// `MultiModalKwargsItems`, and vice versa.
 ///
@@ -137,6 +137,17 @@ impl EngineCoreRequest {
        }
        Ok(())
    }
+
+    /// Extract large request tensors into ordered auxiliary frames.
+    pub(crate) fn extract_aux_frames(&mut self, threshold: usize) -> Vec<Bytes> {
+        let mut aux_frames = Vec::new();
+        if let Some(features) = &mut self.mm_features {
+            for feature in features {
+                feature.extract_aux_frames(&mut aux_frames, threshold);
+            }
+        }
+        aux_frames
+    }
 }

 #[cfg(test)]
@@ -144,9 +155,15 @@ mod tests {
    use rmpv::Value;

    use super::*;
+    use crate::protocol::multimodal::{
+        MmBatchedField, MmFeatureSpec, MmField, MmFieldElem, MmKwargValue, PlaceholderRange,
+    };
    use crate::protocol::sampling::EngineCoreSamplingParams;
+    use crate::protocol::tensor::{WireArrayData, WireTensor};
    use crate::protocol::{decode_value, encode_msgpack};

+    const AUX_FRAME_THRESHOLD: usize = 256;
+
    #[test]
    fn engine_core_request_serializes_as_full_array() {
        let request = EngineCoreRequest {
@@ -175,4 +192,84 @@ mod tests {
        assert_eq!(array[10], Value::Nil);
        assert_eq!(array[11], Value::from(7));
    }
+
+    #[test]
+    fn engine_core_request_extracts_large_nested_tensors_in_wire_order() {
+        let inline = vec![1_u8; AUX_FRAME_THRESHOLD - 1];
+        let first_aux = vec![2_u8; AUX_FRAME_THRESHOLD];
+        let second_aux = vec![3_u8; AUX_FRAME_THRESHOLD + 1];
+        let first_aux_ptr = first_aux.as_ptr();
+        let second_aux_ptr = second_aux.as_ptr();
+        let mut request = EngineCoreRequest {
+            mm_features: Some(vec![MmFeatureSpec {
+                data: Some(BTreeMap::from([
+                    (
+                        "inline".to_string(),
+                        MmFieldElem {
+                            data: Some(MmKwargValue::Tensor(WireTensor::from_raw(
+                                "uint8",
+                                vec![inline.len()],
+                                inline,
+                            ))),
+                            field: MmField::Batched(MmBatchedField { keep_on_cpu: false }),
+                        },
+                    ),
+                    (
+                        "nested".to_string(),
+                        MmFieldElem {
+                            data: Some(MmKwargValue::List(vec![
+                                MmKwargValue::Int(7),
+                                MmKwargValue::Tensor(WireTensor::from_raw(
+                                    "uint8",
+                                    vec![first_aux.len()],
+                                    first_aux,
+                                )),
+                            ])),
+                            field: MmField::Batched(MmBatchedField { keep_on_cpu: false }),
+                        },
+                    ),
+                ])),
+                modality: "image".to_string(),
+                identifier: "id".to_string(),
+                mm_position: PlaceholderRange {
+                    offset: 0,
+                    length: second_aux.len(),
+                    is_embed: Some(WireTensor::from_raw(
+                        "bool",
+                        vec![second_aux.len()],
+                        second_aux,
+                    )),
+                },
+                mm_hash: None,
+            }]),
+            ..EngineCoreRequest::default()
+        };
+
+        let aux_frames = request.extract_aux_frames(AUX_FRAME_THRESHOLD);
+
+        assert_eq!(aux_frames.len(), 2);
+        assert_eq!(aux_frames[0].as_ptr(), first_aux_ptr);
+        assert_eq!(aux_frames[1].as_ptr(), second_aux_ptr);
+        let feature = &request.mm_features.as_ref().unwrap()[0];
+        let MmKwargValue::Tensor(inline) =
+            feature.data.as_ref().unwrap()["inline"].data.as_ref().unwrap()
+        else {
+            panic!("expected inline tensor");
+        };
+        assert!(matches!(inline.data, WireArrayData::RawView(_)));
+        let MmKwargValue::List(nested) =
+            feature.data.as_ref().unwrap()["nested"].data.as_ref().unwrap()
+        else {
+            panic!("expected nested tensor list");
+        };
+        let MmKwargValue::Tensor(nested_tensor) = &nested[1] else {
+            panic!("expected nested tensor");
+        };
+        assert_eq!(nested_tensor.data, WireArrayData::AuxIndex(1));
+        assert_eq!(
+            feature.mm_position.is_embed.as_ref().unwrap().data,
+            WireArrayData::AuxIndex(2)
+        );
+        assert!(request.extract_aux_frames(AUX_FRAME_THRESHOLD).is_empty());
+    }
 }
@@ -1,7 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-use bytemuck::allocation::pod_collect_to_vec;
+use bytemuck::{Pod, cast_slice};
+use bytes::Bytes;
 use enum_as_inner::EnumAsInner;
 use half::{bf16, f16};
 use rmpv::Value;
@@ -20,6 +21,21 @@ struct MsgpackExtRef<'a>((i8, ByteSlice<'a>));

 struct ByteSlice<'a>(&'a [u8]);

+struct PodVec<T: Pod>(Vec<T>);
+
+impl<T: Pod> AsRef<[u8]> for PodVec<T> {
+    fn as_ref(&self) -> &[u8] {
+        cast_slice(&self.0)
+    }
+}
+
+fn bytes_from_pod_vec<T>(data: Vec<T>) -> Bytes
+where
+    T: Pod + Send + 'static,
+{
+    Bytes::from_owner(PodVec(data))
+}
+
 impl Serialize for ByteSlice<'_> {
    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
@@ -55,58 +71,63 @@ pub struct WireNdArray {

 impl WireNdArray {
    /// Build a float32 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_f32(shape: Vec<usize>, data: impl AsRef<[f32]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> {
        validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "float32".to_string(),
+        Ok(Self::from_raw_bytes(
+            "float32",
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
    }

    /// Build a float16 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_f16(shape: Vec<usize>, data: impl AsRef<[f16]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> {
        validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "float16".to_string(),
+        Ok(Self::from_raw_bytes(
+            "float16",
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
    }

    /// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_bf16(shape: Vec<usize>, data: impl AsRef<[bf16]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> {
        validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "bfloat16".to_string(),
+        Ok(Self::from_raw_bytes(
+            "bfloat16",
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
    }

    /// Build an int64 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_i64(shape: Vec<usize>, data: impl AsRef<[i64]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> {
        validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "int64".to_string(),
+        Ok(Self::from_raw_bytes(
+            "int64",
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
    }

    /// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_u32(shape: Vec<usize>, data: impl AsRef<[u32]>) -> Result<Self, String> {
-        let data = data.as_ref();
+    ///
+    /// Takes ownership of the backing buffer without copying its data.
+    pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> {
        validate_element_count(&shape, data.len())?;
-        Ok(Self {
-            dtype: "uint32".to_string(),
+        Ok(Self::from_raw_bytes(
+            "uint32",
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(data)),
-        })
+            bytes_from_pod_vec(data),
+        ))
    }

    /// Build a bool tensor/ndarray backed by raw-view bytes.
@@ -118,7 +139,9 @@ impl WireNdArray {
        Ok(Self {
            dtype: "bool".to_string(),
            shape,
-            data: WireArrayData::RawView(data.iter().map(|value| u8::from(*value)).collect()),
+            data: WireArrayData::RawView(Bytes::from(
+                data.into_iter().map(u8::from).collect::<Vec<_>>(),
+            )),
        })
    }

@@ -127,12 +150,22 @@ impl WireNdArray {
    /// Use this as an escape hatch when the caller already owns bytes that
    /// match the requested `dtype` and `shape`.
    pub fn from_raw(dtype: impl Into<String>, shape: Vec<usize>, data: Vec<u8>) -> Self {
+        Self::from_raw_bytes(dtype, shape, Bytes::from(data))
+    }
+
+    /// Build a tensor/ndarray from an owned immutable raw-view buffer.
+    pub fn from_raw_bytes(dtype: impl Into<String>, shape: Vec<usize>, data: Bytes) -> Self {
        Self {
            dtype: dtype.into(),
            shape,
            data: WireArrayData::RawView(data),
        }
    }
+
+    /// Move a sufficiently large inline buffer into the ordered auxiliary-frame list.
+    pub(crate) fn extract_aux_frame(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        self.data.extract_aux_frame(aux_frames, threshold);
+    }
 }

 /// Validate that the number of elements implied by the shape matches the length
@@ -170,7 +203,25 @@ pub enum WireArrayData {
    /// stored.
    AuxIndex(usize),
    /// The raw bytes of this array/tensor.
-    RawView(Vec<u8>),
+    RawView(Bytes),
+}
+
+impl WireArrayData {
+    /// Replace a sufficiently large raw view with its one-based auxiliary-frame index.
+    fn extract_aux_frame(&mut self, aux_frames: &mut Vec<Bytes>, threshold: usize) {
+        let Self::RawView(bytes) = self else {
+            return;
+        };
+        if bytes.len() < threshold {
+            return;
+        }
+
+        let index = aux_frames.len() + 1;
+        let bytes = std::mem::replace(self, Self::AuxIndex(index))
+            .into_raw_view()
+            .expect("raw view was matched above");
+        aux_frames.push(bytes);
+    }
 }

 impl<'de> Deserialize<'de> for WireArrayData {
@@ -180,7 +231,9 @@ impl<'de> Deserialize<'de> for WireArrayData {
    {
        let value = Value::deserialize(deserializer)?;
        match value {
-            Value::Ext(tag, bytes) if tag == CUSTOM_TYPE_RAW_VIEW => Ok(Self::RawView(bytes)),
+            Value::Ext(tag, bytes) if tag == CUSTOM_TYPE_RAW_VIEW => {
+                Ok(Self::RawView(Bytes::from(bytes)))
+            }
            Value::Ext(tag, _) => Err(serde::de::Error::custom(format!(
                "unsupported extension type code {tag}"
            ))),
@@ -201,9 +254,6 @@ impl Serialize for WireArrayData {
    where
        S: Serializer,
    {
-        // TODO: outbound request serialization currently only supports inline
-        // raw-view bytes. Emitting aux frames needs transport-level plumbing;
-        // serializing `AuxIndex` here only preserves an already-built reference.
        match self {
            Self::AuxIndex(index) => serializer.serialize_u64(*index as u64),
            Self::RawView(bytes) => {
@@ -221,7 +271,7 @@ mod tests {
    fn raw_view_serializes_as_msgpack_ext() {
        let bytes = vec![1, 2, 3, 4];
        let encoded =
-            rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone())).expect("encode");
+            rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone().into())).expect("encode");
        let expected = rmp_serde::to_vec_named(&Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone()))
            .expect("encode expected");

@@ -234,11 +284,15 @@ mod tests {

    #[test]
    fn constructors_build_raw_view_tensors() {
-        let f32_tensor = WireNdArray::from_f32(vec![2], vec![1.0, 2.5]).unwrap();
+        let f32_data = vec![1.0, 2.5];
+        let f32_data_ptr = f32_data.as_ptr().cast::<u8>();
+        let f32_tensor = WireNdArray::from_f32(vec![2], f32_data).unwrap();
        assert_eq!(f32_tensor.dtype, "float32");
        assert_eq!(f32_tensor.shape, vec![2]);
+        let f32_raw_view = f32_tensor.data.into_raw_view().expect("raw view");
+        assert_eq!(f32_raw_view.as_ptr(), f32_data_ptr);
        assert_eq!(
-            f32_tensor.data.into_raw_view().expect("raw view"),
+            f32_raw_view,
            [1.0_f32, 2.5].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
        );

@@ -258,15 +312,15 @@ mod tests {
        let i64_tensor = WireNdArray::from_i64(vec![1], vec![-7]).unwrap();
        assert_eq!(i64_tensor.dtype, "int64");
        assert_eq!(
-            i64_tensor.data.into_raw_view().expect("raw view"),
-            (-7_i64).to_ne_bytes()
+            i64_tensor.data.into_raw_view().expect("raw view").as_ref(),
+            (-7_i64).to_ne_bytes().as_ref()
        );

        let u32_tensor = WireNdArray::from_u32(vec![1], vec![42]).unwrap();
        assert_eq!(u32_tensor.dtype, "uint32");
        assert_eq!(
-            u32_tensor.data.into_raw_view().expect("raw view"),
-            42_u32.to_ne_bytes()
+            u32_tensor.data.into_raw_view().expect("raw view").as_ref(),
+            42_u32.to_ne_bytes().as_ref()
        );

        let bool_tensor = WireNdArray::from_bool(vec![2], vec![false, true]).unwrap();
@@ -32,7 +32,7 @@ use crate::protocol::output::{
 use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType};
 use crate::protocol::sampling::EngineCoreSamplingParams;
 use crate::protocol::stats::SchedulerStats;
-use crate::protocol::tensor::WireTensor;
+use crate::protocol::tensor::{WireArrayData, WireTensor};
 use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
 use crate::test_utils::{
    IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets,
@@ -1728,6 +1728,86 @@ async fn client_decodes_multipart_logprob_outputs() {
    client.shutdown().await.unwrap();
 }

+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn client_sends_large_multimodal_tensor_as_aux_frame() {
+    init_tracing();
+    let ipc = IpcNamespace::new().unwrap();
+    let handshake_address = ipc.handshake_endpoint();
+    let engine_id = b"engine-multimodal-aux".to_vec();
+    let tensor_data = (0..64).map(|value| value as f32).collect::<Vec<_>>();
+    let expected_bytes =
+        tensor_data.iter().flat_map(|value| value.to_ne_bytes()).collect::<Vec<_>>();
+    let mut request = sample_multimodal_request();
+    request.mm_features.as_mut().unwrap()[0]
+        .data
+        .as_mut()
+        .unwrap()
+        .get_mut("pixel_values")
+        .unwrap()
+        .data = Some(MmKwargValue::Tensor(
+        WireTensor::from_f32(vec![64], tensor_data).unwrap(),
+    ));
+
+    let (shutdown_tx, engine_task) = spawn_mock_engine_task(
+        handshake_address.clone(),
+        engine_id.clone(),
+        move |dealer, push| {
+            Box::pin(async move {
+                let add = recv_engine_message(dealer).await;
+                assert_eq!(add.len(), 3);
+                assert_eq!(add[0].as_ref(), &[0x00]);
+                let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
+                let MmKwargValue::Tensor(tensor) =
+                    request.mm_features.as_ref().unwrap()[0].data.as_ref().unwrap()["pixel_values"]
+                        .data
+                        .as_ref()
+                        .unwrap()
+                else {
+                    panic!("expected tensor");
+                };
+                assert_eq!(tensor.data, WireArrayData::AuxIndex(1));
+                assert_eq!(add[2].as_ref(), expected_bytes);
+
+                send_outputs(
+                    push,
+                    RequestBatchOutputs {
+                        outputs: vec![request_output(
+                            "req-mm",
+                            vec![],
+                            Some(EngineCoreFinishReason::Length),
+                        )],
+                        finished_requests: Some(BTreeSet::from(["req-mm".to_string()])),
+                        ..Default::default()
+                    }
+                    .into(),
+                )
+                .await;
+            })
+        },
+    );
+
+    let client = connect_client_with_ipc(
+        handshake_test_config(
+            handshake_address,
+            1,
+            "test-model",
+            Duration::from_secs(2),
+            0,
+            None,
+        ),
+        &ipc,
+    )
+    .await;
+
+    let outputs = client.call(request).await.unwrap().collect::<Vec<_>>().await;
+    assert_eq!(outputs.len(), 1);
+    assert!(outputs[0].is_ok());
+
+    let _ = shutdown_tx.send(());
+    engine_task.await.unwrap();
+    client.shutdown().await.unwrap();
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() {
    init_tracing();
@@ -512,14 +512,14 @@ pub async fn send_message(
    input_send: &mut RouterSendHalf,
    engine_id: &EngineId,
    request_type: Bytes,
-    payload: Vec<u8>,
+    payload: Bytes,
+    aux_frames: Vec<Bytes>,
 ) -> Result<()> {
-    let message = ZmqMessage::try_from(vec![
-        engine_id.to_frame(),
-        request_type,
-        Bytes::from(payload),
-    ])
-    .expect("router messages must contain identity and payload");
+    let mut frames = Vec::with_capacity(3 + aux_frames.len());
+    frames.extend([engine_id.to_frame(), request_type, payload]);
+    frames.extend(aux_frames);
+    let message =
+        ZmqMessage::try_from(frames).expect("router messages must contain identity and payload");

    trace!(
        ?engine_id,
@@ -91,6 +91,7 @@ impl ManagedEngineArgs {
        disable_log_stats: bool,
        shutdown_timeout: u64,
        handshake_port: u16,
+        limit_mm_per_prompt: Option<String>,
    ) -> ManagedEngineConfig {
        let mut python_args = self.python_args;
        // Manually forward some args to the Python engine.
@@ -126,6 +127,10 @@ impl ManagedEngineArgs {
            python_args.push("--data-parallel-size-local".to_string());
            python_args.push(data_parallel_size_local.to_string());
        }
+        if let Some(limit_mm_per_prompt) = limit_mm_per_prompt {
+            python_args.push("--limit-mm-per-prompt".to_string());
+            python_args.push(limit_mm_per_prompt);
+        }

        ManagedEngineConfig {
            python: self.python,
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+use std::collections::HashMap;
 use std::time::Duration;

 use anyhow::{Context, Result, bail};
@@ -70,6 +71,7 @@ async fn main() -> Result<()> {
        language_model_only: false,
        chat_template: None,
        default_chat_template_kwargs: None,
+        limit_mm_per_prompt: HashMap::new(),
        chat_template_content_format: ChatTemplateContentFormatOption::Auto,
        max_logprobs: None,
        api_server_options: ApiServerOptions::default(),
@@ -10,6 +10,7 @@ use axum::http::{HeaderName, HeaderValue, Method};
 use educe::Educe;
 use serde::Serialize;
 use serde_json::Value;
+use vllm_chat::multimodal::MmLimitPerPrompt;
 use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection};
 use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode};

@@ -184,6 +185,9 @@ pub struct Config {
    pub chat_template: Option<String>,
    /// Server-default keyword arguments merged into every chat-template render.
    pub default_chat_template_kwargs: Option<HashMap<String, Value>>,
+    /// Maximum number of input items allowed per prompt for each modality.
+    /// Unspecified modalities are unlimited.
+    pub limit_mm_per_prompt: MmLimitPerPrompt,
    /// How to serialize `message.content` for chat-template rendering.
    pub chat_template_content_format: ChatTemplateContentFormatOption,
    /// Optional maximum number of top log probabilities accepted by the
@@ -101,6 +101,7 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
                .default_chat_template_kwargs
                .clone()
                .unwrap_or_default(),
+            limit_mm_per_prompt: config.limit_mm_per_prompt.clone(),
        },
    )
    .await
@@ -577,6 +577,12 @@ fn render_fake_content(content: &ChatContent, placeholder: &str) -> vllm_chat::R
 }

 fn qwen_multimodal_model_info() -> vllm_chat::multimodal::MultimodalModelInfo {
+    qwen_multimodal_model_info_with_limits(std::collections::HashMap::new())
+}
+
+fn qwen_multimodal_model_info_with_limits(
+    limit_mm_per_prompt: vllm_chat::multimodal::MmLimitPerPrompt,
+) -> vllm_chat::multimodal::MultimodalModelInfo {
    let config_path = std::env::temp_dir().join(format!(
        "vllm-server-qwen-config-{}.json",
        uuid::Uuid::new_v4()
@@ -594,6 +600,7 @@ fn qwen_multimodal_model_info() -> vllm_chat::multimodal::MultimodalModelInfo {
            ..Default::default()
        },
        Arc::new(fake_chat_tokenizer()),
+        limit_mm_per_prompt,
    )
    .expect("load multimodal info")
    .expect("qwen multimodal info is registered");
@@ -2307,6 +2314,74 @@ async fn non_stream_chat_image_url_reaches_engine_mm_features() {
    assert_eq!(json["choices"][0]["message"]["content"], "hi");
 }

+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn non_stream_chat_rejects_when_image_count_exceeds_limit_mm_per_prompt() {
+    // The request is rejected by `--limit-mm-per-prompt` validation before
+    // ever reaching the engine, so the mock engine task is never awaited.
+    let (chat, _engine_task) = test_models_with_engine_outputs_and_backend(
+        b"engine-openai-mm-limit",
+        default_stream_output_specs(),
+        Arc::new(FakeChatBackend::with_multimodal_model_info(
+            qwen_multimodal_model_info_with_limits(std::collections::HashMap::from([(
+                vllm_chat::multimodal::MmLimitModality::Image,
+                vllm_chat::multimodal::MmLimitSpec::Count(1),
+            )])),
+        )),
+    )
+    .await;
+    let app = build_router(Arc::new(AppState::new(
+        vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
+        chat,
+    )));
+
+    let response = app
+        .clone()
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/chat/completions")
+                .header("content-type", "application/json")
+                .body(Body::from(
+                    json!({
+                        "model": "Qwen/Qwen1.5-0.5B-Chat",
+                        "stream": false,
+                        "messages": [{
+                            "role": "user",
+                            "content": [
+                                {"type": "text", "text": "describe "},
+                                {
+                                    "type": "image_url",
+                                    "image_url": {
+                                        "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII="
+                                    }
+                                },
+                                {
+                                    "type": "image_url",
+                                    "image_url": {
+                                        "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII="
+                                    }
+                                }
+                            ]
+                        }]
+                    })
+                    .to_string(),
+                ))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+
+    assert_eq!(response.status(), StatusCode::BAD_REQUEST);
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
+    assert_eq!(json["error"]["type"], "invalid_request_error");
+    assert_eq!(
+        json["error"]["message"],
+        "At most 1 image(s) may be provided in one prompt."
+    );
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn non_stream_chat_includes_logprobs_and_prompt_logprobs() {
@@ -1122,7 +1122,7 @@ if _is_cuda() or _is_hip():
    # copying the relevant .py files from the source repository.
    ext_modules.append(CMakeExtension(name="vllm.triton_kernels", optional=True))

-if sys.version_info >= (3, 11):
+if not _is_xpu() and sys.version_info >= (3, 11):
    ext_modules.append(CMakeExtension(name="vllm.spinloop"))
    ext_modules.append(CMakeExtension(name="vllm.fs_io_C"))

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the EAGLE draft ``max_position_embeddings`` override (#48894).
+
+EAGLE drafts share the target's positional space, but some draft
+checkpoints (e.g. ``yuhuili/EAGLE3-LLaMA3.1-Instruct-8B``) ship a
+``max_position_embeddings`` (2048) far smaller than the target's context.
+That value sizes the draft's rotary ``cos_sin_cache`` while the proposer
+feeds positions up to the target's ``max_model_len``, so the cache gather
+goes out of bounds — a device-side assert under torch.compile and silent
+garbage reads in eager mode. ``SpeculativeConfig`` must raise the draft's
+value to the target's ``max_model_len``, with a log, for the eagle/eagle3
+methods only.
+"""
+
+import logging
+
+import pytest
+from transformers import PretrainedConfig
+
+from vllm.config.model import ModelConfig
+from vllm.config.parallel import ParallelConfig
+from vllm.config.speculative import SpeculativeConfig
+
+# All repos are public; only config/tokenizer-config files are fetched.
+EAGLE3_DRAFT = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"  # max_position_embeddings=2048
+LLAMA3_TARGET = "unsloth/Meta-Llama-3.1-8B-Instruct"  # max_position_embeddings=131072
+AR_MODEL = "JackFram/llama-68m"  # max_position_embeddings=2048
+
+_LOGGER = "vllm.config.speculative"
+_OVERRIDE_MSG = "Overriding draft model max_position_embeddings"
+
+
+@pytest.fixture
+def vllm_caplog(caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch):
+    """Make caplog see vLLM logger records (vLLM sets propagate=False)."""
+    monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True)
+    with caplog.at_level(logging.INFO, logger=_LOGGER):
+        yield caplog
+
+
+def _override_logged(caplog: pytest.LogCaptureFixture) -> bool:
+    return any(_OVERRIDE_MSG in record.getMessage() for record in caplog.records)
+
+
+@pytest.mark.cpu_test
+def test_override_raises_smaller_value(vllm_caplog: pytest.LogCaptureFixture):
+    hf_config = PretrainedConfig(max_position_embeddings=2048)
+    SpeculativeConfig._maybe_override_draft_max_position_embeddings(
+        hf_config, target_max_model_len=8192
+    )
+    assert hf_config.max_position_embeddings == 8192
+    assert _override_logged(vllm_caplog)
+
+
+@pytest.mark.cpu_test
+def test_override_keeps_sufficient_value(vllm_caplog: pytest.LogCaptureFixture):
+    hf_config = PretrainedConfig(max_position_embeddings=8192)
+    SpeculativeConfig._maybe_override_draft_max_position_embeddings(
+        hf_config, target_max_model_len=8192
+    )
+    assert hf_config.max_position_embeddings == 8192
+    assert not _override_logged(vllm_caplog)
+
+
+@pytest.mark.cpu_test
+def test_override_ignores_missing_attribute(vllm_caplog: pytest.LogCaptureFixture):
+    hf_config = PretrainedConfig()
+    hf_config.__dict__.pop("max_position_embeddings", None)
+    SpeculativeConfig._maybe_override_draft_max_position_embeddings(
+        hf_config, target_max_model_len=8192
+    )
+    assert not hasattr(hf_config, "max_position_embeddings")
+    assert not _override_logged(vllm_caplog)
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize("method", ["eagle", "eagle3"])
+def test_eagle_draft_inherits_target_max_model_len(
+    method: str, vllm_caplog: pytest.LogCaptureFixture
+):
+    target_model_config = ModelConfig(LLAMA3_TARGET)
+    assert target_model_config.max_model_len > 2048
+    speculative_config = SpeculativeConfig(
+        target_model_config=target_model_config,
+        target_parallel_config=ParallelConfig(),
+        model=EAGLE3_DRAFT,
+        method=method,
+        num_speculative_tokens=3,
+    )
+    draft_hf_config = speculative_config.draft_model_config.hf_config
+    assert draft_hf_config.max_position_embeddings == target_model_config.max_model_len
+    assert _override_logged(vllm_caplog)
+
+
+@pytest.mark.cpu_test
+def test_independent_draft_model_keeps_its_own_limit(
+    vllm_caplog: pytest.LogCaptureFixture,
+):
+    """An independent AR draft may genuinely have a smaller context than the
+    target; its max_position_embeddings must not be resized."""
+    target_model_config = ModelConfig(
+        AR_MODEL, hf_overrides={"max_position_embeddings": 8192}
+    )
+    assert target_model_config.max_model_len == 8192
+    speculative_config = SpeculativeConfig(
+        target_model_config=target_model_config,
+        target_parallel_config=ParallelConfig(),
+        model=AR_MODEL,
+        method="draft_model",
+        num_speculative_tokens=3,
+    )
+    draft_hf_config = speculative_config.draft_model_config.hf_config
+    assert draft_hf_config.max_position_embeddings == 2048
+    assert not _override_logged(vllm_caplog)
@@ -2,9 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Equivalence test for ``precopy_mamba_align_fused_kernel``.

-The V2 "align" pre-copy must migrate mamba state across block boundaries with
-byte-identical semantics to the V1 copy specs (``get_conv_copy_spec`` /
-``get_temporal_copy_spec``):
+The fused "align" pre-copy must migrate mamba state across block boundaries
+with byte-identical semantics to the scalar V1 copy specs
+(``get_conv_copy_spec`` / ``get_temporal_copy_spec``):

 * conv state (SD layout, conv_width > 0): shift the sliding window by
  ``token_bias`` tokens -- ``state[bt[src_col], token_bias:]`` ->
@@ -14,20 +14,27 @@ byte-identical semantics to the V1 copy specs (``get_conv_copy_spec`` /
  ``state[bt[dst_col]]``.

 The kernel must also no-op when ``src_col < 0`` (fresh request) or
-``src_col == dst_col`` (no boundary crossed).
+``src_col == dst_col`` (no boundary crossed). V2 callers pass an explicit
+``idx_mapping``; V1 align preprocessing launches in batch order with
+``idx_mapping=None``.
 """

 from __future__ import annotations

+from types import SimpleNamespace
+
+import numpy as np
 import torch

+from vllm.model_executor.layers.mamba import mamba_utils as layer_mamba_utils
 from vllm.platforms import current_platform
+from vllm.v1.worker import mamba_utils as worker_mamba_utils
 from vllm.v1.worker.mamba_utils import precopy_mamba_align_fused_kernel

 try:
    import pytest

-    pytestmark = pytest.mark.skipif(
+    _cuda_required = pytest.mark.skipif(
        not current_platform.is_cuda(),
        reason="precopy_mamba_align_fused_kernel needs CUDA/Triton",
    )
@@ -35,6 +42,9 @@ try:
 except ModuleNotFoundError:  # allow running directly as ``python <thisfile>``
    pytest = None

+    def _cuda_required(fn):
+        return fn
+
    def _parametrize(_name, _values):
        def _deco(fn):
            return fn
@@ -49,13 +59,20 @@ SSM_SHAPE = (4, 16, 16)
 MAX_COLS = 8


-def _build_state(num_blocks, device):
+def _build_state(num_blocks, device, conv_state_dim_first):
    """Per-layer (conv SD [nb, width, dim] bf16, ssm [nb, *shape] fp32) pools."""
    convs, ssms = [], []
    for _ in range(NUM_LAYERS):
+        conv_shape = (
+            (num_blocks, CONV_DIM, CONV_WIDTH)
+            if conv_state_dim_first
+            else (num_blocks, CONV_WIDTH, CONV_DIM)
+        )
        convs.append(
            torch.randn(
-                num_blocks, CONV_WIDTH, CONV_DIM, dtype=torch.bfloat16, device=device
+                *conv_shape,
+                dtype=torch.bfloat16,
+                device=device,
            )
        )
        ssms.append(
@@ -64,7 +81,7 @@ def _build_state(num_blocks, device):
    return convs, ssms


-def _build_meta(convs, ssms, device):
+def _build_meta(convs, ssms, device, conv_state_dim_first):
    """Flattened per-(layer, state-type) metadata, ordered conv, ssm per layer."""
    n = NUM_LAYERS * 2
    base = torch.zeros(n, dtype=torch.int64, device=device)
@@ -82,8 +99,14 @@ def _build_meta(convs, ssms, device):
        base[i] = conv.data_ptr()
        blk_stride[i] = conv.stride(0) * conv.element_size()
        elem[i] = conv.element_size()
-        width[i] = conv.size(1)
-        inner[i] = conv.stride(1)
+        if conv_state_dim_first:
+            width[i] = conv.size(2)
+            inner[i] = 1
+            drc[i] = conv.size(1)
+            drs[i] = conv.stride(1) * conv.element_size()
+        else:
+            width[i] = conv.size(1)
+            inner[i] = conv.stride(1)
        i += 1
        # ssm (temporal): width = 0, inner = elems per block
        base[i] = ssm.data_ptr()
@@ -95,7 +118,7 @@ def _build_meta(convs, ssms, device):
    return base, blk_stride, elem, inner, width, group, drc, drs


-def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs):
+def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs, conv_dim_first):
    """Apply the V1 copy semantics on clones, reading from the pre-copy state."""
    conv_pre = [c.clone() for c in convs]
    ssm_pre = [s.clone() for s in ssms]
@@ -108,14 +131,24 @@ def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs):
        sblk, dblk = int(bt[r, sc]), int(bt[r, dc])
        tblk = int(bt[r, sc + tb])  # temporal src column shifted by bias
        for layer in range(NUM_LAYERS):
-            conv_ref[layer][dblk, : CONV_WIDTH - tb] = conv_pre[layer][sblk, tb:]
+            if conv_dim_first:
+                conv_ref[layer][dblk, :, : CONV_WIDTH - tb] = conv_pre[layer][
+                    sblk, :, tb:
+                ]
+            else:
+                conv_ref[layer][dblk, : CONV_WIDTH - tb] = conv_pre[layer][sblk, tb:]
            ssm_ref[layer][dblk] = ssm_pre[layer][tblk]
    return conv_ref, ssm_ref


+@_parametrize("conv_state_dim_first", [False, True])
@_parametrize("num_reqs", [1, 4, 16])
@_parametrize("token_bias", [0, 1, 2])
-def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
+@_parametrize("has_idx_mapping", [True, False])
+@_cuda_required
+def test_precopy_matches_v1_copy_specs(
+    num_reqs, token_bias, has_idx_mapping, conv_state_dim_first
+):
    device = torch.device("cuda")
    torch.manual_seed(0)
    # Distinct physical block per (req, col) so copies never alias.
@@ -136,13 +169,20 @@ def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
    if num_reqs >= 2:
        dst_col[1] = 1  # src_col == dst_col -> no copy

-    convs, ssms = _build_state(num_blocks, device)
+    convs, ssms = _build_state(num_blocks, device, conv_state_dim_first)
    conv_ref, ssm_ref = _reference(
-        convs, ssms, bt.cpu(), src_col.cpu(), dst_col.cpu(), bias.cpu(), num_reqs
+        convs,
+        ssms,
+        bt.cpu(),
+        src_col.cpu(),
+        dst_col.cpu(),
+        bias.cpu(),
+        num_reqs,
+        conv_state_dim_first,
    )

    base, blk_stride, elem, inner, width, group, drc, drs = _build_meta(
-        convs, ssms, device
+        convs, ssms, device, conv_state_dim_first
    )
    bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device)
    idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
@@ -161,10 +201,11 @@ def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
        group,
        drc,
        drs,
-        idx_mapping,
+        idx_mapping if has_idx_mapping else None,
        num_reqs,
        COPY_BLOCK_SIZE=1024,
-        CONV_STATE_DIM_FIRST=False,
+        CONV_STATE_DIM_FIRST=conv_state_dim_first,
+        HAS_IDX_MAPPING=has_idx_mapping,
    )
    torch.accelerator.synchronize()

@@ -173,8 +214,204 @@ def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
        torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0)


+def test_ds_conv_copy_spec_reproduces_multi_accept_assert(monkeypatch):
+    monkeypatch.setattr(
+        layer_mamba_utils,
+        "is_conv_state_dim_first",
+        lambda: True,
+    )
+    state = torch.empty((2, CONV_DIM, CONV_WIDTH), dtype=torch.bfloat16)
+
+    with pytest.raises(AssertionError, match="num_accepted_tokens > 1"):
+        layer_mamba_utils.get_conv_copy_spec(
+            state=state,
+            block_ids=[0, 1],
+            cur_block_idx=0,
+            num_accepted_tokens=3,
+        )
+
+
+class _FakeCpuGpuBuffer:
+    def __init__(self, n):
+        self.np = np.zeros(n, dtype=np.int32)
+        self.gpu = object()
+        self.copy_sizes = []
+
+    def copy_to_gpu(self, n=None):
+        self.copy_sizes.append(n)
+        return self.gpu
+
+
+class _FakePrecopyContext:
+    def __init__(self, n):
+        self.is_initialized = True
+        self.mamba_group_ids = [0]
+        self.mamba_state_idx_buf = _FakeCpuGpuBuffer(n)
+        self.precopy_src_col_buf = _FakeCpuGpuBuffer(n)
+        self.precopy_token_bias_buf = _FakeCpuGpuBuffer(n)
+        self.calls = []
+
+    def initialize_from_forward_context(self, *args, **kwargs):
+        raise AssertionError("test context is pre-initialized")
+
+    def run_fused_precopy(
+        self,
+        *,
+        num_reqs,
+        state_idx_gpu,
+        src_col_gpu,
+        token_bias_gpu,
+        idx_mapping,
+    ):
+        self.calls.append(
+            {
+                "num_reqs": num_reqs,
+                "state_idx": self.mamba_state_idx_buf.np[:num_reqs].copy(),
+                "src_col": self.precopy_src_col_buf.np[:num_reqs].copy(),
+                "token_bias": self.precopy_token_bias_buf.np[:num_reqs].copy(),
+                "idx_mapping": idx_mapping,
+            }
+        )
+
+
+def _make_preprocess_case(token_bias):
+    req_ids = ["fresh", "same", "cross_a", "cross_b"]
+    scheduler_output = SimpleNamespace(
+        finished_req_ids=set(),
+        preempted_req_ids=set(),
+        scheduled_cached_reqs=SimpleNamespace(resumed_req_ids=set()),
+        num_scheduled_tokens={
+            "fresh": 1,
+            "same": 1,
+            "cross_a": 1,
+            "cross_b": 2,
+        },
+    )
+    input_batch = SimpleNamespace(
+        req_ids=req_ids,
+        num_accepted_tokens_cpu=np.array(
+            [token_bias + 1, token_bias + 1, token_bias + 1, 2],
+            dtype=np.int32,
+        ),
+    )
+    requests = {
+        "fresh": SimpleNamespace(req_id="fresh", num_computed_tokens=0),
+        "same": SimpleNamespace(req_id="same", num_computed_tokens=5),
+        "cross_a": SimpleNamespace(req_id="cross_a", num_computed_tokens=8),
+        "cross_b": SimpleNamespace(req_id="cross_b", num_computed_tokens=7),
+    }
+    mamba_state_idx = {"same": 1, "cross_a": 0, "cross_b": 1}
+    return scheduler_output, input_batch, requests, mamba_state_idx
+
+
+@_parametrize("token_bias", [1, 2])
+def test_preprocess_fused_align_matches_scalar_bookkeeping(monkeypatch, token_bias):
+    block_size = 4
+    mamba_spec = SimpleNamespace(block_size=block_size, num_speculative_blocks=1)
+    cache_config = SimpleNamespace(enable_prefix_caching=True)
+    kv_cache_config = SimpleNamespace()
+    scalar_copy_calls = []
+
+    def fake_collect(
+        copy_bufs,
+        kv_cache_config,
+        mamba_state_copy_funcs,
+        mamba_group_ids,
+        src_block_idx,
+        dest_block_idx,
+        accept_token_bias,
+        req_state,
+        forward_context,
+    ):
+        scalar_copy_calls.append(
+            (
+                req_state.req_id,
+                src_block_idx,
+                dest_block_idx,
+                accept_token_bias,
+            )
+        )
+
+    monkeypatch.setattr(worker_mamba_utils, "collect_mamba_copy_meta", fake_collect)
+    monkeypatch.setattr(
+        worker_mamba_utils, "do_mamba_copy_block", lambda copy_bufs: None
+    )
+
+    scalar_case = _make_preprocess_case(token_bias)
+    fused_case = _make_preprocess_case(token_bias)
+
+    scalar_copy_bufs = SimpleNamespace(
+        mamba_group_ids=[0],
+        mamba_spec=mamba_spec,
+        offset=0,
+    )
+    worker_mamba_utils.preprocess_mamba(
+        scheduler_output=scalar_case[0],
+        kv_cache_config=kv_cache_config,
+        cache_config=cache_config,
+        mamba_state_idx=scalar_case[3],
+        input_batch=scalar_case[1],
+        requests=scalar_case[2],
+        forward_context={},
+        mamba_state_copy_funcs=(),
+        copy_bufs=scalar_copy_bufs,
+    )
+
+    ctx = _FakePrecopyContext(len(fused_case[1].req_ids))
+    fused_copy_bufs = SimpleNamespace(
+        mamba_group_ids=[0],
+        mamba_spec=mamba_spec,
+        offset=0,
+    )
+    worker_mamba_utils.preprocess_mamba(
+        scheduler_output=fused_case[0],
+        kv_cache_config=kv_cache_config,
+        cache_config=cache_config,
+        mamba_state_idx=fused_case[3],
+        input_batch=fused_case[1],
+        requests=fused_case[2],
+        forward_context={},
+        mamba_state_copy_funcs=(),
+        copy_bufs=fused_copy_bufs,
+        align_ctx=ctx,
+    )
+
+    assert fused_case[3] == scalar_case[3]
+    np.testing.assert_array_equal(
+        fused_case[1].num_accepted_tokens_cpu,
+        scalar_case[1].num_accepted_tokens_cpu,
+    )
+    assert scalar_copy_calls == [
+        ("cross_a", 0, 2, token_bias),
+        ("cross_b", 1, 2, 1),
+    ]
+    assert len(ctx.calls) == 1
+    call = ctx.calls[0]
+    assert call["num_reqs"] == len(fused_case[1].req_ids)
+    assert call["idx_mapping"] is None
+    np.testing.assert_array_equal(call["state_idx"], np.array([0, 1, 2, 2]))
+    np.testing.assert_array_equal(call["src_col"], np.array([-1, -1, 0, 1]))
+    np.testing.assert_array_equal(call["token_bias"], np.array([0, 0, token_bias, 1]))
+    fused_copy_calls = [
+        (req_id, int(src), int(dst), int(bias))
+        for req_id, src, dst, bias in zip(
+            fused_case[1].req_ids,
+            call["src_col"],
+            call["state_idx"],
+            call["token_bias"],
+        )
+        if int(src) != -1 and int(src) != int(dst)
+    ]
+    assert fused_copy_calls == scalar_copy_calls
+
+
 if __name__ == "__main__":
    for nr in (1, 4, 16):
        for tb in (0, 1, 2):
-            test_precopy_matches_v1_copy_specs(nr, tb)
-            print(f"OK num_reqs={nr} token_bias={tb}")
+            for mapping in (True, False):
+                for dim_first in (False, True):
+                    test_precopy_matches_v1_copy_specs(nr, tb, mapping, dim_first)
+                    print(
+                        f"OK num_reqs={nr} token_bias={tb} "
+                        f"has_idx_mapping={mapping} conv_dim_first={dim_first}"
+                    )
@@ -27,8 +27,8 @@ from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import (  # no
    _AITER_SWIGLU_BETA,
    AiterMxfp8Experts,
 )
-from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (  # noqa: E402
-    Mxfp8NativeTritonExperts,
+from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (  # noqa: E402
+    Mxfp8EmulationTritonExperts,
 )
 from vllm.model_executor.layers.fused_moe.modular_kernel import (  # noqa: E402
    FusedMoEActivationFormat,
@@ -150,16 +150,20 @@ def test_explicit_moe_backend_aiter():

 def test_gfx950_picks_aiter():
    """Auto-select on real ROCm hardware with flydsl usable -> FlyDSL wins."""
-    with _flydsl_installed(True):
+    # NOTE: Fp8MoeBackend.AITER_MXFP8 does not require VLLM_ROCM_USE_AITER=1
+    with (
+        patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=True),
+        _flydsl_installed(True),
+    ):
        backend, experts_cls = select_mxfp8_moe_backend(_config())
    assert backend is Fp8MoeBackend.AITER_MXFP8
    assert experts_cls is AiterMxfp8Experts


-def test_gfx942_picks_triton():
+def test_gfx942_picks_emulation():
    """flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton
    dot_scaled backend wins instead."""
-    with _flydsl_installed(False):
+    with patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=False):
        backend, experts_cls = select_mxfp8_moe_backend(_config())
-    assert backend is Fp8MoeBackend.TRITON_MXFP8
-    assert experts_cls is Mxfp8NativeTritonExperts
+    assert backend is Fp8MoeBackend.EMULATION
+    assert experts_cls is Mxfp8EmulationTritonExperts
@@ -15,7 +15,9 @@ from vllm.config import PoolerConfig
    ["Qwen/Qwen3-Embedding-0.6B"],
 )
@torch.inference_mode
-def test_embed_models(hf_runner, vllm_runner, model: str):
+def test_embed_models(hf_runner, vllm_runner, monkeypatch, model: str):
+    # Keep token_embed on MRV1 when sequence pooling becomes MRV2 by default.
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
    chunk_size = 10
    n_prompt_tokens = [55, 56, 57]
    token_prompts = [[1024 + i for i in range(n)] for n in n_prompt_tokens]
@@ -104,3 +106,44 @@ def test_last_pool_score_chunked_prefill_matches_unchunked(vllm_runner, model: s
    assert chunked == pytest.approx(unchunked, abs=5e-2), (
        f"chunked score {chunked} diverged from unchunked {unchunked}"
    )
+
+
+@torch.inference_mode
+def test_sequence_embed_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None:
+    model = "Qwen/Qwen3-Embedding-0.6B"
+    chunk_size = 10
+    token_prompts = [[1024 + i for i in range(n)] for n in (25, 27)]
+    prompts = [TokensPrompt(prompt_token_ids=t) for t in token_prompts]
+
+    with hf_runner(model, auto_cls=AutoModel) as hf_model:
+        hf_outputs = []
+        for token_prompt in token_prompts:
+            inputs = hf_model.wrap_device({"input_ids": torch.tensor([token_prompt])})
+            output = hf_model.model(inputs["input_ids"])
+            embedding = torch.nn.functional.normalize(
+                output.last_hidden_state.float()[0, -1], dim=0
+            )
+            hf_outputs.append(embedding.cpu().tolist())
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    with vllm_runner(
+        model,
+        runner="pooling",
+        pooler_config=PoolerConfig(task="embed"),
+        max_model_len=64,
+        max_num_batched_tokens=chunk_size,
+        max_num_seqs=2,
+        gpu_memory_utilization=0.25,
+        enforce_eager=True,
+        enable_chunked_prefill=True,
+    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
+        vllm_outputs = vllm_model.embed(prompts)
+
+    check_embeddings_close(
+        embeddings_0_lst=hf_outputs,
+        embeddings_1_lst=vllm_outputs,
+        name_0="hf",
+        name_1="vllm",
+        tol=1e-2,
+    )
@@ -49,3 +49,60 @@ def test_models(
            vllm_output,
            rtol=2e-3 if dtype == "float" else 1e-2,
        )
+
+
+@pytest.mark.core_model
+def test_bert_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None:
+    model = "cross-encoder/ms-marco-TinyBERT-L-2-v2"
+    score_inputs = (
+        "What is the capital of France?",
+        [
+            "Paris.",
+            "Paris is the capital and largest city of France.",
+            "William Shakespeare wrote Hamlet in the early seventeenth century.",
+        ],
+    )
+    prompt_batches = [
+        ["short input"],
+        [
+            "short input",
+            "a longer input that exercises mixed sequence lengths",
+        ],
+    ]
+
+    with hf_runner(
+        model, dtype="half", auto_cls=AutoModelForSequenceClassification
+    ) as hf_model:
+        # HfRunner uses problem_type to preserve the model's
+        # sbert_ce_default_activation_function=Identity raw logits.
+        hf_model.config.problem_type = "regression"
+        hf_outputs = [hf_model.classify(prompts) for prompts in prompt_batches]
+
+    text_1, text_2 = score_inputs
+    text_pairs = [[text_1, document] for document in text_2]
+    with hf_runner(model, dtype="half", is_cross_encoder=True) as hf_model:
+        hf_scores = hf_model.predict(text_pairs).tolist()
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    with vllm_runner(
+        model,
+        runner="pooling",
+        dtype="half",
+        max_model_len=64,
+    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
+        vllm_outputs = [vllm_model.classify(prompts) for prompts in prompt_batches]
+        vllm_scores = vllm_model.score(*score_inputs)
+
+    for hf_batch, vllm_batch in zip(hf_outputs, vllm_outputs):
+        hf_tensor = torch.tensor(hf_batch)
+        vllm_tensor = torch.tensor(vllm_batch)
+        assert vllm_tensor.shape == hf_tensor.shape
+        assert torch.allclose(vllm_tensor, hf_tensor, rtol=1e-2, atol=1e-4)
+
+    assert torch.allclose(
+        torch.tensor(vllm_scores),
+        torch.tensor(hf_scores),
+        rtol=1e-2,
+        atol=1e-4,
+    )
@@ -5,6 +5,7 @@ import pytest
 import torch
 from transformers import AutoModel

+from vllm import PoolingParams
 from vllm.config import PoolerConfig

 from ...utils import check_embeddings_close
@@ -132,6 +133,7 @@ def test_encoder_only_model_runner_v2_attention(
            task="embed", seq_pooling_type="LAST", use_activation=True
        ),
    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
        vllm_outputs = vllm_model.embed(prompts)

    check_embeddings_close(
@@ -141,3 +143,77 @@ def test_encoder_only_model_runner_v2_attention(
        name_1="vllm",
        tol=1e-2,
    )
+
+
+@pytest.mark.core_model
+def test_encoder_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None:
+    model = "sentence-transformers/all-MiniLM-L6-v2"
+    prompt_batches = [
+        ["short input"],
+        [
+            "short input",
+            "a longer input that exercises mixed sequence lengths",
+        ],
+    ]
+
+    with hf_runner(model, is_sentence_transformer=True) as hf_model:
+        hf_outputs = [hf_model.encode(prompts) for prompts in prompt_batches]
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    with vllm_runner(
+        model,
+        runner="pooling",
+        max_model_len=64,
+    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
+        vllm_outputs = [vllm_model.embed(prompts) for prompts in prompt_batches]
+
+    for hf_batch, vllm_batch in zip(hf_outputs, vllm_outputs):
+        check_embeddings_close(
+            embeddings_0_lst=hf_batch,
+            embeddings_1_lst=vllm_batch,
+            name_0="hf",
+            name_1="vllm",
+            tol=1e-2,
+        )
+
+
+@pytest.mark.core_model
+def test_matryoshka_dimensions_model_runner_v2(
+    hf_runner, vllm_runner, monkeypatch
+) -> None:
+    model = "Snowflake/snowflake-arctic-embed-m-v1.5"
+    prompts = ["short input", "a longer input for a different output width"]
+    dimensions = [None, 256]
+
+    with hf_runner(model, is_sentence_transformer=True) as hf_model:
+        hf_outputs = hf_model.encode(prompts)
+
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    with vllm_runner(
+        model,
+        runner="pooling",
+        max_model_len=64,
+        gpu_memory_utilization=0.25,
+    ) as vllm_model:
+        assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
+        vllm_outputs = vllm_model.embed(
+            prompts,
+            pooling_params=[PoolingParams(dimensions=d) for d in dimensions],
+        )
+
+    expected_outputs = []
+    for output, dimension in zip(hf_outputs, dimensions):
+        output = torch.as_tensor(output)
+        if dimension is not None:
+            output = torch.nn.functional.normalize(output[:dimension], dim=0)
+        expected_outputs.append(output.tolist())
+
+    assert [len(output) for output in vllm_outputs] == [768, 256]
+    check_embeddings_close(
+        embeddings_0_lst=expected_outputs,
+        embeddings_1_lst=vllm_outputs,
+        name_0="hf",
+        name_1="vllm",
+        tol=1e-2,
+    )
@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+from unittest.mock import MagicMock, patch
+
+import numpy as np
 import pytest
 import torch
 import torch.nn as nn
@@ -9,8 +12,13 @@ from vllm.model_executor.models.bert import (
    BertMLMHead,
    SPLADESparsePooler,
 )
+from vllm.platforms import current_platform
 from vllm.pooling_params import PoolingParams
+from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
+from vllm.v1.worker.gpu.input_batch import InputBatch
+from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
+from vllm.v1.worker.gpu.states import RequestState

 # ---------------------------------------------------------------------
 # Functional test: SPLADE formula correctness (no HF download needed)
@@ -91,3 +99,70 @@ def test_splade_pooler_matches_reference_formula(B, T, H, V):
        rtol=1e-4,
        atol=1e-4,
    )
+
+
+def test_pooling_runner_gathers_required_token_ids() -> None:
+    runner = PoolingRunner.__new__(PoolingRunner)
+    pooling_params = PoolingParams(task="embed", requires_token_ids=True)
+    runner.pooling_params = {1: pooling_params, 3: pooling_params}
+    runner.pooling_states = {1: PoolingStates(), 3: PoolingStates()}
+    runner.prompt_token_ids = {
+        1: torch.tensor([101, 102]),
+        3: torch.tensor([101, 11, 102]),
+    }
+
+    input_batch = MagicMock(spec=InputBatch)
+    input_batch.idx_mapping_np = np.array([3, 1], dtype=np.int32)
+    input_batch.num_reqs = 2
+    req_states = MagicMock(spec=RequestState)
+    req_states.prompt_len = MagicMock(np=np.array([0, 2, 0, 3], dtype=np.int32))
+    metadata = runner._get_pooling_metadata(
+        input_batch, req_states, torch.device(current_platform.device_type)
+    )
+
+    expected = torch.tensor([[101, 11, 102], [101, 102, 0]])
+    assert metadata.prompt_token_ids_cpu is not None
+    assert metadata.prompt_token_ids is not None
+    assert metadata.prompt_token_ids_cpu.is_pinned() == PIN_MEMORY
+    torch.testing.assert_close(
+        metadata.prompt_lens, torch.tensor([3, 2], dtype=torch.int32)
+    )
+    torch.testing.assert_close(metadata.prompt_token_ids_cpu, expected)
+    torch.testing.assert_close(metadata.prompt_token_ids.cpu(), expected)
+
+
+def test_pooling_runner_stores_only_required_token_ids() -> None:
+    runner = PoolingRunner.__new__(PoolingRunner)
+    runner.model = MagicMock()
+    runner.supported_tasks = frozenset({"embed"})
+    runner.pooling_params = {}
+    runner.pooling_states = {}
+    runner.prompt_token_ids = {}
+
+    runner.add_request(1, PoolingParams(task="embed"), [101, 102])
+    runner.add_request(
+        2,
+        PoolingParams(task="embed", requires_token_ids=True),
+        [101, 11, 102],
+    )
+
+    assert 1 not in runner.prompt_token_ids
+    torch.testing.assert_close(runner.prompt_token_ids[2], torch.tensor([101, 11, 102]))
+
+
+def test_pooling_runner_rejects_unsupported_selected_task() -> None:
+    model = MagicMock()
+    model.pooler.get_supported_tasks.return_value = {
+        "embed",
+        "embed&token_classify",
+        "token_classify",
+    }
+    vllm_config = MagicMock()
+    vllm_config.scheduler_config.max_num_seqs = 2
+    vllm_config.model_config.get_pooling_task.return_value = "embed&token_classify"
+
+    with (
+        patch.object(PoolingRunner, "get_supported_tasks", return_value=["embed"]),
+        pytest.raises(ValueError, match="selects 'embed&token_classify'"),
+    ):
+        PoolingRunner(model, vllm_config)
@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import pytest

+from vllm.exceptions import VLLMValidationError
+
 MODEL_NAME = "sentence-transformers/all-MiniLM-L12-v2"
 max_model_len = 128

@@ -60,7 +62,7 @@ def test_bigger_truncation_size(
    truncate_prompt_tokens = max_model_len + 1

    with (
-        pytest.raises(ValueError),
+        pytest.raises(VLLMValidationError),
        vllm_runner(
            model_name, runner="pooling", max_model_len=max_model_len
        ) as vllm_model,
@@ -497,6 +497,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
    "Qwen2MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen1.5-MoE-A2.7B-Chat"),
    "Qwen3ForCausalLM": _HfExamplesInfo("Qwen/Qwen3-8B"),
    "Qwen3MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"),
+    "Qwen3_5ForCausalLM": _HfExamplesInfo("codecho/Qwen3.5-0.8B-text-only"),
+    "Qwen3_5MoeForCausalLM": _HfExamplesInfo("codecho/Qwen3.5-35B-A3B-text-only"),
    "MellumForCausalLM": _HfExamplesInfo("JetBrains/Mellum2-12B-A2.5B-Base"),
    "Qwen3NextForCausalLM": _HfExamplesInfo(
        "Qwen/Qwen3-Next-80B-A3B-Instruct",
@@ -1024,10 +1026,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
        "moonshotai/Kimi-K2.5",
        trust_remote_code=True,
    ),
-    "KimiK3ForConditionalGeneration": _HfExamplesInfo(
-        "moonshotai/Kimi-K3",
-        trust_remote_code=True,
-    ),
    "KimiVLForConditionalGeneration": _HfExamplesInfo(
        "moonshotai/Kimi-VL-A3B-Instruct",
        extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"},
@@ -1468,14 +1466,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
        is_available_online=False,
        use_original_num_layers=True,  # DSpark has >1 draft block
    ),
-    "K3DSparkModel": _HfExamplesInfo(
-        "moonshotai/Kimi-K3",
-        speculative_model="Inferact/Kimi-K3-DSpark",
-        use_original_num_layers=True,  # DSpark has >1 draft block
-        trust_remote_code=True,
-        # FIXME: Investigate the NVML failure in CI.
-        is_available_online=False,
-    ),
    "Qwen3DSparkModel": _HfExamplesInfo(
        "Qwen/Qwen3-8B",
        speculative_model="deepseek-ai/dspark_qwen3_8b_block7",
@@ -1674,12 +1664,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
        trust_remote_code=True,
        max_model_len=4096,
    ),
-    "KimiK3MTPModel": _HfExamplesInfo(
-        "moonshotai/Kimi-K3",
-        speculative_model="moonshotai/Kimi-K3",
-        trust_remote_code=True,
-        is_available_online=False,
-    ),
    "LongCatFlashMTPModel": _HfExamplesInfo(
        "meituan-longcat/LongCat-Flash-Chat",
        trust_remote_code=True,
@@ -1,153 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-import torch.nn as nn
-
-from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
-from vllm.model_executor.models.registry import ModelRegistry
-from vllm.models.kimi_k3.nvidia import dspark_mla
-from vllm.models.kimi_k3.nvidia.dspark_mla import (
-    K3DSparkForCausalLM,
-    K3DSparkModel,
-    ReplicatedDSparkMarkovHead,
-)
-
-
-def test_dspark_mla_uses_compile_free_model_entrypoint():
-    assert ModelRegistry._try_load_model_cls("K3DSparkModel") is K3DSparkForCausalLM
-    assert not issubclass(K3DSparkModel, TorchCompileWithNoGuardsWrapper)
-
-
-@pytest.mark.parametrize(
-    ("checkpoint_name", "runtime_name", "shard_id"),
-    [
-        (
-            "layers.0.self_attn.q_a_proj.weight",
-            "model.layers.0.self_attn.fused_qkv_a_proj.weight",
-            0,
-        ),
-        (
-            "layers.0.self_attn.kv_a_proj_with_mqa.weight",
-            "model.layers.0.self_attn.fused_qkv_a_proj.weight",
-            1,
-        ),
-        (
-            "layers.0.mlp.gate_proj.weight",
-            "model.layers.0.mlp.gate_up_proj.weight",
-            0,
-        ),
-        (
-            "layers.0.mlp.up_proj.weight",
-            "model.layers.0.mlp.gate_up_proj.weight",
-            1,
-        ),
-        ("context_proj.weight", "model.context_proj.weight", None),
-    ],
-)
-def test_dspark_mla_checkpoint_weight_mapping(checkpoint_name, runtime_name, shard_id):
-    assert K3DSparkForCausalLM.hf_to_vllm_mapper._map_name_with_shard(
-        checkpoint_name
-    ) == (runtime_name, shard_id)
-
-
-def test_dspark_mla_shares_frozen_target_weights_and_skips_training_head():
-    assert not K3DSparkForCausalLM.has_own_embed_tokens
-    assert not K3DSparkForCausalLM.has_own_lm_head
-    assert set(K3DSparkForCausalLM.checkpoint_skip_substrs) == {
-        "confidence_head",
-        "embed_tokens",
-        "lm_head",
-    }
-
-
-@pytest.mark.cpu_test
-def test_dspark_markov_head_replication_is_opt_in(
-    monkeypatch: pytest.MonkeyPatch,
-):
-    from vllm.model_executor.layers import logits_processor, vocab_parallel_embedding
-
-    monkeypatch.setattr(
-        vocab_parallel_embedding, "get_tensor_model_parallel_rank", lambda: 3
-    )
-    monkeypatch.setattr(
-        vocab_parallel_embedding,
-        "get_tensor_model_parallel_world_size",
-        lambda: 8,
-    )
-    monkeypatch.setattr(
-        logits_processor,
-        "get_current_vllm_config",
-        lambda: SimpleNamespace(model_config=None),
-    )
-
-    sharded = DSparkMarkovHead(128, 128, 8, prefix="markov_head")
-    assert sharded.markov_w1.tp_size == 8
-    assert sharded.markov_w2.tp_size == 8
-
-    replicated = ReplicatedDSparkMarkovHead(128, 128, 8, prefix="markov_head")
-    assert isinstance(replicated, DSparkMarkovHead)
-    assert replicated.markov_w1.tp_size == 1
-    assert replicated.markov_w2.tp_size == 1
-    assert replicated.markov_w1.weight.shape == (128, 8)
-    assert replicated.markov_w2.weight.shape == (128, 8)
-
-    def fail_collective(*args, **kwargs):
-        raise AssertionError("replicated Markov head must not invoke TP collectives")
-
-    monkeypatch.setattr(
-        vocab_parallel_embedding,
-        "tensor_model_parallel_all_reduce",
-        fail_collective,
-    )
-    logits_processor = LogitsProcessor(128)
-    monkeypatch.setattr(logits_processor, "_gather_logits", fail_collective)
-
-    markov_embed = replicated.embed(torch.tensor([1, 2]))
-    bias = replicated.bias(markov_embed, logits_processor)
-    assert markov_embed.shape == (2, 8)
-    assert bias.shape == (2, 128)
-
-
-@pytest.mark.cpu_test
-def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch):
-    markov_head_calls = []
-
-    class DummyModule(nn.Module):
-        def __init__(self, *args, **kwargs):
-            super().__init__()
-
-    def make_markov_head(*args, **kwargs):
-        markov_head_calls.append((args, kwargs))
-        return DummyModule()
-
-    monkeypatch.setattr(dspark_mla, "get_draft_quant_config", lambda _: None)
-    monkeypatch.setattr(dspark_mla, "ReplicatedLinear", DummyModule)
-    monkeypatch.setattr(dspark_mla, "RMSNorm", DummyModule)
-    monkeypatch.setattr(dspark_mla, "K3DSparkDecoderLayer", DummyModule)
-    monkeypatch.setattr(dspark_mla, "ReplicatedDSparkMarkovHead", make_markov_head)
-
-    config = SimpleNamespace(
-        target_hidden_size=16,
-        num_target_layers=2,
-        hidden_size=8,
-        rms_norm_eps=1e-6,
-        num_hidden_layers=1,
-        vocab_size=128,
-        draft_vocab_size=128,
-        markov_rank=4,
-    )
-    vllm_config = SimpleNamespace(
-        speculative_config=SimpleNamespace(
-            draft_model_config=SimpleNamespace(hf_config=config)
-        )
-    )
-
-    K3DSparkModel(vllm_config=vllm_config, start_layer_id=0, prefix="model")
-
-    assert len(markov_head_calls) == 1
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import uuid
+from io import BytesIO
 from pathlib import Path

 import numpy as np
@@ -10,6 +11,7 @@ from PIL import Image, ImageDraw

 from vllm.multimodal.hasher import MultiModalHasher
 from vllm.multimodal.media.base import MediaWithBytes
+from vllm.multimodal.media.image import ImageMediaIO
 from vllm.multimodal.parse import MultiModalDataParser

 pytestmark = pytest.mark.cpu_test
@@ -134,3 +136,37 @@ def test_hash_image_exif_id():
    assert hasher.hash_kwargs(image=image1) == hasher.hash_kwargs(image=id.bytes)
    # second image has non-UUID in ImageID, so it should hash to the image data
    assert hasher.hash_kwargs(image=image2) == hasher.hash_kwargs(image=image2a)
+
+
+def _rgba_png_bytes() -> bytes:
+    image = Image.new("RGBA", (8, 8), (255, 0, 0, 128))
+    buf = BytesIO()
+    image.save(buf, format="PNG")
+    return buf.getvalue()
+
+
+def test_hash_collision_media_io_config():
+    data = _rgba_png_bytes()
+    white = ImageMediaIO(rgba_background_color=(255, 255, 255)).load_bytes(data)
+    black = ImageMediaIO(rgba_background_color=(0, 0, 0)).load_bytes(data)
+    white2 = ImageMediaIO(rgba_background_color=(255, 255, 255)).load_bytes(data)
+    keep = ImageMediaIO(image_mode=None).load_bytes(data)
+
+    hasher = MultiModalHasher
+    assert hasher.hash_kwargs(image=white) != hasher.hash_kwargs(image=black)
+    assert hasher.hash_kwargs(image=white) != hasher.hash_kwargs(image=keep)
+    assert hasher.hash_kwargs(image=white) == hasher.hash_kwargs(image=white2)
+
+
+def test_hash_media_io_noop_config_preserves_hash():
+    image = Image.new("RGB", (8, 8), (0, 128, 255))
+    buf = BytesIO()
+    image.save(buf, format="PNG")
+    data = buf.getvalue()
+
+    loaded = ImageMediaIO().load_bytes(data)
+    assert loaded.io_config is None
+
+    plain = MediaWithBytes(loaded.media, data)
+    hasher = MultiModalHasher
+    assert hasher.hash_kwargs(image=loaded) == hasher.hash_kwargs(image=plain)
@@ -14,6 +14,7 @@ from transformers import AutoVideoProcessor
 from transformers.video_utils import VideoMetadata

 from vllm.assets.base import get_vllm_public_assets
+from vllm.models.minimax_m3.common.mm_preprocess import MiniMaxM3VideoBackend
 from vllm.multimodal.video import (
    PYNVVIDEOCODEC_DECODER_CACHE_SIZE,
    PYNVVIDEOCODEC_VIDEO_BACKEND,
@@ -388,6 +389,12 @@ def test_cosmos3_edge_uses_qwen3_vl_video_backend():
            {"fps": 2},
            id="qwen2_5_vl",
        ),
+        pytest.param(
+            "MiniMaxAI/MiniMax-M3",
+            MiniMaxM3VideoBackend,
+            None,
+            id="minimax_m3_vl",
+        ),
    ],
 )
 def test_video_processor_from_model_repo(
@@ -1235,26 +1235,6 @@ def test_vllm_config_defaults_are_none():
                assert getattr(config.compilation_config, k) is None


-def test_validate_mamba_align_subblock_prefill():
-    """Align mode permits configured prefill chunks smaller than a block."""
-    config = SimpleNamespace(
-        cache_config=SimpleNamespace(
-            block_size=11392,
-            mamba_cache_mode="align",
-        ),
-        parallel_config=SimpleNamespace(
-            decode_context_parallel_size=1,
-        ),
-        scheduler_config=SimpleNamespace(
-            max_num_batched_tokens=8192,
-            long_prefill_token_threshold=4096,
-            disable_chunked_mm_input=False,
-        ),
-    )
-
-    VllmConfig.validate_block_size(config)
-
-
@pytest.mark.parametrize(
    ("model_id", "compilation_config", "optimization_level"),
    [
@@ -1,165 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import json
-
-import pytest
-
-from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
-from vllm.transformers_utils.config import get_config
-from vllm.transformers_utils.configs.k3_dspark import K3DSparkConfig
-
-
-def _write_dspark_config(path, **overrides):
-    path.mkdir()
-    config = {
-        "architectures": ["K3DSparkModel"],
-        "model_type": "k3_dspark",
-        "hidden_size": 7168,
-        "intermediate_size": 14336,
-        "num_hidden_layers": 5,
-        "num_attention_heads": 64,
-        "num_key_value_heads": 64,
-        "q_lora_rank": 1536,
-        "kv_lora_rank": 512,
-        "qk_nope_head_dim": 128,
-        "qk_rope_head_dim": 64,
-        "v_head_dim": 128,
-        "vocab_size": 163840,
-        "rms_norm_eps": 1e-5,
-        "max_position_embeddings": 32768,
-        "rope_theta": 50000.0,
-        "num_target_layers": 5,
-        "target_hidden_size": 7168,
-        "target_num_hidden_layers": 93,
-        "target_layer_ids": [2, 23, 47, 71, 89],
-        "markov_rank": 256,
-        "draft_vocab_size": 163840,
-        "torch_dtype": "bfloat16",
-    }
-    config.update(overrides)
-    (path / "config.json").write_text(json.dumps(config))
-
-
-def _write_target_config(path):
-    path.mkdir()
-    config = {
-        "architectures": ["LlamaForCausalLM"],
-        "model_type": "llama",
-        "hidden_size": 7168,
-        "intermediate_size": 14336,
-        "num_hidden_layers": 93,
-        "num_attention_heads": 56,
-        "num_key_value_heads": 8,
-        "vocab_size": 163840,
-        "max_position_embeddings": 32768,
-        "torch_dtype": "bfloat16",
-    }
-    (path / "config.json").write_text(json.dumps(config))
-
-
-def test_dspark_mla_config_loads_from_local_json(tmp_path):
-    draft_path = tmp_path / "draft"
-    _write_dspark_config(draft_path)
-
-    config = get_config(draft_path, trust_remote_code=False)
-
-    assert isinstance(config, K3DSparkConfig)
-    assert config.model_type == "k3_dspark"
-    assert config.architectures == ["K3DSparkModel"]
-    assert config.hidden_act == "silu"
-    assert config.rope_parameters == {
-        "rope_type": "default",
-        "rope_theta": 50000.0,
-    }
-    assert config.n_routed_experts == 0
-    assert config.draft_vocab_size == config.vocab_size
-
-
-@pytest.mark.parametrize(
-    "overrides",
-    [
-        {"mla_use_nope": True},
-        {"mla_use_output_gate": True},
-        {"mla_use_qk_norm": True},
-        {"dspark_bonus_anchor": True},
-        {"q_lora_rank": None},
-        {"draft_vocab_size": 8192},
-        {"target_layer_ids": []},
-        {"num_target_layers": 4},
-    ],
-)
-def test_dspark_mla_rejects_unsupported_checkpoint_options(tmp_path, overrides):
-    draft_path = tmp_path / "draft"
-    _write_dspark_config(draft_path, **overrides)
-
-    with pytest.raises(ValueError, match="MLA DSpark"):
-        get_config(draft_path, trust_remote_code=False)
-
-
-def test_dspark_mla_uses_latent_kv_geometry(tmp_path):
-    draft_path = tmp_path / "draft"
-    _write_dspark_config(draft_path)
-    model_config = ModelConfig(
-        model=str(draft_path),
-        tokenizer_mode="skip",
-        runner="draft",
-        max_model_len=32768,
-    )
-
-    assert model_config.is_deepseek_mla
-    assert model_config.use_mla
-    assert model_config.get_head_size() == 576
-    # external_launcher skips ParallelConfig's local-GPU-count check so the
-    # config logic can be exercised at TP8 on a single-GPU test node.
-    parallel_config = ParallelConfig(
-        tensor_parallel_size=8, distributed_executor_backend="external_launcher"
-    )
-    assert model_config.get_num_kv_heads(parallel_config) == 1
-    assert model_config.get_num_attention_heads(parallel_config) == 8
-    assert model_config.get_num_experts() == 0
-
-
-def test_dspark_mla_speculative_config_preserves_architecture(tmp_path):
-    target_path = tmp_path / "target"
-    draft_path = tmp_path / "draft"
-    _write_target_config(target_path)
-    _write_dspark_config(draft_path)
-    target_config = ModelConfig(
-        model=str(target_path), tokenizer_mode="skip", max_model_len=32768
-    )
-    speculative_config = SpeculativeConfig(
-        model=str(draft_path),
-        method="dspark",
-        num_speculative_tokens=8,
-        target_model_config=target_config,
-        target_parallel_config=ParallelConfig(),
-    )
-
-    assert speculative_config.parallel_drafting
-    assert speculative_config.draft_model_config.architectures == ["K3DSparkModel"]
-    assert speculative_config.draft_model_config.hf_config.model_type == "k3_dspark"
-    assert speculative_config.draft_model_config.use_mla
-
-
-def test_dspark_mla_rejects_decode_context_parallelism(tmp_path):
-    target_path = tmp_path / "target"
-    draft_path = tmp_path / "draft"
-    _write_target_config(target_path)
-    _write_dspark_config(draft_path)
-    target_config = ModelConfig(
-        model=str(target_path), tokenizer_mode="skip", max_model_len=32768
-    )
-
-    with pytest.raises(ValueError, match="does not currently support decode context"):
-        SpeculativeConfig(
-            model=str(draft_path),
-            method="dspark",
-            num_speculative_tokens=8,
-            target_model_config=target_config,
-            target_parallel_config=ParallelConfig(
-                tensor_parallel_size=2,
-                decode_context_parallel_size=2,
-                distributed_executor_backend="external_launcher",
-            ),
-        )
@@ -73,11 +73,7 @@ DEVICE_TYPE = current_platform.device_type
 def test_mla_kv_cache_spec_uses_layer_cache_dtype(
    cache_dtype: str, expected_quant_mode: KVQuantMode
 ):
-    layer = SimpleNamespace(
-        kv_cache_dtype=cache_dtype,
-        head_size=576,
-        non_causal_multi_token_decode=False,
-    )
+    layer = SimpleNamespace(kv_cache_dtype=cache_dtype, head_size=576)
    vllm_config = SimpleNamespace(
        cache_config=SimpleNamespace(block_size=64), model_config=None
    )
@@ -1,124 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm.model_executor.layers.attention.mla_attention import (
-    MLACommonMetadata,
-    MLACommonMetadataBuilder,
-    QueryLenSupport,
-)
-from vllm.v1.attention.backend import CommonAttentionMetadata
-from vllm.v1.kv_cache_interface import MLAAttentionSpec
-
-
-class _NonCausalMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
-    supports_non_causal_multi_token_decode = True
-
-
-def _metadata(
-    query_start_loc: list[int], num_tokens: int | None = None
-) -> CommonAttentionMetadata:
-    num_reqs = len(query_start_loc) - 1
-    num_tokens = query_start_loc[-1] if num_tokens is None else num_tokens
-    return CommonAttentionMetadata(
-        query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32),
-        query_start_loc_cpu=torch.tensor(query_start_loc, dtype=torch.int32),
-        seq_lens=torch.arange(1, num_reqs + 1, dtype=torch.int32) * 100 + 8,
-        num_reqs=num_reqs,
-        num_actual_tokens=num_tokens,
-        max_query_len=max(
-            end - start for start, end in zip(query_start_loc, query_start_loc[1:])
-        ),
-        max_seq_len=num_reqs * 100 + 8,
-        block_table_tensor=torch.arange(num_reqs * 3, dtype=torch.int32).view(
-            num_reqs, 3
-        ),
-        slot_mapping=torch.arange(num_tokens),
-        causal=False,
-        seq_lens_cpu_upper_bound=None,
-    )
-
-
-def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder:
-    builder = object.__new__(_NonCausalMLAMetadataBuilder)
-    builder.device = torch.device("cpu")
-    builder.reorder_batch_threshold = 1
-    builder.query_len_support = QueryLenSupport.SINGLE_ONLY
-    builder.non_causal_multi_token_decode = marked
-    builder.dcp_world_size = 1
-    builder.metadata_cls = MLACommonMetadata
-    builder.model_config = SimpleNamespace(
-        dtype=torch.bfloat16, get_head_size=lambda: 576
-    )
-    return builder
-
-
-def test_noncausal_block_uses_decode_without_cpu_lengths():
-    common_metadata = _metadata([0, 8, 16])
-    metadata = _builder().build(0, common_metadata)
-
-    assert metadata.num_decodes == 2
-    assert metadata.num_decode_tokens == 16
-    assert metadata.num_prefills == 0
-    assert metadata.prefill is None
-    assert metadata.decode is not None
-    assert metadata.decode.block_table.shape == (2, 3)
-    assert metadata.decode.seq_lens.shape == (2,)
-    assert torch.equal(metadata.decode.block_table, common_metadata.block_table_tensor)
-    assert torch.equal(metadata.decode.seq_lens, common_metadata.seq_lens)
-    assert not metadata.causal
-
-
-def test_noncausal_support_is_explicit_and_uniform():
-    with pytest.raises(ValueError, match="explicitly supported"):
-        _builder(marked=False).build(0, _metadata([0, 8, 16]))
-    with pytest.raises(ValueError, match="uniform query block"):
-        _builder().build(0, _metadata([0, 3, 8]))
-
-
-def test_noncausal_block_allows_trailing_cudagraph_padding():
-    common_metadata = _metadata([0, 8, 16, 16], num_tokens=24)
-    common_metadata.seq_lens[-1] = 0
-
-    metadata = _builder().build(0, common_metadata)
-
-    assert metadata.num_decodes == 3
-    assert metadata.num_decode_tokens == 24
-    assert metadata.decode is not None
-    assert metadata.decode.seq_lens.tolist() == [108, 208, 0]
-
-
-def test_noncausal_block_rejects_non_trailing_padding():
-    with pytest.raises(ValueError, match="uniform query block"):
-        _builder().build(0, _metadata([0, 8, 8, 16], num_tokens=24))
-
-
-def test_noncausal_decode_metadata_keeps_live_request_buffers():
-    common_metadata = _metadata([0, 8, 16])
-    metadata = _builder().build(0, common_metadata)
-
-    assert metadata.decode is not None
-    assert metadata.decode.seq_lens.data_ptr() == common_metadata.seq_lens.data_ptr()
-    assert (
-        metadata.decode.block_table.data_ptr()
-        == common_metadata.block_table_tensor.data_ptr()
-    )
-
-
-def test_mla_cache_marker_is_preserved_and_cannot_be_mixed():
-    kwargs = {
-        "block_size": 64,
-        "num_kv_heads": 1,
-        "head_size": 576,
-        "dtype": torch.bfloat16,
-    }
-    marked = MLAAttentionSpec(**kwargs, non_causal_multi_token_decode=True)
-    unmarked = MLAAttentionSpec(**kwargs)
-
-    assert MLAAttentionSpec.merge([marked, marked]).non_causal_multi_token_decode
-    with pytest.raises(AssertionError, match="non-causal decode mode"):
-        MLAAttentionSpec.merge([marked, unmarked])
@@ -40,8 +40,6 @@ def test_mamba_align_split_partial_tail_schedule():
    hash_block_size = 32
    mock = SimpleNamespace(
        cache_config=SimpleNamespace(block_size=block_size),
-        max_num_scheduled_tokens=8192,
-        scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
        use_eagle=False,
        hash_block_size=hash_block_size,
        mamba_partial_cache_hit=True,
@@ -77,78 +75,6 @@ def test_mamba_align_split_partial_tail_schedule():
    assert split(self=mock, request=req2, num_new_tokens=1000) == 512


-def test_mamba_align_split_when_block_exceeds_scheduling_budget():
-    """Sub-block chunks make progress only when no step can fit a full block."""
-    block_size = 11392
-    token_budget = 8192
-    prompt_length = 30000
-    mock = SimpleNamespace(
-        cache_config=SimpleNamespace(block_size=block_size),
-        max_num_scheduled_tokens=token_budget,
-        scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
-        use_eagle=False,
-        hash_block_size=32,
-        mamba_partial_cache_hit=False,
-    )
-    req = make_request("0", [0] * prompt_length, 32, sha256)
-    split = Scheduler._mamba_block_aligned_split
-
-    mock.max_num_scheduled_tokens = block_size
-    assert split(self=mock, request=req, num_new_tokens=token_budget) == 0
-    mock.max_num_scheduled_tokens = token_budget
-
-    scheduled_chunks = []
-    while req.num_computed_tokens < prompt_length:
-        num_new_tokens = min(token_budget, prompt_length - req.num_computed_tokens)
-        num_scheduled_tokens = split(
-            self=mock,
-            request=req,
-            num_new_tokens=num_new_tokens,
-        )
-        assert 0 < num_scheduled_tokens <= token_budget
-        scheduled_chunks.append(num_scheduled_tokens)
-        req.num_computed_tokens += num_scheduled_tokens
-
-    assert scheduled_chunks == [8192, 3200, 8192, 3200, 7216]
-
-
-def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():
-    """A long-prefill cap below the block size permits sub-block progress."""
-    block_size = 512
-    token_budget = 8192
-    long_prefill_threshold = 384
-    prompt_length = 1300
-    mock = SimpleNamespace(
-        cache_config=SimpleNamespace(block_size=block_size),
-        max_num_scheduled_tokens=token_budget,
-        scheduler_config=SimpleNamespace(
-            long_prefill_token_threshold=long_prefill_threshold
-        ),
-        use_eagle=False,
-        hash_block_size=32,
-        mamba_partial_cache_hit=False,
-    )
-    req = make_request("0", [0] * prompt_length, 32, sha256)
-    split = Scheduler._mamba_block_aligned_split
-
-    scheduled_chunks = []
-    while req.num_computed_tokens < prompt_length:
-        num_new_tokens = min(
-            long_prefill_threshold,
-            prompt_length - req.num_computed_tokens,
-        )
-        num_scheduled_tokens = split(
-            self=mock,
-            request=req,
-            num_new_tokens=num_new_tokens,
-        )
-        assert 0 < num_scheduled_tokens <= long_prefill_threshold
-        scheduled_chunks.append(num_scheduled_tokens)
-        req.num_computed_tokens += num_scheduled_tokens
-
-    assert scheduled_chunks == [384, 128, 384, 128, 276]
-
-
 def test_hybrid_mamba_align_partial_hash_hit():
    hash_block_size = 2
    mamba_block_size = 2 * hash_block_size
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections import deque
+from collections import defaultdict, deque
+from collections.abc import Callable
 from unittest.mock import Mock

 import pytest
@@ -329,46 +330,311 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
    assert not scheduler.running


-def test_no_placeholder_underflow_on_discarded_spec_frame():
-    num_spec = 5
+class PipelinedEngine:
+    """Drive a real AsyncScheduler like EngineCore.step_with_batch_queue:
+    schedule until the batch queue is full, then process the oldest step's
+    output. Async PP runs pp_size+1 concurrent batches, so up to pp_size
+    steps are in flight at each schedule() call -- the window in which
+    preemption must handle output that has not yet returned. (Single-GPU e2e
+    tests can never create this window: at PP=1, exactly one step is in
+    flight and it is processed before a preempted request can resume.)
+
+    The model runner is emulated with the V2 runner's own bookkeeping, from
+    only what the scheduler serializes to it: slots flushed on
+    preempted_req_ids, resumed requests re-added from the NewRequestData
+    snapshot, sampling when a step reaches the end of the runner's own view
+    of the sequence. This makes preemption races observable: a stale token
+    delivered after a resume is scheduled extends the scheduler's sequence
+    but not the runner's.
+
+    Every sample emits a globally unique token tagged with its sampled
+    position, so tests can assert exact delivery.
+    """
+
+    def __init__(
+        self,
+        scheduler: AsyncScheduler,
+        queue_size: int,
+        accept_drafts: Callable[[int, str, int], int] | None = None,
+    ):
+        self.scheduler = scheduler
+        self.queue_size = queue_size
+        self.accept_drafts = accept_drafts
+        # In-flight steps: (scheduler_output, new_reqs snapshot) in FIFO order.
+        self.queue: deque[tuple[SchedulerOutput, list[tuple[str, int, int]]]] = deque()
+        # Runner-side request state: req_id -> [seq_len, num_computed] as the
+        # runner sees them (its own sampled tokens, not the scheduler's).
+        self.runner_view: dict[str, list[int]] = {}
+        # All tokens the fake runner ever sampled, per request, in order.
+        self.emitted: dict[str, list[int]] = defaultdict(list)
+        # Sequence position each (globally unique) token was sampled for.
+        self.emitted_position: dict[int, int] = {}
+        self.step_idx = 0
+        self._next_token = 1000
+
+    def _schedule(self) -> bool:
+        scheduler_output = self.scheduler.schedule()
+        self.step_idx += 1
+        # Snapshot what NewRequestData serializes at schedule time (both new
+        # and resumed requests for the V2 runner).
+        new_reqs = [
+            (r.req_id, len(r.prefill_token_ids), r.num_computed_tokens)
+            for r in scheduler_output.scheduled_new_reqs
+        ]
+        # Enqueue empty steps too (the engine executes them), so the runner
+        # still observes their preempted/finished request ids in step order.
+        self.queue.appendleft((scheduler_output, new_reqs))
+        return True
+
+    def _process_oldest_step(self) -> None:
+        scheduler_output, new_reqs = self.queue.pop()
+        # Worker-side state updates, in step order: flush preempted/finished
+        # slots, then (re-)add new/resumed requests.
+        for req_id in scheduler_output.preempted_req_ids or ():
+            self.runner_view.pop(req_id, None)
+        for req_id in scheduler_output.finished_req_ids or ():
+            self.runner_view.pop(req_id, None)
+        for req_id, seq_len, num_computed in new_reqs:
+            self.runner_view[req_id] = [seq_len, num_computed]
+
+        req_ids = list(scheduler_output.num_scheduled_tokens.keys())
+        sampled_token_ids: list[list[int]] = []
+        for req_id in req_ids:
+            num_scheduled = scheduler_output.num_scheduled_tokens[req_id]
+            view = self.runner_view.get(req_id)
+            if view is None:
+                # Slot already flushed (request finished/aborted mid-flight).
+                sampled_token_ids.append([])
+                continue
+            seq_len, num_computed = view
+            end = num_computed + num_scheduled
+            if end < seq_len:
+                # Partial prefill by the runner's own bookkeeping: no sample.
+                view[1] = end
+                sampled_token_ids.append([])
+                continue
+            drafts = scheduler_output.scheduled_spec_decode_tokens.get(req_id, ())
+            num_accepted = (
+                min(self.accept_drafts(self.step_idx, req_id, len(drafts)), len(drafts))
+                if drafts and self.accept_drafts
+                else 0
+            )
+            num_rejected = len(drafts) - num_accepted
+            tokens = list(range(self._next_token, self._next_token + 1 + num_accepted))
+            self._next_token += 1 + num_accepted
+            self.emitted[req_id].extend(tokens)
+            sampled_token_ids.append(tokens)
+            # Rejected drafts roll back computed; the sampled tokens extend
+            # the runner's sequence.
+            view[1] = end - num_rejected
+            view[0] = view[1] + 1
+            for offset, token in enumerate(tokens):
+                self.emitted_position[token] = view[0] - len(tokens) + offset
+        model_runner_output = ModelRunnerOutput(
+            req_ids=req_ids,
+            req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)},
+            sampled_token_ids=sampled_token_ids,
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+        )
+        self.scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    def run(
+        self,
+        max_steps: int = 2000,
+        before_step: Callable[[int, "PipelinedEngine"], None] | None = None,
+    ) -> None:
+        for i in range(max_steps):
+            if not self.scheduler.has_requests() and not self.queue:
+                return
+            if before_step is not None:
+                before_step(i, self)
+            scheduled = (
+                self.scheduler.has_requests()
+                and len(self.queue) < self.queue_size
+                and self._schedule()
+            )
+            if scheduled and len(self.queue) < self.queue_size:
+                # Queue not yet full: the engine returns without blocking.
+                continue
+            if self.queue:
+                self._process_oldest_step()
+        raise AssertionError("engine loop did not converge")
+
+
+def _create_async_pp_scheduler(
+    num_spec: int, pp_size: int = 3, num_blocks: int = 5
+) -> AsyncScheduler:
    scheduler = create_scheduler(
        async_scheduling=True,
-        num_speculative_tokens=num_spec,
-        speculative_method="ngram_gpu",
+        num_speculative_tokens=num_spec or None,
+        speculative_method="ngram_gpu" if num_spec else None,
+        use_v2_model_runner=True,
+        num_blocks=num_blocks,
+        block_size=16,
+        max_num_batched_tokens=512,
    )
-    req = create_requests(num_requests=1, max_tokens=20)[0]
-    req.num_computed_tokens = req.num_tokens
-    scheduler.requests[req.request_id] = req
-    scheduler.running.append(req)
-    req.status = RequestStatus.RUNNING
+    # Emulate PP at the scheduler level; constructing with
+    # pipeline_parallel_size>1 requires that many visible GPUs. Drive with
+    # queue_size=pp_size+1 (V2 async PP runs pp_size+1 concurrent batches).
+    scheduler.pp_size = pp_size
+    scheduler.use_pp = pp_size > 1
+    return scheduler

-    req.num_output_placeholders = 1
-    req.async_tokens_to_discard = num_spec
-    computed_before = req.num_computed_tokens

-    scheduler_output = SchedulerOutput(
-        scheduled_new_reqs=[],
-        scheduled_cached_reqs=CachedRequestData.make_empty(),
-        num_scheduled_tokens={req.request_id: num_spec + 1},
-        total_num_scheduled_tokens=num_spec + 1,
-        scheduled_encoder_inputs={},
-        scheduled_spec_decode_tokens={req.request_id: [10] * num_spec},
-        num_common_prefix_blocks=[],
-        finished_req_ids=set(),
-        free_encoder_mm_hashes=[],
+def _assert_ordered_subset(delivered: list[int], emitted: list[int]) -> None:
+    """Delivered tokens must be an order-preserving subset of the emitted
+    tokens with no duplicates (tokens are globally unique)."""
+    it = iter(emitted)
+    for token in delivered:
+        assert token in it, f"token {token} delivered out of order or twice"
+
+
+def _assert_positions_consistent(req, engine: PipelinedEngine) -> None:
+    """The i-th delivered output token must be one the runner sampled for
+    exactly sequence position prompt_len + i: catches a preempted request's
+    stale output landing on a position the resumed request resampled (or
+    vice versa), which token-stream equality alone cannot see."""
+    for i, token in enumerate(req.output_token_ids):
+        expected = req.num_prompt_tokens + i
+        actual = engine.emitted_position[token]
+        assert actual == expected, (
+            f"output {i} of {req.request_id}: token sampled for position "
+            f"{actual}, delivered as position {expected}"
+        )
+
+
+@pytest.mark.parametrize("num_spec", [0, 3])
+def test_kv_pressure_preemption_with_inflight_output(num_spec: int):
+    """KV-pressure preemption of requests with in-flight async output.
+
+    PP=3 + async scheduling (batch queue of 4), a block pool small enough
+    that decodes contend and preempt mid-flight, and staggered arrivals so
+    the batch queue actually pipelines. A preempted request's in-flight steps
+    still return: their tokens must be delivered exactly once, their stale
+    spec-rejection counts must not corrupt the rolled-back counters, and the
+    resume must not resample a position that output later delivers.
+
+    Regression for the num_output_placeholders underflow EngineCore crash:
+    with the fix reverted, the num_spec=3 variant fails with exactly
+    ``assert request.num_output_placeholders >= 0`` when a stale spec output
+    returns after the preempted request was resumed and sampled.
+    """
+    max_tokens = 24
+    scheduler = _create_async_pp_scheduler(num_spec)
+    requests = create_requests(
+        num_requests=8, num_tokens=8, max_tokens=max_tokens, ignore_eos=True
    )
-    model_runner_output = ModelRunnerOutput(
-        req_ids=[req.request_id],
-        req_id_to_index={req.request_id: 0},
-        sampled_token_ids=[[999]],
-        logprobs=None,
-        prompt_logprobs_dict={},
-        pooler_output=[],
+    pending = list(requests)
+    for _ in range(2):
+        scheduler.add_request(pending.pop(0))
+
+    # Observe that the scenario under test actually occurs.
+    preempts_with_inflight_output = 0
+    orig_preempt = scheduler._preempt_request
+
+    def counting_preempt(request, timestamp, **kwargs):
+        nonlocal preempts_with_inflight_output
+        if request.num_in_flight_tokens > 0:
+            preempts_with_inflight_output += 1
+        return orig_preempt(request, timestamp, **kwargs)
+
+    scheduler._preempt_request = counting_preempt
+
+    def add_requests(step: int, engine: PipelinedEngine):
+        if pending:
+            scheduler.add_request(pending.pop(0))
+
+    engine = PipelinedEngine(
+        scheduler,
+        queue_size=4,
+        # Deterministically vary spec acceptance so stale outputs carry
+        # nonzero rejection counts.
+        accept_drafts=lambda step, req_id, n: (step + int(req_id)) % (n + 1),
    )
+    engine.run(before_step=add_requests)

-    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert preempts_with_inflight_output > 0, (
+        "test did not exercise preemption with in-flight output"
+    )
+    for req in requests:
+        assert req.is_finished()
+        assert req.num_output_tokens == max_tokens
+        # Lossless: delivered tokens are exactly the sampled tokens, in order
+        # (the excluded tail was emitted after the request finished).
+        emitted = engine.emitted[req.request_id]
+        assert list(req.output_token_ids) == emitted[:max_tokens]
+        _assert_positions_consistent(req, engine)

-    assert req.num_output_placeholders == 1
-    assert req.num_computed_tokens == computed_before
-    assert req.async_tokens_to_discard == num_spec - 1
-    assert req.status == RequestStatus.RUNNING
+
+@pytest.mark.parametrize("pp_size", [1, 3])
+def test_reset_prefix_cache_with_inflight_output_under_kv_pressure(pp_size: int):
+    """reset_prefix_cache(reset_running_requests=True) resumes requests in
+    the same step it preempts them, so in-flight output must be dropped (the
+    resume resamples those positions).
+
+    pp_size=1: regression for the frame-based discard this fix replaces,
+    which with spec decode drained one *token* count per output frame and
+    over-discarded, corrupting the fresh frames after the resume.
+    pp_size=3: back-to-back resets, so the second re-preempts requests whose
+    dropped stale share is still in flight -- it must be recorded once (not
+    accumulated) and stay dropped.
+    """
+    max_tokens = 24
+    scheduler = _create_async_pp_scheduler(num_spec=3, pp_size=pp_size)
+    requests = create_requests(
+        num_requests=8, num_tokens=8, max_tokens=max_tokens, ignore_eos=True
+    )
+    pending = list(requests)
+    for _ in range(2):
+        scheduler.add_request(pending.pop(0))
+
+    # Observe re-preemptions with an undrained stale share (the
+    # double-count hazard).
+    repreempts_with_stale = 0
+    orig_preempt = scheduler._preempt_request
+
+    def counting_preempt(request, timestamp, **kwargs):
+        nonlocal repreempts_with_stale
+        if getattr(request, "num_stale_output_tokens", 0) > 0:
+            repreempts_with_stale += 1
+        return orig_preempt(request, timestamp, **kwargs)
+
+    scheduler._preempt_request = counting_preempt
+
+    resets = 0
+    reset_steps = {6, 14} if pp_size == 1 else {6, 7, 18, 19}
+
+    def before_step(step: int, engine: PipelinedEngine):
+        nonlocal resets
+        if pending:
+            scheduler.add_request(pending.pop(0))
+        if step in reset_steps and (engine.queue or scheduler.running):
+            scheduler.reset_prefix_cache(reset_running_requests=True)
+            resets += 1
+
+    engine = PipelinedEngine(
+        scheduler,
+        queue_size=pp_size + 1,
+        accept_drafts=lambda step, req_id, n: (step + int(req_id)) % (n + 1),
+    )
+    engine.run(before_step=before_step)
+
+    assert resets > 0, "test did not exercise reset_prefix_cache"
+    if pp_size > 1:
+        # The re-preempt-while-stale-pending window needs pipeline depth.
+        assert repreempts_with_stale > 0, (
+            "test did not exercise re-preemption with an undrained stale share"
+        )
+    for req in requests:
+        assert req.is_finished()
+        assert req.num_output_tokens == max_tokens
+        # Dropped tokens are never delivered; order must be preserved with
+        # no duplicates.
+        _assert_ordered_subset(
+            list(req.output_token_ids), engine.emitted[req.request_id]
+        )
+        _assert_positions_consistent(req, engine)
+        # All stale shares fully drained by the end.
+        assert getattr(req, "num_stale_output_tokens", 0) == 0
@@ -1078,8 +1078,6 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():
    # Create minimal mock with just the needed attributes
    mock = SimpleNamespace(
        cache_config=SimpleNamespace(block_size=block_size),
-        max_num_scheduled_tokens=3 * block_size,
-        scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
        use_eagle=False,
        hash_block_size=block_size,
        mamba_partial_cache_hit=False,
@@ -1107,7 +1107,7 @@ def test_preemption_re_records_prefix_cache_query():
    request = create_requests(num_requests=1)[0]
    scheduler.add_request(request)

-    scheduler.schedule()
+    scheduler_output = scheduler.schedule()
    stats = scheduler.kv_cache_manager.prefix_cache_stats
    assert stats is not None
    assert (stats.requests, stats.preempted_requests) == (1, 0)
@@ -1116,6 +1116,21 @@ def test_preemption_re_records_prefix_cache_query():
    scheduler._preempt_request(request, 0.0)
    assert request.status == RequestStatus.PREEMPTED

+    scheduler.update_from_output(
+        scheduler_output,
+        ModelRunnerOutput(
+            req_ids=[request.request_id],
+            req_id_to_index={request.request_id: 0},
+            sampled_token_ids=[[1000]],
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+        ),
+    )
+    assert request.num_stale_output_tokens == 0
+    stats = scheduler.kv_cache_manager.prefix_cache_stats
+    assert stats is not None
+
    scheduler.schedule()
    assert request.status == RequestStatus.RUNNING
    assert stats.preempted_requests == 1
@@ -333,54 +333,6 @@ def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream):
    assert torch.equal(x, torch.full((4,), 15.0, device="cuda"))


-def test_eager_attention_inside_multistream_overlap(cuda_capture_stream):
-    """Handle an eager attention break inside a multi-stream overlap region."""
-    from vllm.compilation.breakable_cudagraph import (
-        BreakableCUDAGraphCapture,
-        eager_break_during_capture,
-    )
-    from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
-
-    x = torch.zeros(1024, device="cuda")
-    output = torch.empty_like(x)
-    aux_stream = torch.cuda.Stream()
-    main_event = torch.cuda.Event()
-    aux_event = torch.cuda.Event()
-
-    @eager_break_during_capture
-    def attention(query: torch.Tensor, out: torch.Tensor) -> None:
-        torch.add(query, 3.0, out=out)
-
-    def attention_frontend() -> torch.Tensor:
-        query = x * 2.0
-        attention_output = torch.empty_like(x)
-        attention(query, attention_output)
-        return attention_output
-
-    cap = BreakableCUDAGraphCapture()
-    with cap:
-        attention_output, gate = maybe_execute_in_parallel(
-            attention_frontend,
-            lambda: x * 5.0,
-            main_event,
-            aux_event,
-            aux_stream,
-        )
-        torch.add(attention_output, gate, out=output)
-
-    assert cap.num_graphs == 2
-    assert cap.num_eager_breaks == 1
-
-    for value in (1.0, 7.0, 19.0):
-        x.fill_(value)
-        cap.replay()
-        cuda_capture_stream.synchronize()
-        torch.testing.assert_close(
-            output,
-            torch.full_like(output, value * 7.0 + 3.0),
-        )
-
-
 # ---------------------------------------------------------------------------
 # Replay ordering
 # ---------------------------------------------------------------------------
@@ -338,6 +338,30 @@ def get_fake_process_mamba_fn(
            assert copy_info[0][-1] == expected_temporal_src
            assert copy_info[1][-1] == expected_temporal_dest

+    def check_fused_copy_info(
+        action: tuple[int, int],
+        align_ctx: mamba_utils.MambaSpecDecodeGPUContext,
+    ):
+        # Align + spec-decode on a hybrid model routes the pre-copy through the
+        # fused kernel (preprocess_mamba -> run_fused_precopy) instead of
+        # do_mamba_copy_block, so copy_info is never populated. Verify from the
+        # fused buffers (req-0 scope, mirroring check_copy_info):
+        #   - the copy DECISION: src_col is -1 iff no pre-copy is scheduled;
+        #   - the DESTINATION column: state_idx == action[1] (curr_state_idx;
+        #     maps directly through block_ids, exactly as check_copy_info's dst).
+        # The source column is NOT asserted here: on the scalar path the source
+        # address is produced by the per-state copy func with an accept-token
+        # bias offset (collect_mamba_copy_meta), so prev_state_idx does not map
+        # to action[0] by plain equality. Source block-level exactness (incl.
+        # the accept-bias) is covered by test_precopy_mamba_align.py.
+        src_col = int(align_ctx.precopy_src_col_buf.np[0])
+        state_idx = int(align_ctx.mamba_state_idx_buf.np[0])
+        if action == (-1, -1):
+            assert src_col == -1
+        else:
+            assert src_col != -1
+            assert state_idx == action[1]
+
    def fake_preprocess_mamba_fn(
        scheduler_output: SchedulerOutput,
        kv_cache_config: KVCacheConfig,
@@ -348,6 +372,7 @@ def get_fake_process_mamba_fn(
        forward_context: dict[str, Any],
        mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
        copy_bufs: mamba_utils.MambaCopyBuffers,
+        align_ctx: mamba_utils.MambaSpecDecodeGPUContext | None = None,
    ):
        nonlocal copy_info
        copy_info = None
@@ -361,14 +386,21 @@ def get_fake_process_mamba_fn(
            forward_context,
            mamba_state_copy_funcs,
            copy_bufs,
+            align_ctx,
        )
        if cur_step_action is not None:
-            check_copy_info(
-                cur_step_action.preprocess_copy_idx,
-                kv_cache_config,
-                forward_context,
-                input_batch,
-            )
+            if align_ctx is not None:
+                check_fused_copy_info(
+                    cur_step_action.preprocess_copy_idx,
+                    align_ctx,
+                )
+            else:
+                check_copy_info(
+                    cur_step_action.preprocess_copy_idx,
+                    kv_cache_config,
+                    forward_context,
+                    input_batch,
+                )
        return ret

    def fake_copy_fn(copy_bufs: mamba_utils.MambaCopyBuffers):
@@ -1398,53 +1398,6 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
    assert stats_after_reset is None


-def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist_init):
-    """reqs_to_send deadlines are stamped with the scheduler process's
-    perf_counter, whose epoch differs across processes and (by boot-time
-    deltas) across nodes. Without rebasing, a P worker on a node whose
-    monotonic clock is ahead of the scheduler's by more than the TTL
-    expires the lease on arrival and reports done_sending before D has
-    read the blocks — the freed blocks can then be reallocated and the
-    remote read pulls another request's data (silent accuracy corruption).
-    The worker must anchor the remaining TTL to its own clock.
-    """
-    vllm_config = create_vllm_config()
-    connector = NixlConnector(
-        vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
-    )
-    connector.connector_worker = FakeNixlConnectorWorker(
-        vllm_config, connector.engine_id, hand_shake_latency=0
-    )
-    worker = connector.connector_worker
-
-    req_id = "req-lease-clock"
-    ttl = 480.0
-    # Simulate a scheduler whose monotonic clock is 10,000 s behind this
-    # worker's (e.g. its node booted much later): the raw deadline is
-    # then already far in the past in this worker's clock domain.
-    scheduler_clock = time.perf_counter() - 10_000.0
-
-    metadata = NixlConnectorMetadata()
-    metadata.reqs_in_batch = {req_id}
-    metadata.reqs_to_send = {req_id: scheduler_clock + ttl}
-    metadata.scheduler_clock = scheduler_clock
-    connector.bind_connector_metadata(metadata)
-    dummy_ctx = ForwardContext(
-        no_compile_layers={},
-        attn_metadata={},
-        slot_mapping={},
-    )
-    connector.start_load_kv(dummy_ctx)
-
-    remaining = worker._reqs_to_send[req_id] - time.perf_counter()
-    assert ttl - 5.0 < remaining <= ttl + 5.0
-
-    # The expiry sweep must not release the request.
-    done_sending, _ = worker.get_finished()
-    assert req_id not in done_sending
-    assert req_id in worker._reqs_to_process
-
-
 def test_kv_connector_stats_aggregation():
    """
    Test KV transfer stats aggregation across TP ranks using
@@ -136,8 +136,12 @@ def test_create_cpu_offloading_spec():


 def test_cpu_spec_sizes_normalized_worker_layout():
+    # The CPU spec now rounds the offloaded row up to the mmap page size
+    # (matching the shared region), so kv_bytes_per_chunk picks up padding
+    # while cpu_page_size_per_worker stays the un-padded per-worker slot.
+    alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
    spec = _create_spec(
-        cpu_bytes_to_use=1920,
+        cpu_bytes_to_use=alignment * 3,
        worker_kv_bytes_per_block=16,
        blocks_per_chunk=2,
        world_size=6,
@@ -147,8 +151,8 @@ def test_cpu_spec_sizes_normalized_worker_layout():

    assert isinstance(spec, CPUOffloadingSpec)
    assert spec.cpu_page_size_per_worker == 32
-    assert spec.kv_bytes_per_chunk == 192
-    assert spec.num_blocks == 10
+    assert spec.kv_bytes_per_chunk == alignment
+    assert spec.num_blocks == 3


 def test_cpu_spec_zero_worker_bytes_produces_empty_cache():
@@ -267,7 +271,9 @@ def test_tiering_spec_create_worker_folds_device_index_for_sharded_layout(monkey

@pytest.mark.parametrize("world_size", [2, 4, 8])
 def test_cpu_spec_replicated_config_preserves_per_rank_sizing(world_size: int):
-    worker_kv_bytes_per_block = 4096
+    # Use a page-multiple per-worker block so the row size is unaffected by the
+    # page-alignment rounding regardless of the host mmap page size.
+    worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
    spec = _create_spec(
        cpu_bytes_to_use=worker_kv_bytes_per_block * world_size * 2,
        worker_kv_bytes_per_block=worker_kv_bytes_per_block,
@@ -282,6 +288,107 @@ def test_cpu_spec_replicated_config_preserves_per_rank_sizing(world_size: int):
    assert spec.num_blocks == 2


+def test_cpu_spec_create_worker_uses_mmap_on_cuda_alike(monkeypatch):
+    import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
+
+    worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
+    spec = _create_spec(
+        cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
+        worker_kv_bytes_per_block=worker_kv_bytes_per_block,
+        world_size=4,
+    )
+    assert isinstance(spec, CPUOffloadingSpec)
+
+    region = MagicMock()
+    region_calls: list[dict[str, Any]] = []
+    worker_calls: list[dict[str, Any]] = []
+
+    def fake_region_ctor(**kwargs):
+        region_calls.append(kwargs)
+        return region
+
+    def fake_worker_ctor(**kwargs):
+        worker_calls.append(kwargs)
+        return MagicMock()
+
+    monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
+    monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor)
+    monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor)
+    monkeypatch.setattr(
+        cpu_spec_module.torch.accelerator, "current_device_index", lambda: 5
+    )
+
+    kv_caches = MagicMock()
+    spec.create_worker(kv_caches)
+
+    # rank folds the physical device index into [0, world_size): 5 % 4 == 1.
+    assert region_calls[0]["rank"] == 1
+    assert region_calls[0]["engine_id"] == "test-engine"
+    assert region_calls[0]["kv_bytes_per_block"] == worker_kv_bytes_per_block * 4
+    assert worker_calls[0]["kv_caches"] is kv_caches
+    assert worker_calls[0]["mmap_region"] is region
+
+
+def test_cpu_spec_create_worker_uses_tensor_path_off_cuda_alike(monkeypatch):
+    import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
+
+    spec = _create_spec(worker_kv_bytes_per_block=4096, world_size=4)
+    assert isinstance(spec, CPUOffloadingSpec)
+
+    region_calls: list[dict[str, Any]] = []
+    worker_calls: list[dict[str, Any]] = []
+
+    def fake_region_ctor(**kwargs):
+        region_calls.append(kwargs)
+        return MagicMock()
+
+    def fake_worker_ctor(**kwargs):
+        worker_calls.append(kwargs)
+        return MagicMock()
+
+    monkeypatch.setattr(
+        cpu_spec_module.current_platform, "is_cuda_alike", lambda: False
+    )
+    monkeypatch.setattr(cpu_spec_module, "SharedOffloadRegion", fake_region_ctor)
+    monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", fake_worker_ctor)
+
+    spec.create_worker(MagicMock())
+
+    # Non-CUDA-alike platforms keep the per-rank pinned-tensor path.
+    assert region_calls == []
+    assert worker_calls[0]["mmap_region"] is None
+
+
+def test_cpu_spec_create_worker_skips_mmap_for_empty_cache(monkeypatch):
+    import vllm.v1.kv_offload.cpu.spec as cpu_spec_module
+
+    # worker_kv_bytes_per_block=0 yields num_blocks=0; a zero-byte region cannot
+    # be mmap'd, so even on CUDA-alike this must fall back to the tensor path.
+    spec = _create_spec(worker_kv_bytes_per_block=0, world_size=4)
+    assert isinstance(spec, CPUOffloadingSpec)
+    assert spec.num_blocks == 0
+
+    region_calls: list[dict[str, Any]] = []
+    worker_calls: list[dict[str, Any]] = []
+
+    monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
+    monkeypatch.setattr(
+        cpu_spec_module,
+        "SharedOffloadRegion",
+        lambda **kwargs: region_calls.append(kwargs),
+    )
+    monkeypatch.setattr(
+        cpu_spec_module,
+        "CPUOffloadingWorker",
+        lambda **kwargs: worker_calls.append(kwargs),
+    )
+
+    spec.create_worker(MagicMock())
+
+    assert region_calls == []
+    assert worker_calls[0]["mmap_region"] is None
+
+
 def test_offloading_spec_has_replicated_layout_default():
    spec = SingleArgExternalOffloadingSpec(_make_offloading_config())
    assert spec.replicated_layout is False
@@ -14,6 +14,7 @@ import torch
 from vllm import LLM, SamplingParams
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
    UnquantizedEmbeddingMethod,
 )

@@ -28,6 +29,7 @@ class _FakeLmHead:
        self.weight = weight
        self.quant_method = object() if quantized else UnquantizedEmbeddingMethod()
        self.shard_indices = shard_indices
+        self.tp_size = 1


 def _build_processor(vocab_size: int) -> LogitsProcessor:
@@ -135,11 +137,59 @@ def test_fp32_head_rejects_quantized_lm_head(default_vllm_config):
        lp._get_logits(torch.randn(4, 16, dtype=torch.bfloat16), lm_head, None)


+def test_replicated_lm_head_skips_tp_communication_and_preserves_processing(
+    default_vllm_config,
+):
+    from unittest import mock
+
+    vocab_size, hidden_size = 12, 8
+    soft_cap, scale = 2.0, 0.5
+    lp = LogitsProcessor(
+        vocab_size,
+        soft_cap=soft_cap,
+        scale=scale,
+    )
+    lp.head_dtype = torch.float32
+
+    hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16)
+    weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)
+    world_size_getter = (
+        "vllm.model_executor.layers.vocab_parallel_embedding."
+        "get_tensor_model_parallel_world_size"
+    )
+    with mock.patch(world_size_getter, return_value=2):
+        lm_head = ParallelLMHead(
+            vocab_size,
+            hidden_size,
+            params_dtype=torch.bfloat16,
+            disable_tp=True,
+        )
+    lm_head.weight_loader(lm_head.weight, weight)
+    assert lm_head.tp_size == 1
+
+    with mock.patch.object(lp, "_gather_logits") as gather_mock:
+        logits = lp(lm_head, hidden_states)
+
+    gather_mock.assert_not_called()
+
+    expected = torch.nn.functional.linear(hidden_states.float(), weight.float())
+    expected = torch.tanh(expected / soft_cap) * soft_cap * scale
+    torch.testing.assert_close(logits, expected)
+
+    all_gather_path = (
+        "vllm.model_executor.layers.logits_processor.tensor_model_parallel_all_gather"
+    )
+    with mock.patch(all_gather_path) as all_gather:
+        top = lp.get_top_tokens(lm_head, hidden_states)
+
+    all_gather.assert_not_called()
+    assert torch.equal(top, expected.argmax(dim=-1))
+
+
 def test_get_top_tokens_honors_head_dtype(default_vllm_config):
    # The spec-decode local-argmax path (get_top_tokens) must run the lm_head
    # in head_dtype too, not just _get_logits.
    import types
-    from unittest import mock

    vocab_size, hidden_size = 64, 16
    lp = _build_processor(vocab_size)
@@ -154,13 +204,7 @@ def test_get_top_tokens_honors_head_dtype(default_vllm_config):
        ),
    )

-    with mock.patch(
-        "vllm.model_executor.layers.logits_processor."
-        "get_tensor_model_parallel_world_size",
-        return_value=1,
-    ):
-        top = lp.get_top_tokens(lm_head, hidden_states, None)
-
+    top = lp.get_top_tokens(lm_head, hidden_states, None)
    expected = torch.nn.functional.linear(hidden_states.float(), weight.float()).argmax(
        dim=-1
    )
@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for SpecDecodeBaseProposer.initialize_attn_backend.
+
+Block tables are stored at kernel-block granularity, so the proposer's
+``block_size`` (used for slot-mapping math) must be the kernel block size,
+not the KV cache manager's block size — the two differ when manager blocks
+are split for the attention kernel. The value must also be deterministic:
+``_draft_attn_layer_names`` is a set, whose iteration order varies across
+processes, so anything derived from iteration order must not leak into
+``block_size``.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+import vllm.v1.spec_decode.llm_base_proposer as llm_base_proposer
+from vllm.v1.spec_decode.eagle import EagleProposer
+
+SCHEDULER_BLOCK_SIZE = 256
+KERNEL_BLOCK_SIZE = 64
+
+
+class _FakeAttentionGroup:
+    def __init__(self, backend, layer_names, kv_cache_spec, kv_cache_group_id):
+        self.backend = backend
+        self.layer_names = list(layer_names)
+        self.kv_cache_spec = kv_cache_spec
+        self.kv_cache_group_id = kv_cache_group_id
+        self.kernel_block_size = None
+
+    def create_metadata_builders(self, vllm_config, device, kernel_block_size=None):
+        self.kernel_block_size = kernel_block_size
+
+    def get_metadata_builder(self):
+        return SimpleNamespace(kv_cache_spec=self.kv_cache_spec)
+
+
+def _make_proposer(
+    monkeypatch: pytest.MonkeyPatch, layer_names: set[str]
+) -> EagleProposer:
+    fake_layers = {}
+    for name in layer_names:
+        backend = SimpleNamespace(full_cls_name=lambda: "FakeBackend")
+        fake_layers[name] = SimpleNamespace(
+            get_attn_backend=lambda backend=backend: backend
+        )
+    monkeypatch.setattr(
+        llm_base_proposer, "get_layers_from_vllm_config", lambda *a, **k: fake_layers
+    )
+    monkeypatch.setattr(llm_base_proposer, "AttentionGroup", _FakeAttentionGroup)
+
+    proposer = EagleProposer.__new__(EagleProposer)
+    proposer.vllm_config = None
+    proposer.device = None
+    proposer._draft_attn_layer_names = set(layer_names)
+    proposer.kv_cache_gid = -1
+    proposer.draft_attn_groups = []
+    proposer.block_size = -1
+    return proposer
+
+
+def _make_kv_cache_config(layer_names: set[str]) -> SimpleNamespace:
+    spec = SimpleNamespace(block_size=SCHEDULER_BLOCK_SIZE)
+    group = SimpleNamespace(layer_names=list(layer_names), kv_cache_spec=spec)
+    return SimpleNamespace(kv_cache_groups=[group])
+
+
+def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch):
+    """The proposer's slot-mapping math runs against the kernel-granularity
+    block table, so block_size must come from kernel_block_sizes."""
+    layer_names = {"draft.0.self_attn.attn"}
+    proposer = _make_proposer(monkeypatch, layer_names)
+
+    proposer.initialize_attn_backend(
+        _make_kv_cache_config(layer_names),
+        kernel_block_sizes=[KERNEL_BLOCK_SIZE],
+    )
+
+    assert proposer.block_size == KERNEL_BLOCK_SIZE
+    assert proposer.block_size != SCHEDULER_BLOCK_SIZE
+    # The metadata builder keeps receiving the kernel block size as well.
+    assert proposer.draft_attn_groups[0].kernel_block_size == KERNEL_BLOCK_SIZE
+
+
+def test_block_size_falls_back_to_kv_cache_spec(monkeypatch: pytest.MonkeyPatch):
+    layer_names = {"draft.0.self_attn.attn"}
+    proposer = _make_proposer(monkeypatch, layer_names)
+
+    proposer.initialize_attn_backend(
+        _make_kv_cache_config(layer_names), kernel_block_sizes=None
+    )
+
+    assert proposer.block_size == SCHEDULER_BLOCK_SIZE
+
+
+def test_draft_layer_iteration_is_deterministic(monkeypatch: pytest.MonkeyPatch):
+    """_draft_attn_layer_names is a set; the attention groups built from it
+    must not depend on its (process-random) iteration order."""
+    layer_names = {"draft.c.attn", "draft.a.attn", "draft.b.attn"}
+    expected_order = sorted(layer_names)
+
+    for insertion_order in (expected_order, expected_order[::-1]):
+        proposer = _make_proposer(monkeypatch, set(insertion_order))
+        proposer.initialize_attn_backend(
+            _make_kv_cache_config(set(insertion_order)),
+            kernel_block_sizes=[KERNEL_BLOCK_SIZE],
+        )
+        assert len(proposer.draft_attn_groups) == 1
+        assert proposer.draft_attn_groups[0].layer_names == expected_order
+        assert proposer.block_size == KERNEL_BLOCK_SIZE
@@ -1,9 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-from types import SimpleNamespace
-
-import pytest
 import torch

 from vllm.platforms import current_platform
@@ -11,7 +8,6 @@ from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.sample.logits_processor import LogitsProcessors
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.spec_decode.llm_base_proposer import (
-    SpecDecodeBaseProposer,
    compute_probs_and_sample_next_token,
 )

@@ -72,21 +68,3 @@ def test_compute_probs_and_sample_next_token_uses_fp64_exponential_race():

    assert torch.equal(actual_ids, expected_ids)
    assert torch.allclose(actual_probs, probs)
-
-
-@pytest.mark.parametrize(
-    ("architecture", "expected"),
-    [
-        ("DeepSeekMTPModel", True),
-        ("KimiK3MTPModel", True),
-        ("MiniMaxM3ForCausalLM", False),
-    ],
-)
-def test_mtp_model_returns_tuple(architecture: str, expected: bool):
-    proposer = object.__new__(SpecDecodeBaseProposer)
-    proposer.method = "mtp"
-    proposer.draft_model_config = SimpleNamespace(
-        hf_config=SimpleNamespace(architectures=[architecture])
-    )
-
-    assert proposer.model_returns_tuple() is expected
@@ -41,6 +41,7 @@ def mock_model_runner_with_req_states():
    runner.sampler = None
    runner.prompt_logprobs_worker = None
    runner.is_last_pp_rank = False
+    runner.pooling_runner = None

    # Mock staged writes — they use Triton kernels that require GPU
    runner.req_states.apply_staged_writes = Mock()
@@ -6,9 +6,9 @@ set -e

 # Default values
 # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake
-DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git"
+DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
 # NOTE: This is currently targeting nv-dev branch due to sm120 support
-DEEPGEMM_GIT_REF="f5a76426fa084087169693fd0cd815223576d6e9"
+DEEPGEMM_GIT_REF="a6b593d2826719dcf4892609af7b84ee23aaf32a"
 WHEEL_DIR=""

 # Parse command line arguments
@@ -133,8 +133,6 @@ def _rocm_aiter_fused_moe_impl(
    bias2: torch.Tensor | None = None,
    moe_sorting_dispatch_policy: int = 0,
    swiglu_limit: float = 0.0,
-    beta: float | None = None,
-    linear_beta: float | None = None,
 ) -> torch.Tensor:
    from aiter import ActivationType, QuantType
    from aiter.fused_moe import fused_moe
@@ -145,12 +143,6 @@ def _rocm_aiter_fused_moe_impl(
    extra_kwargs: dict = {}
    if gate_mode and rocm_aiter_ops.fused_moe_supports_gate_mode():
        extra_kwargs["gate_mode"] = gate_mode
-    if (
-        getattr(ActivationType, "Situv2", None) is not None
-        and activation == ActivationType.Situv2
-    ):
-        extra_kwargs["beta"] = beta
-        extra_kwargs["linear_beta"] = linear_beta

    return fused_moe(
        hidden_states,
@@ -201,8 +193,6 @@ def _rocm_aiter_fused_moe_fake(
    bias2: torch.Tensor | None = None,
    moe_sorting_dispatch_policy: int = 0,
    swiglu_limit: float = 0.0,
-    beta: float | None = None,
-    linear_beta: float | None = None,
 ) -> torch.Tensor:
    if output_dtype is not None:
        return torch.empty_like(hidden_states, dtype=output_dtype)
@@ -1569,7 +1559,6 @@ class rocm_aiter_ops:
            "silu": ActivationType.Silu,
            "gelu": ActivationType.Gelu,
            "swiglu": ActivationType.Swiglu,
-            "situ": getattr(ActivationType, "Situv2", None),
        }
        return mapping.get(name)

@@ -2166,8 +2155,6 @@ class rocm_aiter_ops:
        bias2: torch.Tensor | None = None,
        moe_sorting_dispatch_policy: int = 0,
        swiglu_limit: float = 0.0,
-        beta: float | None = None,
-        linear_beta: float | None = None,
    ) -> torch.Tensor:
        return torch.ops.vllm.rocm_aiter_fused_moe(
            hidden_states,
@@ -2192,8 +2179,6 @@ class rocm_aiter_ops:
            bias2,
            moe_sorting_dispatch_policy,
            swiglu_limit,
-            beta,
-            linear_beta,
        )

    @staticmethod
@@ -2391,17 +2391,6 @@ def topk_softmax(
    e_score_correction_bias: torch.Tensor | None = None,
    is_padding: torch.Tensor | None = None,
 ) -> None:
-    if current_platform.is_xpu():
-        # TODO: Remove after vllm-xpu-kernels supports is_padding.
-        torch.ops._moe_C.topk_softmax(
-            topk_weights,
-            topk_ids,
-            token_expert_indices,
-            gating_output,
-            renormalize,
-            e_score_correction_bias,
-        )
-        return
    torch.ops._moe_C.topk_softmax(
        topk_weights,
        topk_ids,
@@ -2447,21 +2436,6 @@ def topk_hash_softplus_sqrt(
    hash_indices_table: torch.Tensor | None = None,
    is_padding: torch.Tensor | None = None,
 ) -> None:
-    if current_platform.is_xpu():
-        # TODO: Remove after vllm-xpu-kernels supports is_padding.
-        torch.ops._moe_C.topk_softplus_sqrt(
-            topk_weights,
-            topk_indices,
-            token_expert_indices,
-            gating_output,
-            renormalize,
-            routed_scaling_factor,
-            e_score_correction_bias,
-            input_tokens,
-            hash_indices_table,
-        )
-        return
-
    torch.ops._moe_C.topk_softplus_sqrt(
        topk_weights,
        topk_indices,
@@ -771,6 +771,7 @@ class CompilationConfig:
        "vllm::qwen_gdn_attention_core",
        "vllm::gdn_attention_core_xpu",
        "vllm::olmo_hybrid_gdn_full_forward",
+        "vllm::kda_attention",
        "vllm::sparse_attn_indexer",
        "vllm::rocm_aiter_sparse_attn_indexer",
        "vllm::deepseek_v4_attention",
@@ -1028,6 +1028,11 @@ class ParallelConfig:
                "Disabled the custom all-reduce kernel because it is not "
                "supported on current platform."
            )
+        if self.nnodes > 1:
+            self.disable_custom_all_reduce = True
+            logger.debug(
+                "Disabled the custom all-reduce since we are running on multi-node."
+            )
        if self.ray_workers_use_nsight and not self.use_ray:
            raise ValueError(
                "Unable to use nsight profiling unless workers run with Ray."
@@ -51,7 +51,6 @@ MTPModelTypes = Literal[
    "minimax_m3_mtp",
    "bailing_hybrid_mtp",
    "mtp",
-    "kimi_k3_mtp",
    "pangu_ultra_moe_mtp",
    "step3p5_mtp",
    "hy_v3_mtp",
@@ -356,16 +355,6 @@ class SpeculativeConfig:
                {"n_predict": n_predict, "architectures": ["OpenPanguMTPModel"]}
            )

-        if hf_config.model_type == "kimi_k3":
-            # Kimi-K3 keeps the text-model fields (incl. the MTP layer count)
-            # nested under ``text_config`` (a KimiLinearConfig).
-            text_config = getattr(hf_config, "text_config", hf_config)
-            n_predict = getattr(text_config, "num_nextn_predict_layers", None)
-            hf_config.model_type = "kimi_k3_mtp"
-            hf_config.update(
-                {"n_predict": n_predict, "architectures": ["KimiK3MTPModel"]}
-            )
-
        if hf_config.architectures[0] == "MiMoForCausalLM":
            hf_config.model_type = "mimo_mtp"
            n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
@@ -921,6 +910,15 @@ class SpeculativeConfig:
                        f"Unsupported speculative method: '{self.method}'"
                    )

+                if self.method in ("eagle", "eagle3"):
+                    # EAGLE drafts share the target's positional space; a
+                    # draft checkpoint with a smaller max_position_embeddings
+                    # than the target under-sizes its rotary cache (#48894).
+                    SpeculativeConfig._maybe_override_draft_max_position_embeddings(
+                        self.draft_model_config.hf_config,
+                        self.target_model_config.max_model_len,
+                    )
+
                # Replace hf_config for EAGLE draft_model
                if self.method in ("eagle", "eagle3", "dflash"):
                    from vllm.transformers_utils.configs.eagle import EAGLEConfig
@@ -945,7 +943,6 @@ class SpeculativeConfig:
                if self.method == "dspark" and (
                    "Qwen3DSparkModel" not in self.draft_model_config.architectures
                    and "Gemma4DSparkModel" not in self.draft_model_config.architectures
-                    and "K3DSparkModel" not in self.draft_model_config.architectures
                ):
                    # DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config
                    # and its weights ship in the target checkpoint.
@@ -975,16 +972,6 @@ class SpeculativeConfig:
                if self.method in ("dflash", "dspark"):
                    self.parallel_drafting = True

-                if (
-                    self.method == "dspark"
-                    and "K3DSparkModel" in self.draft_model_config.architectures
-                    and self.target_parallel_config.decode_context_parallel_size > 1
-                ):
-                    raise ValueError(
-                        "MLA DSpark does not currently support decode context "
-                        "parallelism; set decode_context_parallel_size=1."
-                    )
-
                if self.num_speculative_tokens is not None and hasattr(
                    self.draft_model_config.hf_config, "num_lookahead_tokens"
                ):
@@ -1152,6 +1139,40 @@ class SpeculativeConfig:
            )
        return result

+    @staticmethod
+    def _maybe_override_draft_max_position_embeddings(
+        draft_hf_config: PretrainedConfig,
+        target_max_model_len: int,
+    ) -> None:
+        """Raise an EAGLE draft's max_position_embeddings up to the target's.
+
+        The proposer feeds the draft positions up to the target's
+        max_model_len, while max_position_embeddings sizes the draft's
+        rotary cos_sin_cache. A smaller checkpoint value (e.g. 2048 for
+        yuhuili/EAGLE3-LLaMA3.1-Instruct-8B) makes that cache gather go
+        out of bounds (#48894).
+
+        Args:
+            draft_hf_config: The draft model's HF config, mutated in place.
+            target_max_model_len: The target model's max_model_len.
+        """
+        draft_max_position_embeddings = getattr(
+            draft_hf_config, "max_position_embeddings", None
+        )
+        if (
+            draft_max_position_embeddings is None
+            or draft_max_position_embeddings >= target_max_model_len
+        ):
+            return
+        logger.info(
+            "Overriding draft model max_position_embeddings from %d to the "
+            "target model's max_model_len (%d); EAGLE drafts share the "
+            "target's positional space.",
+            draft_max_position_embeddings,
+            target_max_model_len,
+        )
+        draft_hf_config.max_position_embeddings = target_max_model_len
+
    @staticmethod
    def _verify_and_get_draft_tp(
        target_parallel_config: ParallelConfig,
@@ -72,34 +72,11 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
        "GraniteMoeForCausalLM",
        "InklingForCausalLM",
        "InklingForConditionalGeneration",
-        "KimiK3ForConditionalGeneration",
        "LongcatFlashNgramForCausalLM",
        "Qwen2MoeForCausalLM",
    }
 )

-# Architectures that default to V1 on ROCm: the V2 runner faults during the
-# profile run. VLLM_USE_V2_MODEL_RUNNER=1 still forces V2.
-# TODO: fix V2 enablement
-ROCM_EXCLUDED_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
-    {
-        "KimiK3ForConditionalGeneration",
-    }
-)
-
-
-@lru_cache
-def default_v2_model_runner_architectures() -> frozenset[str]:
-    """Architectures defaulting to the V2 model runner on this platform."""
-    from vllm.platforms import current_platform
-
-    if current_platform.is_rocm():
-        return (
-            DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES
-            - ROCM_EXCLUDED_V2_MODEL_RUNNER_ARCHITECTURES
-        )
-    return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES
-

 class OptimizationLevel(IntEnum):
    """Optimization level enum."""
@@ -642,20 +619,16 @@ class VllmConfig:
        if model_config.runner_type != "generate":
            return False

-        architectures = getattr(model_config, "architectures", [])
-        default_architectures = default_v2_model_runner_architectures()
-        is_default_v2_architecture = any(
-            arch in default_architectures for arch in architectures
-        )
-
-        if getattr(model_config, "is_hybrid", False) and (
-            not is_default_v2_architecture
-        ):
+        if getattr(model_config, "is_hybrid", False):
            return False

        if getattr(model_config, "is_attention_free", False):
            return False
-        return is_default_v2_architecture or not model_config.is_moe
+        architectures = getattr(model_config, "architectures", [])
+        return (
+            any(arch in DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES for arch in architectures)
+            or not model_config.is_moe
+        )

    @property
    def needs_dp_coordinator(self) -> bool:
@@ -1223,9 +1196,6 @@ class VllmConfig:
                    "DeepSeekV4MTPModel",
                    "InklingForCausalLM",
                    "InklingForConditionalGeneration",
-                    "KimiK3ForConditionalGeneration",
-                    "KimiK3MTPModel",
-                    "KimiLinearForCausalLM",
                    "MiniMaxM3SparseForCausalLM",
                    "MiniMaxM3SparseForConditionalGeneration",
                )
@@ -2307,6 +2277,14 @@ class VllmConfig:

        # Mamba cache align-mode constraints
        if self.cache_config.mamba_cache_mode == "align":
+            assert block_size <= self.scheduler_config.max_num_batched_tokens, (
+                "In Mamba cache align mode, block_size "
+                f"({block_size}) must be <= "
+                "max_num_batched_tokens "
+                f"({self.scheduler_config.max_num_batched_tokens})."
+            )
+            if self.scheduler_config.long_prefill_token_threshold > 0:
+                assert self.scheduler_config.long_prefill_token_threshold >= block_size
            assert not self.scheduler_config.disable_chunked_mm_input, (
                "Chunked MM input is required because we need the flexibility "
                "to schedule a multiple of block_size tokens even if they are "
@@ -340,18 +340,6 @@ class CudaCommunicator(DeviceCommunicatorBase):
            torch.distributed.all_reduce(out, group=self.device_group)
        return out

-    def custom_all_gather(self, input_: torch.Tensor) -> torch.Tensor | None:
-        ca_comm = self.ca_comm
-        if ca_comm is None:
-            return None
-        return ca_comm.custom_all_gather(input_.contiguous())
-
-    def custom_reduce_scatter(self, input_: torch.Tensor) -> torch.Tensor | None:
-        ca_comm = self.ca_comm
-        if ca_comm is None:
-            return None
-        return ca_comm.custom_reduce_scatter(input_.contiguous())
-
    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
        # Route uniform dim-0 all-gathers through NVLS symmetric memory when
        # enabled (mirrors reduce_scatter); otherwise fall back to the
@@ -25,11 +25,6 @@ except Exception:
    # For CPUs
    custom_ar = False

-try:
-    import torch.distributed._symmetric_memory as torch_symm_mem
-except ImportError:
-    torch_symm_mem = None
-
 logger = init_logger(__name__)


@@ -54,17 +49,7 @@ from vllm.distributed.utils import is_weak_contiguous  # noqa: E402


 class CustomAllreduce:
-    _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8, 16]
-    _DEFAULT_ALL_GATHER_MAX_SIZE = 2 * 1024 * 1024
-    _DEFAULT_MNNVL_ALL_GATHER_MAX_SIZES = {
-        2: 8 * 1024 * 1024,
-        4: 4 * 1024 * 1024,
-        6: 2 * 1024 * 1024,
-        8: 2 * 1024 * 1024,
-        16: 2 * 1024 * 1024,
-    }
-    _DEFAULT_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024
-    _DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024
+    _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8]

    # max_size: max supported allreduce size
    def __init__(
@@ -72,10 +57,6 @@ class CustomAllreduce:
        group: ProcessGroup,
        device: int | str | torch.device,
        max_size=8192 * 1024,
-        max_all_gather_size=_DEFAULT_ALL_GATHER_MAX_SIZE,
-        max_mnnvl_all_gather_size=None,
-        max_reduce_scatter_size=_DEFAULT_REDUCE_SCATTER_MAX_SIZE,
-        max_mnnvl_reduce_scatter_size=_DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE,
        symm_mem_enabled=False,
    ) -> None:
        """
@@ -89,20 +70,7 @@ class CustomAllreduce:
        are in the same node.
        """
        self._IS_CAPTURING = False
-        self._ptr = 0
        self.disabled = True
-        self.mnnvl_buffer = None
-        self.mnnvl_handle = None
-        self.mnnvl_peer_buffers: list[torch.Tensor] | None = None
-        self.mnnvl_multicast_ptr = 0
-        self.mnnvl_buffer_size = 0
-        self.mnnvl_lamport_ag_local_ptr = 0
-        self.mnnvl_lamport_ag_multicast_ptr = 0
-        self.mnnvl_lamport_rs_local_ptr = 0
-        self.mnnvl_lamport_epochs = None
-        self.mnnvl_lamport_ag_epoch_ptr = 0
-        self.mnnvl_lamport_rs_epoch_ptr = 0
-        self.mnnvl_only = False

        if not custom_ar:
            # disable because of missing custom allreduce library
@@ -119,8 +87,13 @@ class CustomAllreduce:
            "CustomAllreduce should be attached to a non-NCCL group."
        )

-        same_node = all(in_the_same_node_as(group, source_rank=0))
-        self.mnnvl_only = not same_node
+        if not all(in_the_same_node_as(group, source_rank=0)):
+            # No need to initialize custom allreduce for multi-node case.
+            logger.warning(
+                "Custom allreduce is disabled because this process group"
+                " spans across nodes."
+            )
+            return

        rank = dist.get_rank(group=self.group)
        self.rank = rank
@@ -153,30 +126,28 @@ class CustomAllreduce:
            and device_capability is not None
        ):
            device_capability_str = device_capability.as_version_str()
-            if (
-                device_capability_str in CUSTOM_ALL_REDUCE_MAX_SIZES
-                and world_size in CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str]
-            ):
+            if device_capability_str in CUSTOM_ALL_REDUCE_MAX_SIZES:
                max_size = min(
                    CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size],
                    max_size,
                )
        # device.index is a visible ordinal, not a logical local ID.
-        fully_connected = False
-        if same_node:
-            physical_device_id = (
-                current_platform.visible_device_id_to_physical_device_id(device.index)
-            )
-            tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
-            gather_list = [
-                torch.tensor([0], dtype=torch.int, device="cpu")
-                for _ in range(world_size)
-            ]
-            dist.all_gather(gather_list, tensor, group=self.group)
-            physical_device_ids = [t.item() for t in gather_list]
-            assert current_platform.is_cuda_alike()
-            fully_connected = current_platform.is_fully_connected(physical_device_ids)
-        if same_node and world_size > 2 and not fully_connected:
+        physical_device_id = current_platform.visible_device_id_to_physical_device_id(
+            device.index
+        )
+        tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
+        gather_list = [
+            torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)
+        ]
+        dist.all_gather(gather_list, tensor, group=self.group)
+        physical_device_ids = [t.item() for t in gather_list]
+
+        # test nvlink first, this will filter out most of the cases
+        # where custom allreduce is not supported
+        # this checks hardware and driver support for NVLink
+        assert current_platform.is_cuda_alike()
+        fully_connected = current_platform.is_fully_connected(physical_device_ids)
+        if world_size > 2 and not fully_connected:
            logger.warning(
                "Custom allreduce is disabled because it's not supported on"
                " more than two PCIe-only GPUs. To silence this warning, "
@@ -187,11 +158,7 @@ class CustomAllreduce:
        # this is expensive to compute at the first time
        # then we cache the result
        # On AMD GPU, p2p is always enabled between XGMI connected GPUs
-        if (
-            same_node
-            and not current_platform.is_rocm()
-            and not _can_p2p(rank, world_size)
-        ):
+        if not current_platform.is_rocm() and not _can_p2p(rank, world_size):
            logger.warning(
                "Custom allreduce is disabled because your platform lacks "
                "GPU P2P capability or P2P test failed. To silence this "
@@ -203,40 +170,21 @@ class CustomAllreduce:
        # Buffers memory are owned by this Python class and passed to C++.
        # Metadata composes of two parts: metadata for synchronization and a
        # temporary buffer for storing intermediate allreduce results.
-        if same_node:
-            self.meta_ptrs = self.create_shared_buffer(
-                ops.meta_size() + max_size, group=group, uncached=True
-            )
-        else:
-            meta_ptr, _ = ops.allocate_shared_buffer_and_handle(ops.meta_size())
-            self.meta_ptrs = [meta_ptr] * world_size
+        self.meta_ptrs = self.create_shared_buffer(
+            ops.meta_size() + max_size, group=group, uncached=True
+        )
        # This is a pre-registered IPC buffer. In eager mode, input tensors
-        # are first copied into this buffer before the operation is performed
-        legacy_buffer_size = max(max_size, max_all_gather_size, max_reduce_scatter_size)
-        if same_node:
-            self.buffer_ptrs = self.create_shared_buffer(
-                legacy_buffer_size,
-                group=group,
-            )
-        else:
-            buffer_ptr, _ = ops.allocate_shared_buffer_and_handle(legacy_buffer_size)
-            self.buffer_ptrs = [buffer_ptr] * world_size
-        # This stores tuples of pointers to IPC buffers from all ranks.
-        # Each registered tuple contains at most 16 addresses.
-        # Allocating 8MB is enough for 65536 such tuples. The largest model uses
-        # fewer than 10000 registered tuples.
+        # are first copied into this buffer before allreduce is performed
+        self.buffer_ptrs = self.create_shared_buffer(max_size, group=group)
+        # This is a buffer for storing the tuples of pointers pointing to
+        # IPC buffers from all ranks. Each registered tuple has size of
+        # 8*world_size bytes where world_size is at most 8. Allocating 8MB
+        # is enough for 131072 such tuples. The largest model I've seen only
+        # needs less than 10000 of registered tuples.
        self.rank_data = torch.empty(
            8 * 1024 * 1024, dtype=torch.uint8, device=self.device
        )
        self.max_size = max_size
-        self.max_all_gather_size = max_all_gather_size
-        if max_mnnvl_all_gather_size is None:
-            max_mnnvl_all_gather_size = self._DEFAULT_MNNVL_ALL_GATHER_MAX_SIZES[
-                world_size
-            ]
-        self.max_mnnvl_all_gather_size = max_mnnvl_all_gather_size
-        self.max_reduce_scatter_size = max_reduce_scatter_size
-        self.max_mnnvl_reduce_scatter_size = max_mnnvl_reduce_scatter_size
        self.rank = rank
        self.world_size = world_size
        self.fully_connected = fully_connected
@@ -244,72 +192,6 @@ class CustomAllreduce:
            self.meta_ptrs, self.rank_data, rank, self.fully_connected
        )
        ops.register_buffer(self._ptr, self.buffer_ptrs)
-        self._init_mnnvl_buffer(
-            max(
-                max_mnnvl_all_gather_size * world_size,
-                max_mnnvl_reduce_scatter_size,
-            )
-        )
-        if not same_node and not self.mnnvl_multicast_ptr:
-            logger.warning(
-                "Custom collectives are disabled because this multi-node "
-                "group does not support MNNVL multicast."
-            )
-            self.close()
-            self.disabled = True
-
-    def _init_mnnvl_buffer(self, stage_size: int) -> None:
-        if torch_symm_mem is None or not current_platform.is_cuda():
-            return
-        try:
-            buffer_size = stage_size * 6
-            buffer = torch_symm_mem.empty(
-                buffer_size, dtype=torch.uint8, device=self.device
-            )
-            handle = torch_symm_mem.rendezvous(buffer, self.group.group_name)
-            if handle.multicast_ptr == 0:
-                return
-            peer_buffers = [
-                handle.get_buffer(
-                    peer,
-                    (buffer_size,),
-                    torch.uint8,
-                    storage_offset=0,
-                )
-                for peer in range(self.world_size)
-            ]
-            ptrs = [peer_buffer.data_ptr() for peer_buffer in peer_buffers]
-            lamport_ag_offset = 0
-            lamport_rs_offset = stage_size * 3
-            lamport_ag_ptrs = [ptr + lamport_ag_offset for ptr in ptrs]
-            lamport_rs_ptrs = [ptr + lamport_rs_offset for ptr in ptrs]
-            ops.register_buffer(self._ptr, lamport_ag_ptrs)
-            ops.register_buffer(self._ptr, lamport_rs_ptrs)
-
-            buffer.view(torch.int32).fill_(-2147483648)
-            epochs = torch.zeros(
-                (2, 32),
-                dtype=torch.int32,
-                device=self.device,
-            )
-            torch.accelerator.synchronize()
-            dist.barrier(group=self.group)
-
-            self.mnnvl_buffer = buffer
-            self.mnnvl_handle = handle
-            self.mnnvl_peer_buffers = peer_buffers
-            self.mnnvl_multicast_ptr = handle.multicast_ptr
-            self.mnnvl_buffer_size = stage_size
-            self.mnnvl_lamport_ag_local_ptr = lamport_ag_ptrs[self.rank]
-            self.mnnvl_lamport_ag_multicast_ptr = (
-                handle.multicast_ptr + lamport_ag_offset
-            )
-            self.mnnvl_lamport_rs_local_ptr = lamport_rs_ptrs[self.rank]
-            self.mnnvl_lamport_epochs = epochs
-            self.mnnvl_lamport_ag_epoch_ptr = epochs[0].data_ptr()
-            self.mnnvl_lamport_rs_epoch_ptr = epochs[1].data_ptr()
-        except RuntimeError as error:
-            logger.debug("MNNVL AG/RS initialization failed: %s", error)

    @contextmanager
    def capture(self):
@@ -346,7 +228,7 @@ class CustomAllreduce:
        ops.register_graph_buffers(self._ptr, handles, offsets)

    def should_custom_ar(self, inp: torch.Tensor):
-        if self.disabled or self.world_size > 8:
+        if self.disabled:
            return False
        inp_size = inp.numel() * inp.element_size()
        # custom allreduce requires input byte size to be multiples of 16
@@ -397,111 +279,6 @@ class CustomAllreduce:
            # latency) compared to the performance gain of using custom kernels
            return self.all_reduce(input, registered=False)

-    def should_custom_all_gather(self, inp: torch.Tensor) -> bool:
-        if self.disabled or not current_platform.is_cuda():
-            return False
-        if self.world_size == 16 and not self.mnnvl_only:
-            return False
-        inp_size = inp.nbytes
-        if inp.dtype not in (
-            torch.float32,
-            torch.float16,
-            torch.bfloat16,
-        ):
-            return False
-        max_size = (
-            self.max_mnnvl_all_gather_size
-            if self.mnnvl_multicast_ptr
-            else self.max_all_gather_size
-        )
-        return (
-            0 < inp_size <= max_size
-            and inp_size % 16 == 0
-            and is_weak_contiguous(inp)
-            and (self.fully_connected or bool(self.mnnvl_multicast_ptr))
-        )
-
-    def custom_all_gather(self, inp: torch.Tensor) -> torch.Tensor | None:
-        if not self.should_custom_all_gather(inp):
-            return None
-        out_shape = (inp.shape[0] * self.world_size,) + inp.shape[1:]
-        if self.mnnvl_multicast_ptr:
-            logger.info_once(
-                "Using the MNNVL Lamport all-gather kernel.",
-                scope="global",
-            )
-            out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device)
-            ops.mnnvl_lamport_all_gather(
-                self._ptr,
-                inp,
-                out,
-                self.mnnvl_lamport_ag_local_ptr,
-                self.mnnvl_lamport_ag_multicast_ptr,
-                self.mnnvl_lamport_ag_epoch_ptr,
-                self.mnnvl_buffer_size,
-            )
-        else:
-            out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device)
-            ops.custom_all_gather(
-                self._ptr,
-                inp,
-                out,
-                self.buffer_ptrs[self.rank],
-                self.max_all_gather_size,
-            )
-        return out
-
-    def should_custom_reduce_scatter(self, inp: torch.Tensor) -> bool:
-        if self.disabled or not current_platform.is_cuda():
-            return False
-        if self.world_size == 16 and not self.mnnvl_only:
-            return False
-        inp_size = inp.nbytes
-        if inp.dtype not in (torch.float32, torch.float16, torch.bfloat16):
-            return False
-        if inp.shape[0] % self.world_size != 0:
-            return False
-        output_size = inp_size // self.world_size
-        max_size = (
-            self.max_mnnvl_reduce_scatter_size
-            if self.mnnvl_multicast_ptr
-            else self.max_reduce_scatter_size
-        )
-        return (
-            0 < inp_size <= max_size
-            and output_size % 16 == 0
-            and is_weak_contiguous(inp)
-            and (self.fully_connected or bool(self.mnnvl_multicast_ptr))
-        )
-
-    def custom_reduce_scatter(self, inp: torch.Tensor) -> torch.Tensor | None:
-        if not self.should_custom_reduce_scatter(inp):
-            return None
-        out_shape = (inp.shape[0] // self.world_size,) + inp.shape[1:]
-        out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device)
-        if self.mnnvl_multicast_ptr:
-            logger.info_once(
-                "Using the MNNVL Lamport reduce-scatter kernel.",
-                scope="global",
-            )
-            ops.mnnvl_lamport_reduce_scatter(
-                self._ptr,
-                inp,
-                out,
-                self.mnnvl_lamport_rs_local_ptr,
-                self.mnnvl_lamport_rs_epoch_ptr,
-                self.mnnvl_buffer_size,
-            )
-        else:
-            ops.custom_reduce_scatter(
-                self._ptr,
-                inp,
-                out,
-                self.buffer_ptrs[self.rank],
-                self.max_reduce_scatter_size,
-            )
-        return out
-
    def close(self):
        if not self.disabled and self._ptr:
            if ops is not None:
@@ -509,10 +286,6 @@ class CustomAllreduce:
            self._ptr = 0
            self.free_shared_buffer(self.meta_ptrs, rank=self.rank)
            self.free_shared_buffer(self.buffer_ptrs, rank=self.rank)
-            self.mnnvl_peer_buffers = None
-            self.mnnvl_handle = None
-            self.mnnvl_buffer = None
-            self.mnnvl_lamport_epochs = None

    def __del__(self):
        self.close()
@@ -275,3 +275,15 @@ class ECConnectorBase(ABC):
            get_finished().
        """
        return False, None
+
+    def has_pending_push_work(self) -> bool:
+        """Return True if the connector has push-mode work that requires
+        the engine main loop to keep stepping (e.g. for EPD,
+        Producer has push work when Xfer is in progress - Consumer
+        is reading it).
+        This mirrors exactly the KV Connector's has_pending_push_work().
+
+        Connectors that don't implement push-based EC transfer should
+        leave this as False.
+        """
+        return False
@@ -418,10 +418,6 @@ class NixlBaseConnectorScheduler:
            self._build_save_meta(meta, scheduler_output)

        meta.reqs_to_send = self._reqs_need_send
-        # Clock reference for reqs_to_send: deadlines above are in this
-        # process's perf_counter domain; workers (possibly on other nodes,
-        # where perf_counter has a different epoch) rebase against this.
-        meta.scheduler_clock = time.perf_counter()
        meta.reqs_in_batch = self._reqs_in_batch
        meta.reqs_not_processed = self._reqs_not_processed

@@ -180,12 +180,6 @@ class NixlConnectorMetadata(KVConnectorMetadata):
        self.reqs_to_recv: dict[ReqId, ReqMeta] = {}
        self.reqs_to_save: dict[ReqId, ReqMeta] = {}
        self.reqs_to_send: dict[ReqId, float] = {}
-        # The scheduler process's time.perf_counter() when this metadata was
-        # built. reqs_to_send deadlines are stamped with the scheduler's
-        # clock, which is NOT comparable across processes (perf_counter is
-        # process/boot-local): workers must rebase the remaining TTL onto
-        # their own clock via this reference. 0.0 = unset (legacy metadata).
-        self.scheduler_clock: float = 0.0
        self.reqs_in_batch: set[ReqId] = set()
        self.reqs_not_processed: set[ReqId] = set()
        # Heartbeat data grouped by remote engine, sent by D worker to P.
@@ -92,20 +92,8 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
            assert req_id not in self._reqs_to_send

        # Add to requests that are waiting to be read and track expiration.
-        # Deadlines are stamped with the scheduler process's perf_counter,
-        # which is not comparable to ours when the worker runs in another
-        # process on another node (perf_counter epochs differ by boot time).
-        # Rebase the remaining TTL onto our clock; broadcast latency only
-        # lengthens the lease, which is the safe direction. A cross-node
-        # epoch gap larger than the TTL otherwise expires the lease on
-        # arrival and the blocks are freed before D reads them.
-        now_local = time.perf_counter()
        for req_id, expiration_time in metadata.reqs_to_send.items():
            if req_id in self._reqs_to_process:
-                if metadata.scheduler_clock:
-                    expiration_time = now_local + (
-                        expiration_time - metadata.scheduler_clock
-                    )
                self._reqs_to_send[req_id] = expiration_time

        # Send heartbeats to P-side engines to keep KV blocks alive while
@@ -191,15 +191,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
        for req_id in metadata.reqs_not_processed:
            self._reqs_to_process.discard(req_id)
            assert req_id not in self._reqs_to_send
-        # Rebase scheduler-clock deadlines onto this worker's clock — see the
-        # equivalent block in pull_worker.start_load_kv for the rationale.
-        now_local = time.perf_counter()
        for req_id, expiration_time in metadata.reqs_to_send.items():
            if req_id in self._reqs_to_process:
-                if metadata.scheduler_clock:
-                    expiration_time = now_local + (
-                        expiration_time - metadata.scheduler_clock
-                    )
                self._reqs_to_send[req_id] = expiration_time

        # Heartbeats still leave from the main thread (base worker behaviour).
@@ -741,7 +741,6 @@ class EngineArgs:

    fail_on_environ_validation: bool = False
    gdn_prefill_backend: Literal["flashinfer", "triton", "cutedsl"] | None = None
-    kda_prefill_backend: Literal["auto", "triton", "flashkda"] | None = None

    def __post_init__(self):
        # support `EngineArgs(compilation_config={...})`
@@ -1638,13 +1637,6 @@ class EngineArgs:
            default=None,
            help="Select GDN prefill backend.",
        )
-        parser.add_argument(
-            "--kda-prefill-backend",
-            dest="kda_prefill_backend",
-            choices=["auto", "triton", "flashkda"],
-            default=None,
-            help="Select KDA prefill backend.",
-        )
        return parser

    @classmethod
@@ -2439,8 +2431,6 @@ class EngineArgs:

        if self.gdn_prefill_backend is not None:
            self.additional_config["gdn_prefill_backend"] = self.gdn_prefill_backend
-        if self.kda_prefill_backend is not None:
-            self.additional_config["kda_prefill_backend"] = self.kda_prefill_backend

        config = VllmConfig(
            model_config=model_config,
@@ -644,7 +644,7 @@ class SpeechToTextBaseServing(GenerateBaseServing):
                        TranscriptionResponseVerbose(
                            text=text,
                            language=request.language,
-                            duration=str(duration_s),
+                            duration=duration_s,
                            segments=total_segments,
                        ),
                    )
@@ -658,7 +658,7 @@ class SpeechToTextBaseServing(GenerateBaseServing):
                        TranslationResponseVerbose(
                            text=text,
                            language=request.language,
-                            duration=str(duration_s),
+                            duration=duration_s,
                            segments=total_segments,
                        ),
                    )
@@ -384,7 +384,7 @@ class TranscriptionSegment(OpenAIBaseModel):


 class TranscriptionResponseVerbose(OpenAIBaseModel):
-    duration: str
+    duration: float
    """The duration of the input audio."""

    language: str
@@ -357,7 +357,7 @@ class TranslationSegment(OpenAIBaseModel):


 class TranslationResponseVerbose(OpenAIBaseModel):
-    duration: str
+    duration: float
    """The duration of the input audio."""

    language: str
@@ -131,7 +131,6 @@ if TYPE_CHECKING:
    VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False
    VLLM_ROCM_USE_AITER_MOE: bool = True
    VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0
-    AITER_SITUV2_A8W4: bool = False
    VLLM_ROCM_USE_AITER_RMSNORM: bool = True
    VLLM_ROCM_USE_AITER_MLA: bool = True
    VLLM_ROCM_USE_AITER_MHA: bool = True
@@ -177,7 +176,6 @@ if TYPE_CHECKING:
    VLLM_RAY_EXTRA_ENV_VARS_TO_COPY: str = ""
    VLLM_MARLIN_USE_ATOMIC_ADD: bool = False
    VLLM_MARLIN_INPUT_DTYPE: Literal["int8", "fp8"] | None = None
-    VLLM_MARLIN_MXFP8_INPUT_QDQ: bool = False
    VLLM_HUMMING_ONLINE_QUANT_CONFIG: dict[str, Any] | None = None
    VLLM_HUMMING_INPUT_QUANT_CONFIG: dict[str, Any] | None = None
    VLLM_HUMMING_USE_F16_ACCUM: bool = False
@@ -274,7 +272,6 @@ if TYPE_CHECKING:
    VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION: bool = False
    VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
    VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
-    VLLM_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD: int = 256
    VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024
    VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
    VLLM_USE_V2_MODEL_RUNNER: bool | None = None
@@ -1007,6 +1004,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
    # Backend for Video IO — selects the frame-sampling algorithm.
    # - "opencv": uniform sampling.
    # - "opencv_dynamic": duration-aware dynamic sampling.
+    # - "identity": returns raw video bytes for model processor to handle.
    #
    # Custom backend implementations can be registered
    # via `@VIDEO_LOADER_REGISTRY.register("my_custom_video_loader")` and
@@ -1213,12 +1211,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
    "VLLM_ROCM_USE_AITER_MOE": lambda: (
        os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1")
    ),
-    # Route K3 SiTU MXFP4 MoE through the a8w4 (fp8 activation) gate/up-
-    # interleaved flydsl kernels instead of the default a16w4 separated path.
-    # Shared with the AITER runtime, which reads the same env var directly.
-    "AITER_SITUV2_A8W4": lambda: (
-        os.getenv("AITER_SITUV2_A8W4", "0").lower() in ("true", "1")
-    ),
    # MoE sorting dispatch policy for AITER fused MoE kernels.
    #   0 = auto (default): single-pass for small batches, multi-pass
    #       for large batches
@@ -1456,11 +1448,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
    "VLLM_MARLIN_INPUT_DTYPE": env_with_choices(
        "VLLM_MARLIN_INPUT_DTYPE", None, ["int8", "fp8"]
    ),
-    # Debug-only: simulate W4A8 activations on W4A16 Marlin by applying
-    # MXFP8 quantize-dequantize to both Marlin GEMM inputs.
-    "VLLM_MARLIN_MXFP8_INPUT_QDQ": lambda: bool(
-        int(os.getenv("VLLM_MARLIN_MXFP8_INPUT_QDQ", "0"))
-    ),
    # The online quantization dtype for humming kernel
    "VLLM_HUMMING_ONLINE_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file(
        os.environ.get("VLLM_HUMMING_ONLINE_QUANT_CONFIG", None)
@@ -1924,14 +1911,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
    "VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD": lambda: int(
        int(os.getenv("VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD", 256))
    ),
-    # Token-count cutoff for overlapping the MoE router gate with the
-    # routed-expert down projection on a separate CUDA stream (latent MoE).
-    # At or below this many tokens the launch-bound decode path benefits from
-    # multi-stream overlap; above it the GEMMs saturate the device and the
-    # cross-stream sync is pure overhead, so it falls back to sequential.
-    "VLLM_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD": lambda: int(
-        os.getenv("VLLM_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD", "256")
-    ),
    # Token-count cutoff for multi-stream overlap of the attention input
    # GEMM with auxiliary GEMMs (e.g. fused_wqa_wkv overlapped with indexer
    # weights / kv-score projections in DeepSeek-V4). At or below this many
@@ -30,9 +30,6 @@ def rms_norm(
    x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
 ) -> Tensor:
    assert variance_size is None
-    if weight is None:
-        # Kernel requires weight tensor, pass ones
-        weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
    output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
    torch.ops._C.rms_norm(output, x, weight, epsilon)
    return output
@@ -58,8 +55,5 @@ def fused_add_rms_norm(
    variance_size: int | None = None,
 ) -> tuple[Tensor, Tensor]:
    assert variance_size is None
-    if weight is None:
-        # Kernel requires weight tensor, pass ones
-        weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
    torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon)
    return x, x_residual
@@ -272,7 +272,6 @@ from vllm.v1.attention.backends.mla.prefill import (
 )
 from vllm.v1.attention.backends.utils import (
    get_dcp_local_seq_lens,
-    get_num_attention_heads_from_layers,
    split_decodes_and_prefills,
 )
 from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs
@@ -376,7 +375,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
        use_sparse: bool = False,
        indexer: object | None = None,
        topk_indices_buffer: torch.Tensor | None = None,
-        non_causal_multi_token_decode: bool = False,
        **extra_impl_args,
    ):
        super().__init__()
@@ -393,7 +391,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
        self.head_size = kv_lora_rank + qk_rope_head_dim
        self.layer_name = prefix
        self.indexer = indexer
-        self.non_causal_multi_token_decode = non_causal_multi_token_decode
        self.num_kv_heads = 1
        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim

@@ -1126,7 +1123,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
            dtype=kv_cache_dtype,
            cache_dtype_str=self.kv_cache_dtype,
            kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
-            non_causal_multi_token_decode=self.non_causal_multi_token_decode,
        )

    def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
@@ -1454,8 +1450,6 @@ class MLACommonMetadata(AttentionMetadata, Generic[D]):
    num_decode_tokens: int
    num_prefills: int

-    causal: bool = True
-
    # The dimension of the attention heads
    head_dim: int | None = None

@@ -1730,9 +1724,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
    # speculative decoding is enabled.
    query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.SINGLE_ONLY

-    # Whether this builder can flatten a non-causal query block into decode rows.
-    supports_non_causal_multi_token_decode: ClassVar[bool] = False
-
    # The threshold for reordering the batch into decode and prefill requests.
    # If > 1, the batch will be reordered such that requests with
    # query length <= threshold are classified as decode requests.
@@ -1833,14 +1824,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
        self.vllm_config = vllm_config
        self.device = device
        self.use_pcp = parallel_config.prefill_context_parallel_size > 1
-        self.non_causal_multi_token_decode = getattr(
-            kv_cache_spec, "non_causal_multi_token_decode", False
-        )

-        # A draft cache group can have a different head count from the target.
-        self.num_heads = get_num_attention_heads_from_layers(
-            vllm_config, layer_names
-        ) or self.model_config.get_num_attention_heads(parallel_config)
+        self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
        self.mla_dims = get_mla_dims(self.model_config)
        self.aot_schedule = current_platform.is_cuda()

@@ -1966,42 +1951,14 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
        seq_lens = common_attn_metadata.seq_lens
        dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens

-        non_causal_decode = common_attn_metadata.causal is False
-        if non_causal_decode:
-            if not (
-                self.supports_non_causal_multi_token_decode
-                and self.non_causal_multi_token_decode
-            ):
-                raise ValueError(
-                    "Non-causal multi-token MLA requires an explicitly supported "
-                    "attention group."
-                )
-            query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
-            num_active_reqs = int(torch.count_nonzero(query_lens > 0))
-            uniform_active_queries = num_active_reqs > 0 and bool(
-                torch.all(query_lens[:num_active_reqs] == query_lens[0])
-            )
-            trailing_graph_padding = bool(torch.all(query_lens[num_active_reqs:] == 0))
-            if not (uniform_active_queries and trailing_graph_padding):
-                raise ValueError(
-                    "Non-causal MLA requires a uniform query block; got query "
-                    f"lengths {query_lens.tolist()}."
-                )
-            # Use exact GPU sequence lengths instead of the prefill path's CPU
-            # context-length upper bounds.
-            num_decodes = num_reqs
-            num_prefills = 0
-            num_decode_tokens = num_tokens
-            num_prefill_tokens = 0
-        else:
-            num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
-                split_decodes_and_prefills(
-                    common_attn_metadata,
-                    decode_threshold=self.reorder_batch_threshold,
-                    require_uniform=(self.query_len_support != QueryLenSupport.VARLEN),
-                    treat_short_extends_as_decodes=not self.use_pcp,
-                )
+        num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
+            split_decodes_and_prefills(
+                common_attn_metadata,
+                decode_threshold=self.reorder_batch_threshold,
+                require_uniform=(self.query_len_support != QueryLenSupport.VARLEN),
+                treat_short_extends_as_decodes=not self.use_pcp,
            )
+        )

        assert num_decodes + num_prefills == num_reqs
        assert num_decode_tokens + num_prefill_tokens == num_tokens
@@ -2092,7 +2049,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
            num_decodes=num_decodes,
            num_decode_tokens=num_decode_tokens,
            num_prefills=num_prefills,
-            causal=not non_causal_decode,
            prefill=prefill_metadata,
            decode=decode_metadata,
        )
@@ -841,7 +841,6 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
            MoEActivation.SILU,
            MoEActivation.GELU,
            MoEActivation.GELU_TANH,
-            MoEActivation.SITU,
            MoEActivation.SWIGLUOAI,
            MoEActivation.SILU_NO_MUL,
            MoEActivation.GELU_NO_MUL,
@@ -8,9 +8,7 @@ from collections.abc import Callable
 import torch

 import vllm._custom_ops as ops
-import vllm.envs as envs
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
-from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.activation import (
    MoEActivation,
    apply_moe_activation,
@@ -55,68 +53,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 from vllm.platforms import current_platform
 from vllm.scalar_type import ScalarType, scalar_types

-logger = init_logger(__name__)
-
-_MXFP8_BLOCK_SIZE = 32
-
-
-def _mxfp8_qdq_for_marlin_input(x: torch.Tensor) -> torch.Tensor:
-    """Simulate MXFP8 activations, then feed BF16 back to W4A16 Marlin.
-
-    This is intentionally a debug/reference path. The Marlin GEMM still receives
-    BF16 activations, but those activations carry the precision loss from an
-    MXFP8 quantize-dequantize step.
-    """
-    if x.dtype != torch.bfloat16:
-        raise ValueError(
-            "VLLM_MARLIN_MXFP8_INPUT_QDQ expects BF16 Marlin inputs, "
-            f"but got {x.dtype}."
-        )
-
-    original_cols = x.size(-1)
-    pad_cols = (-original_cols) % _MXFP8_BLOCK_SIZE
-
-    qdq_input = x.contiguous()
-    if pad_cols:
-        padded_shape = (*qdq_input.shape[:-1], original_cols + pad_cols)
-        padded = torch.zeros(
-            padded_shape,
-            dtype=qdq_input.dtype,
-            device=qdq_input.device,
-        )
-        padded[..., :original_cols] = qdq_input
-        qdq_input = padded
-
-    qdq_shape = qdq_input.shape
-    num_blocks = qdq_input.size(-1) // _MXFP8_BLOCK_SIZE
-    qdq_blocks = qdq_input.float().view(
-        *qdq_shape[:-1],
-        num_blocks,
-        _MXFP8_BLOCK_SIZE,
-    )
-
-    amax = qdq_blocks.abs().amax(dim=-1)
-    amax = amax.clamp(min=torch.finfo(torch.float32).tiny)
-    fp8_max = torch.finfo(torch.float8_e4m3fn).max
-    scale_exp = torch.ceil(torch.log2(amax / fp8_max)) + 127.0
-    scale_exp = scale_exp.clamp(0, 254).to(torch.uint8)
-    descale = torch.exp2(scale_exp.float() - 127.0)
-
-    qdq_fp8 = (qdq_blocks / descale.unsqueeze(-1)).view(qdq_shape)
-    qdq_fp8 = qdq_fp8.to(torch.float8_e4m3fn)
-    qdq_output = qdq_fp8.float().view(
-        *qdq_shape[:-1],
-        num_blocks,
-        _MXFP8_BLOCK_SIZE,
-    )
-    qdq_output = qdq_output * descale.unsqueeze(-1)
-    qdq_output = qdq_output.view(qdq_shape).to(torch.bfloat16)
-
-    if pad_cols:
-        qdq_output = qdq_output[..., :original_cols]
-
-    return qdq_output.contiguous()
-

 def _fused_marlin_moe(
    hidden_states: torch.Tensor,
@@ -167,19 +103,6 @@ def _fused_marlin_moe(
    if workspace is None:
        workspace = marlin_make_workspace_new(hidden_states.device, 4)

-    simulate_mxfp8_input_qdq = envs.VLLM_MARLIN_MXFP8_INPUT_QDQ
-    if simulate_mxfp8_input_qdq:
-        if input_dtype is not None:
-            raise ValueError(
-                "VLLM_MARLIN_MXFP8_INPUT_QDQ simulates W4A8 on the W4A16 "
-                "Marlin path; do not combine it with VLLM_MARLIN_INPUT_DTYPE."
-            )
-        logger.warning_once(
-            "Using W4A16 Marlin with MXFP8 activation QDQ simulation. "
-            "Both Marlin GEMM inputs are quantized to MXFP8 and dequantized "
-            "back to BF16 before GEMM. This is for accuracy comparison only."
-        )
-
    if intermediate_cache13 is None:
        intermediate_cache13 = torch.empty(
            (M * num_topk * max(w13_num_shards * N, K),),
@@ -210,8 +133,6 @@ def _fused_marlin_moe(
            a_scales1 = a_scales1 * input_global_scale1
    elif input_dtype == torch.float8_e4m3fn:
        gate_up_input, a_scales1 = marlin_quant_input(hidden_states, input_dtype)
-    elif simulate_mxfp8_input_qdq:
-        gate_up_input = _mxfp8_qdq_for_marlin_input(hidden_states)

    intermediate_cache1 = ops.moe_wna16_marlin_gemm(
        gate_up_input,
@@ -270,8 +191,6 @@ def _fused_marlin_moe(
        intermediate_cache2, a_scales2 = marlin_quant_input(
            intermediate_cache2, input_dtype
        )
-    elif simulate_mxfp8_input_qdq:
-        intermediate_cache2 = _mxfp8_qdq_for_marlin_input(intermediate_cache2)

    output = ops.moe_wna16_marlin_gemm(
        intermediate_cache2,
@@ -752,10 +671,10 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):

    @staticmethod
    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
-        # One-sided FI-NVL all2all pairs with MarlinExperts fine (the
-        # compressed-tensors MXFP4 path runs this exact combo); only the
-        # two-sided kernels are unsupported here.
-        return not moe_parallel_config.use_fi_nvl_two_sided_kernels
+        return not (
+            moe_parallel_config.use_fi_nvl_two_sided_kernels
+            or moe_parallel_config.use_fi_nvl_one_sided_kernels
+        )

    @property
    def quant_type_id(self) -> int:
@@ -1149,8 +1068,11 @@ class BatchedMarlinExperts(MarlinExpertsBase):
                return

            num_experts, max_num_tokens = hidden_states.shape[:2]
-            beta = 1.0 if activation_situ_beta is None else activation_situ_beta
+            beta = activation_situ_beta
            linear_beta = activation_situ_linear_beta
+            assert beta is not None, (
+                "SITU requires activation_situ_beta from FusedMoEConfig"
+            )
            torch.ops._C.masked_situ_and_mul(
                act_output.view(num_experts, max_num_tokens, -1),
                act_input.view(num_experts, max_num_tokens, -1),
@@ -5,7 +5,6 @@ from functools import lru_cache

 import torch

-import vllm.envs as envs
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
@@ -263,8 +262,6 @@ def rocm_aiter_fused_experts(
    elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
        activation_method = rocm_aiter_ops.get_aiter_activation_type("swiglu")
        activation_interleave = False
-    elif activation == MoEActivation.SITU:
-        activation_method = rocm_aiter_ops.get_aiter_activation_type("situ")
    else:
        raise ValueError(f"Unsupported activation: {activation}")

@@ -359,12 +356,10 @@ def rocm_aiter_fused_experts(
        # https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1099
        # TODO: Revisit this once we bump AITER to 0.1.15 with padding fixes
        # for CK/FlyDSL MoE GEMM e.g. https://github.com/ROCm/aiter/pull/3401
-        # SITU's A16W4 FlyDSL kernel pads per gate/up half; pass through unrounded.
-        if activation != MoEActivation.SITU:
-            hidden_pad = hidden_pad // 128 * 128
-            intermediate_pad = (
-                intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1)
-            )
+        hidden_pad = hidden_pad // 128 * 128
+        intermediate_pad = (
+            intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1)
+        )

        # https://github.com/ROCm/aiter/pull/3123 specialized the AITER stage1 GEMMs
        # for interleaved vs separated gate and up weights.
@@ -375,15 +370,7 @@ def rocm_aiter_fused_experts(
        from aiter.ops.flydsl.moe_common import GateMode

        gate_mode = ""
-        if activation == MoEActivation.SITU:
-            # a8w4 (AITER_SITUV2_A8W4=1) uses the gate/up-interleaved (_gui_)
-            # fp8 flydsl kernels; default a16w4 SiTU stays separated.
-            gate_mode = (
-                GateMode.INTERLEAVE.value
-                if envs.AITER_SITUV2_A8W4
-                else GateMode.SEPARATED.value
-            )
-        elif quant_config.use_mxfp4_w4a16:
+        if quant_config.use_mxfp4_w4a16:
            gate_mode = GateMode.INTERLEAVE.value
        elif activation_interleave is not None:
            gate_mode = (
@@ -414,8 +401,6 @@ def rocm_aiter_fused_experts(
            bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,
            bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,
            moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
-            beta=moe_config.activation_situ_beta,
-            linear_beta=moe_config.activation_situ_linear_beta,
        )


@@ -140,7 +140,6 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
            MoEActivation.SILU,
            MoEActivation.GELU,
            MoEActivation.GELU_TANH,
-            MoEActivation.SITU,
            MoEActivation.SWIGLUOAI,
            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
            MoEActivation.SWIGLUSTEP,
@@ -315,75 +315,6 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
        output = (M, self.hidden_dim_unpadded)
        return (workspace1, workspace2, output)

-    def _max_supported_tokens(self, top_k: int, global_num_experts: int) -> int:
-        """Max tokens per kernel call before the batched-GEMM grid overflows.
-
-        The TRTLLM-Gen batched GEMM launches a static grid whose batch (Y)
-        dimension is ``getMaxNumCtasInBatchDim(num_tokens, top_k, num_experts,
-        tileTokensDim)`` and must stay <= 65535. Solving that for num_tokens
-        with the smallest tile the kernel may pick (tileTokensDim=8, the runner
-        default) gives a bound that is safe regardless of the tactic selected.
-        Without it, large batches (e.g. Kimi-K3 top_k=16, EP16 profiling with
-        131072 gathered tokens) overflow the grid and the GEMM launch fails.
-        """
-        MAX_GRID_Y = 65535
-        MIN_TILE_TOKENS_DIM = 8
-        max_tokens = (MAX_GRID_Y - global_num_experts) * MIN_TILE_TOKENS_DIM // top_k
-        return max(1, min(300000, max_tokens))
-
-    def _invoke_kernel(
-        self,
-        output: torch.Tensor,
-        x_quant: torch.Tensor,
-        x_scale: torch.Tensor | None,
-        topk_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        w1: torch.Tensor,
-        w2: torch.Tensor,
-        activation: MoEActivation,
-        global_num_experts: int,
-        local_num_experts: int,
-        local_expert_offset: int,
-        topk: int,
-    ) -> None:
-        from flashinfer import trtllm_fp4_block_scale_routed_moe
-
-        packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
-        trtllm_fp4_block_scale_routed_moe(
-            topk_ids=packed_tensor,
-            routing_bias=None,
-            hidden_states=x_quant,
-            hidden_states_scale=x_scale,
-            gemm1_weights=w1,
-            gemm1_weights_scale=self.w1_scale,
-            gemm1_bias=self.w1_bias,
-            gemm1_alpha=self.gemm1_alpha,
-            gemm1_beta=self.gemm1_beta,
-            gemm1_clamp_limit=self.gemm1_clamp_limit,
-            gemm2_weights=w2,
-            gemm2_weights_scale=self.w2_scale,
-            gemm2_bias=self.w2_bias,
-            output1_scale_scalar=None,
-            output1_scale_gate_scalar=None,
-            output2_scale_scalar=None,
-            num_experts=global_num_experts,
-            top_k=topk,
-            n_group=None,
-            topk_group=None,
-            intermediate_size=self.intermediate_size_per_partition,
-            local_expert_offset=local_expert_offset,
-            local_num_experts=local_num_experts,
-            routed_scaling_factor=None,
-            # Modular kernel receives pre-routed tokens, so routing is already
-            # done. Use Renormalize as a safe default the TRTLLM kernel supports.
-            routing_method_type=RoutingMethodType.Renormalize,
-            do_finalize=True,
-            enable_pdl=True,
-            activation_type=self._flashinfer_activation_type(activation),
-            output=output,
-            tune_max_num_tokens=max(self.max_capture_size, 1),
-        )
-
    def apply(
        self,
        output: torch.Tensor,
@@ -404,6 +335,7 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
    ):
        topk = topk_ids.size(-1)
        local_num_experts = w1.size(0)
+        intermediate_size = self.intermediate_size_per_partition
        local_expert_offset = self.moe_config.ep_rank * local_num_experts

        if a1q_scale is not None:
@@ -414,27 +346,49 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
            x_quant = hidden_states
            x_scale = None

+        # Pack topk ids and weights into format expected by the kernel.
+        packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
+
        assert self.w1_scale is not None
        assert self.w2_scale is not None
+        kwargs = {
+            "topk_ids": packed_tensor,
+            "routing_bias": None,
+            "hidden_states": x_quant,
+            "hidden_states_scale": x_scale,
+            "gemm1_weights": w1,
+            "gemm1_weights_scale": self.w1_scale,
+            "gemm1_bias": self.w1_bias,
+            "gemm1_alpha": self.gemm1_alpha,
+            "gemm1_beta": self.gemm1_beta,
+            "gemm1_clamp_limit": self.gemm1_clamp_limit,
+            "gemm2_weights": w2,
+            "gemm2_weights_scale": self.w2_scale,
+            "gemm2_bias": self.w2_bias,
+            "output1_scale_scalar": None,
+            "output1_scale_gate_scalar": None,
+            "output2_scale_scalar": None,
+            "num_experts": global_num_experts,
+            "top_k": topk,
+            "n_group": None,
+            "topk_group": None,
+            "intermediate_size": intermediate_size,
+            "local_expert_offset": local_expert_offset,
+            "local_num_experts": local_num_experts,
+            "routed_scaling_factor": None,
+            # Modular kernel receives pre-routed tokens, so routing
+            # is already done. Use Renormalize as a safe default that
+            # the TRTLLM C++ kernel supports.
+            "routing_method_type": RoutingMethodType.Renormalize,
+            "do_finalize": True,
+            "enable_pdl": True,
+            "activation_type": self._flashinfer_activation_type(activation),
+            "output": output,
+            "tune_max_num_tokens": max(self.max_capture_size, 1),
+        }

-        # Chunk tokens so the batched-GEMM grid stays within CUDA limits.
-        M = x_quant.size(0)
-        chunk_size = self._max_supported_tokens(topk, global_num_experts)
-        for start in range(0, M, chunk_size):
-            end = min(start + chunk_size, M)
-            self._invoke_kernel(
-                output[start:end],
-                x_quant[start:end],
-                None if x_scale is None else x_scale[start:end],
-                topk_ids[start:end],
-                topk_weights[start:end],
-                w1,
-                w2,
-                activation,
-                global_num_experts,
-                local_num_experts,
-                local_expert_offset,
-                topk,
-            )
+        from flashinfer import trtllm_fp4_block_scale_routed_moe
+
+        trtllm_fp4_block_scale_routed_moe(**kwargs)

        return output
@@ -648,7 +648,7 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(
        else:
            hidden_size = round_up(hidden_size, 256)
    elif backend in TRTLLM_BACKENDS:
-        intermediate_size = round_up(intermediate_size, 128)
+        intermediate_size = round_up(intermediate_size, 256)
        hidden_size = round_up(hidden_size, 256)
    elif backend in (
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
@@ -6,6 +6,11 @@ import vllm.envs as envs
 from vllm.config import get_current_vllm_config
 from vllm.distributed import tensor_model_parallel_all_reduce
 from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
+    _AR_RESIDUAL_RMS_NORM,
+    _can_use_flashinfer,
+    flashinfer_trtllm_fused_allreduce_norm,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.utils.torch_utils import aux_stream, current_stream

@@ -219,12 +224,6 @@ class LatentMoERunner(MoERunner):
        norm: RMSNorm,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer."""
-        from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
-            _AR_RESIDUAL_RMS_NORM,
-            _can_use_flashinfer,
-            flashinfer_trtllm_fused_allreduce_norm,
-        )
-
        if self.moe_config.tp_size == 1:
            return norm(hidden_states)

@@ -7,7 +7,6 @@ import torch.nn.functional as F

 from vllm.config import get_current_vllm_config
 from vllm.distributed import (
-    get_tensor_model_parallel_world_size,
    tensor_model_parallel_all_gather,
    tensor_model_parallel_gather,
 )
@@ -144,9 +143,8 @@ class LogitsProcessor(PluggableLayer):
        # Get the logits for the next tokens.
        logits = self._apply_head(lm_head, hidden_states, embedding_bias)

-        # Gather logits only when the head itself is TP-sharded. A replicated
-        # head already produces the complete vocabulary on every rank.
-        if not getattr(lm_head, "replicated", False):
+        # Gather logits for TP
+        if lm_head.tp_size > 1:
            logits = self._gather_logits(logits)

        # Remove paddings in vocab (if any).
@@ -171,11 +169,7 @@ class LogitsProcessor(PluggableLayer):
                "The local argmax reduction optimization is not supported for "
                "non-positive logit scaling factors."
            )
-        tp_size = (
-            1
-            if getattr(lm_head, "replicated", False)
-            else get_tensor_model_parallel_world_size()
-        )
+        tp_size = lm_head.tp_size

        logits = self._apply_head(lm_head, hidden_states, embedding_bias)
        if self.soft_cap is not None:
@@ -1,33 +1,32 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-from collections.abc import Callable
-
 import torch
 from einops import rearrange
 from torch import nn
-from torch.nn.parameter import Parameter

-from vllm.compilation.breakable_cudagraph import eager_break_during_capture
-from vllm.config import VllmConfig
-from vllm.distributed import divide, get_tensor_model_parallel_rank
-from vllm.forward_context import get_forward_context
+from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.distributed import (
+    divide,
+)
+from vllm.forward_context import ForwardContext, get_forward_context
+from vllm.logger import init_logger
 from vllm.model_executor.custom_op import PluggableLayer
 from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-    sharded_weight_loader,
-)
-from vllm.model_executor.parameter import BasevLLMParameter
+from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
 from vllm.model_executor.utils import set_weight_attrs
-from vllm.platforms import current_platform
-from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
+from vllm.third_party.flash_linear_attention.ops.kda import (
+    FusedRMSNormGated,
+    chunk_kda_with_fused_gate,
+    fused_kda_gate,
+    fused_recurrent_kda,
+)
 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
+from vllm.utils.torch_utils import direct_register_custom_op
 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata

 from ...linear import (
    ColumnParallelLinear,
-    MergedColumnParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
 )
@@ -37,118 +36,49 @@ from ..mamba_utils import (
    is_conv_state_dim_first,
 )
 from ..ops.causal_conv1d import causal_conv1d_fn, causal_conv1d_update
-from ..ops.gather_initial_states import gather_initial_states

-# Empirical lower bound for the KDA gate to avoid numerical underflow.
-_KDA_GATE_LOGBOUND_MIN = -5.0
+logger = init_logger(__name__)


-def a_log_weight_loader(
-    shard_axis: int,
-) -> Callable[[torch.Tensor, torch.Tensor], None]:
-    """Load KDA A_log stored as either old 4D or current 1D weights."""
-
-    def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
-        tp_rank = get_tensor_model_parallel_rank()
-        shard_size = param.data.shape[shard_axis]
-        start_idx = tp_rank * shard_size
-
-        if loaded_weight.dim() == 4:
-            assert loaded_weight.shape[:2] == (1, 1), (
-                f"Expected old A_log shape (1, 1, H, 1), got {loaded_weight.shape}"
-            )
-            assert loaded_weight.shape[-1] == 1, (
-                f"Expected old A_log last dim to be 1, got {loaded_weight.shape}"
-            )
-            loaded_weight = loaded_weight.view(loaded_weight.shape[2])
-
-        loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size)
-        return default_weight_loader(param, loaded_weight)
-
-    return loader
+def kda_attention(
+    q_proj_states: torch.Tensor,
+    k_proj_states: torch.Tensor,
+    v_proj_states: torch.Tensor,
+    g1: torch.Tensor,
+    beta: torch.Tensor,
+    core_attn_out: torch.Tensor,
+    layer_name: str,
+) -> None:
+    forward_context: ForwardContext = get_forward_context()
+    self = forward_context.no_compile_layers[layer_name]
+    self._forward(
+        q_proj_states=q_proj_states,
+        k_proj_states=k_proj_states,
+        v_proj_states=v_proj_states,
+        g1=g1,
+        beta=beta,
+        core_attn_out=core_attn_out,
+    )


-def _make_fused_conv1d_weight_loader(
-    dims: list[int],
-    tp_size: int,
-    tp_rank: int,
-) -> Callable[..., None]:
-    sharded_dims = [dim // tp_size for dim in dims]
-
-    def weight_loader(
-        param: torch.Tensor,
-        loaded_weight: torch.Tensor,
-        loaded_shard_id: int,
-    ) -> None:
-        if loaded_weight.dim() == 2:
-            loaded_weight = loaded_weight.unsqueeze(1)
-        shard_size = sharded_dims[loaded_shard_id]
-        source_start = tp_rank * shard_size
-        target_start = sum(sharded_dims[:loaded_shard_id])
-        loaded_shard = loaded_weight[source_start : source_start + shard_size]
-        param.data[target_start : target_start + shard_size].copy_(loaded_shard)
-
-    return weight_loader
+def kda_attention_fake(
+    q_proj_states: torch.Tensor,
+    k_proj_states: torch.Tensor,
+    v_proj_states: torch.Tensor,
+    g1: torch.Tensor,
+    beta: torch.Tensor,
+    core_attn_out: torch.Tensor,
+    layer_name: str,
+) -> None:
+    return


-class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear):
-    """Merged projection with one output replicated across TP ranks.
-
-    The replicated shard is represented as ``size * tp_size`` so the merged
-    parameter reserves ``size`` local rows on every rank. Loading that shard
-    from rank zero then gives every rank the complete checkpoint weight.
-    """
-
-    def __init__(
-        self,
-        input_size: int,
-        output_sizes: list[int],
-        replicated_shard_id: int,
-        tp_size: int,
-        **kwargs,
-    ) -> None:
-        self.replicated_shard_id = replicated_shard_id
-        output_sizes = output_sizes.copy()
-        output_sizes[replicated_shard_id] *= tp_size
-        super().__init__(input_size, output_sizes, **kwargs)
-
-    def weight_loader(
-        self,
-        param: Parameter,
-        loaded_weight: torch.Tensor,
-        loaded_shard_id: tuple[int, ...] | int | None = None,
-    ) -> None:
-        tp_rank = self.tp_rank
-        param_tp_rank = getattr(param, "tp_rank", None)
-        if loaded_shard_id == self.replicated_shard_id:
-            self.tp_rank = 0
-            if param_tp_rank is not None:
-                param.tp_rank = 0
-        try:
-            super().weight_loader(param, loaded_weight, loaded_shard_id)
-        finally:
-            self.tp_rank = tp_rank
-            if param_tp_rank is not None:
-                param.tp_rank = param_tp_rank
-
-    def weight_loader_v2(
-        self,
-        param: BasevLLMParameter,
-        loaded_weight: torch.Tensor,
-        loaded_shard_id: tuple[int, ...] | int | None = None,
-    ) -> None:
-        tp_rank = self.tp_rank
-        param_tp_rank = getattr(param, "tp_rank", None)
-        if loaded_shard_id == self.replicated_shard_id:
-            self.tp_rank = 0
-            if param_tp_rank is not None:
-                param.tp_rank = 0
-        try:
-            super().weight_loader_v2(param, loaded_weight, loaded_shard_id)
-        finally:
-            self.tp_rank = tp_rank
-            if param_tp_rank is not None:
-                param.tp_rank = param_tp_rank
+direct_register_custom_op(
+    op_name="kda_attention",
+    op_func=kda_attention,
+    mutates_args=["core_attn_out"],
+    fake_impl=kda_attention_fake,
+)


@PluggableLayer.register("kimi_gated_delta_net_attention")
@@ -166,11 +96,7 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
        self,
    ) -> tuple[tuple[int, ...], tuple[int, ...]]:
        return MambaStateShapeCalculator.kda_state_shape(
-            self.tp_size,
-            self.num_heads,
-            self.head_dim,
-            conv_kernel_size=self.conv_size,
-            num_spec=self.num_spec,
+            self.tp_size, self.num_heads, self.head_dim, conv_kernel_size=self.conv_size
        )

    def __init__(
@@ -188,142 +114,122 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
        assert self.num_heads % self.tp_size == 0
        self.local_num_heads = divide(self.num_heads, self.tp_size)

-        self.projection_size = self.head_dim * self.num_heads
-        self.local_projection_size = divide(self.projection_size, self.tp_size)
+        projection_size = self.head_dim * self.num_heads
        self.conv_size = kda_config["short_conv_kernel_size"]
-        self.use_full_rank_gate = kda_config.get("use_full_rank_gate", False)

-        if self.use_full_rank_gate:
-            # Keep f_a before the narrow beta shard, then pad each TP-local row
-            # to select the aligned BF16 GEMM path. The padding also avoids an
-            # Inductor correctness issue seen with the row-strided G view.
-            qkvg_output_sizes = [self.projection_size] * 4
-            in_proj_output_sizes = qkvg_output_sizes + [
-                self.head_dim,
-                self.num_heads,
-            ]
-            local_output_size = (
-                4 * self.local_projection_size + self.head_dim + self.local_num_heads
-            )
-            self.in_proj_padding = -local_output_size % 16
-            if self.in_proj_padding:
-                in_proj_output_sizes.append(self.in_proj_padding * self.tp_size)
-        else:
-            in_proj_output_sizes = [self.projection_size] * 3 + [
-                self.num_heads,
-                self.head_dim,
-            ]
-            self.in_proj_padding = 0
-        self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear(
+        self.q_proj = ColumnParallelLinear(
            self.hidden_size,
-            in_proj_output_sizes,
-            replicated_shard_id=4,
-            tp_size=self.tp_size,
+            projection_size,
            bias=False,
            quant_config=self.quant_config,
-            prefix=f"{prefix}.in_proj_qkvgfab",
+            prefix=f"{prefix}.q_proj",
+        )
+        self.k_proj = ColumnParallelLinear(
+            self.hidden_size,
+            projection_size,
+            bias=False,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.k_proj",
+        )
+        self.v_proj = ColumnParallelLinear(
+            self.hidden_size,
+            projection_size,
+            bias=False,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.v_proj",
+        )
+
+        self.f_a_proj = ReplicatedLinear(
+            self.hidden_size,
+            self.head_dim,
+            bias=False,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.f_a_proj",
        )
-        if self.in_proj_padding:
-            self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_()

        self.f_b_proj = ColumnParallelLinear(
            self.head_dim,
-            self.projection_size,
+            projection_size,
            bias=False,
            quant_config=self.quant_config,
            prefix=f"{prefix}.f_b_proj",
        )
        self.dt_bias = nn.Parameter(
-            torch.empty(self.local_projection_size, dtype=torch.float32)
+            torch.empty(divide(projection_size, self.tp_size), dtype=torch.float32)
        )

        set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})

-        # One packed parameter and cache let decode run a single conv update.
-        # Prefill slices them back into Q/K/V to obtain dense outputs cheaply.
-        self.conv1d = ColumnParallelLinear(
+        self.b_proj = ColumnParallelLinear(
+            self.hidden_size,
+            self.num_heads,
+            bias=False,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.b_proj",
+        )
+
+        self.q_conv1d = ColumnParallelLinear(
            input_size=self.conv_size,
-            output_size=3 * self.projection_size,
+            output_size=projection_size,
            bias=False,
            params_dtype=torch.float32,
-            prefix=f"{prefix}.conv1d",
+            prefix=f"{prefix}.q_conv1d",
        )
-        self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
-        delattr(self.conv1d.weight, "weight_loader")
-        set_weight_attrs(
-            self.conv1d.weight,
-            {
-                "weight_loader": _make_fused_conv1d_weight_loader(
-                    [self.projection_size] * 3,
-                    self.tp_size,
-                    self.tp_rank,
-                )
-            },
+        self.k_conv1d = ColumnParallelLinear(
+            input_size=self.conv_size,
+            output_size=projection_size,
+            bias=False,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.k_conv1d",
        )
+        self.v_conv1d = ColumnParallelLinear(
+            input_size=self.conv_size,
+            output_size=projection_size,
+            bias=False,
+            params_dtype=torch.float32,
+            prefix=f"{prefix}.v_conv1d",
+        )
+        # unsqueeze to fit conv1d weights shape into the linear weights shape.
+        # Can't do this in `weight_loader` since it already exists in
+        # `ColumnParallelLinear` and `set_weight_attrs`
+        # doesn't allow to override it
+        self.q_conv1d.weight.data = self.q_conv1d.weight.data.unsqueeze(1)
+        self.k_conv1d.weight.data = self.k_conv1d.weight.data.unsqueeze(1)
+        self.v_conv1d.weight.data = self.v_conv1d.weight.data.unsqueeze(1)

        self.A_log = nn.Parameter(
-            torch.empty(self.local_num_heads, dtype=torch.float32)
+            torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32)
        )
-        set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)})
+        set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(2)})

-        self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None)
-        if self.gate_lower_bound is not None:
-            assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, (
-                "KDA gate lower bound must be in "
-                f"[{_KDA_GATE_LOGBOUND_MIN}, 0). "
-                f"Got {self.gate_lower_bound}."
-            )
-        self.use_safe_gate = self.gate_lower_bound is not None
-        additional_config = vllm_config.additional_config
-        backend = (
-            additional_config.get("kda_prefill_backend", "auto")
-            if isinstance(additional_config, dict)
-            else "auto"
+        self.g_a_proj = ReplicatedLinear(
+            self.hidden_size,
+            self.head_dim,
+            bias=False,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.g_a_proj",
        )
-        backend = "triton" if backend == "auto" else backend
-        assert backend == "triton", (
-            "The shared Kimi GDN layer only supports the Triton KDA "
-            f"prefill backend, got {backend!r}."
+        self.g_b_proj = ColumnParallelLinear(
+            self.head_dim,
+            projection_size,
+            bias=False,
+            quant_config=self.quant_config,
+            prefix=f"{prefix}.g_b_proj",
        )
-        if not self.use_full_rank_gate:
-            self.g_a_proj = ReplicatedLinear(
-                self.hidden_size,
-                self.head_dim,
-                bias=False,
-                quant_config=self.quant_config,
-                prefix=f"{prefix}.g_a_proj",
-            )
-            self.g_b_proj = ColumnParallelLinear(
-                self.head_dim,
-                self.projection_size,
-                bias=False,
-                quant_config=self.quant_config,
-                prefix=f"{prefix}.g_b_proj",
-            )
        self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
        self.o_proj = RowParallelLinear(
-            self.projection_size,
+            projection_size,
            self.hidden_size,
            bias=False,
            quant_config=self.quant_config,
            prefix=f"{prefix}.o_proj",
        )

-        compilation_config = vllm_config.compilation_config
+        compilation_config = get_current_vllm_config().compilation_config
        if prefix in compilation_config.static_forward_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        compilation_config.static_forward_context[prefix] = self

-    def rearrange_mixed_qkv(
-        self, mixed_qkv: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        seq_len = mixed_qkv.shape[0]
-        qkv = mixed_qkv.view(seq_len, 3, self.local_num_heads, self.head_dim)
-        # Materialize all three row-strided inputs with one token-major to
-        # QKV-major permutation. Each unbound tensor is then contiguous.
-        qkv = qkv.permute(1, 0, 2, 3).contiguous().unsqueeze(1)
-        return qkv.unbind(0)
-
    def forward(
        self,
        hidden_states: torch.Tensor,
@@ -331,57 +237,42 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
        output: torch.Tensor,
    ) -> None:
        num_tokens = hidden_states.size(0)
-        projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
-        if self.use_full_rank_gate:
-            split_sizes = [
-                3 * self.local_projection_size,
-                self.local_projection_size,
-                self.head_dim,
-                self.local_num_heads,
-            ]
-            if self.in_proj_padding:
-                split_sizes.append(self.in_proj_padding)
-            projected = projected_qkvgfab.split(split_sizes, dim=-1)
-            mixed_qkv, g_proj_states, f_a, beta = projected[:4]
-        else:
-            mixed_qkv, beta, f_a = projected_qkvgfab.split(
-                [
-                    3 * self.local_projection_size,
-                    self.local_num_heads,
-                    self.head_dim,
-                ],
-                dim=-1,
-            )
-            g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
+        q = self.q_proj(hidden_states)[0]
+        k = self.k_proj(hidden_states)[0]
+        v = self.v_proj(hidden_states)[0]

-        g1 = self.f_b_proj(f_a)[0]
+        beta = self.b_proj(hidden_states)[0].float().sigmoid()
+        g1 = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
        beta = beta.unsqueeze(0)
        g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim)

+        g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
        g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)

-        core_attn_out = torch.empty(
+        core_attn_out = torch.zeros(
            (1, num_tokens, self.local_num_heads, self.head_dim),
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )
-
-        self._forward(
-            mixed_qkv=mixed_qkv,
-            g1=g1,
-            g2=g2,
-            beta=beta,
-            core_attn_out=core_attn_out,
+        torch.ops.vllm.kda_attention(
+            q,
+            k,
+            v,
+            g1,
+            beta,
+            core_attn_out,
+            self.prefix,
        )
+        core_attn_out = self.o_norm(core_attn_out, g2)
        core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
        output[:] = self.o_proj(core_attn_out)[0]

-    @eager_break_during_capture
    def _forward(
        self,
-        mixed_qkv: torch.Tensor,
+        q_proj_states: torch.Tensor,
+        k_proj_states: torch.Tensor,
+        v_proj_states: torch.Tensor,
        g1: torch.Tensor,
-        g2: torch.Tensor,
        beta: torch.Tensor,
        core_attn_out: torch.Tensor,
    ) -> None:
@@ -389,247 +280,165 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
        attn_metadata_raw = forward_context.attn_metadata

        if attn_metadata_raw is None:
+            #     # V1 profile run
            return

-        # Vendor-specific KDA kernels: AMD/ROCm and NVIDIA keep their own copies
-        # under kimi_k3/{amd,nvidia}/ops so each can diverge independently.
-        if current_platform.is_rocm():
-            from vllm.models.kimi_k3.amd.ops.third_party.kda import (
-                chunk_kda_with_fused_gate,
-                fused_recurrent_kda,
-                fused_recurrent_kda_packed_decode,
-            )
-        else:
-            from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
-                chunk_kda_with_fused_gate,
-                fused_recurrent_kda,
-                fused_recurrent_kda_packed_decode,
-            )
-
        assert isinstance(attn_metadata_raw, dict)
        attn_metadata_narrowed = attn_metadata_raw[self.prefix]
        assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata)
-        m = attn_metadata_narrowed
-        has_initial_state = m.has_initial_state
-        non_spec_query_start_loc = m.non_spec_query_start_loc
-        non_spec_state_indices_tensor = m.non_spec_state_indices_tensor
-        spec_sequence_masks = m.spec_sequence_masks
-        spec_token_indx = m.spec_token_indx
-        non_spec_token_indx = m.non_spec_token_indx
-        spec_state_indices_tensor = m.spec_state_indices_tensor
-        spec_query_start_loc = m.spec_query_start_loc
-        num_accepted_tokens = m.num_accepted_tokens
-        num_actual_tokens = m.num_actual_tokens
-        mixed_qkv = mixed_qkv[:num_actual_tokens]
+        has_initial_state = attn_metadata_narrowed.has_initial_state
+        non_spec_query_start_loc = attn_metadata_narrowed.non_spec_query_start_loc
+        non_spec_state_indices_tensor = (
+            attn_metadata_narrowed.non_spec_state_indices_tensor
+        )  # noqa: E501
+        num_actual_tokens = attn_metadata_narrowed.num_actual_tokens
+        constant_caches = self.kv_cache
+
+        q_proj_states = q_proj_states[:num_actual_tokens]
+        k_proj_states = k_proj_states[:num_actual_tokens]
+        v_proj_states = v_proj_states[:num_actual_tokens]
        g1 = g1[:, :num_actual_tokens]
        beta = beta[:, :num_actual_tokens]

-        constant_caches = self.kv_cache
-
-        conv_state, recurrent_state = constant_caches
+        (conv_state, recurrent_state) = constant_caches
        # conv_state must be (..., dim, width-1) for the conv kernels.
        # DS layout stores it that way directly; SD layout needs a transpose.
        if not is_conv_state_dim_first():
            conv_state = conv_state.transpose(-1, -2)

-        conv_weights = self.conv1d.weight.view(
-            self.conv1d.weight.size(0), self.conv1d.weight.size(2)
-        )
-        q_conv_weight, k_conv_weight, v_conv_weight = conv_weights.split(
-            self.local_projection_size, dim=0
-        )
-        q_conv_state, k_conv_state, v_conv_state = conv_state.split(
-            self.local_projection_size, dim=-2
-        )
+        conv_state_q, conv_state_k, conv_state_v = conv_state.chunk(3, dim=-2)

-        # Split tokens into the multi-query spec-decode part and the remaining
-        # (prefill / plain decode) part.
-        if spec_sequence_masks is not None:
-            if m.num_prefills == 0 and m.num_decodes == 0:
-                mixed_qkv_spec = mixed_qkv
-                g1_spec, beta_spec = g1, beta
-                mixed_qkv_ns = g1_ns = beta_ns = None
-            else:
-                mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
-                g1_spec = g1.index_select(1, spec_token_indx)
-                beta_spec = beta.index_select(1, spec_token_indx)
-                mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx)
-                g1_ns = g1.index_select(1, non_spec_token_indx)
-                beta_ns = beta.index_select(1, non_spec_token_indx)
-        else:
-            mixed_qkv_spec = g1_spec = beta_spec = None
-            mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta
-
-        # ---------- spec-decode multi-query path ----------
-        core_attn_out_spec = None
-        if spec_sequence_masks is not None:
-            assert spec_state_indices_tensor is not None
-            assert spec_query_start_loc is not None
-            spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes]
-            spec_max_query_len = spec_state_indices_tensor.size(-1)
-
-            # Sibling beta and, for full-rank gates, output-gate views remain
-            # live, so write the convolution output separately.
-            spec_conv_out = torch.empty(
-                mixed_qkv_spec.shape,
-                dtype=mixed_qkv_spec.dtype,
-                device=mixed_qkv_spec.device,
-            )
-            mixed_qkv_spec = causal_conv1d_update(
-                mixed_qkv_spec,
-                conv_state,
-                conv_weights,
-                self.conv1d.bias,
+        q_conv_weights = self.q_conv1d.weight.view(
+            self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2)
+        )
+        k_conv_weights = self.k_conv1d.weight.view(
+            self.k_conv1d.weight.size(0), self.k_conv1d.weight.size(2)
+        )
+        v_conv_weights = self.v_conv1d.weight.view(
+            self.v_conv1d.weight.size(0), self.v_conv1d.weight.size(2)
+        )
+        if attn_metadata_narrowed.num_prefills > 0:
+            q_proj_states = q_proj_states.transpose(0, 1)
+            k_proj_states = k_proj_states.transpose(0, 1)
+            v_proj_states = v_proj_states.transpose(0, 1)
+            q = causal_conv1d_fn(
+                q_proj_states,
+                q_conv_weights,
+                self.q_conv1d.bias,
                activation="silu",
-                conv_state_indices=spec_conv_indices,
-                num_accepted_tokens=num_accepted_tokens,
-                query_start_loc=spec_query_start_loc,
-                max_query_len=spec_max_query_len,
-                validate_data=False,
-                out=spec_conv_out,
-            )
-            q_spec, k_spec, v_spec = (
-                rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
-                for x in mixed_qkv_spec.split(self.local_projection_size, dim=-1)
-            )
-            spec_cu_seqlens = spec_query_start_loc[: m.num_spec_decodes + 1]
-            # Spec-only batches write directly into core_attn_out.
-            spec_out = (
-                core_attn_out[:, : q_spec.shape[1]]
-                if m.num_prefills == 0 and m.num_decodes == 0
-                else None
-            )
-            core_attn_out_spec, _ = fused_recurrent_kda(
-                q=q_spec,
-                k=k_spec,
-                v=v_spec,
-                raw_g=g1_spec,
-                raw_beta=beta_spec,
-                A_log=self.A_log,
-                dt_bias=self.dt_bias,
-                lower_bound=self.gate_lower_bound,
-                initial_state=recurrent_state,
-                cu_seqlens=spec_cu_seqlens,
-                ssm_state_indices=spec_state_indices_tensor,
-                num_accepted_tokens=num_accepted_tokens,
-                out=spec_out,
-            )
-
-        # ---------- non-spec path (prefill or plain decode) ----------
-        core_attn_out_non_spec = None
-        if mixed_qkv_ns is not None:
-            assert g1_ns is not None and beta_ns is not None
-            if m.num_prefills > 0:
-                q_ns, k_ns, v_ns = mixed_qkv_ns.split(
-                    self.local_projection_size, dim=-1
-                )
-
-                # Packed prefill conv would require copying V solely to make
-                # it dense for KDA. Separate calls accept the strided inputs
-                # and produce dense Q/K/V without that extra traffic.
-                # TODO: Use packed conv once every KDA prefill backend accepts
-                # row-strided Q/K/V directly.
-                def _prefill_conv(
-                    x: torch.Tensor,
-                    state: torch.Tensor,
-                    weight: torch.Tensor,
-                ) -> torch.Tensor:
-                    return causal_conv1d_fn(
-                        x.transpose(0, 1),
-                        weight,
-                        None,
-                        activation="silu",
-                        conv_states=state,
-                        has_initial_state=has_initial_state,
-                        cache_indices=non_spec_state_indices_tensor,
-                        query_start_loc=non_spec_query_start_loc,
-                        metadata=m,
-                    ).transpose(0, 1)
-
-                q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight)
-                k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight)
-                v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight)
-                q_ns, k_ns, v_ns = (
-                    rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
-                    for x in (q_ns, k_ns, v_ns)
-                )
-
-                assert non_spec_state_indices_tensor is not None
-                assert has_initial_state is not None
-                initial_state = gather_initial_states(
-                    recurrent_state,
-                    non_spec_state_indices_tensor,
-                    has_initial_state,
-                )
-                (
-                    core_attn_out_non_spec,
-                    last_recurrent_state,
-                ) = chunk_kda_with_fused_gate(
-                    q=q_ns,
-                    k=k_ns,
-                    v=v_ns,
-                    raw_g=g1_ns,
-                    raw_beta=beta_ns,
-                    A_log=self.A_log,
-                    g_bias=self.dt_bias,
-                    lower_bound=self.gate_lower_bound,
-                    initial_state=initial_state,
-                    output_final_state=True,
-                    use_qk_l2norm_in_kernel=True,
-                    cu_seqlens=non_spec_query_start_loc,
-                )
-                # Init cache
-                recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state
-
-            else:
-                # pure-decode non-spec batch
-                assert non_spec_state_indices_tensor is not None
-                decode_conv_indices = non_spec_state_indices_tensor[
-                    : mixed_qkv_ns.size(0)
-                ]
-                # Sibling beta and, for full-rank gates, output-gate views
-                # remain live, so write the conv output separately.
-                packed_conv_out = torch.empty(
-                    mixed_qkv_ns.shape,
-                    dtype=mixed_qkv_ns.dtype,
-                    device=mixed_qkv_ns.device,
-                )
-                mixed_qkv_ns = causal_conv1d_update(
-                    mixed_qkv_ns,
-                    conv_state,
-                    conv_weights,
-                    self.conv1d.bias,
-                    activation="silu",
-                    conv_state_indices=decode_conv_indices,
-                    validate_data=True,
-                    out=packed_conv_out,
-                )
-                core_attn_out_non_spec, _ = fused_recurrent_kda_packed_decode(
-                    mixed_qkv=mixed_qkv_ns,
-                    raw_g=g1_ns,
-                    raw_beta=beta_ns,
-                    A_log=self.A_log,
-                    dt_bias=self.dt_bias,
-                    lower_bound=self.gate_lower_bound,
-                    initial_state=recurrent_state,
-                    state_indices=decode_conv_indices,
-                )
-
-        # ---------- merge spec and non-spec outputs ----------
-        if core_attn_out_spec is not None and core_attn_out_non_spec is not None:
-            # Mixed batches require indexed placement in the original order.
-            merged = torch.empty(
-                (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
-                dtype=core_attn_out_spec.dtype,
-                device=core_attn_out_spec.device,
-            )
-            merged.index_copy_(1, spec_token_indx, core_attn_out_spec)
-            merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
-            core_attn_out[0, :num_actual_tokens] = merged[0, :num_actual_tokens]
-        elif core_attn_out_non_spec is not None:
-            core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
-                0, :num_actual_tokens
-            ]
+                conv_states=conv_state_q,
+                has_initial_state=has_initial_state,
+                cache_indices=non_spec_state_indices_tensor,
+                query_start_loc=non_spec_query_start_loc,
+                metadata=attn_metadata_narrowed,
+            ).transpose(0, 1)
+            k = causal_conv1d_fn(
+                k_proj_states,
+                k_conv_weights,
+                self.k_conv1d.bias,
+                activation="silu",
+                conv_states=conv_state_k,
+                has_initial_state=has_initial_state,
+                cache_indices=non_spec_state_indices_tensor,
+                query_start_loc=non_spec_query_start_loc,
+                metadata=attn_metadata_narrowed,
+            ).transpose(0, 1)
+            v = causal_conv1d_fn(
+                v_proj_states,
+                v_conv_weights,
+                self.v_conv1d.bias,
+                activation="silu",
+                conv_states=conv_state_v,
+                has_initial_state=has_initial_state,
+                cache_indices=non_spec_state_indices_tensor,
+                query_start_loc=non_spec_query_start_loc,
+                metadata=attn_metadata_narrowed,
+            ).transpose(0, 1)
        else:
-            assert core_attn_out_spec is not None
-        core_attn_out.copy_(self.o_norm(core_attn_out, g2))
+            assert non_spec_state_indices_tensor is not None
+            decode_conv_indices = non_spec_state_indices_tensor[
+                : attn_metadata_narrowed.num_actual_tokens
+            ]
+            q = causal_conv1d_update(
+                q_proj_states,
+                conv_state_q,
+                q_conv_weights,
+                self.q_conv1d.bias,
+                activation="silu",
+                conv_state_indices=decode_conv_indices,
+                validate_data=True,
+            )
+            k = causal_conv1d_update(
+                k_proj_states,
+                conv_state_k,
+                k_conv_weights,
+                self.k_conv1d.bias,
+                activation="silu",
+                conv_state_indices=decode_conv_indices,
+                validate_data=True,
+            )
+            v = causal_conv1d_update(
+                v_proj_states,
+                conv_state_v,
+                v_conv_weights,
+                self.v_conv1d.bias,
+                activation="silu",
+                conv_state_indices=decode_conv_indices,
+                validate_data=True,
+            )
+
+        q, k, v = map(
+            lambda x: rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim), (q, k, v)
+        )
+
+        if attn_metadata_narrowed.num_prefills > 0:
+            assert non_spec_state_indices_tensor is not None
+            assert has_initial_state is not None
+            zero_idx = non_spec_state_indices_tensor[~has_initial_state]
+            recurrent_state[zero_idx] = 0
+            initial_state = recurrent_state[non_spec_state_indices_tensor].contiguous()
+            (
+                core_attn_out_non_spec,
+                last_recurrent_state,
+            ) = chunk_kda_with_fused_gate(
+                q=q,
+                k=k,
+                v=v,
+                raw_g=g1,
+                beta=beta,
+                A_log=self.A_log,
+                g_bias=self.dt_bias,
+                initial_state=initial_state,
+                output_final_state=True,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=non_spec_query_start_loc,
+            )
+            # Init cache
+            recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state
+        else:
+            assert non_spec_query_start_loc is not None
+            g1 = fused_kda_gate(
+                rearrange(g1, "1 n h d -> n (h d)"),
+                self.A_log,
+                self.head_dim,
+                g_bias=self.dt_bias,
+            ).unsqueeze(0)
+            (
+                core_attn_out_non_spec,
+                last_recurrent_state,
+            ) = fused_recurrent_kda(
+                q=q,
+                k=k,
+                v=v,
+                g=g1,
+                beta=beta,
+                initial_state=recurrent_state,
+                use_qk_l2norm_in_kernel=True,
+                cu_seqlens=non_spec_query_start_loc[
+                    : attn_metadata_narrowed.num_decodes + 1
+                ],
+                ssm_state_indices=non_spec_state_indices_tensor,
+            )
+        core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
+            0, :num_actual_tokens
+        ]
@@ -288,7 +288,7 @@ class MambaStateShapeCalculator:

        conv_dim = proj_size + 2 * proj_k_size
        conv_state_shape = cls._orient_conv_shape(
-            divide(conv_dim, tp_world_size), conv_kernel_size - 1 + num_spec
+            divide(conv_dim, tp_world_size), conv_kernel_size - 1
        )
        recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim)
        return (conv_state_shape, recurrent_state_shape)
@@ -8,12 +8,11 @@
 import numpy as np
 import torch

-from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID


-@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"])
+@triton.jit()
 def _causal_conv1d_fwd_kernel(  # continuous batching
    # Pointers to matrices
    x_ptr,  # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences
@@ -34,7 +33,7 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
    o_ptr,  # (dim, seqlen) - actually pointing to x_ptr
    # Matrix dimensions
    dim: tl.constexpr,
-    num_cache_lines,  # added to support vLLM larger cache lines
+    num_cache_lines: tl.constexpr,  # added to support vLLM larger cache lines
    # Strides
    stride_x_dim: tl.constexpr,  # stride to get to next feature-value,
    stride_x_token: tl.int64,  # stride to get to next token (same feature-index, same sequence-index)
@@ -59,7 +58,6 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
    NP2_STATELEN: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
-    launch_pdl: tl.constexpr,
 ):
    conv_states_ptr = initial_states_ptr
    conv_state_indices_ptr = cache_indices_ptr
@@ -70,9 +68,6 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
        KERNEL_WIDTH - 1
    )  # can be passed via argument if it's not the same as this value

-    if launch_pdl:
-        tl.extra.cuda.gdc_wait()
-
    # one program handles one chunk in a single sequence
    # rather than mixing sequences - to make updating initial_states across sequences efficiently

@@ -84,8 +79,6 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
    idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)

    if idx_seq == pad_slot_id:
-        if launch_pdl:
-            tl.extra.cuda.gdc_launch_dependents()
        return

    sequence_start_index = tl.load(query_start_loc_ptr + idx_seq)
@@ -143,8 +136,6 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
    if HAS_NULL_BLOCK:  # noqa
        if conv_states_input_coord == null_block_id:
            # not processing as this is a null block (padding)
-            if launch_pdl:
-                tl.extra.cuda.gdc_launch_dependents()
            return
    conv_states_base = (
        conv_states_ptr
@@ -417,10 +408,6 @@ def _causal_conv1d_fwd_kernel(  # continuous batching
        w_ptrs = w_base + (3 * stride_w_width)  # [BLOCK_N] tensor
        w_col3 = tl.load(w_ptrs, mask_w, other=0.0)
    mask_x_1d = idx_feats < dim
-
-    if launch_pdl:
-        tl.extra.cuda.gdc_launch_dependents()
-
    for idx_token in range(segment_len):
        acc = acc_preload

@@ -754,12 +741,11 @@ def causal_conv1d_fn(
        BLOCK_M=BLOCK_M,
        BLOCK_N=256,
        num_stages=2,
-        launch_pdl=current_platform.is_arch_support_pdl(),
    )
    return out.to(original_x_dtype)


-@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"])
+@triton.jit()
 def _causal_conv1d_update_kernel(
    # Pointers to matrices
    x_ptr,  # (batch, dim, seqlen)
@@ -777,7 +763,7 @@ def _causal_conv1d_update_kernel(
    dim: tl.constexpr,
    seqlen: tl.constexpr,
    state_len: tl.constexpr,
-    num_cache_lines,  # added to support vLLM larger cache lines
+    num_cache_lines: tl.constexpr,  # added to support vLLM larger cache lines
    # Strides
    stride_x_seq: tl.constexpr,
    stride_x_dim: tl.constexpr,
@@ -803,16 +789,10 @@ def _causal_conv1d_update_kernel(
    NP2_STATELEN: tl.constexpr,
    HAS_NULL_BLOCK: tl.constexpr,
    BLOCK_N: tl.constexpr,
-    launch_pdl: tl.constexpr,
 ):
-    if launch_pdl:
-        tl.extra.cuda.gdc_wait()
-
    # ruff: noqa: E501
    idx_seq = tl.program_id(0)
    if idx_seq >= batch:
-        if launch_pdl:
-            tl.extra.cuda.gdc_launch_dependents()
        return

    # [BLOCK_N,] elements along the feature-dimension (channel)
@@ -834,8 +814,6 @@ def _causal_conv1d_update_kernel(
    if HAS_NULL_BLOCK:  # noqa
        if conv_states_input_coord == null_block_id:
            # not processing as this is not the actual sequence
-            if launch_pdl:
-                tl.extra.cuda.gdc_launch_dependents()
            return

    if IS_VARLEN:
@@ -853,8 +831,6 @@ def _causal_conv1d_update_kernel(
        o_offset = idx_seq * stride_o_seq

    if query_start_index == query_end_index:
-        if launch_pdl:
-            tl.extra.cuda.gdc_launch_dependents()
        return

    if IS_SPEC_DECODING:
@@ -993,9 +969,6 @@ def _causal_conv1d_update_kernel(
    mask_x_1d = idx_feats < dim

    # STEP 5: compute each token
-    if launch_pdl:
-        tl.extra.cuda.gdc_launch_dependents()
-
    for idx_token in tl.range(seqlen):
        acc = acc_preload

@@ -1272,13 +1245,14 @@ def causal_conv1d_update(
        NP2_STATELEN=np2_statelen,
        HAS_NULL_BLOCK=null_block_id is not None,
        BLOCK_N=256,
-        launch_pdl=current_platform.is_arch_support_pdl(),
    )
    if unsqueeze:
        out = out.squeeze(-1)
    return out.to(original_x_dtype)


+from vllm.platforms import current_platform  # noqa: E402
+
 if current_platform.is_cpu():
    from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
        causal_conv1d_fn_cpu,
@@ -28,7 +28,6 @@ class MLAModules:
    is_sparse: bool
    topk_indices_buffer: torch.Tensor | None
    indexer_rotary_emb: torch.nn.Module | None = None
-    g_proj: torch.nn.Module | None = None


 # --8<-- [start:multi_head_latent_attention]
@@ -67,7 +66,6 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        skip_topk: bool = False,
-        non_causal_multi_token_decode: bool = False,
        allow_short_prefill_indexer_scoring_skip: bool = False,
    ) -> None:
        super().__init__()
@@ -91,7 +89,6 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
        self.indexer = mla_modules.indexer
        self.indexer_rope_emb = mla_modules.indexer_rotary_emb
        self.is_sparse = mla_modules.is_sparse
-        self.g_proj = mla_modules.g_proj

        # Whether to skip top-k token selection computation in this layer.
        # When True, the indexer will not be called, and the layer will reuse
@@ -123,7 +120,6 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
            use_sparse=self.is_sparse,
            indexer=self.indexer,
            topk_indices_buffer=mla_modules.topk_indices_buffer,
-            non_causal_multi_token_decode=non_causal_multi_token_decode,
        )
        indexer_op = getattr(self.indexer, "indexer_op", None)
        if indexer_op is not None and hasattr(
@@ -220,7 +216,4 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
            q_dcp_replicated=q_dcp_replicated,
        )

-        if self.g_proj is not None:
-            attn_out = attn_out * self.g_proj(hidden_states)[0].sigmoid()
-
        return self.o_proj(attn_out)[0]
@@ -46,28 +46,8 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
        super().__init__(moe)
        self.group_size = 32
        self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
-        force_marlin = moe.moe_backend == "marlin"
-        cutlass_supports_activation = CutlassExpertsMxfp4._supports_activation(
-            moe.activation
-        )
-        if force_marlin:
-            logger.info_once(
-                "Using MarlinExperts for MXFP4 MoE because "
-                "--moe-backend=marlin was requested."
-            )
-        elif not cutlass_supports_activation:
-            logger.warning_once(
-                "CutlassExpertsMxfp4 does not support %s activation; "
-                "falling back to MarlinExperts for MXFP4 MoE.",
-                moe.activation.value,
-            )
-        # Use CUTLASS only when auto-selected and the activation is supported.
-        # Otherwise fall back to weight-only Marlin, which supports SITU.
-        self.use_cutlass_mxfp4 = (
-            not force_marlin
-            and CutlassExpertsMxfp4._supports_current_device()
-            and cutlass_supports_activation
-        )
+        # use cutlass if supported, otherwise fallback to marlin for weight-only FP4
+        self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
        self.experts_cls: type[mk.FusedMoEExperts]
        if self.use_cutlass_mxfp4:
            logger.info_once("Using CutlassExpertsMxfp4 for MXFP4 MoE")
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
stefankoncarevicandGitHub	381b691620	[ROCm][CI] Fix Kimi K3 KDA on ROCm (#50262 ) Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>	2026-07-29 15:47:50 +00:00
Michael Goin GitHub OpenAI Codex	d6247d7173	[Spec Decode][Perf] Replicate DSpark Markov head across TP ranks (#49731 ) Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>	2026-07-29 11:43:51 -04:00
Nick Hill GitHub Chris Eastwood Yifan Qiao	a0c092ee72	[BugFix] Fix `num_output_placeholders` preemption underflow (#48245 ) Signed-off-by: Chris Eastwood <chris.eastwood@pwn4g3.dev> Signed-off-by: Nick Hill <nickhill123@gmail.com> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Chris Eastwood <chris.eastwood@pwn4g3.dev> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>	2026-07-29 06:41:54 -07:00
Taneem IbrahimandGitHub	43eaefba5a	[ModelRunner V2] Enable sequence pooling for embedding and classification models (#48791 ) Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>	2026-07-29 06:30:45 -07:00
christianandGitHub	e0cfa52d22	[Bugfix][Frontend] Return transcription and translation verbose as float (#49073 ) Signed-off-by: Lucas Christian <lucaschgf7@gmail.com>	2026-07-29 13:16:13 +00:00
Reid GitHub Bugen Zhao	242c591d5a	[Rust Frontend] Send multimodal tensors in auxiliary frames (#49341 ) Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-29 12:21:09 +00:00
stefankoncarevicandGitHub	625871b52c	[CI][Test] Fix pooling truncation test after VLLMError hierarchy change (#50241 ) Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>	2026-07-29 12:15:56 +00:00
Yintong Lu GitHub Yongqi Wang Kunshang Ji	72297d859b	[XPU] Route weightless RMSNorm to _C dispatch (#47121 ) Signed-off-by: Yintong Lu <yintong.lu@intel.com> Co-authored-by: Yongqi Wang <yongqi.wang@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>	2026-07-29 18:59:58 +08:00
sungsoo ha GitHub Claude Opus 4.8 Jiangyun Zhu	f51193b9ae	[Kernel][Mamba] Fused-kernel support for align-mode DS-conv state migration with num_accepted_tokens > 1 (#49291 ) Signed-off-by: Sungsoo Ha <sungsooh@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>	2026-07-29 18:54:52 +08:00
Guan-Ming ChiuandGitHub	aeaa50a71c	[Bugfix][Multimodal] Include media IO config in MM cache hash (#49975 ) Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>	2026-07-29 10:48:40 +00:00
Chang Guo GitHub Cursor Agent	542a8fad6d	[KV Offload] Move CPUOffloadingSpec onto SharedOffloadRegion (#50094 ) Signed-off-by: Change72 <changg@nvidia.com> Co-authored-by: Cursor Agent <cursor-agent@cursor.com>	2026-07-29 13:05:50 +03:00
Minh VuandGitHub	9a4e5f9539	[CI/Perf] Fix malformed serving benchmark config (#43538 ) Signed-off-by: Minh Vu <vuhoangminh97@gmail.com>	2026-07-29 09:52:42 +00:00
Maria GuevaraandGitHub	c44e191b01	[Rust Frontend] Add --limit-mm-per-prompt support (#49604 ) Signed-off-by: Maria Guevara <kawaiiplush14@gmail.com>	2026-07-29 17:21:41 +08:00
fxmarty-amdandGitHub	5b14019576	[CI] Fix MXFP8 MOE backend selection tests on gfx942 (#50222 ) Signed-off-by: Felix Marty <Felix.Marty@amd.com>	2026-07-29 17:17:11 +08:00
omerpaz95andGitHub	dad7a6383b	[EC Connector] Add has_pending_push_work (#49582 ) Signed-off-by: omerpaz95 <omerpaz95@gmail.com>	2026-07-29 11:04:22 +02:00
Yan Ma GitHub Kunshang Ji Cyrus Leung	5b29c958c7	[XPU] upgrade to torch 2.13 (#48677 ) Signed-off-by: Yan Ma <yan.ma@intel.com> Signed-off-by: Kunshang Ji <kunshang.ji@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-07-29 01:24:56 -07:00
rongfu.leng GitHub Cyrus Leung	df2735ea2e	[Misc][Minimax-M3]add default video_processor (#50092 ) Signed-off-by: rongfu.leng <lenronfu@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-07-29 01:24:51 -07:00
Jared Wen GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Cyrus Leung	32e657e689	[BugFix] eagle draft max position embeddings (#49343 ) Signed-off-by: JaredforReal <w13431838023@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-07-29 01:24:47 -07:00
Perkz Zheng GitHub Perkz Zheng	ad5d29db70	[Model] Support Qwen3.5 text-only dense and MoE models (#50210 ) Signed-off-by: Perkz Zheng <PerkzZheng@users.noreply.github.com> Co-authored-by: Perkz Zheng <PerkzZheng@users.noreply.github.com>	2026-07-29 08:21:57 +00:00