disable k3dspark test

Signed-off-by: Isotr0py <Isotr0py@outlook.com>
fix k3dspark test
2026-07-29 15:44:17 +00:00 · 2026-07-29 13:34:36 +00:00 · 2026-07-29 13:26:41 +00:00 · 2026-07-29 11:39:37 +00:00 · 2026-07-29 11:33:54 +00:00 · 2026-07-29 11:15:29 +00:00
114 changed files with 3205 additions and 2176 deletions
@@ -369,7 +369,7 @@ export HF_TOKEN ZE_AFFINITY_MASK
    -e CMDS \
    --name "${container_name}" \
    "${IMAGE}" \
-    bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"' \
+    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"' \
    >/dev/null
 } 9>/tmp/docker-pull.lock

@@ -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/deepseek-ai/DeepGEMM.git")
+  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git")
  # NOTE: This is currently targeting nv-dev branch due to sm120 support
-  set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a")
+  set(_DEEPGEMM_UPSTREAM_TAG "f5a76426fa084087169693fd0cd815223576d6e9")

  set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
  if(NOT _deepgemm_fc_root)
@@ -63,6 +63,21 @@ 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 && \
@@ -85,6 +100,22 @@ 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"]

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


-ENV LD_LIBRARY_PATH=/opt/venv/lib:/usr/local/lib
+
+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
 CMD ["/bin/bash"]

 ######################### UCX + NIXL BUILD STAGE #########################
@@ -184,7 +216,8 @@ 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.2
+    uv pip install triton-xpu==3.7.1 && \
+    uv pip uninstall oneccl oneccl-devel

 # 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.2 --extra-index-url https://download.pytorch.org/whl/xpu
+    pip install triton-xpu==3.7.1 --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.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).
+        - 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).

 - Finally, build and install vLLM XPU backend:

@@ -547,6 +547,7 @@ 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,6 +155,8 @@ 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.10
-tilelang==0.1.9
+apache-tvm-ffi==0.1.11
+tilelang==0.1.12
 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.10
+apache-tvm-ffi==0.1.11
    # 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==2026.0.0
+dpcpp-cpp-rt==2025.3.2
    # 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.18.0
+impi-rt==2021.17.2
    # via
    #   oneccl
    #   torch
 iniconfig==2.3.0
    # via pytest
-intel-cmplr-lib-rt==2026.0.0
+intel-cmplr-lib-rt==2025.3.2
    # via
    #   intel-sycl-rt
    #   torch
-intel-cmplr-lib-ur==2026.0.0
+intel-cmplr-lib-ur==2025.3.2
    # via
    #   intel-openmp
    #   intel-sycl-rt
    #   torch
-intel-cmplr-lic-rt==2026.0.0
+intel-cmplr-lic-rt==2025.3.2
    # via
    #   intel-opencl-rt
    #   intel-sycl-rt
    #   torch
-intel-opencl-rt==2026.0.0
+intel-opencl-rt==2025.3.2
    # via
    #   dpcpp-cpp-rt
    #   onemkl-sycl-blas
@@ -282,14 +282,14 @@ intel-opencl-rt==2026.0.0
    #   onemkl-sycl-rng
    #   onemkl-sycl-sparse
    #   torch
-intel-openmp==2026.0.0
+intel-openmp==2025.3.2
    # via
    #   dpcpp-cpp-rt
    #   mkl
    #   torch
-intel-pti==0.17.0
+intel-pti==0.16.0
    # via torch
-intel-sycl-rt==2026.0.0
+intel-sycl-rt==2025.3.2
    # 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==2026.0.0
+mkl==2025.3.1
    # via
    #   onemkl-sycl-blas
    #   onemkl-sycl-dft
@@ -453,28 +453,28 @@ numpy==2.2.6
    #   torchvision
    #   transformers
    #   xgrammar
-oneccl==2022.0.0
+oneccl==2021.17.2
    # via
    #   oneccl-devel
    #   torch
-oneccl-devel==2022.0.0
+oneccl-devel==2021.17.2
    # via torch
-onemkl-license==2026.0.0
+onemkl-license==2025.3.1
    # via
    #   mkl
    #   torch
-onemkl-sycl-blas==2026.0.0
+onemkl-sycl-blas==2025.3.1
    # via
    #   onemkl-sycl-lapack
    #   onemkl-sycl-sparse
    #   torch
-onemkl-sycl-dft==2026.0.0
+onemkl-sycl-dft==2025.3.1
    # via torch
-onemkl-sycl-lapack==2026.0.0
+onemkl-sycl-lapack==2025.3.1
    # via torch
-onemkl-sycl-rng==2026.0.0
+onemkl-sycl-rng==2025.3.1
    # via torch
-onemkl-sycl-sparse==2026.0.0
+onemkl-sycl-sparse==2025.3.1
    # via torch
 openai==2.44.0
    # via
@@ -719,8 +719,6 @@ pyyaml==6.0.3
    #   timm
    #   transformers
    #   uvicorn
-pyzes==0.1.1
-    # via torch
 pyzmq==27.1.0
    # via
    #   -c requirements/common.txt
@@ -873,14 +871,14 @@ tabledata==1.3.4
    # via pytablewriter
 tabulate==0.10.0
    # via sacrebleu
-tbb==2023.0.0
+tbb==2022.3.1
    # via
    #   intel-opencl-rt
    #   mkl
    #   torch
 tblib==3.1.0
    # via -r requirements/test/xpu.in
-tcmlib==1.5.0
+tcmlib==1.4.1
    # via
    #   tbb
    #   torch
@@ -912,7 +910,7 @@ tokenizers==0.22.2
    #   -c requirements/common.txt
    #   -r requirements/test/../common.txt
    #   transformers
-torch==2.13.0+xpu
+torch==2.12.0+xpu
    # via
    #   -c requirements/xpu.txt
    #   accelerate
@@ -922,7 +920,7 @@ torch==2.13.0+xpu
    #   timm
    #   torchvision
    #   xgrammar
-torchvision==0.28.0+xpu
+torchvision==0.27.0+xpu
    # via timm
 tqdm==4.67.3
    # via
@@ -948,7 +946,7 @@ transformers==5.14.1
    #   xgrammar
 triton==3.7.1
    # via xgrammar
-triton-xpu==3.7.2
+triton-xpu==3.7.1
    # via torch
 typepy==1.3.4
    # via
@@ -1003,7 +1001,7 @@ typing-inspection==0.4.2
    #   mcp
    #   pydantic
    #   pydantic-settings
-umf==1.1.0
+umf==1.0.3
    # 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.13.0
+torch==2.12.0
 torchaudio
 torchvision
 torchcodec >= 0.14 # Required for the torchcodec video decoding backend

-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
+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
@@ -57,7 +57,6 @@ 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());
@@ -232,7 +231,6 @@ 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::{MmLimitPerPrompt, MultimodalModelInfo};
+use crate::multimodal::MultimodalModelInfo;
 use crate::output::DynChatOutputProcessor;
 use crate::renderer::DynChatRenderer;
 use crate::request::ChatRequest;
@@ -74,9 +74,6 @@ 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,8 +23,6 @@ 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}`")]
@@ -89,8 +87,7 @@ impl Error {
            Self::Text(error) => error.is_request_validation_error(),
            Self::UnsupportedMultimodalRenderer
            | Self::UnsupportedMultimodalContent(_)
-            | Self::UnsupportedModality { .. }
-            | Self::MmLimitExceeded { .. } => true,
+            | Self::UnsupportedModality { .. } => 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::{BTreeMap, HashMap, HashSet};
+use std::collections::HashSet;
 use std::fs;
 use std::path::Path;
 use std::sync::{Arc, LazyLock};
@@ -24,7 +24,6 @@ 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;
@@ -53,71 +52,6 @@ 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.
@@ -353,7 +287,6 @@ 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) => {
@@ -386,12 +319,7 @@ impl MultimodalModelInfo {
            tokenizer: TokenizerResolver(tokenizer),
        };

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

    /// Resolve multimodal support from an assembled context and parsed
@@ -400,7 +328,6 @@ 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,
@@ -429,7 +356,6 @@ impl MultimodalModelInfo {
            video,
            audio,
            media_connector,
-            limit_mm_per_prompt,
        }))
    }

@@ -641,55 +567,7 @@ 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.
    ///
@@ -706,7 +584,8 @@ impl MultimodalModelInfo {
        if media_parts_len == 0 {
            return Ok(Vec::new());
        }
-        self.validate_mm_limits(&media_parts)?;
+        // TODO: enforce per-modality item-count limits, aligned with the
+        // engine's `--limit-mm-per-prompt` semantics.
        let fetched = self.fetch_media(media_parts).await?;

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

-    fn test_info_with_limits(
+    fn test_info(
        model_type: &str,
        config: serde_json::Value,
        tokenizer: TestTokenizer,
-        limit_mm_per_prompt: MmLimitPerPrompt,
    ) -> MultimodalModelInfo {
        let context = MultimodalModelContext {
            model_id: format!("{model_type}-test"),
@@ -891,20 +769,11 @@ 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",
@@ -914,19 +783,16 @@ mod tests {
        test_info("llama4", config, llama4_tokenizer())
    }

-    fn qwen3_vl_config() -> serde_json::Value {
-        serde_json::json!({
+    pub(super) fn qwen3_vl_info() -> MultimodalModelInfo {
+        let config = 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}
-        })
-    }
-
-    pub(super) fn qwen3_vl_info() -> MultimodalModelInfo {
-        test_info("qwen3_vl", qwen3_vl_config(), qwen3_vl_tokenizer())
+        });
+        test_info("qwen3_vl", config, qwen3_vl_tokenizer())
    }

    #[test]
@@ -987,123 +853,4 @@ 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,7 +107,6 @@ mod tests {
            context,
            PreProcessorConfig::default(),
            PreProcessorConfig::default(),
-            HashMap::new(),
        )
        .unwrap()
        .expect("Inkling multimodal support")
@@ -127,7 +126,6 @@ mod tests {
            context,
            PreProcessorConfig::default(),
            PreProcessorConfig::default(),
-            HashMap::new(),
        )
        .unwrap()
        .expect("Qwen3-ASR multimodal support")
@@ -207,7 +207,6 @@ mod tests {
                Some("qwen3_vl".to_string()),
                files,
                Arc::new(qwen3_vl_tokenizer()),
-                std::collections::HashMap::new(),
            )
        };

@@ -23,7 +23,6 @@ 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};
@@ -187,17 +186,6 @@ 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
@@ -360,18 +348,6 @@ 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()
@@ -423,7 +399,6 @@ 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,
@@ -476,7 +451,6 @@ 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,
@@ -692,7 +666,6 @@ impl ServeArgs {
            self.runtime.disable_log_stats,
            self.runtime.shutdown_timeout,
            handshake_port,
-            self.runtime.limit_mm_per_prompt_json(),
        )
    }
 }
@@ -64,7 +64,6 @@ 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,
@@ -763,7 +762,6 @@ 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,
@@ -1119,24 +1117,6 @@ 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"])
@@ -1351,7 +1331,6 @@ 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,
@@ -1495,7 +1474,6 @@ 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 {
@@ -1580,7 +1558,6 @@ 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 {
@@ -1683,7 +1660,6 @@ 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,6 +299,11 @@ 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>,
@@ -91,7 +91,6 @@ 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.
@@ -127,10 +126,6 @@ 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,7 +1,6 @@
 // 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};
@@ -71,7 +70,6 @@ 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,7 +10,6 @@ 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};

@@ -185,9 +184,6 @@ 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,7 +101,6 @@ 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,12 +577,6 @@ 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()
@@ -600,7 +594,6 @@ fn qwen_multimodal_model_info_with_limits(
            ..Default::default()
        },
        Arc::new(fake_chat_tokenizer()),
-        limit_mm_per_prompt,
    )
    .expect("load multimodal info")
    .expect("qwen multimodal info is registered");
@@ -2314,74 +2307,6 @@ 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 not _is_xpu() and sys.version_info >= (3, 11):
+if sys.version_info >= (3, 11):
    ext_modules.append(CMakeExtension(name="vllm.spinloop"))
    ext_modules.append(CMakeExtension(name="vllm.fs_io_C"))

@@ -1,115 +0,0 @@
-# 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)
@@ -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_emulation_moe import (  # noqa: E402
-    Mxfp8EmulationTritonExperts,
+from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (  # noqa: E402
+    Mxfp8NativeTritonExperts,
 )
 from vllm.model_executor.layers.fused_moe.modular_kernel import (  # noqa: E402
    FusedMoEActivationFormat,
@@ -150,20 +150,16 @@ def test_explicit_moe_backend_aiter():

 def test_gfx950_picks_aiter():
    """Auto-select on real ROCm hardware with flydsl usable -> FlyDSL wins."""
-    # 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),
-    ):
+    with _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_emulation():
+def test_gfx942_picks_triton():
    """flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton
    dot_scaled backend wins instead."""
-    with patch(f"{_AITER_MOD}.current_platform.supports_mx", return_value=False):
+    with _flydsl_installed(False):
        backend, experts_cls = select_mxfp8_moe_backend(_config())
-    assert backend is Fp8MoeBackend.EMULATION
-    assert experts_cls is Mxfp8EmulationTritonExperts
+    assert backend is Fp8MoeBackend.TRITON_MXFP8
+    assert experts_cls is Mxfp8NativeTritonExperts
@@ -497,8 +497,6 @@ _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",
@@ -1026,6 +1024,10 @@ _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"},
@@ -1466,6 +1468,14 @@ _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",
@@ -1664,6 +1674,12 @@ _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,
@@ -0,0 +1,153 @@
+# 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
@@ -187,7 +187,3 @@ def test_merge_multimodal_embeddings_no_sync():
        _merge_multimodal_embeddings(
            inputs_embeds, multimodal_embeddings, is_multimodal
        )
-
-
-def test_ci_retry_workflow():
-    pytest.fail("Intentional failure for /ci retry workflow validation")
@@ -14,7 +14,6 @@ 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,
@@ -389,12 +388,6 @@ 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,6 +1235,26 @@ 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"),
    [
@@ -0,0 +1,165 @@
+# 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,7 +73,11 @@ 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)
+    layer = SimpleNamespace(
+        kv_cache_dtype=cache_dtype,
+        head_size=576,
+        non_causal_multi_token_decode=False,
+    )
    vllm_config = SimpleNamespace(
        cache_config=SimpleNamespace(block_size=64), model_config=None
    )
@@ -0,0 +1,124 @@
+# 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,6 +40,8 @@ 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,
@@ -75,6 +77,78 @@ 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
@@ -1078,6 +1078,8 @@ 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,
@@ -333,6 +333,54 @@ 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
 # ---------------------------------------------------------------------------
@@ -1398,6 +1398,53 @@ 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
@@ -1,112 +0,0 @@
-# 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,6 +1,9 @@
 # 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
@@ -8,6 +11,7 @@ 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,
 )

@@ -68,3 +72,21 @@ 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
@@ -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/deepseek-ai/DeepGEMM.git"
+DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git"
 # NOTE: This is currently targeting nv-dev branch due to sm120 support
-DEEPGEMM_GIT_REF="a6b593d2826719dcf4892609af7b84ee23aaf32a"
+DEEPGEMM_GIT_REF="f5a76426fa084087169693fd0cd815223576d6e9"
 WHEEL_DIR=""

 # Parse command line arguments
@@ -133,6 +133,8 @@ 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
@@ -143,6 +145,12 @@ 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,
@@ -193,6 +201,8 @@ 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)
@@ -1559,6 +1569,7 @@ class rocm_aiter_ops:
            "silu": ActivationType.Silu,
            "gelu": ActivationType.Gelu,
            "swiglu": ActivationType.Swiglu,
+            "situ": getattr(ActivationType, "Situv2", None),
        }
        return mapping.get(name)

@@ -2155,6 +2166,8 @@ 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,
@@ -2179,6 +2192,8 @@ class rocm_aiter_ops:
            bias2,
            moe_sorting_dispatch_policy,
            swiglu_limit,
+            beta,
+            linear_beta,
        )

    @staticmethod
@@ -2391,6 +2391,17 @@ 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,
@@ -2436,6 +2447,21 @@ 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,7 +771,6 @@ 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,11 +1028,6 @@ 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,6 +51,7 @@ MTPModelTypes = Literal[
    "minimax_m3_mtp",
    "bailing_hybrid_mtp",
    "mtp",
+    "kimi_k3_mtp",
    "pangu_ultra_moe_mtp",
    "step3p5_mtp",
    "hy_v3_mtp",
@@ -355,6 +356,16 @@ 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)
@@ -910,15 +921,6 @@ 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
@@ -943,6 +945,7 @@ 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.
@@ -972,6 +975,16 @@ 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"
                ):
@@ -1139,40 +1152,6 @@ 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,11 +72,34 @@ 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."""
@@ -619,16 +642,20 @@ class VllmConfig:
        if model_config.runner_type != "generate":
            return False

-        if getattr(model_config, "is_hybrid", 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
+        ):
            return False

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

    @property
    def needs_dp_coordinator(self) -> bool:
@@ -1196,6 +1223,9 @@ class VllmConfig:
                    "DeepSeekV4MTPModel",
                    "InklingForCausalLM",
                    "InklingForConditionalGeneration",
+                    "KimiK3ForConditionalGeneration",
+                    "KimiK3MTPModel",
+                    "KimiLinearForCausalLM",
                    "MiniMaxM3SparseForCausalLM",
                    "MiniMaxM3SparseForConditionalGeneration",
                )
@@ -2277,14 +2307,6 @@ 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,6 +340,18 @@ 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,6 +25,11 @@ 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__)


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


 class CustomAllreduce:
-    _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8]
+    _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

    # max_size: max supported allreduce size
    def __init__(
@@ -57,6 +72,10 @@ 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:
        """
@@ -70,7 +89,20 @@ 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
@@ -87,13 +119,8 @@ class CustomAllreduce:
            "CustomAllreduce should be attached to a non-NCCL group."
        )

-        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
+        same_node = all(in_the_same_node_as(group, source_rank=0))
+        self.mnnvl_only = not same_node

        rank = dist.get_rank(group=self.group)
        self.rank = rank
@@ -126,28 +153,30 @@ 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:
+            if (
+                device_capability_str in CUSTOM_ALL_REDUCE_MAX_SIZES
+                and world_size in CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str]
+            ):
                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.
-        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:
+        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:
            logger.warning(
                "Custom allreduce is disabled because it's not supported on"
                " more than two PCIe-only GPUs. To silence this warning, "
@@ -158,7 +187,11 @@ 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 not current_platform.is_rocm() and not _can_p2p(rank, world_size):
+        if (
+            same_node
+            and 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 "
@@ -170,21 +203,40 @@ 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.
-        self.meta_ptrs = self.create_shared_buffer(
-            ops.meta_size() + max_size, group=group, uncached=True
-        )
+        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
        # This is a pre-registered IPC buffer. In eager mode, input tensors
-        # 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.
+        # 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.
        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
@@ -192,6 +244,72 @@ 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):
@@ -228,7 +346,7 @@ class CustomAllreduce:
        ops.register_graph_buffers(self._ptr, handles, offsets)

    def should_custom_ar(self, inp: torch.Tensor):
-        if self.disabled:
+        if self.disabled or self.world_size > 8:
            return False
        inp_size = inp.numel() * inp.element_size()
        # custom allreduce requires input byte size to be multiples of 16
@@ -279,6 +397,111 @@ 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:
@@ -286,6 +509,10 @@ 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,15 +275,3 @@ 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,6 +418,10 @@ 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,6 +180,12 @@ 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,8 +92,20 @@ 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,8 +191,15 @@ 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,6 +741,7 @@ 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={...})`
@@ -1637,6 +1638,13 @@ 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
@@ -2431,6 +2439,8 @@ 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,
@@ -131,6 +131,7 @@ 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
@@ -176,6 +177,7 @@ 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
@@ -272,6 +274,7 @@ 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
@@ -1004,7 +1007,6 @@ 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
@@ -1211,6 +1213,12 @@ 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
@@ -1448,6 +1456,11 @@ 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)
@@ -1911,6 +1924,14 @@ 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
@@ -272,6 +272,7 @@ 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
@@ -375,6 +376,7 @@ 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__()
@@ -391,6 +393,7 @@ 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

@@ -1123,6 +1126,7 @@ 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):
@@ -1450,6 +1454,8 @@ 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

@@ -1724,6 +1730,9 @@ 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.
@@ -1824,8 +1833,14 @@ 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
+        )

-        self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
+        # 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.mla_dims = get_mla_dims(self.model_config)
        self.aot_schedule = current_platform.is_cuda()

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

-        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,
+        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,
+                )
            )
-        )

        assert num_decodes + num_prefills == num_reqs
        assert num_decode_tokens + num_prefill_tokens == num_tokens
@@ -2049,6 +2092,7 @@ 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,6 +841,7 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
            MoEActivation.SILU,
            MoEActivation.GELU,
            MoEActivation.GELU_TANH,
+            MoEActivation.SITU,
            MoEActivation.SWIGLUOAI,
            MoEActivation.SILU_NO_MUL,
            MoEActivation.GELU_NO_MUL,
@@ -8,7 +8,9 @@ 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,
@@ -53,6 +55,68 @@ 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,
@@ -103,6 +167,19 @@ 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),),
@@ -133,6 +210,8 @@ 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,
@@ -191,6 +270,8 @@ 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,
@@ -671,10 +752,10 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):

    @staticmethod
    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
-        return not (
-            moe_parallel_config.use_fi_nvl_two_sided_kernels
-            or moe_parallel_config.use_fi_nvl_one_sided_kernels
-        )
+        # 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

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

            num_experts, max_num_tokens = hidden_states.shape[:2]
-            beta = activation_situ_beta
+            beta = 1.0 if activation_situ_beta is None else 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,6 +5,7 @@ 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
@@ -262,6 +263,8 @@ 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}")

@@ -356,10 +359,12 @@ 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
-        hidden_pad = hidden_pad // 128 * 128
-        intermediate_pad = (
-            intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1)
-        )
+        # 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)
+            )

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

        gate_mode = ""
-        if quant_config.use_mxfp4_w4a16:
+        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:
            gate_mode = GateMode.INTERLEAVE.value
        elif activation_interleave is not None:
            gate_mode = (
@@ -401,6 +414,8 @@ 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,6 +140,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
            MoEActivation.SILU,
            MoEActivation.GELU,
            MoEActivation.GELU_TANH,
+            MoEActivation.SITU,
            MoEActivation.SWIGLUOAI,
            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
            MoEActivation.SWIGLUSTEP,
@@ -315,6 +315,75 @@ 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,
@@ -335,7 +404,6 @@ 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:
@@ -346,49 +414,27 @@ 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),
-        }

-        from flashinfer import trtllm_fp4_block_scale_routed_moe
-
-        trtllm_fp4_block_scale_routed_moe(**kwargs)
+        # 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,
+            )

        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, 256)
+        intermediate_size = round_up(intermediate_size, 128)
        hidden_size = round_up(hidden_size, 256)
    elif backend in (
        Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
@@ -6,11 +6,6 @@ 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

@@ -224,6 +219,12 @@ 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)

@@ -144,8 +144,10 @@ class LogitsProcessor(PluggableLayer):
        # Get the logits for the next tokens.
        logits = self._apply_head(lm_head, hidden_states, embedding_bias)

-        # Gather logits for TP
-        logits = self._gather_logits(logits)
+        # 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):
+            logits = self._gather_logits(logits)

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

        logits = self._apply_head(lm_head, hidden_states, embedding_bias)
        if self.soft_cap is not None:
@@ -1,32 +1,33 @@
 # 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.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.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.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 sharded_weight_loader
-from vllm.model_executor.utils import set_weight_attrs
-from vllm.third_party.flash_linear_attention.ops.kda import (
-    FusedRMSNormGated,
-    chunk_kda_with_fused_gate,
-    fused_kda_gate,
-    fused_recurrent_kda,
+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.utils import set_weight_attrs
+from vllm.platforms import current_platform
+from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
 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,
 )
@@ -36,49 +37,118 @@ 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

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


-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 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_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
+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


-direct_register_custom_op(
-    op_name="kda_attention",
-    op_func=kda_attention,
-    mutates_args=["core_attn_out"],
-    fake_impl=kda_attention_fake,
-)
+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


@PluggableLayer.register("kimi_gated_delta_net_attention")
@@ -96,7 +166,11 @@ 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
+            self.tp_size,
+            self.num_heads,
+            self.head_dim,
+            conv_kernel_size=self.conv_size,
+            num_spec=self.num_spec,
        )

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

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

-        self.q_proj = ColumnParallelLinear(
+        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.hidden_size,
-            projection_size,
+            in_proj_output_sizes,
+            replicated_shard_id=4,
+            tp_size=self.tp_size,
            bias=False,
            quant_config=self.quant_config,
-            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",
+            prefix=f"{prefix}.in_proj_qkvgfab",
        )
+        if self.in_proj_padding:
+            self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_()

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

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

-        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(
+        # 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(
            input_size=self.conv_size,
-            output_size=projection_size,
+            output_size=3 * self.projection_size,
            bias=False,
            params_dtype=torch.float32,
-            prefix=f"{prefix}.q_conv1d",
+            prefix=f"{prefix}.conv1d",
        )
-        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.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.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(1, 1, self.local_num_heads, 1, dtype=torch.float32)
+            torch.empty(self.local_num_heads, dtype=torch.float32)
        )
-        set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(2)})
+        set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)})

-        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.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_b_proj = ColumnParallelLinear(
-            self.head_dim,
-            projection_size,
-            bias=False,
-            quant_config=self.quant_config,
-            prefix=f"{prefix}.g_b_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}."
        )
+        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(
-            projection_size,
+            self.projection_size,
            self.hidden_size,
            bias=False,
            quant_config=self.quant_config,
            prefix=f"{prefix}.o_proj",
        )

-        compilation_config = get_current_vllm_config().compilation_config
+        compilation_config = 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,
@@ -237,42 +331,57 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
        output: torch.Tensor,
    ) -> None:
        num_tokens = hidden_states.size(0)
-        q = self.q_proj(hidden_states)[0]
-        k = self.k_proj(hidden_states)[0]
-        v = self.v_proj(hidden_states)[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]

-        beta = self.b_proj(hidden_states)[0].float().sigmoid()
-        g1 = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
+        g1 = self.f_b_proj(f_a)[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.zeros(
+        core_attn_out = torch.empty(
            (1, num_tokens, self.local_num_heads, self.head_dim),
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )
-        torch.ops.vllm.kda_attention(
-            q,
-            k,
-            v,
-            g1,
-            beta,
-            core_attn_out,
-            self.prefix,
+
+        self._forward(
+            mixed_qkv=mixed_qkv,
+            g1=g1,
+            g2=g2,
+            beta=beta,
+            core_attn_out=core_attn_out,
        )
-        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,
-        q_proj_states: torch.Tensor,
-        k_proj_states: torch.Tensor,
-        v_proj_states: torch.Tensor,
+        mixed_qkv: torch.Tensor,
        g1: torch.Tensor,
+        g2: torch.Tensor,
        beta: torch.Tensor,
        core_attn_out: torch.Tensor,
    ) -> None:
@@ -280,165 +389,247 @@ 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)
-        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]
+        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]
        g1 = g1[:, :num_actual_tokens]
        beta = beta[:, :num_actual_tokens]

-        (conv_state, recurrent_state) = constant_caches
+        constant_caches = self.kv_cache
+
+        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_state_q, conv_state_k, conv_state_v = conv_state.chunk(3, dim=-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
+        )

-        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_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)
+        # 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:
-            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,
-            )
+            mixed_qkv_spec = g1_spec = beta_spec = None
+            mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta

-        q, k, v = map(
-            lambda x: rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim), (q, k, v)
-        )
+        # ---------- 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)

-        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,
+            # 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,
+                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,
-                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,
+                dt_bias=self.dt_bias,
+                lower_bound=self.gate_lower_bound,
                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,
+                cu_seqlens=spec_cu_seqlens,
+                ssm_state_indices=spec_state_indices_tensor,
+                num_accepted_tokens=num_accepted_tokens,
+                out=spec_out,
            )
-        core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
-            0, :num_actual_tokens
-        ]
+
+        # ---------- 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
+            ]
+        else:
+            assert core_attn_out_spec is not None
+        core_attn_out.copy_(self.o_norm(core_attn_out, g2))
@@ -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
+            divide(conv_dim, tp_world_size), conv_kernel_size - 1 + num_spec
        )
        recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim)
        return (conv_state_shape, recurrent_state_shape)
@@ -8,11 +8,12 @@
 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()
+@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"])
 def _causal_conv1d_fwd_kernel(  # continuous batching
    # Pointers to matrices
    x_ptr,  # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences
@@ -33,7 +34,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: tl.constexpr,  # added to support vLLM larger cache lines
+    num_cache_lines,  # 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)
@@ -58,6 +59,7 @@ 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
@@ -68,6 +70,9 @@ 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

@@ -79,6 +84,8 @@ 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)
@@ -136,6 +143,8 @@ 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
@@ -408,6 +417,10 @@ 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

