[New model] Kimi K3 (#50000)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Ziming Huang <zelda.huanghuang@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Summer Yang <girasoleyang@gmail.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com> Co-authored-by: Bowen Wang <abmfy@icloud.com> Co-authored-by: gnovack <novackgm@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: xiaozhoupy <peiyuanzhou1994@gmail.com> Co-authored-by: Roy Wang <yasong.wang@inferact.ai> Co-authored-by: Jeff (Junze) Ma <93145857+majunze2001@users.noreply.github.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com>
Jiangyun Zhu committed
Jul 30, 2026 at 18:49 UTC
aeeb36b1f17145975c6713242f2447bb8b98782b
82 files changed
+2931
-1349
cmake/external_projects/deepgemm.cmake
+3
-3
@@ -28,9 +28,9 @@ if(DEEPGEMM_SRC_DIR)
28
message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
29
else()
30
# Keep in sync with tools/install_deepgemm.sh
31
- set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
32
- # NOTE: This is currently targeting nv-dev branch due to sm120 support
33
- set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a")
31
+ set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git")
32
+ # TODO: switch to nv_dev branch after it support situ
33
+ set(_DEEPGEMM_UPSTREAM_TAG "f5a76426fa084087169693fd0cd815223576d6e9")
34
35
set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
36
if(NOT _deepgemm_fc_root)
docs/models/supported_models.md
+1
@@ -547,6 +547,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
547
| `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ |
548
| `KimiAudioForConditionalGeneration` | Kimi-Audio | T + A<sup>+</sup> | `moonshotai/Kimi-Audio-7B-Instruct` | | ✅︎ |
549
| `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I<sup>+</sup> | `moonshotai/Kimi-K2.5` | | ✅︎ |
550
+| `KimiK3ForConditionalGeneration` | Kimi-K3 | T + I<sup>+</sup> | `moonshotai/Kimi-K3` | | ✅︎ |
551
| `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
552
| `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I<sup>+</sup> | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ |
553
| `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I<sup>+</sup> | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |
pyproject.toml
+2
@@ -155,6 +155,8 @@ view_seperator = "view_seperator"
155
inverse_std_variences = "inverse_std_variences"
156
157
[tool.typos.default.extend-words]
158
+Hel = "Hel"
159
+wether = "wether"
160
iy = "iy"
161
indx = "indx"
162
# intel cpu features
requirements/cuda.txt
+2
-2
@@ -16,8 +16,8 @@ PyNvVideoCodec==2.0.4
16
--extra-index-url https://flashinfer.ai/whl/
17
flashinfer-python==0.6.15.post1
18
flashinfer-cubin==0.6.15.post1
19
-apache-tvm-ffi==0.1.10
20
-tilelang==0.1.9
19
+apache-tvm-ffi==0.1.11
20
+tilelang==0.1.12
21
nvidia-cudnn-frontend>=1.19.1
22
# Required for LLM_NVTX_SCOPES_FOR_PROFILING=1
23
nvtx==0.2.15
requirements/test/cuda.txt
+1
-1
@@ -45,7 +45,7 @@ anyio==4.14.1
45
# sse-starlette
46
# starlette
47
# watchfiles
48
-apache-tvm-ffi==0.1.10
48
+apache-tvm-ffi==0.1.11
49
# via
50
# -c requirements/cuda.txt
51
# xgrammar
tests/kernels/moe/test_deepgemm.py
+32
@@ -25,6 +25,9 @@ from vllm.model_executor.layers.fused_moe.config import (
25
FusedMoEQuantDesc,
26
fp8_w8a8_moe_quant_config,
27
)
28
+from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
29
+ deepgemm_moe_permute,
30
+)
31
from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
32
TritonOrDeepGemmExperts,
33
)
@@ -41,6 +44,35 @@ from vllm.utils.deep_gemm import (
44
BLOCK_SIZE = [128, 128]
45
46
47
+@pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels")
48
+def test_deepgemm_moe_permute_initializes_padding_scales(workspace_init):
49
+ hidden_states = torch.randn(2, 128, device="cuda", dtype=torch.bfloat16)
50
+ activations, scales = per_token_group_quant_fp8(
51
+ hidden_states,
52
+ group_size=128,
53
+ use_ue8m0=True,
54
+ )
55
+ topk_ids = torch.tensor([[0], [1]], device="cuda", dtype=torch.int64)
56
+
57
+ _, permuted_scales, expert_ids, _, _ = deepgemm_moe_permute(
58
+ aq=activations,
59
+ aq_scale=scales,
60
+ topk_ids=topk_ids,
61
+ local_num_experts=2,
62
+ expert_map=None,
63
+ expert_tokens_meta=None,
64
+ )
65
+
66
+ padding = expert_ids < 0
67
+ assert padding.any()
68
+ torch.testing.assert_close(
69
+ permuted_scales[padding],
70
+ torch.zeros_like(permuted_scales[padding]),
71
+ rtol=0,
72
+ atol=0,
73
+ )
74
+
75
+
76
def make_block_quant_fp8_weights(
77
e: int,
78
n: int,
tests/models/registry.py
+18
@@ -1026,6 +1026,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
1026
"moonshotai/Kimi-K2.5",
1027
trust_remote_code=True,
1028
),
1029
+ "KimiK3ForConditionalGeneration": _HfExamplesInfo(
1030
+ "moonshotai/Kimi-K3",
1031
+ trust_remote_code=True,
1032
+ ),
1033
"KimiVLForConditionalGeneration": _HfExamplesInfo(
1034
"moonshotai/Kimi-VL-A3B-Instruct",
1035
extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"},
@@ -1466,6 +1470,14 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1470
is_available_online=False,
1471
use_original_num_layers=True, # DSpark has >1 draft block
1472
),
1473
+ "K3DSparkModel": _HfExamplesInfo(
1474
+ "moonshotai/Kimi-K3",
1475
+ speculative_model="Inferact/Kimi-K3-DSpark",
1476
+ use_original_num_layers=True, # DSpark has >1 draft block
1477
+ trust_remote_code=True,
1478
+ # FIXME: Investigate the NVML failure in CI.
1479
+ is_available_online=False,
1480
+ ),
1481
"Qwen3DSparkModel": _HfExamplesInfo(
1482
"Qwen/Qwen3-8B",
1483
speculative_model="deepseek-ai/dspark_qwen3_8b_block7",
@@ -1664,6 +1676,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1676
trust_remote_code=True,
1677
max_model_len=4096,
1678
),
1679
+ "KimiK3MTPModel": _HfExamplesInfo(
1680
+ "moonshotai/Kimi-K3",
1681
+ speculative_model="moonshotai/Kimi-K3",
1682
+ trust_remote_code=True,
1683
+ is_available_online=False,
1684
+ ),
1685
"LongCatFlashMTPModel": _HfExamplesInfo(
1686
"meituan-longcat/LongCat-Flash-Chat",
1687
trust_remote_code=True,
tests/models/test_dspark_mla.py
new
+143
@@ -0,0 +1,143 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+import torch.nn as nn
9
+
10
+from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
11
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
12
+from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
13
+from vllm.model_executor.models.registry import ModelRegistry
14
+from vllm.models.kimi_k3.nvidia import dspark_mla
15
+from vllm.models.kimi_k3.nvidia.dspark_mla import K3DSparkForCausalLM, K3DSparkModel
16
+
17
+
18
+def test_dspark_mla_uses_compile_free_model_entrypoint():
19
+ assert ModelRegistry._try_load_model_cls("K3DSparkModel") is K3DSparkForCausalLM
20
+ assert not issubclass(K3DSparkModel, TorchCompileWithNoGuardsWrapper)
21
+
22
+
23
+@pytest.mark.parametrize(
24
+ ("checkpoint_name", "runtime_name", "shard_id"),
25
+ [
26
+ (
27
+ "layers.0.self_attn.q_a_proj.weight",
28
+ "model.layers.0.self_attn.fused_qkv_a_proj.weight",
29
+ 0,
30
+ ),
31
+ (
32
+ "layers.0.self_attn.kv_a_proj_with_mqa.weight",
33
+ "model.layers.0.self_attn.fused_qkv_a_proj.weight",
34
+ 1,
35
+ ),
36
+ (
37
+ "layers.0.mlp.gate_proj.weight",
38
+ "model.layers.0.mlp.gate_up_proj.weight",
39
+ 0,
40
+ ),
41
+ (
42
+ "layers.0.mlp.up_proj.weight",
43
+ "model.layers.0.mlp.gate_up_proj.weight",
44
+ 1,
45
+ ),
46
+ ("context_proj.weight", "model.context_proj.weight", None),
47
+ ],
48
+)
49
+def test_dspark_mla_checkpoint_weight_mapping(checkpoint_name, runtime_name, shard_id):
50
+ assert K3DSparkForCausalLM.hf_to_vllm_mapper._map_name_with_shard(
51
+ checkpoint_name
52
+ ) == (runtime_name, shard_id)
53
+
54
+
55
+def test_dspark_mla_shares_frozen_target_weights_and_skips_training_head():
56
+ assert not K3DSparkForCausalLM.has_own_embed_tokens
57
+ assert not K3DSparkForCausalLM.has_own_lm_head
58
+ assert set(K3DSparkForCausalLM.checkpoint_skip_substrs) == {
59
+ "confidence_head",
60
+ "embed_tokens",
61
+ "lm_head",
62
+ }
63
+
64
+
65
+@pytest.mark.cpu_test
66
+def test_dspark_markov_head_is_replicated(
67
+ monkeypatch: pytest.MonkeyPatch,
68
+):
69
+ from vllm.model_executor.layers import logits_processor, vocab_parallel_embedding
70
+
71
+ monkeypatch.setattr(
72
+ vocab_parallel_embedding, "get_tensor_model_parallel_rank", lambda: 3
73
+ )
74
+ monkeypatch.setattr(
75
+ vocab_parallel_embedding,
76
+ "get_tensor_model_parallel_world_size",
77
+ lambda: 8,
78
+ )
79
+ monkeypatch.setattr(
80
+ logits_processor,
81
+ "get_current_vllm_config",
82
+ lambda: SimpleNamespace(model_config=None),
83
+ )
84
+
85
+ head = DSparkMarkovHead(128, 128, 8, prefix="markov_head")
86
+ assert head.markov_w2.tp_size == 1
87
+ assert head.markov_w1.weight.shape == (128, 8)
88
+ assert head.markov_w2.weight.shape == (128, 8)
89
+
90
+ def fail_collective(*args, **kwargs):
91
+ raise AssertionError("replicated Markov head must not invoke TP collectives")
92
+
93
+ monkeypatch.setattr(
94
+ vocab_parallel_embedding,
95
+ "tensor_model_parallel_all_reduce",
96
+ fail_collective,
97
+ )
98
+ logits_processor = LogitsProcessor(128)
99
+ monkeypatch.setattr(logits_processor, "_gather_logits", fail_collective)
100
+
101
+ markov_embed = head.embed(torch.tensor([1, 2]))
102
+ bias = head.bias(markov_embed, logits_processor)
103
+ assert markov_embed.shape == (2, 8)
104
+ assert bias.shape == (2, 128)
105
+
106
+
107
+@pytest.mark.cpu_test
108
+def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch):
109
+ markov_head_calls = []
110
+
111
+ class DummyModule(nn.Module):
112
+ def __init__(self, *args, **kwargs):
113
+ super().__init__()
114
+
115
+ def make_markov_head(*args, **kwargs):
116
+ markov_head_calls.append((args, kwargs))
117
+ return DummyModule()
118
+
119
+ monkeypatch.setattr(dspark_mla, "get_draft_quant_config", lambda _: None)
120
+ monkeypatch.setattr(dspark_mla, "ReplicatedLinear", DummyModule)
121
+ monkeypatch.setattr(dspark_mla, "RMSNorm", DummyModule)
122
+ monkeypatch.setattr(dspark_mla, "K3DSparkDecoderLayer", DummyModule)
123
+ monkeypatch.setattr(dspark_mla, "DSparkMarkovHead", make_markov_head)
124
+
125
+ config = SimpleNamespace(
126
+ target_hidden_size=16,
127
+ num_target_layers=2,
128
+ hidden_size=8,
129
+ rms_norm_eps=1e-6,
130
+ num_hidden_layers=1,
131
+ vocab_size=128,
132
+ draft_vocab_size=128,
133
+ markov_rank=4,
134
+ )
135
+ vllm_config = SimpleNamespace(
136
+ speculative_config=SimpleNamespace(
137
+ draft_model_config=SimpleNamespace(hf_config=config)
138
+ )
139
+ )
140
+
141
+ K3DSparkModel(vllm_config=vllm_config, start_layer_id=0, prefix="model")
142
+
143
+ assert len(markov_head_calls) == 1
tests/test_config.py
+20
@@ -1235,6 +1235,26 @@ def test_vllm_config_defaults_are_none():
1235
assert getattr(config.compilation_config, k) is None
1236
1237
1238
+def test_validate_mamba_align_subblock_prefill():
1239
+ """Align mode permits configured prefill chunks smaller than a block."""
1240
+ config = SimpleNamespace(
1241
+ cache_config=SimpleNamespace(
1242
+ block_size=11392,
1243
+ mamba_cache_mode="align",
1244
+ ),
1245
+ parallel_config=SimpleNamespace(
1246
+ decode_context_parallel_size=1,
1247
+ ),
1248
+ scheduler_config=SimpleNamespace(
1249
+ max_num_batched_tokens=8192,
1250
+ long_prefill_token_threshold=4096,
1251
+ disable_chunked_mm_input=False,
1252
+ ),
1253
+ )
1254
+
1255
+ VllmConfig.validate_block_size(config)
1256
+
1257
+
1258
@pytest.mark.parametrize(
1259
("model_id", "compilation_config", "optimization_level"),
1260
[
tests/transformers_utils/test_dspark_mla_config.py
new
+165
@@ -0,0 +1,165 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import json
5
+
6
+import pytest
7
+
8
+from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
9
+from vllm.transformers_utils.config import get_config
10
+from vllm.transformers_utils.configs.k3_dspark import K3DSparkConfig
11
+
12
+
13
+def _write_dspark_config(path, **overrides):
14
+ path.mkdir()
15
+ config = {
16
+ "architectures": ["K3DSparkModel"],
17
+ "model_type": "k3_dspark",
18
+ "hidden_size": 7168,
19
+ "intermediate_size": 14336,
20
+ "num_hidden_layers": 5,
21
+ "num_attention_heads": 64,
22
+ "num_key_value_heads": 64,
23
+ "q_lora_rank": 1536,
24
+ "kv_lora_rank": 512,
25
+ "qk_nope_head_dim": 128,
26
+ "qk_rope_head_dim": 64,
27
+ "v_head_dim": 128,
28
+ "vocab_size": 163840,
29
+ "rms_norm_eps": 1e-5,
30
+ "max_position_embeddings": 32768,
31
+ "rope_theta": 50000.0,
32
+ "num_target_layers": 5,
33
+ "target_hidden_size": 7168,
34
+ "target_num_hidden_layers": 93,
35
+ "target_layer_ids": [2, 23, 47, 71, 89],
36
+ "markov_rank": 256,
37
+ "draft_vocab_size": 163840,
38
+ "torch_dtype": "bfloat16",
39
+ }
40
+ config.update(overrides)
41
+ (path / "config.json").write_text(json.dumps(config))
42
+
43
+
44
+def _write_target_config(path):
45
+ path.mkdir()
46
+ config = {
47
+ "architectures": ["LlamaForCausalLM"],
48
+ "model_type": "llama",
49
+ "hidden_size": 7168,
50
+ "intermediate_size": 14336,
51
+ "num_hidden_layers": 93,
52
+ "num_attention_heads": 56,
53
+ "num_key_value_heads": 8,
54
+ "vocab_size": 163840,
55
+ "max_position_embeddings": 32768,
56
+ "torch_dtype": "bfloat16",
57
+ }
58
+ (path / "config.json").write_text(json.dumps(config))
59
+
60
+
61
+def test_dspark_mla_config_loads_from_local_json(tmp_path):
62
+ draft_path = tmp_path / "draft"
63
+ _write_dspark_config(draft_path)
64
+
65
+ config = get_config(draft_path, trust_remote_code=False)
66
+
67
+ assert isinstance(config, K3DSparkConfig)
68
+ assert config.model_type == "k3_dspark"
69
+ assert config.architectures == ["K3DSparkModel"]
70
+ assert config.hidden_act == "silu"
71
+ assert config.rope_parameters == {
72
+ "rope_type": "default",
73
+ "rope_theta": 50000.0,
74
+ }
75
+ assert config.n_routed_experts == 0
76
+ assert config.draft_vocab_size == config.vocab_size
77
+
78
+
79
+@pytest.mark.parametrize(
80
+ "overrides",
81
+ [
82
+ {"mla_use_nope": True},
83
+ {"mla_use_output_gate": True},
84
+ {"mla_use_qk_norm": True},
85
+ {"dspark_bonus_anchor": True},
86
+ {"q_lora_rank": None},
87
+ {"draft_vocab_size": 8192},
88
+ {"target_layer_ids": []},
89
+ {"num_target_layers": 4},
90
+ ],
91
+)
92
+def test_dspark_mla_rejects_unsupported_checkpoint_options(tmp_path, overrides):
93
+ draft_path = tmp_path / "draft"
94
+ _write_dspark_config(draft_path, **overrides)
95
+
96
+ with pytest.raises(ValueError, match="MLA DSpark"):
97
+ get_config(draft_path, trust_remote_code=False)
98
+
99
+
100
+def test_dspark_mla_uses_latent_kv_geometry(tmp_path):
101
+ draft_path = tmp_path / "draft"
102
+ _write_dspark_config(draft_path)
103
+ model_config = ModelConfig(
104
+ model=str(draft_path),
105
+ tokenizer_mode="skip",
106
+ runner="draft",
107
+ max_model_len=32768,
108
+ )
109
+
110
+ assert model_config.is_deepseek_mla
111
+ assert model_config.use_mla
112
+ assert model_config.get_head_size() == 576
113
+ # external_launcher skips ParallelConfig's local-GPU-count check so the
114
+ # config logic can be exercised at TP8 on a single-GPU test node.
115
+ parallel_config = ParallelConfig(
116
+ tensor_parallel_size=8, distributed_executor_backend="external_launcher"
117
+ )
118
+ assert model_config.get_num_kv_heads(parallel_config) == 1
119
+ assert model_config.get_num_attention_heads(parallel_config) == 8
120
+ assert model_config.get_num_experts() == 0
121
+
122
+
123
+def test_dspark_mla_speculative_config_preserves_architecture(tmp_path):
124
+ target_path = tmp_path / "target"
125
+ draft_path = tmp_path / "draft"
126
+ _write_target_config(target_path)
127
+ _write_dspark_config(draft_path)
128
+ target_config = ModelConfig(
129
+ model=str(target_path), tokenizer_mode="skip", max_model_len=32768
130
+ )
131
+ speculative_config = SpeculativeConfig(
132
+ model=str(draft_path),
133
+ method="dspark",
134
+ num_speculative_tokens=8,
135
+ target_model_config=target_config,
136
+ target_parallel_config=ParallelConfig(),
137
+ )
138
+
139
+ assert speculative_config.parallel_drafting
140
+ assert speculative_config.draft_model_config.architectures == ["K3DSparkModel"]
141
+ assert speculative_config.draft_model_config.hf_config.model_type == "k3_dspark"
142
+ assert speculative_config.draft_model_config.use_mla
143
+
144
+
145
+def test_dspark_mla_rejects_decode_context_parallelism(tmp_path):
146
+ target_path = tmp_path / "target"
147
+ draft_path = tmp_path / "draft"
148
+ _write_target_config(target_path)
149
+ _write_dspark_config(draft_path)
150
+ target_config = ModelConfig(
151
+ model=str(target_path), tokenizer_mode="skip", max_model_len=32768
152
+ )
153
+
154
+ with pytest.raises(ValueError, match="does not currently support decode context"):
155
+ SpeculativeConfig(
156
+ model=str(draft_path),
157
+ method="dspark",
158
+ num_speculative_tokens=8,
159
+ target_model_config=target_config,
160
+ target_parallel_config=ParallelConfig(
161
+ tensor_parallel_size=2,
162
+ decode_context_parallel_size=2,
163
+ distributed_executor_backend="external_launcher",
164
+ ),
165
+ )
tests/v1/attention/test_mla_backends.py
+5
-1
@@ -73,7 +73,11 @@ DEVICE_TYPE = current_platform.device_type
73
def test_mla_kv_cache_spec_uses_layer_cache_dtype(
74
cache_dtype: str, expected_quant_mode: KVQuantMode
75
):
76
- layer = SimpleNamespace(kv_cache_dtype=cache_dtype, head_size=576)
76
+ layer = SimpleNamespace(
77
+ kv_cache_dtype=cache_dtype,
78
+ head_size=576,
79
+ non_causal_multi_token_decode=False,
80
+ )
81
vllm_config = SimpleNamespace(
82
cache_config=SimpleNamespace(block_size=64), model_config=None
83
)
tests/v1/attention/test_mla_noncausal.py
new
+126
@@ -0,0 +1,126 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.model_executor.layers.attention.mla_attention import (
10
+ MLACommonMetadata,
11
+ MLACommonMetadataBuilder,
12
+ QueryLenSupport,
13
+)
14
+from vllm.v1.attention.backend import CommonAttentionMetadata
15
+from vllm.v1.kv_cache_interface import MLAAttentionSpec
16
+
17
+
18
+class _NonCausalMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
19
+ supports_non_causal_multi_token_decode = True
20
+
21
+
22
+def _metadata(
23
+ query_start_loc: list[int], num_tokens: int | None = None
24
+) -> CommonAttentionMetadata:
25
+ num_reqs = len(query_start_loc) - 1
26
+ num_tokens = query_start_loc[-1] if num_tokens is None else num_tokens
27
+ return CommonAttentionMetadata(
28
+ query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32),
29
+ query_start_loc_cpu=torch.tensor(query_start_loc, dtype=torch.int32),
30
+ seq_lens=torch.arange(1, num_reqs + 1, dtype=torch.int32) * 100 + 8,
31
+ num_reqs=num_reqs,
32
+ num_actual_tokens=num_tokens,
33
+ max_query_len=max(
34
+ end - start for start, end in zip(query_start_loc, query_start_loc[1:])
35
+ ),
36
+ max_seq_len=num_reqs * 100 + 8,
37
+ block_table_tensor=torch.arange(num_reqs * 3, dtype=torch.int32).view(
38
+ num_reqs, 3
39
+ ),
40
+ slot_mapping=torch.arange(num_tokens),
41
+ causal=False,
42
+ seq_lens_cpu_upper_bound=None,
43
+ )
44
+
45
+
46
+def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder:
47
+ builder = object.__new__(_NonCausalMLAMetadataBuilder)
48
+ builder.device = torch.device("cpu")
49
+ builder.reorder_batch_threshold = 1
50
+ builder.query_len_support = QueryLenSupport.SINGLE_ONLY
51
+ builder.non_causal_multi_token_decode = marked
52
+ builder.dcp_world_size = 1
53
+ builder.metadata_cls = MLACommonMetadata
54
+ builder.model_config = SimpleNamespace(
55
+ dtype=torch.bfloat16, get_head_size=lambda: 576
56
+ )
57
+ return builder
58
+
59
+
60
+def test_noncausal_block_uses_decode_without_cpu_lengths():
61
+ common_metadata = _metadata([0, 8, 16])
62
+ metadata = _builder().build(0, common_metadata)
63
+
64
+ assert metadata.num_decodes == 2
65
+ assert metadata.num_decode_tokens == 16
66
+ assert metadata.num_prefills == 0
67
+ assert metadata.prefill is None
68
+ assert metadata.decode is not None
69
+ assert metadata.decode.block_table.shape == (2, 3)
70
+ assert metadata.decode.seq_lens.shape == (2,)
71
+ assert torch.equal(metadata.decode.block_table, common_metadata.block_table_tensor)
72
+ assert torch.equal(metadata.decode.seq_lens, common_metadata.seq_lens)
73
+ assert not metadata.causal
74
+
75
+
76
+def test_noncausal_support_is_explicit_and_uniform():
77
+ with pytest.raises(ValueError, match="explicitly supported"):
78
+ _builder(marked=False).build(0, _metadata([0, 8, 16]))
79
+ with pytest.raises(ValueError, match="uniform query block"):
80
+ _builder().build(0, _metadata([0, 3, 8]))
81
+
82
+
83
+def test_noncausal_block_allows_trailing_cudagraph_padding():
84
+ common_metadata = _metadata([0, 8, 16, 16], num_tokens=24)
85
+ common_metadata.seq_lens[-1] = 0
86
+
87
+ metadata = _builder().build(0, common_metadata)
88
+
89
+ assert metadata.num_decodes == 3
90
+ assert metadata.num_decode_tokens == 24
91
+ assert metadata.decode is not None
92
+ assert metadata.decode.seq_lens.tolist() == [108, 208, 0]
93
+
94
+
95
+def test_noncausal_block_rejects_non_trailing_padding():
96
+ with pytest.raises(ValueError, match="uniform query block"):
97
+ _builder().build(0, _metadata([0, 8, 8, 16], num_tokens=24))
98
+
99
+
100
+def test_noncausal_decode_metadata_keeps_live_request_buffers():
101
+ common_metadata = _metadata([0, 8, 16])
102
+ metadata = _builder().build(0, common_metadata)
103
+
104
+ assert metadata.decode is not None
105
+ assert metadata.decode.seq_lens.data_ptr() == common_metadata.seq_lens.data_ptr()
106
+ assert (
107
+ metadata.decode.block_table.data_ptr()
108
+ == common_metadata.block_table_tensor.data_ptr()
109
+ )
110
+
111
+
112
+def test_mla_cache_marker_is_promoted_to_group_capability():
113
+ kwargs = {
114
+ "block_size": 64,
115
+ "num_kv_heads": 1,
116
+ "head_size": 576,
117
+ "dtype": torch.bfloat16,
118
+ }
119
+ marked = MLAAttentionSpec(**kwargs, non_causal_multi_token_decode=True)
120
+ unmarked = MLAAttentionSpec(**kwargs)
121
+
122
+ assert MLAAttentionSpec.merge([marked, marked]).non_causal_multi_token_decode
123
+ assert MLAAttentionSpec.merge([marked, unmarked]).non_causal_multi_token_decode
124
+ assert not MLAAttentionSpec.merge(
125
+ [unmarked, unmarked]
126
+ ).non_causal_multi_token_decode
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+74
@@ -40,6 +40,8 @@ def test_mamba_align_split_partial_tail_schedule():
40
hash_block_size = 32
41
mock = SimpleNamespace(
42
cache_config=SimpleNamespace(block_size=block_size),
43
+ max_num_scheduled_tokens=8192,
44
+ scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
45
use_eagle=False,
46
hash_block_size=hash_block_size,
47
mamba_partial_cache_hit=True,
@@ -75,6 +77,78 @@ def test_mamba_align_split_partial_tail_schedule():
77
assert split(self=mock, request=req2, num_new_tokens=1000) == 512
78
79
80
+def test_mamba_align_split_when_block_exceeds_scheduling_budget():
81
+ """Sub-block chunks make progress only when no step can fit a full block."""
82
+ block_size = 11392
83
+ token_budget = 8192
84
+ prompt_length = 30000
85
+ mock = SimpleNamespace(
86
+ cache_config=SimpleNamespace(block_size=block_size),
87
+ max_num_scheduled_tokens=token_budget,
88
+ scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
89
+ use_eagle=False,
90
+ hash_block_size=32,
91
+ mamba_partial_cache_hit=False,
92
+ )
93
+ req = make_request("0", [0] * prompt_length, 32, sha256)
94
+ split = Scheduler._mamba_block_aligned_split
95
+
96
+ mock.max_num_scheduled_tokens = block_size
97
+ assert split(self=mock, request=req, num_new_tokens=token_budget) == 0
98
+ mock.max_num_scheduled_tokens = token_budget
99
+
100
+ scheduled_chunks = []
101
+ while req.num_computed_tokens < prompt_length:
102
+ num_new_tokens = min(token_budget, prompt_length - req.num_computed_tokens)
103
+ num_scheduled_tokens = split(
104
+ self=mock,
105
+ request=req,
106
+ num_new_tokens=num_new_tokens,
107
+ )
108
+ assert 0 < num_scheduled_tokens <= token_budget
109
+ scheduled_chunks.append(num_scheduled_tokens)
110
+ req.num_computed_tokens += num_scheduled_tokens
111
+
112
+ assert scheduled_chunks == [8192, 3200, 8192, 3200, 7216]
113
+
114
+
115
+def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():
116
+ """A long-prefill cap below the block size permits sub-block progress."""
117
+ block_size = 512
118
+ token_budget = 8192
119
+ long_prefill_threshold = 384
120
+ prompt_length = 1300
121
+ mock = SimpleNamespace(
122
+ cache_config=SimpleNamespace(block_size=block_size),
123
+ max_num_scheduled_tokens=token_budget,
124
+ scheduler_config=SimpleNamespace(
125
+ long_prefill_token_threshold=long_prefill_threshold
126
+ ),
127
+ use_eagle=False,
128
+ hash_block_size=32,
129
+ mamba_partial_cache_hit=False,
130
+ )
131
+ req = make_request("0", [0] * prompt_length, 32, sha256)
132
+ split = Scheduler._mamba_block_aligned_split
133
+
134
+ scheduled_chunks = []
135
+ while req.num_computed_tokens < prompt_length:
136
+ num_new_tokens = min(
137
+ long_prefill_threshold,
138
+ prompt_length - req.num_computed_tokens,
139
+ )
140
+ num_scheduled_tokens = split(
141
+ self=mock,
142
+ request=req,
143
+ num_new_tokens=num_new_tokens,
144
+ )
145
+ assert 0 < num_scheduled_tokens <= long_prefill_threshold
146
+ scheduled_chunks.append(num_scheduled_tokens)
147
+ req.num_computed_tokens += num_scheduled_tokens
148
+
149
+ assert scheduled_chunks == [384, 128, 384, 128, 276]
150
+
151
+
152
def test_hybrid_mamba_align_partial_hash_hit():
153
hash_block_size = 2
154
mamba_block_size = 2 * hash_block_size
tests/v1/core/test_prefix_caching.py
+2
@@ -1078,6 +1078,8 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():
1078
# Create minimal mock with just the needed attributes
1079
mock = SimpleNamespace(
1080
cache_config=SimpleNamespace(block_size=block_size),
1081
+ max_num_scheduled_tokens=3 * block_size,
1082
+ scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
1083
use_eagle=False,
1084
hash_block_size=block_size,
1085
mamba_partial_cache_hit=False,
tests/v1/cudagraph/test_breakable_cudagraph.py
+48
@@ -333,6 +333,54 @@ def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream):
333
assert torch.equal(x, torch.full((4,), 15.0, device="cuda"))
334
335
336
+def test_eager_attention_inside_multistream_overlap(cuda_capture_stream):
337
+ """Handle an eager attention break inside a multi-stream overlap region."""
338
+ from vllm.compilation.breakable_cudagraph import (
339
+ BreakableCUDAGraphCapture,
340
+ eager_break_during_capture,
341
+ )
342
+ from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
343
+
344
+ x = torch.zeros(1024, device="cuda")
345
+ output = torch.empty_like(x)
346
+ aux_stream = torch.cuda.Stream()
347
+ main_event = torch.cuda.Event()
348
+ aux_event = torch.cuda.Event()
349
+
350
+ @eager_break_during_capture
351
+ def attention(query: torch.Tensor, out: torch.Tensor) -> None:
352
+ torch.add(query, 3.0, out=out)
353
+
354
+ def attention_frontend() -> torch.Tensor:
355
+ query = x * 2.0
356
+ attention_output = torch.empty_like(x)
357
+ attention(query, attention_output)
358
+ return attention_output
359
+
360
+ cap = BreakableCUDAGraphCapture()
361
+ with cap:
362
+ attention_output, gate = maybe_execute_in_parallel(
363
+ attention_frontend,
364
+ lambda: x * 5.0,
365
+ main_event,
366
+ aux_event,
367
+ aux_stream,
368
+ )
369
+ torch.add(attention_output, gate, out=output)
370
+
371
+ assert cap.num_graphs == 2
372
+ assert cap.num_eager_breaks == 1
373
+
374
+ for value in (1.0, 7.0, 19.0):
375
+ x.fill_(value)
376
+ cap.replay()
377
+ cuda_capture_stream.synchronize()
378
+ torch.testing.assert_close(
379
+ output,
380
+ torch.full_like(output, value * 7.0 + 3.0),
381
+ )
382
+
383
+
384
# ---------------------------------------------------------------------------
385
# Replay ordering
386
# ---------------------------------------------------------------------------
tests/v1/spec_decode/test_llm_base_proposer_sampling.py
+22
@@ -1,6 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
import torch
8
9
from vllm.platforms import current_platform
@@ -8,6 +11,7 @@ from vllm.utils.torch_utils import set_random_seed
11
from vllm.v1.sample.logits_processor import LogitsProcessors
12
from vllm.v1.sample.metadata import SamplingMetadata
13
from vllm.v1.spec_decode.llm_base_proposer import (
14
+ SpecDecodeBaseProposer,
15
compute_probs_and_sample_next_token,
16
)
17
@@ -68,3 +72,21 @@ def test_compute_probs_and_sample_next_token_uses_fp64_exponential_race():
72
73
assert torch.equal(actual_ids, expected_ids)
74
assert torch.allclose(actual_probs, probs)
75
+
76
+
77
+@pytest.mark.parametrize(
78
+ ("architecture", "expected"),
79
+ [
80
+ ("DeepSeekMTPModel", True),
81
+ ("KimiK3MTPModel", True),
82
+ ("MiniMaxM3ForCausalLM", False),
83
+ ],
84
+)
85
+def test_mtp_model_returns_tuple(architecture: str, expected: bool):
86
+ proposer = object.__new__(SpecDecodeBaseProposer)
87
+ proposer.method = "mtp"
88
+ proposer.draft_model_config = SimpleNamespace(
89
+ hf_config=SimpleNamespace(architectures=[architecture])
90
+ )
91
+
92
+ assert proposer.model_returns_tuple() is expected
tools/install_deepgemm.sh
+2
-2
@@ -6,9 +6,9 @@ set -e
6
7
# Default values
8
# Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake
9
-DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git"
9
+DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git"
10
# NOTE: This is currently targeting nv-dev branch due to sm120 support
11
-DEEPGEMM_GIT_REF="a6b593d2826719dcf4892609af7b84ee23aaf32a"
11
+DEEPGEMM_GIT_REF="f5a76426fa084087169693fd0cd815223576d6e9"
12
WHEEL_DIR=""
13
14
# Parse command line arguments
vllm/_aiter_ops.py
+15
@@ -133,6 +133,8 @@ def _rocm_aiter_fused_moe_impl(
133
bias2: torch.Tensor | None = None,
134
moe_sorting_dispatch_policy: int = 0,
135
swiglu_limit: float = 0.0,
136
+ beta: float | None = None,
137
+ linear_beta: float | None = None,
138
) -> torch.Tensor:
139
from aiter import ActivationType, QuantType
140
from aiter.fused_moe import fused_moe
@@ -143,6 +145,12 @@ def _rocm_aiter_fused_moe_impl(
145
extra_kwargs: dict = {}
146
if gate_mode and rocm_aiter_ops.fused_moe_supports_gate_mode():
147
extra_kwargs["gate_mode"] = gate_mode
148
+ if (
149
+ getattr(ActivationType, "Situv2", None) is not None
150
+ and activation == ActivationType.Situv2
151
+ ):
152
+ extra_kwargs["beta"] = beta
153
+ extra_kwargs["linear_beta"] = linear_beta
154
155
return fused_moe(
156
hidden_states,
@@ -193,6 +201,8 @@ def _rocm_aiter_fused_moe_fake(
201
bias2: torch.Tensor | None = None,
202
moe_sorting_dispatch_policy: int = 0,
203
swiglu_limit: float = 0.0,
204
+ beta: float | None = None,
205
+ linear_beta: float | None = None,
206
) -> torch.Tensor:
207
if output_dtype is not None:
208
return torch.empty_like(hidden_states, dtype=output_dtype)
@@ -1618,6 +1628,7 @@ class rocm_aiter_ops:
1628
"silu": ActivationType.Silu,
1629
"gelu": ActivationType.Gelu,
1630
"swiglu": ActivationType.Swiglu,
1631
+ "situ": getattr(ActivationType, "Situv2", None),
1632
}
1633
return mapping.get(name)
1634
@@ -2222,6 +2233,8 @@ class rocm_aiter_ops:
2233
bias2: torch.Tensor | None = None,
2234
moe_sorting_dispatch_policy: int = 0,
2235
swiglu_limit: float = 0.0,
2236
+ beta: float | None = None,
2237
+ linear_beta: float | None = None,
2238
) -> torch.Tensor:
2239
return torch.ops.vllm.rocm_aiter_fused_moe(
2240
hidden_states,
@@ -2246,6 +2259,8 @@ class rocm_aiter_ops:
2259
bias2,
2260
moe_sorting_dispatch_policy,
2261
swiglu_limit,
2262
+ beta,
2263
+ linear_beta,
2264
)
2265
2266
@staticmethod
vllm/config/compilation.py
-1
@@ -771,7 +771,6 @@ class CompilationConfig:
771
"vllm::qwen_gdn_attention_core",
772
"vllm::gdn_attention_core_xpu",
773
"vllm::olmo_hybrid_gdn_full_forward",
774
- "vllm::kda_attention",
774
"vllm::sparse_attn_indexer",
775
"vllm::rocm_aiter_sparse_attn_indexer",
776
"vllm::deepseek_v4_attention",
vllm/config/parallel.py
-5
@@ -1028,11 +1028,6 @@ class ParallelConfig:
1028
"Disabled the custom all-reduce kernel because it is not "
1029
"supported on current platform."
1030
)
1031
- if self.nnodes > 1:
1032
- self.disable_custom_all_reduce = True
1033
- logger.debug(
1034
- "Disabled the custom all-reduce since we are running on multi-node."
1035
- )
1031
if self.ray_workers_use_nsight and not self.use_ray:
1032
raise ValueError(
1033
"Unable to use nsight profiling unless workers run with Ray."
vllm/config/speculative.py
+22
@@ -51,6 +51,7 @@ MTPModelTypes = Literal[
51
"minimax_m3_mtp",
52
"bailing_hybrid_mtp",
53
"mtp",
54
+ "kimi_k3_mtp",
55
"pangu_ultra_moe_mtp",
56
"step3p5_mtp",
57
"hy_v3_mtp",
@@ -355,6 +356,16 @@ class SpeculativeConfig:
356
{"n_predict": n_predict, "architectures": ["OpenPanguMTPModel"]}
357
)
358
359
+ if hf_config.model_type == "kimi_k3":
360
+ # Kimi-K3 keeps the text-model fields (incl. the MTP layer count)
361
+ # nested under ``text_config`` (a KimiLinearConfig).
362
+ text_config = getattr(hf_config, "text_config", hf_config)
363
+ n_predict = getattr(text_config, "num_nextn_predict_layers", None)
364
+ hf_config.model_type = "kimi_k3_mtp"
365
+ hf_config.update(
366
+ {"n_predict": n_predict, "architectures": ["KimiK3MTPModel"]}
367
+ )
368
+
369
if hf_config.architectures[0] == "MiMoForCausalLM":
370
hf_config.model_type = "mimo_mtp"
371
n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
@@ -943,6 +954,7 @@ class SpeculativeConfig:
954
if self.method == "dspark" and (
955
"Qwen3DSparkModel" not in self.draft_model_config.architectures
956
and "Gemma4DSparkModel" not in self.draft_model_config.architectures
957
+ and "K3DSparkModel" not in self.draft_model_config.architectures
958
):
959
# DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config
960
# and its weights ship in the target checkpoint.
@@ -972,6 +984,16 @@ class SpeculativeConfig:
984
if self.method in ("dflash", "dspark"):
985
self.parallel_drafting = True
986
987
+ if (
988
+ self.method == "dspark"
989
+ and "K3DSparkModel" in self.draft_model_config.architectures
990
+ and self.target_parallel_config.decode_context_parallel_size > 1
991
+ ):
992
+ raise ValueError(
993
+ "MLA DSpark does not currently support decode context "
994
+ "parallelism; set decode_context_parallel_size=1."
995
+ )
996
+
997
if self.num_speculative_tokens is not None and hasattr(
998
self.draft_model_config.hf_config, "num_lookahead_tokens"
999
):
vllm/config/vllm.py
+36
-14
@@ -72,11 +72,34 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
72
"GraniteMoeForCausalLM",
73
"InklingForCausalLM",
74
"InklingForConditionalGeneration",
75
+ "KimiK3ForConditionalGeneration",
76
"LongcatFlashNgramForCausalLM",
77
"Qwen2MoeForCausalLM",
78
}
79
)
80
81
+# Architectures that default to V1 on ROCm: the V2 runner faults during the
82
+# profile run. VLLM_USE_V2_MODEL_RUNNER=1 still forces V2.
83
+# TODO: fix V2 enablement
84
+ROCM_EXCLUDED_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
85
+ {
86
+ "KimiK3ForConditionalGeneration",
87
+ }
88
+)
89
+
90
+
91
+@lru_cache
92
+def default_v2_model_runner_architectures() -> frozenset[str]:
93
+ """Architectures defaulting to the V2 model runner on this platform."""
94
+ from vllm.platforms import current_platform
95
+
96
+ if current_platform.is_rocm():
97
+ return (
98
+ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES
99
+ - ROCM_EXCLUDED_V2_MODEL_RUNNER_ARCHITECTURES
100
+ )
101
+ return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES
102
+
103
104
class OptimizationLevel(IntEnum):
105
"""Optimization level enum."""
@@ -619,16 +642,20 @@ class VllmConfig:
642
if model_config.runner_type != "generate":
643
return False
644
622
- if getattr(model_config, "is_hybrid", False):
645
+ architectures = getattr(model_config, "architectures", [])
646
+ default_architectures = default_v2_model_runner_architectures()
647
+ is_default_v2_architecture = any(
648
+ arch in default_architectures for arch in architectures
649
+ )
650
+
651
+ if getattr(model_config, "is_hybrid", False) and (
652
+ not is_default_v2_architecture
653
+ ):
654
return False
655
656
if getattr(model_config, "is_attention_free", False):
657
return False
627
- architectures = getattr(model_config, "architectures", [])
628
- return (
629
- any(arch in DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES for arch in architectures)
630
- or not model_config.is_moe
631
- )
658
+ return is_default_v2_architecture or not model_config.is_moe
659
660
@property
661
def needs_dp_coordinator(self) -> bool:
@@ -1196,6 +1223,9 @@ class VllmConfig:
1223
"DeepSeekV4MTPModel",
1224
"InklingForCausalLM",
1225
"InklingForConditionalGeneration",
1226
+ "KimiK3ForConditionalGeneration",
1227
+ "KimiK3MTPModel",
1228
+ "KimiLinearForCausalLM",
1229
"MiniMaxM3SparseForCausalLM",
1230
"MiniMaxM3SparseForConditionalGeneration",
1231
)
@@ -2277,14 +2307,6 @@ class VllmConfig:
2307
2308
# Mamba cache align-mode constraints
2309
if self.cache_config.mamba_cache_mode == "align":
2280
- assert block_size <= self.scheduler_config.max_num_batched_tokens, (
2281
- "In Mamba cache align mode, block_size "
2282
- f"({block_size}) must be <= "
2283
- "max_num_batched_tokens "
2284
- f"({self.scheduler_config.max_num_batched_tokens})."
2285
- )
2286
- if self.scheduler_config.long_prefill_token_threshold > 0:
2287
- assert self.scheduler_config.long_prefill_token_threshold >= block_size
2310
assert not self.scheduler_config.disable_chunked_mm_input, (
2311
"Chunked MM input is required because we need the flexibility "
2312
"to schedule a multiple of block_size tokens even if they are "
vllm/distributed/device_communicators/cuda_communicator.py
+12
@@ -340,6 +340,18 @@ class CudaCommunicator(DeviceCommunicatorBase):
340
torch.distributed.all_reduce(out, group=self.device_group)
341
return out
342
343
+ def custom_all_gather(self, input_: torch.Tensor) -> torch.Tensor | None:
344
+ ca_comm = self.ca_comm
345
+ if ca_comm is None:
346
+ return None
347
+ return ca_comm.custom_all_gather(input_.contiguous())
348
+
349
+ def custom_reduce_scatter(self, input_: torch.Tensor) -> torch.Tensor | None:
350
+ ca_comm = self.ca_comm
351
+ if ca_comm is None:
352
+ return None
353
+ return ca_comm.custom_reduce_scatter(input_.contiguous())
354
+
355
def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
356
# Route uniform dim-0 all-gathers through NVLS symmetric memory when
357
# enabled (mirrors reduce_scatter); otherwise fall back to the
vllm/distributed/device_communicators/custom_all_reduce.py
+264
-37
@@ -25,6 +25,11 @@ except Exception:
25
# For CPUs
26
custom_ar = False
27
28
+try:
29
+ import torch.distributed._symmetric_memory as torch_symm_mem
30
+except ImportError:
31
+ torch_symm_mem = None
32
+
33
logger = init_logger(__name__)
34
35
@@ -49,7 +54,17 @@ from vllm.distributed.utils import is_weak_contiguous # noqa: E402
54
55
56
class CustomAllreduce:
52
- _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8]
57
+ _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8, 16]
58
+ _DEFAULT_ALL_GATHER_MAX_SIZE = 2 * 1024 * 1024
59
+ _DEFAULT_MNNVL_ALL_GATHER_MAX_SIZES = {
60
+ 2: 8 * 1024 * 1024,
61
+ 4: 4 * 1024 * 1024,
62
+ 6: 2 * 1024 * 1024,
63
+ 8: 2 * 1024 * 1024,
64
+ 16: 2 * 1024 * 1024,
65
+ }
66
+ _DEFAULT_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024
67
+ _DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024
68
69
# max_size: max supported allreduce size
70
def __init__(
@@ -57,6 +72,10 @@ class CustomAllreduce:
72
group: ProcessGroup,
73
device: int | str | torch.device,
74
max_size=8192 * 1024,
75
+ max_all_gather_size=_DEFAULT_ALL_GATHER_MAX_SIZE,
76
+ max_mnnvl_all_gather_size=None,
77
+ max_reduce_scatter_size=_DEFAULT_REDUCE_SCATTER_MAX_SIZE,
78
+ max_mnnvl_reduce_scatter_size=_DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE,
79
symm_mem_enabled=False,
80
) -> None:
81
"""
@@ -70,7 +89,20 @@ class CustomAllreduce:
89
are in the same node.
90
"""
91
self._IS_CAPTURING = False
92
+ self._ptr = 0
93
self.disabled = True
94
+ self.mnnvl_buffer = None
95
+ self.mnnvl_handle = None
96
+ self.mnnvl_peer_buffers: list[torch.Tensor] | None = None
97
+ self.mnnvl_multicast_ptr = 0
98
+ self.mnnvl_buffer_size = 0
99
+ self.mnnvl_lamport_ag_local_ptr = 0
100
+ self.mnnvl_lamport_ag_multicast_ptr = 0
101
+ self.mnnvl_lamport_rs_local_ptr = 0
102
+ self.mnnvl_lamport_epochs = None
103
+ self.mnnvl_lamport_ag_epoch_ptr = 0
104
+ self.mnnvl_lamport_rs_epoch_ptr = 0
105
+ self.mnnvl_only = False
106
107
if not custom_ar:
108
# disable because of missing custom allreduce library
@@ -87,13 +119,8 @@ class CustomAllreduce:
119
"CustomAllreduce should be attached to a non-NCCL group."
120
)
121
90
- if not all(in_the_same_node_as(group, source_rank=0)):
91
- # No need to initialize custom allreduce for multi-node case.
92
- logger.warning(
93
- "Custom allreduce is disabled because this process group"
94
- " spans across nodes."
95
- )
96
- return
122
+ same_node = all(in_the_same_node_as(group, source_rank=0))
123
+ self.mnnvl_only = not same_node
124
125
rank = dist.get_rank(group=self.group)
126
self.rank = rank
@@ -126,28 +153,30 @@ class CustomAllreduce:
153
and device_capability is not None
154
):
155
device_capability_str = device_capability.as_version_str()
129
- if device_capability_str in CUSTOM_ALL_REDUCE_MAX_SIZES:
156
+ if (
157
+ device_capability_str in CUSTOM_ALL_REDUCE_MAX_SIZES
158
+ and world_size in CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str]
159
+ ):
160
max_size = min(
161
CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size],
162
max_size,
163
)
164
# device.index is a visible ordinal, not a logical local ID.
135
- physical_device_id = current_platform.visible_device_id_to_physical_device_id(
136
- device.index
137
- )
138
- tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
139
- gather_list = [
140
- torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)
141
- ]
142
- dist.all_gather(gather_list, tensor, group=self.group)
143
- physical_device_ids = [t.item() for t in gather_list]
144
-
145
- # test nvlink first, this will filter out most of the cases
146
- # where custom allreduce is not supported
147
- # this checks hardware and driver support for NVLink
148
- assert current_platform.is_cuda_alike()
149
- fully_connected = current_platform.is_fully_connected(physical_device_ids)
150
- if world_size > 2 and not fully_connected:
165
+ fully_connected = False
166
+ if same_node:
167
+ physical_device_id = (
168
+ current_platform.visible_device_id_to_physical_device_id(device.index)
169
+ )
170
+ tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
171
+ gather_list = [
172
+ torch.tensor([0], dtype=torch.int, device="cpu")
173
+ for _ in range(world_size)
174
+ ]
175
+ dist.all_gather(gather_list, tensor, group=self.group)
176
+ physical_device_ids = [t.item() for t in gather_list]
177
+ assert current_platform.is_cuda_alike()
178
+ fully_connected = current_platform.is_fully_connected(physical_device_ids)
179
+ if same_node and world_size > 2 and not fully_connected:
180
logger.warning(
181
"Custom allreduce is disabled because it's not supported on"
182
" more than two PCIe-only GPUs. To silence this warning, "
@@ -158,7 +187,11 @@ class CustomAllreduce:
187
# this is expensive to compute at the first time
188
# then we cache the result
189
# On AMD GPU, p2p is always enabled between XGMI connected GPUs
161
- if not current_platform.is_rocm() and not _can_p2p(rank, world_size):
190
+ if (
191
+ same_node
192
+ and not current_platform.is_rocm()
193
+ and not _can_p2p(rank, world_size)
194
+ ):
195
logger.warning(
196
"Custom allreduce is disabled because your platform lacks "
197
"GPU P2P capability or P2P test failed. To silence this "
@@ -170,21 +203,40 @@ class CustomAllreduce:
203
# Buffers memory are owned by this Python class and passed to C++.
204
# Metadata composes of two parts: metadata for synchronization and a
205
# temporary buffer for storing intermediate allreduce results.
173
- self.meta_ptrs = self.create_shared_buffer(
174
- ops.meta_size() + max_size, group=group, uncached=True
175
- )
206
+ if same_node:
207
+ self.meta_ptrs = self.create_shared_buffer(
208
+ ops.meta_size() + max_size, group=group, uncached=True
209
+ )
210
+ else:
211
+ meta_ptr, _ = ops.allocate_shared_buffer_and_handle(ops.meta_size())
212
+ self.meta_ptrs = [meta_ptr] * world_size
213
# This is a pre-registered IPC buffer. In eager mode, input tensors
177
- # are first copied into this buffer before allreduce is performed
178
- self.buffer_ptrs = self.create_shared_buffer(max_size, group=group)
179
- # This is a buffer for storing the tuples of pointers pointing to
180
- # IPC buffers from all ranks. Each registered tuple has size of
181
- # 8*world_size bytes where world_size is at most 8. Allocating 8MB
182
- # is enough for 131072 such tuples. The largest model I've seen only
183
- # needs less than 10000 of registered tuples.
214
+ # are first copied into this buffer before the operation is performed
215
+ legacy_buffer_size = max(max_size, max_all_gather_size, max_reduce_scatter_size)
216
+ if same_node:
217
+ self.buffer_ptrs = self.create_shared_buffer(
218
+ legacy_buffer_size,
219
+ group=group,
220
+ )
221
+ else:
222
+ buffer_ptr, _ = ops.allocate_shared_buffer_and_handle(legacy_buffer_size)
223
+ self.buffer_ptrs = [buffer_ptr] * world_size
224
+ # This stores tuples of pointers to IPC buffers from all ranks.
225
+ # Each registered tuple contains at most 16 addresses.
226
+ # Allocating 8MB is enough for 65536 such tuples. The largest model uses
227
+ # fewer than 10000 registered tuples.
228
self.rank_data = torch.empty(
229
8 * 1024 * 1024, dtype=torch.uint8, device=self.device
230
)
231
self.max_size = max_size
232
+ self.max_all_gather_size = max_all_gather_size
233
+ if max_mnnvl_all_gather_size is None:
234
+ max_mnnvl_all_gather_size = self._DEFAULT_MNNVL_ALL_GATHER_MAX_SIZES[
235
+ world_size
236
+ ]
237
+ self.max_mnnvl_all_gather_size = max_mnnvl_all_gather_size
238
+ self.max_reduce_scatter_size = max_reduce_scatter_size
239
+ self.max_mnnvl_reduce_scatter_size = max_mnnvl_reduce_scatter_size
240
self.rank = rank
241
self.world_size = world_size
242
self.fully_connected = fully_connected
@@ -192,6 +244,72 @@ class CustomAllreduce:
244
self.meta_ptrs, self.rank_data, rank, self.fully_connected
245
)
246
ops.register_buffer(self._ptr, self.buffer_ptrs)
247
+ self._init_mnnvl_buffer(
248
+ max(
249
+ max_mnnvl_all_gather_size * world_size,
250
+ max_mnnvl_reduce_scatter_size,
251
+ )
252
+ )
253
+ if not same_node and not self.mnnvl_multicast_ptr:
254
+ logger.warning(
255
+ "Custom collectives are disabled because this multi-node "
256
+ "group does not support MNNVL multicast."
257
+ )
258
+ self.close()
259
+ self.disabled = True
260
+
261
+ def _init_mnnvl_buffer(self, stage_size: int) -> None:
262
+ if torch_symm_mem is None or not current_platform.is_cuda():
263
+ return
264
+ try:
265
+ buffer_size = stage_size * 6
266
+ buffer = torch_symm_mem.empty(
267
+ buffer_size, dtype=torch.uint8, device=self.device
268
+ )
269
+ handle = torch_symm_mem.rendezvous(buffer, self.group.group_name)
270
+ if handle.multicast_ptr == 0:
271
+ return
272
+ peer_buffers = [
273
+ handle.get_buffer(
274
+ peer,
275
+ (buffer_size,),
276
+ torch.uint8,
277
+ storage_offset=0,
278
+ )
279
+ for peer in range(self.world_size)
280
+ ]
281
+ ptrs = [peer_buffer.data_ptr() for peer_buffer in peer_buffers]
282
+ lamport_ag_offset = 0
283
+ lamport_rs_offset = stage_size * 3
284
+ lamport_ag_ptrs = [ptr + lamport_ag_offset for ptr in ptrs]
285
+ lamport_rs_ptrs = [ptr + lamport_rs_offset for ptr in ptrs]
286
+ ops.register_buffer(self._ptr, lamport_ag_ptrs)
287
+ ops.register_buffer(self._ptr, lamport_rs_ptrs)
288
+
289
+ buffer.view(torch.int32).fill_(-2147483648)
290
+ epochs = torch.zeros(
291
+ (2, 32),
292
+ dtype=torch.int32,
293
+ device=self.device,
294
+ )
295
+ torch.accelerator.synchronize()
296
+ dist.barrier(group=self.group)
297
+
298
+ self.mnnvl_buffer = buffer
299
+ self.mnnvl_handle = handle
300
+ self.mnnvl_peer_buffers = peer_buffers
301
+ self.mnnvl_multicast_ptr = handle.multicast_ptr
302
+ self.mnnvl_buffer_size = stage_size
303
+ self.mnnvl_lamport_ag_local_ptr = lamport_ag_ptrs[self.rank]
304
+ self.mnnvl_lamport_ag_multicast_ptr = (
305
+ handle.multicast_ptr + lamport_ag_offset
306
+ )
307
+ self.mnnvl_lamport_rs_local_ptr = lamport_rs_ptrs[self.rank]
308
+ self.mnnvl_lamport_epochs = epochs
309
+ self.mnnvl_lamport_ag_epoch_ptr = epochs[0].data_ptr()
310
+ self.mnnvl_lamport_rs_epoch_ptr = epochs[1].data_ptr()
311
+ except RuntimeError as error:
312
+ logger.debug("MNNVL AG/RS initialization failed: %s", error)
313
314
@contextmanager
315
def capture(self):
@@ -228,7 +346,7 @@ class CustomAllreduce:
346
ops.register_graph_buffers(self._ptr, handles, offsets)
347
348
def should_custom_ar(self, inp: torch.Tensor):
231
- if self.disabled:
349
+ if self.disabled or self.world_size > 8:
350
return False
351
inp_size = inp.numel() * inp.element_size()
352
# custom allreduce requires input byte size to be multiples of 16
@@ -279,6 +397,111 @@ class CustomAllreduce:
397
# latency) compared to the performance gain of using custom kernels
398
return self.all_reduce(input, registered=False)
399
400
+ def should_custom_all_gather(self, inp: torch.Tensor) -> bool:
401
+ if self.disabled or not current_platform.is_cuda():
402
+ return False
403
+ if self.world_size == 16 and not self.mnnvl_only:
404
+ return False
405
+ inp_size = inp.nbytes
406
+ if inp.dtype not in (
407
+ torch.float32,
408
+ torch.float16,
409
+ torch.bfloat16,
410
+ ):
411
+ return False
412
+ max_size = (
413
+ self.max_mnnvl_all_gather_size
414
+ if self.mnnvl_multicast_ptr
415
+ else self.max_all_gather_size
416
+ )
417
+ return (
418
+ 0 < inp_size <= max_size
419
+ and inp_size % 16 == 0
420
+ and is_weak_contiguous(inp)
421
+ and (self.fully_connected or bool(self.mnnvl_multicast_ptr))
422
+ )
423
+
424
+ def custom_all_gather(self, inp: torch.Tensor) -> torch.Tensor | None:
425
+ if not self.should_custom_all_gather(inp):
426
+ return None
427
+ out_shape = (inp.shape[0] * self.world_size,) + inp.shape[1:]
428
+ if self.mnnvl_multicast_ptr:
429
+ logger.info_once(
430
+ "Using the MNNVL Lamport all-gather kernel.",
431
+ scope="global",
432
+ )
433
+ out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device)
434
+ ops.mnnvl_lamport_all_gather(
435
+ self._ptr,
436
+ inp,
437
+ out,
438
+ self.mnnvl_lamport_ag_local_ptr,
439
+ self.mnnvl_lamport_ag_multicast_ptr,
440
+ self.mnnvl_lamport_ag_epoch_ptr,
441
+ self.mnnvl_buffer_size,
442
+ )
443
+ else:
444
+ out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device)
445
+ ops.custom_all_gather(
446
+ self._ptr,
447
+ inp,
448
+ out,
449
+ self.buffer_ptrs[self.rank],
450
+ self.max_all_gather_size,
451
+ )
452
+ return out
453
+
454
+ def should_custom_reduce_scatter(self, inp: torch.Tensor) -> bool:
455
+ if self.disabled or not current_platform.is_cuda():
456
+ return False
457
+ if self.world_size == 16 and not self.mnnvl_only:
458
+ return False
459
+ inp_size = inp.nbytes
460
+ if inp.dtype not in (torch.float32, torch.float16, torch.bfloat16):
461
+ return False
462
+ if inp.shape[0] % self.world_size != 0:
463
+ return False
464
+ output_size = inp_size // self.world_size
465
+ max_size = (
466
+ self.max_mnnvl_reduce_scatter_size
467
+ if self.mnnvl_multicast_ptr
468
+ else self.max_reduce_scatter_size
469
+ )
470
+ return (
471
+ 0 < inp_size <= max_size
472
+ and output_size % 16 == 0
473
+ and is_weak_contiguous(inp)
474
+ and (self.fully_connected or bool(self.mnnvl_multicast_ptr))
475
+ )
476
+
477
+ def custom_reduce_scatter(self, inp: torch.Tensor) -> torch.Tensor | None:
478
+ if not self.should_custom_reduce_scatter(inp):
479
+ return None
480
+ out_shape = (inp.shape[0] // self.world_size,) + inp.shape[1:]
481
+ out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device)
482
+ if self.mnnvl_multicast_ptr:
483
+ logger.info_once(
484
+ "Using the MNNVL Lamport reduce-scatter kernel.",
485
+ scope="global",
486
+ )
487
+ ops.mnnvl_lamport_reduce_scatter(
488
+ self._ptr,
489
+ inp,
490
+ out,
491
+ self.mnnvl_lamport_rs_local_ptr,
492
+ self.mnnvl_lamport_rs_epoch_ptr,
493
+ self.mnnvl_buffer_size,
494
+ )
495
+ else:
496
+ ops.custom_reduce_scatter(
497
+ self._ptr,
498
+ inp,
499
+ out,
500
+ self.buffer_ptrs[self.rank],
501
+ self.max_reduce_scatter_size,
502
+ )
503
+ return out
504
+
505
def close(self):
506
if not self.disabled and self._ptr:
507
if ops is not None:
@@ -286,6 +509,10 @@ class CustomAllreduce:
509
self._ptr = 0
510
self.free_shared_buffer(self.meta_ptrs, rank=self.rank)
511
self.free_shared_buffer(self.buffer_ptrs, rank=self.rank)
512
+ self.mnnvl_peer_buffers = None
513
+ self.mnnvl_handle = None
514
+ self.mnnvl_buffer = None
515
+ self.mnnvl_lamport_epochs = None
516
517
def __del__(self):
518
self.close()
vllm/engine/arg_utils.py
+10
@@ -741,6 +741,7 @@ class EngineArgs:
741
742
fail_on_environ_validation: bool = False
743
gdn_prefill_backend: Literal["flashinfer", "triton", "cutedsl"] | None = None
744
+ kda_prefill_backend: Literal["auto", "triton", "flashkda"] | None = None
745
746
def __post_init__(self):
747
# support `EngineArgs(compilation_config={...})`
@@ -1637,6 +1638,13 @@ class EngineArgs:
1638
default=None,
1639
help="Select GDN prefill backend.",
1640
)
1641
+ parser.add_argument(
1642
+ "--kda-prefill-backend",
1643
+ dest="kda_prefill_backend",
1644
+ choices=["auto", "triton", "flashkda"],
1645
+ default=None,
1646
+ help="Select KDA prefill backend.",
1647
+ )
1648
return parser
1649
1650
@classmethod
@@ -2431,6 +2439,8 @@ class EngineArgs:
2439
2440
if self.gdn_prefill_backend is not None:
2441
self.additional_config["gdn_prefill_backend"] = self.gdn_prefill_backend
2442
+ if self.kda_prefill_backend is not None:
2443
+ self.additional_config["kda_prefill_backend"] = self.kda_prefill_backend
2444
2445
config = VllmConfig(
2446
model_config=model_config,
vllm/envs.py
+7
-7
@@ -131,6 +131,7 @@ if TYPE_CHECKING:
131
VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False
132
VLLM_ROCM_USE_AITER_MOE: bool = True
133
VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0
134
+ AITER_SITUV2_A8W4: bool = False
135
VLLM_ROCM_USE_AITER_RMSNORM: bool = True
136
VLLM_ROCM_USE_AITER_MLA: bool = True
137
VLLM_ROCM_USE_AITER_MHA: bool = True
@@ -269,7 +270,6 @@ if TYPE_CHECKING:
270
VLLM_NCCL_INCLUDE_PATH: str | None = None
271
VLLM_GC_DEBUG: str = ""
272
VLLM_DEBUG_WORKSPACE: bool = False
272
- VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION: bool = False
273
VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
274
VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
275
VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024
@@ -1004,7 +1004,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
1004
# Backend for Video IO — selects the frame-sampling algorithm.
1005
# - "opencv": uniform sampling.
1006
# - "opencv_dynamic": duration-aware dynamic sampling.
1007
- # - "identity": returns raw video bytes for model processor to handle.
1007
#
1008
# Custom backend implementations can be registered
1009
# via `@VIDEO_LOADER_REGISTRY.register("my_custom_video_loader")` and
@@ -1211,6 +1210,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
1210
"VLLM_ROCM_USE_AITER_MOE": lambda: (
1211
os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1")
1212
),
1213
+ # Route K3 SiTU MXFP4 MoE through the a8w4 (fp8 activation) gate/up-
1214
+ # interleaved flydsl kernels instead of the default a16w4 separated path.
1215
+ # Shared with the AITER runtime, which reads the same env var directly.
1216
+ "AITER_SITUV2_A8W4": lambda: (
1217
+ os.getenv("AITER_SITUV2_A8W4", "0").lower() in ("true", "1")
1218
+ ),
1219
# MoE sorting dispatch policy for AITER fused MoE kernels.
1220
# 0 = auto (default): single-pass for small batches, multi-pass
1221
# for large batches
@@ -1895,11 +1900,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
1900
# Debug workspace allocations.
1901
# logging of workspace resize operations.
1902
"VLLM_DEBUG_WORKSPACE": lambda: bool(int(os.getenv("VLLM_DEBUG_WORKSPACE", "0"))),
1898
- # Enable the experimental Kimi K3 latent-MoE tail fusion.
1899
- # Currently supported only on SM100 with TP=8/16 and BF16.
1900
- "VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION": lambda: bool(
1901
- int(os.getenv("VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION", "0"))
1902
- ),
1903
# Disables parallel execution of shared_experts via separate cuda stream
1904
"VLLM_DISABLE_SHARED_EXPERTS_STREAM": lambda: bool(
1905
int(os.getenv("VLLM_DISABLE_SHARED_EXPERTS_STREAM", "0"))
vllm/model_executor/layers/attention/mla_attention.py
+52
-8
@@ -272,6 +272,7 @@ from vllm.v1.attention.backends.mla.prefill import (
272
)
273
from vllm.v1.attention.backends.utils import (
274
get_dcp_local_seq_lens,
275
+ get_num_attention_heads_from_layers,
276
split_decodes_and_prefills,
277
)
278
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):
376
use_sparse: bool = False,
377
indexer: object | None = None,
378
topk_indices_buffer: torch.Tensor | None = None,
379
+ non_causal_multi_token_decode: bool = False,
380
**extra_impl_args,
381
):
382
super().__init__()
@@ -391,6 +393,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
393
self.head_size = kv_lora_rank + qk_rope_head_dim
394
self.layer_name = prefix
395
self.indexer = indexer
396
+ self.non_causal_multi_token_decode = non_causal_multi_token_decode
397
self.num_kv_heads = 1
398
self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
399
@@ -1123,6 +1126,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1126
dtype=kv_cache_dtype,
1127
cache_dtype_str=self.kv_cache_dtype,
1128
kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
1129
+ non_causal_multi_token_decode=self.non_causal_multi_token_decode,
1130
)
1131
1132
def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
@@ -1450,6 +1454,8 @@ class MLACommonMetadata(AttentionMetadata, Generic[D]):
1454
num_decode_tokens: int
1455
num_prefills: int
1456
1457
+ causal: bool = True
1458
+
1459
# The dimension of the attention heads
1460
head_dim: int | None = None
1461
@@ -1724,6 +1730,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1730
# speculative decoding is enabled.
1731
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.SINGLE_ONLY
1732
1733
+ # Whether this builder can flatten a non-causal query block into decode rows.
1734
+ supports_non_causal_multi_token_decode: ClassVar[bool] = False
1735
+
1736
# The threshold for reordering the batch into decode and prefill requests.
1737
# If > 1, the batch will be reordered such that requests with
1738
# query length <= threshold are classified as decode requests.
@@ -1824,8 +1833,14 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1833
self.vllm_config = vllm_config
1834
self.device = device
1835
self.use_pcp = parallel_config.prefill_context_parallel_size > 1
1836
+ self.non_causal_multi_token_decode = getattr(
1837
+ kv_cache_spec, "non_causal_multi_token_decode", False
1838
+ )
1839
1828
- self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
1840
+ # A draft cache group can have a different head count from the target.
1841
+ self.num_heads = get_num_attention_heads_from_layers(
1842
+ vllm_config, layer_names
1843
+ ) or self.model_config.get_num_attention_heads(parallel_config)
1844
self.mla_dims = get_mla_dims(self.model_config)
1845
self.aot_schedule = current_platform.is_cuda()
1846
@@ -1951,14 +1966,42 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1966
seq_lens = common_attn_metadata.seq_lens
1967
dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens
1968
1954
- num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
1955
- split_decodes_and_prefills(
1956
- common_attn_metadata,
1957
- decode_threshold=self.reorder_batch_threshold,
1958
- require_uniform=(self.query_len_support != QueryLenSupport.VARLEN),
1959
- treat_short_extends_as_decodes=not self.use_pcp,
1969
+ non_causal_decode = common_attn_metadata.causal is False
1970
+ if non_causal_decode:
1971
+ if not (
1972
+ self.supports_non_causal_multi_token_decode
1973
+ and self.non_causal_multi_token_decode
1974
+ ):
1975
+ raise ValueError(
1976
+ "Non-causal multi-token MLA requires an explicitly supported "
1977
+ "attention group."
1978
+ )
1979
+ query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
1980
+ num_active_reqs = int(torch.count_nonzero(query_lens > 0))
1981
+ uniform_active_queries = num_active_reqs > 0 and bool(
1982
+ torch.all(query_lens[:num_active_reqs] == query_lens[0])
1983
+ )
1984
+ trailing_graph_padding = bool(torch.all(query_lens[num_active_reqs:] == 0))
1985
+ if not (uniform_active_queries and trailing_graph_padding):
1986
+ raise ValueError(
1987
+ "Non-causal MLA requires a uniform query block; got query "
1988
+ f"lengths {query_lens.tolist()}."
1989
+ )
1990
+ # Use exact GPU sequence lengths instead of the prefill path's CPU
1991
+ # context-length upper bounds.
1992
+ num_decodes = num_reqs
1993
+ num_prefills = 0
1994
+ num_decode_tokens = num_tokens
1995
+ num_prefill_tokens = 0
1996
+ else:
1997
+ num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
1998
+ split_decodes_and_prefills(
1999
+ common_attn_metadata,
2000
+ decode_threshold=self.reorder_batch_threshold,
2001
+ require_uniform=(self.query_len_support != QueryLenSupport.VARLEN),
2002
+ treat_short_extends_as_decodes=not self.use_pcp,
2003
+ )
2004
)
1961
- )
2005
2006
assert num_decodes + num_prefills == num_reqs
2007
assert num_decode_tokens + num_prefill_tokens == num_tokens
@@ -2049,6 +2092,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
2092
num_decodes=num_decodes,
2093
num_decode_tokens=num_decode_tokens,
2094
num_prefills=num_prefills,
2095
+ causal=not non_causal_decode,
2096
prefill=prefill_metadata,
2097
decode=decode_metadata,
2098
)
vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
+1
-1
@@ -505,7 +505,7 @@ def deepgemm_moe_permute(
505
dtype=torch.int32,
506
)
507
else:
508
- aq_scale_out = torch.empty((M_sum, sf_k), device=device, dtype=torch.float32)
508
+ aq_scale_out = torch.zeros((M_sum, sf_k), device=device, dtype=torch.float32)
509
510
# DeepGEMM uses negative values in m_indices (here expert_ids) to mark
511
# completely invalid / padded blocks that should be skipped. We always
vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py
+1
@@ -204,6 +204,7 @@ def persistent_masked_m_silu_mul_quant(
204
dtype=ys_dtype,
205
device=y.device,
206
)
207
+ y_s.zero_()
208
209
ceil_ue8m0 = quant_scale_fmt in [
210
DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0,
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+5
-8
@@ -671,10 +671,10 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
671
672
@staticmethod
673
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
674
- return not (
675
- moe_parallel_config.use_fi_nvl_two_sided_kernels
676
- or moe_parallel_config.use_fi_nvl_one_sided_kernels
677
- )
674
+ # One-sided FI-NVL all2all pairs with MarlinExperts fine (the
675
+ # compressed-tensors MXFP4 path runs this exact combo); only the
676
+ # two-sided kernels are unsupported here.
677
+ return not moe_parallel_config.use_fi_nvl_two_sided_kernels
678
679
@property
680
def quant_type_id(self) -> int:
@@ -1068,11 +1068,8 @@ class BatchedMarlinExperts(MarlinExpertsBase):
1068
return
1069
1070
num_experts, max_num_tokens = hidden_states.shape[:2]
1071
- beta = activation_situ_beta
1071
+ beta = 1.0 if activation_situ_beta is None else activation_situ_beta
1072
linear_beta = activation_situ_linear_beta
1073
- assert beta is not None, (
1074
- "SITU requires activation_situ_beta from FusedMoEConfig"
1075
- )
1073
torch.ops._C.masked_situ_and_mul(
1074
act_output.view(num_experts, max_num_tokens, -1),
1075
act_input.view(num_experts, max_num_tokens, -1),
vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
+20
-5
@@ -5,6 +5,7 @@ from functools import lru_cache
5
6
import torch
7
8
+import vllm.envs as envs
9
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10
from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
@@ -262,6 +263,8 @@ def rocm_aiter_fused_experts(
263
elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
264
activation_method = rocm_aiter_ops.get_aiter_activation_type("swiglu")
265
activation_interleave = False
266
+ elif activation == MoEActivation.SITU:
267
+ activation_method = rocm_aiter_ops.get_aiter_activation_type("situ")
268
else:
269
raise ValueError(f"Unsupported activation: {activation}")
270
@@ -356,10 +359,12 @@ def rocm_aiter_fused_experts(
359
# https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1099
360
# TODO: Revisit this once we bump AITER to 0.1.15 with padding fixes
361
# for CK/FlyDSL MoE GEMM e.g. https://github.com/ROCm/aiter/pull/3401
359
- hidden_pad = hidden_pad // 128 * 128
360
- intermediate_pad = (
361
- intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1)
362
- )
362
+ # SITU's A16W4 FlyDSL kernel pads per gate/up half; pass through unrounded.
363
+ if activation != MoEActivation.SITU:
364
+ hidden_pad = hidden_pad // 128 * 128
365
+ intermediate_pad = (
366
+ intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1)
367
+ )
368
369
# https://github.com/ROCm/aiter/pull/3123 specialized the AITER stage1 GEMMs
370
# for interleaved vs separated gate and up weights.
@@ -370,7 +375,15 @@ def rocm_aiter_fused_experts(
375
from aiter.ops.flydsl.moe_common import GateMode
376
377
gate_mode = ""
373
- if quant_config.use_mxfp4_w4a16:
378
+ if activation == MoEActivation.SITU:
379
+ # a8w4 (AITER_SITUV2_A8W4=1) uses the gate/up-interleaved (_gui_)
380
+ # fp8 flydsl kernels; default a16w4 SiTU stays separated.
381
+ gate_mode = (
382
+ GateMode.INTERLEAVE.value
383
+ if envs.AITER_SITUV2_A8W4
384
+ else GateMode.SEPARATED.value
385
+ )
386
+ elif quant_config.use_mxfp4_w4a16:
387
gate_mode = GateMode.INTERLEAVE.value
388
elif activation_interleave is not None:
389
gate_mode = (
@@ -401,6 +414,8 @@ def rocm_aiter_fused_experts(
414
bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,
415
bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,
416
moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
417
+ beta=moe_config.activation_situ_beta,
418
+ linear_beta=moe_config.activation_situ_linear_beta,
419
)
420
421
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
+95
-48
@@ -15,6 +15,10 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
15
TopKWeightAndReduceNoOP,
16
)
17
from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights
18
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
19
+ activation_to_flashinfer_int,
20
+ has_flashinfer_situ_activation,
21
+)
22
from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
QuantKey,
24
kMxfp4Static,
@@ -133,19 +137,16 @@ class TrtLlmMxfp4ExpertsBase:
137
138
@staticmethod
139
def _supports_activation(activation: MoEActivation) -> bool:
140
+ if activation == MoEActivation.SITU:
141
+ return has_flashinfer_situ_activation()
142
return activation in (
143
MoEActivation.SWIGLUOAI,
144
MoEActivation.SILU,
139
- MoEActivation.SITU,
145
)
146
147
@staticmethod
148
def _flashinfer_activation_type(activation: MoEActivation) -> int:
144
- from flashinfer.fused_moe.core import ActivationType
145
-
146
- if activation == MoEActivation.SITU:
147
- return ActivationType.Situ.value
148
- return ActivationType.Swiglu.value
149
+ return activation_to_flashinfer_int(activation)
150
151
@staticmethod
152
def activation_format() -> mk.FusedMoEActivationFormat:
@@ -315,6 +316,75 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
316
output = (M, self.hidden_dim_unpadded)
317
return (workspace1, workspace2, output)
318
319
+ def _max_supported_tokens(self, top_k: int, global_num_experts: int) -> int:
320
+ """Max tokens per kernel call before the batched-GEMM grid overflows.
321
+
322
+ The TRTLLM-Gen batched GEMM launches a static grid whose batch (Y)
323
+ dimension is ``getMaxNumCtasInBatchDim(num_tokens, top_k, num_experts,
324
+ tileTokensDim)`` and must stay <= 65535. Solving that for num_tokens
325
+ with the smallest tile the kernel may pick (tileTokensDim=8, the runner
326
+ default) gives a bound that is safe regardless of the tactic selected.
327
+ Without it, large batches (e.g. Kimi-K3 top_k=16, EP16 profiling with
328
+ 131072 gathered tokens) overflow the grid and the GEMM launch fails.
329
+ """
330
+ MAX_GRID_Y = 65535
331
+ MIN_TILE_TOKENS_DIM = 8
332
+ max_tokens = (MAX_GRID_Y - global_num_experts) * MIN_TILE_TOKENS_DIM // top_k
333
+ return max(1, min(300000, max_tokens))
334
+
335
+ def _invoke_kernel(
336
+ self,
337
+ output: torch.Tensor,
338
+ x_quant: torch.Tensor,
339
+ x_scale: torch.Tensor | None,
340
+ topk_ids: torch.Tensor,
341
+ topk_weights: torch.Tensor,
342
+ w1: torch.Tensor,
343
+ w2: torch.Tensor,
344
+ activation: MoEActivation,
345
+ global_num_experts: int,
346
+ local_num_experts: int,
347
+ local_expert_offset: int,
348
+ topk: int,
349
+ ) -> None:
350
+ from flashinfer import trtllm_fp4_block_scale_routed_moe
351
+
352
+ packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
353
+ trtllm_fp4_block_scale_routed_moe(
354
+ topk_ids=packed_tensor,
355
+ routing_bias=None,
356
+ hidden_states=x_quant,
357
+ hidden_states_scale=x_scale,
358
+ gemm1_weights=w1,
359
+ gemm1_weights_scale=self.w1_scale,
360
+ gemm1_bias=self.w1_bias,
361
+ gemm1_alpha=self.gemm1_alpha,
362
+ gemm1_beta=self.gemm1_beta,
363
+ gemm1_clamp_limit=self.gemm1_clamp_limit,
364
+ gemm2_weights=w2,
365
+ gemm2_weights_scale=self.w2_scale,
366
+ gemm2_bias=self.w2_bias,
367
+ output1_scale_scalar=None,
368
+ output1_scale_gate_scalar=None,
369
+ output2_scale_scalar=None,
370
+ num_experts=global_num_experts,
371
+ top_k=topk,
372
+ n_group=None,
373
+ topk_group=None,
374
+ intermediate_size=self.intermediate_size_per_partition,
375
+ local_expert_offset=local_expert_offset,
376
+ local_num_experts=local_num_experts,
377
+ routed_scaling_factor=None,
378
+ # Modular kernel receives pre-routed tokens, so routing is already
379
+ # done. Use Renormalize as a safe default the TRTLLM kernel supports.
380
+ routing_method_type=RoutingMethodType.Renormalize,
381
+ do_finalize=True,
382
+ enable_pdl=True,
383
+ activation_type=self._flashinfer_activation_type(activation),
384
+ output=output,
385
+ tune_max_num_tokens=max(self.max_capture_size, 1),
386
+ )
387
+
388
def apply(
389
self,
390
output: torch.Tensor,
@@ -335,7 +405,6 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
405
):
406
topk = topk_ids.size(-1)
407
local_num_experts = w1.size(0)
338
- intermediate_size = self.intermediate_size_per_partition
408
local_expert_offset = self.moe_config.ep_rank * local_num_experts
409
410
if a1q_scale is not None:
@@ -346,49 +415,27 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
415
x_quant = hidden_states
416
x_scale = None
417
349
- # Pack topk ids and weights into format expected by the kernel.
350
- packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
351
-
418
assert self.w1_scale is not None
419
assert self.w2_scale is not None
354
- kwargs = {
355
- "topk_ids": packed_tensor,
356
- "routing_bias": None,
357
- "hidden_states": x_quant,
358
- "hidden_states_scale": x_scale,
359
- "gemm1_weights": w1,
360
- "gemm1_weights_scale": self.w1_scale,
361
- "gemm1_bias": self.w1_bias,
362
- "gemm1_alpha": self.gemm1_alpha,
363
- "gemm1_beta": self.gemm1_beta,
364
- "gemm1_clamp_limit": self.gemm1_clamp_limit,
365
- "gemm2_weights": w2,
366
- "gemm2_weights_scale": self.w2_scale,
367
- "gemm2_bias": self.w2_bias,
368
- "output1_scale_scalar": None,
369
- "output1_scale_gate_scalar": None,
370
- "output2_scale_scalar": None,
371
- "num_experts": global_num_experts,
372
- "top_k": topk,
373
- "n_group": None,
374
- "topk_group": None,
375
- "intermediate_size": intermediate_size,
376
- "local_expert_offset": local_expert_offset,
377
- "local_num_experts": local_num_experts,
378
- "routed_scaling_factor": None,
379
- # Modular kernel receives pre-routed tokens, so routing
380
- # is already done. Use Renormalize as a safe default that
381
- # the TRTLLM C++ kernel supports.
382
- "routing_method_type": RoutingMethodType.Renormalize,
383
- "do_finalize": True,
384
- "enable_pdl": True,
385
- "activation_type": self._flashinfer_activation_type(activation),
386
- "output": output,
387
- "tune_max_num_tokens": max(self.max_capture_size, 1),
388
- }
420
390
- from flashinfer import trtllm_fp4_block_scale_routed_moe
391
-
392
- trtllm_fp4_block_scale_routed_moe(**kwargs)
421
+ # Chunk tokens so the batched-GEMM grid stays within CUDA limits.
422
+ M = x_quant.size(0)
423
+ chunk_size = self._max_supported_tokens(topk, global_num_experts)
424
+ for start in range(0, M, chunk_size):
425
+ end = min(start + chunk_size, M)
426
+ self._invoke_kernel(
427
+ output[start:end],
428
+ x_quant[start:end],
429
+ None if x_scale is None else x_scale[start:end],
430
+ topk_ids[start:end],
431
+ topk_weights[start:end],
432
+ w1,
433
+ w2,
434
+ activation,
435
+ global_num_experts,
436
+ local_num_experts,
437
+ local_expert_offset,
438
+ topk,
439
+ )
440
441
return output
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+4
-3
@@ -22,6 +22,7 @@ from vllm.model_executor.layers.fused_moe.utils import (
22
)
23
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
24
activation_to_flashinfer_int,
25
+ has_flashinfer_situ_activation,
26
)
27
from vllm.model_executor.layers.quantization.utils.quant_utils import (
28
QuantKey,
@@ -225,15 +226,15 @@ class TrtLlmNvFp4ExpertsBase:
226
227
@staticmethod
228
def _supports_activation(activation: MoEActivation) -> bool:
228
- """Supports SiLU, RELU^2 non-gated, GELU, clamped SwiGLU-OAI, and the
229
- private SituGLU (SITU) cubin."""
229
+ """Supports SITU only when the installed FlashInfer exposes it."""
230
+ if activation == MoEActivation.SITU:
231
+ return has_flashinfer_situ_activation()
232
return activation in [
233
MoEActivation.SILU,
234
MoEActivation.RELU2_NO_MUL,
235
MoEActivation.GELU,
236
MoEActivation.GELU_TANH,
237
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
236
- MoEActivation.SITU,
238
]
239
240
@staticmethod
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+1
-1
@@ -648,7 +648,7 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(
648
else:
649
hidden_size = round_up(hidden_size, 256)
650
elif backend in TRTLLM_BACKENDS:
651
- intermediate_size = round_up(intermediate_size, 256)
651
+ intermediate_size = round_up(intermediate_size, 128)
652
hidden_size = round_up(hidden_size, 256)
653
elif backend in (
654
Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
vllm/model_executor/layers/fused_moe/prepare_finalize/batched.py
+1
-1
@@ -102,7 +102,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
102
num_local_experts, self.max_num_tokens, hidden_dim
103
)
104
105
- b_a1_scale = torch.empty(scale_shape, dtype=torch.float32, device=a1.device)
105
+ b_a1_scale = torch.zeros(scale_shape, dtype=torch.float32, device=a1.device)
106
else:
107
assert quant_config.a1_scale is None
108
b_a1_scale = None
vllm/model_executor/layers/fused_moe/runner/latent_moe_runner.py
+22
-16
@@ -6,11 +6,6 @@ import vllm.envs as envs
6
from vllm.config import get_current_vllm_config
7
from vllm.distributed import tensor_model_parallel_all_reduce
8
from vllm.logger import init_logger
9
-from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
10
- _AR_RESIDUAL_RMS_NORM,
11
- _can_use_flashinfer,
12
- flashinfer_trtllm_fused_allreduce_norm,
13
-)
9
from vllm.model_executor.layers.layernorm import RMSNorm
10
from vllm.utils.torch_utils import aux_stream, current_stream
11
@@ -82,25 +77,30 @@ class LatentMoERunner(MoERunner):
77
)
78
self._k3_latent_moe_tail_op = op
79
85
- def _get_zero_residual(self, hidden_states: torch.Tensor) -> torch.Tensor:
80
+ def _get_zero_residual(
81
+ self,
82
+ hidden_states: torch.Tensor,
83
+ max_token_num: int,
84
+ ) -> torch.Tensor:
85
"""Read-only zero ``residual_in`` for the fused AR+RMSNorm kernel.
86
87
flashinfer requires a residual buffer even when there is no residual to
88
add.
89
"""
91
- numel = hidden_states.numel()
90
buf = getattr(self, "_zero_residual", None)
93
- if (
94
- buf is None
95
- or buf.dtype != hidden_states.dtype
96
- or buf.device != hidden_states.device
97
- or buf.numel() < numel
98
- ):
91
+ if buf is None:
92
buf = torch.zeros(
100
- numel, dtype=hidden_states.dtype, device=hidden_states.device
93
+ max_token_num * hidden_states.shape[-1],
94
+ dtype=hidden_states.dtype,
95
+ device=hidden_states.device,
96
)
97
self._zero_residual = buf
103
- return buf[:numel].view_as(hidden_states)
98
+
99
+ assert buf.dtype == hidden_states.dtype
100
+ assert buf.device == hidden_states.device
101
+ assert hidden_states.numel() <= buf.numel()
102
+
103
+ return buf[: hidden_states.numel()].view_as(hidden_states)
104
105
def _use_fused_path(self) -> bool:
106
# The fused path merges the latent and shared reductions into one
@@ -224,6 +224,12 @@ class LatentMoERunner(MoERunner):
224
norm: RMSNorm,
225
) -> tuple[torch.Tensor, torch.Tensor]:
226
"""All-reduce + add residual + (standard) RMSNorm, fused via flashinfer."""
227
+ from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
228
+ _AR_RESIDUAL_RMS_NORM,
229
+ _can_use_flashinfer,
230
+ flashinfer_trtllm_fused_allreduce_norm,
231
+ )
232
+
233
if self.moe_config.tp_size == 1:
234
return norm(hidden_states)
235
@@ -238,7 +244,7 @@ class LatentMoERunner(MoERunner):
244
# buffer and the normalized result into norm_out.
245
flashinfer_trtllm_fused_allreduce_norm(
246
allreduce_in=hidden_states,
241
- residual=self._get_zero_residual(hidden_states),
247
+ residual=self._get_zero_residual(hidden_states, max_token_num),
248
rms_gamma=norm.weight,
249
rms_eps=norm.variance_epsilon,
250
world_size=self.moe_config.tp_size,
vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py
+482
-291
@@ -1,32 +1,33 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from collections.abc import Callable
5
+
6
import torch
7
from einops import rearrange
8
from torch import nn
9
+from torch.nn.parameter import Parameter
10
8
-from vllm.config import VllmConfig, get_current_vllm_config
9
-from vllm.distributed import (
10
- divide,
11
-)
12
-from vllm.forward_context import ForwardContext, get_forward_context
13
-from vllm.logger import init_logger
11
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
12
+from vllm.config import VllmConfig
13
+from vllm.distributed import divide, get_tensor_model_parallel_rank
14
+from vllm.forward_context import get_forward_context
15
from vllm.model_executor.custom_op import PluggableLayer
16
from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
16
-from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
17
-from vllm.model_executor.utils import set_weight_attrs
18
-from vllm.third_party.flash_linear_attention.ops.kda import (
19
- FusedRMSNormGated,
20
- chunk_kda_with_fused_gate,
21
- fused_kda_gate,
22
- fused_recurrent_kda,
17
+from vllm.model_executor.model_loader.weight_utils import (
18
+ default_weight_loader,
19
+ sharded_weight_loader,
20
)
21
+from vllm.model_executor.parameter import BasevLLMParameter
22
+from vllm.model_executor.utils import set_weight_attrs
23
+from vllm.platforms import current_platform
24
+from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
25
from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
25
-from vllm.utils.torch_utils import direct_register_custom_op
26
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
27
28
from ...linear import (
29
ColumnParallelLinear,
30
+ MergedColumnParallelLinear,
31
ReplicatedLinear,
32
RowParallelLinear,
33
)
@@ -36,49 +37,118 @@ from ..mamba_utils import (
37
is_conv_state_dim_first,
38
)
39
from ..ops.causal_conv1d import causal_conv1d_fn, causal_conv1d_update
40
+from ..ops.gather_initial_states import gather_initial_states
41
40
-logger = init_logger(__name__)
41
-
42
-
43
-def kda_attention(
44
- q_proj_states: torch.Tensor,
45
- k_proj_states: torch.Tensor,
46
- v_proj_states: torch.Tensor,
47
- g1: torch.Tensor,
48
- beta: torch.Tensor,
49
- core_attn_out: torch.Tensor,
50
- layer_name: str,
51
-) -> None:
52
- forward_context: ForwardContext = get_forward_context()
53
- self = forward_context.no_compile_layers[layer_name]
54
- self._forward(
55
- q_proj_states=q_proj_states,
56
- k_proj_states=k_proj_states,
57
- v_proj_states=v_proj_states,
58
- g1=g1,
59
- beta=beta,
60
- core_attn_out=core_attn_out,
61
- )
62
-
63
-
64
-def kda_attention_fake(
65
- q_proj_states: torch.Tensor,
66
- k_proj_states: torch.Tensor,
67
- v_proj_states: torch.Tensor,
68
- g1: torch.Tensor,
69
- beta: torch.Tensor,
70
- core_attn_out: torch.Tensor,
71
- layer_name: str,
72
-) -> None:
73
- return
74
-
75
-
76
-direct_register_custom_op(
77
- op_name="kda_attention",
78
- op_func=kda_attention,
79
- mutates_args=["core_attn_out"],
80
- fake_impl=kda_attention_fake,
81
-)
42
+# Empirical lower bound for the KDA gate to avoid numerical underflow.
43
+_KDA_GATE_LOGBOUND_MIN = -5.0
44
+
45
+
46
+def a_log_weight_loader(
47
+ shard_axis: int,
48
+) -> Callable[[torch.Tensor, torch.Tensor], None]:
49
+ """Load KDA A_log stored as either old 4D or current 1D weights."""
50
+
51
+ def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
52
+ tp_rank = get_tensor_model_parallel_rank()
53
+ shard_size = param.data.shape[shard_axis]
54
+ start_idx = tp_rank * shard_size
55
+
56
+ if loaded_weight.dim() == 4:
57
+ assert loaded_weight.shape[:2] == (1, 1), (
58
+ f"Expected old A_log shape (1, 1, H, 1), got {loaded_weight.shape}"
59
+ )
60
+ assert loaded_weight.shape[-1] == 1, (
61
+ f"Expected old A_log last dim to be 1, got {loaded_weight.shape}"
62
+ )
63
+ loaded_weight = loaded_weight.view(loaded_weight.shape[2])
64
+
65
+ loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size)
66
+ return default_weight_loader(param, loaded_weight)
67
+
68
+ return loader
69
+
70
+
71
+def _make_fused_conv1d_weight_loader(
72
+ dims: list[int],
73
+ tp_size: int,
74
+ tp_rank: int,
75
+) -> Callable[..., None]:
76
+ sharded_dims = [dim // tp_size for dim in dims]
77
+
78
+ def weight_loader(
79
+ param: torch.Tensor,
80
+ loaded_weight: torch.Tensor,
81
+ loaded_shard_id: int,
82
+ ) -> None:
83
+ if loaded_weight.dim() == 2:
84
+ loaded_weight = loaded_weight.unsqueeze(1)
85
+ shard_size = sharded_dims[loaded_shard_id]
86
+ source_start = tp_rank * shard_size
87
+ target_start = sum(sharded_dims[:loaded_shard_id])
88
+ loaded_shard = loaded_weight[source_start : source_start + shard_size]
89
+ param.data[target_start : target_start + shard_size].copy_(loaded_shard)
90
+
91
+ return weight_loader
92
+
93
+
94
+class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear):
95
+ """Merged projection with one output replicated across TP ranks.
96
+
97
+ The replicated shard is represented as ``size * tp_size`` so the merged
98
+ parameter reserves ``size`` local rows on every rank. Loading that shard
99
+ from rank zero then gives every rank the complete checkpoint weight.
100
+ """
101
+
102
+ def __init__(
103
+ self,
104
+ input_size: int,
105
+ output_sizes: list[int],
106
+ replicated_shard_id: int,
107
+ tp_size: int,
108
+ **kwargs,
109
+ ) -> None:
110
+ self.replicated_shard_id = replicated_shard_id
111
+ output_sizes = output_sizes.copy()
112
+ output_sizes[replicated_shard_id] *= tp_size
113
+ super().__init__(input_size, output_sizes, **kwargs)
114
+
115
+ def weight_loader(
116
+ self,
117
+ param: Parameter,
118
+ loaded_weight: torch.Tensor,
119
+ loaded_shard_id: tuple[int, ...] | int | None = None,
120
+ ) -> None:
121
+ tp_rank = self.tp_rank
122
+ param_tp_rank = getattr(param, "tp_rank", None)
123
+ if loaded_shard_id == self.replicated_shard_id:
124
+ self.tp_rank = 0
125
+ if param_tp_rank is not None:
126
+ param.tp_rank = 0
127
+ try:
128
+ super().weight_loader(param, loaded_weight, loaded_shard_id)
129
+ finally:
130
+ self.tp_rank = tp_rank
131
+ if param_tp_rank is not None:
132
+ param.tp_rank = param_tp_rank
133
+
134
+ def weight_loader_v2(
135
+ self,
136
+ param: BasevLLMParameter,
137
+ loaded_weight: torch.Tensor,
138
+ loaded_shard_id: tuple[int, ...] | int | None = None,
139
+ ) -> None:
140
+ tp_rank = self.tp_rank
141
+ param_tp_rank = getattr(param, "tp_rank", None)
142
+ if loaded_shard_id == self.replicated_shard_id:
143
+ self.tp_rank = 0
144
+ if param_tp_rank is not None:
145
+ param.tp_rank = 0
146
+ try:
147
+ super().weight_loader_v2(param, loaded_weight, loaded_shard_id)
148
+ finally:
149
+ self.tp_rank = tp_rank
150
+ if param_tp_rank is not None:
151
+ param.tp_rank = param_tp_rank
152
153
154
@PluggableLayer.register("kimi_gated_delta_net_attention")
@@ -96,7 +166,11 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
166
self,
167
) -> tuple[tuple[int, ...], tuple[int, ...]]:
168
return MambaStateShapeCalculator.kda_state_shape(
99
- self.tp_size, self.num_heads, self.head_dim, conv_kernel_size=self.conv_size
169
+ self.tp_size,
170
+ self.num_heads,
171
+ self.head_dim,
172
+ conv_kernel_size=self.conv_size,
173
+ num_spec=self.num_spec,
174
)
175
176
def __init__(
@@ -114,122 +188,142 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
188
assert self.num_heads % self.tp_size == 0
189
self.local_num_heads = divide(self.num_heads, self.tp_size)
190
117
- projection_size = self.head_dim * self.num_heads
191
+ self.projection_size = self.head_dim * self.num_heads
192
+ self.local_projection_size = divide(self.projection_size, self.tp_size)
193
self.conv_size = kda_config["short_conv_kernel_size"]
119
-
120
- self.q_proj = ColumnParallelLinear(
121
- self.hidden_size,
122
- projection_size,
123
- bias=False,
124
- quant_config=self.quant_config,
125
- prefix=f"{prefix}.q_proj",
126
- )
127
- self.k_proj = ColumnParallelLinear(
128
- self.hidden_size,
129
- projection_size,
130
- bias=False,
131
- quant_config=self.quant_config,
132
- prefix=f"{prefix}.k_proj",
133
- )
134
- self.v_proj = ColumnParallelLinear(
135
- self.hidden_size,
136
- projection_size,
137
- bias=False,
138
- quant_config=self.quant_config,
139
- prefix=f"{prefix}.v_proj",
140
- )
141
-
142
- self.f_a_proj = ReplicatedLinear(
194
+ self.use_full_rank_gate = kda_config.get("use_full_rank_gate", False)
195
+
196
+ if self.use_full_rank_gate:
197
+ # Keep f_a before the narrow beta shard, then pad each TP-local row
198
+ # to select the aligned BF16 GEMM path. The padding also avoids an
199
+ # Inductor correctness issue seen with the row-strided G view.
200
+ qkvg_output_sizes = [self.projection_size] * 4
201
+ in_proj_output_sizes = qkvg_output_sizes + [
202
+ self.head_dim,
203
+ self.num_heads,
204
+ ]
205
+ local_output_size = (
206
+ 4 * self.local_projection_size + self.head_dim + self.local_num_heads
207
+ )
208
+ self.in_proj_padding = -local_output_size % 16
209
+ if self.in_proj_padding:
210
+ in_proj_output_sizes.append(self.in_proj_padding * self.tp_size)
211
+ else:
212
+ in_proj_output_sizes = [self.projection_size] * 3 + [
213
+ self.num_heads,
214
+ self.head_dim,
215
+ ]
216
+ self.in_proj_padding = 0
217
+ self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear(
218
self.hidden_size,
144
- self.head_dim,
219
+ in_proj_output_sizes,
220
+ replicated_shard_id=4,
221
+ tp_size=self.tp_size,
222
bias=False,
223
quant_config=self.quant_config,
147
- prefix=f"{prefix}.f_a_proj",
224
+ prefix=f"{prefix}.in_proj_qkvgfab",
225
)
226
+ if self.in_proj_padding:
227
+ self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_()
228
229
self.f_b_proj = ColumnParallelLinear(
230
self.head_dim,
152
- projection_size,
231
+ self.projection_size,
232
bias=False,
233
quant_config=self.quant_config,
234
prefix=f"{prefix}.f_b_proj",
235
)
236
self.dt_bias = nn.Parameter(
158
- torch.empty(divide(projection_size, self.tp_size), dtype=torch.float32)
237
+ torch.empty(self.local_projection_size, dtype=torch.float32)
238
)
239
240
set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
241
163
- self.b_proj = ColumnParallelLinear(
164
- self.hidden_size,
165
- self.num_heads,
166
- bias=False,
167
- quant_config=self.quant_config,
168
- prefix=f"{prefix}.b_proj",
169
- )
170
-
171
- self.q_conv1d = ColumnParallelLinear(
242
+ # One packed parameter and cache let decode run a single conv update.
243
+ # Prefill slices them back into Q/K/V to obtain dense outputs cheaply.
244
+ self.conv1d = ColumnParallelLinear(
245
input_size=self.conv_size,
173
- output_size=projection_size,
246
+ output_size=3 * self.projection_size,
247
bias=False,
248
params_dtype=torch.float32,
176
- prefix=f"{prefix}.q_conv1d",
249
+ prefix=f"{prefix}.conv1d",
250
)
178
- self.k_conv1d = ColumnParallelLinear(
179
- input_size=self.conv_size,
180
- output_size=projection_size,
181
- bias=False,
182
- params_dtype=torch.float32,
183
- prefix=f"{prefix}.k_conv1d",
251
+ self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
252
+ delattr(self.conv1d.weight, "weight_loader")
253
+ set_weight_attrs(
254
+ self.conv1d.weight,
255
+ {
256
+ "weight_loader": _make_fused_conv1d_weight_loader(
257
+ [self.projection_size] * 3,
258
+ self.tp_size,
259
+ self.tp_rank,
260
+ )
261
+ },
262
)
185
- self.v_conv1d = ColumnParallelLinear(
186
- input_size=self.conv_size,
187
- output_size=projection_size,
188
- bias=False,
189
- params_dtype=torch.float32,
190
- prefix=f"{prefix}.v_conv1d",
191
- )
192
- # unsqueeze to fit conv1d weights shape into the linear weights shape.
193
- # Can't do this in `weight_loader` since it already exists in
194
- # `ColumnParallelLinear` and `set_weight_attrs`
195
- # doesn't allow to override it
196
- self.q_conv1d.weight.data = self.q_conv1d.weight.data.unsqueeze(1)
197
- self.k_conv1d.weight.data = self.k_conv1d.weight.data.unsqueeze(1)
198
- self.v_conv1d.weight.data = self.v_conv1d.weight.data.unsqueeze(1)
263
264
self.A_log = nn.Parameter(
201
- torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32)
265
+ torch.empty(self.local_num_heads, dtype=torch.float32)
266
)
203
- set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(2)})
204
-
205
- self.g_a_proj = ReplicatedLinear(
206
- self.hidden_size,
207
- self.head_dim,
208
- bias=False,
209
- quant_config=self.quant_config,
210
- prefix=f"{prefix}.g_a_proj",
267
+ set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)})
268
+
269
+ self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None)
270
+ if self.gate_lower_bound is not None:
271
+ assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, (
272
+ "KDA gate lower bound must be in "
273
+ f"[{_KDA_GATE_LOGBOUND_MIN}, 0). "
274
+ f"Got {self.gate_lower_bound}."
275
+ )
276
+ self.use_safe_gate = self.gate_lower_bound is not None
277
+ additional_config = vllm_config.additional_config
278
+ backend = (
279
+ additional_config.get("kda_prefill_backend", "auto")
280
+ if isinstance(additional_config, dict)
281
+ else "auto"
282
)
212
- self.g_b_proj = ColumnParallelLinear(
213
- self.head_dim,
214
- projection_size,
215
- bias=False,
216
- quant_config=self.quant_config,
217
- prefix=f"{prefix}.g_b_proj",
283
+ backend = "triton" if backend == "auto" else backend
284
+ assert backend == "triton", (
285
+ "The shared Kimi GDN layer only supports the Triton KDA "
286
+ f"prefill backend, got {backend!r}."
287
)
288
+ if not self.use_full_rank_gate:
289
+ self.g_a_proj = ReplicatedLinear(
290
+ self.hidden_size,
291
+ self.head_dim,
292
+ bias=False,
293
+ quant_config=self.quant_config,
294
+ prefix=f"{prefix}.g_a_proj",
295
+ )
296
+ self.g_b_proj = ColumnParallelLinear(
297
+ self.head_dim,
298
+ self.projection_size,
299
+ bias=False,
300
+ quant_config=self.quant_config,
301
+ prefix=f"{prefix}.g_b_proj",
302
+ )
303
self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
304
self.o_proj = RowParallelLinear(
221
- projection_size,
305
+ self.projection_size,
306
self.hidden_size,
307
bias=False,
308
quant_config=self.quant_config,
309
prefix=f"{prefix}.o_proj",
310
)
311
228
- compilation_config = get_current_vllm_config().compilation_config
312
+ compilation_config = vllm_config.compilation_config
313
if prefix in compilation_config.static_forward_context:
314
raise ValueError(f"Duplicate layer name: {prefix}")
315
compilation_config.static_forward_context[prefix] = self
316
317
+ def rearrange_mixed_qkv(
318
+ self, mixed_qkv: torch.Tensor
319
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
320
+ seq_len = mixed_qkv.shape[0]
321
+ qkv = mixed_qkv.view(seq_len, 3, self.local_num_heads, self.head_dim)
322
+ # Materialize all three row-strided inputs with one token-major to
323
+ # QKV-major permutation. Each unbound tensor is then contiguous.
324
+ qkv = qkv.permute(1, 0, 2, 3).contiguous().unsqueeze(1)
325
+ return qkv.unbind(0)
326
+
327
def forward(
328
self,
329
hidden_states: torch.Tensor,
@@ -237,42 +331,57 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
331
output: torch.Tensor,
332
) -> None:
333
num_tokens = hidden_states.size(0)
240
- q = self.q_proj(hidden_states)[0]
241
- k = self.k_proj(hidden_states)[0]
242
- v = self.v_proj(hidden_states)[0]
334
+ projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
335
+ if self.use_full_rank_gate:
336
+ split_sizes = [
337
+ 3 * self.local_projection_size,
338
+ self.local_projection_size,
339
+ self.head_dim,
340
+ self.local_num_heads,
341
+ ]
342
+ if self.in_proj_padding:
343
+ split_sizes.append(self.in_proj_padding)
344
+ projected = projected_qkvgfab.split(split_sizes, dim=-1)
345
+ mixed_qkv, g_proj_states, f_a, beta = projected[:4]
346
+ else:
347
+ mixed_qkv, beta, f_a = projected_qkvgfab.split(
348
+ [
349
+ 3 * self.local_projection_size,
350
+ self.local_num_heads,
351
+ self.head_dim,
352
+ ],
353
+ dim=-1,
354
+ )
355
+ g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
356
244
- beta = self.b_proj(hidden_states)[0].float().sigmoid()
245
- g1 = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
357
+ g1 = self.f_b_proj(f_a)[0]
358
beta = beta.unsqueeze(0)
359
g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim)
360
249
- g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
361
g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
362
252
- core_attn_out = torch.zeros(
363
+ core_attn_out = torch.empty(
364
(1, num_tokens, self.local_num_heads, self.head_dim),
365
dtype=hidden_states.dtype,
366
device=hidden_states.device,
367
)
257
- torch.ops.vllm.kda_attention(
258
- q,
259
- k,
260
- v,
261
- g1,
262
- beta,
263
- core_attn_out,
264
- self.prefix,
368
+
369
+ self._forward(
370
+ mixed_qkv=mixed_qkv,
371
+ g1=g1,
372
+ g2=g2,
373
+ beta=beta,
374
+ core_attn_out=core_attn_out,
375
)
266
- core_attn_out = self.o_norm(core_attn_out, g2)
376
core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
377
output[:] = self.o_proj(core_attn_out)[0]
378
379
+ @eager_break_during_capture
380
def _forward(
381
self,
272
- q_proj_states: torch.Tensor,
273
- k_proj_states: torch.Tensor,
274
- v_proj_states: torch.Tensor,
382
+ mixed_qkv: torch.Tensor,
383
g1: torch.Tensor,
384
+ g2: torch.Tensor,
385
beta: torch.Tensor,
386
core_attn_out: torch.Tensor,
387
) -> None:
@@ -280,165 +389,247 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
389
attn_metadata_raw = forward_context.attn_metadata
390
391
if attn_metadata_raw is None:
283
- # # V1 profile run
392
return
393
394
+ # Vendor-specific KDA kernels: AMD/ROCm and NVIDIA keep their own copies
395
+ # under kimi_k3/{amd,nvidia}/ops so each can diverge independently.
396
+ if current_platform.is_rocm():
397
+ from vllm.models.kimi_k3.amd.ops.third_party.kda import (
398
+ chunk_kda_with_fused_gate,
399
+ fused_recurrent_kda,
400
+ fused_recurrent_kda_packed_decode,
401
+ )
402
+ else:
403
+ from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
404
+ chunk_kda_with_fused_gate,
405
+ fused_recurrent_kda,
406
+ fused_recurrent_kda_packed_decode,
407
+ )
408
+
409
assert isinstance(attn_metadata_raw, dict)
410
attn_metadata_narrowed = attn_metadata_raw[self.prefix]
411
assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata)
289
- has_initial_state = attn_metadata_narrowed.has_initial_state
290
- non_spec_query_start_loc = attn_metadata_narrowed.non_spec_query_start_loc
291
- non_spec_state_indices_tensor = (
292
- attn_metadata_narrowed.non_spec_state_indices_tensor
293
- ) # noqa: E501
294
- num_actual_tokens = attn_metadata_narrowed.num_actual_tokens
295
- constant_caches = self.kv_cache
296
-
297
- q_proj_states = q_proj_states[:num_actual_tokens]
298
- k_proj_states = k_proj_states[:num_actual_tokens]
299
- v_proj_states = v_proj_states[:num_actual_tokens]
412
+ m = attn_metadata_narrowed
413
+ has_initial_state = m.has_initial_state
414
+ non_spec_query_start_loc = m.non_spec_query_start_loc
415
+ non_spec_state_indices_tensor = m.non_spec_state_indices_tensor
416
+ spec_sequence_masks = m.spec_sequence_masks
417
+ spec_token_indx = m.spec_token_indx
418
+ non_spec_token_indx = m.non_spec_token_indx
419
+ spec_state_indices_tensor = m.spec_state_indices_tensor
420
+ spec_query_start_loc = m.spec_query_start_loc
421
+ num_accepted_tokens = m.num_accepted_tokens
422
+ num_actual_tokens = m.num_actual_tokens
423
+ mixed_qkv = mixed_qkv[:num_actual_tokens]
424
g1 = g1[:, :num_actual_tokens]
425
beta = beta[:, :num_actual_tokens]
426
303
- (conv_state, recurrent_state) = constant_caches
427
+ constant_caches = self.kv_cache
428
+
429
+ conv_state, recurrent_state = constant_caches
430
# conv_state must be (..., dim, width-1) for the conv kernels.
431
# DS layout stores it that way directly; SD layout needs a transpose.
432
if not is_conv_state_dim_first():
433
conv_state = conv_state.transpose(-1, -2)
434
309
- conv_state_q, conv_state_k, conv_state_v = conv_state.chunk(3, dim=-2)
310
-
311
- q_conv_weights = self.q_conv1d.weight.view(
312
- self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2)
435
+ conv_weights = self.conv1d.weight.view(
436
+ self.conv1d.weight.size(0), self.conv1d.weight.size(2)
437
)
314
- k_conv_weights = self.k_conv1d.weight.view(
315
- self.k_conv1d.weight.size(0), self.k_conv1d.weight.size(2)
438
+ q_conv_weight, k_conv_weight, v_conv_weight = conv_weights.split(
439
+ self.local_projection_size, dim=0
440
)
317
- v_conv_weights = self.v_conv1d.weight.view(
318
- self.v_conv1d.weight.size(0), self.v_conv1d.weight.size(2)
441
+ q_conv_state, k_conv_state, v_conv_state = conv_state.split(
442
+ self.local_projection_size, dim=-2
443
)
320
- if attn_metadata_narrowed.num_prefills > 0:
321
- q_proj_states = q_proj_states.transpose(0, 1)
322
- k_proj_states = k_proj_states.transpose(0, 1)
323
- v_proj_states = v_proj_states.transpose(0, 1)
324
- q = causal_conv1d_fn(
325
- q_proj_states,
326
- q_conv_weights,
327
- self.q_conv1d.bias,
328
- activation="silu",
329
- conv_states=conv_state_q,
330
- has_initial_state=has_initial_state,
331
- cache_indices=non_spec_state_indices_tensor,
332
- query_start_loc=non_spec_query_start_loc,
333
- metadata=attn_metadata_narrowed,
334
- ).transpose(0, 1)
335
- k = causal_conv1d_fn(
336
- k_proj_states,
337
- k_conv_weights,
338
- self.k_conv1d.bias,
339
- activation="silu",
340
- conv_states=conv_state_k,
341
- has_initial_state=has_initial_state,
342
- cache_indices=non_spec_state_indices_tensor,
343
- query_start_loc=non_spec_query_start_loc,
344
- metadata=attn_metadata_narrowed,
345
- ).transpose(0, 1)
346
- v = causal_conv1d_fn(
347
- v_proj_states,
348
- v_conv_weights,
349
- self.v_conv1d.bias,
350
- activation="silu",
351
- conv_states=conv_state_v,
352
- has_initial_state=has_initial_state,
353
- cache_indices=non_spec_state_indices_tensor,
354
- query_start_loc=non_spec_query_start_loc,
355
- metadata=attn_metadata_narrowed,
356
- ).transpose(0, 1)
444
+
445
+ # Split tokens into the multi-query spec-decode part and the remaining
446
+ # (prefill / plain decode) part.
447
+ if spec_sequence_masks is not None:
448
+ if m.num_prefills == 0 and m.num_decodes == 0:
449
+ mixed_qkv_spec = mixed_qkv
450
+ g1_spec, beta_spec = g1, beta
451
+ mixed_qkv_ns = g1_ns = beta_ns = None
452
+ else:
453
+ mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
454
+ g1_spec = g1.index_select(1, spec_token_indx)
455
+ beta_spec = beta.index_select(1, spec_token_indx)
456
+ mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx)
457
+ g1_ns = g1.index_select(1, non_spec_token_indx)
458
+ beta_ns = beta.index_select(1, non_spec_token_indx)
459
else:
358
- assert non_spec_state_indices_tensor is not None
359
- decode_conv_indices = non_spec_state_indices_tensor[
360
- : attn_metadata_narrowed.num_actual_tokens
361
- ]
362
- q = causal_conv1d_update(
363
- q_proj_states,
364
- conv_state_q,
365
- q_conv_weights,
366
- self.q_conv1d.bias,
367
- activation="silu",
368
- conv_state_indices=decode_conv_indices,
369
- validate_data=True,
460
+ mixed_qkv_spec = g1_spec = beta_spec = None
461
+ mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta
462
+
463
+ # ---------- spec-decode multi-query path ----------
464
+ core_attn_out_spec = None
465
+ if spec_sequence_masks is not None:
466
+ assert spec_state_indices_tensor is not None
467
+ assert spec_query_start_loc is not None
468
+ spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes]
469
+ spec_max_query_len = spec_state_indices_tensor.size(-1)
470
+
471
+ # Sibling beta and, for full-rank gates, output-gate views remain
472
+ # live, so write the convolution output separately.
473
+ spec_conv_out = torch.empty(
474
+ mixed_qkv_spec.shape,
475
+ dtype=mixed_qkv_spec.dtype,
476
+ device=mixed_qkv_spec.device,
477
)
371
- k = causal_conv1d_update(
372
- k_proj_states,
373
- conv_state_k,
374
- k_conv_weights,
375
- self.k_conv1d.bias,
478
+ mixed_qkv_spec = causal_conv1d_update(
479
+ mixed_qkv_spec,
480
+ conv_state,
481
+ conv_weights,
482
+ self.conv1d.bias,
483
activation="silu",
377
- conv_state_indices=decode_conv_indices,
378
- validate_data=True,
484
+ conv_state_indices=spec_conv_indices,
485
+ num_accepted_tokens=num_accepted_tokens,
486
+ query_start_loc=spec_query_start_loc,
487
+ max_query_len=spec_max_query_len,
488
+ validate_data=False,
489
+ out=spec_conv_out,
490
)
380
- v = causal_conv1d_update(
381
- v_proj_states,
382
- conv_state_v,
383
- v_conv_weights,
384
- self.v_conv1d.bias,
385
- activation="silu",
386
- conv_state_indices=decode_conv_indices,
387
- validate_data=True,
491
+ q_spec, k_spec, v_spec = (
492
+ rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
493
+ for x in mixed_qkv_spec.split(self.local_projection_size, dim=-1)
494
)
389
-
390
- q, k, v = map(
391
- lambda x: rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim), (q, k, v)
392
- )
393
-
394
- if attn_metadata_narrowed.num_prefills > 0:
395
- assert non_spec_state_indices_tensor is not None
396
- assert has_initial_state is not None
397
- zero_idx = non_spec_state_indices_tensor[~has_initial_state]
398
- recurrent_state[zero_idx] = 0
399
- initial_state = recurrent_state[non_spec_state_indices_tensor].contiguous()
400
- (
401
- core_attn_out_non_spec,
402
- last_recurrent_state,
403
- ) = chunk_kda_with_fused_gate(
404
- q=q,
405
- k=k,
406
- v=v,
407
- raw_g=g1,
408
- beta=beta,
409
- A_log=self.A_log,
410
- g_bias=self.dt_bias,
411
- initial_state=initial_state,
412
- output_final_state=True,
413
- use_qk_l2norm_in_kernel=True,
414
- cu_seqlens=non_spec_query_start_loc,
495
+ spec_cu_seqlens = spec_query_start_loc[: m.num_spec_decodes + 1]
496
+ # Spec-only batches write directly into core_attn_out.
497
+ spec_out = (
498
+ core_attn_out[:, : q_spec.shape[1]]
499
+ if m.num_prefills == 0 and m.num_decodes == 0
500
+ else None
501
)
416
- # Init cache
417
- recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state
418
- else:
419
- assert non_spec_query_start_loc is not None
420
- g1 = fused_kda_gate(
421
- rearrange(g1, "1 n h d -> n (h d)"),
422
- self.A_log,
423
- self.head_dim,
424
- g_bias=self.dt_bias,
425
- ).unsqueeze(0)
426
- (
427
- core_attn_out_non_spec,
428
- last_recurrent_state,
429
- ) = fused_recurrent_kda(
430
- q=q,
431
- k=k,
432
- v=v,
433
- g=g1,
434
- beta=beta,
502
+ core_attn_out_spec, _ = fused_recurrent_kda(
503
+ q=q_spec,
504
+ k=k_spec,
505
+ v=v_spec,
506
+ raw_g=g1_spec,
507
+ raw_beta=beta_spec,
508
+ A_log=self.A_log,
509
+ dt_bias=self.dt_bias,
510
+ lower_bound=self.gate_lower_bound,
511
initial_state=recurrent_state,
436
- use_qk_l2norm_in_kernel=True,
437
- cu_seqlens=non_spec_query_start_loc[
438
- : attn_metadata_narrowed.num_decodes + 1
439
- ],
440
- ssm_state_indices=non_spec_state_indices_tensor,
512
+ cu_seqlens=spec_cu_seqlens,
513
+ ssm_state_indices=spec_state_indices_tensor,
514
+ num_accepted_tokens=num_accepted_tokens,
515
+ out=spec_out,
516
)
442
- core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
443
- 0, :num_actual_tokens
444
- ]
517
+
518
+ # ---------- non-spec path (prefill or plain decode) ----------
519
+ core_attn_out_non_spec = None
520
+ if mixed_qkv_ns is not None:
521
+ assert g1_ns is not None and beta_ns is not None
522
+ if m.num_prefills > 0:
523
+ q_ns, k_ns, v_ns = mixed_qkv_ns.split(
524
+ self.local_projection_size, dim=-1
525
+ )
526
+
527
+ # Packed prefill conv would require copying V solely to make
528
+ # it dense for KDA. Separate calls accept the strided inputs
529
+ # and produce dense Q/K/V without that extra traffic.
530
+ # TODO: Use packed conv once every KDA prefill backend accepts
531
+ # row-strided Q/K/V directly.
532
+ def _prefill_conv(
533
+ x: torch.Tensor,
534
+ state: torch.Tensor,
535
+ weight: torch.Tensor,
536
+ ) -> torch.Tensor:
537
+ return causal_conv1d_fn(
538
+ x.transpose(0, 1),
539
+ weight,
540
+ None,
541
+ activation="silu",
542
+ conv_states=state,
543
+ has_initial_state=has_initial_state,
544
+ cache_indices=non_spec_state_indices_tensor,
545
+ query_start_loc=non_spec_query_start_loc,
546
+ metadata=m,
547
+ ).transpose(0, 1)
548
+
549
+ q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight)
550
+ k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight)
551
+ v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight)
552
+ q_ns, k_ns, v_ns = (
553
+ rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
554
+ for x in (q_ns, k_ns, v_ns)
555
+ )
556
+
557
+ assert non_spec_state_indices_tensor is not None
558
+ assert has_initial_state is not None
559
+ initial_state = gather_initial_states(
560
+ recurrent_state,
561
+ non_spec_state_indices_tensor,
562
+ has_initial_state,
563
+ )
564
+ (
565
+ core_attn_out_non_spec,
566
+ last_recurrent_state,
567
+ ) = chunk_kda_with_fused_gate(
568
+ q=q_ns,
569
+ k=k_ns,
570
+ v=v_ns,
571
+ raw_g=g1_ns,
572
+ raw_beta=beta_ns,
573
+ A_log=self.A_log,
574
+ g_bias=self.dt_bias,
575
+ lower_bound=self.gate_lower_bound,
576
+ initial_state=initial_state,
577
+ output_final_state=True,
578
+ use_qk_l2norm_in_kernel=True,
579
+ cu_seqlens=non_spec_query_start_loc,
580
+ )
581
+ # Init cache
582
+ recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state
583
+
584
+ else:
585
+ # pure-decode non-spec batch
586
+ assert non_spec_state_indices_tensor is not None
587
+ decode_conv_indices = non_spec_state_indices_tensor[
588
+ : mixed_qkv_ns.size(0)
589
+ ]
590
+ # Sibling beta and, for full-rank gates, output-gate views
591
+ # remain live, so write the conv output separately.
592
+ packed_conv_out = torch.empty(
593
+ mixed_qkv_ns.shape,
594
+ dtype=mixed_qkv_ns.dtype,
595
+ device=mixed_qkv_ns.device,
596
+ )
597
+ mixed_qkv_ns = causal_conv1d_update(
598
+ mixed_qkv_ns,
599
+ conv_state,
600
+ conv_weights,
601
+ self.conv1d.bias,
602
+ activation="silu",
603
+ conv_state_indices=decode_conv_indices,
604
+ validate_data=True,
605
+ out=packed_conv_out,
606
+ )
607
+ core_attn_out_non_spec, _ = fused_recurrent_kda_packed_decode(
608
+ mixed_qkv=mixed_qkv_ns,
609
+ raw_g=g1_ns,
610
+ raw_beta=beta_ns,
611
+ A_log=self.A_log,
612
+ dt_bias=self.dt_bias,
613
+ lower_bound=self.gate_lower_bound,
614
+ initial_state=recurrent_state,
615
+ state_indices=decode_conv_indices,
616
+ )
617
+
618
+ # ---------- merge spec and non-spec outputs ----------
619
+ if core_attn_out_spec is not None and core_attn_out_non_spec is not None:
620
+ # Mixed batches require indexed placement in the original order.
621
+ merged = torch.empty(
622
+ (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
623
+ dtype=core_attn_out_spec.dtype,
624
+ device=core_attn_out_spec.device,
625
+ )
626
+ merged.index_copy_(1, spec_token_indx, core_attn_out_spec)
627
+ merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
628
+ core_attn_out[0, :num_actual_tokens] = merged[0, :num_actual_tokens]
629
+ elif core_attn_out_non_spec is not None:
630
+ core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
631
+ 0, :num_actual_tokens
632
+ ]
633
+ else:
634
+ assert core_attn_out_spec is not None
635
+ core_attn_out.copy_(self.o_norm(core_attn_out, g2))
vllm/model_executor/layers/mamba/mamba_utils.py
+1
-1
@@ -288,7 +288,7 @@ class MambaStateShapeCalculator:
288
289
conv_dim = proj_size + 2 * proj_k_size
290
conv_state_shape = cls._orient_conv_shape(
291
- divide(conv_dim, tp_world_size), conv_kernel_size - 1
291
+ divide(conv_dim, tp_world_size), conv_kernel_size - 1 + num_spec
292
)
293
recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim)
294
return (conv_state_shape, recurrent_state_shape)
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+32
-6
@@ -8,11 +8,12 @@
8
import numpy as np
9
import torch
10
11
+from vllm.platforms import current_platform
12
from vllm.triton_utils import tl, triton
13
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID
14
15
15
-@triton.jit()
16
+@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"])
17
def _causal_conv1d_fwd_kernel( # continuous batching
18
# Pointers to matrices
19
x_ptr, # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences
@@ -33,7 +34,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching
34
o_ptr, # (dim, seqlen) - actually pointing to x_ptr
35
# Matrix dimensions
36
dim: tl.constexpr,
36
- num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines
37
+ num_cache_lines, # added to support vLLM larger cache lines
38
# Strides
39
stride_x_dim: tl.constexpr, # stride to get to next feature-value,
40
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
59
NP2_STATELEN: tl.constexpr,
60
BLOCK_M: tl.constexpr,
61
BLOCK_N: tl.constexpr,
62
+ launch_pdl: tl.constexpr,
63
):
64
conv_states_ptr = initial_states_ptr
65
conv_state_indices_ptr = cache_indices_ptr
@@ -68,6 +70,9 @@ def _causal_conv1d_fwd_kernel( # continuous batching
70
KERNEL_WIDTH - 1
71
) # can be passed via argument if it's not the same as this value
72
73
+ if launch_pdl:
74
+ tl.extra.cuda.gdc_wait()
75
+
76
# one program handles one chunk in a single sequence
77
# rather than mixing sequences - to make updating initial_states across sequences efficiently
78
@@ -79,6 +84,8 @@ def _causal_conv1d_fwd_kernel( # continuous batching
84
idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
85
86
if idx_seq == pad_slot_id:
87
+ if launch_pdl:
88
+ tl.extra.cuda.gdc_launch_dependents()
89
return
90
91
sequence_start_index = tl.load(query_start_loc_ptr + idx_seq)
@@ -136,6 +143,8 @@ def _causal_conv1d_fwd_kernel( # continuous batching
143
if HAS_NULL_BLOCK: # noqa
144
if conv_states_input_coord == null_block_id:
145
# not processing as this is a null block (padding)
146
+ if launch_pdl:
147
+ tl.extra.cuda.gdc_launch_dependents()
148
return
149
conv_states_base = (
150
conv_states_ptr
@@ -408,6 +417,10 @@ def _causal_conv1d_fwd_kernel( # continuous batching
417
w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor
418
w_col3 = tl.load(w_ptrs, mask_w, other=0.0)
419
mask_x_1d = idx_feats < dim
420
+
421
+ if launch_pdl:
422
+ tl.extra.cuda.gdc_launch_dependents()
423
+
424
for idx_token in range(segment_len):
425
acc = acc_preload
426
@@ -741,11 +754,12 @@ def causal_conv1d_fn(
754
BLOCK_M=BLOCK_M,
755
BLOCK_N=256,
756
num_stages=2,
757
+ launch_pdl=current_platform.is_arch_support_pdl(),
758
)
759
return out.to(original_x_dtype)
760
761
748
-@triton.jit()
762
+@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"])
763
def _causal_conv1d_update_kernel(
764
# Pointers to matrices
765
x_ptr, # (batch, dim, seqlen)
@@ -763,7 +777,7 @@ def _causal_conv1d_update_kernel(
777
dim: tl.constexpr,
778
seqlen: tl.constexpr,
779
state_len: tl.constexpr,
766
- num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines
780
+ num_cache_lines, # added to support vLLM larger cache lines
781
# Strides
782
stride_x_seq: tl.constexpr,
783
stride_x_dim: tl.constexpr,
@@ -789,10 +803,16 @@ def _causal_conv1d_update_kernel(
803
NP2_STATELEN: tl.constexpr,
804
HAS_NULL_BLOCK: tl.constexpr,
805
BLOCK_N: tl.constexpr,
806
+ launch_pdl: tl.constexpr,
807
):
808
+ if launch_pdl:
809
+ tl.extra.cuda.gdc_wait()
810
+
811
# ruff: noqa: E501
812
idx_seq = tl.program_id(0)
813
if idx_seq >= batch:
814
+ if launch_pdl:
815
+ tl.extra.cuda.gdc_launch_dependents()
816
return
817
818
# [BLOCK_N,] elements along the feature-dimension (channel)
@@ -814,6 +834,8 @@ def _causal_conv1d_update_kernel(
834
if HAS_NULL_BLOCK: # noqa
835
if conv_states_input_coord == null_block_id:
836
# not processing as this is not the actual sequence
837
+ if launch_pdl:
838
+ tl.extra.cuda.gdc_launch_dependents()
839
return
840
841
if IS_VARLEN:
@@ -831,6 +853,8 @@ def _causal_conv1d_update_kernel(
853
o_offset = idx_seq * stride_o_seq
854
855
if query_start_index == query_end_index:
856
+ if launch_pdl:
857
+ tl.extra.cuda.gdc_launch_dependents()
858
return
859
860
if IS_SPEC_DECODING:
@@ -969,6 +993,9 @@ def _causal_conv1d_update_kernel(
993
mask_x_1d = idx_feats < dim
994
995
# STEP 5: compute each token
996
+ if launch_pdl:
997
+ tl.extra.cuda.gdc_launch_dependents()
998
+
999
for idx_token in tl.range(seqlen):
1000
acc = acc_preload
1001
@@ -1245,14 +1272,13 @@ def causal_conv1d_update(
1272
NP2_STATELEN=np2_statelen,
1273
HAS_NULL_BLOCK=null_block_id is not None,
1274
BLOCK_N=256,
1275
+ launch_pdl=current_platform.is_arch_support_pdl(),
1276
)
1277
if unsqueeze:
1278
out = out.squeeze(-1)
1279
return out.to(original_x_dtype)
1280
1281
1254
-from vllm.platforms import current_platform # noqa: E402
1255
-
1282
if current_platform.is_cpu():
1283
from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
1284
causal_conv1d_fn_cpu,
vllm/model_executor/layers/mla.py
+7
@@ -28,6 +28,7 @@ class MLAModules:
28
is_sparse: bool
29
topk_indices_buffer: torch.Tensor | None
30
indexer_rotary_emb: torch.nn.Module | None = None
31
+ g_proj: torch.nn.Module | None = None
32
33
34
# --8<-- [start:multi_head_latent_attention]
@@ -66,6 +67,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
67
quant_config: QuantizationConfig | None = None,
68
prefix: str = "",
69
skip_topk: bool = False,
70
+ non_causal_multi_token_decode: bool = False,
71
allow_short_prefill_indexer_scoring_skip: bool = False,
72
) -> None:
73
super().__init__()
@@ -89,6 +91,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
91
self.indexer = mla_modules.indexer
92
self.indexer_rope_emb = mla_modules.indexer_rotary_emb
93
self.is_sparse = mla_modules.is_sparse
94
+ self.g_proj = mla_modules.g_proj
95
96
# Whether to skip top-k token selection computation in this layer.
97
# When True, the indexer will not be called, and the layer will reuse
@@ -120,6 +123,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
123
use_sparse=self.is_sparse,
124
indexer=self.indexer,
125
topk_indices_buffer=mla_modules.topk_indices_buffer,
126
+ non_causal_multi_token_decode=non_causal_multi_token_decode,
127
)
128
indexer_op = getattr(self.indexer, "indexer_op", None)
129
if indexer_op is not None and hasattr(
@@ -216,4 +220,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
220
q_dcp_replicated=q_dcp_replicated,
221
)
222
223
+ if self.g_proj is not None:
224
+ attn_out = attn_out * self.g_proj(hidden_states)[0].sigmoid()
225
+
226
return self.o_proj(attn_out)[0]
vllm/model_executor/layers/quantization/mxfp4.py
+104
-4
@@ -3,6 +3,7 @@
3
4
import torch
5
6
+import vllm.envs as envs
7
from vllm.logger import init_logger
8
from vllm.model_executor.layers.attention import Attention
9
from vllm.model_executor.layers.fused_moe import (
@@ -17,6 +18,7 @@ from vllm.model_executor.layers.fused_moe import modular_kernel as mk
18
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
19
TRITON_BACKENDS,
20
Mxfp4MoeBackend,
21
+ backend_to_kernel_cls,
22
convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,
23
convert_weight_to_mxfp4_moe_kernel_format,
24
make_mxfp4_moe_kernel,
@@ -68,8 +70,11 @@ class Mxfp4Config(QuantizationConfig):
70
def get_config_filenames(cls) -> list[str]:
71
return []
72
71
- # TODO (zyongye) This is only temporaty fallback.
72
- # We should have `Mxfp4MoEMethod` after this migration is complete.
73
+ def _make_moe_method(self, moe: FusedMoEConfig) -> FusedMoEMethodBase:
74
+ """MoE method for RoutedExperts. Subclasses override to pick a
75
+ checkpoint-specific kernel family."""
76
+ return Mxfp4MoEMethod(moe)
77
+
78
def get_quant_method(
79
self, layer: torch.nn.Module, prefix: str
80
) -> "QuantizeMethodBase | None":
@@ -86,7 +91,7 @@ class Mxfp4Config(QuantizationConfig):
91
)
92
return UnquantizedLinearMethod()
93
elif isinstance(layer, RoutedExperts):
89
- return GptOssMxfp4MoEMethod(layer.moe_config)
94
+ return self._make_moe_method(layer.moe_config)
95
elif isinstance(layer, Attention):
96
logger.debug_once(
97
"MXFP4 attention layer is not implemented. "
@@ -131,6 +136,9 @@ class GptOssMxfp4Config(Mxfp4Config):
136
return None
137
return "gpt_oss_mxfp4"
138
139
+ def _make_moe_method(self, moe: FusedMoEConfig) -> FusedMoEMethodBase:
140
+ return GptOssMxfp4MoEMethod(moe)
141
+
142
143
class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
144
"""MXFP4 MoE quantization method."""
@@ -475,13 +483,45 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
483
)
484
485
486
+def _use_k3_situ_aiter(moe: FusedMoEConfig) -> bool:
487
+ """Whether Kimi-K3's SiTU MXFP4 MoE should use the AITER A16W4 kernel.
488
+
489
+ K3 is weight-only MXFP4 (W4A16) with SiTU activation, which the generic
490
+ MXFP4 backend selector does not cover; route it to AITER on gfx950.
491
+ """
492
+ from vllm.platforms import current_platform
493
+
494
+ if not current_platform.is_rocm():
495
+ return False
496
+ from vllm._aiter_ops import rocm_aiter_ops
497
+ from vllm.model_executor.layers.fused_moe.activation import MoEActivation
498
+ from vllm.platforms.rocm import on_gfx950
499
+
500
+ return (
501
+ rocm_aiter_ops.is_fused_moe_enabled()
502
+ and on_gfx950()
503
+ and moe.activation == MoEActivation.SITU
504
+ and moe.activation_situ_linear_beta is not None
505
+ and rocm_aiter_ops.get_aiter_activation_type("situ") is not None
506
+ )
507
+
508
+
509
class Mxfp4MoEMethod(FusedMoEMethodBase):
510
"""MXFP4 MoE quantization method."""
511
512
def __init__(self, moe: FusedMoEConfig):
513
super().__init__(moe)
514
self.weight_dtype = "mxfp4"
484
- self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(moe)
515
+ self.is_k3_situ_aiter = _use_k3_situ_aiter(moe)
516
+ self.experts_cls: type[mk.FusedMoEExperts] | None
517
+ if self.is_k3_situ_aiter:
518
+ self.mxfp4_backend = Mxfp4MoeBackend.AITER_MXFP4_BF16
519
+ self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0]
520
+ logger.info_once("Using AITER_MXFP4_BF16 for Kimi-K3 SiTU MXFP4 MoE.")
521
+ else:
522
+ self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(
523
+ moe
524
+ )
525
526
self.max_capture_size = moe.max_capture_size
527
@@ -523,6 +563,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
563
act_dtype=act_dtype,
564
moe_parallel_config=moe_parallel_config,
565
)
566
+ if self.is_k3_situ_aiter:
567
+ # K3's AITER A16W4 kernel handles K3's native intermediate size
568
+ # (moe_intermediate 3072; e.g. 384/partition at TP8); the generic
569
+ # 256 round-up would inflate weights and OOM.
570
+ return hidden_size, intermediate_size_per_partition
571
return mxfp4_round_up_hidden_size_and_intermediate_size(
572
self.mxfp4_backend, hidden_size, intermediate_size_per_partition
573
)
@@ -724,7 +769,58 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
769
layer=layer,
770
)
771
772
+ def _setup_kernel_k3_situ(self, layer: RoutedExperts) -> None:
773
+ # K3's AITER A16W4 kernel wants the separated ([gate_all, up_all])
774
+ # stage-1 layout, unlike the interleaved gpt-oss/DeepSeek path in
775
+ # convert_weight_to_mxfp4_moe_kernel_format. Preshuffle once here.
776
+ from aiter.utility.fp4_utils import e8m0_shuffle
777
+
778
+ from vllm._aiter_ops import rocm_aiter_ops
779
+
780
+ fp4_dtype = torch.float4_e2m1fn_x2
781
+ e8m0_dtype = torch.float8_e8m0fnu
782
+ num_experts = layer.w13_weight.shape[0]
783
+
784
+ # a8w4 (AITER_SITUV2_A8W4=1) uses the gate/up-interleaved (_gui_) fp8
785
+ # flydsl kernels, which need w13 weight+scale in interleave layout.
786
+ # Default a16w4 keeps the separated layout.
787
+ guinterleave = envs.AITER_SITUV2_A8W4
788
+ w13 = rocm_aiter_ops.shuffle_weight_a16w4(
789
+ layer.w13_weight.data.view(fp4_dtype), 16, guinterleave
790
+ )
791
+ w2 = rocm_aiter_ops.shuffle_weight_a16w4(
792
+ layer.w2_weight.data.view(fp4_dtype), 16, False
793
+ )
794
+ w13_scale_raw = layer.w13_weight_scale.data.view(e8m0_dtype)
795
+ w2_scale_raw = layer.w2_weight_scale.data.view(e8m0_dtype)
796
+ w13_scale = rocm_aiter_ops.shuffle_scale_a16w4(
797
+ w13_scale_raw.view(-1, w13_scale_raw.shape[-1]), num_experts, guinterleave
798
+ )
799
+ w2_scale = e8m0_shuffle(w2_scale_raw.view(-1, w2_scale_raw.shape[-1]))
800
+
801
+ replace_parameter(layer, "w13_weight", w13)
802
+ replace_parameter(layer, "w2_weight", w2)
803
+ replace_parameter(layer, "w13_weight_scale", w13_scale)
804
+ replace_parameter(layer, "w2_weight_scale", w2_scale)
805
+ layer.w13_weight.is_shuffled = True
806
+ layer.w2_weight.is_shuffled = True
807
+
808
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
809
+ if self.moe_quant_config is not None and self.experts_cls is not None:
810
+ self.moe_kernel = make_mxfp4_moe_kernel(
811
+ moe_quant_config=self.moe_quant_config,
812
+ moe_config=self.moe,
813
+ mxfp4_backend=self.mxfp4_backend,
814
+ experts_cls=self.experts_cls,
815
+ routing_tables=layer._expert_routing_tables(),
816
+ layer=layer,
817
+ )
818
+
819
def process_weights_after_loading(self, layer):
820
+ if self.is_k3_situ_aiter:
821
+ self._setup_kernel_k3_situ(layer)
822
+ return
823
+
824
w13 = layer.w13_weight
825
w2 = layer.w2_weight
826
w13_scale = layer.w13_weight_scale
@@ -819,4 +915,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
915
global_num_experts=layer.global_num_experts,
916
expert_map=layer.expert_map,
917
apply_router_weight_on_input=layer.apply_router_weight_on_input,
918
+ num_expert_group=layer.num_expert_group,
919
+ topk_group=layer.topk_group,
920
+ e_score_correction_bias=layer.e_score_correction_bias,
921
+ routed_scaling_factor=layer.routed_scaling_factor,
922
)
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+15
-1
@@ -18,20 +18,34 @@ def activation_to_flashinfer_int(activation: MoEActivation) -> int:
18
return activation_to_flashinfer_type(activation).value
19
20
21
+def has_flashinfer_situ_activation() -> bool:
22
+ try:
23
+ from flashinfer.fused_moe.core import ActivationType
24
+ except ImportError:
25
+ return False
26
+ return hasattr(ActivationType, "Situ")
27
+
28
+
29
def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType":
30
from flashinfer.fused_moe.core import ActivationType
31
32
+ if activation == MoEActivation.SITU:
33
+ situ = getattr(ActivationType, "Situ", None)
34
+ if situ is None:
35
+ raise ValueError("The installed FlashInfer does not support SITU")
36
+ return situ
37
+
38
# silu and gelu are mapped to their gated versions SwiGLU and GeGLU respectively
39
ACTIVATION_TO_FI_ACTIVATION = {
40
MoEActivation.SILU_NO_MUL: ActivationType.Silu,
41
MoEActivation.GELU_NO_MUL: ActivationType.Gelu,
42
MoEActivation.SILU: ActivationType.Swiglu,
43
# SwiGLU-OAI uses Swiglu; the OAI alpha/beta/clamp come from gemm1_* args.
44
+ MoEActivation.SWIGLUOAI: ActivationType.Swiglu,
45
MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu,
46
MoEActivation.GELU: ActivationType.Geglu,
47
MoEActivation.GELU_TANH: ActivationType.Geglu,
48
MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,
34
- MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu,
49
}
50
return ACTIVATION_TO_FI_ACTIVATION[activation]
51
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
+5
-1
@@ -474,7 +474,11 @@ def prepare_moe_fp4_layer_for_marlin(
474
475
group_size = 16 if is_nvfp4 else 32
476
477
- e = layer.moe_config.num_experts
477
+ # Use the per-rank (local) expert count: under expert parallelism the
478
+ # w13_weight/w2_weight tensors only hold this rank's experts (created with
479
+ # local_num_experts), whereas moe_config.num_experts is the global count.
480
+ # With no EP the two are equal, so the non-EP path is unchanged.
481
+ e = layer.moe_config.num_local_experts
482
k = layer.moe_config.hidden_dim
483
n = layer.moe_config.intermediate_size_per_partition
484
vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
+4
-3
@@ -41,8 +41,8 @@ def _mxfp8_e4m3_quantize_torch(
41
) -> tuple[torch.Tensor, torch.Tensor]:
42
"""Naive MXFP8 quantization.
43
For each block of 32 elements along the last dimension, compute a
44
- shared e8m0 scale (the biased exponent of the block-wise amax)
45
- and quantize each element to float8_e4m3fn.
44
+ shared e8m0 scale that fits the block-wise amax into the finite
45
+ float8_e4m3fn range, and quantize each element to float8_e4m3fn.
46
47
Returns (quantized_values [same shape, fp8], scales uint8).
48
Scale shape depends on is_sf_swizzled_layout:
@@ -58,7 +58,8 @@ def _mxfp8_e4m3_quantize_torch(
58
59
amax = x_blocked.abs().amax(dim=-1)
60
amax = amax.clamp(min=torch.finfo(torch.float32).tiny)
61
- scale_biased = torch.floor(torch.log2(amax)) + 127.0
61
+ fp8_max = torch.finfo(MXFP8_VALUE_DTYPE).max
62
+ scale_biased = torch.ceil(torch.log2(amax / fp8_max)) + 127.0
63
scale_biased = scale_biased.clamp(0, 254)
64
scales_uint8 = scale_biased.to(torch.uint8)
65
vllm/model_executor/model_loader/utils.py
+7
-7
@@ -16,10 +16,9 @@ import vllm.envs as envs
16
from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
17
from vllm.logger import init_logger
18
from vllm.model_executor.layers.attention import (
19
- Attention,
20
- MLAAttention,
19
MMEncoderAttention,
20
)
21
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
22
from vllm.model_executor.layers.hpc import HpcModule
23
from vllm.model_executor.layers.quantization.base_config import (
24
QuantizationConfig,
@@ -122,12 +121,13 @@ def process_weights_after_loading(
121
# the caching allocator, which starves the OS on UMA devices.
122
release_device_memory_under_pressure(target_device)
123
125
- # Initialize post-load attention weights for Attention, MLA, and MM encoder.
126
- # NOTE: Happens after other modules so we can easily decompress weights.
124
+ # Initialize post-load attention weights for any attention layer and MM
125
+ # encoder. NOTE: Happens after other modules so we can easily decompress
126
+ # weights.
127
for _, module in model.named_modules():
128
- if isinstance(
129
- module, (Attention, MLAAttention, MMEncoderAttention)
130
- ) and hasattr(module, "process_weights_after_loading"):
128
+ if isinstance(module, (AttentionLayerBase, MMEncoderAttention)) and hasattr(
129
+ module, "process_weights_after_loading"
130
+ ):
131
# TODO(lucas): see if there is a way to unify the signatures
132
# of process_weights_after_loading
133
with device_loading_context(module, target_device):
vllm/model_executor/models/config.py
+35
@@ -339,6 +339,39 @@ class DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig):
339
)
340
341
342
+class KimiK3ForConditionalGenerationConfig(VerifyAndUpdateConfig):
343
+ """Route MXFP4-checkpointed Kimi-K3 MoE experts to the MXFP4 interface.
344
+
345
+ Kimi-K3 ships its routed experts as compressed-tensors
346
+ ``mxfp4-pack-quantized`` (``quant_method="compressed-tensors"``), which
347
+ lands them on ``CompressedTensorsW4A4Mxfp4MoEMethod`` and its narrow kernel
348
+ selection. Rewriting ``quant_method`` to ``"mxfp4"`` selects ``Mxfp4Config``
349
+ (hence ``Mxfp4MoEMethod``) with its full backend set, while any non-MXFP4
350
+ checkpoint is left untouched. Covers both the main model and the MTP draft.
351
+
352
+ ``model_arch_config.quantization_config`` is a separate dict, snapshotted in
353
+ ``ModelConfig.__init__`` before this hook runs, and it is what
354
+ ``_verify_quantization`` reads when resolving the quant method. Patch it
355
+ alongside the hf configs so the rewrite lands before resolution; otherwise
356
+ the main model still resolves to compressed-tensors.
357
+ """
358
+
359
+ @staticmethod
360
+ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
361
+ for cfg in (
362
+ model_config.hf_config,
363
+ model_config.hf_text_config,
364
+ model_config.model_arch_config,
365
+ ):
366
+ quant_config = getattr(cfg, "quantization_config", None)
367
+ if (
368
+ isinstance(quant_config, dict)
369
+ and quant_config.get("quant_method") == "compressed-tensors"
370
+ and quant_config.get("format") == "mxfp4-pack-quantized"
371
+ ):
372
+ quant_config["quant_method"] = "mxfp4"
373
+
374
+
375
class GptOssForCausalLMConfig(VerifyAndUpdateConfig):
376
@staticmethod
377
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -884,6 +917,8 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
917
"JambaForSequenceClassification": JambaForSequenceClassificationConfig,
918
"JinaForRanking": JinaForRankingConfig,
919
"JinaVLForRanking": JinaVLForSequenceClassificationConfig,
920
+ "KimiK3ForConditionalGeneration": KimiK3ForConditionalGenerationConfig,
921
+ "KimiK3MTPModel": KimiK3ForConditionalGenerationConfig,
922
"LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig,
923
"LlamaBidirectionalModel": LlamaBidirectionalConfig,
924
"LlamaNemotronVLForSequenceClassification": LlamaNemotronVLConfig,
vllm/model_executor/models/deepseek_v2.py
+2
@@ -983,6 +983,7 @@ class DeepseekV2MLAAttention(nn.Module):
983
topk_indices_buffer: torch.Tensor | None = None,
984
input_size: int | None = None,
985
reduce_results: bool = True,
986
+ non_causal_multi_token_decode: bool = False,
987
) -> None:
988
super().__init__()
989
self.hidden_size = hidden_size
@@ -1178,6 +1179,7 @@ class DeepseekV2MLAAttention(nn.Module):
1179
# the V1 proposer. A frozen True would leave the draft reading a
1180
# never-written topk buffer.
1181
skip_topk=_skip_topk and not is_mtp_layer,
1182
+ non_causal_multi_token_decode=non_causal_multi_token_decode,
1183
# Do not skip scoring for MTP layers: their top-k buffer may be
1184
# reused by later draft iterations through index sharing.
1185
allow_short_prefill_indexer_scoring_skip=not is_mtp_layer,
vllm/model_executor/models/interfaces.py
+1
-1
@@ -1407,7 +1407,7 @@ class EagleModelMixin:
1407
aux_hidden_states: list[torch.Tensor],
1408
layer_idx: int,
1409
hidden_states: torch.Tensor,
1410
- residual: torch.Tensor,
1410
+ residual: torch.Tensor | None,
1411
) -> list[torch.Tensor]:
1412
if layer_idx in self.aux_hidden_state_layers:
1413
value = hidden_states + residual if residual is not None else hidden_states
vllm/model_executor/models/kimi_k25_vit.py
+204
-25
@@ -34,6 +34,7 @@ from vllm.model_executor.models.vision import (
34
is_vit_use_data_parallel,
35
run_dp_sharded_mrope_vision_model,
36
)
37
+from vllm.platforms import current_platform
38
from vllm.transformers_utils.configs.kimi_k25 import KimiK25VisionConfig
39
40
logger = init_logger(__name__)
@@ -154,9 +155,7 @@ class Learnable2DInterpPosEmbDivided_fixed(nn.Module):
155
def reset_parameters(self):
156
nn.init.normal_(self.weight)
157
157
- def forward(
158
- self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]]
159
- ) -> torch.Tensor:
158
+ def get_pos_embeds(self, grid_thws: torch.Tensor | list[list[int]]) -> torch.Tensor:
159
pos_embs = []
160
grid_thw_list = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist()
161
for t, h, w in grid_thw_list:
@@ -179,8 +178,12 @@ class Learnable2DInterpPosEmbDivided_fixed(nn.Module):
178
179
pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1]))
180
182
- out = x + torch.cat(pos_embs)
183
- return out
181
+ return torch.cat(pos_embs)
182
+
183
+ def forward(
184
+ self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]]
185
+ ) -> torch.Tensor:
186
+ return x + self.get_pos_embeds(grid_thws)
187
188
189
class MoonVision3dPatchEmbed(nn.Module):
@@ -195,6 +198,8 @@ class MoonVision3dPatchEmbed(nn.Module):
198
pos_emb_width: int = 14,
199
pos_emb_time: int = 4,
200
pos_emb_type: str = "divided_fixed",
201
+ patch_embed_proj_bias: bool = True,
202
+ pos_emb_interpolation_mode: str = "bicubic",
203
):
204
super().__init__()
205
assert isinstance(patch_size, int | Sequence), (
@@ -208,7 +213,11 @@ class MoonVision3dPatchEmbed(nn.Module):
213
self.patch_size = patch_size
214
215
self.proj = nn.Conv2d(
211
- in_dim, out_dim, kernel_size=patch_size, stride=patch_size
216
+ in_dim,
217
+ out_dim,
218
+ kernel_size=patch_size,
219
+ stride=patch_size,
220
+ bias=patch_embed_proj_bias,
221
)
222
223
if pos_emb_type == "divided_fixed":
@@ -217,17 +226,37 @@ class MoonVision3dPatchEmbed(nn.Module):
226
width=pos_emb_width,
227
num_frames=pos_emb_time,
228
dim=out_dim,
229
+ interpolation_mode=pos_emb_interpolation_mode,
230
)
231
else:
232
raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}")
233
234
def forward(
225
- self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]]
235
+ self,
236
+ x: torch.Tensor,
237
+ grid_thws: torch.Tensor | list[list[int]] | None,
238
+ *,
239
+ pos_embeds: torch.Tensor | None = None,
240
) -> torch.Tensor:
227
- x = self.proj(x).view(x.size(0), -1)
228
- # apply positional embedding
229
- x = self.pos_emb(x, grid_thws)
230
- return x
241
+ x = self._proj(x).view(x.size(0), -1)
242
+ if pos_embeds is not None:
243
+ return x + pos_embeds
244
+ assert grid_thws is not None
245
+ return self.pos_emb(x, grid_thws)
246
+
247
+ def _proj(self, x: torch.Tensor) -> torch.Tensor:
248
+ # MIOpen conv2d intermittently fails under load on ROCm; use aiter Triton.
249
+ if current_platform.is_rocm() and x.dtype in (torch.float16, torch.bfloat16):
250
+ from aiter.ops.triton.conv.conv2d import conv2d
251
+
252
+ return conv2d(
253
+ x,
254
+ self.proj.weight,
255
+ self.proj.bias,
256
+ stride=self.patch_size,
257
+ layout="nchw",
258
+ )
259
+ return self.proj(x)
260
261
262
class Rope2DPosEmbRepeated(nn.Module):
@@ -342,6 +371,14 @@ class MLP2(nn.Module):
371
return x
372
373
374
+def _make_vision_norm(norm_type: str, hidden_dim: int) -> nn.Module:
375
+ if norm_type == "layernorm":
376
+ return nn.LayerNorm(hidden_dim)
377
+ if norm_type == "rmsnorm":
378
+ return nn.RMSNorm(hidden_dim)
379
+ raise NotImplementedError(f"Not support norm_type: {norm_type}")
380
+
381
+
382
class MoonViTEncoderLayer(nn.Module):
383
"""Single encoder layer for MoonViT with TP/DP support."""
384
@@ -355,23 +392,33 @@ class MoonViTEncoderLayer(nn.Module):
392
*,
393
activation=F.gelu,
394
attn_bias: bool = False,
395
+ qkv_hidden_size: int | None = None,
396
+ norm_type: str = "layernorm",
397
+ mlp_type: str = "mlp2",
398
+ linear_bias: bool = True,
399
):
400
super().__init__()
401
self.use_data_parallel = is_vit_use_data_parallel(num_heads)
402
403
self.num_heads = num_heads
404
self.hidden_dim = hidden_dim
364
- self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads
405
+ self.qkv_hidden_size = (
406
+ hidden_dim if qkv_hidden_size is None else qkv_hidden_size
407
+ )
408
+ self.hidden_size_per_attention_head = self.qkv_hidden_size // self.num_heads
409
self.tp_size = (
410
1 if self.use_data_parallel else get_tensor_model_parallel_world_size()
411
)
412
self.num_attention_heads_per_partition = divide(num_heads, self.tp_size)
413
370
- self.norm0 = nn.LayerNorm(hidden_dim)
371
- self.norm1 = nn.LayerNorm(hidden_dim)
414
+ self.norm0 = _make_vision_norm(norm_type, hidden_dim)
415
+ self.norm1 = _make_vision_norm(norm_type, hidden_dim)
416
+ if mlp_type != "mlp2":
417
+ raise NotImplementedError(f"Not support mlp_type: {mlp_type}")
418
self.mlp = MLP2(
419
[hidden_dim, mlp_dim, hidden_dim],
420
activation,
421
+ bias=linear_bias,
422
quant_config=quant_config,
423
prefix=f"{prefix}.mlp",
424
use_data_parallel=self.use_data_parallel,
@@ -387,7 +434,7 @@ class MoonViTEncoderLayer(nn.Module):
434
disable_tp=self.use_data_parallel,
435
)
436
self.wo = RowParallelLinear(
390
- hidden_dim,
437
+ self.qkv_hidden_size,
438
hidden_dim,
439
bias=attn_bias,
440
quant_config=quant_config,
@@ -493,8 +540,9 @@ class MoonViT3dEncoder(nn.Module):
540
f'video_attn_type must be "spatial_temporal", got {video_attn_type}'
541
)
542
self.video_attn_type = video_attn_type
543
+ qkv_hidden_size = block_cfg.get("qkv_hidden_size") or block_cfg["hidden_dim"]
544
self.rope_2d = Rope2DPosEmbRepeated(
497
- block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512
545
+ qkv_hidden_size // block_cfg["num_heads"], 512, 512
546
)
547
self.blocks = nn.ModuleList(
548
[
@@ -506,13 +554,17 @@ class MoonViT3dEncoder(nn.Module):
554
for layer_idx in range(num_layers)
555
]
556
)
509
- self.final_layernorm = nn.LayerNorm(hidden_dim)
557
+ self.final_layernorm = _make_vision_norm(
558
+ block_cfg.get("norm_type", "layernorm"), hidden_dim
559
+ )
560
561
def prepare_encoder_metadata(
562
self,
563
grid_thw_list: list[list[int]],
564
*,
565
device: torch.device,
566
+ max_batch_size: int | None = None,
567
+ max_seqlen_override: int | None = None,
568
) -> dict[str, torch.Tensor | None]:
569
metadata: dict[str, torch.Tensor | None] = {}
570
metadata["rope_freqs_cis"] = self.rope_2d.get_freqs_cis(
@@ -524,15 +576,30 @@ class MoonViT3dEncoder(nn.Module):
576
cu_seqlens = np.concatenate(
577
[np.zeros(1, dtype=np.int32), lengths.cumsum(dtype=np.int32)]
578
)
579
+ if max_batch_size is not None:
580
+ num_seqs = len(cu_seqlens) - 1
581
+ if num_seqs < max_batch_size:
582
+ cu_seqlens = np.concatenate(
583
+ [
584
+ cu_seqlens,
585
+ np.full(
586
+ max_batch_size - num_seqs,
587
+ cu_seqlens[-1],
588
+ dtype=np.int32,
589
+ ),
590
+ ]
591
+ )
592
593
attn_backend = self.blocks[0].attn.attn_backend
594
metadata["sequence_lengths"] = MMEncoderAttention.maybe_compute_seq_lens(
595
attn_backend, cu_seqlens, device
596
)
532
- metadata["max_seqlen"] = torch.tensor(
533
- MMEncoderAttention.compute_max_seqlen(attn_backend, cu_seqlens),
534
- dtype=torch.int32,
597
+ max_seqlen = (
598
+ max_seqlen_override
599
+ if max_seqlen_override is not None
600
+ else MMEncoderAttention.compute_max_seqlen(attn_backend, cu_seqlens)
601
)
602
+ metadata["max_seqlen"] = torch.tensor(max_seqlen, dtype=torch.int32)
603
metadata["cu_seqlens"] = MMEncoderAttention.maybe_recompute_cu_seqlens(
604
attn_backend,
605
cu_seqlens,
@@ -545,11 +612,12 @@ class MoonViT3dEncoder(nn.Module):
612
def forward(
613
self,
614
hidden_states: torch.Tensor,
548
- grid_thws: torch.Tensor | list[list[int]],
615
+ grid_thws: torch.Tensor | list[list[int]] | None,
616
*,
617
encoder_metadata: dict[str, torch.Tensor | None] | None = None,
618
) -> torch.Tensor:
619
if encoder_metadata is None:
620
+ assert grid_thws is not None
621
grid_thw_list = (
622
grid_thws if isinstance(grid_thws, list) else grid_thws.tolist()
623
)
@@ -604,6 +672,34 @@ def tpool_patch_merger(
672
return outputs
673
674
675
+def build_image_merge_gather_idx(
676
+ grid_thws: list[list[int]] | list[tuple[int, int, int]],
677
+ merge_kernel_size: tuple[int, int],
678
+) -> np.ndarray:
679
+ """Build packed spatial-merge indices for image-only CUDA graphs."""
680
+ kh, kw = merge_kernel_size
681
+ parts: list[np.ndarray] = []
682
+ offset = 0
683
+ for t, h, w in grid_thws:
684
+ if t != 1:
685
+ raise ValueError("Image encoder CUDA graphs require grid T == 1")
686
+ idx = np.arange(h * w, dtype=np.int64).reshape(h, w)
687
+ idx = idx.reshape(h // kh, kh, w // kw, kw)
688
+ parts.append(idx.transpose(0, 2, 1, 3).reshape(-1, kh * kw) + offset)
689
+ offset += h * w
690
+ if not parts:
691
+ return np.empty((0, kh * kw), dtype=np.int64)
692
+ return np.concatenate(parts)
693
+
694
+
695
+def tpool_patch_merger_packed(
696
+ x: torch.Tensor,
697
+ merge_gather_idx: torch.Tensor,
698
+) -> torch.Tensor:
699
+ """Apply the image-only spatial merge using precomputed tensor indices."""
700
+ return x[merge_gather_idx]
701
+
702
+
703
class MoonViT3dPretrainedModel(nn.Module):
704
"""Main vision tower model.
705
@@ -630,6 +726,10 @@ class MoonViT3dPretrainedModel(nn.Module):
726
pos_emb_width=config.init_pos_emb_width,
727
pos_emb_time=config.init_pos_emb_time,
728
pos_emb_type=config.pos_emb_type,
729
+ patch_embed_proj_bias=getattr(config, "patch_embed_proj_bias", True),
730
+ pos_emb_interpolation_mode=getattr(
731
+ config, "pos_emb_interpolation_mode", "bicubic"
732
+ ),
733
)
734
735
self.encoder = MoonViT3dEncoder(
@@ -638,9 +738,15 @@ class MoonViT3dPretrainedModel(nn.Module):
738
block_cfg={
739
"num_heads": config.num_attention_heads,
740
"hidden_dim": config.hidden_size,
741
+ "qkv_hidden_size": getattr(config, "qkv_hidden_size", None),
742
"mlp_dim": config.intermediate_size,
642
- "activation": get_act_fn("gelu_pytorch_tanh"),
643
- "attn_bias": True,
743
+ "activation": get_act_fn(
744
+ getattr(config, "activation_func", "gelu_pytorch_tanh")
745
+ ),
746
+ "attn_bias": getattr(config, "attn_bias", True),
747
+ "norm_type": getattr(config, "norm_type", "layernorm"),
748
+ "mlp_type": getattr(config, "mlp_type", "mlp2"),
749
+ "linear_bias": getattr(config, "linear_bias", True),
750
},
751
video_attn_type=config.video_attn_type,
752
quant_config=quant_config,
@@ -650,7 +756,7 @@ class MoonViT3dPretrainedModel(nn.Module):
756
def forward(
757
self,
758
pixel_values: torch.Tensor,
653
- grid_thws: torch.Tensor | list[list[int]],
759
+ grid_thws: torch.Tensor | list[list[int]] | None,
760
*,
761
encoder_metadata: dict[str, torch.Tensor | None] | None = None,
762
) -> torch.Tensor:
@@ -662,6 +768,22 @@ class MoonViT3dPretrainedModel(nn.Module):
768
Returns:
769
torch.Tensor: The output tokens.
770
"""
771
+ if encoder_metadata is not None and "pos_embeds" in encoder_metadata:
772
+ hidden_states = self.patch_embed(
773
+ pixel_values,
774
+ None,
775
+ pos_embeds=encoder_metadata["pos_embeds"],
776
+ )
777
+ hidden_states = self.encoder(
778
+ hidden_states,
779
+ None,
780
+ encoder_metadata=encoder_metadata,
781
+ )
782
+ merge_gather_idx = encoder_metadata["merge_gather_idx"]
783
+ assert merge_gather_idx is not None
784
+ return tpool_patch_merger_packed(hidden_states, merge_gather_idx)
785
+
786
+ assert grid_thws is not None
787
grid_thw_list = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist()
788
if encoder_metadata is None:
789
encoder_metadata = self.encoder.prepare_encoder_metadata(
@@ -685,13 +807,40 @@ class MoonViT3dPretrainedModel(nn.Module):
807
808
return hidden_states
809
810
+ def prepare_encoder_cudagraph_metadata(
811
+ self,
812
+ grid_thw_list: list[list[int]],
813
+ *,
814
+ max_batch_size: int,
815
+ max_seqlen_override: int | None = None,
816
+ device: torch.device,
817
+ ) -> dict[str, torch.Tensor | None]:
818
+ """Precompute fixed-buffer metadata for image encoder CUDA graphs."""
819
+ grid_thw_list = [list(map(int, grid)) for grid in grid_thw_list]
820
+ metadata = self.encoder.prepare_encoder_metadata(
821
+ grid_thw_list,
822
+ device=device,
823
+ max_batch_size=max_batch_size,
824
+ max_seqlen_override=max_seqlen_override,
825
+ )
826
+ metadata["pos_embeds"] = self.patch_embed.pos_emb.get_pos_embeds(
827
+ grid_thw_list
828
+ ).to(device=device)
829
+ merge_gather_idx = build_image_merge_gather_idx(
830
+ grid_thw_list, self.merge_kernel_size
831
+ )
832
+ metadata["merge_gather_idx"] = torch.from_numpy(merge_gather_idx).to(
833
+ device=device, non_blocking=True
834
+ )
835
+ return metadata
836
+
837
838
@torch.inference_mode()
839
def mm_projector_forward(mm_projector: torch.nn.Module, vt_output: list[torch.Tensor]):
840
"""Apply MM projector to vision tower outputs."""
841
num_embedding_list = [x.shape[0] for x in vt_output]
842
batched = torch.cat(vt_output, dim=0)
694
- projector_dtype = mm_projector.pre_norm.weight.dtype
843
+ projector_dtype = next(mm_projector.parameters()).dtype
844
if batched.dtype != projector_dtype:
845
batched = batched.to(projector_dtype)
846
proj_out = mm_projector(batched)
@@ -747,11 +896,34 @@ class KimiK25MultiModalProjector(nn.Module):
896
):
897
super().__init__()
898
self.use_data_parallel = use_data_parallel
899
+ self.mm_projector_type = getattr(config, "mm_projector_type", "patchmerger")
900
901
# Hidden size after patch merging
902
merge_h, merge_w = config.merge_kernel_size
903
self.hidden_size = config.hidden_size * merge_h * merge_w
904
905
+ if self.mm_projector_type == "patchmergerv2":
906
+ self.linear_1 = ReplicatedLinear(
907
+ self.hidden_size,
908
+ self.hidden_size,
909
+ bias=False,
910
+ quant_config=quant_config,
911
+ prefix=f"{prefix}.linear_1",
912
+ )
913
+ self.linear_2 = ReplicatedLinear(
914
+ self.hidden_size,
915
+ getattr(config, "text_hidden_size", config.mm_hidden_size),
916
+ bias=False,
917
+ quant_config=quant_config,
918
+ prefix=f"{prefix}.linear_2",
919
+ )
920
+ self.post_norm = torch.nn.RMSNorm(
921
+ getattr(config, "text_hidden_size", config.mm_hidden_size),
922
+ eps=config.projector_ln_eps,
923
+ )
924
+ self.act = GELUActivation()
925
+ return
926
+
927
self.pre_norm = torch.nn.LayerNorm(config.hidden_size, eps=1e-5)
928
self.linear_1 = ReplicatedLinear(
929
self.hidden_size,
@@ -770,6 +942,13 @@ class KimiK25MultiModalProjector(nn.Module):
942
self.act = GELUActivation()
943
944
def forward(self, image_features: torch.Tensor) -> torch.Tensor:
945
+ if self.mm_projector_type == "patchmergerv2":
946
+ hidden_states = image_features.view(image_features.shape[0], -1)
947
+ hidden_states, _ = self.linear_1(hidden_states)
948
+ hidden_states = self.act(hidden_states)
949
+ hidden_states, _ = self.linear_2(hidden_states)
950
+ return self.post_norm(hidden_states)
951
+
952
hidden_states = self.pre_norm(image_features).view(-1, self.hidden_size)
953
hidden_states, _ = self.linear_1(hidden_states)
954
hidden_states = self.act(hidden_states)
vllm/model_executor/models/kimi_linear.py
deleted
-646
@@ -1,646 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from collections.abc import Iterable
5
-
6
-import torch
7
-from torch import nn
8
-
9
-from vllm.compilation.decorators import support_torch_compile
10
-from vllm.config import CacheConfig, VllmConfig
11
-from vllm.distributed import (
12
- get_pp_group,
13
- get_tensor_model_parallel_world_size,
14
-)
15
-from vllm.logger import init_logger
16
-from vllm.model_executor.layers.activation import SiluAndMul
17
-from vllm.model_executor.layers.fused_moe import (
18
- FusedMoE,
19
- fused_moe_make_expert_params_mapping,
20
-)
21
-from vllm.model_executor.layers.layernorm import RMSNorm
22
-from vllm.model_executor.layers.linear import (
23
- ColumnParallelLinear,
24
- MergedColumnParallelLinear,
25
- ReplicatedLinear,
26
- RowParallelLinear,
27
-)
28
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
29
-from vllm.model_executor.layers.mamba.gdn.kimi_gdn_linear_attn import (
30
- KimiGatedDeltaNetAttention,
31
-)
32
-from vllm.model_executor.layers.mamba.mamba_utils import (
33
- MambaStateCopyFunc,
34
- MambaStateCopyFuncCalculator,
35
- MambaStateDtypeCalculator,
36
- MambaStateShapeCalculator,
37
-)
38
-from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
39
-from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
40
-from vllm.model_executor.layers.vocab_parallel_embedding import (
41
- ParallelLMHead,
42
- VocabParallelEmbedding,
43
-)
44
-from vllm.model_executor.model_loader.weight_utils import (
45
- default_weight_loader,
46
- maybe_remap_kv_scale_name,
47
-)
48
-from vllm.sequence import IntermediateTensors
49
-from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
50
-
51
-from .interfaces import HasInnerState, IsHybrid, MixtureOfExperts, SupportsPP
52
-from .utils import (
53
- AutoWeightsLoader,
54
- PPMissingLayer,
55
- get_spec_layer_idx_from_weight_name,
56
- is_pp_missing_parameter,
57
- make_layers,
58
- maybe_prefix,
59
-)
60
-
61
-logger = init_logger(__name__)
62
-
63
-
64
-class KimiMLP(nn.Module):
65
- def __init__(
66
- self,
67
- hidden_size: int,
68
- intermediate_size: int,
69
- hidden_act: str,
70
- quant_config: QuantizationConfig | None = None,
71
- reduce_results: bool = True,
72
- prefix: str = "",
73
- ) -> None:
74
- super().__init__()
75
-
76
- self.gate_up_proj = MergedColumnParallelLinear(
77
- hidden_size,
78
- [intermediate_size] * 2,
79
- bias=False,
80
- quant_config=quant_config,
81
- prefix=f"{prefix}.gate_up_proj",
82
- )
83
- self.down_proj = RowParallelLinear(
84
- intermediate_size,
85
- hidden_size,
86
- bias=False,
87
- quant_config=quant_config,
88
- reduce_results=reduce_results,
89
- prefix=f"{prefix}.down_proj",
90
- )
91
- if hidden_act != "silu":
92
- raise ValueError(
93
- f"Unsupported activation: {hidden_act}. Only silu is supported for now."
94
- )
95
- self.act_fn = SiluAndMul()
96
-
97
- def forward(self, x):
98
- gate_up, _ = self.gate_up_proj(x)
99
- x = self.act_fn(gate_up)
100
- x, _ = self.down_proj(x)
101
- return x
102
-
103
-
104
-class KimiMoE(nn.Module):
105
- def __init__(
106
- self,
107
- config: KimiLinearConfig,
108
- quant_config: QuantizationConfig | None = None,
109
- prefix: str = "",
110
- layer_idx: int = 0,
111
- ):
112
- super().__init__()
113
- hidden_size = config.hidden_size
114
- intermediate_size = config.intermediate_size
115
- moe_intermediate_size = config.moe_intermediate_size
116
- num_experts = config.num_experts
117
- moe_renormalize = config.moe_renormalize
118
- self.tp_size = get_tensor_model_parallel_world_size()
119
- self.routed_scaling_factor = config.routed_scaling_factor
120
- self.num_shared_experts = config.num_shared_experts
121
- self.layer_idx = layer_idx
122
-
123
- if config.hidden_act != "silu":
124
- raise ValueError(
125
- f"Unsupported activation: {config.hidden_act}. "
126
- "Only silu is supported for now."
127
- )
128
-
129
- # Gate always runs at half / full precision for now.
130
- self.gate = ReplicatedLinear(
131
- hidden_size,
132
- num_experts,
133
- bias=False,
134
- quant_config=None,
135
- prefix=f"{prefix}.gate",
136
- )
137
-
138
- self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts))
139
-
140
- if self.num_shared_experts is not None:
141
- intermediate_size = moe_intermediate_size * self.num_shared_experts
142
- self.shared_experts = KimiMLP(
143
- hidden_size=config.hidden_size,
144
- intermediate_size=intermediate_size,
145
- hidden_act=config.hidden_act,
146
- quant_config=quant_config,
147
- reduce_results=False,
148
- prefix=f"{prefix}.shared_experts",
149
- )
150
- else:
151
- self.shared_experts = None
152
-
153
- self.experts = FusedMoE(
154
- shared_experts=self.shared_experts,
155
- num_experts=num_experts,
156
- top_k=config.num_experts_per_token,
157
- hidden_size=hidden_size,
158
- intermediate_size=moe_intermediate_size,
159
- renormalize=moe_renormalize,
160
- quant_config=quant_config,
161
- use_grouped_topk=config.use_grouped_topk,
162
- num_expert_group=config.num_expert_group,
163
- topk_group=config.topk_group,
164
- prefix=f"{prefix}.experts",
165
- scoring_func=config.moe_router_activation_func,
166
- e_score_correction_bias=self.gate.e_score_correction_bias,
167
- routed_scaling_factor=self.routed_scaling_factor,
168
- )
169
-
170
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
171
- num_tokens, hidden_size = hidden_states.shape
172
- hidden_states = hidden_states.view(-1, hidden_size)
173
- router_logits, _ = self.gate(hidden_states)
174
- final_hidden_states = self.experts(
175
- hidden_states=hidden_states, router_logits=router_logits
176
- )
177
- return final_hidden_states.view(num_tokens, hidden_size)
178
-
179
-
180
-class KimiMLAAttention(nn.Module):
181
- """
182
- Main reference: DeepseekV2 vllm Implementation
183
- """
184
-
185
- def __init__(
186
- self,
187
- config: KimiLinearConfig,
188
- hidden_size: int,
189
- num_heads: int,
190
- qk_nope_head_dim: int,
191
- qk_rope_head_dim: int,
192
- v_head_dim: int,
193
- q_lora_rank: int | None,
194
- kv_lora_rank: int,
195
- use_nope: bool = False,
196
- cache_config: CacheConfig | None = None,
197
- quant_config: QuantizationConfig | None = None,
198
- prefix: str = "",
199
- **kwargs,
200
- ) -> None:
201
- super().__init__()
202
- self.hidden_size = hidden_size
203
- self.qk_nope_head_dim = qk_nope_head_dim
204
- self.qk_rope_head_dim = qk_rope_head_dim
205
- self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
206
- self.v_head_dim = v_head_dim
207
- self.q_lora_rank = q_lora_rank
208
- self.kv_lora_rank = kv_lora_rank
209
- self.num_heads = num_heads
210
- tp_size = get_tensor_model_parallel_world_size()
211
- self.num_local_heads = num_heads // tp_size
212
- self.scaling = self.qk_head_dim**-0.5
213
- self.use_nope = use_nope
214
- assert self.use_nope is True
215
- assert self.q_lora_rank is None
216
- assert num_heads % tp_size == 0
217
- self.kv_a_proj_with_mqa = ReplicatedLinear(
218
- self.hidden_size,
219
- self.kv_lora_rank + self.qk_rope_head_dim,
220
- bias=False,
221
- quant_config=quant_config,
222
- prefix=f"{prefix}.kv_a_proj_with_mqa",
223
- )
224
- self.q_proj = ColumnParallelLinear(
225
- self.hidden_size,
226
- self.num_heads * self.qk_head_dim,
227
- bias=False,
228
- quant_config=quant_config,
229
- prefix=f"{prefix}.q_proj",
230
- )
231
- self.kv_a_layernorm = RMSNorm(
232
- self.kv_lora_rank,
233
- eps=config.rms_norm_eps,
234
- )
235
- self.kv_b_proj = ColumnParallelLinear(
236
- self.kv_lora_rank,
237
- self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
238
- bias=False,
239
- quant_config=quant_config,
240
- prefix=f"{prefix}.kv_b_proj",
241
- )
242
- self.o_proj = RowParallelLinear(
243
- self.num_heads * self.v_head_dim,
244
- self.hidden_size,
245
- bias=False,
246
- quant_config=quant_config,
247
- prefix=f"{prefix}.o_proj",
248
- )
249
-
250
- mla_modules = MLAModules(
251
- kv_a_layernorm=self.kv_a_layernorm,
252
- kv_b_proj=self.kv_b_proj,
253
- rotary_emb=None,
254
- o_proj=self.o_proj,
255
- fused_qkv_a_proj=None,
256
- kv_a_proj_with_mqa=self.kv_a_proj_with_mqa,
257
- q_a_layernorm=None,
258
- q_b_proj=None,
259
- q_proj=self.q_proj,
260
- indexer=None,
261
- is_sparse=False,
262
- topk_indices_buffer=None,
263
- )
264
- self.mla_attn = MultiHeadLatentAttentionWrapper(
265
- self.hidden_size,
266
- self.num_local_heads,
267
- self.scaling,
268
- self.qk_nope_head_dim,
269
- self.qk_rope_head_dim,
270
- self.v_head_dim,
271
- self.q_lora_rank,
272
- self.kv_lora_rank,
273
- mla_modules,
274
- cache_config,
275
- quant_config,
276
- prefix,
277
- )
278
-
279
- def forward(
280
- self,
281
- positions: torch.Tensor,
282
- hidden_states: torch.Tensor,
283
- output: torch.Tensor,
284
- ) -> None:
285
- output[:] = self.mla_attn(positions, hidden_states)
286
-
287
-
288
-class KimiDecoderLayer(nn.Module):
289
- def __init__(
290
- self,
291
- config: KimiLinearConfig,
292
- vllm_config: VllmConfig,
293
- prefix: str = "",
294
- ) -> None:
295
- super().__init__()
296
- self.hidden_size = config.hidden_size
297
-
298
- self.is_moe = config.is_moe
299
- layer_idx = int(prefix.rsplit(".", 1)[1])
300
- model_config = vllm_config.model_config
301
- cache_config = vllm_config.cache_config
302
- quant_config = vllm_config.quant_config
303
-
304
- if config.is_kda_layer(layer_idx):
305
- self.self_attn = KimiGatedDeltaNetAttention(
306
- config,
307
- vllm_config,
308
- prefix=f"{prefix}.self_attn",
309
- )
310
- else:
311
- self.self_attn = KimiMLAAttention(
312
- layer_idx=layer_idx,
313
- hidden_size=self.hidden_size,
314
- num_heads=config.num_attention_heads,
315
- quant_config=quant_config,
316
- cache_config=cache_config,
317
- model_config=model_config,
318
- prefix=f"{prefix}.self_attn",
319
- config=config,
320
- qk_nope_head_dim=config.qk_nope_head_dim,
321
- qk_rope_head_dim=config.qk_rope_head_dim,
322
- v_head_dim=config.v_head_dim,
323
- q_lora_rank=config.q_lora_rank,
324
- kv_lora_rank=config.kv_lora_rank,
325
- use_nope=config.mla_use_nope,
326
- )
327
-
328
- if (
329
- self.is_moe
330
- and config.num_experts is not None
331
- and layer_idx >= config.first_k_dense_replace
332
- and layer_idx % config.moe_layer_freq == 0
333
- ):
334
- self.block_sparse_moe = KimiMoE(
335
- config=config,
336
- quant_config=quant_config,
337
- prefix=f"{prefix}.block_sparse_moe",
338
- )
339
- self.mlp = self.block_sparse_moe
340
- else:
341
- self.mlp = KimiMLP(
342
- hidden_size=self.hidden_size,
343
- intermediate_size=config.intermediate_size,
344
- hidden_act=config.hidden_act,
345
- quant_config=quant_config,
346
- prefix=f"{prefix}.mlp",
347
- )
348
- self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
349
- self.post_attention_layernorm = RMSNorm(
350
- config.hidden_size, eps=config.rms_norm_eps
351
- )
352
-
353
- def forward(
354
- self,
355
- positions: torch.Tensor,
356
- hidden_states: torch.Tensor,
357
- residual: torch.Tensor | None,
358
- **kwargs,
359
- ) -> tuple[torch.Tensor, torch.Tensor]:
360
- # Self Attention
361
- if residual is None:
362
- residual = hidden_states
363
- hidden_states = self.input_layernorm(hidden_states)
364
- else:
365
- hidden_states, residual = self.input_layernorm(hidden_states, residual)
366
-
367
- attn_output = torch.empty_like(hidden_states)
368
- self.self_attn(
369
- hidden_states=hidden_states,
370
- positions=positions,
371
- output=attn_output,
372
- )
373
- hidden_states = attn_output
374
-
375
- # Fully Connected
376
- hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
377
- hidden_states = self.mlp(hidden_states)
378
- return hidden_states, residual
379
-
380
-
381
-@support_torch_compile
382
-class KimiLinearModel(nn.Module):
383
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
384
- super().__init__()
385
-
386
- config = vllm_config.model_config.hf_text_config
387
- self.config = config
388
-
389
- self.vocab_size = config.vocab_size
390
-
391
- if get_pp_group().is_first_rank:
392
- self.embed_tokens = VocabParallelEmbedding(
393
- config.vocab_size,
394
- config.hidden_size,
395
- prefix=f"{prefix}.embed_tokens",
396
- )
397
- else:
398
- self.embed_tokens = PPMissingLayer()
399
-
400
- def get_layer(prefix: str):
401
- return KimiDecoderLayer(
402
- config,
403
- vllm_config,
404
- prefix,
405
- )
406
-
407
- self.start_layer, self.end_layer, self.layers = make_layers(
408
- config.num_hidden_layers,
409
- get_layer,
410
- prefix=f"{prefix}.layers",
411
- )
412
-
413
- if get_pp_group().is_last_rank:
414
- self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
415
- else:
416
- self.norm = PPMissingLayer()
417
-
418
- world_size = get_tensor_model_parallel_world_size()
419
- assert config.num_attention_heads % world_size == 0, (
420
- "num_attention_heads must be divisible by world_size"
421
- )
422
-
423
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
424
- return self.embed_tokens(input_ids)
425
-
426
- def forward(
427
- self,
428
- input_ids: torch.Tensor | None,
429
- positions: torch.Tensor,
430
- intermediate_tensors: IntermediateTensors | None,
431
- inputs_embeds: torch.Tensor | None = None,
432
- **kwargs,
433
- ) -> torch.Tensor:
434
- if get_pp_group().is_first_rank:
435
- if inputs_embeds is not None:
436
- hidden_states = inputs_embeds
437
- else:
438
- hidden_states = self.embed_input_ids(input_ids)
439
- residual = None
440
- else:
441
- assert intermediate_tensors is not None
442
- hidden_states = intermediate_tensors["hidden_states"]
443
- residual = intermediate_tensors["residual"]
444
-
445
- for _, layer in enumerate(self.layers[self.start_layer : self.end_layer]):
446
- hidden_states, residual = layer(
447
- positions=positions,
448
- hidden_states=hidden_states,
449
- residual=residual,
450
- )
451
-
452
- if not get_pp_group().is_last_rank:
453
- return IntermediateTensors(
454
- {"hidden_states": hidden_states, "residual": residual}
455
- )
456
-
457
- hidden_states, _ = self.norm(hidden_states, residual)
458
- return hidden_states
459
-
460
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
461
- stacked_params_mapping = [
462
- # (param_name, shard_name, shard_id)
463
- (".gate_up_proj", ".gate_proj", 0),
464
- (".gate_up_proj", ".up_proj", 1),
465
- ]
466
- if self.config.is_moe:
467
- # Params for weights, fp8 weight scales, fp8 activation scales
468
- # (param_name, weight_name, expert_id, shard_id)
469
- expert_params_mapping = fused_moe_make_expert_params_mapping(
470
- self,
471
- ckpt_gate_proj_name="w1",
472
- ckpt_down_proj_name="w2",
473
- ckpt_up_proj_name="w3",
474
- num_experts=self.config.num_experts,
475
- )
476
- else:
477
- expert_params_mapping = []
478
- params_dict = dict(self.named_parameters())
479
- loaded_params: set[str] = set()
480
- for args in weights:
481
- name, loaded_weight = args[:2]
482
- kwargs = args[2] if len(args) > 2 else {}
483
- if "rotary_emb.inv_freq" in name:
484
- continue
485
-
486
- spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
487
- if spec_layer is not None:
488
- continue # skip spec decode layers for main model
489
- if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
490
- # Models trained using ColossalAI may include these tensors in
491
- # the checkpoint. Skip them.
492
- continue
493
- for param_name, weight_name, shard_id in stacked_params_mapping:
494
- if weight_name not in name:
495
- continue
496
- # We have mlp.experts[0].gate_proj in the checkpoint.
497
- # Since we handle the experts below in expert_params_mapping,
498
- # we need to skip here BEFORE we update the name, otherwise
499
- # name will be updated to mlp.experts[0].gate_up_proj, which
500
- # will then be updated below in expert_params_mapping
501
- # for mlp.experts[0].gate_gate_up_proj, which breaks load.
502
- if ("mlp.experts." in name) and name not in params_dict:
503
- continue
504
- name = name.replace(weight_name, param_name)
505
- # Skip loading extra bias for GPTQ models.
506
- if name.endswith(".bias") and name not in params_dict:
507
- continue
508
- if is_pp_missing_parameter(name, self):
509
- continue
510
- param = params_dict[name]
511
- weight_loader = param.weight_loader
512
- weight_loader(param, loaded_weight, shard_id)
513
- break
514
- else:
515
- for idx, (param_name, weight_name, expert_id, shard_id) in enumerate(
516
- expert_params_mapping
517
- ):
518
- if weight_name not in name:
519
- continue
520
- name = name.replace(weight_name, param_name)
521
- if is_pp_missing_parameter(name, self):
522
- continue
523
- param = params_dict[name]
524
- weight_loader = param.weight_loader
525
- weight_loader(
526
- param,
527
- loaded_weight,
528
- name,
529
- expert_id=expert_id,
530
- shard_id=shard_id,
531
- )
532
- break
533
- else:
534
- # Skip loading extra bias for GPTQ models.
535
- if (
536
- name.endswith(".bias")
537
- and name not in params_dict
538
- and not self.config.is_linear_attn
539
- ): # noqa: E501
540
- continue
541
- # Remapping the name of FP8 kv-scale.
542
- name = maybe_remap_kv_scale_name(name, params_dict)
543
- if name is None:
544
- continue
545
- if is_pp_missing_parameter(name, self):
546
- continue
547
-
548
- param = params_dict[name]
549
- weight_loader = getattr(
550
- param, "weight_loader", default_weight_loader
551
- )
552
- weight_loader(param, loaded_weight, **kwargs)
553
- loaded_params.add(name)
554
- return loaded_params
555
-
556
-
557
-class KimiLinearForCausalLM(
558
- nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid
559
-):
560
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
561
- super().__init__()
562
- self.model_config = vllm_config.model_config
563
- self.vllm_config = vllm_config
564
- self.config = self.model_config.hf_config
565
- quant_config = vllm_config.quant_config
566
- self.quant_config = quant_config
567
- self.model = KimiLinearModel(
568
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
569
- )
570
- if get_pp_group().is_last_rank:
571
- self.lm_head = ParallelLMHead(
572
- self.config.vocab_size,
573
- self.config.hidden_size,
574
- quant_config=quant_config,
575
- prefix=maybe_prefix(prefix, "lm_head"),
576
- )
577
- else:
578
- self.lm_head = PPMissingLayer()
579
- logit_scale = getattr(self.config, "logit_scale", 1.0)
580
- self.logits_processor = LogitsProcessor(
581
- self.config.vocab_size, scale=logit_scale
582
- )
583
-
584
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
585
- return self.model.embed_input_ids(input_ids)
586
-
587
- def forward(
588
- self,
589
- input_ids: torch.Tensor | None,
590
- positions: torch.Tensor,
591
- intermediate_tensors: IntermediateTensors | None = None,
592
- inputs_embeds: torch.Tensor | None = None,
593
- **kwargs,
594
- ) -> torch.Tensor | IntermediateTensors:
595
- hidden_states = self.model(
596
- input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
597
- )
598
- return hidden_states
599
-
600
- @classmethod
601
- def get_mamba_state_dtype_from_config(
602
- cls,
603
- vllm_config: "VllmConfig",
604
- ) -> tuple[torch.dtype, torch.dtype]:
605
- return MambaStateDtypeCalculator.kda_state_dtype(
606
- vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
607
- )
608
-
609
- @classmethod
610
- def get_mamba_state_shape_from_config(
611
- cls, vllm_config: "VllmConfig"
612
- ) -> tuple[tuple[int, ...], tuple[int, ...]]:
613
- parallel_config = vllm_config.parallel_config
614
- hf_config = vllm_config.model_config.hf_config
615
- tp_size = parallel_config.tensor_parallel_size
616
- num_spec = (
617
- vllm_config.speculative_config.num_speculative_tokens
618
- if vllm_config.speculative_config
619
- else 0
620
- )
621
- return MambaStateShapeCalculator.kda_state_shape(
622
- tp_size,
623
- hf_config.linear_attn_config["num_heads"],
624
- hf_config.linear_attn_config["head_dim"],
625
- conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"],
626
- num_spec=num_spec,
627
- )
628
-
629
- @classmethod
630
- def get_mamba_state_copy_func(
631
- cls,
632
- ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
633
- return MambaStateCopyFuncCalculator.kda_state_copy_func()
634
-
635
- def compute_logits(
636
- self,
637
- hidden_states: torch.Tensor,
638
- ) -> torch.Tensor | None:
639
- return self.logits_processor(self.lm_head, hidden_states)
640
-
641
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
642
- loader = AutoWeightsLoader(
643
- self,
644
- skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
645
- )
646
- return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+13
-1
@@ -137,7 +137,10 @@ _TEXT_GENERATION_MODELS = {
137
"IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"),
138
"Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"),
139
"JambaForCausalLM": ("jamba", "JambaForCausalLM"),
140
- "KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"),
140
+ "KimiLinearForCausalLM": (
141
+ "vllm.models.kimi_k3",
142
+ "KimiLinearForCausalLM",
143
+ ),
144
"Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"),
145
"Lfm2MoeForCausalLM": ("lfm2_moe", "Lfm2MoeForCausalLM"),
146
"LagunaForCausalLM": ("laguna", "LagunaForCausalLM"),
@@ -460,6 +463,10 @@ _MULTIMODAL_MODELS = {
463
),
464
"KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"),
465
"KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"),
466
+ "KimiK3ForConditionalGeneration": (
467
+ "vllm.models.kimi_k3",
468
+ "KimiK3ForConditionalGeneration",
469
+ ),
470
"MoonshotKimiaForCausalLM": ("kimi_audio", "KimiAudioForConditionalGeneration"),
471
"MossTranscribeDiarizeForConditionalGeneration": (
472
"moss_transcribe_diarize",
@@ -609,6 +616,10 @@ _SPECULATIVE_DECODING_MODELS = {
616
"DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
617
"DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"),
618
"Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
619
+ "K3DSparkModel": (
620
+ "vllm.models.kimi_k3.nvidia.dspark_mla",
621
+ "K3DSparkForCausalLM",
622
+ ),
623
"DFlashLagunaForCausalLM": ("laguna_dflash", "DFlashLagunaForCausalLM"),
624
"Gemma4DSparkModel": ("gemma4_dspark", "Gemma4DSparkForCausalLM"),
625
"PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
@@ -649,6 +660,7 @@ _SPECULATIVE_DECODING_MODELS = {
660
"Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"),
661
"Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"),
662
"HYV3MTPModel": ("hy_v3_mtp", "HYV3MTP"),
663
+ "KimiK3MTPModel": ("vllm.models.kimi_k3", "KimiK3MTP"),
664
# Temporarily disabled.
665
# # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.
666
# "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),
vllm/model_executor/warmup/deep_gemm_warmup.py
+2
-2
@@ -234,7 +234,7 @@ def _deepgemm_fp8_gemm_nt_warmup(
234
235
device = w.device
236
a1q = torch.empty((max_tokens, k), device=device, dtype=torch.float8_e4m3fn)
237
- a1q_scales = torch.empty(
237
+ a1q_scales = torch.zeros(
238
(max_tokens, k // block_m), device=device, dtype=torch.float32
239
)
240
out = torch.empty((max_tokens, n), device=device, dtype=torch.bfloat16)
@@ -336,7 +336,7 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
336
def _warmup(w: torch.Tensor, w_scale: torch.Tensor):
337
_, n, k = w.size()
338
a1q = torch.empty((MAX_M, k), device=device, dtype=torch.float8_e4m3fn)
339
- a1q_scales = torch.empty(
339
+ a1q_scales = torch.zeros(
340
(MAX_M, k // block_m), device=device, dtype=torch.float32
341
)
342
out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16)
vllm/model_executor/warmup/kernel_warmup.py
+4
@@ -28,6 +28,9 @@ from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
28
deepseek_v4_sparse_mla_attention_warmup,
29
flashinfer_sparse_mla_decode_autotune_warmup,
30
)
31
+from vllm.model_executor.warmup.kimi_k3_triton_warmup import (
32
+ kimi_k3_triton_warmup,
33
+)
34
from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
35
from vllm.model_executor.warmup.sparse_mla_triton_warmup import (
36
sparse_mla_triton_warmup,
@@ -123,6 +126,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
126
127
# Run next so input-prep kernels JIT against pristine runner state.
128
if worker.vllm_config.kernel_config.enable_jit_warmup:
129
+ kimi_k3_triton_warmup(worker)
130
fa4_cutedsl_warmup(worker)
131
sparse_mla_triton_warmup(worker)
132
vllm/model_executor/warmup/kimi_k3_triton_warmup.py
new
+182
@@ -0,0 +1,182 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Warm up Kimi-K3 Triton kernels."""
4
+
5
+from __future__ import annotations
6
+
7
+from typing import TYPE_CHECKING
8
+
9
+import torch
10
+
11
+from vllm.logger import init_logger
12
+from vllm.platforms import current_platform
13
+
14
+if TYPE_CHECKING:
15
+ from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
16
+ from vllm.v1.worker.gpu_worker import Worker
17
+
18
+logger = init_logger(__name__)
19
+
20
+
21
+def _get_kda_layer(worker: Worker) -> KimiK3DeltaAttention | None:
22
+ from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
23
+
24
+ compilation_config = getattr(
25
+ worker.model_runner,
26
+ "compilation_config",
27
+ None,
28
+ )
29
+ static_context = getattr(compilation_config, "static_forward_context", None)
30
+ if not isinstance(static_context, dict):
31
+ return None
32
+ return next(
33
+ (
34
+ layer
35
+ for layer in static_context.values()
36
+ if isinstance(layer, KimiK3DeltaAttention)
37
+ ),
38
+ None,
39
+ )
40
+
41
+
42
+def _warm_attn_res(worker: Worker) -> None:
43
+ from vllm.models.kimi_k3.nvidia.ops.attn_res import (
44
+ attn_res,
45
+ get_attn_res_triton_warmup_profiles,
46
+ )
47
+
48
+ config = worker.model_config.hf_text_config
49
+ block_size = getattr(config, "attn_res_block_size", None)
50
+ if block_size is None:
51
+ return
52
+
53
+ hidden_size = int(config.hidden_size)
54
+ max_blocks = (int(config.num_hidden_layers) + block_size - 1) // block_size
55
+ if max_blocks < 2:
56
+ return
57
+
58
+ dtype = worker.model_config.dtype
59
+ device = torch.device("cuda")
60
+ eps = float(config.rms_norm_eps)
61
+ prefix = torch.zeros((1, hidden_size), dtype=dtype, device=device)
62
+ delta = torch.zeros_like(prefix)
63
+ blocks = torch.zeros(
64
+ (1, max_blocks, hidden_size),
65
+ dtype=dtype,
66
+ device=device,
67
+ )
68
+ norm_weight = torch.zeros(hidden_size, dtype=dtype, device=device)
69
+ qk_weight = torch.zeros_like(norm_weight)
70
+ output_norm_weight = torch.zeros_like(norm_weight)
71
+
72
+ for (
73
+ num_blocks,
74
+ has_delta,
75
+ block_write_idx,
76
+ apply_output_norm,
77
+ ) in get_attn_res_triton_warmup_profiles(max_blocks):
78
+ attn_res(
79
+ prefix,
80
+ delta if has_delta else None,
81
+ blocks,
82
+ norm_weight,
83
+ qk_weight,
84
+ output_norm_weight if apply_output_norm else None,
85
+ num_blocks=num_blocks,
86
+ block_write_idx=block_write_idx,
87
+ eps=eps,
88
+ output_norm_eps=eps if apply_output_norm else 0.0,
89
+ )
90
+
91
+
92
+def _warm_recurrent_kda(
93
+ layer: KimiK3DeltaAttention,
94
+ input_dtype: torch.dtype,
95
+) -> None:
96
+ from vllm.models.kimi_k3.nvidia.ops.third_party.kda.fused_recurrent import (
97
+ fused_recurrent_kda,
98
+ get_fused_recurrent_kda_fwd_warmup_profiles,
99
+ )
100
+
101
+ num_speculative_tokens = int(layer.num_spec)
102
+ # fused_recurrent_kda_fwd_kernel is only used by speculative decode.
103
+ if num_speculative_tokens <= 0:
104
+ return
105
+
106
+ kv_cache = layer.kv_cache
107
+ if not isinstance(kv_cache, (list, tuple)) or len(kv_cache) < 2:
108
+ return
109
+ state = kv_cache[1]
110
+ if not isinstance(state, torch.Tensor) or not state.numel():
111
+ return
112
+
113
+ logger.info("Warming up Kimi-K3 speculative KDA kernels.")
114
+ h = int(layer.local_num_heads)
115
+ d = int(layer.head_dim)
116
+ tokens_per_sequence = num_speculative_tokens + 1
117
+ for num_sequences in get_fused_recurrent_kda_fwd_warmup_profiles(h):
118
+ num_tokens = num_sequences * tokens_per_sequence
119
+ packed_qkv = torch.empty(
120
+ (num_tokens, 3 * h * d),
121
+ dtype=input_dtype,
122
+ device=state.device,
123
+ )
124
+ q, k, v = (
125
+ tensor.view(1, num_tokens, h, d)
126
+ for tensor in packed_qkv.split(h * d, dim=-1)
127
+ )
128
+ fused_recurrent_kda(
129
+ q=q,
130
+ k=k,
131
+ v=v,
132
+ raw_g=torch.empty(
133
+ (1, num_tokens, h, d),
134
+ dtype=input_dtype,
135
+ device=state.device,
136
+ ),
137
+ raw_beta=torch.empty(
138
+ (1, num_tokens, h),
139
+ dtype=input_dtype,
140
+ device=state.device,
141
+ ),
142
+ A_log=layer.A_log,
143
+ dt_bias=layer.dt_bias,
144
+ lower_bound=layer.gate_lower_bound,
145
+ initial_state=state[:1],
146
+ cu_seqlens=torch.arange(
147
+ 0,
148
+ num_tokens + 1,
149
+ tokens_per_sequence,
150
+ dtype=torch.int32,
151
+ device=state.device,
152
+ ),
153
+ ssm_state_indices=torch.zeros(
154
+ (num_sequences, tokens_per_sequence),
155
+ dtype=torch.int32,
156
+ device=state.device,
157
+ ),
158
+ num_accepted_tokens=torch.ones(
159
+ num_sequences,
160
+ dtype=torch.int32,
161
+ device=state.device,
162
+ ),
163
+ out=torch.empty(
164
+ (1, num_tokens, h, d),
165
+ dtype=input_dtype,
166
+ device=state.device,
167
+ ),
168
+ )
169
+
170
+
171
+@torch.inference_mode()
172
+def kimi_k3_triton_warmup(worker: Worker) -> None:
173
+ """Warm Kimi-K3 Triton kernels reachable by this server."""
174
+ if not current_platform.is_cuda():
175
+ return
176
+
177
+ layer = _get_kda_layer(worker)
178
+ if layer is None:
179
+ return
180
+
181
+ _warm_attn_res(worker)
182
+ _warm_recurrent_kda(layer, worker.model_config.dtype)
vllm/models/deepseek_v4/attention.py
+1
-1
@@ -22,9 +22,9 @@ from vllm.model_executor.layers.linear import (
22
RowParallelLinear,
23
)
24
from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
25
+from vllm.models.common.ops import fused_q_kv_rmsnorm
26
from vllm.models.deepseek_v4.common.ops import (
27
fused_indexer_q_rope_quant,
27
- fused_q_kv_rmsnorm,
28
)
29
from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
30
vllm/models/deepseek_v4/common/ops/__init__.py
-2
@@ -11,7 +11,6 @@ from .cache_utils import (
11
from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant
12
from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant
13
from .fused_mtp_input_rmsnorm import fused_mtp_input_rmsnorm, mtp_shared_head_rmsnorm
14
-from .fused_qk_rmsnorm import fused_q_kv_rmsnorm
14
from .save_partial_states import save_partial_states
15
16
__all__ = [
@@ -23,7 +22,6 @@ __all__ = [
22
"fused_indexer_q_rope_quant",
23
"fused_inv_rope_fp8_quant",
24
"fused_mtp_input_rmsnorm",
26
- "fused_q_kv_rmsnorm",
25
"mtp_shared_head_rmsnorm",
26
"quantize_and_insert_k_cache",
27
"save_partial_states",
vllm/models/deepseek_v4/common/ops/fused_qk_rmsnorm.py
deleted
-96
@@ -1,96 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import torch
4
-
5
-from vllm.triton_utils import tl, triton
6
-
7
-
8
-@triton.jit
9
-def _fused_q_kv_rmsnorm_kernel(
10
- q_ptr,
11
- q_out_ptr,
12
- q_weight_ptr,
13
- q_in_stride,
14
- q_out_stride,
15
- kv_ptr,
16
- kv_out_ptr,
17
- kv_weight_ptr,
18
- kv_in_stride,
19
- kv_out_stride,
20
- eps,
21
- Q_SIZE: tl.constexpr,
22
- KV_SIZE: tl.constexpr,
23
- BLOCK_SIZE: tl.constexpr,
24
-):
25
- # num_tokens goes on grid-x (max 2**31 - 1); task goes on grid-y.
26
- # CUDA's grid-y/z are capped at 65535, so putting num_tokens there crashes
27
- # the launch at max-num-batched-tokens >= 65536 with "invalid argument".
28
- # int64: q_in_stride can be ~24K (128 heads × 192) and overflows int32
29
- # past num_tokens ~87K under large chunked prefill.
30
- token_idx = tl.program_id(0).to(tl.int64)
31
- pid_task = tl.program_id(1)
32
-
33
- if pid_task == 0:
34
- SIZE = Q_SIZE
35
- row_in = q_ptr + token_idx * q_in_stride
36
- weight_ptr = q_weight_ptr
37
- row_out = q_out_ptr + token_idx * q_out_stride
38
- else:
39
- SIZE = KV_SIZE
40
- row_in = kv_ptr + token_idx * kv_in_stride
41
- weight_ptr = kv_weight_ptr
42
- row_out = kv_out_ptr + token_idx * kv_out_stride
43
-
44
- # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's
45
- # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which
46
- # keep x, rrms, and w all in fp32 and perform a single cast at store.
47
- block = tl.arange(0, BLOCK_SIZE)
48
- mask = block < SIZE
49
- x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32)
50
- variance = tl.sum(x * x, axis=0) / SIZE
51
- rrms = tl.rsqrt(variance + eps)
52
- w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32)
53
- y = x * rrms * w
54
- tl.store(row_out + block, y.to(row_out.dtype.element_ty), mask=mask)
55
-
56
-
57
-def fused_q_kv_rmsnorm(
58
- qr: torch.Tensor,
59
- kv: torch.Tensor,
60
- q_weight: torch.Tensor,
61
- kv_weight: torch.Tensor,
62
- eps: float,
63
-) -> tuple[torch.Tensor, torch.Tensor]:
64
- assert qr.ndim == 2 and kv.ndim == 2
65
- assert qr.shape[0] == kv.shape[0], (
66
- f"token dim mismatch: qr={qr.shape}, kv={kv.shape}"
67
- )
68
- assert qr.stride(-1) == 1 and kv.stride(-1) == 1
69
- assert q_weight.is_contiguous() and kv_weight.is_contiguous()
70
-
71
- q_size = qr.shape[1]
72
- kv_size = kv.shape[1]
73
- num_tokens = qr.shape[0]
74
- qr_out = torch.empty_like(qr)
75
- kv_out = torch.empty_like(kv)
76
- if num_tokens == 0:
77
- return qr_out, kv_out
78
-
79
- block_size = triton.next_power_of_2(max(q_size, kv_size))
80
- _fused_q_kv_rmsnorm_kernel[(num_tokens, 2)](
81
- qr,
82
- qr_out,
83
- q_weight,
84
- qr.stride(0),
85
- qr_out.stride(0),
86
- kv,
87
- kv_out,
88
- kv_weight,
89
- kv.stride(0),
90
- kv_out.stride(0),
91
- eps,
92
- Q_SIZE=q_size,
93
- KV_SIZE=kv_size,
94
- BLOCK_SIZE=block_size,
95
- )
96
- return qr_out, kv_out
vllm/models/kimi_k3/nvidia/dspark_mla.py
+1
-21
@@ -25,26 +25,6 @@ from vllm.models.kimi_k3.nvidia.model import KimiMLP
25
from vllm.utils.torch_utils import is_quantized_kv_cache
26
27
28
-class ReplicatedDSparkMarkovHead(DSparkMarkovHead):
29
- """DSpark Markov head with full weights on every TP rank."""
30
-
31
- def __init__(
32
- self,
33
- vocab_size: int,
34
- draft_vocab_size: int,
35
- markov_rank: int,
36
- prefix: str,
37
- ) -> None:
38
- # TODO: Remove this mypy workaround once the K3 PR is fully merged.
39
- super().__init__( # type: ignore[call-arg]
40
- vocab_size,
41
- draft_vocab_size,
42
- markov_rank,
43
- prefix,
44
- replicated=True,
45
- )
46
-
47
-
28
class K3DSparkDecoderLayer(nn.Module):
29
def __init__(
30
self,
@@ -148,7 +128,7 @@ class K3DSparkModel(nn.Module):
128
]
129
)
130
self.final_norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
151
- self.markov_head = ReplicatedDSparkMarkovHead(
131
+ self.markov_head = DSparkMarkovHead(
132
self.config.vocab_size,
133
self.config.draft_vocab_size,
134
self.config.markov_rank,
vllm/models/kimi_k3/nvidia/model.py
+16
-3
@@ -121,6 +121,13 @@ from ..common.mm_preprocess import (
121
122
logger = init_logger(__name__)
123
124
+# Token-count cutoff for overlapping the MoE router gate with the routed-expert
125
+# down projection on a separate CUDA stream (latent MoE). At or below this many
126
+# tokens the launch-bound decode path benefits from multi-stream overlap; above
127
+# it the GEMMs saturate the device and the cross-stream sync is pure overhead,
128
+# so it falls back to sequential.
129
+_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD = 256
130
+
131
132
class KimiMLP(nn.Module):
133
def __init__(
@@ -513,7 +520,7 @@ class KimiMoE(nn.Module):
520
)
521
# Auxiliary CUDA stream to overlap the router gate with the routed
522
# down projection on decode-sized batches (gated by
516
- # VLLM_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD).
523
+ # _ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD).
524
self._down_proj_stream: torch.cuda.Stream | None = aux_stream()
525
self._down_proj_events = (torch.cuda.Event(), torch.cuda.Event())
526
else:
@@ -546,7 +553,13 @@ class KimiMoE(nn.Module):
553
activation_linear_beta=activation_situ_linear_beta,
554
)
555
else:
549
- enable_tail_fusion = envs.VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION
556
+ # The tail-fusion kernels are tcgen05-based, so they require an
557
+ # SM100 NVIDIA device; the runner falls back to the default latent
558
+ # MoE path everywhere else.
559
+ enable_tail_fusion = (
560
+ current_platform.is_cuda()
561
+ and current_platform.is_device_capability_family(100)
562
+ )
563
self.experts = FusedMoE(
564
shared_experts=self.shared_experts,
565
num_experts=num_experts,
@@ -645,7 +658,7 @@ class KimiMoE(nn.Module):
658
self._down_proj_events[0],
659
self._down_proj_events[1],
660
self._down_proj_stream
648
- if num_tokens <= envs.VLLM_ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD
661
+ if num_tokens <= _ROUTED_DOWN_PROJ_STREAM_TOKEN_THRESHOLD
662
else None,
663
)
664
)
vllm/platforms/interface.py
+6
@@ -880,6 +880,12 @@ class Platform:
880
),
881
cache_config.block_size,
882
)
883
+ if model_config.use_mla:
884
+ # TRTLLM/FlashInfer MLA decode kernels require the physical
885
+ # number of kernel blocks to be aligned to 128 / kernel_block_size.
886
+ # For hybrid MLA/Mamba models, make the manager block size a
887
+ # multiple of 128 so split kernel blocks keep that invariant.
888
+ kernel_block_alignment_size = max(kernel_block_alignment_size, 128)
889
890
if cache_config.mamba_cache_mode == "all":
891
# With prefix caching, align to mamba chunk size for kernel perf
vllm/third_party/flash_linear_attention/ops/fused_recurrent.py
+1
-1
@@ -12,7 +12,7 @@ import torch
12
13
from vllm.triton_utils import tl, triton
14
15
-from .op import exp
15
+from .op import exp, log
16
17
18
@triton.heuristics(
vllm/transformers_utils/config.py
+2
@@ -87,6 +87,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
87
deepseek_vl_v2="DeepseekVLV2Config",
88
deepseek_v32="DeepseekV3Config",
89
deepseek_v4="DeepseekV4Config",
90
+ k3_dspark="K3DSparkConfig",
91
flex_olmo="FlexOlmoConfig",
92
fireredlid="FireRedLIDConfig",
93
funaudiochat="FunAudioChatConfig",
@@ -100,6 +101,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
101
kimi_linear="KimiLinearConfig",
102
kimi_vl="KimiVLConfig",
103
kimi_k25="KimiK25Config",
104
+ kimi_k3="KimiK3Config",
105
RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct)
106
RefinedWebModel="RWConfig", # For tiiuae/falcon-7b(-instruct)
107
mlp_speculator="MLPSpeculatorConfig",
vllm/transformers_utils/configs/__init__.py
+6
@@ -35,6 +35,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
35
"DiffusionGemmaTextConfig": "vllm.transformers_utils.configs.diffusion_gemma",
36
"DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2",
37
"DeepseekV4Config": "vllm.transformers_utils.configs.deepseek_v4",
38
+ "K3DSparkConfig": "vllm.transformers_utils.configs.k3_dspark",
39
"DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr",
40
"EAGLEConfig": "vllm.transformers_utils.configs.eagle",
41
"FireRedLIDConfig": "vllm.transformers_utils.configs.fireredlid",
@@ -72,6 +73,8 @@ _CLASS_TO_MODULE: dict[str, str] = {
73
"KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear",
74
"KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl",
75
"KimiK25Config": "vllm.transformers_utils.configs.kimi_k25",
76
+ "KimiK3Config": "vllm.transformers_utils.configs.kimi_k3",
77
+ "KimiK3VisionConfig": "vllm.transformers_utils.configs.kimi_k3",
78
"NemotronConfig": "vllm.transformers_utils.configs.nemotron",
79
"NemotronHConfig": "vllm.transformers_utils.configs.nemotron_h",
80
"OlmoHybridConfig": "vllm.transformers_utils.configs.olmo_hybrid",
@@ -124,6 +127,7 @@ __all__ = [
127
"DeepseekVLV2Config",
128
"DeepseekV3Config",
129
"DeepseekV4Config",
130
+ "K3DSparkConfig",
131
"DotsOCRConfig",
132
"EAGLEConfig",
133
"FlexOlmoConfig",
@@ -156,6 +160,8 @@ __all__ = [
160
"KimiLinearConfig",
161
"KimiVLConfig",
162
"KimiK25Config",
163
+ "KimiK3Config",
164
+ "KimiK3VisionConfig",
165
"NemotronConfig",
166
"NemotronHConfig",
167
"OlmoHybridConfig",
vllm/transformers_utils/configs/k3_dspark.py
new
+77
@@ -0,0 +1,77 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from transformers import DeepseekV2Config
5
+
6
+
7
+class K3DSparkConfig(DeepseekV2Config):
8
+ """Configuration for a dense MLA DSpark draft model."""
9
+
10
+ model_type = "k3_dspark"
11
+ has_no_defaults_at_init = True
12
+
13
+ def __init__(
14
+ self,
15
+ mla_use_nope: bool = False,
16
+ mla_use_output_gate: bool = False,
17
+ mla_use_qk_norm: bool = False,
18
+ rope_theta: float = 50000.0,
19
+ **kwargs,
20
+ ) -> None:
21
+ # DeepseekV2Config defaults to a MoE topology. Zero these fields so
22
+ # generic vLLM config logic also recognizes this draft as dense.
23
+ kwargs.setdefault("n_routed_experts", 0)
24
+ kwargs.setdefault("n_shared_experts", 0)
25
+ kwargs.setdefault("num_experts_per_tok", 0)
26
+
27
+ rope_parameters = kwargs.get("rope_parameters")
28
+ if rope_parameters is None:
29
+ kwargs["rope_parameters"] = {
30
+ "rope_type": "default",
31
+ "rope_theta": rope_theta,
32
+ }
33
+ else:
34
+ rope_parameters = dict(rope_parameters)
35
+ rope_parameters.setdefault("rope_type", "default")
36
+ rope_parameters.setdefault("rope_theta", rope_theta)
37
+ kwargs["rope_parameters"] = rope_parameters
38
+
39
+ super().__init__(**kwargs)
40
+ self.mla_use_nope = mla_use_nope
41
+ self.mla_use_output_gate = mla_use_output_gate
42
+ self.mla_use_qk_norm = mla_use_qk_norm
43
+
44
+ unsupported = [
45
+ name
46
+ for name in (
47
+ "mla_use_nope",
48
+ "mla_use_output_gate",
49
+ "mla_use_qk_norm",
50
+ "dspark_bonus_anchor",
51
+ )
52
+ if getattr(self, name, False)
53
+ ]
54
+ if self.q_lora_rank is None:
55
+ unsupported.append("q_lora_rank=None")
56
+ if unsupported:
57
+ raise ValueError(
58
+ "MLA DSpark does not support " + ", ".join(unsupported) + "."
59
+ )
60
+
61
+ self.draft_vocab_size = (
62
+ getattr(self, "draft_vocab_size", None) or self.vocab_size
63
+ )
64
+ if self.draft_vocab_size != self.vocab_size:
65
+ raise ValueError(
66
+ "MLA DSpark requires draft_vocab_size to equal vocab_size when "
67
+ "sharing the target embedding and LM head."
68
+ )
69
+
70
+ target_layer_ids = getattr(self, "target_layer_ids", None)
71
+ if not target_layer_ids or getattr(self, "num_target_layers", None) != len(
72
+ target_layer_ids
73
+ ):
74
+ raise ValueError(
75
+ "MLA DSpark requires non-empty target_layer_ids and a matching "
76
+ "num_target_layers."
77
+ )
vllm/transformers_utils/configs/kimi_linear.py
+17
-4
@@ -3,10 +3,6 @@
3
4
from transformers.configuration_utils import PretrainedConfig
5
6
-from vllm.logger import init_logger
7
-
8
-logger = init_logger(__name__)
9
-
6
7
class KimiLinearConfig(PretrainedConfig):
8
model_type = "kimi_linear"
@@ -49,8 +45,16 @@ class KimiLinearConfig(PretrainedConfig):
45
qk_rope_head_dim: int | None = None,
46
v_head_dim: int | None = None,
47
mla_use_nope: bool | None = False,
48
+ mla_use_output_gate: bool | None = False,
49
num_nextn_predict_layers: int = 0,
50
linear_attn_config: dict | None = None,
51
+ attn_res_block_size: int | None = None,
52
+ latent_moe_use_norm: bool = False,
53
+ activation_situ_beta: float | None = None,
54
+ activation_situ_linear_beta: float | None = None,
55
+ max_position_embeddings: int = 4096,
56
+ routed_expert_hidden_size: int | None = None,
57
+ topk_method: str = "noaux_tc",
58
**kwargs,
59
):
60
self.model_type = model_type
@@ -86,6 +90,7 @@ class KimiLinearConfig(PretrainedConfig):
90
self.qk_rope_head_dim = qk_rope_head_dim
91
self.v_head_dim = v_head_dim
92
self.mla_use_nope = mla_use_nope
93
+ self.mla_use_output_gate = mla_use_output_gate
94
# moe config
95
self.num_experts = num_experts
96
self.num_experts_per_token = num_experts_per_token
@@ -102,6 +107,14 @@ class KimiLinearConfig(PretrainedConfig):
107
self.topk_group = topk_group
108
self.num_nextn_predict_layers = num_nextn_predict_layers
109
110
+ self.attn_res_block_size = attn_res_block_size
111
+ self.latent_moe_use_norm = latent_moe_use_norm
112
+ self.activation_situ_beta = activation_situ_beta
113
+ self.activation_situ_linear_beta = activation_situ_linear_beta
114
+ self.max_position_embeddings = max_position_embeddings
115
+ self.routed_expert_hidden_size = routed_expert_hidden_size
116
+ self.topk_method = topk_method
117
+
118
if linear_attn_config is not None:
119
assert linear_attn_config["kda_layers"] is not None
120
assert linear_attn_config["full_attn_layers"] is not None
vllm/transformers_utils/model_arch_config_convertor.py
+1
@@ -267,6 +267,7 @@ class ModelArchConfigConvertorBase:
267
"deepseek_v32",
268
"deepseek_v4",
269
"deepseek_mtp",
270
+ "k3_dspark",
271
"glm_moe_dsa",
272
"glm4_moe_lite",
273
"glm4_moe_lite_mtp",
vllm/transformers_utils/processors/__init__.py
+2
@@ -27,6 +27,7 @@ __all__ = [
27
"IsaacProcessor",
28
"KimiAudioProcessor",
29
"KimiK25Processor",
30
+ "KimiK3Processor",
31
"MiMoOmniProcessor",
32
"MiniCPMOProcessor",
33
"MiniCPMVProcessor",
@@ -65,6 +66,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
66
"IsaacProcessor": "vllm.transformers_utils.processors.isaac",
67
"KimiAudioProcessor": "vllm.transformers_utils.processors.kimi_audio",
68
"KimiK25Processor": "vllm.transformers_utils.processors.kimi_k25",
69
+ "KimiK3Processor": "vllm.transformers_utils.processors.kimi_k3",
70
"MiMoOmniProcessor": "vllm.transformers_utils.processors.mimo_v2_omni",
71
"MiniCPMOProcessor": "vllm.transformers_utils.processors.minicpmo",
72
"MiniCPMVProcessor": "vllm.transformers_utils.processors.minicpmv",
vllm/utils/multi_stream_utils.py
+11
-3
@@ -27,8 +27,9 @@ def maybe_execute_in_parallel(
27
"""Run two functions potentially in parallel on separate CUDA streams.
28
29
When aux_stream is provided, fn0 runs on the current (default) stream and
30
- fn1 runs on aux_stream, synchronized via CUDA events. When aux_stream is
31
- None, both functions execute sequentially on the current stream.
30
+ fn1 runs on aux_stream, synchronized via CUDA events. When aux_stream is
31
+ None or a breakable CUDA graph capture is active, both functions execute
32
+ sequentially on the current stream.
33
34
This design follows TensorRT-LLM's maybe_execute_in_parallel pattern
35
(tensorrt_llm/_torch/modules/multi_stream_utils.py).
@@ -39,11 +40,18 @@ def maybe_execute_in_parallel(
40
event0: CUDA event recorded before fn0 so aux_stream can wait.
41
event1: CUDA event recorded after fn1 so default stream can wait.
42
aux_stream: The second CUDA stream for fn1.
42
- Multi-stream is disabled when aux_stream is None.
43
+ Multi-stream is disabled when aux_stream is None or a breakable
44
+ CUDA graph capture is active.
45
46
Returns:
47
Tuple of (fn0_result, fn1_result).
48
"""
49
+ if aux_stream is not None:
50
+ from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
51
+
52
+ if BreakableCUDAGraphCapture.is_active():
53
+ aux_stream = None
54
+
55
if aux_stream is not None:
56
event0.record()
57
result0 = fn0()
vllm/v1/attention/backends/mla/flashinfer_mla.py
+119
-3
@@ -5,7 +5,12 @@ from typing import ClassVar
5
6
import torch
7
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
8
+from flashinfer.utils import (
9
+ get_device_sm_count,
10
+ get_trtllm_gen_multi_ctas_kv_counter_bytes,
11
+)
12
13
+from vllm.config import get_current_vllm_config
14
from vllm.config.cache import CacheDType
15
from vllm.logger import init_logger
16
from vllm.model_executor.layers.attention.mla_attention import (
@@ -27,6 +32,37 @@ from vllm.v1.attention.backends.utils import KVCacheLayoutType
32
33
logger = init_logger(__name__)
34
35
+
36
+def _trtllm_gen_mla_decode_supports_num_heads(num_heads: int) -> bool:
37
+ """True if trtllm-gen's MLA decode kernel supports this query head count.
38
+
39
+ The kernel groups Q heads into CTAs of ``min(num_heads, tileSizeQ)`` and
40
+ requires ``num_heads`` divisible by that, else raises "The
41
+ numHeadsQ/numHeadsKv is not supported" (flashinfer fmhaKernels.cuh).
42
+ ``tileSizeQ`` is 8/16 for ``num_heads <= 8``/``<= 32`` (SwapsMmaAb) else 64;
43
+ treat ``> 32`` as tile 64 (safe bound). E.g. 96/24 -> False, 48/64/128 -> True.
44
+ """
45
+ if num_heads <= 8:
46
+ tile = 8
47
+ elif num_heads <= 32:
48
+ tile = 16
49
+ else:
50
+ tile = 64
51
+ return num_heads % min(num_heads, tile) == 0
52
+
53
+
54
+def _select_mla_decode_backend(num_heads: int) -> str | None:
55
+ """cute-dsl for head counts trtllm-gen cannot tile, else None (=> auto)."""
56
+ if not _trtllm_gen_mla_decode_supports_num_heads(num_heads):
57
+ logger.warning_once(
58
+ "trtllm-gen MLA decode does not support num_heads=%d "
59
+ "(query/kv head ratio); falling back to the cute-dsl backend.",
60
+ num_heads,
61
+ )
62
+ return "cute-dsl"
63
+ return None
64
+
65
+
66
FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
67
FLASHINFER_MLA_LSE_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024
68
@@ -48,9 +84,39 @@ def _get_workspace_buffer(return_lse: bool) -> torch.Tensor:
84
return _fi_workspace
85
86
87
+_fi_multi_ctas_kv_counter: torch.Tensor | None = None
88
+
89
+
90
+def _get_multi_ctas_kv_counter_buffer(
91
+ min_bytes: int, device: torch.device
92
+) -> torch.Tensor:
93
+ """Persistent, zero-initialized trtllm-gen multi-CTA-KV counter buffer.
94
+
95
+ trtllm-gen's multi-CTA-KV MLA decode kernel resets these semaphores to zero
96
+ at the end of every launch, so the buffer only needs zeroing once. The
97
+ public ``trtllm_batch_decode_with_kv_cache_mla`` entry point builds a fresh
98
+ runner per call, so without a caller-owned buffer it re-allocates and
99
+ re-zeros this counter on every decode step (a tiny ``FillFunctor<uint8>``
100
+ launch right before the FMHA). Owning it here and passing it in removes that
101
+ per-step launch. ``min_bytes`` is sized to the worst-case batch so the
102
+ buffer is allocated once and never reallocated after CUDA-graph capture.
103
+ """
104
+ global _fi_multi_ctas_kv_counter
105
+ if (
106
+ _fi_multi_ctas_kv_counter is None
107
+ or _fi_multi_ctas_kv_counter.numel() < min_bytes
108
+ ):
109
+ _fi_multi_ctas_kv_counter = torch.zeros(
110
+ min_bytes, dtype=torch.uint8, device=device
111
+ )
112
+ return _fi_multi_ctas_kv_counter
113
+
114
+
115
class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
116
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
117
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
118
+ # Non-causal DSpark blocks are flattened to single-token rows in forward_mqa.
119
+ supports_non_causal_multi_token_decode: ClassVar[bool] = True
120
121
122
class FlashInferMLABackend(MLACommonBackend):
@@ -91,6 +157,10 @@ class FlashInferMLABackend(MLACommonBackend):
157
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
158
return capability.major == 10
159
160
+ @classmethod
161
+ def supports_non_causal(cls) -> bool:
162
+ return True
163
+
164
@classmethod
165
def supports_combination(
166
cls,
@@ -179,6 +249,16 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
249
250
self.bmm1_scale: float | None = None
251
self.bmm2_scale: float | None = None
252
+ # Worst-case decode batch for the persistent trtllm-gen multi-CTA-KV
253
+ # counter buffer (see _get_multi_ctas_kv_counter_buffer). Captured here
254
+ # (config is in scope during construction) so the byte size can be
255
+ # resolved once on the first decode and never grows after CUDA-graph
256
+ # capture. The impl has no _vllm_config, hence get_current_vllm_config().
257
+ _sched = get_current_vllm_config().scheduler_config
258
+ self._mla_counter_max_batch: int = (
259
+ _sched.max_num_batched_tokens or _sched.max_num_seqs
260
+ )
261
+ self._mla_counter_bytes: int | None = None
262
263
def forward_mqa(
264
self,
@@ -194,8 +274,20 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
274
q_nope, q_pe = q
275
q = torch.cat([q_nope, q_pe], dim=-1)
276
277
+ block_table = attn_metadata.decode.block_table
278
+ seq_lens = attn_metadata.decode.seq_lens
279
+
280
+ if not attn_metadata.causal:
281
+ # Non-causal DSpark block: flatten to single-token decode rows with
282
+ # per-row context seq_lens (trtllm-gen has no causal flag and would
283
+ # otherwise mask the block causally).
284
+ query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes
285
+ q = q.unsqueeze(1)
286
+ if query_len > 1:
287
+ block_table = block_table.repeat_interleave(query_len, dim=0)
288
+ seq_lens = seq_lens.repeat_interleave(query_len)
289
# trtllm API requires extra dimension q_len_per_request for MTP
198
- if attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:
290
+ elif attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:
291
logger.warning_once(
292
"""FlashInferMLAImpl got a query of uneven length.
293
This usually indicates an issue in batch reordering
@@ -217,6 +309,29 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
309
310
return_lse = self.need_to_return_lse_for_decode
311
workspace_buffer = _get_workspace_buffer(return_lse)
312
+ # trtllm-gen rejects MLA head counts it can't tile (e.g. 96);
313
+ # fall back to cute-dsl for those.
314
+ decode_backend = _select_mla_decode_backend(self.num_heads)
315
+ extra_kwargs = {}
316
+ if decode_backend:
317
+ extra_kwargs["backend"] = decode_backend
318
+ elif kv_c_and_k_pe_cache.shape[-2] in (32, 64):
319
+ # The auto path can dispatch to trtllm-gen, whose multi-CTA-KV decode
320
+ # kernel self-resets its semaphore counter after each launch (so it
321
+ # only needs zeroing once). Pass a persistent counter buffer to skip
322
+ # the per-step re-allocate + re-zero the public entry point would
323
+ # otherwise do. Guarded to configs where a trtllm-gen runner is
324
+ # eligible (page/block size in {32, 64}); the arg is rejected when
325
+ # only a cute-dsl runner can run.
326
+ if self._mla_counter_bytes is None:
327
+ self._mla_counter_bytes = get_trtllm_gen_multi_ctas_kv_counter_bytes(
328
+ self._mla_counter_max_batch,
329
+ self.num_heads,
330
+ get_device_sm_count(q.device),
331
+ )
332
+ extra_kwargs["multi_ctas_kv_counter_buffer"] = (
333
+ _get_multi_ctas_kv_counter_buffer(self._mla_counter_bytes, q.device)
334
+ )
335
kernel_out = trtllm_batch_decode_with_kv_cache_mla(
336
query=q,
337
kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
@@ -224,12 +339,13 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
339
qk_nope_head_dim=self.qk_nope_head_dim,
340
kv_lora_rank=self.kv_lora_rank,
341
qk_rope_head_dim=self.qk_rope_head_dim,
227
- block_tables=attn_metadata.decode.block_table,
228
- seq_lens=attn_metadata.decode.seq_lens,
342
+ block_tables=block_table,
343
+ seq_lens=seq_lens,
344
max_seq_len=attn_metadata.max_seq_len,
345
bmm1_scale=self.bmm1_scale,
346
bmm2_scale=self.bmm2_scale,
347
return_lse=return_lse,
348
+ **extra_kwargs,
349
)
350
if return_lse:
351
o, lse = kernel_out
vllm/v1/attention/backends/mla/prefill/base.py
+11
@@ -73,6 +73,17 @@ class MLAPrefillBackend(ABC):
73
pass. Overridden by backends that support it."""
74
return False
75
76
+ def supports_out(self) -> bool:
77
+ """Whether `run_prefill_new_tokens` honors a caller-provided `out`
78
+ tensor of shape `[num_tokens, num_heads, v_head_dim]`, writing the
79
+ final result into it in place.
80
+
81
+ When True, callers may pass `out` and skip the post-hoc
82
+ slice/flatten/copy. False for backends that ignore `out` or emit a
83
+ padded (`qk_head_dim`) output. Overridden by backends that support it.
84
+ """
85
+ return False
86
+
87
@classmethod
88
def validate_configuration(
89
cls,
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+5
@@ -419,6 +419,11 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
419
return attn_out, lse
420
return attn_out
421
422
+ def supports_out(self) -> bool:
423
+ # A padded V produces a qk_head_dim output that cannot be written into
424
+ # a v_head_dim `out`; only the unpadded path honors `out`.
425
+ return not self.requires_v_padding
426
+
427
def run_prefill_new_tokens(
428
self,
429
q: torch.Tensor,
vllm/v1/attention/backends/mla/prefill/flashinfer.py
+15
-4
@@ -108,15 +108,21 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
108
return self._global_hyperparameters
109
110
from vllm.model_executor.layers.attention.mla_attention import (
111
- MLAAttention,
111
MLACommonImpl,
112
)
113
+ from vllm.model_executor.layers.attention_layer_base import (
114
+ AttentionLayerBase,
115
+ )
116
117
+ # Match any layer with an MLA impl, not just the MLAAttention wrapper:
118
+ # fused MLA modules (Kimi-K3's MultiHeadLatentAttention) register a
119
+ # different layer type. Keying on impl also excludes linear/KDA layers.
120
forward_context = self.vllm_config.compilation_config.static_forward_context
121
layer_names = [
122
name
123
for name, layer in forward_context.items()
119
- if isinstance(layer, MLAAttention)
124
+ if isinstance(layer, AttentionLayerBase)
125
+ and isinstance(getattr(layer, "impl", None), MLACommonImpl)
126
]
127
128
self._global_hyperparameters = infer_global_hyperparameters(
@@ -191,6 +197,10 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
197
o_data_type=prefill_metadata.output_dtype,
198
)
199
200
+ def supports_out(self) -> bool:
201
+ # Planned with head_dim_vo == v_head_dim, so the output is unpadded.
202
+ return True
203
+
204
def run_prefill_new_tokens(
205
self,
206
q: torch.Tensor,
@@ -206,12 +216,13 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
216
q=q,
217
k=k,
218
v=v,
219
+ out=out,
220
return_lse=return_softmax_lse,
221
)
222
223
if isinstance(ret, tuple):
224
# Convert from (q_len, num_heads) to (num_heads, q_len)
214
- return ret[0], ret[1].transpose(0, 1).contiguous()
225
+ return ret[0], ret[1].transpose(0, 1)
226
return ret
227
228
def run_prefill_context_chunk(
@@ -229,4 +240,4 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
240
)
241
242
# Convert from (q_len, num_heads) to (num_heads, q_len)
232
- return attn_out, lse.transpose(0, 1).contiguous()
243
+ return attn_out, lse.transpose(0, 1)
vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py
+12
-7
@@ -118,6 +118,10 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
118
prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1]
119
)
120
121
+ def supports_out(self) -> bool:
122
+ # Output head dim is v_head_dim (unpadded); `out` supported since 0.1.8.
123
+ return True
124
+
125
def run_prefill_new_tokens(
126
self,
127
q: torch.Tensor,
@@ -129,11 +133,11 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
133
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
134
from tokenspeed_mla import tokenspeed_mla_prefill
135
132
- # `v` arrives as the second half of `kv_nope.split(...)` in
133
- # mla_attention.forward_mha — a non-contiguous view of `kv_nope` along
134
- # dim=-1. The kernel does `v.reshape(1, total_kv, h_k, 1, d_v)` which
135
- # would silently copy on a non-contiguous tensor; force contiguity here
136
- # so the copy (if any) happens once outside the kernel call.
136
+ # `v` arrives as the second half of `kv_nope.split(...)` — a
137
+ # non-contiguous view of `kv_nope` along dim=-1. The kernel wraps inputs
138
+ # via `from_dlpack`, which preserves strides, but the FMHA kernel reads
139
+ # `v` assuming contiguous storage and silently produces wrong output on
140
+ # a strided view (verified on tokenspeed-mla 0.1.8). Force contiguity.
141
v = v.contiguous()
142
143
ret = tokenspeed_mla_prefill(
@@ -148,11 +152,12 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
152
is_causal=True,
153
return_lse=return_softmax_lse,
154
enable_pdl=False,
155
+ out=out,
156
)
157
158
if isinstance(ret, tuple):
159
# Convert from (q_len, num_heads) to (num_heads, q_len)
155
- return ret[0], ret[1].transpose(0, 1).contiguous()
160
+ return ret[0], ret[1].transpose(0, 1)
161
return ret
162
163
def run_prefill_context_chunk(
@@ -188,4 +193,4 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
193
)
194
195
# Convert from (q_len, num_heads) to (num_heads, q_len)
191
- return attn_out, lse.transpose(0, 1).contiguous()
196
+ return attn_out, lse.transpose(0, 1)
vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py
+14
-9
@@ -91,6 +91,10 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
91
prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1]
92
)
93
94
+ def supports_out(self) -> bool:
95
+ # Output head dim is v.shape[-1] == v_head_dim, so `out` is unpadded.
96
+ return True
97
+
98
def run_prefill_new_tokens(
99
self,
100
q: torch.Tensor,
@@ -102,13 +106,14 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
106
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
107
from flashinfer.prefill import trtllm_ragged_attention_deepseek
108
105
- out = torch.empty(
106
- q.shape[0],
107
- q.shape[1],
108
- v.shape[2],
109
- device=q.device,
110
- dtype=self._prefill_metadata.output_dtype,
111
- )
109
+ if out is None:
110
+ out = torch.empty(
111
+ q.shape[0],
112
+ q.shape[1],
113
+ v.shape[2],
114
+ device=q.device,
115
+ dtype=self._prefill_metadata.output_dtype,
116
+ )
117
118
ret = trtllm_ragged_attention_deepseek(
119
query=q,
@@ -133,7 +138,7 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
138
139
if isinstance(ret, tuple):
140
# Convert from (q_len, num_heads) to (num_heads, q_len)
136
- return ret[0], ret[1].transpose(0, 1).contiguous()
141
+ return ret[0], ret[1].transpose(0, 1)
142
return ret
143
144
def run_prefill_context_chunk(
@@ -182,4 +187,4 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
187
)
188
189
# Convert from (q_len, num_heads) to (num_heads, q_len)
185
- return attn_out, lse.transpose(0, 1).contiguous()
190
+ return attn_out, lse.transpose(0, 1)
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+157
-15
@@ -31,6 +31,32 @@ from vllm.v1.kv_cache_interface import AttentionSpec
31
logger = init_logger(__name__)
32
33
34
+@functools.lru_cache(maxsize=1)
35
+def _get_mla_gluon():
36
+ """Load the small-head Gluon MLA entry point."""
37
+ unified_module = "aiter.ops.triton.gluon.mla_gluon"
38
+ try:
39
+ from aiter.ops.triton.gluon.mla_gluon import mla_gluon
40
+
41
+ return mla_gluon
42
+ except ModuleNotFoundError as unified_import_error:
43
+ if not unified_module.startswith(unified_import_error.name or ""):
44
+ raise
45
+ legacy_module = "aiter.ops.triton.gluon.mla_decode_gluon"
46
+ try:
47
+ from aiter.ops.triton.gluon.mla_decode_gluon import mla_decode_gluon
48
+
49
+ return mla_decode_gluon
50
+ except ModuleNotFoundError as legacy_import_error:
51
+ if not legacy_module.startswith(legacy_import_error.name or ""):
52
+ raise
53
+ raise RuntimeError(
54
+ "ROCM_AITER_MLA requires an AITER build with the small-head "
55
+ "Gluon MLA kernel (mla_gluon or mla_decode_gluon) when decode "
56
+ "heads are fewer than 16."
57
+ ) from unified_import_error
58
+
59
+
60
@functools.lru_cache(maxsize=1)
61
def _fp8_mla_prefill_supported() -> bool:
62
"""Auto-detect FP8 MLA prefill via mla_prefill_ps_asm_fwd + mla_reduce_v1.
@@ -102,6 +128,10 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
128
attn_out_dtype: torch.dtype = torch.bfloat16
129
# The max query output length: int
130
max_qo_len: int | None = None
131
+ # Minimum KV length used by Gluon to choose a safe split count.
132
+ min_kv_seq_len: int = 1
133
+ # Small-head decode uses Gluon (avoids padding to 16).
134
+ use_gluon_decode: bool = False
135
# Whether persistent MLA metadata was computed (only for qseqlen=1)
136
has_persistent_metadata: bool = False
137
@@ -241,7 +271,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
271
device=device,
272
)
273
244
- self._fp8_prefill_enabled = _fp8_mla_prefill_supported()
274
+ # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads.
275
+ self._fp8_prefill_enabled = (
276
+ _fp8_mla_prefill_supported() and self.num_heads % 16 == 0
277
+ )
278
if self._fp8_prefill_enabled:
279
max_prefill_qlen = min(
280
vllm_config.model_config.max_model_len,
@@ -469,6 +502,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
502
)
503
qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
504
max_qo_len = qo_len.max().item()
505
+ use_gluon_decode = AiterMLAHelper.use_gluon_decode(
506
+ self.num_heads, int(max_qo_len)
507
+ )
508
509
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
510
self.paged_kv_indices.fill_(-1)
@@ -518,7 +554,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
554
# so forward_mqa can skip passing it (falling back to the kernel
555
# computing its own metadata internally, like v0.18.0).
556
has_persistent_metadata = False
521
- if max_qo_len == 1:
557
+ if max_qo_len == 1 and not use_gluon_decode:
558
from aiter import get_mla_metadata_v1
559
560
get_mla_metadata_v1(
@@ -553,6 +589,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
589
qo_indptr=qo_indptr,
590
dcp_tot_seq_lens=dcp_tot_seq_lens_device,
591
max_qo_len=max_qo_len,
592
+ use_gluon_decode=use_gluon_decode,
593
attn_out_dtype=self.decode_attn_out_dtype,
594
has_persistent_metadata=has_persistent_metadata,
595
)
@@ -646,17 +683,16 @@ class AiterMLAHelper:
683
@staticmethod
684
def check_num_heads_validity(num_heads: int):
685
assert AiterMLAHelper.is_valid_num_heads(num_heads), (
649
- f"Aiter MLA requires that num_heads be multiples or divisors of 16, "
650
- f"but provided {num_heads} number of heads.\n"
686
+ "ROCM AITER MLA requires 1-15 heads for Gluon decode or a multiple "
687
+ f"of 16 heads for persistent decode, but got {num_heads}.\n"
688
f"Try adjusting tensor_parallel_size value."
689
)
690
691
@staticmethod
692
def is_valid_num_heads(num_heads: int) -> bool:
656
- return (
657
- num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0
658
- if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
659
- else AiterMLAHelper._AITER_MIN_MLA_HEADS % num_heads == 0
693
+ return num_heads > 0 and (
694
+ num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS
695
+ or num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0
696
)
697
698
@staticmethod
@@ -681,6 +717,10 @@ class AiterMLAHelper:
717
else o[:, :: AiterMLAHelper._AITER_MIN_MLA_HEADS // num_heads, :]
718
)
719
720
+ @staticmethod
721
+ def use_gluon_decode(num_heads: int, max_qo_len: int) -> bool:
722
+ return num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS and max_qo_len == 1
723
+
724
725
class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
726
def __init__(
@@ -726,7 +766,10 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
766
767
# FP8 MLA prefill kernel imports (lazy, only when enabled).
768
# Auto-enabled on gfx950 when AITER ships the kernels.
729
- self._fp8_prefill_enabled = _fp8_mla_prefill_supported()
769
+ # FP8 MLA prefill (kn_mla_reduce_v1) only supports 16-aligned heads.
770
+ self._fp8_prefill_enabled = (
771
+ _fp8_mla_prefill_supported() and self.num_heads % 16 == 0
772
+ )
773
if self._fp8_prefill_enabled:
774
from aiter import mla_prefill_ps_asm_fwd, mla_reduce_v1
775
@@ -917,7 +960,106 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
960
) -> tuple[torch.Tensor, torch.Tensor | None]:
961
assert kv_c_and_k_pe_cache.numel() > 0
962
assert attn_metadata.decode is not None
920
- assert attn_metadata.decode.max_qo_len is not None
963
+
964
+ decode = attn_metadata.decode
965
+ assert decode.max_qo_len is not None
966
+ assert decode.paged_kv_indptr is not None
967
+ assert decode.paged_kv_indices is not None
968
+ if decode.use_gluon_decode:
969
+ if type(q) is tuple:
970
+ q_nope, q_pe = q
971
+ else:
972
+ q_nope, q_pe = torch.split(
973
+ q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
974
+ )
975
+ B, num_q_heads, _ = q_nope.shape
976
+ o = torch.empty(
977
+ B,
978
+ num_q_heads,
979
+ self.kv_lora_rank,
980
+ dtype=decode.attn_out_dtype,
981
+ device=q_nope.device,
982
+ )
983
+ kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1])
984
+ mla_gluon = _get_mla_gluon()
985
+ mla_gluon(
986
+ q_nope=q_nope,
987
+ q_pe=q_pe,
988
+ kv_c=kv_buffer,
989
+ o=o,
990
+ page_table=decode.paged_kv_indices,
991
+ seq_info=decode.paged_kv_indptr,
992
+ sm_scale=self.scale,
993
+ k_pe=None,
994
+ kv_pe_offset=self.kv_lora_rank,
995
+ use_2d_view=False,
996
+ kv_scale=1.0,
997
+ min_kv_seq_len=decode.min_kv_seq_len,
998
+ )
999
+ return o, None
1000
+
1001
+ # 12-head (<16) non-causal multi-token verify (DSpark): the asm path has
1002
+ # no gqa<16, qseqlen>1 kernel. Flatten each verify token to a qseqlen=1
1003
+ # gluon decode over the same committed prefix (non-causal), mirroring the
1004
+ # TRITON_MLA / sparse-backend flatten but on the fast gluon kernel. Each
1005
+ # request's KV metadata is repeated max_qo_len times.
1006
+ if (
1007
+ self.num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS
1008
+ and int(decode.max_qo_len) > 1
1009
+ ):
1010
+ qlen = int(decode.max_qo_len)
1011
+ if type(q) is tuple:
1012
+ q_nope, q_pe = q
1013
+ else:
1014
+ q_nope, q_pe = torch.split(
1015
+ q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
1016
+ )
1017
+ B, num_q_heads, _ = q_nope.shape
1018
+ o = torch.empty(
1019
+ B,
1020
+ num_q_heads,
1021
+ self.kv_lora_rank,
1022
+ dtype=decode.attn_out_dtype,
1023
+ device=q_nope.device,
1024
+ )
1025
+ kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1])
1026
+ # Expand per-request paged-KV to per-verify-token: each request's KV
1027
+ # block is repeated qlen times (row r*qlen+t attends request r's
1028
+ # committed prefix). Fully vectorized (no host loop).
1029
+ old_indptr = decode.paged_kv_indptr
1030
+ per_req_len = old_indptr[1:] - old_indptr[:-1]
1031
+ dev = q_nope.device
1032
+ row_req = torch.arange(per_req_len.shape[0], device=dev).repeat_interleave(
1033
+ qlen
1034
+ )
1035
+ row_len = per_req_len[row_req]
1036
+ new_indptr = torch.cat([old_indptr.new_zeros(1), row_len.cumsum(0)]).to(
1037
+ torch.int32
1038
+ )
1039
+ total = int(new_indptr[-1].item())
1040
+ within = torch.arange(total, device=dev, dtype=torch.int64) - new_indptr[
1041
+ :-1
1042
+ ].to(torch.int64).repeat_interleave(row_len)
1043
+ src = (
1044
+ old_indptr[row_req].to(torch.int64).repeat_interleave(row_len) + within
1045
+ )
1046
+ new_indices = decode.paged_kv_indices[src]
1047
+ mla_gluon = _get_mla_gluon()
1048
+ mla_gluon(
1049
+ q_nope=q_nope,
1050
+ q_pe=q_pe,
1051
+ kv_c=kv_buffer,
1052
+ o=o,
1053
+ page_table=new_indices,
1054
+ seq_info=new_indptr,
1055
+ sm_scale=self.scale,
1056
+ k_pe=None,
1057
+ kv_pe_offset=self.kv_lora_rank,
1058
+ use_2d_view=False,
1059
+ kv_scale=1.0,
1060
+ min_kv_seq_len=int(per_req_len.min()),
1061
+ )
1062
+ return o, None
1063
1064
if type(q) is tuple:
1065
q = torch.cat(q, dim=-1)
@@ -960,11 +1102,11 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1102
kv_buffer,
1103
o,
1104
self.scale,
963
- attn_metadata.decode.qo_indptr,
964
- attn_metadata.decode.max_qo_len,
965
- attn_metadata.decode.paged_kv_indptr,
966
- attn_metadata.decode.paged_kv_indices,
967
- attn_metadata.decode.paged_kv_last_page_len,
1105
+ decode.qo_indptr,
1106
+ decode.max_qo_len,
1107
+ decode.paged_kv_indptr,
1108
+ decode.paged_kv_indices,
1109
+ decode.paged_kv_last_page_len,
1110
**mla_kwargs,
1111
)
1112
vllm/v1/attention/backends/mla/triton_mla.py
+33
-2
@@ -51,9 +51,17 @@ class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
51
_cudagraph_support: ClassVar[AttentionCGSupport] = (
52
AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
53
)
54
+ # Non-causal DSpark block is flattened to one decode row per query token in
55
+ # forward_mqa, so no intra-block causal masking is required.
56
+ supports_non_causal_multi_token_decode: ClassVar[bool] = True
57
58
def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
59
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
60
+ # Only the non-causal DSpark draft group serves multi-token blocks via
61
+ # the decode path; raise its reorder threshold to the spec block length
62
+ # so full-cudagraph capture admits it. Causal usage stays single-token.
63
+ if getattr(self, "non_causal_multi_token_decode", False):
64
+ self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
65
self._reserve_attn_logits_workspace()
66
67
def _reserve_attn_logits_workspace(self) -> None:
@@ -68,6 +76,10 @@ class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
76
return
77
# Decode reorder threshold is 1, so decode tokens <= max_num_seqs.
78
B = self.vllm_config.scheduler_config.max_num_seqs
79
+ # Non-causal DSpark draft flattens each request's block to query_len
80
+ # decode rows; cover max_num_seqs * block_len rows.
81
+ if getattr(self, "non_causal_multi_token_decode", False):
82
+ B *= self.reorder_batch_threshold
83
# DCP all-gathers the query heads before forward_mqa.
84
q_num_heads = self.num_heads * self.dcp_world_size
85
max_splits = _compute_num_kv_splits(
@@ -132,6 +144,13 @@ class TritonMLABackend(MLACommonBackend):
144
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
145
return True
146
147
+ @classmethod
148
+ def supports_non_causal(cls) -> bool:
149
+ # DSpark non-causal blocks are flattened to single-token decode rows in
150
+ # TritonMLAImpl.forward_mqa (decode_attention_fwd has no causal flag /
151
+ # no intra-block masking). Enables the non-causal AMD MLA path.
152
+ return True
153
+
154
155
class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
156
can_return_lse_for_decode: bool = True
@@ -264,6 +283,18 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
283
kv_c_cache = kv_c_and_k_pe_cache[..., : self.kv_lora_rank]
284
PAGE_SIZE = kv_c_and_k_pe_cache.size(1)
285
286
+ block_table = attn_metadata.decode.block_table
287
+ seq_lens = attn_metadata.decode.seq_lens
288
+ if not attn_metadata.causal:
289
+ # Non-causal DSpark block: flatten to one decode row per query token.
290
+ # Each row attends to the same committed KV prefix (per-row seq_lens)
291
+ # and never to sibling block tokens = non-causal block semantics.
292
+ # Mirrors FlashInferMLA's non-causal path.
293
+ query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes
294
+ if query_len > 1:
295
+ block_table = block_table.repeat_interleave(query_len, dim=0)
296
+ seq_lens = seq_lens.repeat_interleave(query_len)
297
+
298
# Run MQA — always pass layer scales. When KV cache is
299
# BF16 the kernel's `if dtype.is_fp8()` check is a no-op.
300
decode_attention_fwd(
@@ -272,8 +303,8 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
303
kv_c_cache,
304
o,
305
lse,
275
- attn_metadata.decode.block_table,
276
- attn_metadata.decode.seq_lens,
306
+ block_table,
307
+ seq_lens,
308
attn_logits,
309
num_kv_splits,
310
self.scale,
vllm/v1/core/kv_cache_utils.py
+24
-3
@@ -1209,6 +1209,27 @@ def _get_kv_cache_groups_uniform_page_size(
1209
for layer_name, layer_spec in kv_cache_spec.items():
1210
same_type_layers[layer_spec].append(layer_name)
1211
1212
+ # Attempt to further merge same-type layers based on whether their KV
1213
+ # cache specs can be merged, to minimize the group count. This benefits
1214
+ # situations where specs share a block layout and differ only in a
1215
+ # property it can reconcile (e.g. full attention layers differing only in
1216
+ # sliding window / attention chunk size).
1217
+ layer_buckets: list[list[str]] = []
1218
+ spec_buckets: list[list[KVCacheSpec]] = []
1219
+ for layer_spec, layer_names in same_type_layers.items():
1220
+ for names, specs in zip(layer_buckets, spec_buckets):
1221
+ try:
1222
+ # A raise means that the specs are incompatible.
1223
+ type(specs[0]).merge([*specs, layer_spec])
1224
+ except (AssertionError, ValueError):
1225
+ continue
1226
+ names.extend(layer_names)
1227
+ specs.append(layer_spec)
1228
+ break
1229
+ else:
1230
+ layer_buckets.append(list(layer_names))
1231
+ spec_buckets.append([layer_spec])
1232
+
1233
# Split each group into smaller groups, to make the number of layers in each
1234
# group identical. Add padding to the last group of each type if necessary.
1235
# E.g., (full.0, full.1), (sw.0, sw.1, sw.2)
@@ -1221,9 +1242,9 @@ def _get_kv_cache_groups_uniform_page_size(
1242
# is the minimum number of layers among all attention types. Need a better
1243
# strategy if we want to support more complex patterns (e.g., 20 full + 30
1244
# sw, where the group size should be 10).
1224
- min_num_layers = min([len(layers) for layers in same_type_layers.values()])
1245
+ min_num_layers = min([len(layers) for layers in layer_buckets])
1246
group_size = min_num_layers
1226
- max_num_layers = max([len(layers) for layers in same_type_layers.values()])
1247
+ max_num_layers = max([len(layers) for layers in layer_buckets])
1248
if max_num_layers < min_num_layers * 1.5:
1249
# If the number of layers is not much larger than the minimum number of
1250
# 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(
1255
# extra layers to one attention type.
1256
group_size = max_num_layers
1257
grouped_layers = []
1237
- for layers in same_type_layers.values():
1258
+ for layers in layer_buckets:
1259
num_padding_layers = group_size - len(layers) % group_size
1260
if num_padding_layers != group_size:
1261
logger.warning(
vllm/v1/core/sched/scheduler.py
+15
-8
@@ -368,10 +368,11 @@ class Scheduler(SchedulerInterface):
368
) -> int:
369
"""Clip a prefill chunk so it ends where Mamba state must be cached.
370
371
- In "align" cache mode the SSM state is only materialized at chunk
372
- ends, so chunk ends are steered onto cacheable positions: block
373
- boundaries by default, plus mandatory early stops (the prompt's
374
- partial-tail hash boundary, a detected shared-prefix junction).
371
+ In "align" cache mode reusable SSM states are materialized at block
372
+ boundaries, plus mandatory early stops (the prompt's partial-tail hash
373
+ boundary, a detected shared-prefix junction). If a block is larger
374
+ than the configured prefill chunk limit, intermediate chunks keep
375
+ private running state until they reach the next cacheable position.
376
"""
377
start = (
378
request.num_computed_tokens
@@ -392,11 +393,17 @@ class Scheduler(SchedulerInterface):
393
last_cache_position = max(last_cache_position - block_size, 0)
394
395
end = start + num_new_tokens
395
- # Until `last_cache_position`, chunk ends must land on block
396
- # boundaries. May yield an empty chunk (budget cannot reach the next
397
- # boundary); the caller then skips the request.
396
+ # Until `last_cache_position`, prefer chunks ending on block
397
+ # boundaries. When a block cannot fit in any configured prefill chunk,
398
+ # allow sub-block progress and re-align at the next reachable boundary.
399
if end < last_cache_position:
399
- end = end // block_size * block_size
400
+ max_prefill_tokens = self.max_num_scheduled_tokens
401
+ long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold
402
+ if long_prefill_threshold > 0:
403
+ max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold)
404
+ aligned_end = end // block_size * block_size
405
+ if aligned_end > start or block_size <= max_prefill_tokens:
406
+ end = aligned_end
407
408
next_block_boundary = (start // block_size + 1) * block_size
409
tail_boundary = (
vllm/v1/kv_cache_interface.py
+13
-1
@@ -385,6 +385,8 @@ class MLAAttentionSpec(FullAttentionSpec):
385
alignment: int | None = None # Default to None for no padding.
386
compress_ratio: int = 1 # Default to 1 for no compression.
387
model_version: str | None = None
388
+ # Marks draft groups that flatten a non-causal query block into decode rows.
389
+ non_causal_multi_token_decode: bool = False
390
391
def __post_init__(self):
392
super().__post_init__()
@@ -434,7 +436,7 @@ class MLAAttentionSpec(FullAttentionSpec):
436
"quantization method, compress ratio, model version, and KV block "
437
"stride indexing."
438
)
437
- return cls(
439
+ merged_spec = cls(
440
block_size=specs[0].block_size,
441
num_kv_heads=specs[0].num_kv_heads,
442
head_size=specs[0].head_size,
@@ -445,7 +447,17 @@ class MLAAttentionSpec(FullAttentionSpec):
447
cache_dtype_str=cache_dtype_str_set.pop(),
448
compress_ratio=compress_ratio_set.pop(),
449
model_version=model_version_set.pop(),
450
+ non_causal_multi_token_decode=any(
451
+ spec.non_causal_multi_token_decode for spec in specs
452
+ ),
453
)
454
+ for spec in specs:
455
+ for f in fields(AttentionSpec):
456
+ assert getattr(spec, f.name) == getattr(merged_spec, f.name), (
457
+ "All attention layers in the same KV cache group must have "
458
+ "the same attention spec."
459
+ )
460
+ return merged_spec
461
462
463
@dataclass(frozen=True, kw_only=True)
vllm/v1/spec_decode/llm_base_proposer.py
+9
-6
@@ -1006,11 +1006,11 @@ class SpecDecodeBaseProposer:
1006
1007
def model_returns_tuple(self) -> bool:
1008
if self.method == "mtp":
1009
- # DeepSeek-family MTP (deepseek_mtp.py) recycles the post-final-
1010
- # norm hidden, so its forward returns (logit_hidden,
1011
- # recycle_hidden). Other MTP families return a single tensor.
1012
- return "DeepSeekMTPModel" in (
1013
- self.draft_model_config.hf_config.architectures or []
1009
+ # These models return separate hidden states for logits and for
1010
+ # feedback into the next draft step.
1011
+ architectures = self.draft_model_config.hf_config.architectures or []
1012
+ return bool(
1013
+ {"DeepSeekMTPModel", "KimiK3MTPModel"}.intersection(architectures)
1014
)
1015
return self.method not in ("mtp", "draft_model", "dflash")
1016
@@ -1386,7 +1386,10 @@ class SpecDecodeBaseProposer:
1386
self.model.config.image_token_index = (
1387
target_model.config.vision_config.image_token_id
1388
)
1389
- elif self.get_model_name(target_model) == "KimiK25ForConditionalGeneration":
1389
+ elif self.get_model_name(target_model) in (
1390
+ "KimiK25ForConditionalGeneration",
1391
+ "KimiK3ForConditionalGeneration",
1392
+ ):
1393
self.model.config.image_token_index = (
1394
target_model.config.media_placeholder_token_id
1395
)
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+25
-7
@@ -344,13 +344,6 @@ class DFlashSpeculator(DraftModelSpeculator):
344
hidden_states = last_hidden_states
345
self.hidden_states[:num_target_tokens].copy_(hidden_states[:num_target_tokens])
346
347
- self._copy_request_inputs(
348
- num_reqs,
349
- input_batch.idx_mapping,
350
- temperature,
351
- seeds,
352
- )
353
-
347
if dummy_run and skip_attn_for_dummy_run:
348
# Memory profiling path: block_tables / kv_cache_config are not initialized.
349
# Since DFlash needs to build its own attention metadata, we must skip the
@@ -385,11 +378,15 @@ class DFlashSpeculator(DraftModelSpeculator):
378
self.sample_indices,
379
self.sample_pos,
380
self.sample_idx_mapping,
381
+ self.temperature,
382
+ self.seeds,
383
input_batch,
384
num_sampled,
385
num_rejected,
386
last_sampled,
387
next_prefill_tokens,
388
+ temperature,
389
+ seeds,
390
self.block_tables.input_block_tables[gid],
391
self.block_tables.kernel_block_sizes[gid],
392
self.parallel_drafting_token_id,
@@ -482,6 +479,8 @@ def _prepare_dflash_inputs_kernel(
479
out_sample_indices_ptr,
480
out_sample_pos_ptr,
481
out_sample_idx_mapping_ptr,
482
+ out_temperature_ptr,
483
+ out_seeds_ptr,
484
# Inputs from target batch
485
target_positions_ptr,
486
target_query_start_loc_ptr,
@@ -490,6 +489,9 @@ def _prepare_dflash_inputs_kernel(
489
next_prefill_tokens_ptr,
490
num_sampled_ptr,
491
num_rejected_ptr,
492
+ # Sampling params
493
+ temperature_ptr,
494
+ seeds_ptr,
495
# Block table for slot mapping lookup.
496
block_table_ptr,
497
block_table_stride,
@@ -585,6 +587,12 @@ def _prepare_dflash_inputs_kernel(
587
# reads up to (context + query), not just the count of accepted
588
# tokens this step.
589
tl.store(out_seq_lens_ptr + req_idx, last_valid_pos + 1 + num_query_per_req)
590
+ # Copy sampling state.
591
+ tl.store(
592
+ out_temperature_ptr + req_state_idx,
593
+ tl.load(temperature_ptr + req_state_idx),
594
+ )
595
+ tl.store(out_seeds_ptr + req_state_idx, tl.load(seeds_ptr + req_state_idx))
596
if req_idx == num_reqs - 1:
597
# Pad per-request buffers to max_num_reqs for CUDA graph safety.
598
last_query_end = num_reqs * num_query_per_req
@@ -626,6 +634,8 @@ def prepare_dflash_inputs(
634
sample_indices: torch.Tensor,
635
sample_pos: torch.Tensor,
636
sample_idx_mapping: torch.Tensor,
637
+ temperature: torch.Tensor,
638
+ seeds: torch.Tensor,
639
input_batch: InputBatch,
640
# [num_reqs]
641
num_sampled: torch.Tensor,
@@ -635,6 +645,10 @@ def prepare_dflash_inputs(
645
last_sampled: torch.Tensor,
646
# [max_num_reqs]
647
next_prefill_tokens: torch.Tensor,
648
+ # [max_num_reqs]
649
+ input_temperature: torch.Tensor,
650
+ # [max_num_reqs]
651
+ input_seeds: torch.Tensor,
652
# [max_num_reqs, max_num_blocks]
653
block_table: torch.Tensor,
654
block_size: int,
@@ -665,6 +679,8 @@ def prepare_dflash_inputs(
679
sample_indices,
680
sample_pos,
681
sample_idx_mapping,
682
+ temperature,
683
+ seeds,
684
input_batch.positions,
685
input_batch.query_start_loc,
686
input_batch.idx_mapping,
@@ -672,6 +688,8 @@ def prepare_dflash_inputs(
688
next_prefill_tokens,
689
num_sampled,
690
num_rejected,
691
+ input_temperature,
692
+ input_seeds,
693
block_table,
694
block_table.stride(0),
695
parallel_drafting_token_id,
vllm/v1/worker/gpu_model_runner.py
+4
@@ -6935,6 +6935,10 @@ class GPUModelRunner(
6935
num_active_loras=desc.num_active_loras,
6936
profile_seq_lens=profile_seq_lens,
6937
)
6938
+ if num_warmups > 0:
6939
+ # Warmups may use auxiliary streams. Ensure all of their work has
6940
+ # completed before beginning CUDA graph capture.
6941
+ torch.accelerator.synchronize()
6942
with (
6943
profiler,
6944
torch.profiler.record_function(