disable k3dspark test

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

cmake/external_projects/deepgemm.cmake

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

docs/models/supported_models.md

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

pyproject.toml

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

requirements/cuda.txt

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

requirements/test/cuda.txt

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

tests/models/registry.py

@@ -1024,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"},
@@ -1464,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",
@@ -1662,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,

tests/models/test_dspark_mla.py

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

tests/test_config.py

@@ -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"),
     [

tests/transformers_utils/test_dspark_mla_config.py

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

tests/v1/attention/test_mla_backends.py

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

tests/v1/attention/test_mla_noncausal.py

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

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

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

tests/v1/core/test_prefix_caching.py

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

tests/v1/cudagraph/test_breakable_cudagraph.py

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

tests/v1/kv_connector/unit/test_nixl_connector.py

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

tests/v1/spec_decode/test_llm_base_proposer_sampling.py

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

tools/install_deepgemm.sh

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

vllm/_aiter_ops.py

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

vllm/config/compilation.py

@@ -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",

vllm/config/parallel.py

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

vllm/config/speculative.py

@@ -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)
@@ -934,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.
@@ -963,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"
                 ):

vllm/config/vllm.py

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

vllm/distributed/device_communicators/cuda_communicator.py

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

vllm/distributed/device_communicators/custom_all_reduce.py

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

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py

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

vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py

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

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py

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

vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py

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

vllm/engine/arg_utils.py

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

vllm/envs.py

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

vllm/model_executor/layers/attention/mla_attention.py

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

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

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

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

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

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

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

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

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

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

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

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

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

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

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

vllm/model_executor/layers/logits_processor.py

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

vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py

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

vllm/model_executor/layers/mamba/mamba_utils.py

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

vllm/model_executor/layers/mamba/ops/causal_conv1d.py

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

vllm/model_executor/layers/mla.py

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

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py

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

vllm/model_executor/layers/quantization/mxfp4.py

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

vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

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

vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py

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

vllm/model_executor/layers/quantization/utils/mxfp8_utils.py

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

vllm/model_executor/layers/vocab_parallel_embedding.py

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

vllm/model_executor/model_loader/utils.py

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

vllm/model_executor/models/config.py

@@ -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:
@@ -870,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,

vllm/model_executor/models/deepseek_v2.py

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

vllm/model_executor/models/interfaces.py

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

vllm/model_executor/models/kimi_k25_vit.py

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

vllm/model_executor/models/kimi_linear.py

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

vllm/model_executor/models/qwen3_dspark.py

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

vllm/model_executor/models/registry.py

@@ -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"),
@@ -458,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",
@@ -607,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"),
@@ -647,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"),

vllm/model_executor/warmup/kernel_warmup.py

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

vllm/model_executor/warmup/kimi_k3_triton_warmup.py

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

vllm/models/deepseek_v4/attention.py

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

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

@@ -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",

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

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

vllm/platforms/interface.py

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

vllm/third_party/flash_linear_attention/ops/fused_recurrent.py

@@ -12,7 +12,7 @@ import torch
 
 from vllm.triton_utils import tl, triton
 
-from .op import exp
+from .op import exp, log
 
 
 @triton.heuristics(

vllm/transformers_utils/config.py

@@ -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",

vllm/transformers_utils/configs/__init__.py

@@ -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",

vllm/transformers_utils/configs/k3_dspark.py

@@ -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."
+            )

vllm/transformers_utils/configs/kimi_linear.py

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

vllm/transformers_utils/model_arch_config_convertor.py

@@ -267,6 +267,7 @@ class ModelArchConfigConvertorBase:
             "deepseek_v32",
             "deepseek_v4",
             "deepseek_mtp",
+            "k3_dspark",
             "glm_moe_dsa",
             "glm4_moe_lite",
             "glm4_moe_lite_mtp",

vllm/transformers_utils/processors/__init__.py

@@ -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",

vllm/utils/multi_stream_utils.py

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

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

@@ -5,7 +5,12 @@ from typing import ClassVar
 
 import torch
 from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
+from flashinfer.utils import (
+    get_device_sm_count,
+    get_trtllm_gen_multi_ctas_kv_counter_bytes,
+)
 