@@ -741,11 +754,12 @@ 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()
+@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"])
 def _causal_conv1d_update_kernel(
    # Pointers to matrices
    x_ptr,  # (batch, dim, seqlen)
@@ -763,7 +777,7 @@ def _causal_conv1d_update_kernel(
    dim: tl.constexpr,
    seqlen: tl.constexpr,
    state_len: tl.constexpr,
-    num_cache_lines: tl.constexpr,  # added to support vLLM larger cache lines
+    num_cache_lines,  # added to support vLLM larger cache lines
    # Strides
    stride_x_seq: tl.constexpr,
    stride_x_dim: tl.constexpr,
@@ -789,10 +803,16 @@ 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)
@@ -814,6 +834,8 @@ 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:
@@ -831,6 +853,8 @@ 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:
@@ -969,6 +993,9 @@ 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

@@ -1245,14 +1272,13 @@ 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,6 +28,7 @@ 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]
@@ -66,6 +67,7 @@ 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__()
@@ -89,6 +91,7 @@ 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
@@ -120,6 +123,7 @@ 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(
@@ -216,4 +220,7 @@ 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,8 +46,28 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
        super().__init__(moe)
        self.group_size = 32
        self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
-        # use cutlass if supported, otherwise fallback to marlin for weight-only FP4
-        self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
+        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
+        )
        self.experts_cls: type[mk.FusedMoEExperts]
        if self.use_cutlass_mxfp4:
            logger.info_once("Using CutlassExpertsMxfp4 for MXFP4 MoE")
@@ -3,6 +3,7 @@

 import torch

+import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import (
@@ -17,6 +18,7 @@ from vllm.model_executor.layers.fused_moe import modular_kernel as mk
 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
    TRITON_BACKENDS,
    Mxfp4MoeBackend,
+    backend_to_kernel_cls,
    convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,
    convert_weight_to_mxfp4_moe_kernel_format,
    make_mxfp4_moe_kernel,
@@ -68,8 +70,11 @@ class Mxfp4Config(QuantizationConfig):
    def get_config_filenames(cls) -> list[str]:
        return []

-    # TODO (zyongye) This is only temporaty fallback.
-    # We should have `Mxfp4MoEMethod` after this migration is complete.
+    def _make_moe_method(self, moe: FusedMoEConfig) -> FusedMoEMethodBase:
+        """MoE method for RoutedExperts. Subclasses override to pick a
+        checkpoint-specific kernel family."""
+        return Mxfp4MoEMethod(moe)
+
    def get_quant_method(
        self, layer: torch.nn.Module, prefix: str
    ) -> "QuantizeMethodBase | None":
@@ -86,7 +91,7 @@ class Mxfp4Config(QuantizationConfig):
            )
            return UnquantizedLinearMethod()
        elif isinstance(layer, RoutedExperts):
-            return GptOssMxfp4MoEMethod(layer.moe_config)
+            return self._make_moe_method(layer.moe_config)
        elif isinstance(layer, Attention):
            logger.debug_once(
                "MXFP4 attention layer is not implemented. "
@@ -131,6 +136,9 @@ class GptOssMxfp4Config(Mxfp4Config):
            return None
        return "gpt_oss_mxfp4"

+    def _make_moe_method(self, moe: FusedMoEConfig) -> FusedMoEMethodBase:
+        return GptOssMxfp4MoEMethod(moe)
+

 class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
    """MXFP4 MoE quantization method."""
@@ -475,13 +483,45 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
        )


+def _use_k3_situ_aiter(moe: FusedMoEConfig) -> bool:
+    """Whether Kimi-K3's SiTU MXFP4 MoE should use the AITER A16W4 kernel.
+
+    K3 is weight-only MXFP4 (W4A16) with SiTU activation, which the generic
+    MXFP4 backend selector does not cover; route it to AITER on gfx950.
+    """
+    from vllm.platforms import current_platform
+
+    if not current_platform.is_rocm():
+        return False
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+    from vllm.platforms.rocm import on_gfx950
+
+    return (
+        rocm_aiter_ops.is_fused_moe_enabled()
+        and on_gfx950()
+        and moe.activation == MoEActivation.SITU
+        and moe.activation_situ_linear_beta is not None
+        and rocm_aiter_ops.get_aiter_activation_type("situ") is not None
+    )
+
+
 class Mxfp4MoEMethod(FusedMoEMethodBase):
    """MXFP4 MoE quantization method."""

    def __init__(self, moe: FusedMoEConfig):
        super().__init__(moe)
        self.weight_dtype = "mxfp4"
-        self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(moe)
+        self.is_k3_situ_aiter = _use_k3_situ_aiter(moe)
+        self.experts_cls: type[mk.FusedMoEExperts] | None
+        if self.is_k3_situ_aiter:
+            self.mxfp4_backend = Mxfp4MoeBackend.AITER_MXFP4_BF16
+            self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0]
+            logger.info_once("Using AITER_MXFP4_BF16 for Kimi-K3 SiTU MXFP4 MoE.")
+        else:
+            self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(
+                moe
+            )

        self.max_capture_size = moe.max_capture_size

@@ -523,6 +563,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
            act_dtype=act_dtype,
            moe_parallel_config=moe_parallel_config,
        )
+        if self.is_k3_situ_aiter:
+            # K3's AITER A16W4 kernel handles K3's native intermediate size
+            # (moe_intermediate 3072; e.g. 384/partition at TP8); the generic
+            # 256 round-up would inflate weights and OOM.
+            return hidden_size, intermediate_size_per_partition
        return mxfp4_round_up_hidden_size_and_intermediate_size(
            self.mxfp4_backend, hidden_size, intermediate_size_per_partition
        )
@@ -724,7 +769,58 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                layer=layer,
            )

+    def _setup_kernel_k3_situ(self, layer: RoutedExperts) -> None:
+        # K3's AITER A16W4 kernel wants the separated ([gate_all, up_all])
+        # stage-1 layout, unlike the interleaved gpt-oss/DeepSeek path in
+        # convert_weight_to_mxfp4_moe_kernel_format. Preshuffle once here.
+        from aiter.utility.fp4_utils import e8m0_shuffle
+
+        from vllm._aiter_ops import rocm_aiter_ops
+
+        fp4_dtype = torch.float4_e2m1fn_x2
+        e8m0_dtype = torch.float8_e8m0fnu
+        num_experts = layer.w13_weight.shape[0]
+
+        # a8w4 (AITER_SITUV2_A8W4=1) uses the gate/up-interleaved (_gui_) fp8
+        # flydsl kernels, which need w13 weight+scale in interleave layout.
+        # Default a16w4 keeps the separated layout.
+        guinterleave = envs.AITER_SITUV2_A8W4
+        w13 = rocm_aiter_ops.shuffle_weight_a16w4(
+            layer.w13_weight.data.view(fp4_dtype), 16, guinterleave
+        )
+        w2 = rocm_aiter_ops.shuffle_weight_a16w4(
+            layer.w2_weight.data.view(fp4_dtype), 16, False
+        )
+        w13_scale_raw = layer.w13_weight_scale.data.view(e8m0_dtype)
+        w2_scale_raw = layer.w2_weight_scale.data.view(e8m0_dtype)
+        w13_scale = rocm_aiter_ops.shuffle_scale_a16w4(
+            w13_scale_raw.view(-1, w13_scale_raw.shape[-1]), num_experts, guinterleave
+        )
+        w2_scale = e8m0_shuffle(w2_scale_raw.view(-1, w2_scale_raw.shape[-1]))
+
+        replace_parameter(layer, "w13_weight", w13)
+        replace_parameter(layer, "w2_weight", w2)
+        replace_parameter(layer, "w13_weight_scale", w13_scale)
+        replace_parameter(layer, "w2_weight_scale", w2_scale)
+        layer.w13_weight.is_shuffled = True
+        layer.w2_weight.is_shuffled = True
+
+        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
+        if self.moe_quant_config is not None and self.experts_cls is not None:
+            self.moe_kernel = make_mxfp4_moe_kernel(
+                moe_quant_config=self.moe_quant_config,
+                moe_config=self.moe,
+                mxfp4_backend=self.mxfp4_backend,
+                experts_cls=self.experts_cls,
+                routing_tables=layer._expert_routing_tables(),
+                layer=layer,
+            )
+
    def process_weights_after_loading(self, layer):
+        if self.is_k3_situ_aiter:
+            self._setup_kernel_k3_situ(layer)
+            return
+
        w13 = layer.w13_weight
        w2 = layer.w2_weight
        w13_scale = layer.w13_weight_scale
@@ -819,4 +915,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
            global_num_experts=layer.global_num_experts,
            expert_map=layer.expert_map,
            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            num_expert_group=layer.num_expert_group,
+            topk_group=layer.topk_group,
+            e_score_correction_bias=layer.e_score_correction_bias,
+            routed_scaling_factor=layer.routed_scaling_factor,
        )
@@ -21,17 +21,20 @@ def activation_to_flashinfer_int(activation: MoEActivation) -> int:
 def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType":
    from flashinfer.fused_moe.core import ActivationType

+    if activation == MoEActivation.SITU:
+        return ActivationType.Situ
+
    # silu and gelu are mapped to their gated versions SwiGLU and GeGLU respectively
    ACTIVATION_TO_FI_ACTIVATION = {
        MoEActivation.SILU_NO_MUL: ActivationType.Silu,
        MoEActivation.GELU_NO_MUL: ActivationType.Gelu,
        MoEActivation.SILU: ActivationType.Swiglu,
        # SwiGLU-OAI uses Swiglu; the OAI alpha/beta/clamp come from gemm1_* args.
+        MoEActivation.SWIGLUOAI: ActivationType.Swiglu,
        MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu,
        MoEActivation.GELU: ActivationType.Geglu,
        MoEActivation.GELU_TANH: ActivationType.Geglu,
        MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,
-        MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu,
    }
    return ACTIVATION_TO_FI_ACTIVATION[activation]

@@ -474,7 +474,11 @@ def prepare_moe_fp4_layer_for_marlin(

    group_size = 16 if is_nvfp4 else 32

-    e = layer.moe_config.num_experts
+    # Use the per-rank (local) expert count: under expert parallelism the
+    # w13_weight/w2_weight tensors only hold this rank's experts (created with
+    # local_num_experts), whereas moe_config.num_experts is the global count.
+    # With no EP the two are equal, so the non-EP path is unchanged.
+    e = layer.moe_config.num_local_experts
    k = layer.moe_config.hidden_dim
    n = layer.moe_config.intermediate_size_per_partition

@@ -41,8 +41,8 @@ def _mxfp8_e4m3_quantize_torch(
 ) -> tuple[torch.Tensor, torch.Tensor]:
    """Naive MXFP8 quantization.
    For each block of 32 elements along the last dimension, compute a
-    shared e8m0 scale (the biased exponent of the block-wise amax)
-    and quantize each element to float8_e4m3fn.
+    shared e8m0 scale that fits the block-wise amax into the finite
+    float8_e4m3fn range, and quantize each element to float8_e4m3fn.

    Returns (quantized_values [same shape, fp8], scales uint8).
    Scale shape depends on is_sf_swizzled_layout:
@@ -58,7 +58,8 @@ def _mxfp8_e4m3_quantize_torch(

    amax = x_blocked.abs().amax(dim=-1)
    amax = amax.clamp(min=torch.finfo(torch.float32).tiny)
-    scale_biased = torch.floor(torch.log2(amax)) + 127.0
+    fp8_max = torch.finfo(MXFP8_VALUE_DTYPE).max
+    scale_biased = torch.ceil(torch.log2(amax / fp8_max)) + 127.0
    scale_biased = scale_biased.clamp(0, 254)
    scales_uint8 = scale_biased.to(torch.uint8)

@@ -232,6 +232,7 @@ class VocabParallelEmbedding(PluggableLayer):
        padding_size: padding size for the vocabulary.
        quant_config: quant config for the layer
        prefix: full name of the layer in the state dict
+        replicated: whether to replicate the embedding on every TP rank.
    """  # noqa: E501

    # --8<-- [end:vocab_parallel_embedding]
@@ -245,12 +246,15 @@ class VocabParallelEmbedding(PluggableLayer):
        padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
+        *,
+        replicated: bool = False,
    ):
        super().__init__()

        # Keep the input dimensions.
-        tp_rank = get_tensor_model_parallel_rank()
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.replicated = replicated
+        tp_rank = get_tensor_model_parallel_rank() if not replicated else 0
+        self.tp_size = get_tensor_model_parallel_world_size() if not replicated else 1
        self.num_embeddings = num_embeddings
        self.padding_size = padding_size
        self.org_vocab_size = org_num_embeddings or num_embeddings
@@ -487,8 +491,12 @@ class VocabParallelEmbedding(PluggableLayer):
        # Mask the output embedding.
        if self.tp_size > 1:
            output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
-        # Reduce across all the model parallel GPUs.
-        output = tensor_model_parallel_all_reduce(output_parallel)
+        # Reduce sharded embeddings across all the model parallel GPUs.
+        output = (
+            output_parallel
+            if self.replicated
+            else tensor_model_parallel_all_reduce(output_parallel)
+        )
        return output

    def extra_repr(self) -> str:
@@ -516,6 +524,7 @@ class ParallelLMHead(VocabParallelEmbedding):
        params_dtype: type of the parameters.
        org_num_embeddings: original vocabulary size (without LoRA).
        padding_size: padding size for the vocabulary.
