@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d4da0c55af

[Model][Frontend] Add Ling 3.0 Flash BF16, MTP, and parser support (#51045)

Signed-off-by: zexplorerhj <zhjoneson@163.com>

zexplorerhj committed Aug 6, 2026 at 00:38 UTC d4da0c55af3aa231b6209bf77871f3ed36eab0d2
15 files changed +2070 -1
benchmarks/kernels/benchmark_moe.py
+1
@@ -777,6 +777,7 @@ def get_model_params(config):
777 intermediate_size = config.moe_intermediate_size
778 hidden_size = config.hidden_size
779 elif architecture in (
780 + "BailingMoeV3ForCausalLM",
781 "Qwen2MoeForCausalLM",
782 "Qwen3MoeForCausalLM",
783 "Qwen3NextForCausalLM",
docs/models/supported_models.md
+1
@@ -348,6 +348,7 @@ th {
348 | `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | ✅︎ | ✅︎ |
349 | `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ |
350 | `BailingMoeV2_5ForCausalLM` | Ling | `inclusionAI/Ling-2.5-1T`, `inclusionAI/Ring-2.5-1T` | | ✅︎ |
351 +| `BailingMoeV3ForCausalLM` | Ling | `inclusionAI/Ling-3.0-flash` | | ✅︎ |
352 | `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ |
353 | `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
354 | `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+8
@@ -213,6 +213,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
213 "BailingMoeV2_5ForCausalLM": _HfExamplesInfo(
214 "inclusionAI/Ring-2.5-1T", trust_remote_code=True
215 ),
216 + "BailingMoeV3ForCausalLM": _HfExamplesInfo(
217 + "inclusionAI/Ling-3.0-flash", trust_remote_code=True
218 + ),
219 "BloomForCausalLM": _HfExamplesInfo(
220 "bigscience/bloom-560m", {"1b": "bigscience/bloomz-1b1"}
221 ),
@@ -1614,6 +1617,11 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1617 trust_remote_code=True,
1618 is_available_online=False,
1619 ),
1620 + "BailingMoeV3MTPModel": _HfExamplesInfo(
1621 + "inclusionAI/Ling-3.0-flash",
1622 + speculative_model="inclusionAI/Ling-3.0-flash",
1623 + trust_remote_code=True,
1624 + ),
1625 "DeepSeekMTPModel": _HfExamplesInfo(
1626 "luccafong/deepseek_mtp_main_random",
1627 speculative_model="luccafong/deepseek_mtp_draft_random",
tests/parser/engine/test_ling3.py new
+105
@@ -0,0 +1,105 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import json
5 +
6 +from tests.parser.engine.replay_harness import MockTokenizer, _test_request
7 +from vllm.reasoning import ReasoningParserManager
8 +from vllm.tool_parsers import ToolParserManager
9 +
10 +
11 +def _tokenizer() -> MockTokenizer:
12 + vocab = {
13 + "<think>": 1,
14 + "</think>": 2,
15 + "<tool_call>": 3,
16 + "</tool_call>": 4,
17 + "<arg_key>": 5,
18 + "</arg_key>": 6,
19 + "<arg_value>": 7,
20 + "</arg_value>": 8,
21 + }
22 + return MockTokenizer(vocab=vocab, tokens=[])
23 +
24 +
25 +def test_ling3_registered():
26 + assert ReasoningParserManager.get_reasoning_parser("ling3") is not None
27 + assert ToolParserManager.get_tool_parser("ling3") is not None
28 +
29 +
30 +def test_ling3_defaults_thinking_on():
31 + parser_cls = ReasoningParserManager.get_reasoning_parser("ling3")
32 + parser = parser_cls(_tokenizer())
33 +
34 + assert parser.reasoning_start_str == "<think>"
35 + assert parser.reasoning_end_str == "</think>"
36 +
37 + reasoning, content = parser.extract_reasoning(
38 + "<think>reason</think>answer", _test_request()
39 + )
40 +
41 + assert reasoning == "reason"
42 + assert content == "answer"
43 +
44 +
45 +def test_ling3_disable_thinking_keeps_reasoning_as_content():
46 + parser_cls = ReasoningParserManager.get_reasoning_parser("ling3")
47 + parser = parser_cls(
48 + _tokenizer(),
49 + chat_template_kwargs={"enable_thinking": False},
50 + )
51 +
52 + reasoning, content = parser.extract_reasoning(
53 + "<think>reason</think>answer", _test_request()
54 + )
55 +
56 + assert reasoning is None
57 + assert content == "<think>reason</think>answer"
58 +
59 +
60 +def test_ling3_enable_thinking_keeps_open_reasoning_as_content():
61 + parser_cls = ReasoningParserManager.get_reasoning_parser("ling3")
62 + parser = parser_cls(
63 + _tokenizer(),
64 + chat_template_kwargs={"enable_thinking": True},
65 + )
66 +
67 + reasoning, content = parser.extract_reasoning(
68 + "<think>only reasoning", _test_request()
69 + )
70 +
71 + assert reasoning is None
72 + assert content == "only reasoning"
73 +
74 +
75 +def test_ling3_tool_call_without_newline():
76 + parser_cls = ToolParserManager.get_tool_parser("ling3")
77 + parser = parser_cls(_tokenizer())
78 +
79 + result = parser.extract_tool_calls(
80 + "<tool_call>get_weather<arg_key>city</arg_key>"
81 + "<arg_value>Beijing</arg_value></tool_call>",
82 + request=_test_request(),
83 + )
84 +
85 + assert result.tools_called
86 + assert result.content is None
87 + assert result.tool_calls[0].function.name == "get_weather"
88 + assert json.loads(result.tool_calls[0].function.arguments) == {
89 + "city": "Beijing",
90 + }
91 +
92 +
93 +def test_ling3_tool_call_without_args():
94 + parser_cls = ToolParserManager.get_tool_parser("ling3")
95 + parser = parser_cls(_tokenizer())
96 +
97 + result = parser.extract_tool_calls(
98 + "<tool_call>ping</tool_call>",
99 + request=_test_request(),
100 + )
101 +
102 + assert result.tools_called
103 + assert result.content is None
104 + assert result.tool_calls[0].function.name == "ping"
105 + assert json.loads(result.tool_calls[0].function.arguments) == {}
vllm/config/speculative.py
+12 -1
@@ -48,6 +48,7 @@ MTPModelTypes = Literal[
48 "qwen3_next_mtp",
49 "qwen3_5_mtp",
50 "longcat_flash_mtp",
51 + "bailing_hybrid_v3_mtp",
52 "minimax_m3_mtp",
53 "bailing_hybrid_mtp",
54 "mtp",
@@ -489,7 +490,9 @@ class SpeculativeConfig:
490 )
491
492 architectures = getattr(hf_config, "architectures", []) or []
492 - if (
493 + if initial_architecture == "BailingMoeV3ForCausalLM":
494 + hf_config.model_type = "bailing_hybrid_v3_mtp"
495 + elif (
496 hf_config.model_type == "bailing_hybrid"
497 or "BailingMoeV2_5ForCausalLM" in architectures
498 ):
@@ -502,6 +505,14 @@ class SpeculativeConfig:
505 "architectures": ["BailingMoeV25MTPModel"],
506 }
507 )
508 + if hf_config.model_type == "bailing_hybrid_v3_mtp":
509 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
510 + hf_config.update(
511 + {
512 + "n_predict": n_predict,
513 + "architectures": ["BailingMoeV3MTPModel"],
514 + }
515 + )
516
517 if hf_config.model_type == "exaone_moe":
518 hf_config.model_type = "exaone_moe_mtp"
vllm/model_executor/layers/fused_moe/configs/E=512,N=192,device_name=NVIDIA_H20-3e.json new
+147
@@ -0,0 +1,147 @@
1 +{
2 + "triton_version": "3.6.0",
3 + "1": {
4 + "BLOCK_SIZE_M": 16,
5 + "BLOCK_SIZE_N": 32,
6 + "BLOCK_SIZE_K": 64,
7 + "GROUP_SIZE_M": 1,
8 + "num_warps": 4,
9 + "num_stages": 4
10 + },
11 + "2": {
12 + "BLOCK_SIZE_M": 16,
13 + "BLOCK_SIZE_N": 64,
14 + "BLOCK_SIZE_K": 64,
15 + "GROUP_SIZE_M": 1,
16 + "num_warps": 4,
17 + "num_stages": 4
18 + },
19 + "4": {
20 + "BLOCK_SIZE_M": 16,
21 + "BLOCK_SIZE_N": 64,
22 + "BLOCK_SIZE_K": 64,
23 + "GROUP_SIZE_M": 32,
24 + "num_warps": 4,
25 + "num_stages": 4
26 + },
27 + "8": {
28 + "BLOCK_SIZE_M": 16,
29 + "BLOCK_SIZE_N": 64,
30 + "BLOCK_SIZE_K": 64,
31 + "GROUP_SIZE_M": 64,
32 + "num_warps": 4,
33 + "num_stages": 4
34 + },
35 + "16": {
36 + "BLOCK_SIZE_M": 16,
37 + "BLOCK_SIZE_N": 128,
38 + "BLOCK_SIZE_K": 64,
39 + "GROUP_SIZE_M": 1,
40 + "num_warps": 4,
41 + "num_stages": 3
42 + },
43 + "24": {
44 + "BLOCK_SIZE_M": 16,
45 + "BLOCK_SIZE_N": 64,
46 + "BLOCK_SIZE_K": 64,
47 + "GROUP_SIZE_M": 1,
48 + "num_warps": 4,
49 + "num_stages": 4
50 + },
51 + "32": {
52 + "BLOCK_SIZE_M": 16,
53 + "BLOCK_SIZE_N": 64,
54 + "BLOCK_SIZE_K": 64,
55 + "GROUP_SIZE_M": 1,
56 + "num_warps": 4,
57 + "num_stages": 3
58 + },
59 + "48": {
60 + "BLOCK_SIZE_M": 16,
61 + "BLOCK_SIZE_N": 128,
62 + "BLOCK_SIZE_K": 64,
63 + "GROUP_SIZE_M": 16,
64 + "num_warps": 4,
65 + "num_stages": 3
66 + },
67 + "64": {
68 + "BLOCK_SIZE_M": 16,
69 + "BLOCK_SIZE_N": 128,
70 + "BLOCK_SIZE_K": 64,
71 + "GROUP_SIZE_M": 16,
72 + "num_warps": 4,
73 + "num_stages": 3
74 + },
75 + "96": {
76 + "BLOCK_SIZE_M": 16,
77 + "BLOCK_SIZE_N": 128,
78 + "BLOCK_SIZE_K": 64,
79 + "GROUP_SIZE_M": 16,
80 + "num_warps": 4,
81 + "num_stages": 3
82 + },
83 + "128": {
84 + "BLOCK_SIZE_M": 16,
85 + "BLOCK_SIZE_N": 128,
86 + "BLOCK_SIZE_K": 64,
87 + "GROUP_SIZE_M": 16,
88 + "num_warps": 4,
89 + "num_stages": 3
90 + },
91 + "256": {
92 + "BLOCK_SIZE_M": 16,
93 + "BLOCK_SIZE_N": 128,
94 + "BLOCK_SIZE_K": 64,
95 + "GROUP_SIZE_M": 1,
96 + "num_warps": 4,
97 + "num_stages": 3
98 + },
99 + "512": {
100 + "BLOCK_SIZE_M": 16,
101 + "BLOCK_SIZE_N": 128,
102 + "BLOCK_SIZE_K": 64,
103 + "GROUP_SIZE_M": 16,
104 + "num_warps": 4,
105 + "num_stages": 3
106 + },
107 + "1024": {
108 + "BLOCK_SIZE_M": 16,
109 + "BLOCK_SIZE_N": 128,
110 + "BLOCK_SIZE_K": 64,
111 + "GROUP_SIZE_M": 16,
112 + "num_warps": 4,
113 + "num_stages": 3
114 + },
115 + "1536": {
116 + "BLOCK_SIZE_M": 32,
117 + "BLOCK_SIZE_N": 128,
118 + "BLOCK_SIZE_K": 64,
119 + "GROUP_SIZE_M": 64,
120 + "num_warps": 4,
121 + "num_stages": 3
122 + },
123 + "2048": {
124 + "BLOCK_SIZE_M": 64,
125 + "BLOCK_SIZE_N": 128,
126 + "BLOCK_SIZE_K": 64,
127 + "GROUP_SIZE_M": 1,
128 + "num_warps": 4,
129 + "num_stages": 2
130 + },
131 + "3072": {
132 + "BLOCK_SIZE_M": 64,
133 + "BLOCK_SIZE_N": 128,
134 + "BLOCK_SIZE_K": 64,
135 + "GROUP_SIZE_M": 1,
136 + "num_warps": 4,
137 + "num_stages": 2
138 + },
139 + "4096": {
140 + "BLOCK_SIZE_M": 64,
141 + "BLOCK_SIZE_N": 128,
142 + "BLOCK_SIZE_K": 64,
143 + "GROUP_SIZE_M": 1,
144 + "num_warps": 4,
145 + "num_stages": 2
146 + }
147 +}
vllm/model_executor/models/bailing_moe_v3.py new
+1289
@@ -0,0 +1,1289 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""vLLM implementation for BailingMoeV3ForCausalLM.
4 +
5 +The HuggingFace reference model mixes MLA full-attention layers with Kimi
6 +Delta Attention linear layers and Bailing MoE blocks. This file keeps the V3
7 +module/weight names aligned with the reference implementation while reusing
8 +vLLM's parallel linear layers, MLA kernel, KDA kernel and fused MoE loader.
9 +"""
10 +
11 +import copy
12 +from collections.abc import Iterable
13 +
14 +import torch
15 +import torch.nn as nn
16 +import torch.nn.functional as F
17 +from einops import rearrange
18 +from transformers.configuration_utils import PretrainedConfig
19 +
20 +from vllm.compilation.decorators import support_torch_compile
21 +from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
22 +from vllm.distributed import (
23 + get_pp_group,
24 + get_tensor_model_parallel_rank,
25 + get_tensor_model_parallel_world_size,
26 +)
27 +from vllm.forward_context import ForwardContext, get_forward_context
28 +from vllm.model_executor.custom_op import PluggableLayer
29 +from vllm.model_executor.layers.activation import SiluAndMul, SwigluStepAndMul
30 +from vllm.model_executor.layers.fused_moe import (
31 + FusedMoEFactory,
32 + fused_moe_make_expert_params_mapping,
33 +)
34 +from vllm.model_executor.layers.layernorm import RMSNorm
35 +from vllm.model_executor.layers.linear import (
36 + ColumnParallelLinear,
37 + MergedColumnParallelLinear,
38 + ReplicatedLinear,
39 + RowParallelLinear,
40 +)
41 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
42 +from vllm.model_executor.layers.mamba.abstract import MambaBase
43 +from vllm.model_executor.layers.mamba.mamba_utils import (
44 + MambaStateCopyFuncCalculator,
45 + MambaStateDtypeCalculator,
46 + MambaStateShapeCalculator,
47 + is_conv_state_dim_first,
48 +)
49 +from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
50 + causal_conv1d_fn,
51 + causal_conv1d_update,
52 +)
53 +from vllm.model_executor.layers.mla import (
54 + MLAModules,
55 + MultiHeadLatentAttentionWrapper,
56 +)
57 +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
58 +from vllm.model_executor.layers.rotary_embedding import get_rope
59 +from vllm.model_executor.layers.vocab_parallel_embedding import (
60 + ParallelLMHead,
61 + VocabParallelEmbedding,
62 +)
63 +from vllm.model_executor.model_loader.weight_utils import (
64 + default_weight_loader,
65 + maybe_remap_kv_scale_name,
66 + sharded_weight_loader,
67 +)
68 +from vllm.model_executor.utils import set_weight_attrs
69 +from vllm.sequence import IntermediateTensors
70 +from vllm.third_party.flash_linear_attention.ops.kda import (
71 + FusedRMSNormGated,
72 + chunk_kda,
73 + fused_kda_gate,
74 + fused_recurrent_kda,
75 + fused_recurrent_kda_fwd,
76 +)
77 +from vllm.utils.torch_utils import direct_register_custom_op
78 +from vllm.v1.attention.backend import AttentionMetadata
79 +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
80 +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
81 +
82 +from .interfaces import HasInnerState, IsHybrid, SupportsPP
83 +from .utils import PPMissingLayer, is_pp_missing_parameter, make_layers, maybe_prefix
84 +
85 +
86 +def bailing_v3_kda_attention(
87 + q_proj_states: torch.Tensor,
88 + k_proj_states: torch.Tensor,
89 + v_proj_states: torch.Tensor,
90 + g1: torch.Tensor,
91 + beta: torch.Tensor,
92 + core_attn_out: torch.Tensor,
93 + layer_name: str,
94 +) -> None:
95 + """Run Bailing V3's KDA state update outside the compiled graph."""
96 + forward_context: ForwardContext = get_forward_context()
97 + layer = forward_context.no_compile_layers[layer_name]
98 + layer._forward(
99 + q_proj_states=q_proj_states,
100 + k_proj_states=k_proj_states,
101 + v_proj_states=v_proj_states,
102 + g1=g1,
103 + beta=beta,
104 + core_attn_out=core_attn_out,
105 + )
106 +
107 +
108 +def bailing_v3_kda_attention_fake(
109 + q_proj_states: torch.Tensor,
110 + k_proj_states: torch.Tensor,
111 + v_proj_states: torch.Tensor,
112 + g1: torch.Tensor,
113 + beta: torch.Tensor,
114 + core_attn_out: torch.Tensor,
115 + layer_name: str,
116 +) -> None:
117 + return
118 +
119 +
120 +direct_register_custom_op(
121 + op_name="bailing_v3_kda_attention",
122 + op_func=bailing_v3_kda_attention,
123 + mutates_args=["core_attn_out"],
124 + fake_impl=bailing_v3_kda_attention_fake,
125 +)
126 +
127 +
128 +def _is_kda_layer(
129 + layer_idx: int, layer_group_size: int, num_hidden_layers: int
130 +) -> bool:
131 + return not (
132 + (layer_idx + 1) % layer_group_size == 0
133 + or layer_idx >= num_hidden_layers // layer_group_size * layer_group_size
134 + )
135 +
136 +
137 +def _get_kda_state_shape_for_config(
138 + vllm_config: VllmConfig,
139 +) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
140 + config = vllm_config.model_config.hf_config
141 + num_spec = (
142 + vllm_config.speculative_config.num_speculative_tokens
143 + if vllm_config.speculative_config
144 + else 0
145 + )
146 + shapes = MambaStateShapeCalculator.kda_state_shape(
147 + vllm_config.parallel_config.tensor_parallel_size,
148 + config.num_attention_heads,
149 + config.head_dim,
150 + conv_kernel_size=config.short_conv_kernel_size,
151 + num_spec=num_spec,
152 + )
153 +
154 + return shapes
155 +
156 +
157 +def _build_rope_parameters(config: PretrainedConfig) -> dict | None:
158 + rope_parameters = copy.deepcopy(getattr(config, "rope_parameters", None)) or {}
159 + if "rope_theta" not in rope_parameters and hasattr(config, "rope_theta"):
160 + rope_parameters["rope_theta"] = config.rope_theta
161 + # BailingMoeV3 stores the MLA rotary width explicitly in qk_rope_head_dim.
162 + # Some checkpoints also carry partial_rotary_factor=0.5, but applying it
163 + # here would shrink the rotary cache to half of the qk_rope_head_dim.
164 + rope_parameters.pop("partial_rotary_factor", None)
165 +
166 + rope_scaling = getattr(config, "rope_scaling", None)
167 + if isinstance(rope_scaling, dict):
168 + rope_scaling = copy.deepcopy(rope_scaling)
169 + if "type" in rope_scaling and "rope_type" not in rope_scaling:
170 + rope_scaling["rope_type"] = rope_scaling.pop("type")
171 + rope_scaling.pop("partial_rotary_factor", None)
172 + rope_parameters.update(rope_scaling)
173 +
174 + return rope_parameters or None
175 +
176 +
177 +def _get_layer_swiglu_limit(limit_list: list | None, layer_idx: int) -> float | None:
178 + if limit_list is None or layer_idx >= len(limit_list):
179 + return None
180 + limit = limit_list[layer_idx]
181 + if limit in (None, 0):
182 + return None
183 + return float(limit)
184 +
185 +
186 +def _load_a_log(param: torch.nn.Parameter, loaded_weight: torch.Tensor) -> None:
187 + tp_rank = get_tensor_model_parallel_rank()
188 + tp_size = get_tensor_model_parallel_world_size()
189 + if loaded_weight.ndim == 1:
190 + shard_size = loaded_weight.shape[0] // tp_size
191 + shard = loaded_weight.narrow(0, tp_rank * shard_size, shard_size)
192 + param.data.copy_(shard.view_as(param.data))
193 + else:
194 + sharded_weight_loader(2)(param, loaded_weight)
195 +
196 +
197 +class _IdentityProjection(nn.Module):
198 + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]:
199 + return x, None
200 +
201 +
202 +class BailingMoeV3MLP(nn.Module):
203 + def __init__(
204 + self,
205 + intermediate_size: int,
206 + config: PretrainedConfig,
207 + quant_config: QuantizationConfig | None = None,
208 + reduce_results: bool | None = True,
209 + prefix: str = "",
210 + swiglu_limit: float | None = None,
211 + ) -> None:
212 + super().__init__()
213 + self.gate_up_proj = MergedColumnParallelLinear(
214 + config.hidden_size,
215 + [intermediate_size] * 2,
216 + bias=config.use_bias,
217 + quant_config=quant_config,
218 + prefix=f"{prefix}.gate_up_proj",
219 + )
220 + self.down_proj = RowParallelLinear(
221 + intermediate_size,
222 + config.hidden_size,
223 + bias=config.use_bias,
224 + quant_config=quant_config,
225 + reduce_results=reduce_results,
226 + prefix=f"{prefix}.down_proj",
227 + )
228 + self.act_fn = (
229 + SwigluStepAndMul(limit=swiglu_limit)
230 + if swiglu_limit not in (None, 0)
231 + else SiluAndMul()
232 + )
233 +
234 + def forward(self, x):
235 + x, _ = self.gate_up_proj(x)
236 + x = self.act_fn(x)
237 + x, _ = self.down_proj(x)
238 + return x
239 +
240 +
241 +class BailingMoeV3MLAAttention(nn.Module):
242 + def __init__(
243 + self,
244 + config: PretrainedConfig,
245 + quant_config: QuantizationConfig | None = None,
246 + layer_id: int = 0,
247 + prefix: str = "attention",
248 + cache_config: CacheConfig | None = None,
249 + ) -> None:
250 + super().__init__()
251 + self.hidden_size = config.hidden_size
252 + self.num_heads = config.num_attention_heads
253 + self.layer_id = layer_id
254 +
255 + self.qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 128)
256 + self.qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 64)
257 + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
258 + self.v_head_dim = getattr(config, "v_head_dim", 128)
259 + self.q_lora_rank = getattr(config, "q_lora_rank", None)
260 + self.kv_lora_rank = getattr(config, "kv_lora_rank", 512)
261 +
262 + tp_size = get_tensor_model_parallel_world_size()
263 + assert self.num_heads % tp_size == 0
264 + self.num_local_heads = self.num_heads // tp_size
265 + self.scaling = self.qk_head_dim**-0.5
266 +
267 + if self.q_lora_rank is None:
268 + self.q_proj = ColumnParallelLinear(
269 + self.hidden_size,
270 + self.num_heads * self.qk_head_dim,
271 + bias=False,
272 + quant_config=quant_config,
273 + prefix=f"{prefix}.q_proj",
274 + )
275 + self.fused_qkv_a_proj = None
276 + self.q_a_layernorm = None
277 + self.q_b_proj = None
278 + self.kv_a_proj_with_mqa = ReplicatedLinear(
279 + self.hidden_size,
280 + self.kv_lora_rank + self.qk_rope_head_dim,
281 + bias=False,
282 + quant_config=quant_config,
283 + prefix=f"{prefix}.kv_a_proj_with_mqa",
284 + )
285 + else:
286 + self.fused_qkv_a_proj = MergedColumnParallelLinear(
287 + self.hidden_size,
288 + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
289 + bias=False,
290 + quant_config=quant_config,
291 + prefix=f"{prefix}.fused_qkv_a_proj",
292 + disable_tp=True,
293 + )
294 + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
295 + self.q_b_proj = ColumnParallelLinear(
296 + self.q_lora_rank,
297 + self.num_heads * self.qk_head_dim,
298 + bias=False,
299 + quant_config=quant_config,
300 + prefix=f"{prefix}.q_b_proj",
301 + )
302 + self.q_proj = None
303 + self.kv_a_proj_with_mqa = None
304 +
305 + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
306 + self.kv_b_proj = ColumnParallelLinear(
307 + self.kv_lora_rank,
308 + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
309 + bias=False,
310 + quant_config=quant_config,
311 + prefix=f"{prefix}.kv_b_proj",
312 + )
313 + self.gated_attention_proj_granularity_type = getattr(
314 + config, "gated_attention_proj_granularity_type", None
315 + )
316 + if self.gated_attention_proj_granularity_type == "head_wise":
317 + g_out = self.num_heads
318 + elif self.gated_attention_proj_granularity_type == "element_wise":
319 + g_out = self.num_heads * self.v_head_dim
320 + else:
321 + g_out = 0
322 + self.g_proj = (
323 + ColumnParallelLinear(
324 + self.hidden_size,
325 + g_out,
326 + bias=False,
327 + quant_config=quant_config,
328 + prefix=f"{prefix}.g_proj",
329 + )
330 + if g_out > 0
331 + else None
332 + )
333 + self.dense = RowParallelLinear(
334 + self.num_heads * self.v_head_dim,
335 + self.hidden_size,
336 + bias=False,
337 + quant_config=quant_config,
338 + prefix=f"{prefix}.dense",
339 + )
340 +
341 + self.rotary_emb = get_rope(
342 + head_size=self.qk_rope_head_dim,
343 + max_position=getattr(config, "max_position_embeddings", 8192),
344 + is_neox_style=False,
345 + rope_parameters=_build_rope_parameters(config),
346 + )
347 + mla_modules = MLAModules(
348 + kv_a_layernorm=self.kv_a_layernorm,
349 + kv_b_proj=self.kv_b_proj,
350 + rotary_emb=self.rotary_emb,
351 + o_proj=_IdentityProjection(),
352 + fused_qkv_a_proj=self.fused_qkv_a_proj,
353 + kv_a_proj_with_mqa=self.kv_a_proj_with_mqa,
354 + q_a_layernorm=self.q_a_layernorm,
355 + q_b_proj=self.q_b_proj,
356 + q_proj=self.q_proj,
357 + indexer=None,
358 + is_sparse=False,
359 + topk_indices_buffer=None,
360 + )
361 + self.mla_attn = MultiHeadLatentAttentionWrapper(
362 + self.hidden_size,
363 + self.num_local_heads,
364 + self.scaling,
365 + self.qk_nope_head_dim,
366 + self.qk_rope_head_dim,
367 + self.v_head_dim,
368 + self.q_lora_rank,
369 + self.kv_lora_rank,
370 + mla_modules,
371 + cache_config,
372 + quant_config,
373 + prefix,
374 + )
375 +
376 + def forward(self, hidden_states: torch.Tensor, positions: torch.Tensor):
377 + attn_out = self.mla_attn(positions, hidden_states)
378 + if self.g_proj is not None:
379 + gate = torch.sigmoid(self.g_proj(hidden_states)[0].float()).to(
380 + hidden_states.dtype
381 + )
382 + if self.gated_attention_proj_granularity_type == "head_wise":
383 + attn_out = attn_out.view(-1, self.num_local_heads, self.v_head_dim)
384 + attn_out = attn_out * gate.unsqueeze(-1)
385 + attn_out = attn_out.reshape(hidden_states.shape[0], -1)
386 + else:
387 + attn_out = attn_out * gate
388 +
389 + return self.dense(attn_out)[0]
390 +
391 +
392 +# --8<-- [start:bailing_moe_v3_kimi_delta_attention]
393 +@PluggableLayer.register("bailing_moe_v3_kimi_delta_attention")
394 +class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase):
395 + # --8<-- [end:bailing_moe_v3_kimi_delta_attention]
396 +
397 + @property
398 + def mamba_type(self) -> MambaAttentionBackendEnum:
399 + return MambaAttentionBackendEnum.GDN_ATTN
400 +
401 + def get_state_dtype(
402 + self,
403 + ) -> tuple[torch.dtype, torch.dtype, torch.dtype, torch.dtype]:
404 + if self.model_config is None or self.cache_config is None:
405 + raise ValueError("model_config and cache_config must be set")
406 + return MambaStateDtypeCalculator.kda_state_dtype(
407 + self.model_config.dtype, self.cache_config.mamba_cache_dtype
408 + )
409 +
410 + def get_state_shape(
411 + self,
412 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
413 + return MambaStateShapeCalculator.kda_state_shape(
414 + self.tp_size,
415 + self.num_heads,
416 + self.head_dim,
417 + conv_kernel_size=self.conv_size,
418 + num_spec=self.num_speculative_tokens,
419 + )
420 +
421 + def __init__(
422 + self,
423 + config: PretrainedConfig,
424 + quant_config: QuantizationConfig | None = None,
425 + layer_id: int = 0,
426 + prefix: str = "attention",
427 + model_config: ModelConfig | None = None,
428 + cache_config: CacheConfig | None = None,
429 + num_speculative_tokens: int = 0,
430 + ) -> None:
431 + super().__init__()
432 + self.tp_size = get_tensor_model_parallel_world_size()
433 + self.hidden_size = config.hidden_size
434 + self.head_dim = config.head_dim
435 + self.num_heads = config.num_attention_heads
436 + self.local_num_heads = self.num_heads // self.tp_size
437 + self.conv_size = config.short_conv_kernel_size
438 + self.layer_id = layer_id
439 + self.prefix = prefix
440 + self.model_config = model_config
441 + self.cache_config = cache_config
442 + self.num_speculative_tokens = num_speculative_tokens
443 + self.safe_gate = getattr(config, "kda_safe_gate", True)
444 + self.lower_bound = getattr(config, "kda_lower_bound", -5.0)
445 + if not getattr(config, "no_kda_lora", False):
446 + raise ValueError("BailingMoeV3 KDA currently expects no_kda_lora=True")
447 +
448 + projection_size = self.head_dim * self.num_heads
449 + self.projection_size_per_partition = projection_size // self.tp_size
450 + self.qkvb_proj = MergedColumnParallelLinear(
451 + self.hidden_size,
452 + [projection_size, projection_size, projection_size, self.num_heads],
453 + bias=False,
454 + quant_config=quant_config,
455 + prefix=f"{prefix}.qkvb_proj",
456 + )
457 + self.f_proj = ColumnParallelLinear(
458 + self.hidden_size,
459 + projection_size,
460 + bias=False,
461 + quant_config=quant_config,
462 + prefix=f"{prefix}.f_proj",
463 + )
464 + self.dt_bias = nn.Parameter(
465 + torch.empty(projection_size // self.tp_size, dtype=torch.float32)
466 + )
467 + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
468 +
469 + self.q_conv1d = ColumnParallelLinear(
470 + input_size=self.conv_size,
471 + output_size=projection_size,
472 + bias=False,
473 + params_dtype=torch.float32,
474 + prefix=f"{prefix}.q_conv1d",
475 + )
476 + self.k_conv1d = ColumnParallelLinear(
477 + input_size=self.conv_size,
478 + output_size=projection_size,
479 + bias=False,
480 + params_dtype=torch.float32,
481 + prefix=f"{prefix}.k_conv1d",
482 + )
483 + self.v_conv1d = ColumnParallelLinear(
484 + input_size=self.conv_size,
485 + output_size=projection_size,
486 + bias=False,
487 + params_dtype=torch.float32,
488 + prefix=f"{prefix}.v_conv1d",
489 + )
490 + self.q_conv1d.weight.data = self.q_conv1d.weight.data.unsqueeze(1)
491 + self.k_conv1d.weight.data = self.k_conv1d.weight.data.unsqueeze(1)
492 + self.v_conv1d.weight.data = self.v_conv1d.weight.data.unsqueeze(1)
493 +
494 + self.A_log = nn.Parameter(
495 + torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32)
496 + )
497 + set_weight_attrs(self.A_log, {"weight_loader": _load_a_log})
498 +
499 + self.g_proj = ColumnParallelLinear(
500 + self.hidden_size,
501 + projection_size,
502 + bias=False,
503 + quant_config=quant_config,
504 + prefix=f"{prefix}.g_proj",
505 + )
506 + self.o_norm = FusedRMSNormGated(
507 + self.head_dim, eps=config.rms_norm_eps, activation="sigmoid"
508 + )
509 + self.o_proj = RowParallelLinear(
510 + projection_size,
511 + self.hidden_size,
512 + bias=False,
513 + quant_config=quant_config,
514 + prefix=f"{prefix}.o_proj",
515 + )
516 +
517 + compilation_config = get_current_vllm_config().compilation_config
518 + if prefix in compilation_config.static_forward_context:
519 + raise ValueError(f"Duplicate layer name: {prefix}")
520 + compilation_config.static_forward_context[prefix] = self
521 + splitting_op = "vllm::bailing_v3_kda_attention"
522 + if (
523 + compilation_config.splitting_ops is not None
524 + and splitting_op not in compilation_config.splitting_ops
525 + ):
526 + compilation_config.splitting_ops.append(splitting_op)
527 +
528 + def forward(
529 + self,
530 + hidden_states: torch.Tensor,
531 + positions: torch.Tensor,
532 + output: torch.Tensor,
533 + ) -> None:
534 + del positions
535 + num_tokens = hidden_states.size(0)
536 + qkvb = self.qkvb_proj(hidden_states)[0]
537 + q, k, v, beta_logits = qkvb.split(
538 + [
539 + self.projection_size_per_partition,
540 + self.projection_size_per_partition,
541 + self.projection_size_per_partition,
542 + self.local_num_heads,
543 + ],
544 + dim=-1,
545 + )
546 + beta = beta_logits.float().sigmoid().unsqueeze(0)
547 + g1 = self.f_proj(hidden_states)[0]
548 + g1 = fused_kda_gate(
549 + g1,
550 + self.A_log,
551 + self.head_dim,
552 + g_bias=self.dt_bias,
553 + lower_bound=self.lower_bound if self.safe_gate else None,
554 + )
555 + g1 = g1.unsqueeze(0)
556 + g2 = rearrange(
557 + self.g_proj(hidden_states)[0], "... (h d) -> ... h d", d=self.head_dim
558 + )
559 +
560 + core_attn_out = torch.zeros(
561 + (1, num_tokens, self.local_num_heads, self.head_dim),
562 + dtype=hidden_states.dtype,
563 + device=hidden_states.device,
564 + )
565 + torch.ops.vllm.bailing_v3_kda_attention(
566 + q, k, v, g1, beta, core_attn_out, self.prefix
567 + )
568 + core_attn_out = self.o_norm(core_attn_out, g2)
569 + core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
570 + output[:] = self.o_proj(core_attn_out)[0]
571 +
572 + def _forward(
573 + self,
574 + q_proj_states: torch.Tensor,
575 + k_proj_states: torch.Tensor,
576 + v_proj_states: torch.Tensor,
577 + g1: torch.Tensor,
578 + beta: torch.Tensor,
579 + core_attn_out: torch.Tensor,
580 + ) -> None:
581 + forward_context = get_forward_context()
582 + attn_metadata: AttentionMetadata = forward_context.attn_metadata
583 + if attn_metadata is None:
584 + return
585 +
586 + assert isinstance(attn_metadata, dict)
587 + attn_metadata = attn_metadata[self.prefix]
588 + assert isinstance(attn_metadata, GDNAttentionMetadata)
589 + has_initial_state = attn_metadata.has_initial_state
590 + spec_query_start_loc = attn_metadata.spec_query_start_loc
591 + query_start_loc = attn_metadata.non_spec_query_start_loc
592 + spec_state_indices = attn_metadata.spec_state_indices_tensor
593 + state_indices = attn_metadata.non_spec_state_indices_tensor
594 + spec_sequence_masks = attn_metadata.spec_sequence_masks
595 + spec_token_indx = attn_metadata.spec_token_indx
596 + non_spec_token_indx = attn_metadata.non_spec_token_indx
597 + num_accepted_tokens = attn_metadata.num_accepted_tokens
598 + num_actual_tokens = attn_metadata.num_actual_tokens
599 + conv_state, recurrent_state = self.kv_cache
600 + recurrent_state_active = recurrent_state[..., : self.head_dim]
601 + if not is_conv_state_dim_first():
602 + conv_state = conv_state.transpose(-1, -2)
603 + conv_state_q, conv_state_k, conv_state_v = conv_state.chunk(3, dim=-2)
604 +
605 + q_proj_states = q_proj_states[:num_actual_tokens]
606 + k_proj_states = k_proj_states[:num_actual_tokens]
607 + v_proj_states = v_proj_states[:num_actual_tokens]
608 + g1 = g1[:, :num_actual_tokens]
609 + beta = beta[:, :num_actual_tokens]
610 +
611 + spec_only = (
612 + spec_sequence_masks is not None
613 + and attn_metadata.num_prefills == 0
614 + and attn_metadata.num_decodes == 0
615 + )
616 + if spec_sequence_masks is not None:
617 + assert spec_token_indx is not None
618 + assert non_spec_token_indx is not None
619 + if spec_only:
620 + q_proj_states_spec = q_proj_states
621 + k_proj_states_spec = k_proj_states
622 + v_proj_states_spec = v_proj_states
623 + g1_spec = g1
624 + beta_spec = beta
625 + q_proj_states_non_spec = None
626 + k_proj_states_non_spec = None
627 + v_proj_states_non_spec = None
628 + g1_non_spec = None
629 + beta_non_spec = None
630 + else:
631 + q_proj_states_spec = q_proj_states.index_select(0, spec_token_indx)
632 + k_proj_states_spec = k_proj_states.index_select(0, spec_token_indx)
633 + v_proj_states_spec = v_proj_states.index_select(0, spec_token_indx)
634 + g1_spec = g1.index_select(1, spec_token_indx)
635 + beta_spec = beta.index_select(1, spec_token_indx)
636 + q_proj_states_non_spec = q_proj_states.index_select(
637 + 0, non_spec_token_indx
638 + )
639 + k_proj_states_non_spec = k_proj_states.index_select(
640 + 0, non_spec_token_indx
641 + )
642 + v_proj_states_non_spec = v_proj_states.index_select(
643 + 0, non_spec_token_indx
644 + )
645 + g1_non_spec = g1.index_select(1, non_spec_token_indx)
646 + beta_non_spec = beta.index_select(1, non_spec_token_indx)
647 + else:
648 + q_proj_states_spec = None
649 + k_proj_states_spec = None
650 + v_proj_states_spec = None
651 + g1_spec = None
652 + beta_spec = None
653 + q_proj_states_non_spec = q_proj_states
654 + k_proj_states_non_spec = k_proj_states
655 + v_proj_states_non_spec = v_proj_states
656 + g1_non_spec = g1
657 + beta_non_spec = beta
658 +
659 + q_conv_weights = self.q_conv1d.weight.view(
660 + self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2)
661 + )
662 + k_conv_weights = self.k_conv1d.weight.view(
663 + self.k_conv1d.weight.size(0), self.k_conv1d.weight.size(2)
664 + )
665 + v_conv_weights = self.v_conv1d.weight.view(
666 + self.v_conv1d.weight.size(0), self.v_conv1d.weight.size(2)
667 + )
668 +
669 + if spec_sequence_masks is not None:
670 + assert q_proj_states_spec is not None
671 + assert k_proj_states_spec is not None
672 + assert v_proj_states_spec is not None
673 + assert spec_state_indices is not None
674 + q_spec = causal_conv1d_update(
675 + q_proj_states_spec,
676 + conv_state_q,
677 + q_conv_weights,
678 + self.q_conv1d.bias,
679 + activation="silu",
680 + conv_state_indices=spec_state_indices[:, 0][
681 + : attn_metadata.num_spec_decodes
682 + ],
683 + num_accepted_tokens=num_accepted_tokens,
684 + query_start_loc=spec_query_start_loc,
685 + max_query_len=spec_state_indices.size(-1),
686 + validate_data=False,
687 + )
688 + k_spec = causal_conv1d_update(
689 + k_proj_states_spec,
690 + conv_state_k,
691 + k_conv_weights,
692 + self.k_conv1d.bias,
693 + activation="silu",
694 + conv_state_indices=spec_state_indices[:, 0][
695 + : attn_metadata.num_spec_decodes
696 + ],
697 + num_accepted_tokens=num_accepted_tokens,
698 + query_start_loc=spec_query_start_loc,
699 + max_query_len=spec_state_indices.size(-1),
700 + validate_data=False,
701 + )
702 + v_spec = causal_conv1d_update(
703 + v_proj_states_spec,
704 + conv_state_v,
705 + v_conv_weights,
706 + self.v_conv1d.bias,
707 + activation="silu",
708 + conv_state_indices=spec_state_indices[:, 0][
709 + : attn_metadata.num_spec_decodes
710 + ],
711 + num_accepted_tokens=num_accepted_tokens,
712 + query_start_loc=spec_query_start_loc,
713 + max_query_len=spec_state_indices.size(-1),
714 + validate_data=False,
715 + )
716 + else:
717 + q_spec = k_spec = v_spec = None
718 +
719 + if attn_metadata.num_prefills > 0:
720 + assert q_proj_states_non_spec is not None
721 + assert k_proj_states_non_spec is not None
722 + assert v_proj_states_non_spec is not None
723 + q = causal_conv1d_fn(
724 + q_proj_states_non_spec.transpose(0, 1),
725 + q_conv_weights,
726 + self.q_conv1d.bias,
727 + activation="silu",
728 + conv_states=conv_state_q,
729 + has_initial_state=has_initial_state,
730 + cache_indices=state_indices,
731 + query_start_loc=query_start_loc,
732 + metadata=attn_metadata,
733 + ).transpose(0, 1)
734 + k = causal_conv1d_fn(
735 + k_proj_states_non_spec.transpose(0, 1),
736 + k_conv_weights,
737 + self.k_conv1d.bias,
738 + activation="silu",
739 + conv_states=conv_state_k,
740 + has_initial_state=has_initial_state,
741 + cache_indices=state_indices,
742 + query_start_loc=query_start_loc,
743 + metadata=attn_metadata,
744 + ).transpose(0, 1)
745 + v = causal_conv1d_fn(
746 + v_proj_states_non_spec.transpose(0, 1),
747 + v_conv_weights,
748 + self.v_conv1d.bias,
749 + activation="silu",
750 + conv_states=conv_state_v,
751 + has_initial_state=has_initial_state,
752 + cache_indices=state_indices,
753 + query_start_loc=query_start_loc,
754 + metadata=attn_metadata,
755 + ).transpose(0, 1)
756 + elif attn_metadata.num_decodes > 0:
757 + assert q_proj_states_non_spec is not None
758 + assert k_proj_states_non_spec is not None
759 + assert v_proj_states_non_spec is not None
760 + assert state_indices is not None
761 + decode_indices = state_indices[:num_actual_tokens]
762 + q = causal_conv1d_update(
763 + q_proj_states_non_spec,
764 + conv_state_q,
765 + q_conv_weights,
766 + self.q_conv1d.bias,
767 + activation="silu",
768 + conv_state_indices=decode_indices,
769 + validate_data=True,
770 + )
771 + k = causal_conv1d_update(
772 + k_proj_states_non_spec,
773 + conv_state_k,
774 + k_conv_weights,
775 + self.k_conv1d.bias,
776 + activation="silu",
777 + conv_state_indices=decode_indices,
778 + validate_data=True,
779 + )
780 + v = causal_conv1d_update(
781 + v_proj_states_non_spec,
782 + conv_state_v,
783 + v_conv_weights,
784 + self.v_conv1d.bias,
785 + activation="silu",
786 + conv_state_indices=decode_indices,
787 + validate_data=True,
788 + )
789 + else:
790 + q = k = v = None
791 +
792 + if q_spec is not None:
793 + q_spec, k_spec, v_spec = map(
794 + lambda x: rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim),
795 + (q_spec, k_spec, v_spec),
796 + )
797 + if q is not None:
798 + q, k, v = map(
799 + lambda x: rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim),
800 + (q, k, v),
801 + )
802 + if spec_sequence_masks is not None:
803 + assert g1_spec is not None
804 + assert beta_spec is not None
805 + assert spec_query_start_loc is not None
806 + assert spec_state_indices is not None
807 + out_spec, _ = fused_recurrent_kda_fwd(
808 + q=q_spec.contiguous(),
809 + k=k_spec.contiguous(),
810 + v=v_spec.contiguous(),
811 + g=g1_spec.contiguous(),
812 + beta=beta_spec.contiguous(),
813 + scale=self.head_dim**-0.5,
814 + initial_state=recurrent_state_active,
815 + inplace_final_state=True,
816 + cu_seqlens=spec_query_start_loc[: attn_metadata.num_spec_decodes + 1],
817 + ssm_state_indices=spec_state_indices,
818 + num_accepted_tokens=num_accepted_tokens,
819 + use_qk_l2norm_in_kernel=True,
820 + )
821 + else:
822 + out_spec = None
823 +
824 + if attn_metadata.num_prefills > 0:
825 + assert q is not None and k is not None and v is not None
826 + assert g1_non_spec is not None and beta_non_spec is not None
827 + assert state_indices is not None
828 + zero_idx = state_indices[~has_initial_state]
829 + recurrent_state[zero_idx] = 0
830 + initial_state = recurrent_state_active[state_indices].contiguous()
831 + out, last_state = chunk_kda(
832 + q=q,
833 + k=k,
834 + v=v,
835 + g=g1_non_spec,
836 + beta=beta_non_spec,
837 + initial_state=initial_state,
838 + output_final_state=True,
839 + use_qk_l2norm_in_kernel=True,
840 + safe_gate=self.safe_gate,
841 + lower_bound=self.lower_bound,
842 + cu_seqlens=query_start_loc,
843 + )
844 + recurrent_state_active[state_indices] = last_state
845 + elif attn_metadata.num_decodes > 0:
846 + assert q is not None and k is not None and v is not None
847 + assert g1_non_spec is not None and beta_non_spec is not None
848 + assert state_indices is not None
849 + out, _ = fused_recurrent_kda(
850 + q=q,
851 + k=k,
852 + v=v,
853 + g=g1_non_spec,
854 + beta=beta_non_spec,
855 + initial_state=recurrent_state_active,
856 + use_qk_l2norm_in_kernel=True,
857 + safe_gate=self.safe_gate,
858 + lower_bound=self.lower_bound,
859 + cu_seqlens=query_start_loc[: attn_metadata.num_decodes + 1],
860 + ssm_state_indices=state_indices,
861 + )
862 + else:
863 + out = None
864 +
865 + if spec_sequence_masks is not None and out is not None:
866 + assert spec_token_indx is not None and non_spec_token_indx is not None
867 + merged_out = torch.empty(
868 + (1, num_actual_tokens, *out_spec.shape[2:]),
869 + dtype=out_spec.dtype,
870 + device=out_spec.device,
871 + )
872 + merged_out.index_copy_(1, spec_token_indx, out_spec)
873 + merged_out.index_copy_(1, non_spec_token_indx, out)
874 + core_attn_out[0, :num_actual_tokens] = merged_out[0, :num_actual_tokens]
875 + elif spec_sequence_masks is not None:
876 + core_attn_out[0, :num_actual_tokens] = out_spec[0, :num_actual_tokens]
877 + else:
878 + core_attn_out[0, :num_actual_tokens] = out[0, :num_actual_tokens]
879 +
880 +
881 +class BailingMoeV3Gate(nn.Module):
882 + def __init__(
883 + self,
884 + config: PretrainedConfig,
885 + params_dtype: torch.dtype | None = None,
886 + ) -> None:
887 + super().__init__()
888 + if params_dtype is None:
889 + params_dtype = torch.float32
890 + self.weight = nn.Parameter(
891 + torch.empty((config.num_experts, config.hidden_size), dtype=params_dtype)
892 + )
893 + self.expert_bias = nn.Parameter(
894 + torch.empty((config.num_experts,), dtype=torch.float32)
895 + )
896 +
897 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
898 + return F.linear(hidden_states.to(self.weight.dtype), self.weight).to(
899 + hidden_states.dtype
900 + )
901 +
902 +
903 +class BailingMoeV3MoE(nn.Module):
904 + def __init__(
905 + self,
906 + config: PretrainedConfig,
907 + quant_config: QuantizationConfig | None = None,
908 + layer_id: int = 0,
909 + prefix: str = "",
910 + ) -> None:
911 + super().__init__()
912 + self.tp_size = get_tensor_model_parallel_world_size()
913 + self.num_experts = config.num_experts
914 + self.top_k = config.num_experts_per_tok
915 + self.hidden_size = config.hidden_size
916 + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
917 + self.gate = BailingMoeV3Gate(config)
918 + self.expert_swiglu_limit = _get_layer_swiglu_limit(
919 + getattr(config, "expert_swiglu_limit_list", None), layer_id
920 + )
921 + self.share_expert_swiglu_limit = _get_layer_swiglu_limit(
922 + getattr(config, "share_expert_swiglu_limit_list", None), layer_id
923 + )
924 + shared_intermediate = (
925 + config.moe_shared_expert_intermediate_size * config.num_shared_experts
926 + )
927 + self.shared_experts = BailingMoeV3MLP(
928 + intermediate_size=shared_intermediate,
929 + config=config,
930 + quant_config=quant_config,
931 + reduce_results=False,
932 + prefix=f"{prefix}.shared_experts",
933 + swiglu_limit=self.share_expert_swiglu_limit,
934 + )
935 + self.experts = FusedMoEFactory(
936 + shared_experts=self.shared_experts,
937 + num_experts=self.num_experts,
938 + top_k=self.top_k,
939 + hidden_size=self.hidden_size,
940 + intermediate_size=config.moe_intermediate_size,
941 + renormalize=True,
942 + quant_config=quant_config,
943 + activation="silu",
944 + swiglu_limit=self.expert_swiglu_limit,
945 + prefix=f"{prefix}.experts",
946 + scoring_func="sigmoid",
947 + routed_scaling_factor=self.routed_scaling_factor,
948 + e_score_correction_bias=self.gate.expert_bias,
949 + num_expert_group=config.n_group,
950 + topk_group=config.topk_group,
951 + use_grouped_topk=True,
952 + router_logits_dtype=torch.float32,
953 + )
954 +
955 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
956 + num_tokens, hidden_size = hidden_states.shape
957 + hidden_states = hidden_states.contiguous().view(-1, hidden_size)
958 + router_logits = self.gate(hidden_states.to(torch.float32))
959 + hidden_states = self.experts(
960 + hidden_states=hidden_states, router_logits=router_logits
961 + )
962 + return hidden_states.view(num_tokens, hidden_size)
963 +
964 +
965 +class BailingMoeV3DecoderLayer(nn.Module):
966 + def __init__(
967 + self,
968 + config: PretrainedConfig,
969 + quant_config: QuantizationConfig | None = None,
970 + layer_id: int = 0,
971 + prefix: str = "layer",
972 + model_config: ModelConfig | None = None,
973 + cache_config: CacheConfig | None = None,
974 + num_speculative_tokens: int = 0,
975 + ) -> None:
976 + super().__init__()
977 + self.layer_id = layer_id
978 + self.hidden_size = config.hidden_size
979 + self.is_kda_layer = _is_kda_layer(
980 + layer_id, config.layer_group_size, config.num_hidden_layers
981 + )
982 + if self.is_kda_layer:
983 + self.self_attn = BailingMoeV3KimiDeltaAttention(
984 + config,
985 + quant_config=quant_config,
986 + layer_id=layer_id,
987 + prefix=f"{prefix}.self_attn",
988 + model_config=model_config,
989 + cache_config=cache_config,
990 + num_speculative_tokens=num_speculative_tokens,
991 + )
992 + else:
993 + self.self_attn = BailingMoeV3MLAAttention(
994 + config,
995 + quant_config=quant_config,
996 + layer_id=layer_id,
997 + prefix=f"{prefix}.self_attn",
998 + cache_config=cache_config,
999 + )
1000 +
1001 + if config.num_experts is not None and layer_id >= config.first_k_dense_replace:
1002 + self.mlp = BailingMoeV3MoE(
1003 + config,
1004 + quant_config=quant_config,
1005 + layer_id=layer_id,
1006 + prefix=f"{prefix}.mlp",
1007 + )
1008 + else:
1009 + self.mlp = BailingMoeV3MLP(
1010 + intermediate_size=config.intermediate_size,
1011 + config=config,
1012 + quant_config=quant_config,
1013 + reduce_results=True,
1014 + prefix=f"{prefix}.mlp",
1015 + )
1016 +
1017 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1018 + self.post_attention_layernorm = RMSNorm(
1019 + config.hidden_size, eps=config.rms_norm_eps
1020 + )
1021 +
1022 + def forward(
1023 + self,
1024 + hidden_states: torch.Tensor,
1025 + positions: torch.Tensor,
1026 + residual: torch.Tensor | None = None,
1027 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
1028 + if residual is None:
1029 + residual = hidden_states
1030 + hidden_states = self.input_layernorm(hidden_states)
1031 + else:
1032 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
1033 +
1034 + if self.is_kda_layer:
1035 + attn_output = torch.zeros_like(hidden_states)
1036 + self.self_attn(hidden_states, positions, attn_output)
1037 + else:
1038 + attn_output = self.self_attn(hidden_states, positions)
1039 +
1040 + hidden_states, residual = self.post_attention_layernorm(attn_output, residual)
1041 + hidden_states = self.mlp(hidden_states)
1042 + return hidden_states, residual
1043 +
1044 +
1045 +@support_torch_compile(
1046 + dynamic_arg_dims={
1047 + "input_ids": 0,
1048 + "positions": -1,
1049 + "intermediate_tensors": 0,
1050 + "inputs_embeds": 0,
1051 + }
1052 +)
1053 +class BailingMoeV3Model(nn.Module):
1054 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
1055 + super().__init__()
1056 + config = vllm_config.model_config.hf_config
1057 + model_config = vllm_config.model_config
1058 + quant_config = vllm_config.quant_config
1059 + cache_config = vllm_config.cache_config
1060 + num_speculative_tokens = (
1061 + vllm_config.speculative_config.num_speculative_tokens
1062 + if vllm_config.speculative_config
1063 + else 0
1064 + )
1065 +
1066 + self.config = config
1067 + self.vocab_size = config.vocab_size
1068 + self.embed_dim = config.hidden_size
1069 + self.num_layers = config.num_hidden_layers
1070 +
1071 + if get_pp_group().is_first_rank:
1072 + self.word_embeddings = VocabParallelEmbedding(
1073 + self.vocab_size,
1074 + self.embed_dim,
1075 + org_num_embeddings=self.vocab_size,
1076 + )
1077 + else:
1078 + self.word_embeddings = PPMissingLayer()
1079 +
1080 + def layer_fn(prefix: str):
1081 + layer_idx = int(prefix.split(".")[-1])
1082 + return BailingMoeV3DecoderLayer(
1083 + config=config,
1084 + quant_config=quant_config,
1085 + layer_id=layer_idx,
1086 + prefix=prefix,
1087 + model_config=model_config,
1088 + cache_config=cache_config,
1089 + num_speculative_tokens=num_speculative_tokens,
1090 + )
1091 +
1092 + self.start_layer, self.end_layer, self.layers = make_layers(
1093 + self.num_layers, layer_fn, prefix=f"{prefix}.layers"
1094 + )
1095 +
1096 + if get_pp_group().is_last_rank:
1097 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1098 + else:
1099 + self.norm = PPMissingLayer()
1100 +
1101 + @property
1102 + def embed_tokens(self) -> nn.Module:
1103 + return self.word_embeddings
1104 +
1105 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1106 + return self.word_embeddings(input_ids)
1107 +
1108 + def forward(
1109 + self,
1110 + input_ids: torch.Tensor | None,
1111 + positions: torch.Tensor,
1112 + intermediate_tensors: IntermediateTensors | None = None,
1113 + inputs_embeds: torch.Tensor | None = None,
1114 + ) -> torch.Tensor | IntermediateTensors:
1115 + if get_pp_group().is_first_rank:
1116 + hidden_states = (
1117 + self.word_embeddings(input_ids)
1118 + if inputs_embeds is None
1119 + else inputs_embeds
1120 + )
1121 + residual = None
1122 + else:
1123 + assert intermediate_tensors is not None
1124 + hidden_states = intermediate_tensors["hidden_states"]
1125 + residual = intermediate_tensors["residual"]
1126 +
1127 + for layer in self.layers[self.start_layer : self.end_layer]:
1128 + hidden_states, residual = layer(hidden_states, positions, residual)
1129 +
1130 + if not get_pp_group().is_last_rank:
1131 + return IntermediateTensors(
1132 + {"hidden_states": hidden_states, "residual": residual}
1133 + )
1134 + if residual is not None:
1135 + hidden_states, _ = self.norm(hidden_states, residual)
1136 + else:
1137 + hidden_states = self.norm(hidden_states)
1138 + return hidden_states
1139 +
1140 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1141 + return fused_moe_make_expert_params_mapping(
1142 + self,
1143 + ckpt_gate_proj_name="gate_proj",
1144 + ckpt_down_proj_name="down_proj",
1145 + ckpt_up_proj_name="up_proj",
1146 + num_experts=self.config.num_experts,
1147 + )
1148 +
1149 +
1150 +class BailingMoeV3ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):
1151 + packed_modules_mapping = {
1152 + "gate_up_proj": ["gate_proj", "up_proj"],
1153 + "qkvb_proj": ["q_proj", "k_proj", "v_proj", "b_proj"],
1154 + }
1155 +
1156 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
1157 + super().__init__()
1158 + config = vllm_config.model_config.hf_config
1159 + quant_config = vllm_config.quant_config
1160 + self.config = config
1161 + self.quant_config = quant_config
1162 + self.model = BailingMoeV3Model(
1163 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
1164 + )
1165 + if get_pp_group().is_last_rank:
1166 + self.lm_head = ParallelLMHead(
1167 + config.vocab_size,
1168 + config.hidden_size,
1169 + quant_config=quant_config,
1170 + prefix=maybe_prefix(prefix, "lm_head"),
1171 + )
1172 + self.logits_processor = LogitsProcessor(config.vocab_size)
1173 + else:
1174 + self.lm_head = PPMissingLayer()
1175 +
1176 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1177 + return self.model.embed_input_ids(input_ids)
1178 +
1179 + def forward(
1180 + self,
1181 + input_ids: torch.Tensor | None,
1182 + positions: torch.Tensor,
1183 + intermediate_tensors: IntermediateTensors | None = None,
1184 + inputs_embeds: torch.Tensor | None = None,
1185 + ) -> torch.Tensor | IntermediateTensors:
1186 + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
1187 +
1188 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
1189 + return self.logits_processor(self.lm_head, hidden_states)
1190 +
1191 + def make_empty_intermediate_tensors(
1192 + self, batch_size: int, dtype: torch.dtype, device: torch.device
1193 + ) -> IntermediateTensors:
1194 + return IntermediateTensors(
1195 + {
1196 + "hidden_states": torch.zeros(
1197 + (batch_size, self.config.hidden_size), dtype=dtype, device=device
1198 + ),
1199 + "residual": torch.zeros(
1200 + (batch_size, self.config.hidden_size), dtype=dtype, device=device
1201 + ),
1202 + }
1203 + )
1204 +
1205 + @classmethod
1206 + def get_mamba_state_shape_from_config(
1207 + cls, vllm_config: VllmConfig
1208 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
1209 + return _get_kda_state_shape_for_config(vllm_config)
1210 +
1211 + @classmethod
1212 + def get_mamba_state_dtype_from_config(
1213 + cls, vllm_config: VllmConfig
1214 + ) -> tuple[torch.dtype, torch.dtype, torch.dtype, torch.dtype]:
1215 + return MambaStateDtypeCalculator.kda_state_dtype(
1216 + vllm_config.model_config.dtype,
1217 + vllm_config.cache_config.mamba_cache_dtype,
1218 + )
1219 +
1220 + @classmethod
1221 + def get_mamba_state_copy_func(cls) -> tuple:
1222 + return MambaStateCopyFuncCalculator.kda_state_copy_func()
1223 +
1224 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1225 + return self.model.get_expert_mapping()
1226 +
1227 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1228 + params_dict = dict(self.named_parameters(remove_duplicate=False))
1229 + loaded_params: set[str] = set()
1230 + stacked_mappings = [
1231 + (".gate_up_proj", ".gate_proj", 0),
1232 + (".gate_up_proj", ".up_proj", 1),
1233 + (".qkvb_proj", ".q_proj", 0),
1234 + (".qkvb_proj", ".k_proj", 1),
1235 + (".qkvb_proj", ".v_proj", 2),
1236 + (".qkvb_proj", ".b_proj", 3),
1237 + (".fused_qkv_a_proj", ".q_a_proj", 0),
1238 + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
1239 + ]
1240 + expert_mappings = list(self.get_expert_mapping())
1241 +
1242 + def load_param(name: str, tensor: torch.Tensor, shard_id=None) -> bool:
1243 + if name not in params_dict or is_pp_missing_parameter(name, self):
1244 + return False
1245 + param = params_dict[name]
1246 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
1247 + if shard_id is None:
1248 + weight_loader(param, tensor)
1249 + elif isinstance(shard_id, int):
1250 + weight_loader(param, tensor, shard_id)
1251 + else:
1252 + weight_loader(
1253 + param, tensor, name, expert_id=shard_id[0], shard_id=shard_id[1]
1254 + )
1255 + loaded_params.add(name)
1256 + return True
1257 +
1258 + def normalize_name(name: str) -> str | None:
1259 + if name.startswith("model.layers."):
1260 + layer_idx = int(name.split("model.layers.")[1].split(".")[0])
1261 + if layer_idx >= self.config.num_hidden_layers:
1262 + return None
1263 + name = name.replace("attention.", "self_attn.")
1264 + return maybe_remap_kv_scale_name(name, params_dict)
1265 +
1266 + for orig_name, weight in weights:
1267 + name = normalize_name(orig_name)
1268 + if name is None:
1269 + continue
1270 + loaded = False
1271 + for param_suf, weight_suf, shard_id in stacked_mappings:
1272 + if weight_suf in name:
1273 + mapped = name.replace(weight_suf, param_suf)
1274 + if load_param(mapped, weight, shard_id):
1275 + loaded = True
1276 + break
1277 + if loaded:
1278 + continue
1279 +
1280 + if ".mlp.experts." in name:
1281 + for param_name, weight_name, expert_id, shard_id in expert_mappings:
1282 + if weight_name in name:
1283 + mapped = name.replace(weight_name, param_name)
1284 + load_param(mapped, weight, (expert_id, shard_id))
1285 + break
1286 + continue
1287 +
1288 + load_param(name, weight)
1289 + return loaded_params
vllm/model_executor/models/bailing_moe_v3_mtp.py new
+389
@@ -0,0 +1,389 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only Bailing MoE v3 MTP model."""
4 +
5 +from collections.abc import Iterable
6 +
7 +import torch
8 +import torch.nn as nn
9 +from transformers.configuration_utils import PretrainedConfig
10 +
11 +from vllm.compilation.decorators import support_torch_compile
12 +from vllm.config import VllmConfig
13 +from vllm.distributed import get_pp_group
14 +from vllm.model_executor.layers.fused_moe import (
15 + fused_moe_make_expert_params_mapping,
16 +)
17 +from vllm.model_executor.layers.layernorm import RMSNorm
18 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
19 +from vllm.model_executor.layers.vocab_parallel_embedding import (
20 + ParallelLMHead,
21 + VocabParallelEmbedding,
22 +)
23 +from vllm.model_executor.model_loader.weight_utils import (
24 + default_weight_loader,
25 + maybe_remap_kv_scale_name,
26 +)
27 +from vllm.model_executor.models.bailing_moe_v3 import (
28 + BailingMoeV3MLAAttention,
29 + BailingMoeV3MoE,
30 +)
31 +from vllm.model_executor.models.interfaces import SupportsPP
32 +from vllm.model_executor.models.utils import (
33 + PPMissingLayer,
34 + is_pp_missing_parameter,
35 + maybe_prefix,
36 +)
37 +from vllm.sequence import IntermediateTensors
38 +
39 +
40 +def _get_draft_hf_config(vllm_config: VllmConfig) -> PretrainedConfig:
41 + speculative_config = vllm_config.speculative_config
42 + if speculative_config is not None:
43 + draft_model_config = speculative_config.draft_model_config
44 + if draft_model_config is not None:
45 + return draft_model_config.hf_config
46 + return vllm_config.model_config.hf_config
47 +
48 +
49 +class BailingMoeV3MTPSharedHead(nn.Module):
50 + def __init__(
51 + self,
52 + config: PretrainedConfig,
53 + prefix: str,
54 + vllm_config: VllmConfig,
55 + ) -> None:
56 + super().__init__()
57 + self.head = ParallelLMHead(
58 + config.vocab_size,
59 + config.hidden_size,
60 + quant_config=vllm_config.quant_config,
61 + prefix=maybe_prefix(prefix, "head"),
62 + )
63 +
64 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
65 + return hidden_states
66 +
67 +
68 +class BailingMoeV3MultiTokenPredictorLayer(nn.Module):
69 + def __init__(
70 + self,
71 + *,
72 + vllm_config: VllmConfig,
73 + prefix: str,
74 + layer_id: int,
75 + ) -> None:
76 + super().__init__()
77 + config = _get_draft_hf_config(vllm_config)
78 + self.config = config
79 + self.layer_id = layer_id
80 + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
81 + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
82 + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
83 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
84 + self.self_attn = BailingMoeV3MLAAttention(
85 + config,
86 + quant_config=vllm_config.quant_config,
87 + layer_id=layer_id,
88 + prefix=maybe_prefix(prefix, "self_attn"),
89 + cache_config=vllm_config.cache_config,
90 + )
91 + self.post_attention_layernorm = RMSNorm(
92 + config.hidden_size, eps=config.rms_norm_eps
93 + )
94 + self.mlp = BailingMoeV3MoE(
95 + config,
96 + quant_config=vllm_config.quant_config,
97 + layer_id=layer_id,
98 + prefix=maybe_prefix(prefix, "mlp"),
99 + )
100 + self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
101 + self.shared_head = BailingMoeV3MTPSharedHead(
102 + config,
103 + maybe_prefix(prefix, "shared_head"),
104 + vllm_config,
105 + )
106 +
107 + def forward(
108 + self,
109 + input_ids: torch.Tensor,
110 + positions: torch.Tensor,
111 + previous_hidden_states: torch.Tensor,
112 + inputs_embeds: torch.Tensor | None = None,
113 + spec_step_index: int = 0,
114 + ) -> torch.Tensor:
115 + del input_ids, spec_step_index
116 + assert inputs_embeds is not None
117 + inputs_embeds = self.enorm(inputs_embeds)
118 + previous_hidden_states = self.hnorm(previous_hidden_states)
119 +
120 + hidden_states = self.eh_proj(
121 + torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
122 + )
123 + residual = hidden_states
124 + hidden_states = self.input_layernorm(hidden_states)
125 +
126 + hidden_states = self.self_attn(hidden_states, positions)
127 + hidden_states = residual + hidden_states
128 +
129 + residual = hidden_states
130 + hidden_states = self.post_attention_layernorm(hidden_states)
131 + hidden_states = self.mlp(hidden_states)
132 + hidden_states = residual + hidden_states.to(residual.device)
133 + return self.final_layernorm(hidden_states)
134 +
135 +
136 +class BailingMoeV3MultiTokenPredictor(nn.Module):
137 + def __init__(
138 + self,
139 + *,
140 + vllm_config: VllmConfig,
141 + prefix: str = "",
142 + ) -> None:
143 + super().__init__()
144 + config = _get_draft_hf_config(vllm_config)
145 + self.mtp_start_layer_idx = config.num_hidden_layers
146 + self.num_mtp_layers = config.num_nextn_predict_layers
147 + self.layers = nn.ModuleDict(
148 + {
149 + str(idx): BailingMoeV3MultiTokenPredictorLayer(
150 + vllm_config=vllm_config,
151 + prefix=f"{prefix}.layers.{idx}",
152 + layer_id=idx,
153 + )
154 + for idx in range(
155 + self.mtp_start_layer_idx,
156 + self.mtp_start_layer_idx + self.num_mtp_layers,
157 + )
158 + }
159 + )
160 + if get_pp_group().is_first_rank:
161 + self.embed_tokens = VocabParallelEmbedding(
162 + config.vocab_size,
163 + config.hidden_size,
164 + org_num_embeddings=config.vocab_size,
165 + prefix=maybe_prefix(prefix, "embed_tokens"),
166 + )
167 + else:
168 + self.embed_tokens = PPMissingLayer()
169 + self.logits_processor = LogitsProcessor(config.vocab_size)
170 +
171 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
172 + return self.embed_tokens(input_ids)
173 +
174 + def forward(
175 + self,
176 + input_ids: torch.Tensor,
177 + positions: torch.Tensor,
178 + previous_hidden_states: torch.Tensor,
179 + inputs_embeds: torch.Tensor | None = None,
180 + spec_step_idx: int = 0,
181 + ) -> torch.Tensor:
182 + if inputs_embeds is None:
183 + inputs_embeds = self.embed_tokens(input_ids)
184 + current_step_idx = spec_step_idx % self.num_mtp_layers
185 + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
186 + input_ids,
187 + positions,
188 + previous_hidden_states,
189 + inputs_embeds,
190 + current_step_idx,
191 + )
192 +
193 + def compute_logits(
194 + self,
195 + hidden_states: torch.Tensor,
196 + spec_step_idx: int = 0,
197 + lm_head: nn.Module | None = None,
198 + ) -> torch.Tensor:
199 + current_step_idx = spec_step_idx % self.num_mtp_layers
200 + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
201 + head = lm_head if lm_head is not None else mtp_layer.shared_head.head
202 + return self.logits_processor(head, mtp_layer.shared_head(hidden_states))
203 +
204 +
205 +@support_torch_compile
206 +class BailingMoeV3MTPModel(nn.Module, SupportsPP):
207 + packed_modules_mapping = {
208 + "gate_up_proj": ["gate_proj", "up_proj"],
209 + "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
210 + }
211 +
212 + def __init__(
213 + self,
214 + *,
215 + vllm_config: VllmConfig,
216 + prefix: str = "",
217 + ) -> None:
218 + super().__init__()
219 + self.config = _get_draft_hf_config(vllm_config)
220 + self.model = BailingMoeV3MultiTokenPredictor(
221 + vllm_config=vllm_config,
222 + prefix=maybe_prefix(prefix, "model"),
223 + )
224 + self.lm_head: nn.Module | None = None
225 +
226 + def share_lm_head(self, lm_head: nn.Module) -> None:
227 + self.lm_head = lm_head
228 +
229 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
230 + return self.model.embed_input_ids(input_ids)
231 +
232 + def forward(
233 + self,
234 + input_ids: torch.Tensor | None,
235 + positions: torch.Tensor,
236 + hidden_states: torch.Tensor,
237 + intermediate_tensors: IntermediateTensors | None = None,
238 + inputs_embeds: torch.Tensor | None = None,
239 + spec_step_idx: int = 0,
240 + ) -> torch.Tensor:
241 + del intermediate_tensors
242 + return self.model(
243 + input_ids,
244 + positions,
245 + hidden_states,
246 + inputs_embeds,
247 + spec_step_idx,
248 + )
249 +
250 + def compute_logits(
251 + self,
252 + hidden_states: torch.Tensor,
253 + spec_step_idx: int = 0,
254 + ) -> torch.Tensor:
255 + return self.model.compute_logits(hidden_states, spec_step_idx, self.lm_head)
256 +
257 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
258 + return fused_moe_make_expert_params_mapping(
259 + self,
260 + ckpt_gate_proj_name="gate_proj",
261 + ckpt_down_proj_name="down_proj",
262 + ckpt_up_proj_name="up_proj",
263 + num_experts=self.config.num_experts,
264 + )
265 +
266 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
267 + stacked_params_mapping = [
268 + (".fused_qkv_a_proj", ".q_a_proj", 0),
269 + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
270 + (".gate_up_proj", ".gate_proj", 0),
271 + (".gate_up_proj", ".up_proj", 1),
272 + ]
273 + expert_params_mapping = list(self.get_expert_mapping())
274 + params_dict = dict(self.named_parameters(remove_duplicate=False))
275 + loaded_params: set[str] = set()
276 + loaded_mtp_layers: set[int] = set()
277 +
278 + def load_param(
279 + name: str,
280 + loaded_weight: torch.Tensor,
281 + shard_id=None,
282 + ) -> bool:
283 + name = maybe_remap_kv_scale_name(name, params_dict)
284 + if name is None:
285 + return False
286 + if name not in params_dict or is_pp_missing_parameter(name, self):
287 + return False
288 +
289 + param = params_dict[name]
290 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
291 + if shard_id is None:
292 + weight_loader(param, loaded_weight)
293 + elif isinstance(shard_id, int):
294 + weight_loader(param, loaded_weight, shard_id)
295 + else:
296 + weight_loader(
297 + param,
298 + loaded_weight,
299 + name,
300 + expert_id=shard_id[0],
301 + shard_id=shard_id[1],
302 + )
303 + loaded_params.add(name)
304 + return True
305 +
306 + def get_spec_layer_idx(name: str) -> int | None:
307 + if not name.startswith("model.layers."):
308 + return None
309 + try:
310 + layer_idx = int(name.split("model.layers.", 1)[1].split(".", 1)[0])
311 + except (IndexError, ValueError):
312 + return None
313 + mtp_idx = layer_idx - self.config.num_hidden_layers
314 + if 0 <= mtp_idx < self.config.num_nextn_predict_layers:
315 + return layer_idx
316 + return None
317 +
318 + def normalize_name(name: str) -> str:
319 + name = name.replace(".attention.", ".self_attn.")
320 + return name.replace(
321 + "mlp.gate.e_score_correction_bias",
322 + "mlp.gate.expert_bias",
323 + )
324 +
325 + def load_lm_head(loaded_weight: torch.Tensor) -> None:
326 + for layer_idx in range(
327 + self.model.mtp_start_layer_idx,
328 + self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
329 + ):
330 + name = f"model.layers.{layer_idx}.shared_head.head.weight"
331 + load_param(name, loaded_weight)
332 +
333 + for name, loaded_weight in weights:
334 + if "rotary_emb.inv_freq" in name:
335 + continue
336 +
337 + if name == "model.word_embeddings.weight":
338 + load_param("model.embed_tokens.weight", loaded_weight)
339 + continue
340 + if name == "lm_head.weight":
341 + load_lm_head(loaded_weight)
342 + continue
343 +
344 + spec_layer = get_spec_layer_idx(name)
345 + if spec_layer is None:
346 + continue
347 + name = normalize_name(name)
348 +
349 + loaded = False
350 + for param_name, weight_name, shard_id in stacked_params_mapping:
351 + if weight_name not in name:
352 + continue
353 + if "mlp.experts." in name and name not in params_dict:
354 + continue
355 + mapped_name = name.replace(weight_name, param_name)
356 + if load_param(mapped_name, loaded_weight, shard_id):
357 + loaded = True
358 + break
359 + if loaded:
360 + loaded_mtp_layers.add(spec_layer)
361 + continue
362 +
363 + if "mlp.experts" in name:
364 + for mapping in expert_params_mapping:
365 + param_name, weight_name, expert_id, shard_id = mapping
366 + if weight_name not in name:
367 + continue
368 + mapped_name = name.replace(weight_name, param_name)
369 + if load_param(mapped_name, loaded_weight, (expert_id, shard_id)):
370 + loaded = True
371 + break
372 + if loaded:
373 + loaded_mtp_layers.add(spec_layer)
374 + continue
375 +
376 + if load_param(name, loaded_weight):
377 + loaded_mtp_layers.add(spec_layer)
378 +
379 + for layer_idx in range(
380 + self.model.mtp_start_layer_idx,
381 + self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
382 + ):
383 + if layer_idx not in loaded_mtp_layers:
384 + raise ValueError(
385 + f"Bailing V3 MTP speculative decoding layer {layer_idx} "
386 + "weights are missing from checkpoint. Use a checkpoint that "
387 + "includes MTP layer weights, or disable speculative decoding."
388 + )
389 + return loaded_params
vllm/model_executor/models/registry.py
+2
@@ -79,6 +79,7 @@ _TEXT_GENERATION_MODELS = {
79 "BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"),
80 "BailingMoeV2ForCausalLM": ("bailing_moe", "BailingMoeV2ForCausalLM"),
81 "BailingMoeV2_5ForCausalLM": ("bailing_moe_linear", "BailingMoeV25ForCausalLM"),
82 + "BailingMoeV3ForCausalLM": ("bailing_moe_v3", "BailingMoeV3ForCausalLM"),
83 "BloomForCausalLM": ("bloom", "BloomForCausalLM"),
84 "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"),
85 "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"),
@@ -641,6 +642,7 @@ _SPECULATIVE_DECODING_MODELS = {
642 "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
643 "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
644 "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"),
645 + "BailingMoeV3MTPModel": ("bailing_moe_v3_mtp", "BailingMoeV3MTPModel"),
646 "MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"),
647 "BailingMoeV25MTPModel": ("bailing_moe_mtp", "BailingMoeV25MTPModel"),
648 "InklingMTPModel": ("vllm.models.inkling", "InklingMTP"),
vllm/parser/ling3.py new
+83
@@ -0,0 +1,83 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Ling3 parser for reasoning and tool calls.
4 +
5 +Ling3 uses the same XML tool-call format as GLM-4.7 and defaults thinking on.
6 +It follows the GLM-style ``<think>`` / ``</think>`` reasoning format and
7 +treats ``<tool_call>`` as an implicit reasoning terminator.
8 +"""
9 +
10 +from __future__ import annotations
11 +
12 +from dataclasses import replace
13 +from typing import TYPE_CHECKING
14 +
15 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
16 +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
17 +from vllm.parser.engine.adapters import make_adapters
18 +from vllm.parser.engine.parser_engine import ParserEngine
19 +from vllm.parser.glm47_moe import (
20 + THINK_END,
21 + THINK_START,
22 + Glm47MoeParser,
23 + glm47_moe_config,
24 +)
25 +
26 +if TYPE_CHECKING:
27 + from vllm.tokenizers import TokenizerLike
28 + from vllm.tool_parsers.abstract_tool_parser import Tool
29 +
30 +
31 +class Ling3Parser(Glm47MoeParser):
32 + """Ling3 parser backed by the GLM XML parser engine."""
33 +
34 + def __init__(
35 + self,
36 + tokenizer: TokenizerLike,
37 + tools: list[Tool] | None = None,
38 + **kwargs,
39 + ) -> None:
40 + chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
41 + thinking = chat_kwargs.get("thinking", None)
42 + enable_thinking = chat_kwargs.get("enable_thinking", None)
43 + self.thinking_enabled = (
44 + True
45 + if thinking is None and enable_thinking is None
46 + else bool(thinking) or bool(enable_thinking)
47 + )
48 + parser_config = replace(
49 + glm47_moe_config(thinking=self.thinking_enabled),
50 + name="ling3",
51 + )
52 + kwargs.setdefault(
53 + "parser_engine_config",
54 + parser_config,
55 + )
56 + ParserEngine.__init__(self, tokenizer, tools, **kwargs)
57 +
58 + @property
59 + def reasoning_start_str(self) -> str:
60 + return THINK_START
61 +
62 + @property
63 + def reasoning_end_str(self) -> str:
64 + return THINK_END
65 +
66 + def extract_reasoning(
67 + self,
68 + model_output: str,
69 + request: ChatCompletionRequest | ResponsesRequest,
70 + ) -> tuple[str | None, str | None]:
71 + if not self.thinking_enabled:
72 + return None, model_output
73 +
74 + reasoning, content = super().extract_reasoning(model_output, request)
75 + if reasoning and not content and "<tool_call>" not in model_output:
76 + return None, reasoning
77 + return reasoning, content
78 +
79 +
80 +(
81 + Ling3ParserReasoningAdapter,
82 + Ling3ParserToolAdapter,
83 +) = make_adapters(Ling3Parser)
vllm/reasoning/__init__.py
+4
@@ -60,6 +60,10 @@ _REASONING_PARSERS_TO_REGISTER = {
60 "glm47_moe_reasoning_parser",
61 "Glm47MoeParserReasoningAdapter",
62 ),
63 + "ling3": (
64 + "ling3_reasoning_parser",
65 + "Ling3ParserReasoningAdapter",
66 + ),
67 "openai_gptoss": (
68 "gptoss_reasoning_parser",
69 "GptOssReasoningParser",
vllm/reasoning/ling3_reasoning_parser.py new
+6
@@ -0,0 +1,6 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.parser.ling3 import Ling3ParserReasoningAdapter
5 +
6 +__all__ = ["Ling3ParserReasoningAdapter"]
vllm/tool_parsers/__init__.py
+4
@@ -58,6 +58,10 @@ _TOOL_PARSERS_TO_REGISTER = {
58 "glm47_moe_tool_parser",
59 "Glm47MoeModelToolParser",
60 ),
61 + "ling3": (
62 + "ling3_tool_parser",
63 + "Ling3ToolParser",
64 + ),
65 "granite-20b-fc": (
66 "granite_20b_fc_tool_parser",
67 "Granite20bFCToolParser",
vllm/tool_parsers/ling3_tool_parser.py new
+11
@@ -0,0 +1,11 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from __future__ import annotations
5 +
6 +from vllm.parser.ling3 import Ling3ParserToolAdapter
7 +
8 +
9 +class Ling3ToolParser(Ling3ParserToolAdapter): # type: ignore[valid-type, misc]
10 + supports_required_and_named = False
11 + structural_tag_model = "glm_4_7"
vllm/transformers_utils/model_arch_config_convertor.py
+8
@@ -289,6 +289,7 @@ class ModelArchConfigConvertorBase:
289 "pangu_ultra_moe_mtp",
290 "bailing_hybrid",
291 "bailing_hybrid_mtp",
292 + "bailing_hybrid_v3_mtp",
293 ):
294 # check is deepseek_v4 model
295 if hasattr(self.hf_text_config, "compress_ratios"):
@@ -571,6 +572,12 @@ class BailingHybridMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
572 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
573
574
575 +class BailingHybridV3MTPModelArchConfigConvertor(
576 + BailingHybridMTPModelArchConfigConvertor
577 +):
578 + pass
579 +
580 +
581 class Qwen3_5MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
582 def get_num_hidden_layers(self) -> int:
583 return getattr(self.hf_text_config, "mtp_num_hidden_layers", 0)
@@ -693,6 +700,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = {
700 "moss_audio": MossAudioModelArchConfigConvertor,
701 "mpt": MPTModelArchConfigConvertor,
702 "nemotron-nas": NemotronNasModelArchConfigConvertor,
703 + "bailing_hybrid_v3_mtp": BailingHybridV3MTPModelArchConfigConvertor,
704 "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor,
705 "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor,
706 "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor,