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

[Feature] Support MTP speculative decoding for Bailing hybrid models (#44880)

Signed-off-by: zc02384840 <zc02384840@antgroup.com> Co-authored-by: zc02384840 <zc02384840@antgroup.com>

alex101-ops committed Jul 6, 2026 at 10:38 UTC f2aaf5915102cd56b3a60f8e6e59c4a7f31268dd
11 files changed +1200 -6
tests/config/test_bailing_mtp_config.py new
+52
@@ -0,0 +1,52 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from transformers import PretrainedConfig
5 +
6 +from vllm.config.speculative import MTPModelTypes, SpeculativeConfig
7 +from vllm.transformers_utils.model_arch_config_convertor import (
8 + BailingHybridMTPModelArchConfigConvertor,
9 +)
10 +
11 +
12 +def _bailing_config() -> PretrainedConfig:
13 + config = PretrainedConfig(
14 + architectures=["BailingMoeV2_5ForCausalLM"],
15 + hidden_size=4096,
16 + kv_lora_rank=512,
17 + num_attention_heads=32,
18 + num_experts=256,
19 + num_hidden_layers=32,
20 + num_key_value_heads=32,
21 + num_nextn_predict_layers=1,
22 + qk_rope_head_dim=64,
23 + vocab_size=157184,
24 + )
25 + config.model_type = "bailing_hybrid"
26 + return config
27 +
28 +
29 +def test_bailing_hybrid_mtp_hf_config_override():
30 + config = _bailing_config()
31 +
32 + overridden = SpeculativeConfig.hf_config_override(config)
33 +
34 + assert overridden.model_type == "bailing_hybrid_mtp"
35 + assert overridden.architectures == ["BailingMoeV25MTPModel"]
36 + assert overridden.n_predict == 1
37 + assert "bailing_hybrid_mtp" in MTPModelTypes.__args__
38 +
39 +
40 +def test_bailing_hybrid_mtp_model_arch_config():
41 + config = _bailing_config()
42 + config.model_type = "bailing_hybrid_mtp"
43 + config.architectures = ["BailingMoeV25MTPModel"]
44 +
45 + model_arch_config = BailingHybridMTPModelArchConfigConvertor(
46 + config, config
47 + ).convert()
48 +
49 + assert model_arch_config.model_type == "bailing_hybrid_mtp"
50 + assert model_arch_config.architectures == ["BailingMoeV25MTPModel"]
51 + assert model_arch_config.total_num_hidden_layers == 1
52 + assert model_arch_config.is_deepseek_mla
tests/models/registry.py
+6
@@ -1557,6 +1557,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1557 use_original_num_layers=True,
1558 ),
1559 # [MTP]
1560 + "BailingMoeV25MTPModel": _HfExamplesInfo(
1561 + "inclusionAI/Ring-2.5-1T",
1562 + speculative_model="inclusionAI/Ring-2.5-1T",
1563 + trust_remote_code=True,
1564 + is_available_online=False,
1565 + ),
1566 "DeepSeekMTPModel": _HfExamplesInfo(
1567 "luccafong/deepseek_mtp_main_random",
1568 speculative_model="luccafong/deepseek_mtp_draft_random",
tests/v1/attention/test_linear_attention_metadata_builder.py new
+188
@@ -0,0 +1,188 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from tests.v1.attention.utils import (
7 + BatchSpec,
8 + create_common_attn_metadata,
9 + create_vllm_config,
10 +)
11 +from vllm.config import CUDAGraphMode, SpeculativeConfig
12 +from vllm.v1.attention.backend import AttentionCGSupport
13 +from vllm.v1.attention.backends.linear_attn import (
14 + BailingLinearAttentionMetadataBuilder,
15 + LinearAttentionMetadataBuilder,
16 +)
17 +from vllm.v1.attention.backends.utils import PAD_SLOT_ID
18 +from vllm.v1.kv_cache_interface import MambaSpec
19 +
20 +BLOCK_SIZE = 16
21 +DEVICE = torch.device("cpu")
22 +
23 +
24 +def _create_mamba_spec(num_speculative_blocks: int = 1) -> MambaSpec:
25 + return MambaSpec(
26 + block_size=BLOCK_SIZE,
27 + shapes=((16, 64),),
28 + dtypes=(torch.float16,),
29 + num_speculative_blocks=num_speculative_blocks,
30 + )
31 +
32 +
33 +def test_bailing_linear_attention_reports_uniform_batch_cudagraph_support():
34 + vllm_config = create_vllm_config(
35 + hf_config_override={
36 + "architectures": ["BailingMoeV2_5ForCausalLM"],
37 + "model_type": "bailing_hybrid",
38 + }
39 + )
40 +
41 + support = BailingLinearAttentionMetadataBuilder.get_cudagraph_support(
42 + vllm_config, _create_mamba_spec()
43 + )
44 +
45 + assert support == AttentionCGSupport.UNIFORM_BATCH
46 +
47 +
48 +def test_non_bailing_linear_attention_keeps_single_token_cudagraph_support():
49 + vllm_config = create_vllm_config(
50 + hf_config_override={
51 + "architectures": ["MiniMaxText01ForCausalLM"],
52 + "model_type": "minimax_text_01",
53 + }
54 + )
55 +
56 + support = LinearAttentionMetadataBuilder.get_cudagraph_support(
57 + vllm_config, _create_mamba_spec()
58 + )
59 +
60 + assert support == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
61 +
62 +
63 +def test_linear_attention_spec_decode_full_graph_metadata_pads_cache_slots():
64 + vllm_config = create_vllm_config(
65 + hf_config_override={
66 + "architectures": ["BailingMoeV2_5ForCausalLM"],
67 + "model_type": "bailing_hybrid",
68 + }
69 + )
70 + vllm_config.speculative_config = SpeculativeConfig(
71 + method="ngram",
72 + num_speculative_tokens=1,
73 + )
74 + vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
75 +
76 + builder = BailingLinearAttentionMetadataBuilder(
77 + kv_cache_spec=_create_mamba_spec(),
78 + layer_names=["model.layers.0.self_attn"],
79 + vllm_config=vllm_config,
80 + device=DEVICE,
81 + )
82 +
83 + common = create_common_attn_metadata(
84 + BatchSpec(seq_lens=[20, 20, 0], query_lens=[2, 2, 0]),
85 + BLOCK_SIZE,
86 + DEVICE,
87 + )
88 + common.block_table_tensor[2].fill_(-1)
89 +
90 + metadata = builder.build(
91 + common_prefix_len=0,
92 + common_attn_metadata=common,
93 + num_accepted_tokens=torch.tensor([1, 2, 1], dtype=torch.int32),
94 + )
95 +
96 + assert metadata.num_decodes == 3
97 + assert metadata.num_prefills == 0
98 + assert metadata.num_decode_tokens == 4
99 + assert metadata.state_indices_tensor_d is not None
100 + assert metadata.state_indices_tensor_d.shape == (3, 2)
101 + assert torch.equal(
102 + metadata.state_indices_tensor_d[2],
103 + torch.full((2,), PAD_SLOT_ID, dtype=torch.int32),
104 + )
105 + assert torch.equal(
106 + metadata.state_indices_tensor[2],
107 + torch.tensor(PAD_SLOT_ID, dtype=torch.int32),
108 + )
109 + assert metadata.query_start_loc_d is not None
110 + assert metadata.query_start_loc_d.tolist() == [0, 2, 4, 4]
111 + assert metadata.num_accepted_tokens is not None
112 + assert metadata.num_accepted_tokens.tolist() == [1, 2, 1]
113 +
114 +
115 +def test_linear_attention_full_graph_metadata_uses_stable_decode_buffers():
116 + vllm_config = create_vllm_config(
117 + hf_config_override={
118 + "architectures": ["BailingMoeV2_5ForCausalLM"],
119 + "model_type": "bailing_hybrid",
120 + }
121 + )
122 + vllm_config.speculative_config = SpeculativeConfig(
123 + method="ngram",
124 + num_speculative_tokens=1,
125 + )
126 + vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
127 +
128 + builder = BailingLinearAttentionMetadataBuilder(
129 + kv_cache_spec=_create_mamba_spec(),
130 + layer_names=["model.layers.0.self_attn"],
131 + vllm_config=vllm_config,
132 + device=DEVICE,
133 + )
134 +
135 + common = create_common_attn_metadata(
136 + BatchSpec(seq_lens=[20, 20, 0], query_lens=[2, 2, 0]),
137 + BLOCK_SIZE,
138 + DEVICE,
139 + arange_block_indices=True,
140 + )
141 + common.block_table_tensor = torch.tensor(
142 + [[10, 11], [12, 13], [-1, -1]],
143 + dtype=torch.int32,
144 + device=DEVICE,
145 + )
146 +
147 + first = builder.build(
148 + common_prefix_len=0,
149 + common_attn_metadata=common,
150 + num_accepted_tokens=torch.tensor([1, 2, 1], dtype=torch.int32),
151 + )
152 + assert first.state_indices_tensor_d is not None
153 + assert first.query_start_loc_d is not None
154 + assert first.num_accepted_tokens is not None
155 + state_ptr = first.state_indices_tensor_d.data_ptr()
156 + query_ptr = first.query_start_loc_d.data_ptr()
157 + accepted_ptr = first.num_accepted_tokens.data_ptr()
158 +
159 + common2 = create_common_attn_metadata(
160 + BatchSpec(seq_lens=[36, 0, 0], query_lens=[2, 0, 0]),
161 + BLOCK_SIZE,
162 + DEVICE,
163 + arange_block_indices=True,
164 + )
165 + common2.block_table_tensor = torch.tensor(
166 + [[20, 21], [-1, -1], [-1, -1]],
167 + dtype=torch.int32,
168 + device=DEVICE,
169 + )
170 + second = builder.build(
171 + common_prefix_len=0,
172 + common_attn_metadata=common2,
173 + num_accepted_tokens=torch.tensor([2, 1, 1], dtype=torch.int32),
174 + )
175 +
176 + assert second.state_indices_tensor_d is not None
177 + assert second.query_start_loc_d is not None
178 + assert second.num_accepted_tokens is not None
179 + assert second.state_indices_tensor_d.data_ptr() == state_ptr
180 + assert second.query_start_loc_d.data_ptr() == query_ptr
181 + assert second.num_accepted_tokens.data_ptr() == accepted_ptr
182 + assert second.state_indices_tensor_d.tolist() == [
183 + [20, 21],
184 + [PAD_SLOT_ID, PAD_SLOT_ID],
185 + [PAD_SLOT_ID, PAD_SLOT_ID],
186 + ]
187 + assert second.query_start_loc_d.tolist() == [0, 2, 2, 2]
188 + assert second.num_accepted_tokens.tolist() == [2, 1, 1]
vllm/config/speculative.py
+16
@@ -46,6 +46,7 @@ MTPModelTypes = Literal[
46 "qwen3_5_mtp",
47 "longcat_flash_mtp",
48 "minimax_m3_mtp",
49 + "bailing_hybrid_mtp",
50 "mtp",
51 "pangu_ultra_moe_mtp",
52 "step3p5_mtp",
@@ -463,6 +464,21 @@ class SpeculativeConfig:
464 {"n_predict": n_predict, "architectures": ["Qwen3NextMTP"]}
465 )
466
467 + architectures = getattr(hf_config, "architectures", []) or []
468 + if (
469 + hf_config.model_type == "bailing_hybrid"
470 + or "BailingMoeV2_5ForCausalLM" in architectures
471 + ):
472 + hf_config.model_type = "bailing_hybrid_mtp"
473 + if hf_config.model_type == "bailing_hybrid_mtp":
474 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
475 + hf_config.update(
476 + {
477 + "n_predict": n_predict,
478 + "architectures": ["BailingMoeV25MTPModel"],
479 + }
480 + )
481 +
482 if hf_config.model_type == "exaone_moe":
483 hf_config.model_type = "exaone_moe_mtp"
484 if hf_config.model_type == "exaone_moe_mtp":
vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py
+322 -3
@@ -30,11 +30,10 @@ from vllm.model_executor.layers.mamba.linear.base import LinearAttention
30 from vllm.model_executor.layers.mamba.linear.minimax_linear_attn import (
31 MiniMaxText01LinearAttention,
32 MiniMaxText01LinearKernel,
33 - clear_linear_attention_cache_for_new_sequences,
33 linear_attention_decode,
35 - linear_attention_prefill_and_mix,
34 )
35 from vllm.model_executor.layers.rotary_embedding import get_rope
36 +from vllm.triton_utils import tl, triton
37 from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
38
39
@@ -57,6 +56,290 @@ def _build_rope_parameters(config: PretrainedConfig) -> dict | None:
56 return rope_parameters or None
57
58
59 +def clear_linear_attention_cache_for_new_sequences(
60 + kv_cache: torch.Tensor,
61 + state_indices_tensor: torch.Tensor,
62 + attn_metadata: LinearAttentionMetadata,
63 +) -> None:
64 + num_prefills = getattr(attn_metadata, "num_prefills", 0)
65 + if num_prefills <= 0:
66 + return
67 +
68 + num_decodes = getattr(attn_metadata, "num_decodes", 0)
69 + prefill_state_indices = getattr(attn_metadata, "state_indices_tensor_p", None)
70 + for prefill_idx in range(num_prefills):
71 + if num_decodes + prefill_idx + 1 >= len(attn_metadata.query_start_loc):
72 + break
73 + q_start = attn_metadata.query_start_loc[num_decodes + prefill_idx]
74 + q_end = attn_metadata.query_start_loc[num_decodes + prefill_idx + 1]
75 + query_len = q_end - q_start
76 + context_len = attn_metadata.seq_lens[num_decodes + prefill_idx] - query_len
77 + if context_len == 0:
78 + if prefill_state_indices is not None:
79 + block_to_clear = prefill_state_indices[prefill_idx]
80 + else:
81 + block_to_clear = state_indices_tensor[num_decodes + prefill_idx]
82 + kv_cache[block_to_clear, ...] = 0
83 +
84 +
85 +def linear_attention_prefill_and_mix(
86 + q: torch.Tensor,
87 + k: torch.Tensor,
88 + v: torch.Tensor,
89 + kv_cache: torch.Tensor,
90 + state_indices_tensor: torch.Tensor,
91 + attn_metadata: LinearAttentionMetadata,
92 + slope_rate: torch.Tensor,
93 + block_size: int,
94 + decode_fn,
95 + prefix_fn,
96 + layer_idx: int | None = None,
97 +) -> torch.Tensor:
98 + hidden = []
99 + req_offset = getattr(attn_metadata, "num_decodes", 0)
100 + prefill_state_indices = getattr(attn_metadata, "state_indices_tensor_p", None)
101 + for _prefill_idx in range(getattr(attn_metadata, "num_prefills", 0)):
102 + if req_offset + _prefill_idx + 1 >= len(attn_metadata.query_start_loc):
103 + break
104 + if prefill_state_indices is not None and _prefill_idx >= len(
105 + prefill_state_indices
106 + ):
107 + break
108 + if prefill_state_indices is None and _prefill_idx >= len(state_indices_tensor):
109 + break
110 + _start = attn_metadata.query_start_loc[req_offset + _prefill_idx]
111 + _end = attn_metadata.query_start_loc[req_offset + _prefill_idx + 1]
112 + if prefill_state_indices is not None:
113 + slot_id = prefill_state_indices[_prefill_idx]
114 + else:
115 + slot_id = state_indices_tensor[req_offset + _prefill_idx]
116 + qs = q[_start:_end].transpose(0, 1).contiguous()
117 + ks = k[_start:_end].transpose(0, 1).contiguous()
118 + vs = v[_start:_end].transpose(0, 1).contiguous()
119 + slice_layer_cache = kv_cache[slot_id, ...]
120 + out_slice = prefix_fn(
121 + qs,
122 + ks,
123 + vs,
124 + slice_layer_cache,
125 + slope_rate,
126 + block_size,
127 + layer_idx=layer_idx,
128 + )
129 + hidden.append(out_slice.contiguous())
130 +
131 + if attn_metadata.num_decode_tokens > 0:
132 + hidden_decode = decode_fn(
133 + q, k, v, kv_cache, state_indices_tensor, attn_metadata
134 + )
135 + hidden.insert(0, hidden_decode)
136 +
137 + if not hidden:
138 + return torch.empty((0, q.size(1) * q.size(2)), device=q.device, dtype=q.dtype)
139 +
140 + hidden = torch.concat(hidden, dim=0).contiguous()
141 + return hidden
142 +
143 +
144 +@triton.jit
145 +def _bailing_linear_attn_decode_spec_step_kernel(
146 + q_ptr,
147 + k_ptr,
148 + v_ptr,
149 + kv_cache_ptr,
150 + slope_rate_ptr,
151 + state_indices_ptr,
152 + query_start_loc_ptr,
153 + num_accepted_tokens_ptr,
154 + output_ptr,
155 + q_start: tl.constexpr,
156 + D: tl.constexpr,
157 + q_b_stride,
158 + q_h_stride,
159 + q_d_stride,
160 + k_b_stride,
161 + k_h_stride,
162 + k_d_stride,
163 + v_b_stride,
164 + v_h_stride,
165 + v_d_stride,
166 + cache_b_stride,
167 + cache_h_stride,
168 + cache_d0_stride,
169 + cache_d1_stride,
170 + state_indices_b_stride,
171 + state_indices_t_stride,
172 + output_b_stride,
173 + output_d_stride,
174 + DRAFT_IDX: tl.constexpr,
175 + STATE_WIDTH: tl.constexpr,
176 + BLOCK_SIZE: tl.constexpr,
177 +):
178 + req_id = tl.program_id(0)
179 + head_id = tl.program_id(1)
180 + block_id = tl.program_id(2)
181 +
182 + req_start = tl.load(query_start_loc_ptr + req_id).to(tl.int64)
183 + req_end = tl.load(query_start_loc_ptr + req_id + 1).to(tl.int64)
184 + query_len = req_end - req_start
185 + if query_len <= DRAFT_IDX:
186 + return
187 +
188 + dst_slot = tl.load(
189 + state_indices_ptr
190 + + req_id * state_indices_b_stride
191 + + DRAFT_IDX * state_indices_t_stride
192 + ).to(tl.int64)
193 + if dst_slot == -1:
194 + return
195 +
196 + if DRAFT_IDX == 0:
197 + accepted_offset = tl.load(num_accepted_tokens_ptr + req_id).to(tl.int64) - 1
198 + accepted_offset = tl.maximum(accepted_offset, 0)
199 + accepted_offset = tl.minimum(accepted_offset, STATE_WIDTH - 1)
200 + src_slot = tl.load(
201 + state_indices_ptr
202 + + req_id * state_indices_b_stride
203 + + accepted_offset * state_indices_t_stride
204 + ).to(tl.int64)
205 + else:
206 + src_slot = tl.load(
207 + state_indices_ptr
208 + + req_id * state_indices_b_stride
209 + + (DRAFT_IDX - 1) * state_indices_t_stride
210 + ).to(tl.int64)
211 + if src_slot == -1:
212 + return
213 +
214 + token_idx = req_start - q_start + DRAFT_IDX
215 + qk_offsets = tl.arange(0, D)
216 + v_offsets = tl.arange(0, BLOCK_SIZE) + block_id * BLOCK_SIZE
217 + qk_mask = qk_offsets < D
218 + v_mask = v_offsets < D
219 + kv_mask = qk_mask[:, None] & v_mask[None, :]
220 +
221 + q = tl.load(
222 + q_ptr + token_idx * q_b_stride + head_id * q_h_stride + qk_offsets * q_d_stride,
223 + mask=qk_mask,
224 + other=0.0,
225 + )
226 + k = tl.load(
227 + k_ptr + token_idx * k_b_stride + head_id * k_h_stride + qk_offsets * k_d_stride,
228 + mask=qk_mask,
229 + other=0.0,
230 + )
231 + v = tl.load(
232 + v_ptr + token_idx * v_b_stride + head_id * v_h_stride + v_offsets * v_d_stride,
233 + mask=v_mask,
234 + other=0.0,
235 + )
236 +
237 + cache_offsets = (
238 + qk_offsets[:, None] * cache_d0_stride + v_offsets[None, :] * cache_d1_stride
239 + )
240 + src_cache_ptr = (
241 + kv_cache_ptr
242 + + src_slot * cache_b_stride
243 + + head_id * cache_h_stride
244 + + cache_offsets
245 + )
246 + dst_cache_ptr = (
247 + kv_cache_ptr
248 + + dst_slot * cache_b_stride
249 + + head_id * cache_h_stride
250 + + cache_offsets
251 + )
252 +
253 + slope = tl.load(slope_rate_ptr + head_id)
254 + decay = tl.exp(-slope)
255 + kv_old = tl.load(src_cache_ptr, mask=kv_mask, other=0.0)
256 + kv_new = k[:, None] * v[None, :] + decay * kv_old
257 +
258 + output = tl.sum(q[:, None].to(tl.float32) * kv_new, axis=0)
259 + tl.store(dst_cache_ptr, kv_new, mask=kv_mask)
260 + tl.store(
261 + output_ptr
262 + + token_idx * output_b_stride
263 + + (head_id * D + v_offsets) * output_d_stride,
264 + output,
265 + mask=v_mask,
266 + )
267 +
268 +
269 +def bailing_linear_attention_decode_spec(
270 + q: torch.Tensor,
271 + k: torch.Tensor,
272 + v: torch.Tensor,
273 + kv_cache: torch.Tensor,
274 + slope_rate: torch.Tensor,
275 + state_indices_tensor: torch.Tensor,
276 + query_start_loc: torch.Tensor,
277 + num_accepted_tokens: torch.Tensor,
278 + q_start: int,
279 + q_end: int | None,
280 + slot_start: int,
281 + slot_end: int | None,
282 + block_size: int,
283 +) -> torch.Tensor:
284 + q_decode = q[q_start:q_end]
285 + k_decode = k[q_start:q_end]
286 + v_decode = v[q_start:q_end]
287 + hidden = torch.empty(
288 + (q_decode.shape[0], q.shape[1] * q.shape[2]),
289 + device=q.device,
290 + dtype=q.dtype,
291 + )
292 + hidden.zero_()
293 +
294 + state_indices_tensor = state_indices_tensor[slot_start:slot_end]
295 + query_start_loc = query_start_loc.to(device=q.device)
296 +
297 + batch_size = state_indices_tensor.shape[0]
298 + num_heads = q_decode.shape[1]
299 + head_dim = q_decode.shape[2]
300 + assert k_decode.shape == (q_decode.shape[0], num_heads, head_dim)
301 + assert v_decode.shape == (q_decode.shape[0], num_heads, head_dim)
302 + state_width = state_indices_tensor.shape[1]
303 +
304 + grid = (batch_size, num_heads, triton.cdiv(head_dim, block_size))
305 + for draft_idx in range(state_width):
306 + _bailing_linear_attn_decode_spec_step_kernel[grid](
307 + q_decode,
308 + k_decode,
309 + v_decode,
310 + kv_cache,
311 + slope_rate,
312 + state_indices_tensor,
313 + query_start_loc,
314 + num_accepted_tokens[:batch_size],
315 + hidden,
316 + q_start,
317 + head_dim,
318 + q_decode.stride(0),
319 + q_decode.stride(1),
320 + q_decode.stride(2),
321 + k_decode.stride(0),
322 + k_decode.stride(1),
323 + k_decode.stride(2),
324 + v_decode.stride(0),
325 + v_decode.stride(1),
326 + v_decode.stride(2),
327 + kv_cache.stride(0),
328 + kv_cache.stride(1),
329 + kv_cache.stride(2),
330 + kv_cache.stride(3),
331 + state_indices_tensor.stride(0),
332 + state_indices_tensor.stride(1),
333 + hidden.stride(0),
334 + hidden.stride(1),
335 + DRAFT_IDX=draft_idx,
336 + STATE_WIDTH=state_width,
337 + BLOCK_SIZE=block_size,
338 + )
339 +
340 + return hidden
341 +
342 +
343 class BailingGroupRMSNormGate(RMSNormGated):
344 def __init__(
345 self,
@@ -208,6 +491,13 @@ class BailingMoELinearAttention(LinearAttention):
491 raise ValueError(f"Duplicate layer name: {prefix}")
492 compilation_config.static_forward_context[prefix] = self
493
494 + def get_attn_backend(self):
495 + from vllm.v1.attention.backends.linear_attn import (
496 + BailingLinearAttentionBackend,
497 + )
498 +
499 + return BailingLinearAttentionBackend
500 +
501 @staticmethod
502 def weight_direct_load(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
503 """Load weight for linear attention layers.
@@ -368,13 +658,42 @@ class BailingMoELinearAttention(LinearAttention):
658
659 def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor, attn_metadata):
660 """Handle decode (single token per sequence)."""
661 + decode_state_indices = getattr(attn_metadata, "state_indices_tensor_d", None)
662 + num_accepted_tokens = getattr(attn_metadata, "num_accepted_tokens", None)
663 + query_start_loc = getattr(attn_metadata, "query_start_loc_d", None)
664 + if (
665 + decode_state_indices is not None
666 + and decode_state_indices.dim() > 1
667 + and num_accepted_tokens is not None
668 + and query_start_loc is not None
669 + ):
670 + return bailing_linear_attention_decode_spec(
671 + q,
672 + k,
673 + v,
674 + kv_cache,
675 + self.tp_slope,
676 + decode_state_indices,
677 + query_start_loc,
678 + num_accepted_tokens,
679 + q_start=0,
680 + q_end=attn_metadata.num_decode_tokens,
681 + slot_start=0,
682 + slot_end=attn_metadata.num_decodes,
683 + block_size=32,
684 + )
685 + decode_state_indices = (
686 + state_indices_tensor
687 + if decode_state_indices is None
688 + else decode_state_indices
689 + )
690 hidden = linear_attention_decode(
691 q,
692 k,
693 v,
694 kv_cache,
695 self.tp_slope,
377 - state_indices_tensor,
696 + decode_state_indices,
697 q_start=0,
698 q_end=attn_metadata.num_decode_tokens,
699 slot_start=0,
vllm/model_executor/models/bailing_moe_linear.py
+4
@@ -533,6 +533,10 @@ class BailingMoeV25Model(nn.Module):
533 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
534 return self.word_embeddings(input_ids)
535
536 + @property
537 + def embed_tokens(self) -> nn.Module:
538 + return self.word_embeddings
539 +
540 def forward(
541 self,
542 input_ids: torch.Tensor | None,
vllm/model_executor/models/bailing_moe_mtp.py new
+380
@@ -0,0 +1,380 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only Bailing MoE v2.5 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_linear import (
28 + BailingMoeV25,
29 + BailingMoeV25MLAAttention,
30 +)
31 +from vllm.sequence import IntermediateTensors
32 +
33 +from .utils import PPMissingLayer, is_pp_missing_parameter, maybe_prefix
34 +
35 +
36 +def _get_draft_hf_config(vllm_config: VllmConfig) -> PretrainedConfig:
37 + speculative_config = vllm_config.speculative_config
38 + if speculative_config is not None:
39 + draft_model_config = speculative_config.draft_model_config
40 + if draft_model_config is not None:
41 + return draft_model_config.hf_config
42 + return vllm_config.model_config.hf_config
43 +
44 +
45 +class BailingMTPSharedHead(nn.Module):
46 + def __init__(
47 + self,
48 + config: PretrainedConfig,
49 + prefix: str,
50 + vllm_config: VllmConfig,
51 + ) -> None:
52 + super().__init__()
53 + self.head = ParallelLMHead(
54 + config.vocab_size,
55 + config.hidden_size,
56 + quant_config=vllm_config.quant_config,
57 + prefix=maybe_prefix(prefix, "head"),
58 + )
59 +
60 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
61 + return hidden_states
62 +
63 +
64 +class BailingMoeV25MultiTokenPredictorLayer(nn.Module):
65 + def __init__(
66 + self,
67 + vllm_config: VllmConfig,
68 + prefix: str,
69 + layer_id: int,
70 + ) -> None:
71 + super().__init__()
72 + config = _get_draft_hf_config(vllm_config)
73 + self.config = config
74 + self.layer_id = layer_id
75 + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
76 + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
77 + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
78 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
79 + self.self_attn = BailingMoeV25MLAAttention(
80 + config,
81 + quant_config=vllm_config.quant_config,
82 + layer_id=layer_id,
83 + prefix=maybe_prefix(prefix, "self_attn"),
84 + cache_config=vllm_config.cache_config,
85 + )
86 + self.post_attention_layernorm = RMSNorm(
87 + config.hidden_size, eps=config.rms_norm_eps
88 + )
89 + self.mlp = BailingMoeV25(
90 + config,
91 + quant_config=vllm_config.quant_config,
92 + layer_id=layer_id,
93 + prefix=maybe_prefix(prefix, "mlp"),
94 + )
95 + self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
96 + self.shared_head = BailingMTPSharedHead(
97 + config,
98 + maybe_prefix(prefix, "shared_head"),
99 + vllm_config,
100 + )
101 +
102 + def forward(
103 + self,
104 + input_ids: torch.Tensor,
105 + positions: torch.Tensor,
106 + previous_hidden_states: torch.Tensor,
107 + inputs_embeds: torch.Tensor | None = None,
108 + spec_step_index: int = 0,
109 + ) -> torch.Tensor:
110 + assert inputs_embeds is not None
111 + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
112 + inputs_embeds = self.enorm(inputs_embeds)
113 + previous_hidden_states = self.hnorm(previous_hidden_states)
114 +
115 + hidden_states = self.eh_proj(
116 + torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
117 + )
118 + residual = hidden_states
119 + hidden_states = self.input_layernorm(hidden_states)
120 +
121 + hidden_states = self.self_attn(hidden_states, positions)
122 + hidden_states = residual + hidden_states
123 +
124 + residual = hidden_states
125 + hidden_states = self.post_attention_layernorm(hidden_states)
126 + hidden_states = self.mlp(hidden_states)
127 + hidden_states = residual + hidden_states.to(residual.device)
128 + return self.final_layernorm(hidden_states)
129 +
130 +
131 +class BailingMoeV25MultiTokenPredictor(nn.Module):
132 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
133 + super().__init__()
134 + config = _get_draft_hf_config(vllm_config)
135 + self.mtp_start_layer_idx = config.num_hidden_layers
136 + self.num_mtp_layers = config.num_nextn_predict_layers
137 + self.layers = nn.ModuleDict(
138 + {
139 + str(idx): BailingMoeV25MultiTokenPredictorLayer(
140 + vllm_config,
141 + f"{prefix}.layers.{idx}",
142 + idx,
143 + )
144 + for idx in range(
145 + self.mtp_start_layer_idx,
146 + self.mtp_start_layer_idx + self.num_mtp_layers,
147 + )
148 + }
149 + )
150 + if get_pp_group().is_first_rank:
151 + self.embed_tokens = VocabParallelEmbedding(
152 + config.vocab_size,
153 + config.hidden_size,
154 + org_num_embeddings=config.vocab_size,
155 + prefix=maybe_prefix(prefix, "embed_tokens"),
156 + )
157 + else:
158 + self.embed_tokens = PPMissingLayer()
159 + self.logits_processor = LogitsProcessor(config.vocab_size)
160 +
161 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
162 + return self.embed_tokens(input_ids)
163 +
164 + def forward(
165 + self,
166 + input_ids: torch.Tensor,
167 + positions: torch.Tensor,
168 + previous_hidden_states: torch.Tensor,
169 + inputs_embeds: torch.Tensor | None = None,
170 + spec_step_idx: int = 0,
171 + ) -> torch.Tensor:
172 + if inputs_embeds is None:
173 + inputs_embeds = self.embed_tokens(input_ids)
174 + current_step_idx = spec_step_idx % self.num_mtp_layers
175 + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
176 + input_ids,
177 + positions,
178 + previous_hidden_states,
179 + inputs_embeds,
180 + current_step_idx,
181 + )
182 +
183 + def compute_logits(
184 + self,
185 + hidden_states: torch.Tensor,
186 + spec_step_idx: int = 0,
187 + lm_head: nn.Module | None = None,
188 + ) -> torch.Tensor:
189 + current_step_idx = spec_step_idx % self.num_mtp_layers
190 + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
191 + head = lm_head if lm_head is not None else mtp_layer.shared_head.head
192 + return self.logits_processor(
193 + head,
194 + mtp_layer.shared_head(hidden_states),
195 + )
196 +
197 +
198 +@support_torch_compile
199 +class BailingMoeV25MTPModel(nn.Module):
200 + packed_modules_mapping = {
201 + "gate_up_proj": ["gate_proj", "up_proj"],
202 + "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
203 + }
204 +
205 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
206 + super().__init__()
207 + self.config = _get_draft_hf_config(vllm_config)
208 + self.lm_head: nn.Module | None = None
209 + self.model = BailingMoeV25MultiTokenPredictor(
210 + vllm_config=vllm_config,
211 + prefix=maybe_prefix(prefix, "model"),
212 + )
213 +
214 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
215 + return self.model.embed_input_ids(input_ids)
216 +
217 + def forward(
218 + self,
219 + input_ids: torch.Tensor | None,
220 + positions: torch.Tensor,
221 + hidden_states: torch.Tensor,
222 + intermediate_tensors: IntermediateTensors | None = None,
223 + inputs_embeds: torch.Tensor | None = None,
224 + spec_step_idx: int = 0,
225 + ) -> torch.Tensor:
226 + return self.model(
227 + input_ids,
228 + positions,
229 + hidden_states,
230 + inputs_embeds,
231 + spec_step_idx,
232 + )
233 +
234 + def compute_logits(
235 + self,
236 + hidden_states: torch.Tensor,
237 + spec_step_idx: int = 0,
238 + ) -> torch.Tensor:
239 + return self.model.compute_logits(hidden_states, spec_step_idx, self.lm_head)
240 +
241 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
242 + return fused_moe_make_expert_params_mapping(
243 + self,
244 + ckpt_gate_proj_name="gate_proj",
245 + ckpt_down_proj_name="down_proj",
246 + ckpt_up_proj_name="up_proj",
247 + num_experts=self.config.num_experts,
248 + num_redundant_experts=0,
249 + )
250 +
251 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
252 + stacked_params_mapping = [
253 + (".fused_qkv_a_proj", ".q_a_proj", 0),
254 + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
255 + (".gate_up_proj", ".gate_proj", 0),
256 + (".gate_up_proj", ".up_proj", 1),
257 + ]
258 + expert_params_mapping = list(self.get_expert_mapping())
259 + params_dict = dict(self.named_parameters(remove_duplicate=False))
260 + loaded_params: set[str] = set()
261 + loaded_mtp_layers: set[int] = set()
262 +
263 + def load_param(
264 + name: str,
265 + loaded_weight: torch.Tensor,
266 + shard_id=None,
267 + ) -> bool:
268 + name = maybe_remap_kv_scale_name(name, params_dict)
269 + if name is None:
270 + return False
271 + if name not in params_dict or is_pp_missing_parameter(name, self):
272 + return False
273 +
274 + param = params_dict[name]
275 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
276 + if shard_id is None:
277 + weight_loader(param, loaded_weight)
278 + elif isinstance(shard_id, int):
279 + weight_loader(param, loaded_weight, shard_id)
280 + else:
281 + weight_loader(
282 + param,
283 + loaded_weight,
284 + name,
285 + expert_id=shard_id[0],
286 + shard_id=shard_id[1],
287 + )
288 + loaded_params.add(name)
289 + return True
290 +
291 + def get_spec_layer_idx(name: str) -> int | None:
292 + if not name.startswith("model.layers."):
293 + return None
294 + try:
295 + layer_idx = int(name.split("model.layers.", 1)[1].split(".", 1)[0])
296 + except (IndexError, ValueError):
297 + return None
298 + mtp_idx = layer_idx - self.config.num_hidden_layers
299 + if 0 <= mtp_idx < self.config.num_nextn_predict_layers:
300 + return layer_idx
301 + return None
302 +
303 + def normalize_name(name: str) -> str:
304 + name = name.replace(".attention.dense", ".self_attn.o_proj")
305 + name = name.replace(".attention.", ".self_attn.")
306 + return name.replace(
307 + "mlp.gate.e_score_correction_bias",
308 + "mlp.gate.expert_bias",
309 + )
310 +
311 + def load_lm_head(loaded_weight: torch.Tensor) -> None:
312 + for layer_idx in range(
313 + self.model.mtp_start_layer_idx,
314 + self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
315 + ):
316 + name = f"model.layers.{layer_idx}.shared_head.head.weight"
317 + load_param(name, loaded_weight)
318 +
319 + for name, loaded_weight in weights:
320 + if "rotary_emb.inv_freq" in name:
321 + continue
322 +
323 + if name == "model.word_embeddings.weight":
324 + load_param("model.embed_tokens.weight", loaded_weight)
325 + continue
326 + if name == "lm_head.weight":
327 + load_lm_head(loaded_weight)
328 + continue
329 +
330 + spec_layer = get_spec_layer_idx(name)
331 + if spec_layer is None:
332 + continue
333 + name = normalize_name(name)
334 +
335 + loaded = False
336 + for param_name, weight_name, shard_id in stacked_params_mapping:
337 + if weight_name not in name:
338 + continue
339 + if "mlp.experts." in name and name not in params_dict:
340 + continue
341 + mapped_name = name.replace(weight_name, param_name)
342 + if load_param(mapped_name, loaded_weight, shard_id):
343 + loaded = True
344 + break
345 + if loaded:
346 + loaded_mtp_layers.add(spec_layer)
347 + continue
348 +
349 + if "mlp.experts" in name:
350 + for mapping in expert_params_mapping:
351 + param_name, weight_name, expert_id, shard_id = mapping
352 + if weight_name not in name:
353 + continue
354 + mapped_name = name.replace(weight_name, param_name)
355 + if load_param(
356 + mapped_name,
357 + loaded_weight,
358 + (expert_id, shard_id),
359 + ):
360 + loaded = True
361 + break
362 + if loaded:
363 + loaded_mtp_layers.add(spec_layer)
364 + continue
365 +
366 + if load_param(name, loaded_weight):
367 + loaded_mtp_layers.add(spec_layer)
368 +
369 + for layer_idx in range(
370 + self.model.mtp_start_layer_idx,
371 + self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
372 + ):
373 + if layer_idx not in loaded_mtp_layers:
374 + raise ValueError(
375 + f"Bailing MTP speculative decoding layer {layer_idx} "
376 + "weights are missing from checkpoint. Use a checkpoint "
377 + "that includes MTP layer weights, or disable speculative "
378 + "decoding."
379 + )
380 + return loaded_params
vllm/model_executor/models/registry.py
+1
@@ -612,6 +612,7 @@ _SPECULATIVE_DECODING_MODELS = {
612 "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
613 "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"),
614 "MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"),
615 + "BailingMoeV25MTPModel": ("bailing_moe_mtp", "BailingMoeV25MTPModel"),
616 "Gemma4MTPModel": ("gemma4_mtp", "Gemma4MTP"),
617 "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),
618 "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),
vllm/transformers_utils/model_arch_config_convertor.py
+7
@@ -271,6 +271,7 @@ class ModelArchConfigConvertorBase:
271 "pangu_ultra_moe",
272 "pangu_ultra_moe_mtp",
273 "bailing_hybrid",
274 + "bailing_hybrid_mtp",
275 ):
276 # check is deepseek_v4 model
277 if hasattr(self.hf_text_config, "compress_ratios"):
@@ -539,6 +540,11 @@ class Qwen3NextMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
540 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
541
542
543 +class BailingHybridMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
544 + def get_num_hidden_layers(self) -> int:
545 + return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
546 +
547 +
548 class Qwen3_5MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
549 def get_num_hidden_layers(self) -> int:
550 return getattr(self.hf_text_config, "mtp_num_hidden_layers", 0)
@@ -633,6 +639,7 @@ class MossAudioModelArchConfigConvertor(ModelArchConfigConvertorBase):
639
640 # hf_config.model_type -> convertor class
641 MODEL_ARCH_CONFIG_CONVERTORS = {
642 + "bailing_hybrid_mtp": BailingHybridMTPModelArchConfigConvertor,
643 "cohere_asr": CohereAsrModelArchConfigConvertor,
644 "dbrx": DbrxModelArchConfigConvertor,
645 "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor,
vllm/v1/attention/backends/linear_attn.py
+212 -1
@@ -4,7 +4,7 @@ from dataclasses import dataclass
4
5 import torch
6
7 -from vllm.config import VllmConfig
7 +from vllm.config import CompilationConfig, VllmConfig
8 from vllm.v1.attention.backend import (
9 AttentionBackend,
10 AttentionCGSupport,
@@ -12,6 +12,7 @@ from vllm.v1.attention.backend import (
12 CommonAttentionMetadata,
13 )
14 from vllm.v1.attention.backends.utils import (
15 + PAD_SLOT_ID,
16 mamba_get_block_table_tensor,
17 split_decodes_and_prefills,
18 )
@@ -91,3 +92,213 @@ class LinearAttentionMetadataBuilder(AttentionMetadataBuilder[LinearAttentionMet
92 state_indices_tensor=state_indices_tensor,
93 )
94 return attn_metadata
95 +
96 +
97 +class BailingLinearAttentionBackend(LinearAttentionBackend):
98 + @staticmethod
99 + def get_name() -> str:
100 + return "BAILING_LINEAR_ATTN"
101 +
102 + @staticmethod
103 + def get_builder_cls() -> type["BailingLinearAttentionMetadataBuilder"]:
104 + return BailingLinearAttentionMetadataBuilder
105 +
106 +
107 +@dataclass
108 +class BailingLinearAttentionMetadata(LinearAttentionMetadata):
109 + state_indices_tensor_d: torch.Tensor | None = None
110 + state_indices_tensor_p: torch.Tensor | None = None
111 + num_accepted_tokens: torch.Tensor | None = None
112 + query_start_loc_d: torch.Tensor | None = None
113 +
114 +
115 +class BailingLinearAttentionMetadataBuilder(LinearAttentionMetadataBuilder):
116 + supports_spec_decode_metadata = True
117 + supports_update_block_table: bool = False
118 +
119 + @classmethod
120 + def get_cudagraph_support(
121 + cls,
122 + vllm_config: VllmConfig,
123 + kv_cache_spec: AttentionSpec,
124 + ) -> AttentionCGSupport:
125 + return AttentionCGSupport.UNIFORM_BATCH
126 +
127 + def __init__(
128 + self,
129 + kv_cache_spec: AttentionSpec,
130 + layer_names: list[str],
131 + vllm_config: VllmConfig,
132 + device: torch.device,
133 + ):
134 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
135 + self.compilation_config: CompilationConfig = vllm_config.compilation_config
136 + self.num_spec_tokens: int = vllm_config.num_speculative_tokens
137 + self.use_spec_decode: bool = self.num_spec_tokens > 0
138 + self.decode_cudagraph_max_bs: int = vllm_config.scheduler_config.max_num_seqs
139 + if self.compilation_config.max_cudagraph_capture_size is not None:
140 + self.decode_cudagraph_max_bs = min(
141 + self.decode_cudagraph_max_bs,
142 + self.compilation_config.max_cudagraph_capture_size,
143 + )
144 + self.decode_state_indices_tensor: torch.Tensor = torch.empty(
145 + (self.decode_cudagraph_max_bs, 1 + self.num_spec_tokens),
146 + dtype=torch.int32,
147 + device=device,
148 + )
149 + self.decode_legacy_state_indices_tensor: torch.Tensor = torch.empty(
150 + (self.decode_cudagraph_max_bs,),
151 + dtype=torch.int32,
152 + device=device,
153 + )
154 + self.decode_query_start_loc: torch.Tensor = torch.empty(
155 + (self.decode_cudagraph_max_bs + 1,),
156 + dtype=torch.int32,
157 + device=device,
158 + )
159 + self.decode_num_accepted_tokens: torch.Tensor = torch.empty(
160 + (self.decode_cudagraph_max_bs,),
161 + dtype=torch.int32,
162 + device=device,
163 + )
164 + self._init_reorder_batch_threshold(1, self.use_spec_decode)
165 +
166 + def build_for_cudagraph_capture(
167 + self,
168 + common_attn_metadata: CommonAttentionMetadata,
169 + ) -> BailingLinearAttentionMetadata:
170 + num_accepted_tokens = None
171 + if self.use_spec_decode:
172 + assert common_attn_metadata.max_query_len <= 1 + self.num_spec_tokens, (
173 + "Bailing linear attention only supports speculative decoding "
174 + "with query length <= 1 + number of speculative tokens."
175 + )
176 + num_accepted_tokens = torch.diff(common_attn_metadata.query_start_loc)
177 + return self.build(
178 + common_prefix_len=0,
179 + common_attn_metadata=common_attn_metadata,
180 + num_accepted_tokens=num_accepted_tokens,
181 + )
182 +
183 + def build( # type: ignore[override]
184 + self,
185 + common_prefix_len: int,
186 + common_attn_metadata: CommonAttentionMetadata,
187 + fast_build: bool = False,
188 + *,
189 + num_accepted_tokens: torch.Tensor | None = None,
190 + num_decode_draft_tokens_cpu: torch.Tensor | None = None,
191 + ) -> BailingLinearAttentionMetadata:
192 + query_start_loc = common_attn_metadata.query_start_loc
193 + seq_lens = common_attn_metadata.seq_lens
194 + num_reqs = common_attn_metadata.num_reqs
195 + use_spec_decode = self.use_spec_decode and num_accepted_tokens is not None
196 +
197 + state_indices_tensor = mamba_get_block_table_tensor(
198 + common_attn_metadata.block_table_tensor,
199 + common_attn_metadata.seq_lens,
200 + self.kv_cache_spec,
201 + self.vllm_config.cache_config.mamba_cache_mode,
202 + )
203 + if state_indices_tensor.dim() == 1:
204 + state_indices_tensor = state_indices_tensor.unsqueeze(-1)
205 +
206 + decode_threshold = self.reorder_batch_threshold if use_spec_decode else 1
207 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
208 + split_decodes_and_prefills(
209 + common_attn_metadata,
210 + decode_threshold=decode_threshold,
211 + )
212 + )
213 + state_indices_tensor_d, state_indices_tensor_p = torch.split(
214 + state_indices_tensor,
215 + [num_decodes, num_prefills],
216 + dim=0,
217 + )
218 + state_indices_tensor_p = state_indices_tensor_p[:, 0]
219 +
220 + query_start_loc_d = None
221 + if use_spec_decode:
222 + assert num_accepted_tokens is not None
223 + state_indices_tensor_d = state_indices_tensor_d[
224 + :, : 1 + self.num_spec_tokens
225 + ]
226 + query_start_loc_d = query_start_loc[: num_decodes + 1]
227 + num_accepted_tokens = num_accepted_tokens[:num_decodes]
228 + else:
229 + state_indices_tensor_d = state_indices_tensor_d[:, 0]
230 + num_accepted_tokens = None
231 +
232 + legacy_state_indices_tensor = state_indices_tensor[:, 0]
233 + cudagraph_mode = self.compilation_config.cudagraph_mode
234 + use_full_cudagraph = (
235 + cudagraph_mode is not None and cudagraph_mode.has_full_cudagraphs()
236 + )
237 + if (
238 + num_prefills == 0
239 + and num_decodes <= self.decode_cudagraph_max_bs
240 + and use_full_cudagraph
241 + ):
242 + padded_bs = num_reqs
243 + is_padded_decode = seq_lens[:num_decodes] == 0
244 + if state_indices_tensor_d.dim() > 1:
245 + state_indices_tensor_d = torch.where(
246 + is_padded_decode.unsqueeze(1),
247 + torch.full_like(state_indices_tensor_d, PAD_SLOT_ID),
248 + state_indices_tensor_d,
249 + )
250 + self.decode_state_indices_tensor[:num_decodes].copy_(
251 + state_indices_tensor_d,
252 + non_blocking=True,
253 + )
254 + state_indices_tensor_d = self.decode_state_indices_tensor[:padded_bs]
255 + state_indices_tensor_d[num_decodes:] = PAD_SLOT_ID
256 +
257 + self.decode_legacy_state_indices_tensor[:num_decodes].copy_(
258 + torch.where(
259 + is_padded_decode,
260 + torch.full_like(
261 + legacy_state_indices_tensor[:num_decodes],
262 + PAD_SLOT_ID,
263 + ),
264 + legacy_state_indices_tensor[:num_decodes],
265 + ),
266 + non_blocking=True,
267 + )
268 + legacy_state_indices_tensor = self.decode_legacy_state_indices_tensor[
269 + :padded_bs
270 + ]
271 + legacy_state_indices_tensor[num_decodes:] = PAD_SLOT_ID
272 + if state_indices_tensor_d.dim() == 1:
273 + state_indices_tensor_d = legacy_state_indices_tensor
274 +
275 + if use_spec_decode and num_accepted_tokens is not None:
276 + assert query_start_loc_d is not None
277 + self.decode_query_start_loc[: num_decodes + 1].copy_(
278 + query_start_loc_d,
279 + non_blocking=True,
280 + )
281 + decode_num_query_tokens = query_start_loc_d[-1]
282 + query_start_loc_d = self.decode_query_start_loc[: padded_bs + 1]
283 + query_start_loc_d[num_decodes + 1 :] = decode_num_query_tokens
284 +
285 + self.decode_num_accepted_tokens[:num_decodes].copy_(
286 + num_accepted_tokens,
287 + non_blocking=True,
288 + )
289 + num_accepted_tokens = self.decode_num_accepted_tokens[:padded_bs]
290 + num_accepted_tokens[num_decodes:] = 1
291 +
292 + return BailingLinearAttentionMetadata(
293 + num_prefills=num_prefills,
294 + num_prefill_tokens=num_prefill_tokens,
295 + num_decodes=num_decodes,
296 + num_decode_tokens=num_decode_tokens,
297 + query_start_loc=query_start_loc,
298 + seq_lens=seq_lens,
299 + state_indices_tensor=legacy_state_indices_tensor,
300 + state_indices_tensor_d=state_indices_tensor_d,
301 + state_indices_tensor_p=state_indices_tensor_p,
302 + num_accepted_tokens=num_accepted_tokens,
303 + query_start_loc_d=query_start_loc_d,
304 + )
vllm/v1/worker/gpu_model_runner.py
+12 -2
@@ -132,6 +132,9 @@ from vllm.v1.attention.backend import (
132 CommonAttentionMetadata,
133 )
134 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
135 +from vllm.v1.attention.backends.linear_attn import (
136 + BailingLinearAttentionMetadataBuilder,
137 +)
138 from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder
139 from vllm.v1.attention.backends.utils import (
140 NULL_BLOCK_ID,
@@ -2442,9 +2445,16 @@ class GPUModelRunner(
2445
2446 extra_attn_metadata_args = {}
2447 if use_spec_decode and isinstance(
2445 - builder, (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder)
2448 + builder,
2449 + (
2450 + Mamba2AttentionMetadataBuilder,
2451 + GDNAttentionMetadataBuilder,
2452 + BailingLinearAttentionMetadataBuilder,
2453 + ),
2454 ):
2447 - assert ubid is None, "UBatching not supported with GDN yet"
2455 + assert ubid is None, (
2456 + "UBatching not supported with GDN or linear attn yet"
2457 + )
2458 extra_attn_metadata_args = dict(
2459 num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs_padded],
2460 num_decode_draft_tokens_cpu=self.num_decode_draft_tokens.cpu[