+        replicated: whether to replicate the head on every TP rank.
    """

    # --8<-- [end:parallel_lm_head]
@@ -530,6 +539,8 @@ class ParallelLMHead(VocabParallelEmbedding):
        padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
+        *,
+        replicated: bool = False,
    ):
        super().__init__(
            num_embeddings,
@@ -539,6 +550,7 @@ class ParallelLMHead(VocabParallelEmbedding):
            padding_size,
            quant_config,
            prefix,
+            replicated=replicated,
        )
        self.quant_config = quant_config
        if bias:
@@ -16,10 +16,9 @@ import vllm.envs as envs
 from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import (
-    Attention,
-    MLAAttention,
    MMEncoderAttention,
 )
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.hpc import HpcModule
 from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig,
@@ -122,12 +121,13 @@ def process_weights_after_loading(
            # the caching allocator, which starves the OS on UMA devices.
            release_device_memory_under_pressure(target_device)

-    # Initialize post-load attention weights for Attention, MLA, and MM encoder.
-    # NOTE: Happens after other modules so we can easily decompress weights.
+    # Initialize post-load attention weights for any attention layer and MM
+    # encoder. NOTE: Happens after other modules so we can easily decompress
+    # weights.
    for _, module in model.named_modules():
-        if isinstance(
-            module, (Attention, MLAAttention, MMEncoderAttention)
-        ) and hasattr(module, "process_weights_after_loading"):
+        if isinstance(module, (AttentionLayerBase, MMEncoderAttention)) and hasattr(
+            module, "process_weights_after_loading"
+        ):
            # TODO(lucas): see if there is a way to unify the signatures
            # of process_weights_after_loading
            with device_loading_context(module, target_device):
@@ -339,6 +339,39 @@ class DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig):
                )


+class KimiK3ForConditionalGenerationConfig(VerifyAndUpdateConfig):
+    """Route MXFP4-checkpointed Kimi-K3 MoE experts to the MXFP4 interface.
+
+    Kimi-K3 ships its routed experts as compressed-tensors
+    ``mxfp4-pack-quantized`` (``quant_method="compressed-tensors"``), which
+    lands them on ``CompressedTensorsW4A4Mxfp4MoEMethod`` and its narrow kernel
+    selection. Rewriting ``quant_method`` to ``"mxfp4"`` selects ``Mxfp4Config``
+    (hence ``Mxfp4MoEMethod``) with its full backend set, while any non-MXFP4
+    checkpoint is left untouched. Covers both the main model and the MTP draft.
+
+    ``model_arch_config.quantization_config`` is a separate dict, snapshotted in
+    ``ModelConfig.__init__`` before this hook runs, and it is what
+    ``_verify_quantization`` reads when resolving the quant method. Patch it
+    alongside the hf configs so the rewrite lands before resolution; otherwise
+    the main model still resolves to compressed-tensors.
+    """
+
+    @staticmethod
+    def verify_and_update_model_config(model_config: "ModelConfig") -> None:
+        for cfg in (
+            model_config.hf_config,
+            model_config.hf_text_config,
+            model_config.model_arch_config,
+        ):
+            quant_config = getattr(cfg, "quantization_config", None)
+            if (
+                isinstance(quant_config, dict)
+                and quant_config.get("quant_method") == "compressed-tensors"
+                and quant_config.get("format") == "mxfp4-pack-quantized"
+            ):
+                quant_config["quant_method"] = "mxfp4"
+
+
 class GptOssForCausalLMConfig(VerifyAndUpdateConfig):
    @staticmethod
    def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -768,20 +801,6 @@ class Qwen3_5ForConditionalGenerationConfig(VerifyAndUpdateConfig):
            )


-class Qwen3_5ForCausalLMConfig(Qwen3_5ForConditionalGenerationConfig):
-    @staticmethod
-    def verify_and_update_config(vllm_config: "VllmConfig") -> None:
-        Qwen3_5ForConditionalGenerationConfig.verify_and_update_config(vllm_config)
-
-        # Text-only Qwen3.5 models use one-dimensional positions. Remove the
-        # M-RoPE fields inherited from the multimodal configuration.
-        hf_text_config = vllm_config.model_config.hf_text_config
-        rope_parameters = getattr(hf_text_config, "rope_parameters", None)
-        if rope_parameters is not None:
-            rope_parameters.pop("mrope_section", None)
-            rope_parameters.pop("mrope_interleaved", None)
-
-
 class ColQwen3_5Config(Qwen3_5ForConditionalGenerationConfig):
    """Apply the attention contract declared by a ColQwen3.5 checkpoint."""

@@ -884,6 +903,8 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
    "JambaForSequenceClassification": JambaForSequenceClassificationConfig,
    "JinaForRanking": JinaForRankingConfig,
    "JinaVLForRanking": JinaVLForSequenceClassificationConfig,
+    "KimiK3ForConditionalGeneration": KimiK3ForConditionalGenerationConfig,
+    "KimiK3MTPModel": KimiK3ForConditionalGenerationConfig,
    "LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig,
    "LlamaBidirectionalModel": LlamaBidirectionalConfig,
    "LlamaNemotronVLForSequenceClassification": LlamaNemotronVLConfig,
@@ -898,9 +919,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
    "Qwen2ForRewardModel": Qwen2ForRewardModelConfig,
    "Qwen3ForSequenceClassification": Qwen3ForSequenceClassificationConfig,
    "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig,
-    "Qwen3_5ForCausalLM": Qwen3_5ForCausalLMConfig,
    "Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,
-    "Qwen3_5MoeForCausalLM": Qwen3_5ForCausalLMConfig,
    "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,
    "UnlimitedOCRForCausalLM": UnlimitedOCRForCausalLMConfig,
    "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig,
@@ -983,6 +983,7 @@ class DeepseekV2MLAAttention(nn.Module):
        topk_indices_buffer: torch.Tensor | None = None,
        input_size: int | None = None,
        reduce_results: bool = True,
+        non_causal_multi_token_decode: bool = False,
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
@@ -1178,6 +1179,7 @@ class DeepseekV2MLAAttention(nn.Module):
            # the V1 proposer. A frozen True would leave the draft reading a
            # never-written topk buffer.
            skip_topk=_skip_topk and not is_mtp_layer,
+            non_causal_multi_token_decode=non_causal_multi_token_decode,
            # Do not skip scoring for MTP layers: their top-k buffer may be
            # reused by later draft iterations through index sharing.
            allow_short_prefill_indexer_scoring_skip=not is_mtp_layer,
@@ -1384,7 +1384,7 @@ class EagleModelMixin:
        aux_hidden_states: list[torch.Tensor],
        layer_idx: int,
        hidden_states: torch.Tensor,
-        residual: torch.Tensor,
+        residual: torch.Tensor | None,
    ) -> list[torch.Tensor]:
        if layer_idx in self.aux_hidden_state_layers:
            value = hidden_states + residual if residual is not None else hidden_states
@@ -34,6 +34,7 @@ from vllm.model_executor.models.vision import (
    is_vit_use_data_parallel,
    run_dp_sharded_mrope_vision_model,
 )
+from vllm.platforms import current_platform
 from vllm.transformers_utils.configs.kimi_k25 import KimiK25VisionConfig

 logger = init_logger(__name__)
@@ -154,9 +155,7 @@ class Learnable2DInterpPosEmbDivided_fixed(nn.Module):
    def reset_parameters(self):
        nn.init.normal_(self.weight)

-    def forward(
-        self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]]
-    ) -> torch.Tensor:
+    def get_pos_embeds(self, grid_thws: torch.Tensor | list[list[int]]) -> torch.Tensor:
        pos_embs = []
        grid_thw_list = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist()
        for t, h, w in grid_thw_list:
@@ -179,8 +178,12 @@ class Learnable2DInterpPosEmbDivided_fixed(nn.Module):

            pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1]))

-        out = x + torch.cat(pos_embs)
-        return out
+        return torch.cat(pos_embs)
+
+    def forward(
+        self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]]
+    ) -> torch.Tensor:
+        return x + self.get_pos_embeds(grid_thws)


 class MoonVision3dPatchEmbed(nn.Module):
@@ -195,6 +198,8 @@ class MoonVision3dPatchEmbed(nn.Module):
        pos_emb_width: int = 14,
        pos_emb_time: int = 4,
        pos_emb_type: str = "divided_fixed",
+        patch_embed_proj_bias: bool = True,
+        pos_emb_interpolation_mode: str = "bicubic",
    ):
        super().__init__()
        assert isinstance(patch_size, int | Sequence), (
@@ -208,7 +213,11 @@ class MoonVision3dPatchEmbed(nn.Module):
        self.patch_size = patch_size

        self.proj = nn.Conv2d(
-            in_dim, out_dim, kernel_size=patch_size, stride=patch_size
+            in_dim,
+            out_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=patch_embed_proj_bias,
        )

        if pos_emb_type == "divided_fixed":
@@ -217,17 +226,37 @@ class MoonVision3dPatchEmbed(nn.Module):
                width=pos_emb_width,
                num_frames=pos_emb_time,
                dim=out_dim,
+                interpolation_mode=pos_emb_interpolation_mode,
            )
        else:
            raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}")

    def forward(
-        self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]]
+        self,
+        x: torch.Tensor,
+        grid_thws: torch.Tensor | list[list[int]] | None,
+        *,
+        pos_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor:
-        x = self.proj(x).view(x.size(0), -1)
-        # apply positional embedding
-        x = self.pos_emb(x, grid_thws)
-        return x
+        x = self._proj(x).view(x.size(0), -1)
+        if pos_embeds is not None:
+            return x + pos_embeds
+        assert grid_thws is not None
+        return self.pos_emb(x, grid_thws)
+
+    def _proj(self, x: torch.Tensor) -> torch.Tensor:
+        # MIOpen conv2d intermittently fails under load on ROCm; use aiter Triton.
+        if current_platform.is_rocm() and x.dtype in (torch.float16, torch.bfloat16):
+            from aiter.ops.triton.conv.conv2d import conv2d
+
+            return conv2d(
+                x,
+                self.proj.weight,
+                self.proj.bias,
+                stride=self.patch_size,
+                layout="nchw",
+            )
+        return self.proj(x)


 class Rope2DPosEmbRepeated(nn.Module):
@@ -342,6 +371,14 @@ class MLP2(nn.Module):
        return x


+def _make_vision_norm(norm_type: str, hidden_dim: int) -> nn.Module:
+    if norm_type == "layernorm":
+        return nn.LayerNorm(hidden_dim)
+    if norm_type == "rmsnorm":
+        return nn.RMSNorm(hidden_dim)
+    raise NotImplementedError(f"Not support norm_type: {norm_type}")
+
+
 class MoonViTEncoderLayer(nn.Module):
    """Single encoder layer for MoonViT with TP/DP support."""

@@ -355,23 +392,33 @@ class MoonViTEncoderLayer(nn.Module):
        *,
        activation=F.gelu,
        attn_bias: bool = False,
+        qkv_hidden_size: int | None = None,
+        norm_type: str = "layernorm",
+        mlp_type: str = "mlp2",
+        linear_bias: bool = True,
    ):
        super().__init__()
        self.use_data_parallel = is_vit_use_data_parallel(num_heads)

        self.num_heads = num_heads
        self.hidden_dim = hidden_dim
-        self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads
+        self.qkv_hidden_size = (
+            hidden_dim if qkv_hidden_size is None else qkv_hidden_size
+        )
+        self.hidden_size_per_attention_head = self.qkv_hidden_size // self.num_heads
        self.tp_size = (
            1 if self.use_data_parallel else get_tensor_model_parallel_world_size()
        )
        self.num_attention_heads_per_partition = divide(num_heads, self.tp_size)

-        self.norm0 = nn.LayerNorm(hidden_dim)
-        self.norm1 = nn.LayerNorm(hidden_dim)
+        self.norm0 = _make_vision_norm(norm_type, hidden_dim)
+        self.norm1 = _make_vision_norm(norm_type, hidden_dim)
+        if mlp_type != "mlp2":
+            raise NotImplementedError(f"Not support mlp_type: {mlp_type}")
        self.mlp = MLP2(
            [hidden_dim, mlp_dim, hidden_dim],
            activation,
+            bias=linear_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.mlp",
            use_data_parallel=self.use_data_parallel,
@@ -387,7 +434,7 @@ class MoonViTEncoderLayer(nn.Module):
            disable_tp=self.use_data_parallel,
        )
        self.wo = RowParallelLinear(
-            hidden_dim,
+            self.qkv_hidden_size,
            hidden_dim,
            bias=attn_bias,
            quant_config=quant_config,
@@ -493,8 +540,9 @@ class MoonViT3dEncoder(nn.Module):
            f'video_attn_type must be "spatial_temporal", got {video_attn_type}'
        )
        self.video_attn_type = video_attn_type
+        qkv_hidden_size = block_cfg.get("qkv_hidden_size") or block_cfg["hidden_dim"]
        self.rope_2d = Rope2DPosEmbRepeated(
-            block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512
+            qkv_hidden_size // block_cfg["num_heads"], 512, 512
        )
        self.blocks = nn.ModuleList(
            [
@@ -506,13 +554,17 @@ class MoonViT3dEncoder(nn.Module):
                for layer_idx in range(num_layers)
            ]
        )
-        self.final_layernorm = nn.LayerNorm(hidden_dim)
+        self.final_layernorm = _make_vision_norm(
+            block_cfg.get("norm_type", "layernorm"), hidden_dim
+        )

    def prepare_encoder_metadata(
        self,
        grid_thw_list: list[list[int]],
        *,
        device: torch.device,
+        max_batch_size: int | None = None,
+        max_seqlen_override: int | None = None,
    ) -> dict[str, torch.Tensor | None]:
        metadata: dict[str, torch.Tensor | None] = {}
        metadata["rope_freqs_cis"] = self.rope_2d.get_freqs_cis(
@@ -524,15 +576,30 @@ class MoonViT3dEncoder(nn.Module):
        cu_seqlens = np.concatenate(
            [np.zeros(1, dtype=np.int32), lengths.cumsum(dtype=np.int32)]
        )
+        if max_batch_size is not None:
+            num_seqs = len(cu_seqlens) - 1
+            if num_seqs < max_batch_size:
+                cu_seqlens = np.concatenate(
+                    [
+                        cu_seqlens,
+                        np.full(
+                            max_batch_size - num_seqs,
+                            cu_seqlens[-1],
+                            dtype=np.int32,
+                        ),
+                    ]
+                )

        attn_backend = self.blocks[0].attn.attn_backend
        metadata["sequence_lengths"] = MMEncoderAttention.maybe_compute_seq_lens(
            attn_backend, cu_seqlens, device
        )
-        metadata["max_seqlen"] = torch.tensor(
-            MMEncoderAttention.compute_max_seqlen(attn_backend, cu_seqlens),
-            dtype=torch.int32,
+        max_seqlen = (
+            max_seqlen_override
+            if max_seqlen_override is not None
+            else MMEncoderAttention.compute_max_seqlen(attn_backend, cu_seqlens)
        )
+        metadata["max_seqlen"] = torch.tensor(max_seqlen, dtype=torch.int32)
        metadata["cu_seqlens"] = MMEncoderAttention.maybe_recompute_cu_seqlens(
            attn_backend,
            cu_seqlens,
@@ -545,11 +612,12 @@ class MoonViT3dEncoder(nn.Module):
    def forward(
        self,
        hidden_states: torch.Tensor,
-        grid_thws: torch.Tensor | list[list[int]],
+        grid_thws: torch.Tensor | list[list[int]] | None,
        *,
        encoder_metadata: dict[str, torch.Tensor | None] | None = None,
    ) -> torch.Tensor:
        if encoder_metadata is None:
+            assert grid_thws is not None
            grid_thw_list = (
                grid_thws if isinstance(grid_thws, list) else grid_thws.tolist()
            )
@@ -604,6 +672,34 @@ def tpool_patch_merger(
    return outputs


+def build_image_merge_gather_idx(
+    grid_thws: list[list[int]] | list[tuple[int, int, int]],
+    merge_kernel_size: tuple[int, int],
+) -> np.ndarray:
+    """Build packed spatial-merge indices for image-only CUDA graphs."""
+    kh, kw = merge_kernel_size
+    parts: list[np.ndarray] = []
+    offset = 0
+    for t, h, w in grid_thws:
+        if t != 1:
+            raise ValueError("Image encoder CUDA graphs require grid T == 1")
+        idx = np.arange(h * w, dtype=np.int64).reshape(h, w)
+        idx = idx.reshape(h // kh, kh, w // kw, kw)
+        parts.append(idx.transpose(0, 2, 1, 3).reshape(-1, kh * kw) + offset)
+        offset += h * w
+    if not parts:
+        return np.empty((0, kh * kw), dtype=np.int64)
+    return np.concatenate(parts)
+
+
+def tpool_patch_merger_packed(
+    x: torch.Tensor,
+    merge_gather_idx: torch.Tensor,
+) -> torch.Tensor:
+    """Apply the image-only spatial merge using precomputed tensor indices."""
+    return x[merge_gather_idx]
+
+
 class MoonViT3dPretrainedModel(nn.Module):
    """Main vision tower model.