+from vllm.config import get_current_vllm_config
 from vllm.config.cache import CacheDType
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention.mla_attention import (
@@ -27,6 +32,37 @@ from vllm.v1.attention.backends.utils import KVCacheLayoutType
 
 logger = init_logger(__name__)
 
+
+def _trtllm_gen_mla_decode_supports_num_heads(num_heads: int) -> bool:
+    """True if trtllm-gen's MLA decode kernel supports this query head count.
+
+    The kernel groups Q heads into CTAs of ``min(num_heads, tileSizeQ)`` and
+    requires ``num_heads`` divisible by that, else raises "The
+    numHeadsQ/numHeadsKv is not supported" (flashinfer fmhaKernels.cuh).
+    ``tileSizeQ`` is 8/16 for ``num_heads <= 8``/``<= 32`` (SwapsMmaAb) else 64;
+    treat ``> 32`` as tile 64 (safe bound). E.g. 96/24 -> False, 48/64/128 -> True.
+    """
+    if num_heads <= 8:
+        tile = 8
+    elif num_heads <= 32:
+        tile = 16
+    else:
+        tile = 64
+    return num_heads % min(num_heads, tile) == 0
+
+
+def _select_mla_decode_backend(num_heads: int) -> str | None:
+    """cute-dsl for head counts trtllm-gen cannot tile, else None (=> auto)."""
+    if not _trtllm_gen_mla_decode_supports_num_heads(num_heads):
+        logger.warning_once(
+            "trtllm-gen MLA decode does not support num_heads=%d "
+            "(query/kv head ratio); falling back to the cute-dsl backend.",
+            num_heads,
+        )
+        return "cute-dsl"
+    return None
+
+
 FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
 FLASHINFER_MLA_LSE_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024
 
@@ -48,9 +84,39 @@ def _get_workspace_buffer(return_lse: bool) -> torch.Tensor:
     return _fi_workspace
 
 
+_fi_multi_ctas_kv_counter: torch.Tensor | None = None
+
+
+def _get_multi_ctas_kv_counter_buffer(
+    min_bytes: int, device: torch.device
+) -> torch.Tensor:
+    """Persistent, zero-initialized trtllm-gen multi-CTA-KV counter buffer.
+
+    trtllm-gen's multi-CTA-KV MLA decode kernel resets these semaphores to zero
+    at the end of every launch, so the buffer only needs zeroing once. The
+    public ``trtllm_batch_decode_with_kv_cache_mla`` entry point builds a fresh
+    runner per call, so without a caller-owned buffer it re-allocates and
+    re-zeros this counter on every decode step (a tiny ``FillFunctor<uint8>``
+    launch right before the FMHA). Owning it here and passing it in removes that
+    per-step launch. ``min_bytes`` is sized to the worst-case batch so the
+    buffer is allocated once and never reallocated after CUDA-graph capture.
+    """
+    global _fi_multi_ctas_kv_counter
+    if (
+        _fi_multi_ctas_kv_counter is None
+        or _fi_multi_ctas_kv_counter.numel() < min_bytes
+    ):
+        _fi_multi_ctas_kv_counter = torch.zeros(
+            min_bytes, dtype=torch.uint8, device=device
+        )
+    return _fi_multi_ctas_kv_counter
+
+
 class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
     query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
+    # Non-causal DSpark blocks are flattened to single-token rows in forward_mqa.
+    supports_non_causal_multi_token_decode: ClassVar[bool] = True
 
 
 class FlashInferMLABackend(MLACommonBackend):
@@ -91,6 +157,10 @@ class FlashInferMLABackend(MLACommonBackend):
     def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
         return capability.major == 10
 
+    @classmethod
+    def supports_non_causal(cls) -> bool:
+        return True
+
     @classmethod
     def supports_combination(
         cls,
@@ -179,6 +249,16 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
 
         self.bmm1_scale: float | None = None
         self.bmm2_scale: float | None = None
+        # Worst-case decode batch for the persistent trtllm-gen multi-CTA-KV
+        # counter buffer (see _get_multi_ctas_kv_counter_buffer). Captured here
+        # (config is in scope during construction) so the byte size can be
+        # resolved once on the first decode and never grows after CUDA-graph
+        # capture. The impl has no _vllm_config, hence get_current_vllm_config().
+        _sched = get_current_vllm_config().scheduler_config
+        self._mla_counter_max_batch: int = (
+            _sched.max_num_batched_tokens or _sched.max_num_seqs
+        )
+        self._mla_counter_bytes: int | None = None
 
     def forward_mqa(
         self,
@@ -194,8 +274,20 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
             q_nope, q_pe = q
             q = torch.cat([q_nope, q_pe], dim=-1)
 
+        block_table = attn_metadata.decode.block_table
+        seq_lens = attn_metadata.decode.seq_lens
+
+        if not attn_metadata.causal:
+            # Non-causal DSpark block: flatten to single-token decode rows with
+            # per-row context seq_lens (trtllm-gen has no causal flag and would
+            # otherwise mask the block causally).
+            query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes
+            q = q.unsqueeze(1)
+            if query_len > 1:
+                block_table = block_table.repeat_interleave(query_len, dim=0)
+                seq_lens = seq_lens.repeat_interleave(query_len)
         # trtllm API requires extra dimension q_len_per_request for MTP
-        if attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:
+        elif attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:
             logger.warning_once(
                 """FlashInferMLAImpl got a query of uneven length.
                 This usually indicates an issue in batch reordering
@@ -217,6 +309,29 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
 
         return_lse = self.need_to_return_lse_for_decode
         workspace_buffer = _get_workspace_buffer(return_lse)
+        # trtllm-gen rejects MLA head counts it can't tile (e.g. 96);
+        # fall back to cute-dsl for those.
+        decode_backend = _select_mla_decode_backend(self.num_heads)
+        extra_kwargs = {}
+        if decode_backend:
+            extra_kwargs["backend"] = decode_backend
+        elif kv_c_and_k_pe_cache.shape[-2] in (32, 64):
+            # The auto path can dispatch to trtllm-gen, whose multi-CTA-KV decode
+            # kernel self-resets its semaphore counter after each launch (so it
+            # only needs zeroing once). Pass a persistent counter buffer to skip
+            # the per-step re-allocate + re-zero the public entry point would
+            # otherwise do. Guarded to configs where a trtllm-gen runner is
+            # eligible (page/block size in {32, 64}); the arg is rejected when
+            # only a cute-dsl runner can run.
+            if self._mla_counter_bytes is None:
+                self._mla_counter_bytes = get_trtllm_gen_multi_ctas_kv_counter_bytes(
+                    self._mla_counter_max_batch,
+                    self.num_heads,
+                    get_device_sm_count(q.device),
+                )
+            extra_kwargs["multi_ctas_kv_counter_buffer"] = (
+                _get_multi_ctas_kv_counter_buffer(self._mla_counter_bytes, q.device)
+            )
         kernel_out = trtllm_batch_decode_with_kv_cache_mla(
             query=q,
             kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
@@ -224,12 +339,13 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
             qk_nope_head_dim=self.qk_nope_head_dim,
             kv_lora_rank=self.kv_lora_rank,
             qk_rope_head_dim=self.qk_rope_head_dim,
-            block_tables=attn_metadata.decode.block_table,
-            seq_lens=attn_metadata.decode.seq_lens,
+            block_tables=block_table,
+            seq_lens=seq_lens,
             max_seq_len=attn_metadata.max_seq_len,
             bmm1_scale=self.bmm1_scale,
             bmm2_scale=self.bmm2_scale,
             return_lse=return_lse,
+            **extra_kwargs,
         )
         if return_lse:
             o, lse = kernel_out

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

@@ -73,6 +73,17 @@ class MLAPrefillBackend(ABC):
         pass. Overridden by backends that support it."""
         return False
 
+    def supports_out(self) -> bool:
+        """Whether `run_prefill_new_tokens` honors a caller-provided `out`
+        tensor of shape `[num_tokens, num_heads, v_head_dim]`, writing the
+        final result into it in place.
+
+        When True, callers may pass `out` and skip the post-hoc
+        slice/flatten/copy. False for backends that ignore `out` or emit a
+        padded (`qk_head_dim`) output. Overridden by backends that support it.
+        """
+        return False
+
     @classmethod
     def validate_configuration(
         cls,

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

@@ -419,6 +419,11 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
             return attn_out, lse
         return attn_out
 
+    def supports_out(self) -> bool:
+        # A padded V produces a qk_head_dim output that cannot be written into
+        # a v_head_dim `out`; only the unpadded path honors `out`.
+        return not self.requires_v_padding
+
     def run_prefill_new_tokens(
         self,
         q: torch.Tensor,

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

@@ -108,15 +108,21 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
             return self._global_hyperparameters
 
         from vllm.model_executor.layers.attention.mla_attention import (
-            MLAAttention,
             MLACommonImpl,
         )
+        from vllm.model_executor.layers.attention_layer_base import (
+            AttentionLayerBase,
+        )
 
+        # Match any layer with an MLA impl, not just the MLAAttention wrapper:
+        # fused MLA modules (Kimi-K3's MultiHeadLatentAttention) register a
+        # different layer type. Keying on impl also excludes linear/KDA layers.
         forward_context = self.vllm_config.compilation_config.static_forward_context
         layer_names = [
             name
             for name, layer in forward_context.items()
-            if isinstance(layer, MLAAttention)
+            if isinstance(layer, AttentionLayerBase)
+            and isinstance(getattr(layer, "impl", None), MLACommonImpl)
         ]
 
         self._global_hyperparameters = infer_global_hyperparameters(
@@ -191,6 +197,10 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
                     o_data_type=prefill_metadata.output_dtype,
                 )
 
+    def supports_out(self) -> bool:
+        # Planned with head_dim_vo == v_head_dim, so the output is unpadded.
+        return True
+
     def run_prefill_new_tokens(
         self,
         q: torch.Tensor,
@@ -206,12 +216,13 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
             q=q,
             k=k,
             v=v,
+            out=out,
             return_lse=return_softmax_lse,
         )
 
         if isinstance(ret, tuple):
             # Convert from (q_len, num_heads) to (num_heads, q_len)
-            return ret[0], ret[1].transpose(0, 1).contiguous()
+            return ret[0], ret[1].transpose(0, 1)
         return ret
 
     def run_prefill_context_chunk(
@@ -229,4 +240,4 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
         )
 
         # Convert from (q_len, num_heads) to (num_heads, q_len)
-        return attn_out, lse.transpose(0, 1).contiguous()
+        return attn_out, lse.transpose(0, 1)

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

@@ -118,6 +118,10 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
             prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1]
         )
 
+    def supports_out(self) -> bool:
+        # Output head dim is v_head_dim (unpadded); `out` supported since 0.1.8.
+        return True
+
     def run_prefill_new_tokens(
         self,
         q: torch.Tensor,
@@ -129,11 +133,11 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         from tokenspeed_mla import tokenspeed_mla_prefill
 
-        # `v` arrives as the second half of `kv_nope.split(...)` in
-        # mla_attention.forward_mha — a non-contiguous view of `kv_nope` along
-        # dim=-1. The kernel does `v.reshape(1, total_kv, h_k, 1, d_v)` which
-        # would silently copy on a non-contiguous tensor; force contiguity here
-        # so the copy (if any) happens once outside the kernel call.
+        # `v` arrives as the second half of `kv_nope.split(...)` — a
+        # non-contiguous view of `kv_nope` along dim=-1. The kernel wraps inputs
+        # via `from_dlpack`, which preserves strides, but the FMHA kernel reads
+        # `v` assuming contiguous storage and silently produces wrong output on
+        # a strided view (verified on tokenspeed-mla 0.1.8). Force contiguity.
         v = v.contiguous()
 
         ret = tokenspeed_mla_prefill(
@@ -148,11 +152,12 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
             is_causal=True,
             return_lse=return_softmax_lse,
             enable_pdl=False,
+            out=out,
         )
 
         if isinstance(ret, tuple):
             # Convert from (q_len, num_heads) to (num_heads, q_len)
-            return ret[0], ret[1].transpose(0, 1).contiguous()
+            return ret[0], ret[1].transpose(0, 1)
         return ret
 
     def run_prefill_context_chunk(
@@ -188,4 +193,4 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
         )
 
         # Convert from (q_len, num_heads) to (num_heads, q_len)
-        return attn_out, lse.transpose(0, 1).contiguous()
+        return attn_out, lse.transpose(0, 1)

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

@@ -91,6 +91,10 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
             prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1]
         )
 
+    def supports_out(self) -> bool:
+        # Output head dim is v.shape[-1] == v_head_dim, so `out` is unpadded.
+        return True
+
     def run_prefill_new_tokens(
         self,
         q: torch.Tensor,
@@ -102,13 +106,14 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         from flashinfer.prefill import trtllm_ragged_attention_deepseek
 
-        out = torch.empty(
-            q.shape[0],
-            q.shape[1],
-            v.shape[2],
-            device=q.device,
-            dtype=self._prefill_metadata.output_dtype,
-        )
+        if out is None:
+            out = torch.empty(
+                q.shape[0],
+                q.shape[1],
+                v.shape[2],
+                device=q.device,
+                dtype=self._prefill_metadata.output_dtype,
+            )
 
         ret = trtllm_ragged_attention_deepseek(
             query=q,
@@ -133,7 +138,7 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
 
         if isinstance(ret, tuple):
             # Convert from (q_len, num_heads) to (num_heads, q_len)
-            return ret[0], ret[1].transpose(0, 1).contiguous()
+            return ret[0], ret[1].transpose(0, 1)
         return ret
 
     def run_prefill_context_chunk(
@@ -182,4 +187,4 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
         )
 
         # Convert from (q_len, num_heads) to (num_heads, q_len)
-        return attn_out, lse.transpose(0, 1).contiguous()
+        return attn_out, lse.transpose(0, 1)

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

@@ -31,6 +31,32 @@ from vllm.v1.kv_cache_interface import AttentionSpec
 logger = init_logger(__name__)
 
 
+@functools.lru_cache(maxsize=1)
+def _get_mla_gluon():
+    """Load the small-head Gluon MLA entry point."""
+    unified_module = "aiter.ops.triton.gluon.mla_gluon"
+    try:
+        from aiter.ops.triton.gluon.mla_gluon import mla_gluon
+
+        return mla_gluon
+    except ModuleNotFoundError as unified_import_error:
+        if not unified_module.startswith(unified_import_error.name or ""):
+            raise
+        legacy_module = "aiter.ops.triton.gluon.mla_decode_gluon"
+        try:
+            from aiter.ops.triton.gluon.mla_decode_gluon import mla_decode_gluon
+
+            return mla_decode_gluon
+        except ModuleNotFoundError as legacy_import_error:
+            if not legacy_module.startswith(legacy_import_error.name or ""):
+                raise
+            raise RuntimeError(
+                "ROCM_AITER_MLA requires an AITER build with the small-head "
+                "Gluon MLA kernel (mla_gluon or mla_decode_gluon) when decode "
+                "heads are fewer than 16."
+            ) from unified_import_error
+
+
 @functools.lru_cache(maxsize=1)
 def _fp8_mla_prefill_supported() -> bool:
     """Auto-detect FP8 MLA prefill via mla_prefill_ps_asm_fwd + mla_reduce_v1.
@@ -102,6 +128,10 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
     attn_out_dtype: torch.dtype = torch.bfloat16
     # The max query output length: int
     max_qo_len: int | None = None
+    # Minimum KV length used by Gluon to choose a safe split count.
+    min_kv_seq_len: int = 1
+    # Small-head decode uses Gluon (avoids padding to 16).
+    use_gluon_decode: bool = False
     # Whether persistent MLA metadata was computed (only for qseqlen=1)
     has_persistent_metadata: bool = False
 
@@ -241,7 +271,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
             device=device,
         )
 
-        self._fp8_prefill_enabled = _fp8_mla_prefill_supported()
+        # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads.
+        self._fp8_prefill_enabled = (
+            _fp8_mla_prefill_supported() and self.num_heads % 16 == 0
+        )
         if self._fp8_prefill_enabled:
             max_prefill_qlen = min(
                 vllm_config.model_config.max_model_len,
@@ -469,6 +502,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         )
         qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
         max_qo_len = qo_len.max().item()
+        use_gluon_decode = AiterMLAHelper.use_gluon_decode(
+            self.num_heads, int(max_qo_len)
+        )
 
         if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
             self.paged_kv_indices.fill_(-1)
@@ -518,7 +554,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         # so forward_mqa can skip passing it (falling back to the kernel
         # computing its own metadata internally, like v0.18.0).
         has_persistent_metadata = False
-        if max_qo_len == 1:
+        if max_qo_len == 1 and not use_gluon_decode:
             from aiter import get_mla_metadata_v1
 
             get_mla_metadata_v1(
@@ -553,6 +589,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
             qo_indptr=qo_indptr,
             dcp_tot_seq_lens=dcp_tot_seq_lens_device,
             max_qo_len=max_qo_len,
+            use_gluon_decode=use_gluon_decode,
             attn_out_dtype=self.decode_attn_out_dtype,
             has_persistent_metadata=has_persistent_metadata,
         )
@@ -646,17 +683,16 @@ class AiterMLAHelper:
     @staticmethod
     def check_num_heads_validity(num_heads: int):
         assert AiterMLAHelper.is_valid_num_heads(num_heads), (
-            f"Aiter MLA requires that num_heads be multiples or divisors of 16, "
-            f"but provided {num_heads} number of heads.\n"
+            "ROCM AITER MLA requires 1-15 heads for Gluon decode or a multiple "
+            f"of 16 heads for persistent decode, but got {num_heads}.\n"
             f"Try adjusting tensor_parallel_size value."
         )
 
     @staticmethod
     def is_valid_num_heads(num_heads: int) -> bool:
-        return (
-            num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0
-            if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
-            else AiterMLAHelper._AITER_MIN_MLA_HEADS % num_heads == 0
+        return num_heads > 0 and (
+            num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS
+            or num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0
         )
 
     @staticmethod
@@ -681,6 +717,10 @@ class AiterMLAHelper:
             else o[:, :: AiterMLAHelper._AITER_MIN_MLA_HEADS // num_heads, :]
         )
 
+    @staticmethod
+    def use_gluon_decode(num_heads: int, max_qo_len: int) -> bool:
+        return num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS and max_qo_len == 1
+
 
 class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
     def __init__(
@@ -726,7 +766,10 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
 
         # FP8 MLA prefill kernel imports (lazy, only when enabled).
         # Auto-enabled on gfx950 when AITER ships the kernels.
-        self._fp8_prefill_enabled = _fp8_mla_prefill_supported()
+        # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads.
+        self._fp8_prefill_enabled = (
+            _fp8_mla_prefill_supported() and self.num_heads % 16 == 0
+        )
         if self._fp8_prefill_enabled:
             from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1
 
@@ -917,7 +960,106 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         assert kv_c_and_k_pe_cache.numel() > 0
         assert attn_metadata.decode is not None
-        assert attn_metadata.decode.max_qo_len is not None
+
+        decode = attn_metadata.decode
+        assert decode.max_qo_len is not None
+        assert decode.paged_kv_indptr is not None
+        assert decode.paged_kv_indices is not None
+        if decode.use_gluon_decode:
+            if type(q) is tuple:
+                q_nope, q_pe = q
+            else:
+                q_nope, q_pe = torch.split(
+                    q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
+                )
+            B, num_q_heads, _ = q_nope.shape
+            o = torch.empty(
+                B,
+                num_q_heads,
+                self.kv_lora_rank,
+                dtype=decode.attn_out_dtype,
+                device=q_nope.device,
+            )
+            kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1])
+            mla_gluon = _get_mla_gluon()
+            mla_gluon(
+                q_nope=q_nope,
+                q_pe=q_pe,
+                kv_c=kv_buffer,
+                o=o,
+                page_table=decode.paged_kv_indices,
+                seq_info=decode.paged_kv_indptr,
+                sm_scale=self.scale,
+                k_pe=None,
+                kv_pe_offset=self.kv_lora_rank,
+                use_2d_view=False,
+                kv_scale=1.0,
+                min_kv_seq_len=decode.min_kv_seq_len,
+            )
+            return o, None
+
+        # 12-head (<16) non-causal multi-token verify (DSpark): the asm path has
+        # no gqa<16, qseqlen>1 kernel. Flatten each verify token to a qseqlen=1
+        # gluon decode over the same committed prefix (non-causal), mirroring the
+        # TRITON_MLA / sparse-backend flatten but on the fast gluon kernel. Each
+        # request's KV metadata is repeated max_qo_len times.
+        if (
+            self.num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS
+            and int(decode.max_qo_len) > 1
+        ):
+            qlen = int(decode.max_qo_len)
+            if type(q) is tuple:
+                q_nope, q_pe = q
+            else:
+                q_nope, q_pe = torch.split(
+                    q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
+                )
+            B, num_q_heads, _ = q_nope.shape
+            o = torch.empty(
+                B,
+                num_q_heads,
+                self.kv_lora_rank,
+                dtype=decode.attn_out_dtype,
+                device=q_nope.device,
+            )
+            kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1])
+            # Expand per-request paged-KV to per-verify-token: each request's KV
+            # block is repeated qlen times (row r*qlen+t attends request r's
+            # committed prefix). Fully vectorized (no host loop).
+            old_indptr = decode.paged_kv_indptr
+            per_req_len = old_indptr[1:] - old_indptr[:-1]
+            dev = q_nope.device
+            row_req = torch.arange(per_req_len.shape[0], device=dev).repeat_interleave(
+                qlen
+            )
+            row_len = per_req_len[row_req]
+            new_indptr = torch.cat([old_indptr.new_zeros(1), row_len.cumsum(0)]).to(
+                torch.int32
+            )
+            total = int(new_indptr[-1].item())
+            within = torch.arange(total, device=dev, dtype=torch.int64) - new_indptr[
+                :-1
+            ].to(torch.int64).repeat_interleave(row_len)
+            src = (
+                old_indptr[row_req].to(torch.int64).repeat_interleave(row_len) + within
+            )
+            new_indices = decode.paged_kv_indices[src]
+            mla_gluon = _get_mla_gluon()
+            mla_gluon(
+                q_nope=q_nope,
+                q_pe=q_pe,
+                kv_c=kv_buffer,
+                o=o,
+                page_table=new_indices,
+                seq_info=new_indptr,
+                sm_scale=self.scale,
+                k_pe=None,
+                kv_pe_offset=self.kv_lora_rank,
+                use_2d_view=False,
+                kv_scale=1.0,
+                min_kv_seq_len=int(per_req_len.min()),
+            )
+            return o, None
 
         if type(q) is tuple:
             q = torch.cat(q, dim=-1)
@@ -960,11 +1102,11 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
             kv_buffer,
             o,
             self.scale,
-            attn_metadata.decode.qo_indptr,
-            attn_metadata.decode.max_qo_len,
-            attn_metadata.decode.paged_kv_indptr,
-            attn_metadata.decode.paged_kv_indices,
-            attn_metadata.decode.paged_kv_last_page_len,
+            decode.qo_indptr,
+            decode.max_qo_len,
+            decode.paged_kv_indptr,
+            decode.paged_kv_indices,
+            decode.paged_kv_last_page_len,
             **mla_kwargs,
         )
 

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

@@ -51,9 +51,17 @@ class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
     _cudagraph_support: ClassVar[AttentionCGSupport] = (
         AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
     )
+    # Non-causal DSpark block is flattened to one decode row per query token in
+    # forward_mqa, so no intra-block causal masking is required.
+    supports_non_causal_multi_token_decode: ClassVar[bool] = True
 
     def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
         super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        # Only the non-causal DSpark draft group serves multi-token blocks via
+        # the decode path; raise its reorder threshold to the spec block length
+        # so full-cudagraph capture admits it. Causal usage stays single-token.
+        if getattr(self, "non_causal_multi_token_decode", False):
+            self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
         self._reserve_attn_logits_workspace()
 
     def _reserve_attn_logits_workspace(self) -> None:
@@ -68,6 +76,10 @@ class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
             return
         # Decode reorder threshold is 1, so decode tokens <= max_num_seqs.
         B = self.vllm_config.scheduler_config.max_num_seqs
+        # Non-causal DSpark draft flattens each request's block to query_len
+        # decode rows; cover max_num_seqs * block_len rows.
+        if getattr(self, "non_causal_multi_token_decode", False):
+            B *= self.reorder_batch_threshold
         # DCP all-gathers the query heads before forward_mqa.
         q_num_heads = self.num_heads * self.dcp_world_size
         max_splits = _compute_num_kv_splits(
@@ -132,6 +144,13 @@ class TritonMLABackend(MLACommonBackend):
     def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
         return True
 
+    @classmethod
+    def supports_non_causal(cls) -> bool:
+        # DSpark non-causal blocks are flattened to single-token decode rows in
+        # TritonMLAImpl.forward_mqa (decode_attention_fwd has no causal flag /
+        # no intra-block masking). Enables the non-causal AMD MLA path.
+        return True
+
 
 class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
     can_return_lse_for_decode: bool = True
@@ -264,6 +283,18 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
         kv_c_cache = kv_c_and_k_pe_cache[..., : self.kv_lora_rank]
         PAGE_SIZE = kv_c_and_k_pe_cache.size(1)
 
+        block_table = attn_metadata.decode.block_table
+        seq_lens = attn_metadata.decode.seq_lens
+        if not attn_metadata.causal:
+            # Non-causal DSpark block: flatten to one decode row per query token.
+            # Each row attends to the same committed KV prefix (per-row seq_lens)
+            # and never to sibling block tokens = non-causal block semantics.
+            # Mirrors FlashInferMLA's non-causal path.
+            query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes
+            if query_len > 1:
+                block_table = block_table.repeat_interleave(query_len, dim=0)
+                seq_lens = seq_lens.repeat_interleave(query_len)
+
         # Run MQA — always pass layer scales. When KV cache is
         # BF16 the kernel's `if dtype.is_fp8()` check is a no-op.
         decode_attention_fwd(
@@ -272,8 +303,8 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
             kv_c_cache,
             o,
             lse,
-            attn_metadata.decode.block_table,
-            attn_metadata.decode.seq_lens,
+            block_table,
+            seq_lens,
             attn_logits,
             num_kv_splits,
             self.scale,

vllm/v1/core/kv_cache_utils.py

@@ -1209,6 +1209,27 @@ def _get_kv_cache_groups_uniform_page_size(
     for layer_name, layer_spec in kv_cache_spec.items():
         same_type_layers[layer_spec].append(layer_name)
 
+    # Attempt to further merge same-type layers based on whether their KV
+    # cache specs can be merged, to minimize the group count. This benefits
+    # situations where specs share a block layout and differ only in a
+    # property it can reconcile (e.g. full attention layers differing only in
+    # sliding window / attention chunk size).
+    layer_buckets: list[list[str]] = []
+    spec_buckets: list[list[KVCacheSpec]] = []
+    for layer_spec, layer_names in same_type_layers.items():
+        for names, specs in zip(layer_buckets, spec_buckets):
+            try:
+                # A raise means that the specs are incompatible.
+                type(specs[0]).merge([*specs, layer_spec])
+            except (AssertionError, ValueError):
+                continue
+            names.extend(layer_names)
+            specs.append(layer_spec)
+            break
+        else:
+            layer_buckets.append(list(layer_names))
+            spec_buckets.append([layer_spec])
+
     # Split each group into smaller groups, to make the number of layers in each
     # group identical. Add padding to the last group of each type if necessary.
     # E.g., (full.0, full.1), (sw.0, sw.1, sw.2)
@@ -1221,9 +1242,9 @@ def _get_kv_cache_groups_uniform_page_size(
     # is the minimum number of layers among all attention types. Need a better
     # strategy if we want to support more complex patterns (e.g., 20 full + 30
     # sw, where the group size should be 10).
-    min_num_layers = min([len(layers) for layers in same_type_layers.values()])
+    min_num_layers = min([len(layers) for layers in layer_buckets])
     group_size = min_num_layers
-    max_num_layers = max([len(layers) for layers in same_type_layers.values()])
+    max_num_layers = max([len(layers) for layers in layer_buckets])
     if max_num_layers < min_num_layers * 1.5:
         # If the number of layers is not much larger than the minimum number of
         # layers, use the maximum number of layers as the group size to avoid
@@ -1234,7 +1255,7 @@ def _get_kv_cache_groups_uniform_page_size(
         # extra layers to one attention type.
         group_size = max_num_layers
     grouped_layers = []
-    for layers in same_type_layers.values():
+    for layers in layer_buckets:
         num_padding_layers = group_size - len(layers) % group_size
         if num_padding_layers != group_size:
             logger.warning(

vllm/v1/core/sched/scheduler.py

@@ -363,10 +363,11 @@ class Scheduler(SchedulerInterface):
     ) -> int:
         """Clip a prefill chunk so it ends where Mamba state must be cached.
 
-        In "align" cache mode the SSM state is only materialized at chunk
-        ends, so chunk ends are steered onto cacheable positions: block
-        boundaries by default, plus mandatory early stops (the prompt's
-        partial-tail hash boundary, a detected shared-prefix junction).
+        In "align" cache mode reusable SSM states are materialized at block
+        boundaries, plus mandatory early stops (the prompt's partial-tail hash
+        boundary, a detected shared-prefix junction). If a block is larger
+        than the configured prefill chunk limit, intermediate chunks keep
+        private running state until they reach the next cacheable position.
         """
         start = (
             request.num_computed_tokens
@@ -387,11 +388,17 @@ class Scheduler(SchedulerInterface):
             last_cache_position = max(last_cache_position - block_size, 0)
 
         end = start + num_new_tokens
-        # Until `last_cache_position`, chunk ends must land on block
-        # boundaries. May yield an empty chunk (budget cannot reach the next
-        # boundary); the caller then skips the request.
+        # Until `last_cache_position`, prefer chunks ending on block
+        # boundaries. When a block cannot fit in any configured prefill chunk,
+        # allow sub-block progress and re-align at the next reachable boundary.
         if end < last_cache_position:
-            end = end // block_size * block_size
+            max_prefill_tokens = self.max_num_scheduled_tokens
+            long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold
+            if long_prefill_threshold > 0:
+                max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold)
+            aligned_end = end // block_size * block_size
+            if aligned_end > start or block_size <= max_prefill_tokens:
+                end = aligned_end
 
         next_block_boundary = (start // block_size + 1) * block_size
         tail_boundary = (

vllm/v1/kv_cache_interface.py

@@ -385,6 +385,8 @@ class MLAAttentionSpec(FullAttentionSpec):
     alignment: int | None = None  # Default to None for no padding.
     compress_ratio: int = 1  # Default to 1 for no compression.
     model_version: str | None = None
+    # Marks draft groups that flatten a non-causal query block into decode rows.
+    non_causal_multi_token_decode: bool = False
 
     def __post_init__(self):
         super().__post_init__()
@@ -424,6 +426,12 @@ class MLAAttentionSpec(FullAttentionSpec):
         compress_ratio_set = set(spec.compress_ratio for spec in specs)
         model_version_set = set(spec.model_version for spec in specs)
         block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs)
+        non_causal_decode_set = {spec.non_causal_multi_token_decode for spec in specs}
+        assert len(non_causal_decode_set) == 1, (
+            "All attention layers in the same KV cache group must agree on the "
+            "non-causal decode mode, since it configures the shared metadata "
+            "builder (e.g. its reorder batch threshold)."
+        )
         assert (
             len(cache_dtype_str_set) == 1
             and len(compress_ratio_set) == 1
@@ -434,7 +442,7 @@ class MLAAttentionSpec(FullAttentionSpec):
             "quantization method, compress ratio, model version, and KV block "
             "stride indexing."
         )
-        return cls(
+        merged_spec = cls(
             block_size=specs[0].block_size,
             num_kv_heads=specs[0].num_kv_heads,
             head_size=specs[0].head_size,
@@ -445,7 +453,15 @@ class MLAAttentionSpec(FullAttentionSpec):
             cache_dtype_str=cache_dtype_str_set.pop(),
             compress_ratio=compress_ratio_set.pop(),
             model_version=model_version_set.pop(),
+            non_causal_multi_token_decode=non_causal_decode_set.pop(),
         )
+        for spec in specs:
+            for f in fields(AttentionSpec):
+                assert getattr(spec, f.name) == getattr(merged_spec, f.name), (
+                    "All attention layers in the same KV cache group must have "
+                    "the same attention spec."
+                )
+        return merged_spec
 
 
 @dataclass(frozen=True, kw_only=True)

vllm/v1/spec_decode/llm_base_proposer.py

@@ -1006,11 +1006,11 @@ class SpecDecodeBaseProposer:
 
     def model_returns_tuple(self) -> bool:
         if self.method == "mtp":
-            # DeepSeek-family MTP (deepseek_mtp.py) recycles the post-final-
-            # norm hidden, so its forward returns (logit_hidden,
-            # recycle_hidden). Other MTP families return a single tensor.
-            return "DeepSeekMTPModel" in (
-                self.draft_model_config.hf_config.architectures or []
+            # These models return separate hidden states for logits and for
+            # feedback into the next draft step.
+            architectures = self.draft_model_config.hf_config.architectures or []
+            return bool(
+                {"DeepSeekMTPModel", "KimiK3MTPModel"}.intersection(architectures)
             )
         return self.method not in ("mtp", "draft_model", "dflash")
 
@@ -1386,7 +1386,10 @@ class SpecDecodeBaseProposer:
                 self.model.config.image_token_index = (
                     target_model.config.vision_config.image_token_id
                 )
-            elif self.get_model_name(target_model) == "KimiK25ForConditionalGeneration":
+            elif self.get_model_name(target_model) in (
+                "KimiK25ForConditionalGeneration",
+                "KimiK3ForConditionalGeneration",
+            ):
                 self.model.config.image_token_index = (
                     target_model.config.media_placeholder_token_id
                 )

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

@@ -344,13 +344,6 @@ class DFlashSpeculator(DraftModelSpeculator):
             hidden_states = last_hidden_states
         self.hidden_states[:num_target_tokens].copy_(hidden_states[:num_target_tokens])
 
-        self._copy_request_inputs(
-            num_reqs,
-            input_batch.idx_mapping,
-            temperature,
-            seeds,
-        )
-
         if dummy_run and skip_attn_for_dummy_run:
             # Memory profiling path: block_tables / kv_cache_config are not initialized.
             # Since DFlash needs to build its own attention metadata, we must skip the
@@ -385,11 +378,15 @@ class DFlashSpeculator(DraftModelSpeculator):
                 self.sample_indices,
                 self.sample_pos,
                 self.sample_idx_mapping,
+                self.temperature,
+                self.seeds,
                 input_batch,
                 num_sampled,
                 num_rejected,
                 last_sampled,
                 next_prefill_tokens,
+                temperature,
+                seeds,
                 self.block_tables.input_block_tables[gid],
                 self.block_tables.kernel_block_sizes[gid],
                 self.parallel_drafting_token_id,
@@ -482,6 +479,8 @@ def _prepare_dflash_inputs_kernel(
     out_sample_indices_ptr,
     out_sample_pos_ptr,
     out_sample_idx_mapping_ptr,
+    out_temperature_ptr,
+    out_seeds_ptr,
     # Inputs from target batch
     target_positions_ptr,
     target_query_start_loc_ptr,
@@ -490,6 +489,9 @@ def _prepare_dflash_inputs_kernel(
     next_prefill_tokens_ptr,
     num_sampled_ptr,
     num_rejected_ptr,
+    # Sampling params
+    temperature_ptr,
+    seeds_ptr,
     # Block table for slot mapping lookup.
     block_table_ptr,
     block_table_stride,
@@ -585,6 +587,12 @@ def _prepare_dflash_inputs_kernel(
         # reads up to (context + query), not just the count of accepted
         # tokens this step.
         tl.store(out_seq_lens_ptr + req_idx, last_valid_pos + 1 + num_query_per_req)
+        # Copy sampling state.
+        tl.store(
+            out_temperature_ptr + req_state_idx,
+            tl.load(temperature_ptr + req_state_idx),
+        )
+        tl.store(out_seeds_ptr + req_state_idx, tl.load(seeds_ptr + req_state_idx))
         if req_idx == num_reqs - 1:
             # Pad per-request buffers to max_num_reqs for CUDA graph safety.
             last_query_end = num_reqs * num_query_per_req
@@ -626,6 +634,8 @@ def prepare_dflash_inputs(
     sample_indices: torch.Tensor,
     sample_pos: torch.Tensor,
     sample_idx_mapping: torch.Tensor,
+    temperature: torch.Tensor,
+    seeds: torch.Tensor,
     input_batch: InputBatch,
     # [num_reqs]
     num_sampled: torch.Tensor,
@@ -635,6 +645,10 @@ def prepare_dflash_inputs(
     last_sampled: torch.Tensor,
     # [max_num_reqs]
     next_prefill_tokens: torch.Tensor,
+    # [max_num_reqs]
+    input_temperature: torch.Tensor,
+    # [max_num_reqs]
+    input_seeds: torch.Tensor,
     # [max_num_reqs, max_num_blocks]
     block_table: torch.Tensor,
     block_size: int,
@@ -665,6 +679,8 @@ def prepare_dflash_inputs(
         sample_indices,
         sample_pos,
         sample_idx_mapping,
+        temperature,
+        seeds,
         input_batch.positions,
         input_batch.query_start_loc,
         input_batch.idx_mapping,
@@ -672,6 +688,8 @@ def prepare_dflash_inputs(
         next_prefill_tokens,
         num_sampled,
         num_rejected,
+        input_temperature,
+        input_seeds,
         block_table,
         block_table.stride(0),
         parallel_drafting_token_id,

vllm/v1/worker/gpu_model_runner.py

@@ -6934,6 +6934,10 @@ class GPUModelRunner(
                 num_active_loras=desc.num_active_loras,
                 profile_seq_lens=profile_seq_lens,
             )
+        if num_warmups > 0:
+            # Warmups may use auxiliary streams. Ensure all of their work has
+            # completed before beginning CUDA graph capture.
+            torch.accelerator.synchronize()
         with (
             profiler,
             torch.profiler.record_function(