@@ -630,6 +726,10 @@ class MoonViT3dPretrainedModel(nn.Module):
            pos_emb_width=config.init_pos_emb_width,
            pos_emb_time=config.init_pos_emb_time,
            pos_emb_type=config.pos_emb_type,
+            patch_embed_proj_bias=getattr(config, "patch_embed_proj_bias", True),
+            pos_emb_interpolation_mode=getattr(
+                config, "pos_emb_interpolation_mode", "bicubic"
+            ),
        )

        self.encoder = MoonViT3dEncoder(
@@ -638,9 +738,15 @@ class MoonViT3dPretrainedModel(nn.Module):
            block_cfg={
                "num_heads": config.num_attention_heads,
                "hidden_dim": config.hidden_size,
+                "qkv_hidden_size": getattr(config, "qkv_hidden_size", None),
                "mlp_dim": config.intermediate_size,
-                "activation": get_act_fn("gelu_pytorch_tanh"),
-                "attn_bias": True,
+                "activation": get_act_fn(
+                    getattr(config, "activation_func", "gelu_pytorch_tanh")
+                ),
+                "attn_bias": getattr(config, "attn_bias", True),
+                "norm_type": getattr(config, "norm_type", "layernorm"),
+                "mlp_type": getattr(config, "mlp_type", "mlp2"),
+                "linear_bias": getattr(config, "linear_bias", True),
            },
            video_attn_type=config.video_attn_type,
            quant_config=quant_config,
@@ -650,7 +756,7 @@ class MoonViT3dPretrainedModel(nn.Module):
    def forward(
        self,
        pixel_values: torch.Tensor,
-        grid_thws: torch.Tensor | list[list[int]],
+        grid_thws: torch.Tensor | list[list[int]] | None,
        *,
        encoder_metadata: dict[str, torch.Tensor | None] | None = None,
    ) -> torch.Tensor:
@@ -662,6 +768,22 @@ class MoonViT3dPretrainedModel(nn.Module):
        Returns:
            torch.Tensor: The output tokens.
        """
+        if encoder_metadata is not None and "pos_embeds" in encoder_metadata:
+            hidden_states = self.patch_embed(
+                pixel_values,
+                None,
+                pos_embeds=encoder_metadata["pos_embeds"],
+            )
+            hidden_states = self.encoder(
+                hidden_states,
+                None,
+                encoder_metadata=encoder_metadata,
+            )
+            merge_gather_idx = encoder_metadata["merge_gather_idx"]
+            assert merge_gather_idx is not None
+            return tpool_patch_merger_packed(hidden_states, merge_gather_idx)
+
+        assert grid_thws is not None
        grid_thw_list = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist()
        if encoder_metadata is None:
            encoder_metadata = self.encoder.prepare_encoder_metadata(
@@ -685,13 +807,40 @@ class MoonViT3dPretrainedModel(nn.Module):

        return hidden_states

+    def prepare_encoder_cudagraph_metadata(
+        self,
+        grid_thw_list: list[list[int]],
+        *,
+        max_batch_size: int,
+        max_seqlen_override: int | None = None,
+        device: torch.device,
+    ) -> dict[str, torch.Tensor | None]:
+        """Precompute fixed-buffer metadata for image encoder CUDA graphs."""
+        grid_thw_list = [list(map(int, grid)) for grid in grid_thw_list]
+        metadata = self.encoder.prepare_encoder_metadata(
+            grid_thw_list,
+            device=device,
+            max_batch_size=max_batch_size,
+            max_seqlen_override=max_seqlen_override,
+        )
+        metadata["pos_embeds"] = self.patch_embed.pos_emb.get_pos_embeds(
+            grid_thw_list
+        ).to(device=device)
+        merge_gather_idx = build_image_merge_gather_idx(
+            grid_thw_list, self.merge_kernel_size
+        )
+        metadata["merge_gather_idx"] = torch.from_numpy(merge_gather_idx).to(
+            device=device, non_blocking=True
+        )
+        return metadata
+

@torch.inference_mode()
 def mm_projector_forward(mm_projector: torch.nn.Module, vt_output: list[torch.Tensor]):
    """Apply MM projector to vision tower outputs."""
    num_embedding_list = [x.shape[0] for x in vt_output]
    batched = torch.cat(vt_output, dim=0)
-    projector_dtype = mm_projector.pre_norm.weight.dtype
+    projector_dtype = next(mm_projector.parameters()).dtype
    if batched.dtype != projector_dtype:
        batched = batched.to(projector_dtype)
    proj_out = mm_projector(batched)
@@ -747,11 +896,34 @@ class KimiK25MultiModalProjector(nn.Module):
    ):
        super().__init__()
        self.use_data_parallel = use_data_parallel
+        self.mm_projector_type = getattr(config, "mm_projector_type", "patchmerger")

        # Hidden size after patch merging
        merge_h, merge_w = config.merge_kernel_size
        self.hidden_size = config.hidden_size * merge_h * merge_w

+        if self.mm_projector_type == "patchmergerv2":
+            self.linear_1 = ReplicatedLinear(
+                self.hidden_size,
+                self.hidden_size,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.linear_1",
+            )
+            self.linear_2 = ReplicatedLinear(
+                self.hidden_size,
+                getattr(config, "text_hidden_size", config.mm_hidden_size),
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.linear_2",
+            )
+            self.post_norm = torch.nn.RMSNorm(
+                getattr(config, "text_hidden_size", config.mm_hidden_size),
+                eps=config.projector_ln_eps,
+            )
+            self.act = GELUActivation()
+            return
+
        self.pre_norm = torch.nn.LayerNorm(config.hidden_size, eps=1e-5)
        self.linear_1 = ReplicatedLinear(
            self.hidden_size,
@@ -770,6 +942,13 @@ class KimiK25MultiModalProjector(nn.Module):
        self.act = GELUActivation()

    def forward(self, image_features: torch.Tensor) -> torch.Tensor:
+        if self.mm_projector_type == "patchmergerv2":
+            hidden_states = image_features.view(image_features.shape[0], -1)
+            hidden_states, _ = self.linear_1(hidden_states)
+            hidden_states = self.act(hidden_states)
+            hidden_states, _ = self.linear_2(hidden_states)
+            return self.post_norm(hidden_states)
+
        hidden_states = self.pre_norm(image_features).view(-1, self.hidden_size)
        hidden_states, _ = self.linear_1(hidden_states)
        hidden_states = self.act(hidden_states)
@@ -1,646 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from collections.abc import Iterable
-
-import torch
-from torch import nn
-
-from vllm.compilation.decorators import support_torch_compile
-from vllm.config import CacheConfig, VllmConfig
-from vllm.distributed import (
-    get_pp_group,
-    get_tensor_model_parallel_world_size,
-)
-from vllm.logger import init_logger
-from vllm.model_executor.layers.activation import SiluAndMul
-from vllm.model_executor.layers.fused_moe import (
-    FusedMoE,
-    fused_moe_make_expert_params_mapping,
-)
-from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
-    MergedColumnParallelLinear,
-    ReplicatedLinear,
-    RowParallelLinear,
-)
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import (
-    KimiGatedDeltaNetAttention,
-)
-from vllm.model_executor.layers.mamba.mamba_utils import (
-    MambaStateCopyFunc,
-    MambaStateCopyFuncCalculator,
-    MambaStateDtypeCalculator,
-    MambaStateShapeCalculator,
-)
-from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
-from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    ParallelLMHead,
-    VocabParallelEmbedding,
-)
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-    maybe_remap_kv_scale_name,
-)
-from vllm.sequence import IntermediateTensors
-from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
-
-from .interfaces import HasInnerState, IsHybrid, MixtureOfExperts, SupportsPP
-from .utils import (
-    AutoWeightsLoader,
-    PPMissingLayer,
-    get_spec_layer_idx_from_weight_name,
-    is_pp_missing_parameter,
-    make_layers,
-    maybe_prefix,
-)
-
-logger = init_logger(__name__)
-
-
-class KimiMLP(nn.Module):
-    def __init__(
-        self,
-        hidden_size: int,
-        intermediate_size: int,
-        hidden_act: str,
-        quant_config: QuantizationConfig | None = None,
-        reduce_results: bool = True,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-
-        self.gate_up_proj = MergedColumnParallelLinear(
-            hidden_size,
-            [intermediate_size] * 2,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.gate_up_proj",
-        )
-        self.down_proj = RowParallelLinear(
-            intermediate_size,
-            hidden_size,
-            bias=False,
-            quant_config=quant_config,
-            reduce_results=reduce_results,
-            prefix=f"{prefix}.down_proj",
-        )
-        if hidden_act != "silu":
-            raise ValueError(
-                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
-            )
-        self.act_fn = SiluAndMul()
-
-    def forward(self, x):
-        gate_up, _ = self.gate_up_proj(x)
-        x = self.act_fn(gate_up)
-        x, _ = self.down_proj(x)
-        return x
-
-
-class KimiMoE(nn.Module):
-    def __init__(
-        self,
-        config: KimiLinearConfig,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-        layer_idx: int = 0,
-    ):
-        super().__init__()
-        hidden_size = config.hidden_size
-        intermediate_size = config.intermediate_size
-        moe_intermediate_size = config.moe_intermediate_size
-        num_experts = config.num_experts
-        moe_renormalize = config.moe_renormalize
-        self.tp_size = get_tensor_model_parallel_world_size()
-        self.routed_scaling_factor = config.routed_scaling_factor
-        self.num_shared_experts = config.num_shared_experts
-        self.layer_idx = layer_idx
-
-        if config.hidden_act != "silu":
-            raise ValueError(
-                f"Unsupported activation: {config.hidden_act}. "
-                "Only silu is supported for now."
-            )
-
-        # Gate always runs at half / full precision for now.
-        self.gate = ReplicatedLinear(
-            hidden_size,
-            num_experts,
-            bias=False,
-            quant_config=None,
-            prefix=f"{prefix}.gate",
-        )
-
-        self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts))
-
-        if self.num_shared_experts is not None:
-            intermediate_size = moe_intermediate_size * self.num_shared_experts
-            self.shared_experts = KimiMLP(
-                hidden_size=config.hidden_size,
-                intermediate_size=intermediate_size,
-                hidden_act=config.hidden_act,
-                quant_config=quant_config,
-                reduce_results=False,
-                prefix=f"{prefix}.shared_experts",
-            )
-        else:
-            self.shared_experts = None
-
-        self.experts = FusedMoE(
-            shared_experts=self.shared_experts,
-            num_experts=num_experts,
-            top_k=config.num_experts_per_token,
-            hidden_size=hidden_size,
-            intermediate_size=moe_intermediate_size,
-            renormalize=moe_renormalize,
-            quant_config=quant_config,
-            use_grouped_topk=config.use_grouped_topk,
-            num_expert_group=config.num_expert_group,
-            topk_group=config.topk_group,
-            prefix=f"{prefix}.experts",
-            scoring_func=config.moe_router_activation_func,
-            e_score_correction_bias=self.gate.e_score_correction_bias,
-            routed_scaling_factor=self.routed_scaling_factor,
-        )
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        num_tokens, hidden_size = hidden_states.shape
-        hidden_states = hidden_states.view(-1, hidden_size)
-        router_logits, _ = self.gate(hidden_states)
-        final_hidden_states = self.experts(
-            hidden_states=hidden_states, router_logits=router_logits
-        )
-        return final_hidden_states.view(num_tokens, hidden_size)
-
-
-class KimiMLAAttention(nn.Module):
-    """
-    Main reference: DeepseekV2 vllm Implementation
-    """
-
-    def __init__(
-        self,
-        config: KimiLinearConfig,
-        hidden_size: int,
-        num_heads: int,
-        qk_nope_head_dim: int,
-        qk_rope_head_dim: int,
-        v_head_dim: int,
-        q_lora_rank: int | None,
-        kv_lora_rank: int,
-        use_nope: bool = False,
-        cache_config: CacheConfig | None = None,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-        **kwargs,
-    ) -> None:
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.qk_nope_head_dim = qk_nope_head_dim
-        self.qk_rope_head_dim = qk_rope_head_dim
-        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
-        self.v_head_dim = v_head_dim
-        self.q_lora_rank = q_lora_rank
-        self.kv_lora_rank = kv_lora_rank
-        self.num_heads = num_heads
-        tp_size = get_tensor_model_parallel_world_size()
-        self.num_local_heads = num_heads // tp_size
-        self.scaling = self.qk_head_dim**-0.5
-        self.use_nope = use_nope
-        assert self.use_nope is True
-        assert self.q_lora_rank is None
-        assert num_heads % tp_size == 0
-        self.kv_a_proj_with_mqa = ReplicatedLinear(
-            self.hidden_size,
-            self.kv_lora_rank + self.qk_rope_head_dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.kv_a_proj_with_mqa",
-        )
-        self.q_proj = ColumnParallelLinear(
-            self.hidden_size,
-            self.num_heads * self.qk_head_dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.q_proj",
-        )
-        self.kv_a_layernorm = RMSNorm(
-            self.kv_lora_rank,
-            eps=config.rms_norm_eps,
-        )
-        self.kv_b_proj = ColumnParallelLinear(
-            self.kv_lora_rank,
-            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.kv_b_proj",
-        )
-        self.o_proj = RowParallelLinear(
-            self.num_heads * self.v_head_dim,
-            self.hidden_size,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.o_proj",
-        )
-
-        mla_modules = MLAModules(
-            kv_a_layernorm=self.kv_a_layernorm,
-            kv_b_proj=self.kv_b_proj,
-            rotary_emb=None,
-            o_proj=self.o_proj,
-            fused_qkv_a_proj=None,
-            kv_a_proj_with_mqa=self.kv_a_proj_with_mqa,
-            q_a_layernorm=None,
-            q_b_proj=None,
-            q_proj=self.q_proj,
-            indexer=None,
-            is_sparse=False,
-            topk_indices_buffer=None,
-        )
-        self.mla_attn = MultiHeadLatentAttentionWrapper(
-            self.hidden_size,
-            self.num_local_heads,
-            self.scaling,
-            self.qk_nope_head_dim,
-            self.qk_rope_head_dim,
-            self.v_head_dim,
-            self.q_lora_rank,
-            self.kv_lora_rank,
-            mla_modules,
-            cache_config,
-            quant_config,
-            prefix,
-        )
-
-    def forward(
-        self,
-        positions: torch.Tensor,
-        hidden_states: torch.Tensor,
-        output: torch.Tensor,
-    ) -> None:
-        output[:] = self.mla_attn(positions, hidden_states)
-
-
-class KimiDecoderLayer(nn.Module):
-    def __init__(
-        self,
-        config: KimiLinearConfig,
-        vllm_config: VllmConfig,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.hidden_size = config.hidden_size
-
-        self.is_moe = config.is_moe
-        layer_idx = int(prefix.rsplit(".", 1)[1])
-        model_config = vllm_config.model_config
-        cache_config = vllm_config.cache_config
-        quant_config = vllm_config.quant_config
-
-        if config.is_kda_layer(layer_idx):
-            self.self_attn = KimiGatedDeltaNetAttention(
-                config,
-                vllm_config,
-                prefix=f"{prefix}.self_attn",
-            )
-        else:
-            self.self_attn = KimiMLAAttention(
-                layer_idx=layer_idx,
-                hidden_size=self.hidden_size,
-                num_heads=config.num_attention_heads,
-                quant_config=quant_config,
-                cache_config=cache_config,
-                model_config=model_config,
-                prefix=f"{prefix}.self_attn",
-                config=config,
-                qk_nope_head_dim=config.qk_nope_head_dim,
-                qk_rope_head_dim=config.qk_rope_head_dim,
-                v_head_dim=config.v_head_dim,
-                q_lora_rank=config.q_lora_rank,
-                kv_lora_rank=config.kv_lora_rank,
-                use_nope=config.mla_use_nope,
-            )
-
-        if (
-            self.is_moe
-            and config.num_experts is not None
-            and layer_idx >= config.first_k_dense_replace
-            and layer_idx % config.moe_layer_freq == 0
-        ):
-            self.block_sparse_moe = KimiMoE(
-                config=config,
-                quant_config=quant_config,
-                prefix=f"{prefix}.block_sparse_moe",
-            )
-            self.mlp = self.block_sparse_moe
-        else:
-            self.mlp = KimiMLP(
-                hidden_size=self.hidden_size,
-                intermediate_size=config.intermediate_size,
-                hidden_act=config.hidden_act,
-                quant_config=quant_config,
-                prefix=f"{prefix}.mlp",
-            )
-        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps
-        )
-
-    def forward(
-        self,
-        positions: torch.Tensor,
-        hidden_states: torch.Tensor,
-        residual: torch.Tensor | None,
-        **kwargs,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        # Self Attention
-        if residual is None:
-            residual = hidden_states
-            hidden_states = self.input_layernorm(hidden_states)
-        else:
-            hidden_states, residual = self.input_layernorm(hidden_states, residual)
-
-        attn_output = torch.empty_like(hidden_states)
-        self.self_attn(
-            hidden_states=hidden_states,
-            positions=positions,
-            output=attn_output,
-        )
-        hidden_states = attn_output
-
-        # Fully Connected
-        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
-        hidden_states = self.mlp(hidden_states)
-        return hidden_states, residual
-
-
-@support_torch_compile
-class KimiLinearModel(nn.Module):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        super().__init__()
-
-        config = vllm_config.model_config.hf_text_config
-        self.config = config
-
-        self.vocab_size = config.vocab_size
-
-        if get_pp_group().is_first_rank:
-            self.embed_tokens = VocabParallelEmbedding(
-                config.vocab_size,
-                config.hidden_size,
-                prefix=f"{prefix}.embed_tokens",
-            )
-        else:
-            self.embed_tokens = PPMissingLayer()
-
-        def get_layer(prefix: str):
-            return KimiDecoderLayer(
-                config,
-                vllm_config,
-                prefix,
-            )
-
-        self.start_layer, self.end_layer, self.layers = make_layers(
-            config.num_hidden_layers,
-            get_layer,
-            prefix=f"{prefix}.layers",
-        )
-
-        if get_pp_group().is_last_rank:
-            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        else:
-            self.norm = PPMissingLayer()
-
-        world_size = get_tensor_model_parallel_world_size()
-        assert config.num_attention_heads % world_size == 0, (
-            "num_attention_heads must be divisible by world_size"
-        )
-
-    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.embed_tokens(input_ids)
-
-    def forward(
-        self,
-        input_ids: torch.Tensor | None,
-        positions: torch.Tensor,
-        intermediate_tensors: IntermediateTensors | None,
-        inputs_embeds: torch.Tensor | None = None,
-        **kwargs,
-    ) -> torch.Tensor:
-        if get_pp_group().is_first_rank:
-            if inputs_embeds is not None:
-                hidden_states = inputs_embeds
-            else:
-                hidden_states = self.embed_input_ids(input_ids)
-            residual = None
-        else:
-            assert intermediate_tensors is not None
-            hidden_states = intermediate_tensors["hidden_states"]
-            residual = intermediate_tensors["residual"]
-
-        for _, layer in enumerate(self.layers[self.start_layer : self.end_layer]):
-            hidden_states, residual = layer(
-                positions=positions,
-                hidden_states=hidden_states,
-                residual=residual,
-            )
-
-        if not get_pp_group().is_last_rank:
-            return IntermediateTensors(
-                {"hidden_states": hidden_states, "residual": residual}
-            )
-
-        hidden_states, _ = self.norm(hidden_states, residual)
-        return hidden_states
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            (".gate_up_proj", ".gate_proj", 0),
-            (".gate_up_proj", ".up_proj", 1),
-        ]
-        if self.config.is_moe:
-            # Params for weights, fp8 weight scales, fp8 activation scales
-            # (param_name, weight_name, expert_id, shard_id)
-            expert_params_mapping = fused_moe_make_expert_params_mapping(
-                self,
-                ckpt_gate_proj_name="w1",
-                ckpt_down_proj_name="w2",
-                ckpt_up_proj_name="w3",
-                num_experts=self.config.num_experts,
-            )
-        else:
-            expert_params_mapping = []
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for args in weights:
-            name, loaded_weight = args[:2]
-            kwargs = args[2] if len(args) > 2 else {}
-            if "rotary_emb.inv_freq" in name:
-                continue
-
-            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
-            if spec_layer is not None:
-                continue  # skip spec decode layers for main model
-            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
-                # Models trained using ColossalAI may include these tensors in
-                # the checkpoint. Skip them.
-                continue
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                # We have mlp.experts[0].gate_proj in the checkpoint.
-                # Since we handle the experts below in expert_params_mapping,
-                # we need to skip here BEFORE we update the name, otherwise
-                # name will be updated to mlp.experts[0].gate_up_proj, which
-                # will then be updated below in expert_params_mapping
-                # for mlp.experts[0].gate_gate_up_proj, which breaks load.
-                if ("mlp.experts." in name) and name not in params_dict:
-                    continue
-                name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if is_pp_missing_parameter(name, self):
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                for idx, (param_name, weight_name, expert_id, shard_id) in enumerate(
-                    expert_params_mapping
-                ):
-                    if weight_name not in name:
-                        continue
-                    name = name.replace(weight_name, param_name)
-                    if is_pp_missing_parameter(name, self):
-                        continue
-                    param = params_dict[name]
-                    weight_loader = param.weight_loader
-                    weight_loader(
-                        param,
-                        loaded_weight,
-                        name,
-                        expert_id=expert_id,
-                        shard_id=shard_id,
-                    )
-                    break
-                else:
-                    # Skip loading extra bias for GPTQ models.
-                    if (
-                        name.endswith(".bias")
-                        and name not in params_dict
-                        and not self.config.is_linear_attn
-                    ):  # noqa: E501
-                        continue
-                    # Remapping the name of FP8 kv-scale.
-                    name = maybe_remap_kv_scale_name(name, params_dict)
-                    if name is None:
-                        continue
-                    if is_pp_missing_parameter(name, self):
-                        continue
-
-                    param = params_dict[name]
-                    weight_loader = getattr(
-                        param, "weight_loader", default_weight_loader
-                    )
-                    weight_loader(param, loaded_weight, **kwargs)
-            loaded_params.add(name)
-        return loaded_params
-
-
-class KimiLinearForCausalLM(
-    nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid
-):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        super().__init__()
-        self.model_config = vllm_config.model_config
-        self.vllm_config = vllm_config
-        self.config = self.model_config.hf_config
-        quant_config = vllm_config.quant_config
-        self.quant_config = quant_config
-        self.model = KimiLinearModel(
-            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
-        )
-        if get_pp_group().is_last_rank:
-            self.lm_head = ParallelLMHead(
-                self.config.vocab_size,
-                self.config.hidden_size,
-                quant_config=quant_config,
-                prefix=maybe_prefix(prefix, "lm_head"),
-            )
-        else:
-            self.lm_head = PPMissingLayer()
-        logit_scale = getattr(self.config, "logit_scale", 1.0)
-        self.logits_processor = LogitsProcessor(
-            self.config.vocab_size, scale=logit_scale
-        )
-
-    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.model.embed_input_ids(input_ids)
-
-    def forward(
-        self,
-        input_ids: torch.Tensor | None,
-        positions: torch.Tensor,
-        intermediate_tensors: IntermediateTensors | None = None,
-        inputs_embeds: torch.Tensor | None = None,
-        **kwargs,
-    ) -> torch.Tensor | IntermediateTensors:
-        hidden_states = self.model(
-            input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
-        )
-        return hidden_states
-
-    @classmethod
-    def get_mamba_state_dtype_from_config(
-        cls,
-        vllm_config: "VllmConfig",
-    ) -> tuple[torch.dtype, torch.dtype]:
-        return MambaStateDtypeCalculator.kda_state_dtype(
-            vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
-        )
-
-    @classmethod
-    def get_mamba_state_shape_from_config(
-        cls, vllm_config: "VllmConfig"
-    ) -> tuple[tuple[int, ...], tuple[int, ...]]:
-        parallel_config = vllm_config.parallel_config
-        hf_config = vllm_config.model_config.hf_config
-        tp_size = parallel_config.tensor_parallel_size
-        num_spec = (
-            vllm_config.speculative_config.num_speculative_tokens
-            if vllm_config.speculative_config
-            else 0
-        )
-        return MambaStateShapeCalculator.kda_state_shape(
-            tp_size,
-            hf_config.linear_attn_config["num_heads"],
-            hf_config.linear_attn_config["head_dim"],
-            conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"],
-            num_spec=num_spec,
-        )
-
-    @classmethod
-    def get_mamba_state_copy_func(
-        cls,
-    ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
-        return MambaStateCopyFuncCalculator.kda_state_copy_func()
-
-    def compute_logits(
-        self,
-        hidden_states: torch.Tensor,
-    ) -> torch.Tensor | None:
-        return self.logits_processor(self.lm_head, hidden_states)
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
-        )
-        return loader.load_weights(weights)
@@ -282,7 +282,6 @@ class Qwen3_5Model(Qwen3NextModel):
 class Qwen3_5ForCausalLMBase(
    nn.Module,
    HasInnerState,
-    IsHybrid,
    SupportsEagle3,
    SupportsLoRA,
    SupportsPP,
@@ -362,45 +361,6 @@ class Qwen3_5ForCausalLMBase(

        return hidden_states

-    @classmethod
-    def get_mamba_state_dtype_from_config(
-        cls,
-        vllm_config: "VllmConfig",
-    ) -> tuple[torch.dtype, torch.dtype]:
-        return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
-            vllm_config.model_config.dtype,
-            vllm_config.cache_config.mamba_cache_dtype,
-            vllm_config.cache_config.mamba_ssm_cache_dtype,
-        )
-
-    @classmethod
-    def get_mamba_state_shape_from_config(
-        cls, vllm_config: "VllmConfig"
-    ) -> tuple[tuple[int, int], tuple[int, int]]:
-        parallel_config = vllm_config.parallel_config
-        hf_config = vllm_config.model_config.hf_text_config
-        tp_size = parallel_config.tensor_parallel_size
-        num_spec = (
-            vllm_config.speculative_config.num_speculative_tokens
-            if vllm_config.speculative_config
-            else 0
-        )
-        return MambaStateShapeCalculator.gated_delta_net_state_shape(
-            tp_size,
-            hf_config.linear_num_key_heads,
-            hf_config.linear_num_value_heads,
-            hf_config.linear_key_head_dim,
-            hf_config.linear_value_head_dim,
-            hf_config.linear_conv_kernel_dim,
-            num_spec,
-        )
-
-    @classmethod
-    def get_mamba_state_copy_func(
-        cls,
-    ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
-        return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
-
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
@@ -39,7 +39,8 @@ class DSparkMarkovHead(nn.Module):
    ``markov_w1[token]`` embeds the previously sampled token (target vocab,
    ``vocab_size``); ``markov_w2`` projects it to a draft-vocab bias
    (``draft_vocab_size``) added to the base draft logits. The two sizes
-    coincide for full-vocab drafts.
+    coincide for full-vocab drafts. Set ``replicated`` to keep both weights
+    complete on every TP rank and avoid collectives in the sequential path.
    """

    def __init__(
@@ -48,14 +49,21 @@ class DSparkMarkovHead(nn.Module):
        draft_vocab_size: int,
        markov_rank: int,
        prefix: str,
+        *,
+        replicated: bool = False,
    ) -> None:
        super().__init__()
-        # TODO(ben): profile for which (if any) it makes sense to replicate or TP-shard
        self.markov_w1 = VocabParallelEmbedding(
-            vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w1")
+            vocab_size,
+            markov_rank,
+            prefix=maybe_prefix(prefix, "markov_w1"),
+            replicated=replicated,
        )
        self.markov_w2 = ParallelLMHead(
-            draft_vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w2")
+            draft_vocab_size,
+            markov_rank,
+            prefix=maybe_prefix(prefix, "markov_w2"),
+            replicated=replicated,
        )

    def embed(self, token_ids: torch.Tensor) -> torch.Tensor:
@@ -137,7 +137,10 @@ _TEXT_GENERATION_MODELS = {
    "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"),
    "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"),
    "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
-    "KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"),
+    "KimiLinearForCausalLM": (
+        "vllm.models.kimi_k3",
+        "KimiLinearForCausalLM",
+    ),
    "Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"),
    "Lfm2MoeForCausalLM": ("lfm2_moe", "Lfm2MoeForCausalLM"),
    "LagunaForCausalLM": ("laguna", "LagunaForCausalLM"),
@@ -195,8 +198,6 @@ _TEXT_GENERATION_MODELS = {
    "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
    "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
    "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
-    "Qwen3_5ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM"),
-    "Qwen3_5MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"),
    "RWForCausalLM": ("falcon", "FalconForCausalLM"),
    "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"),
    "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"),
@@ -460,6 +461,10 @@ _MULTIMODAL_MODELS = {
    ),
    "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"),
    "KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"),
+    "KimiK3ForConditionalGeneration": (
+        "vllm.models.kimi_k3",
+        "KimiK3ForConditionalGeneration",
+    ),
    "MoonshotKimiaForCausalLM": ("kimi_audio", "KimiAudioForConditionalGeneration"),
    "MossTranscribeDiarizeForConditionalGeneration": (
        "moss_transcribe_diarize",
@@ -609,6 +614,10 @@ _SPECULATIVE_DECODING_MODELS = {
    "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
    "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"),
    "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
+    "K3DSparkModel": (
+        "vllm.models.kimi_k3.nvidia.dspark_mla",
+        "K3DSparkForCausalLM",
+    ),
    "DFlashLagunaForCausalLM": ("laguna_dflash", "DFlashLagunaForCausalLM"),
    "Gemma4DSparkModel": ("gemma4_dspark", "Gemma4DSparkForCausalLM"),
    "PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
@@ -649,6 +658,7 @@ _SPECULATIVE_DECODING_MODELS = {
    "Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"),
    "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"),
    "HYV3MTPModel": ("hy_v3_mtp", "HYV3MTP"),
+    "KimiK3MTPModel": ("vllm.models.kimi_k3", "KimiK3MTP"),
    # Temporarily disabled.
    # # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.
    # "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
@@ -28,6 +28,9 @@ from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
    deepseek_v4_sparse_mla_attention_warmup,
    flashinfer_sparse_mla_decode_autotune_warmup,
 )
+from vllm.model_executor.warmup.kimi_k3_triton_warmup import (
+    kimi_k3_triton_warmup,
+)
 from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
 from vllm.model_executor.warmup.sparse_mla_triton_warmup import (
    sparse_mla_triton_warmup,
@@ -123,6 +126,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):

    # Run next so input-prep kernels JIT against pristine runner state.
    if worker.vllm_config.kernel_config.enable_jit_warmup:
+        kimi_k3_triton_warmup(worker)
        fa4_cutedsl_warmup(worker)
        sparse_mla_triton_warmup(worker)

@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Warm up Kimi-K3 Triton kernels."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+if TYPE_CHECKING:
+    from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
+    from vllm.v1.worker.gpu_worker import Worker
+
+logger = init_logger(__name__)
+
+
+def _get_kda_layer(worker: Worker) -> KimiK3DeltaAttention | None:
+    from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
+
+    compilation_config = getattr(
+        worker.model_runner,
+        "compilation_config",
+        None,
+    )
+    static_context = getattr(compilation_config, "static_forward_context", None)
+    if not isinstance(static_context, dict):
+        return None
+    return next(
+        (
+            layer
+            for layer in static_context.values()
+            if isinstance(layer, KimiK3DeltaAttention)
+        ),
+        None,
+    )
+
+
+def _warm_attn_res(worker: Worker) -> None:
+    from vllm.models.kimi_k3.nvidia.ops.attn_res import (
+        attn_res,
+        get_attn_res_triton_warmup_profiles,
+    )
+
+    config = worker.model_config.hf_text_config
+    block_size = getattr(config, "attn_res_block_size", None)
+    if block_size is None:
+        return
+
+    hidden_size = int(config.hidden_size)
+    max_blocks = (int(config.num_hidden_layers) + block_size - 1) // block_size
+    if max_blocks < 2:
+        return
+
+    dtype = worker.model_config.dtype
+    device = torch.device("cuda")
+    eps = float(config.rms_norm_eps)
+    prefix = torch.zeros((1, hidden_size), dtype=dtype, device=device)
+    delta = torch.zeros_like(prefix)
+    blocks = torch.zeros(
+        (1, max_blocks, hidden_size),
+        dtype=dtype,
+        device=device,
+    )
+    norm_weight = torch.zeros(hidden_size, dtype=dtype, device=device)
+    qk_weight = torch.zeros_like(norm_weight)
+    output_norm_weight = torch.zeros_like(norm_weight)
+
+    for (
+        num_blocks,
+        has_delta,
+        block_write_idx,
+        apply_output_norm,
+    ) in get_attn_res_triton_warmup_profiles(max_blocks):
+        attn_res(
+            prefix,
+            delta if has_delta else None,
+            blocks,
+            norm_weight,
+            qk_weight,
+            output_norm_weight if apply_output_norm else None,
+            num_blocks=num_blocks,
+            block_write_idx=block_write_idx,
+            eps=eps,
+            output_norm_eps=eps if apply_output_norm else 0.0,
+        )
+
+
+def _warm_recurrent_kda(
+    layer: KimiK3DeltaAttention,
+    input_dtype: torch.dtype,
+) -> None:
+    from vllm.models.kimi_k3.nvidia.ops.third_party.kda.fused_recurrent import (
+        fused_recurrent_kda,
+        get_fused_recurrent_kda_fwd_warmup_profiles,
+    )
+
+    num_speculative_tokens = int(layer.num_spec)
+    # fused_recurrent_kda_fwd_kernel is only used by speculative decode.
+    if num_speculative_tokens <= 0:
+        return
+
+    kv_cache = layer.kv_cache
+    if not isinstance(kv_cache, (list, tuple)) or len(kv_cache) < 2:
+        return
+    state = kv_cache[1]
+    if not isinstance(state, torch.Tensor) or not state.numel():
+        return
+
+    logger.info("Warming up Kimi-K3 speculative KDA kernels.")
+    h = int(layer.local_num_heads)
+    d = int(layer.head_dim)
+    tokens_per_sequence = num_speculative_tokens + 1
+    for num_sequences in get_fused_recurrent_kda_fwd_warmup_profiles(h):
+        num_tokens = num_sequences * tokens_per_sequence
+        packed_qkv = torch.empty(
+            (num_tokens, 3 * h * d),
+            dtype=input_dtype,
+            device=state.device,
+        )
+        q, k, v = (
+            tensor.view(1, num_tokens, h, d)
+            for tensor in packed_qkv.split(h * d, dim=-1)
+        )
+        fused_recurrent_kda(
+            q=q,
+            k=k,
+            v=v,
+            raw_g=torch.empty(
+                (1, num_tokens, h, d),
+                dtype=input_dtype,
+                device=state.device,
+            ),
+            raw_beta=torch.empty(
+                (1, num_tokens, h),
+                dtype=input_dtype,
+                device=state.device,
+            ),
+            A_log=layer.A_log,
+            dt_bias=layer.dt_bias,
+            lower_bound=layer.gate_lower_bound,
+            initial_state=state[:1],
+            cu_seqlens=torch.arange(
+                0,
+                num_tokens + 1,
+                tokens_per_sequence,
+                dtype=torch.int32,
+                device=state.device,
+            ),
+            ssm_state_indices=torch.zeros(
+                (num_sequences, tokens_per_sequence),
+                dtype=torch.int32,
+                device=state.device,
+            ),
+            num_accepted_tokens=torch.ones(
+                num_sequences,
+                dtype=torch.int32,
+                device=state.device,
+            ),
+            out=torch.empty(
+                (1, num_tokens, h, d),
+                dtype=input_dtype,
+                device=state.device,
+            ),
+        )
+
+
+@torch.inference_mode()
+def kimi_k3_triton_warmup(worker: Worker) -> None:
+    """Warm Kimi-K3 Triton kernels reachable by this server."""
+    if not current_platform.is_cuda():
+        return
+
+    layer = _get_kda_layer(worker)
+    if layer is None:
+        return
+
+    _warm_attn_res(worker)
+    _warm_recurrent_kda(layer, worker.model_config.dtype)
@@ -22,9 +22,9 @@ from vllm.model_executor.layers.linear import (
    RowParallelLinear,
 )
 from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
+from vllm.models.common.ops import fused_q_kv_rmsnorm
 from vllm.models.deepseek_v4.common.ops import (
    fused_indexer_q_rope_quant,
-    fused_q_kv_rmsnorm,
 )
 from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE

@@ -11,7 +11,6 @@ from .cache_utils import (
 from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant
 from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant
 from .fused_mtp_input_rmsnorm import fused_mtp_input_rmsnorm, mtp_shared_head_rmsnorm
-from .fused_qk_rmsnorm import fused_q_kv_rmsnorm
 from .save_partial_states import save_partial_states

 __all__ = [
@@ -23,7 +22,6 @@ __all__ = [
    "fused_indexer_q_rope_quant",
    "fused_inv_rope_fp8_quant",
    "fused_mtp_input_rmsnorm",
-    "fused_q_kv_rmsnorm",
    "mtp_shared_head_rmsnorm",
    "quantize_and_insert_k_cache",
    "save_partial_states",
@@ -1,96 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import torch
-
-from vllm.triton_utils import tl, triton
-
-
-@triton.jit
-def _fused_q_kv_rmsnorm_kernel(
-    q_ptr,
-    q_out_ptr,
-    q_weight_ptr,
-    q_in_stride,
-    q_out_stride,
-    kv_ptr,
-    kv_out_ptr,
-    kv_weight_ptr,
-    kv_in_stride,
-    kv_out_stride,
-    eps,
-    Q_SIZE: tl.constexpr,
-    KV_SIZE: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-):
-    # num_tokens goes on grid-x (max 2**31 - 1); task goes on grid-y.
-    # CUDA's grid-y/z are capped at 65535, so putting num_tokens there crashes
-    # the launch at max-num-batched-tokens >= 65536 with "invalid argument".
-    # int64: q_in_stride can be ~24K (128 heads × 192) and overflows int32
-    # past num_tokens ~87K under large chunked prefill.
-    token_idx = tl.program_id(0).to(tl.int64)
-    pid_task = tl.program_id(1)
-
-    if pid_task == 0:
-        SIZE = Q_SIZE
-        row_in = q_ptr + token_idx * q_in_stride
-        weight_ptr = q_weight_ptr
-        row_out = q_out_ptr + token_idx * q_out_stride
-    else:
-        SIZE = KV_SIZE
-        row_in = kv_ptr + token_idx * kv_in_stride
-        weight_ptr = kv_weight_ptr
-        row_out = kv_out_ptr + token_idx * kv_out_stride
-
-    # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's
-    # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which
-    # keep x, rrms, and w all in fp32 and perform a single cast at store.
-    block = tl.arange(0, BLOCK_SIZE)
-    mask = block < SIZE
-    x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32)
-    variance = tl.sum(x * x, axis=0) / SIZE
-    rrms = tl.rsqrt(variance + eps)
-    w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32)
-    y = x * rrms * w
-    tl.store(row_out + block, y.to(row_out.dtype.element_ty), mask=mask)
-
-
-def fused_q_kv_rmsnorm(
-    qr: torch.Tensor,
-    kv: torch.Tensor,
-    q_weight: torch.Tensor,
-    kv_weight: torch.Tensor,
-    eps: float,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    assert qr.ndim == 2 and kv.ndim == 2
-    assert qr.shape[0] == kv.shape[0], (
-        f"token dim mismatch: qr={qr.shape}, kv={kv.shape}"
-    )
-    assert qr.stride(-1) == 1 and kv.stride(-1) == 1
-    assert q_weight.is_contiguous() and kv_weight.is_contiguous()
-
-    q_size = qr.shape[1]
-    kv_size = kv.shape[1]
-    num_tokens = qr.shape[0]
-    qr_out = torch.empty_like(qr)
-    kv_out = torch.empty_like(kv)
-    if num_tokens == 0:
-        return qr_out, kv_out
-
-    block_size = triton.next_power_of_2(max(q_size, kv_size))
-    _fused_q_kv_rmsnorm_kernel[(num_tokens, 2)](
-        qr,
-        qr_out,
-        q_weight,
-        qr.stride(0),
-        qr_out.stride(0),
-        kv,
-        kv_out,
-        kv_weight,
-        kv.stride(0),
-        kv_out.stride(0),
-        eps,
-        Q_SIZE=q_size,
-        KV_SIZE=kv_size,
-        BLOCK_SIZE=block_size,
-    )
-    return qr_out, kv_out
@@ -3,9 +3,8 @@

 import math
 from collections.abc import Mapping, Sequence
-from typing import Any, Literal, cast
+from typing import cast

-import numpy.typing as npt
 import torch
 from transformers import BatchFeature
 from transformers.video_utils import VideoMetadata
@@ -470,39 +469,10 @@ class MiniMaxM3VLMultiModalProcessor(
        ]


-@VIDEO_LOADER_REGISTRY.register(
-    name="minimax_m3_vl",
-    video_processor="MiniMaxM3VLVideoProcessor",
-)
+# TODO(Isotr0py): Tie with MinimaxVideoProcessor
+# after https://github.com/vllm-project/vllm/pull/44126
+@VIDEO_LOADER_REGISTRY.register("minimax_m3_vl")
 class MiniMaxM3VideoBackend(VideoBackend):
-    @classmethod
-    def load_bytes(
-        cls,
-        data: bytes,
-        num_frames: int = -1,
-        fps: int = 1,
-        max_duration: int = 300,
-        frame_recovery: bool = False,
-        *,
-        backend: Literal[
-            "opencv",
-            "pyav",
-            "torchcodec",
-            "pynvvideocodec",
-            "deepstream",
-        ] = "opencv",
-        **kwargs,
-    ) -> tuple[npt.NDArray, dict[str, Any]]:
-        return super().load_bytes(
-            data,
-            num_frames=num_frames,
-            fps=fps,
-            max_duration=max_duration,
-            frame_recovery=frame_recovery,
-            backend=backend,
-            **kwargs,
-        )
-
    @classmethod
    def compute_frames_index_to_sample(
        cls,
@@ -513,6 +483,7 @@ class MiniMaxM3VideoBackend(VideoBackend):
        total_frames = source.total_frames_num
        video_fps = source.original_fps
        fps = target.fps
+
        if total_frames <= 0 or video_fps <= 0 or fps <= 0:
            return [0] if total_frames > 0 else []

@@ -532,9 +503,8 @@ class MiniMaxM3VideoBackend(VideoBackend):
                break
            indices.append(target_frame)
            prev_kept_ts = target_frame / video_fps
-        # Because HF sample_frames includes the last frame,
-        # we will use HF as the standard.
-        last_frame_idx = total_frames
+
+        last_frame_idx = total_frames - 1
        last_ts = last_frame_idx / video_fps
        if indices and indices[-1] != last_frame_idx and last_ts - prev_kept_ts > eps:
            indices.append(last_frame_idx)
@@ -880,6 +880,12 @@ class Platform:
                ),
                cache_config.block_size,
            )
+            if model_config.use_mla:
+                # TRTLLM/FlashInfer MLA decode kernels require the physical
+                # number of kernel blocks to be aligned to 128 / kernel_block_size.
+                # For hybrid MLA/Mamba models, make the manager block size a
+                # multiple of 128 so split kernel blocks keep that invariant.
+                kernel_block_alignment_size = max(kernel_block_alignment_size, 128)

        if cache_config.mamba_cache_mode == "all":
            # With prefix caching, align to mamba chunk size for kernel perf
@@ -12,7 +12,7 @@ import torch

 from vllm.triton_utils import tl, triton

-from .op import exp
+from .op import exp, log


@triton.heuristics(
@@ -87,6 +87,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
    deepseek_vl_v2="DeepseekVLV2Config",
    deepseek_v32="DeepseekV3Config",
    deepseek_v4="DeepseekV4Config",
+    k3_dspark="K3DSparkConfig",
    flex_olmo="FlexOlmoConfig",
    fireredlid="FireRedLIDConfig",
    funaudiochat="FunAudioChatConfig",
@@ -100,6 +101,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
    kimi_linear="KimiLinearConfig",
    kimi_vl="KimiVLConfig",
    kimi_k25="KimiK25Config",
+    kimi_k3="KimiK3Config",
    RefinedWeb="RWConfig",  # For tiiuae/falcon-40b(-instruct)
    RefinedWebModel="RWConfig",  # For tiiuae/falcon-7b(-instruct)
    mlp_speculator="MLPSpeculatorConfig",
@@ -124,9 +126,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
    qwen3_asr="Qwen3ASRConfig",
    qwen3_next="Qwen3NextConfig",
    qwen3_5="Qwen3_5Config",
-    qwen3_5_text="Qwen3_5TextConfig",
    qwen3_5_moe="Qwen3_5MoeConfig",
-    qwen3_5_moe_text="Qwen3_5MoeTextConfig",
    laguna="LagunaConfig",
    lfm2_moe="Lfm2MoeConfig",
    **{"unlimited-ocr": "UnlimitedOCRConfig"},
@@ -35,6 +35,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
    "DiffusionGemmaTextConfig": "vllm.transformers_utils.configs.diffusion_gemma",
    "DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2",
    "DeepseekV4Config": "vllm.transformers_utils.configs.deepseek_v4",
+    "K3DSparkConfig": "vllm.transformers_utils.configs.k3_dspark",
    "DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr",
    "EAGLEConfig": "vllm.transformers_utils.configs.eagle",
    "FireRedLIDConfig": "vllm.transformers_utils.configs.fireredlid",
@@ -72,6 +73,8 @@ _CLASS_TO_MODULE: dict[str, str] = {
    "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear",
    "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl",
    "KimiK25Config": "vllm.transformers_utils.configs.kimi_k25",
+    "KimiK3Config": "vllm.transformers_utils.configs.kimi_k3",
+    "KimiK3VisionConfig": "vllm.transformers_utils.configs.kimi_k3",
    "NemotronConfig": "vllm.transformers_utils.configs.nemotron",
    "NemotronHConfig": "vllm.transformers_utils.configs.nemotron_h",
    "OlmoHybridConfig": "vllm.transformers_utils.configs.olmo_hybrid",
@@ -124,6 +127,7 @@ __all__ = [
    "DeepseekVLV2Config",
    "DeepseekV3Config",
    "DeepseekV4Config",
+    "K3DSparkConfig",
    "DotsOCRConfig",
    "EAGLEConfig",
    "FlexOlmoConfig",
@@ -156,6 +160,8 @@ __all__ = [
    "KimiLinearConfig",
    "KimiVLConfig",
    "KimiK25Config",
+    "KimiK3Config",
+    "KimiK3VisionConfig",
    "NemotronConfig",
    "NemotronHConfig",
    "OlmoHybridConfig",
@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from transformers import DeepseekV2Config
+
+
+class K3DSparkConfig(DeepseekV2Config):
+    """Configuration for a dense MLA DSpark draft model."""
+
+    model_type = "k3_dspark"
+    has_no_defaults_at_init = True
+
+    def __init__(
+        self,
+        mla_use_nope: bool = False,
+        mla_use_output_gate: bool = False,
+        mla_use_qk_norm: bool = False,
+        rope_theta: float = 50000.0,
+        **kwargs,
+    ) -> None:
+        # DeepseekV2Config defaults to a MoE topology. Zero these fields so
+        # generic vLLM config logic also recognizes this draft as dense.
+        kwargs.setdefault("n_routed_experts", 0)
+        kwargs.setdefault("n_shared_experts", 0)
+        kwargs.setdefault("num_experts_per_tok", 0)
+
+        rope_parameters = kwargs.get("rope_parameters")
+        if rope_parameters is None:
+            kwargs["rope_parameters"] = {
+                "rope_type": "default",
+                "rope_theta": rope_theta,
+            }
+        else:
+            rope_parameters = dict(rope_parameters)
+            rope_parameters.setdefault("rope_type", "default")
+            rope_parameters.setdefault("rope_theta", rope_theta)
+            kwargs["rope_parameters"] = rope_parameters
+
+        super().__init__(**kwargs)
+        self.mla_use_nope = mla_use_nope
+        self.mla_use_output_gate = mla_use_output_gate
+        self.mla_use_qk_norm = mla_use_qk_norm
+
+        unsupported = [
+            name
+            for name in (
+                "mla_use_nope",
+                "mla_use_output_gate",
+                "mla_use_qk_norm",
+                "dspark_bonus_anchor",
+            )
+            if getattr(self, name, False)
+        ]
+        if self.q_lora_rank is None:
+            unsupported.append("q_lora_rank=None")
+        if unsupported:
+            raise ValueError(
+                "MLA DSpark does not support " + ", ".join(unsupported) + "."
+            )
+
+        self.draft_vocab_size = (
+            getattr(self, "draft_vocab_size", None) or self.vocab_size
+        )
+        if self.draft_vocab_size != self.vocab_size:
+            raise ValueError(
+                "MLA DSpark requires draft_vocab_size to equal vocab_size when "
+                "sharing the target embedding and LM head."
+            )
+
+        target_layer_ids = getattr(self, "target_layer_ids", None)
+        if not target_layer_ids or getattr(self, "num_target_layers", None) != len(
+            target_layer_ids
+        ):
+            raise ValueError(
+                "MLA DSpark requires non-empty target_layer_ids and a matching "
+                "num_target_layers."
+            )
@@ -3,10 +3,6 @@

 from transformers.configuration_utils import PretrainedConfig

-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
-

 class KimiLinearConfig(PretrainedConfig):
    model_type = "kimi_linear"
@@ -49,8 +45,16 @@ class KimiLinearConfig(PretrainedConfig):
        qk_rope_head_dim: int | None = None,
        v_head_dim: int | None = None,
        mla_use_nope: bool | None = False,
+        mla_use_output_gate: bool | None = False,
        num_nextn_predict_layers: int = 0,
        linear_attn_config: dict | None = None,
+        attn_res_block_size: int | None = None,
+        latent_moe_use_norm: bool = False,
+        activation_situ_beta: float | None = None,
+        activation_situ_linear_beta: float | None = None,
+        max_position_embeddings: int = 4096,
+        routed_expert_hidden_size: int | None = None,
+        topk_method: str = "noaux_tc",
        **kwargs,
    ):
        self.model_type = model_type
@@ -86,6 +90,7 @@ class KimiLinearConfig(PretrainedConfig):
        self.qk_rope_head_dim = qk_rope_head_dim
        self.v_head_dim = v_head_dim
        self.mla_use_nope = mla_use_nope
+        self.mla_use_output_gate = mla_use_output_gate
        # moe config
        self.num_experts = num_experts
        self.num_experts_per_token = num_experts_per_token
@@ -102,6 +107,14 @@ class KimiLinearConfig(PretrainedConfig):
        self.topk_group = topk_group
        self.num_nextn_predict_layers = num_nextn_predict_layers

+        self.attn_res_block_size = attn_res_block_size
+        self.latent_moe_use_norm = latent_moe_use_norm
+        self.activation_situ_beta = activation_situ_beta
+        self.activation_situ_linear_beta = activation_situ_linear_beta
+        self.max_position_embeddings = max_position_embeddings
+        self.routed_expert_hidden_size = routed_expert_hidden_size
+        self.topk_method = topk_method
+
        if linear_attn_config is not None:
            assert linear_attn_config["kda_layers"] is not None
            assert linear_attn_config["full_attn_layers"] is not None
@@ -267,6 +267,7 @@ class ModelArchConfigConvertorBase:
            "deepseek_v32",
            "deepseek_v4",
            "deepseek_mtp",
+            "k3_dspark",
            "glm_moe_dsa",
            "glm4_moe_lite",
            "glm4_moe_lite_mtp",
@@ -27,6 +27,7 @@ __all__ = [
    "IsaacProcessor",
    "KimiAudioProcessor",
    "KimiK25Processor",
+    "KimiK3Processor",
    "MiMoOmniProcessor",
    "MiniCPMOProcessor",
    "MiniCPMVProcessor",
@@ -65,6 +66,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
    "IsaacProcessor": "vllm.transformers_utils.processors.isaac",
    "KimiAudioProcessor": "vllm.transformers_utils.processors.kimi_audio",
    "KimiK25Processor": "vllm.transformers_utils.processors.kimi_k25",
+    "KimiK3Processor": "vllm.transformers_utils.processors.kimi_k3",
    "MiMoOmniProcessor": "vllm.transformers_utils.processors.mimo_v2_omni",
    "MiniCPMOProcessor": "vllm.transformers_utils.processors.minicpmo",
    "MiniCPMVProcessor": "vllm.transformers_utils.processors.minicpmv",
@@ -27,8 +27,9 @@ def maybe_execute_in_parallel(
    """Run two functions potentially in parallel on separate CUDA streams.

    When aux_stream is provided, fn0 runs on the current (default) stream and
-    fn1 runs on aux_stream, synchronized via CUDA events.  When aux_stream is
-    None, both functions execute sequentially on the current stream.
+    fn1 runs on aux_stream, synchronized via CUDA events. When aux_stream is
+    None or a breakable CUDA graph capture is active, both functions execute
+    sequentially on the current stream.

    This design follows TensorRT-LLM's maybe_execute_in_parallel pattern
    (tensorrt_llm/_torch/modules/multi_stream_utils.py).
@@ -39,11 +40,18 @@ def maybe_execute_in_parallel(
        event0: CUDA event recorded before fn0 so aux_stream can wait.
        event1: CUDA event recorded after fn1 so default stream can wait.
        aux_stream: The second CUDA stream for fn1.
-            Multi-stream is disabled when aux_stream is None.
+            Multi-stream is disabled when aux_stream is None or a breakable
+            CUDA graph capture is active.

    Returns:
        Tuple of (fn0_result, fn1_result).
    """
+    if aux_stream is not None:
+        from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+        if BreakableCUDAGraphCapture.is_active():
+            aux_stream = None
+
    if aux_stream is not None:
        event0.record()
        result0 = fn0()
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Isotr0py	6d04bdf3e3	disable k3dspark test Signed-off-by: Isotr0py <Isotr0py@outlook.com>	2026-07-29 15:44:17 +00:00
Isotr0py	f4161ac80f	fix k3dspark test Signed-off-by: Isotr0py <Isotr0py@outlook.com>	2026-07-29 13:34:36 +00:00
Isotr0py	0e56986060	fix k3dspark registry Signed-off-by: Isotr0py <Isotr0py@outlook.com>	2026-07-29 13:26:41 +00:00
Jee Jee Li	7ace0c2987	fix mla testing Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>	2026-07-29 11:39:37 +00:00
Jee Jee Li	f4e1889b76	fix mla testing Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>	2026-07-29 11:33:54 +00:00
Isotr0py	342266bcca	fix pre-commit Signed-off-by: Isotr0py <Isotr0py@outlook.com>	2026-07-29 11:15:29 +00:00
Isotr0py	69e3307048	lazy import Signed-off-by: Isotr0py <Isotr0py@outlook.com>	2026-07-29 11:03:40 +00:00
Jee Jee Li	6d0ddf8d13	fix apache-tvm-ffi Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>	2026-07-29 11:02:43 +00:00
zjy0516	f88a24e62e	remove fla Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-07-29 09:58:05 +00:00
zjy0516	5f0a79c58c	fix is_private Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-07-29 09:03:59 +00:00
Bugen Zhao	2e9fbde1e4	cleanup undesired changes Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-29 08:34:40 +00:00
zjy0516andCodex	1d1786f93a	Merge branch 'main' into kimi-k3 Co-authored-by: Codex <codex@openai.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-07-29 08:14:32 +00:00
zjy0516andCodex	947ed2b899	Restore default SITU beta for Marlin MoE Co-authored-by: Codex <codex@openai.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-07-29 06:58:15 +00:00
zjy0516andCodex	22fb22a394	Merge branch 'main' into kimi-k3 Co-authored-by: Codex <codex@openai.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-07-29 06:41:19 +00:00
zjy0516andCodex	0498dc7ea9	Fix pre-commit checks after main merge Co-authored-by: Codex <codex@openai.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-07-28 04:45:59 +00:00
zjy0516	faddf3b16f	Merge branch 'main' into kimi-k3 Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> # Conflicts: # docs/design/attention_backends.md # vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py # vllm/v1/worker/gpu/cudagraph_utils.py # vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py	2026-07-28 04:09:25 +00:00
Bugen Zhao	6dc76a9ade	fix rust unit test Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-28 02:07:08 +00:00
zjy0516	658f2f56e5	update dependency Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-07-27 16:46:04 +00:00
Bugen Zhao	b4b6407ac6	update parser behavior change Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-27 16:36:49 +00:00
zjy0516	9ea963c33e	update .dockerignore Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-07-27 16:26:12 +00:00
Bugen Zhao	3b226f1e61	revert unnecessary frontend changes to minimize diff Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-27 16:20:34 +00:00
zjy0516	fb704bae60	update test Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-07-27 16:06:26 +00:00
zjy0516	e91272629c	update test Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-07-27 16:02:08 +00:00
zjy0516	31776d0c97	update dependency Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-07-27 15:45:55 +00:00
Isotr0py	a43ab53c88	update cargo.toml to public llm-multimodal Signed-off-by: Isotr0py <Isotr0py@outlook.com>	2026-07-27 15:26:56 +00:00
+11 zjy0516 Isotr0py Thien Tran Bugen Zhao Yongye Zhu Ziming Huang Roger Wang Isotr0py aoshen02 Woosuk Kwon Jee Jee Li aoshen02 Summer Yang Kevin H. Luu Bowen Wang gnovack Nick Hill Yifan Qiao xiaozhoupy Roy Wang Jeff Ma	f68f4fddea	kimi-k3 Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Ziming Huang <zelda.huanghuang@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com> Co-authored-by: Summer Yang <girasoleyang@gmail.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com> Co-authored-by: khluu <khluu000@gmail.com> Co-authored-by: Bowen Wang <abmfy@icloud.com> Co-authored-by: gnovack <novackgm@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: xiaozhoupy <peiyuanzhou1994@gmail.com> Co-authored-by: Roy Wang <yasong.wang@inferact.ai> Co-authored-by: Jeff (Junze) Ma <93145857+majunze2001@users.noreply.github.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-07-27 07:01:52 +00:00