@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ab87f85231

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 7, 16:30 UTC

[Model] Ring 2.5 (#35102)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

Jiangyun Zhu committed Feb 26, 2026 at 18:17 UTC ab87f85231b07b107f16ddb9e985deb0d83975ae

8 files changed +1407 -70

docs/models/supported_models.md

@@ -372,6 +372,7 @@ th {

372 | `BaiChuanForCausalLM` | Baichuan2, Baichuan | `baichuan-inc/Baichuan2-13B-Chat`, `baichuan-inc/Baichuan-7B`, etc. | ✅︎ | ✅︎ |

373 | `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | ✅︎ | ✅︎ |

374 | `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ |

376 | `BambaForCausalLM` | Bamba | `ibm-ai-platform/Bamba-9B-fp8`, `ibm-ai-platform/Bamba-9B` | ✅︎ | ✅︎ |

378 | `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |

tests/models/registry.py

@@ -206,6 +206,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

206 "BailingMoeV2ForCausalLM": _HfExamplesInfo(

207 "inclusionAI/Ling-mini-2.0", trust_remote_code=True

208 ),

209 + "BailingMoeV2_5ForCausalLM": _HfExamplesInfo(

210 + "inclusionAI/Ring-2.5-1T", trust_remote_code=True

211 + ),

212 "BambaForCausalLM": _HfExamplesInfo(

213 "ibm-ai-platform/Bamba-9B-v1",

214 extras={"tiny": "hmellor/tiny-random-BambaForCausalLM"},

vllm/model_executor/layers/fla/ops/layernorm_guard.py

+30 -5

@@ -84,6 +84,7 @@ def layer_norm_fwd_kernel(

84 HAS_Z: tl.constexpr,

85 NORM_BEFORE_GATE: tl.constexpr,

86 IS_RMS_NORM: tl.constexpr,

87 + ACTIVATION: tl.constexpr,

88 ):

89 # Map the program id to the starting row of X and Y it should compute.

90 row_start = tl.program_id(0) * ROWS_PER_BLOCK

@@ -112,7 +113,10 @@ def layer_norm_fwd_kernel(

113 if HAS_Z and not NORM_BEFORE_GATE:

114 Z_base = Z + rows[:, None] * stride_z_row + col_offsets

115 z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32)

115 - x *= z * tl.sigmoid(z)

116 + if ACTIVATION == "swish" or ACTIVATION == "silu":

117 + x *= z * tl.sigmoid(z)

118 + elif ACTIVATION == "sigmoid":

119 + x *= tl.sigmoid(z)

120

121 # Compute mean and variance per row (reduce along axis 1)

122 if not IS_RMS_NORM:

@@ -155,7 +159,10 @@ def layer_norm_fwd_kernel(

159 if HAS_Z and NORM_BEFORE_GATE:

160 Z_base = Z + rows[:, None] * stride_z_row + col_offsets

161 z = tl.load(Z_base, mask=mask, other=0.0).to(tl.float32)

158 - y *= z * tl.sigmoid(z)

162 + if ACTIVATION == "swish" or ACTIVATION == "silu":

163 + y *= z * tl.sigmoid(z)

164 + elif ACTIVATION == "sigmoid":

165 + y *= tl.sigmoid(z)

166

167 # Write output

168 tl.store(Y_base, y, mask=mask)

@@ -178,6 +185,7 @@ def layer_norm_fwd(

185 group_size: int = None,

186 norm_before_gate: bool = True,

187 is_rms_norm: bool = False,

188 + activation: str = "swish",

189 ):

190 M, N = x.shape

191 if group_size is None:

@@ -232,9 +240,12 @@ def layer_norm_fwd(

240 eps,

241 BLOCK_N=BLOCK_N,

242 ROWS_PER_BLOCK=rows_per_block,

243 + HAS_BIAS=bias is not None,

244 + HAS_Z=z is not None,

245 NORM_BEFORE_GATE=norm_before_gate,

246 IS_RMS_NORM=is_rms_norm,

247 num_warps=num_warps,

248 + ACTIVATION=activation,

249 )

250 return out, mean, rstd

251

@@ -252,6 +263,7 @@ class LayerNormFn(torch.autograd.Function):

263 group_size=None,

264 norm_before_gate=True,

265 is_rms_norm=False,

266 + activation: str = "swish",

267 ):

268 """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))"""

269

@@ -277,6 +289,7 @@ class LayerNormFn(torch.autograd.Function):

289 group_size=group_size,

290 norm_before_gate=norm_before_gate,

291 is_rms_norm=is_rms_norm,

292 + activation=activation,

293 )

294 ctx.save_for_backward(x, weight, bias, mean, rstd, z)

295 ctx.x_shape_og = x_shape_og

@@ -284,6 +297,7 @@ class LayerNormFn(torch.autograd.Function):

297 ctx.group_size = group_size

298 ctx.norm_before_gate = norm_before_gate

299 ctx.is_rms_norm = is_rms_norm

300 + ctx.activation = activation

301 return y.reshape(x_shape_og)

302

303

@@ -296,17 +310,25 @@ def layernorm_fn(

310 group_size=None,

311 norm_before_gate=True,

312 is_rms_norm=False,

313 + activation: str = "swish",

314 ):

315 return LayerNormFn.apply(

301 - x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm

316 + x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm, activation

317 )

318

319

320 def rmsnorm_fn(

306 - x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True

321 + x,

322 + weight,

323 + bias,

324 + z=None,

325 + eps=1e-6,

326 + group_size=None,

327 + norm_before_gate=True,

328 + activation: str = "swish",

329 ):

330 return LayerNormFn.apply(

309 - x, weight, bias, z, eps, group_size, norm_before_gate, True

331 + x, weight, bias, z, eps, group_size, norm_before_gate, True, activation

332 )

333

334

@@ -359,6 +381,7 @@ class RMSNormGated(nn.Module):

381 norm_before_gate: bool = False,

382 device: torch.device | None = None,

383 dtype: torch.dtype | None = None,

384 + activation: str = "swish",

385 ):

386 """If group_size is not None, we do GroupNorm with each group having group_size elements.

387 group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group).

@@ -366,6 +389,7 @@ class RMSNormGated(nn.Module):

389 factory_kwargs = {"device": device, "dtype": dtype}

390 super().__init__()

391 self.eps = eps

392 + self.activation = activation

393 self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))

394 self.register_parameter("bias", None)

395 self.group_size = group_size

@@ -385,4 +409,5 @@ class RMSNormGated(nn.Module):

409 eps=self.eps,

410 group_size=self.group_size,

411 norm_before_gate=self.norm_before_gate,

412 + activation=self.activation,

413 )

vllm/model_executor/layers/layernorm.py

@@ -592,6 +592,7 @@ class RMSNormGated(CustomOp):

592 eps=self.eps,

593 group_size=self.group_size,

594 norm_before_gate=self.norm_before_gate,

595 + activation=self.activation,

596 )

597

598

vllm/model_executor/layers/mamba/linear_attn.py

+124 -65

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import math

5 +from collections.abc import Callable

7 import torch

8 import torch.nn.functional as F

@@ -43,7 +44,6 @@ class MiniMaxText01RMSNormTP(CustomOp):

45 self.weight.weight_loader = self.weight_loader

46 self.variance_epsilon = eps

46 - return

48 @staticmethod

49 def weight_loader(

@@ -56,7 +56,6 @@ class MiniMaxText01RMSNormTP(CustomOp):

56 shard_size = loaded_weight.shape[0] // tp_world

57 shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)

58 param.data.copy_(loaded_weight[shard])

59 - return

60 def _forward(

61 self,

@@ -102,6 +101,101 @@ class MiniMaxText01RMSNormTP(CustomOp):

101 return q, k

102

103

104 +def clear_linear_attention_cache_for_new_sequences(

105 + kv_cache: torch.Tensor,

106 + state_indices_tensor: torch.Tensor,

107 + attn_metadata: LinearAttentionMetadata,

108 +) -> None:

109 + num_prefills = getattr(attn_metadata, "num_prefills", 0)

110 + if num_prefills <= 0:

111 + return

112 +

113 + num_decode_tokens = getattr(attn_metadata, "num_decode_tokens", 0)

114 + for prefill_idx in range(num_prefills):

115 + q_start = attn_metadata.query_start_loc[num_decode_tokens + prefill_idx]

116 + q_end = attn_metadata.query_start_loc[num_decode_tokens + prefill_idx + 1]

117 + query_len = q_end - q_start

118 + context_len = (

119 + attn_metadata.seq_lens[num_decode_tokens + prefill_idx] - query_len

120 + )

121 + if context_len == 0:

122 + block_to_clear = state_indices_tensor[num_decode_tokens + prefill_idx]

123 + kv_cache[block_to_clear, ...] = 0

124 +

125 +

126 +def linear_attention_decode(

127 + q: torch.Tensor,

128 + k: torch.Tensor,

129 + v: torch.Tensor,

130 + kv_cache: torch.Tensor,

131 + slope_rate: torch.Tensor,

132 + state_indices_tensor: torch.Tensor,

133 + q_start: int = 0,

134 + q_end: int | None = None,

135 + slot_start: int = 0,

136 + slot_end: int | None = None,

137 + block_size: int = 32,

138 +) -> torch.Tensor:

139 + q = q[q_start:q_end].unsqueeze(2).contiguous()

140 + k = k[q_start:q_end].unsqueeze(2).contiguous()

141 + v = v[q_start:q_end].unsqueeze(2).contiguous()

142 + slot_id = state_indices_tensor[slot_start:slot_end]

143 + return linear_decode_forward_triton(

144 + q, k, v, kv_cache, slope_rate, slot_id, block_size

145 + )

146 +

147 +

148 +def linear_attention_prefill_and_mix(

149 + q: torch.Tensor,

150 + k: torch.Tensor,

151 + v: torch.Tensor,

152 + kv_cache: torch.Tensor,

153 + state_indices_tensor: torch.Tensor,

154 + attn_metadata: LinearAttentionMetadata,

155 + slope_rate: torch.Tensor,

156 + block_size: int,

157 + decode_fn: Callable[..., torch.Tensor],

158 + prefix_fn: Callable[..., torch.Tensor],

159 + layer_idx: int | None = None,

160 +) -> torch.Tensor:

161 + hidden = []

162 + for _prefill_idx in range(getattr(attn_metadata, "num_prefills", 0)):

163 + if _prefill_idx >= len(attn_metadata.query_start_loc):

164 + break

165 + if _prefill_idx >= len(state_indices_tensor):

166 + break

167 + offset = attn_metadata.num_decode_tokens

168 + _start = attn_metadata.query_start_loc[offset + _prefill_idx]

169 + _end = attn_metadata.query_start_loc[offset + _prefill_idx + 1]

170 + slot_id = state_indices_tensor[offset + _prefill_idx]

171 + qs = q[_start:_end].transpose(0, 1).contiguous()

172 + ks = k[_start:_end].transpose(0, 1).contiguous()

173 + vs = v[_start:_end].transpose(0, 1).contiguous()

174 + slice_layer_cache = kv_cache[slot_id, ...]

175 + out_slice = prefix_fn(

176 + qs,

177 + ks,

178 + vs,

179 + slice_layer_cache,

180 + slope_rate,

181 + block_size,

182 + layer_idx=layer_idx,

183 + )

184 + hidden.append(out_slice.contiguous())

185 +

186 + if attn_metadata.num_decode_tokens > 0:

187 + hidden_decode = decode_fn(

188 + q, k, v, kv_cache, state_indices_tensor, attn_metadata

189 + )

190 + hidden.insert(0, hidden_decode)

191 +

192 + if not hidden:

193 + return torch.empty((0, q.size(-1)), device=q.device, dtype=q.dtype)

194 +

195 + hidden = torch.concat(hidden, dim=0).contiguous()

196 + return hidden

197 +

198 +

199 class MiniMaxText01LinearKernel:

200 @staticmethod

201 def jit_linear_forward_prefix(

@@ -258,50 +352,33 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):

352 def _prefill_and_mix_infer(

353 self, q, k, v, kv_cache, state_indices_tensor, attn_metadata

354 ):

261 - hidden = []

262 - for _prefill_idx in range(getattr(attn_metadata, "num_prefills", 0)):

263 - if _prefill_idx >= len(attn_metadata.query_start_loc):

264 - break

265 - if _prefill_idx >= len(state_indices_tensor):

266 - break

267 - offset = attn_metadata.num_decode_tokens

268 - _start = attn_metadata.query_start_loc[offset + _prefill_idx]

269 - _end = attn_metadata.query_start_loc[offset + _prefill_idx + 1]

270 - slot_id = state_indices_tensor[offset + _prefill_idx]

271 - qs = q[_start:_end].transpose(0, 1).contiguous()

272 - ks = k[_start:_end].transpose(0, 1).contiguous()

273 - vs = v[_start:_end].transpose(0, 1).contiguous()

274 - slice_layer_cache = kv_cache[slot_id, ...]

275 -

276 - out_slice = MiniMaxText01LinearKernel.jit_linear_forward_prefix(

277 - qs,

278 - ks,

279 - vs,

280 - slice_layer_cache,

281 - self.tp_slope,

282 - self.BLOCK,

283 - layer_idx=self.layer_idx,

284 - )

285 - hidden.append(out_slice.contiguous())

286 - if attn_metadata.num_decode_tokens > 0:

287 - hidden_decode = self._decode_infer(

288 - q, k, v, kv_cache, state_indices_tensor, attn_metadata

289 - )

290 - hidden.insert(0, hidden_decode)

291 -

292 - if not hidden:

293 - return torch.empty((0, q.size(-1)), device=q.device, dtype=q.dtype)

294 -

295 - hidden = torch.concat(hidden, dim=0).contiguous()

296 - return hidden

355 + return linear_attention_prefill_and_mix(

356 + q=q,

357 + k=k,

358 + v=v,

359 + kv_cache=kv_cache,

360 + state_indices_tensor=state_indices_tensor,

361 + attn_metadata=attn_metadata,

362 + slope_rate=self.tp_slope,

363 + block_size=self.BLOCK,

364 + decode_fn=self._decode_infer,

365 + prefix_fn=MiniMaxText01LinearKernel.jit_linear_forward_prefix,

366 + layer_idx=self.layer_idx,

367 + )

368

369 def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor, attn_metadata):

299 - q = q[: attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()

300 - k = k[: attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()

301 - v = v[: attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()

302 - slot_id = state_indices_tensor[: attn_metadata.num_decodes]

303 - hidden = linear_decode_forward_triton(

304 - q, k, v, kv_cache, self.tp_slope, slot_id, 32

370 + hidden = linear_attention_decode(

371 + q,

372 + k,

373 + v,

374 + kv_cache,

375 + self.tp_slope,

376 + state_indices_tensor,

377 + q_start=0,

378 + q_end=attn_metadata.num_decode_tokens,

379 + slot_start=0,

380 + slot_end=attn_metadata.num_decodes,

381 + block_size=32,

382 )

383 return hidden

384

@@ -338,27 +415,9 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):

415 if attn_metadata is not None:

416 kv_cache = self.kv_cache[forward_context.virtual_engine][0]

417 state_indices_tensor = attn_metadata.state_indices_tensor

341 -

342 - num_prefills = getattr(attn_metadata, "num_prefills", 0)

343 - if num_prefills > 0:

344 - num_decode_tokens = getattr(attn_metadata, "num_decode_tokens", 0)

345 - for prefill_idx in range(num_prefills):

346 - q_start = attn_metadata.query_start_loc[

347 - num_decode_tokens + prefill_idx

348 - ]

349 - q_end = attn_metadata.query_start_loc[

350 - num_decode_tokens + prefill_idx + 1

351 - ]

352 - query_len = q_end - q_start

353 - context_len = (

354 - attn_metadata.seq_lens[num_decode_tokens + prefill_idx]

355 - - query_len

356 - )

357 - if context_len == 0:

358 - block_to_clear = state_indices_tensor[

359 - num_decode_tokens + prefill_idx

360 - ]

361 - kv_cache[block_to_clear, ...] = 0

418 + clear_linear_attention_cache_for_new_sequences(

419 + kv_cache, state_indices_tensor, attn_metadata

420 + )

421

422 decode_only = getattr(attn_metadata, "num_prefills", 0) == 0

423 if attn_metadata is None:

vllm/model_executor/models/bailing_moe_linear.py new

+1246

@@ -0,0 +1,1246 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import copy

4 +from collections.abc import Iterable

5 +

6 +import torch

7 +import torch.nn as nn

8 +import torch.nn.functional as F

9 +from transformers.configuration_utils import PretrainedConfig

10 +

11 +from vllm.compilation.decorators import support_torch_compile

12 +from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config

13 +from vllm.distributed import (

14 + get_pp_group,

15 + get_tensor_model_parallel_rank,

16 + get_tensor_model_parallel_world_size,

17 +)

18 +from vllm.forward_context import get_forward_context

19 +from vllm.logger import init_logger

20 +from vllm.model_executor.layers.fla.ops.layernorm_guard import (

21 + RMSNormGated,

22 + layernorm_fn,

23 +)

24 +from vllm.model_executor.layers.fused_moe import FusedMoE, SharedFusedMoE

25 +from vllm.model_executor.layers.layernorm import RMSNorm

26 +from vllm.model_executor.layers.linear import (

27 + ColumnParallelLinear,

28 + MergedColumnParallelLinear,

29 + QKVParallelLinear,

30 + ReplicatedLinear,

31 + RowParallelLinear,

32 +)

33 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

34 +from vllm.model_executor.layers.mamba.abstract import MambaBase

35 +from vllm.model_executor.layers.mamba.linear_attn import (

36 + MiniMaxText01LinearAttention,

37 + MiniMaxText01LinearKernel,

38 + MiniMaxText01RMSNormTP,

39 + clear_linear_attention_cache_for_new_sequences,

40 + linear_attention_decode,

41 + linear_attention_prefill_and_mix,

42 +)

43 +from vllm.model_executor.layers.mamba.mamba_utils import (

44 + MambaStateCopyFuncCalculator,

45 + MambaStateDtypeCalculator,

46 + MambaStateShapeCalculator,

47 +)

48 +from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper

49 +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

50 +from vllm.model_executor.layers.rotary_embedding import get_rope

51 +from vllm.model_executor.layers.vocab_parallel_embedding import (

52 + ParallelLMHead,

53 + VocabParallelEmbedding,

54 +)

55 +from vllm.model_executor.model_loader.weight_utils import (

56 + default_weight_loader,

57 + maybe_remap_kv_scale_name,

58 +)

59 +from vllm.model_executor.models.bailing_moe import BailingMLP

60 +from vllm.sequence import IntermediateTensors

61 +from vllm.v1.attention.backend import AttentionMetadata

62 +from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata

63 +

64 +from .interfaces import HasInnerState, IsHybrid, SupportsPP

65 +from .utils import (

66 + AutoWeightsLoader,

67 + PPMissingLayer,

68 + is_pp_missing_parameter,

69 + make_layers,

70 + maybe_prefix,

71 +)

72 +

73 +logger = init_logger(__name__)

74 +

75 +

76 +def is_linear_layer(layer_idx, layer_group_size):

77 + if layer_idx is None:

78 + return False

79 + if layer_group_size > 0:

80 + return (layer_idx + 1) % layer_group_size != 0

81 + else:

82 + return False

83 +

84 +

85 +def _build_rope_parameters(config: PretrainedConfig) -> dict | None:

86 + rope_parameters = copy.deepcopy(getattr(config, "rope_parameters", None)) or {}

87 + if "rope_theta" not in rope_parameters and hasattr(config, "rope_theta"):

88 + rope_parameters["rope_theta"] = config.rope_theta

89 + if "partial_rotary_factor" not in rope_parameters and hasattr(

90 + config, "partial_rotary_factor"

91 + ):

92 + rope_parameters["partial_rotary_factor"] = config.partial_rotary_factor

93 +

94 + rope_scaling = getattr(config, "rope_scaling", None)

95 + if isinstance(rope_scaling, dict):

96 + rope_scaling = copy.deepcopy(rope_scaling)

97 + if "type" in rope_scaling and "rope_type" not in rope_scaling:

98 + rope_scaling["rope_type"] = rope_scaling.pop("type")

99 + rope_parameters.update(rope_scaling)

100 +

101 + return rope_parameters or None

102 +

103 +

104 +class BailingMoeV25MLAAttention(nn.Module):

105 + """

106 + MLA Attention for BailingMoeV2.5 full attention layers.

107 + """

108 +

109 + def __init__(

110 + self,

111 + config: PretrainedConfig,

112 + quant_config: QuantizationConfig | None = None,

113 + layer_id: int = 0,

114 + prefix: str = "attention",

115 + cache_config: CacheConfig | None = None,

116 + ) -> None:

117 + super().__init__()

118 + self.hidden_size = config.hidden_size

119 + self.num_heads = config.num_attention_heads

120 + self.layer_id = layer_id

121 + self.prefix = prefix

122 +

123 + # MLA dimensions

124 + self.qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 128)

125 + self.qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 64)

126 + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim

127 + self.v_head_dim = getattr(config, "v_head_dim", 128)

128 +

129 + # LoRA ranks

130 + self.q_lora_rank = getattr(config, "q_lora_rank", None)

131 + self.kv_lora_rank = getattr(config, "kv_lora_rank", 512)

132 +

133 + tp_size = get_tensor_model_parallel_world_size()

134 + assert self.num_heads % tp_size == 0

135 + self.num_local_heads = self.num_heads // tp_size

136 +

137 + self.scaling = self.qk_head_dim**-0.5

138 +

139 + # KV projections

140 + self.kv_a_layernorm = RMSNorm(

141 + self.kv_lora_rank,

142 + eps=config.rms_norm_eps,

143 + )

144 + self.kv_b_proj = ColumnParallelLinear(

145 + self.kv_lora_rank,

146 + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),

147 + bias=False,

148 + quant_config=quant_config,

149 + prefix=f"{prefix}.kv_b_proj",

150 + )

151 +

152 + # Output projection

153 + self.o_proj = RowParallelLinear(

154 + self.num_heads * self.v_head_dim,

155 + self.hidden_size,

156 + bias=False,

157 + quant_config=quant_config,

158 + prefix=f"{prefix}.o_proj",

159 + )

160 +

161 + if self.q_lora_rank is not None:

162 + # Use fused_qkv_a_proj when q_lora_rank is set

163 + self.fused_qkv_a_proj = MergedColumnParallelLinear(

164 + self.hidden_size,

165 + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],

166 + bias=False,

167 + quant_config=quant_config,

168 + prefix=f"{prefix}.fused_qkv_a_proj",

169 + disable_tp=True,

170 + )

171 + self.q_a_layernorm = RMSNorm(

172 + self.q_lora_rank,

173 + eps=config.rms_norm_eps,

174 + )

175 + self.q_b_proj = ColumnParallelLinear(

176 + self.q_lora_rank,

177 + self.num_heads * self.qk_head_dim,

178 + bias=False,

179 + quant_config=quant_config,

180 + prefix=f"{prefix}.q_b_proj",

181 + )

182 + self.q_proj = None

183 + self.kv_a_proj_with_mqa = None

184 + else:

185 + # Direct projections when no q_lora_rank

186 + self.q_proj = ColumnParallelLinear(

187 + self.hidden_size,

188 + self.num_heads * self.qk_head_dim,

189 + bias=False,

190 + quant_config=quant_config,

191 + prefix=f"{prefix}.q_proj",

192 + )

193 + self.kv_a_proj_with_mqa = ReplicatedLinear(

194 + self.hidden_size,

195 + self.kv_lora_rank + self.qk_rope_head_dim,

196 + bias=False,

197 + quant_config=quant_config,

198 + prefix=f"{prefix}.kv_a_proj_with_mqa",

199 + )

200 + self.fused_qkv_a_proj = None

201 + self.q_a_layernorm = None

202 + self.q_b_proj = None

203 +

204 + rope_parameters = _build_rope_parameters(config)

205 + max_position = getattr(config, "max_position_embeddings", 8192)

206 + self.rotary_emb = get_rope(

207 + head_size=self.qk_rope_head_dim,

208 + max_position=max_position,

209 + is_neox_style=False,

210 + rope_parameters=rope_parameters or None,

211 + dtype=torch.float32,

212 + )

213 +

214 + # Build MLAModules for MultiHeadLatentAttentionWrapper

215 + mla_modules = MLAModules(

216 + kv_a_layernorm=self.kv_a_layernorm,

217 + kv_b_proj=self.kv_b_proj,

218 + rotary_emb=self.rotary_emb,

219 + o_proj=self.o_proj,

220 + fused_qkv_a_proj=self.fused_qkv_a_proj,

221 + kv_a_proj_with_mqa=self.kv_a_proj_with_mqa,

222 + q_a_layernorm=self.q_a_layernorm,

223 + q_b_proj=self.q_b_proj,

224 + q_proj=self.q_proj,

225 + indexer=None,

226 + is_sparse=False,

227 + topk_indices_buffer=None,

228 + )

229 +

230 + self.mla_attn = MultiHeadLatentAttentionWrapper(

231 + self.hidden_size,

232 + self.num_local_heads,

233 + self.scaling,

234 + self.qk_nope_head_dim,

235 + self.qk_rope_head_dim,

236 + self.v_head_dim,

237 + self.q_lora_rank,

238 + self.kv_lora_rank,

239 + mla_modules,

240 + cache_config,

241 + quant_config,

242 + prefix,

243 + )

244 +

245 + def forward(

246 + self,

247 + hidden_states: torch.Tensor,

248 + positions: torch.Tensor,

249 + ) -> torch.Tensor:

250 + """Forward pass for MLA attention."""

251 + return self.mla_attn(positions, hidden_states)

252 +

253 +

254 +class BailingMoEGate(nn.Module):

255 + def __init__(

256 + self,

257 + config: PretrainedConfig,

258 + params_dtype: torch.dtype | None = None,

259 + prefix: str = "",

260 + ):

261 + super().__init__()

262 + if params_dtype is None:

263 + params_dtype = torch.get_default_dtype()

264 + self.params_dtype = params_dtype

265 + self.weight = nn.Parameter(

266 + torch.empty(

267 + (config.num_experts, config.hidden_size),

268 + dtype=self.params_dtype,

269 + ),

270 + )

271 + if getattr(config, "moe_router_enable_expert_bias", False):

272 + self.expert_bias = nn.Parameter(

273 + torch.empty((config.num_experts,), dtype=torch.float32),

274 + )

275 + else:

276 + self.expert_bias = None

277 +

278 + def forward(self, hidden_states):

279 + logits = F.linear(hidden_states.to(self.weight.dtype), self.weight, None).to(

280 + hidden_states.dtype

281 + )

282 + return logits

283 +

284 +

285 +class BailingMoeV25(nn.Module):

286 + """Bailing MoE v2.5 - standalone implementation for linear attention model."""

287 +

288 + def __init__(

289 + self,

290 + config: PretrainedConfig,

291 + quant_config: QuantizationConfig | None = None,

292 + layer_id: int = 0,

293 + prefix: str = "",

294 + ):

295 + super().__init__()

296 +

297 + self.layer_id = layer_id

298 + self.tp_size = get_tensor_model_parallel_world_size()

299 + self.tp_rank = get_tensor_model_parallel_rank()

300 + self.num_experts = config.num_experts

301 + self.top_k = config.num_experts_per_tok

302 + norm_topk_prob = getattr(config, "norm_topk_prob", None)

303 + # Ring-2.5 reference implementations normalize routing weights by default.

304 + self.norm_expert_prob = True if norm_topk_prob is None else bool(norm_topk_prob)

305 + self.hidden_size = config.hidden_size

306 + self.quant_config = quant_config

307 + self.num_shared_experts = config.num_shared_experts

308 + self.score_function = getattr(config, "score_function", None)

309 + self.n_group = getattr(config, "n_group", None)

310 + self.topk_group = getattr(config, "topk_group", None)

311 + self.use_grouped_topk = self.n_group is not None and self.topk_group is not None

312 + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)

313 +

314 + router_dtype = getattr(config, "router_dtype", None)

315 + if router_dtype is None or router_dtype == "fp32":

316 + self.router_dtype = torch.float32

317 + else:

318 + self.router_dtype = torch.bfloat16

319 +

320 + # Gate for routing

321 + self.gate = BailingMoEGate(

322 + config=config,

323 + params_dtype=self.router_dtype,

324 + prefix=f"{prefix}.gate",

325 + )

326 + correction_bias = (

327 + self.gate.expert_bias if self.gate.expert_bias is not None else None

328 + )

329 + if self.score_function is not None:

330 + assert (self.score_function == "softmax" and correction_bias is None) or (

331 + self.score_function == "sigmoid" and correction_bias is not None

332 + ), (

333 + "score_function and correction_bias should be "

334 + "(softmax, None) or (sigmoid, not None)"

335 + )

336 +

337 + # Shared experts (using BailingMLP)

338 + if self.num_shared_experts > 0:

339 + if hasattr(config, "moe_shared_expert_intermediate_size"):

340 + intermediate_size = config.moe_shared_expert_intermediate_size

341 + else:

342 + intermediate_size = config.moe_intermediate_size

343 + intermediate_size *= config.num_shared_experts

344 + self.shared_experts = BailingMLP(

345 + intermediate_size=intermediate_size,

346 + config=config,

347 + quant_config=quant_config,

348 + reduce_results=False,

349 + prefix=f"{prefix}.shared_experts",

350 + )

351 + else:

352 + self.shared_experts = None

353 +

354 + # Routed experts using SharedFusedMoE

355 + self.experts = SharedFusedMoE(

356 + shared_experts=self.shared_experts,

357 + num_experts=self.num_experts,

358 + top_k=self.top_k,

359 + hidden_size=self.hidden_size,

360 + intermediate_size=config.moe_intermediate_size,

361 + reduce_results=False,

362 + renormalize=self.norm_expert_prob,

363 + quant_config=quant_config,

364 + prefix=f"{prefix}.experts",

365 + scoring_func=self.score_function,

366 + e_score_correction_bias=correction_bias,

367 + num_expert_group=self.n_group,

368 + topk_group=self.topk_group,

369 + use_grouped_topk=self.use_grouped_topk,

370 + router_logits_dtype=self.router_dtype,

371 + )

372 +

373 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

374 + num_tokens, hidden_size = hidden_states.shape

375 + # Ensure contiguous token-major layout before router/projections.

376 + hidden_states = hidden_states.contiguous().view(-1, hidden_size)

377 +

378 + # router_logits: (num_tokens, n_experts)

379 + router_logits = self.gate(hidden_states.to(self.router_dtype))

380 + router_logits = router_logits.to(hidden_states.dtype)

381 +

382 + final_hidden_states = self.experts(

383 + hidden_states=hidden_states, router_logits=router_logits

384 + )

385 +

386 + # Handle tuple return from SharedFusedMoE

387 + if self.shared_experts is not None:

388 + shared_output, final_hidden_states = final_hidden_states

389 + else:

390 + shared_output = None

391 +

392 + final_hidden_states *= self.routed_scaling_factor

393 +

394 + if shared_output is not None:

395 + final_hidden_states = final_hidden_states + shared_output

396 +

397 + if self.tp_size > 1:

398 + final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(

399 + final_hidden_states

400 + )

401 +

402 + return final_hidden_states.view(num_tokens, hidden_size)

403 +

404 +

405 +BailingRMSNormTP = MiniMaxText01RMSNormTP

406 +

407 +

408 +class BailingGroupRMSNormGate(RMSNormGated):

409 + def __init__(

410 + self,

411 + hidden_size,

412 + eps=1e-5,

413 + group_size=None,

414 + norm_before_gate=True,

415 + device=None,

416 + dtype=None,

417 + ):

418 + super().__init__(

419 + hidden_size,

420 + eps=eps,

421 + group_size=group_size,

422 + norm_before_gate=norm_before_gate,

423 + device=device,

424 + dtype=dtype,

425 + activation="sigmoid",

426 + )

427 + # Add custom weight loader for TP sharding

428 + self.weight.weight_loader = self._weight_loader

429 +

430 + @staticmethod

431 + def _weight_loader(param: torch.nn.Parameter, loaded_weight: torch.Tensor) -> None:

432 + """Load weight with TP sharding."""

433 + tp_size = get_tensor_model_parallel_world_size()

434 + tp_rank = get_tensor_model_parallel_rank()

435 + shard_size = loaded_weight.shape[0] // tp_size

436 + shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)

437 + param.data.copy_(loaded_weight[shard].contiguous())

438 +

439 +

440 +class BailingMoELinearAttention(nn.Module, MambaBase):

441 + """

442 + Bailing MoE Linear Attention implementation using minimax backend.

443 +

444 + This implements the linear attention mechanism from sglang, adapted for vLLM's

445 + v1 engine with MambaBase interface support.

446 + """

447 +

448 + @property

449 + def mamba_type(self) -> str:

450 + return "linear_attention"

451 +

452 + def get_state_shape(self) -> tuple[tuple[int, ...], ...]:

453 + """Return state shape for linear attention cache.

454 +

455 + Must match the calculation in get_mamba_state_shape_from_config.

456 + """

457 + return MambaStateShapeCalculator.linear_attention_state_shape(

458 + num_heads=self.total_num_heads,

459 + tp_size=self.tp_size,

460 + head_dim=self.head_dim,

461 + )

462 +

463 + def get_state_dtype(self) -> tuple[torch.dtype, ...]:

464 + """Return state dtype for linear attention cache.

465 +

466 + Must match the calculation in get_mamba_state_dtype_from_config.

467 + """

468 + return MambaStateDtypeCalculator.linear_attention_state_dtype(

469 + self.model_config.dtype,

470 + self.cache_config.mamba_cache_dtype,

471 + )

472 +

473 + def __init__(

474 + self,

475 + config: PretrainedConfig,

476 + quant_config: QuantizationConfig | None = None,

477 + layer_id: int = 0,

478 + prefix: str = "linear_attn",

479 + model_config: ModelConfig | None = None,

480 + cache_config: CacheConfig | None = None,

481 + ):

482 + super().__init__()

483 +

484 + self.layer_id = layer_id

485 + self.hidden_size = config.hidden_size

486 + self.total_num_heads = config.num_attention_heads

487 + self.total_kv_heads = config.num_attention_heads # MHA

488 + self.tp_size = get_tensor_model_parallel_world_size()

489 + self.tp_rank = get_tensor_model_parallel_rank()

490 + self.model_config = model_config

491 + self.cache_config = cache_config

492 + self.prefix = prefix

493 +

494 + self.head_dim = (

495 + config.head_dim

496 + if hasattr(config, "head_dim")

497 + else config.hidden_size // self.total_num_heads

498 + )

499 +

500 + self.hidden_inner_size = self.head_dim * self.total_num_heads

501 + self.scaling = self.head_dim**-0.5

502 +

503 + assert self.total_num_heads % self.tp_size == 0

504 + self.tp_heads = self.total_num_heads // self.tp_size

505 +

506 + self.max_position_embeddings = config.max_position_embeddings

507 + self.rope_theta = getattr(config, "rope_theta", 600000)

508 +

509 + self.tp_kv_heads = self.total_kv_heads // self.tp_size

510 + self.q_size_per_rank = self.head_dim * self.tp_heads

511 + self.kv_size_per_rank = self.head_dim * self.tp_kv_heads

512 +

513 + self.use_qk_norm = getattr(config, "use_qk_norm", False)

514 + self.linear_backend = "minimax"

515 + self.linear_scale = self.linear_backend == "minimax"

516 + self.linear_rope = getattr(config, "linear_rope", True)

517 + if hasattr(config, "use_linear_silu"):

518 + self.linear_silu = config.use_linear_silu

519 + elif hasattr(config, "linear_silu"):

520 + self.linear_silu = config.linear_silu

521 + else:

522 + self.linear_silu = False

523 +

524 + # Block size for lightning attention

525 + self.BLOCK = getattr(config, "block", 256)

526 +

527 + self.query_key_value = QKVParallelLinear(

528 + self.hidden_size,

529 + self.head_dim,

530 + self.total_num_heads,

531 + self.total_num_heads, # MHA: kv_heads = num_heads

532 + bias=(config.use_bias or config.use_qkv_bias),

533 + quant_config=quant_config,

534 + prefix=f"{prefix}.query_key_value",

535 + )

536 +

537 + if self.use_qk_norm:

538 + self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

539 + self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

540 +

541 + self.g_proj = ColumnParallelLinear(

542 + self.hidden_size,

543 + self.hidden_inner_size,

544 + bias=False,

545 + quant_config=quant_config,

546 + prefix=f"{prefix}.g_proj",

547 + )

548 + self.dense = RowParallelLinear(

549 + self.hidden_inner_size,

550 + self.hidden_size,

551 + bias=config.use_bias,

552 + quant_config=quant_config,

553 + prefix=f"{prefix}.dense",

554 + reduce_results=True,

555 + )

556 +

557 + self.group_norm_size = getattr(config, "group_norm_size", 1)

558 + self.rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-5))

559 + assert self.tp_size <= self.group_norm_size, (

560 + "tp_size must be <= group_norm_size for local rms norm"

561 + )

562 + assert self.group_norm_size % self.tp_size == 0, (

563 + "group_norm_size must be divisible by tp_size"

564 + )

565 +

566 + # When group_norm_size == 1, group_size equals hidden_size // tp_size

567 + self.g_norm = BailingGroupRMSNormGate(

568 + hidden_size=self.hidden_inner_size // self.tp_size,

569 + eps=self.rms_norm_eps,

570 + group_size=(

571 + self.hidden_inner_size // self.group_norm_size

572 + if self.group_norm_size > 1

573 + else self.hidden_inner_size // self.tp_size

574 + ),

575 + )

576 +

577 + # use fp32 rotary embedding

578 + rope_parameters = _build_rope_parameters(config)

579 +

580 + self.rotary_emb = get_rope(

581 + self.head_dim,

582 + max_position=self.max_position_embeddings,

583 + is_neox_style=True,

584 + dtype=torch.float32,

585 + rope_parameters=rope_parameters or None,

586 + )

587 +

588 + # Build slope tensor for linear attention decay

589 + num_hidden_layers = config.num_hidden_layers

590 + slope_rate = MiniMaxText01LinearAttention._build_slope_tensor(

591 + self.total_num_heads

592 + )

593 + if num_hidden_layers <= 1:

594 + self.slope_rate = slope_rate * (1 + 1e-5)

595 + else:

596 + self.slope_rate = slope_rate * (

597 + 1 - layer_id / (num_hidden_layers - 1) + 1e-5

598 + )

599 + self.tp_slope = self.slope_rate[

600 + self.tp_rank * self.tp_heads : (self.tp_rank + 1) * self.tp_heads

601 + ].contiguous()

602 +

603 + # Register for compilation

604 + compilation_config = get_current_vllm_config().compilation_config

605 + if prefix in compilation_config.static_forward_context:

606 + raise ValueError(f"Duplicate layer name: {prefix}")

607 + compilation_config.static_forward_context[prefix] = self

608 +

609 + @staticmethod

610 + def weight_direct_load(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:

611 + """Load weight for linear attention layers.

612 +

613 + For FP8 quantized parameters, we need to use the weight_loader if available,

614 + as it handles special cases like tensor parallelism sharding.

615 + """

616 + # Check if param has a weight_loader (for vLLM ModelWeightParameter)

617 + weight_loader = getattr(param, "weight_loader", None)

618 + if weight_loader is not None:

619 + # Use the weight_loader which handles TP sharding and quantization

620 + weight_loader(param, loaded_weight)

621 + else:

622 + # Fall back to direct copy for standard tensors

623 + assert param.size() == loaded_weight.size(), (

624 + f"Shape mismatch: {param.shape} vs {loaded_weight.shape}"

625 + )

626 + param.data.copy_(loaded_weight)

627 +

628 + def forward(

629 + self,

630 + hidden_states: torch.Tensor,

631 + output: torch.Tensor,

632 + positions: torch.Tensor,

633 + ) -> None:

634 + """Forward method called by torch.ops.vllm.linear_attention"""

635 + torch.ops.vllm.linear_attention(

636 + hidden_states,

637 + output,

638 + positions,

639 + self.prefix,

640 + )

641 +

642 + def _forward(

643 + self,

644 + hidden_states: torch.Tensor,

645 + output: torch.Tensor,

646 + positions: torch.Tensor,

647 + ) -> None:

648 + """Actual forward implementation."""

649 + forward_context = get_forward_context()

650 + attn_metadata: AttentionMetadata = forward_context.attn_metadata

651 + if attn_metadata is not None:

652 + assert isinstance(attn_metadata, dict)

653 + attn_metadata = attn_metadata[self.prefix]

654 + assert isinstance(attn_metadata, LinearAttentionMetadata)

655 + num_actual_tokens = (

656 + attn_metadata.num_prefill_tokens + attn_metadata.num_decode_tokens

657 + )

658 + else:

659 + num_actual_tokens = hidden_states.shape[0]

660 +

661 + # QKV projection

662 + qkv, _ = self.query_key_value(hidden_states[:num_actual_tokens])

663 +

664 + # use rotary_emb support fp32

665 + qkv = qkv.to(torch.float32)

666 + if self.linear_silu:

667 + qkv = F.silu(qkv)

668 +

669 + # Split q, k, v

670 + q, k, v = torch.split(

671 + qkv,

672 + [self.q_size_per_rank, self.kv_size_per_rank, self.kv_size_per_rank],

673 + dim=-1,

674 + )

675 +

676 + # Apply QK norm if needed

677 + if self.use_qk_norm:

678 + q = q.reshape(-1, self.tp_heads, self.head_dim)

679 + k = k.reshape(-1, self.tp_kv_heads, self.head_dim)

680 + q = layernorm_fn(

681 + q,

682 + self.query_layernorm.weight.data,

683 + bias=None,

684 + eps=self.rms_norm_eps,

685 + is_rms_norm=True,

686 + )

687 + k = layernorm_fn(

688 + k,

689 + self.key_layernorm.weight.data,

690 + bias=None,

691 + eps=self.rms_norm_eps,

692 + is_rms_norm=True,

693 + )

694 + q = q.reshape(-1, self.q_size_per_rank)

695 + k = k.reshape(-1, self.kv_size_per_rank)

696 +

697 + # Apply rotary embeddings

698 + if self.linear_rope:

699 + q, k = self.rotary_emb(positions[:num_actual_tokens], q, k)

700 +

701 + # Reshape to [batch, heads, seq_len, head_dim]

702 + q = q.view((qkv.shape[0], self.tp_heads, self.head_dim))

703 + k = k.view((qkv.shape[0], self.tp_kv_heads, self.head_dim))

704 + v = v.view((qkv.shape[0], self.tp_kv_heads, self.head_dim))

705 +

706 + # Apply scaling if using minimax backend

707 + if self.linear_scale:

708 + q = q * self.scaling

709 +

710 + # Get KV cache and state indices

711 + if attn_metadata is not None:

712 + kv_cache = self.kv_cache[forward_context.virtual_engine][0]

713 + state_indices_tensor = attn_metadata.state_indices_tensor

714 + clear_linear_attention_cache_for_new_sequences(

715 + kv_cache, state_indices_tensor, attn_metadata

716 + )

717 +

718 + # Compute attention

719 + decode_only = getattr(attn_metadata, "num_prefills", 0) == 0

720 + if attn_metadata is None:

721 + hidden = torch.empty(

722 + (q.shape[0], q.shape[1] * q.shape[2]), device=q.device, dtype=q.dtype

723 + )

724 + else:

725 + if not decode_only:

726 + hidden = self._prefill_and_mix_infer(

727 + q, k, v, kv_cache, state_indices_tensor, attn_metadata

728 + )

729 + else:

730 + hidden = self._decode_infer(

731 + q, k, v, kv_cache, state_indices_tensor, attn_metadata

732 + )

733 +

734 + # Apply group norm and gate (matching SGLang behavior)

735 + gate, _ = self.g_proj(hidden_states[:num_actual_tokens])

736 +

737 + if self.group_norm_size > 1:

738 + hidden = self.g_norm(hidden, gate)

739 + else:

740 + hidden = self.g_norm(hidden)

741 + hidden = F.sigmoid(gate) * hidden

742 +

743 + hidden = hidden.to(hidden_states.dtype)

744 +

745 + # Output projection

746 + dense_out, _ = self.dense(hidden)

747 + output[:num_actual_tokens] = dense_out

748 +

749 + def _prefill_and_mix_infer(

750 + self, q, k, v, kv_cache, state_indices_tensor, attn_metadata

751 + ):

752 + """Handle prefill (mixed with decode if any)."""

753 + return linear_attention_prefill_and_mix(

754 + q=q,

755 + k=k,

756 + v=v,

757 + kv_cache=kv_cache,

758 + state_indices_tensor=state_indices_tensor,

759 + attn_metadata=attn_metadata,

760 + slope_rate=self.tp_slope,

761 + block_size=self.BLOCK,

762 + decode_fn=self._decode_infer,

763 + prefix_fn=MiniMaxText01LinearKernel.jit_linear_forward_prefix,

764 + layer_idx=self.layer_id,

765 + )

766 +

767 + def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor, attn_metadata):

768 + """Handle decode (single token per sequence)."""

769 + num_prefill_tokens = attn_metadata.num_prefill_tokens

770 + num_prefills = attn_metadata.num_prefills

771 + hidden = linear_attention_decode(

772 + q,

773 + k,

774 + v,

775 + kv_cache,

776 + self.tp_slope,

777 + state_indices_tensor,

778 + q_start=num_prefill_tokens,

779 + q_end=None,

780 + slot_start=num_prefills,

781 + slot_end=None,

782 + block_size=32,

783 + )

784 + return hidden

785 +

786 +

787 +class BailingMoeV25DecoderLayer(nn.Module):

788 + """Decoder layer supporting both linear and full attention."""

789 +

790 + def __init__(

791 + self,

792 + config: PretrainedConfig,

793 + quant_config: QuantizationConfig | None = None,

794 + layer_id: int = 0,

795 + prefix: str = "layer",

796 + model_config: ModelConfig | None = None,

797 + cache_config: CacheConfig | None = None,

798 + ) -> None:

799 + super().__init__()

800 + self.layer_id = layer_id

801 + self.hidden_size = config.hidden_size

802 +

803 + # Determine attention type (0 = linear, 1 = full)

804 + self.attention_type = getattr(config, "attention_type", 1)

805 +

806 + if self.attention_type == 0: # Linear attention

807 + self.self_attn = BailingMoELinearAttention(

808 + config,

809 + quant_config=quant_config,

810 + layer_id=layer_id,

811 + prefix=f"{prefix}.self_attn",

812 + model_config=model_config,

813 + cache_config=cache_config,

814 + )

815 + else: # Full attention

816 + self.self_attn = BailingMoeV25MLAAttention(

817 + config,

818 + quant_config=quant_config,

819 + layer_id=layer_id,

820 + prefix=f"{prefix}.self_attn",

821 + cache_config=cache_config,

822 + )

823 +

824 + # MLP/MoE

825 + is_moe_layer = config.num_experts > 1 and layer_id >= getattr(

826 + config, "first_k_dense_replace", 0

827 + )

828 +

829 + if is_moe_layer:

830 + self.mlp = BailingMoeV25(

831 + config,

832 + quant_config=quant_config,

833 + layer_id=layer_id,

834 + prefix=f"{prefix}.mlp",

835 + )

836 + else:

837 + self.mlp = BailingMLP(

838 + intermediate_size=config.intermediate_size,

839 + config=config,

840 + quant_config=quant_config,

841 + reduce_results=True,

842 + prefix=f"{prefix}.mlp",

843 + )

844 +

845 + # Layer norms

846 + rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-5))

847 + self.input_layernorm = RMSNorm(self.hidden_size, eps=rms_norm_eps)

848 + self.post_attention_layernorm = RMSNorm(self.hidden_size, eps=rms_norm_eps)

849 +

850 + def forward(

851 + self,

852 + hidden_states: torch.Tensor,

853 + positions: torch.Tensor,

854 + attn_metadata: AttentionMetadata | None = None,

855 + residual: torch.Tensor | None = None,

856 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

857 + # Input layernorm

858 + if residual is None:

859 + residual = hidden_states

860 + hidden_states = self.input_layernorm(hidden_states)

861 + else:

862 + hidden_states, residual = self.input_layernorm(hidden_states, residual)

863 +

864 + # Self attention

865 + if self.attention_type == 0:

866 + # Linear attention uses output tensor

867 + self_attention_output = torch.zeros_like(hidden_states)

868 + self.self_attn(

869 + hidden_states=hidden_states,

870 + output=self_attention_output,

871 + positions=positions,

872 + )

873 + else:

874 + # Full attention

875 + self_attention_output = self.self_attn(hidden_states, positions)

876 +

877 + hidden_states, residual = self.post_attention_layernorm(

878 + self_attention_output, residual

879 + )

880 + hidden_states = self.mlp(hidden_states)

881 + return hidden_states, residual

882 +

883 +

884 +@support_torch_compile(

885 + dynamic_arg_dims={

886 + "input_ids": 0,

887 + "positions": -1,

888 + "intermediate_tensors": 0,

889 + "inputs_embeds": 0,

890 + }

891 +)

892 +class BailingMoeV25Model(nn.Module):

893 + """Bailing MoE v2.5 Model with hybrid attention support."""

894 +

895 + def __init__(

896 + self,

897 + *,

898 + vllm_config: VllmConfig,

899 + prefix: str = "",

900 + ):

901 + super().__init__()

902 + config = vllm_config.model_config.hf_config

903 + model_config = vllm_config.model_config

904 + quant_config = vllm_config.quant_config

905 + cache_config = vllm_config.cache_config

906 +

907 + self.config = config

908 + self.vocab_size = config.vocab_size

909 + self.embed_dim = config.hidden_size

910 +

911 + # Determine layer types based on layer_group_size

912 + self.layer_group_size = getattr(config, "layer_group_size", 1)

913 + self.num_layers = config.num_hidden_layers

914 +

915 + # decoder_attention_types: 0 = linear, 1 = full

916 + self.decoder_attention_types = [

917 + 0 if is_linear_layer(i, self.layer_group_size) else 1

918 + for i in range(self.num_layers)

919 + ]

920 +

921 + # Embeddings

922 + if get_pp_group().is_first_rank:

923 + self.word_embeddings = VocabParallelEmbedding(

924 + self.vocab_size,

925 + self.embed_dim,

926 + org_num_embeddings=self.vocab_size,

927 + )

928 + else:

929 + from vllm.model_executor.models.utils import PPMissingLayer

930 +

931 + self.word_embeddings = PPMissingLayer()

932 +

933 + # Layers

934 + def layer_fn(prefix):

935 + layer_idx = int(prefix.split(".")[-1])

936 + layer_config = copy.deepcopy(config)

937 + layer_config.attention_type = self.decoder_attention_types[layer_idx]

938 +

939 + return BailingMoeV25DecoderLayer(

940 + config=layer_config,

941 + quant_config=quant_config,

942 + layer_id=layer_idx,

943 + prefix=prefix,

944 + model_config=model_config,

945 + cache_config=cache_config,

946 + )

947 +

948 + self.start_layer, self.end_layer, self.layers = make_layers(

949 + self.num_layers, layer_fn, prefix=f"{prefix}.layers"

950 + )

951 +

952 + # Final norm

953 + norm_kwargs = {}

954 + if hasattr(config, "rms_norm_eps"):

955 + norm_kwargs["eps"] = config.rms_norm_eps

956 + if get_pp_group().is_last_rank:

957 + self.norm = RMSNorm(config.hidden_size, **norm_kwargs)

958 + else:

959 + from vllm.model_executor.models.utils import PPMissingLayer

960 +

961 + self.norm = PPMissingLayer()

962 +

963 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

964 + return self.word_embeddings(input_ids)

965 +

966 + def forward(

967 + self,

968 + input_ids: torch.Tensor | None,

969 + positions: torch.Tensor,

970 + intermediate_tensors: IntermediateTensors | None = None,

971 + inputs_embeds: torch.Tensor | None = None,

972 + ) -> torch.Tensor:

973 + forward_context = get_forward_context()

974 + attn_metadata = forward_context.attn_metadata

975 +

976 + if get_pp_group().is_first_rank:

977 + if inputs_embeds is None:

978 + hidden_states = self.word_embeddings(input_ids)

979 + else:

980 + hidden_states = inputs_embeds

981 + residual = None

982 + else:

983 + assert intermediate_tensors is not None

984 + hidden_states = intermediate_tensors["hidden_states"]

985 + residual = intermediate_tensors["residual"]

986 +

987 + for layer in self.layers[self.start_layer : self.end_layer]:

988 + hidden_states, residual = layer(

989 + hidden_states=hidden_states,

990 + positions=positions,

991 + attn_metadata=attn_metadata,

992 + residual=residual,

993 + )

994 +

995 + if not get_pp_group().is_last_rank:

996 + return IntermediateTensors(

997 + {"hidden_states": hidden_states, "residual": residual}

998 + )

999 + else:

1000 + if residual is not None:

1001 + hidden_states, _ = self.norm(hidden_states, residual)

1002 + else:

1003 + hidden_states = self.norm(hidden_states)

1004 + return hidden_states

1005 +

1006 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

1007 + """Get expert parameter mapping for MoE layers."""

1008 + return FusedMoE.make_expert_params_mapping(

1009 + self,

1010 + ckpt_gate_proj_name="gate_proj",

1011 + ckpt_down_proj_name="down_proj",

1012 + ckpt_up_proj_name="up_proj",

1013 + num_experts=self.config.num_experts,

1014 + num_redundant_experts=0,

1015 + )

1016 +

1017 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

1018 + """Load checkpoint weights with simplified mapping."""

1019 +

1020 + params_dict = dict(self.named_parameters(remove_duplicate=False))

1021 + loaded_params: set[str] = set()

1022 +

1023 + # Stacked parameter mappings (fused projections)

1024 + stacked_mappings = [

1025 + (".fused_qkv_a_proj", ".q_a_proj", 0),

1026 + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),

1027 + (".gate_up_proj", ".gate_proj", 0),

1028 + (".gate_up_proj", ".up_proj", 1),

1029 + ]

1030 +

1031 + # Expert parameter mappings from FusedMoE

1032 + expert_mappings = list(self.get_expert_mapping())

1033 +

1034 + def load_param(name: str, tensor: torch.Tensor, shard_id=None) -> bool:

1035 + """Load a single parameter."""

1036 + if name not in params_dict or is_pp_missing_parameter(name, self):

1037 + return False

1038 + if name.endswith(".bias") and name not in params_dict:

1039 + return False

1040 +

1041 + param = params_dict[name]

1042 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

1043 +

1044 + if shard_id is None:

1045 + weight_loader(param, tensor)

1046 + elif isinstance(shard_id, int):

1047 + weight_loader(param, tensor, shard_id)

1048 + else:

1049 + # Expert param: (expert_id, shard_id)

1050 + weight_loader(

1051 + param, tensor, name, expert_id=shard_id[0], shard_id=shard_id[1]

1052 + )

1053 +

1054 + loaded_params.add(name)

1055 + return True

1056 +

1057 + def normalize_name(name: str) -> str | None:

1058 + """Normalize checkpoint name to model parameter name."""

1059 + # Skip special weights

1060 + if name.startswith("model.mtp"):

1061 + return None

1062 + # Remove 'model.' prefix if present

1063 + # (e.g., 'model.layers.0...' -> 'layers.0...')

1064 + name = name.removeprefix("model.")

1065 + # Map attention.dense based on layer type

1066 + if "attention.dense" in name:

1067 + layer_idx = (

1068 + int(name.split("layers.")[1].split(".")[0])

1069 + if "layers." in name

1070 + else 0

1071 + )

1072 + attn_name = (

1073 + "self_attn.dense"

1074 + if is_linear_layer(layer_idx, self.config.layer_group_size)

1075 + else "self_attn.o_proj"

1076 + )

1077 + name = name.replace("attention.dense", attn_name)

1078 +

1079 + # Standard mappings

1080 + name = name.replace("attention.", "self_attn.")

1081 + name = name.replace(

1082 + "mlp.gate.e_score_correction_bias", "mlp.gate.expert_bias"

1083 + )

1084 +

1085 + return maybe_remap_kv_scale_name(name, params_dict)

1086 +

1087 + for orig_name, weight in weights:

1088 + norm_name = normalize_name(orig_name)

1089 + if norm_name is None:

1090 + continue

1091 +

1092 + # Try stacked mappings

1093 + loaded = False

1094 + for param_suf, weight_suf, shard_id in stacked_mappings:

1095 + if weight_suf not in norm_name:

1096 + continue

1097 + mapped = norm_name.replace(weight_suf, param_suf).replace(

1098 + "attention.", "self_attn."

1099 + )

1100 + if load_param(mapped, weight, shard_id):

1101 + loaded = True

1102 + break

1103 + if loaded:

1104 + continue

1105 +

1106 + # Handle expert weights

1107 + if "mlp.experts" in norm_name:

1108 + # Expert bias

1109 + if (

1110 + "mlp.experts.e_score_correction_bias" in norm_name

1111 + or "mlp.experts.expert_bias" in norm_name

1112 + ):

1113 + alt = norm_name.replace(

1114 + "mlp.experts.e_score_correction_bias", "mlp.gate.expert_bias"

1115 + ).replace("mlp.experts.expert_bias", "mlp.gate.expert_bias")

1116 + if load_param(alt, weight) or load_param(norm_name, weight):

1117 + continue

1118 +

1119 + # Routed experts

1120 + for param_name, weight_name, expert_id, shard_id in expert_mappings:

1121 + if weight_name not in norm_name:

1122 + continue

1123 + mapped = norm_name.replace(weight_name, param_name)

1124 + if load_param(mapped, weight, (expert_id, shard_id)):

1125 + break

1126 + continue

1127 +

1128 + # General parameters

1129 + load_param(norm_name, weight)

1130 +

1131 + return loaded_params

1132 +

1133 +

1134 +class BailingMoeV25ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):

1135 + """Bailing MoE v2.5 For CausalLM."""

1136 +

1137 + packed_modules_mapping = {

1138 + "gate_up_proj": ["gate_proj", "up_proj"],

1139 + }

1140 +

1141 + def __init__(

1142 + self,

1143 + *,

1144 + vllm_config: VllmConfig,

1145 + prefix: str = "",

1146 + ) -> None:

1147 + super().__init__()

1148 + config = vllm_config.model_config.hf_config

1149 + quant_config = vllm_config.quant_config

1150 +

1151 + self.config = config

1152 + self.quant_config = quant_config

1153 +

1154 + self.model = BailingMoeV25Model(

1155 + vllm_config=vllm_config,

1156 + prefix=maybe_prefix(prefix, "model"),

1157 + )

1158 +

1159 + if get_pp_group().is_last_rank:

1160 + self.lm_head = ParallelLMHead(

1161 + config.vocab_size,

1162 + config.hidden_size,

1163 + quant_config=quant_config,

1164 + )

1165 + self.logits_processor = LogitsProcessor(config.vocab_size)

1166 + else:

1167 + self.lm_head = PPMissingLayer()

1168 +

1169 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

1170 + return self.model.embed_input_ids(input_ids)

1171 +

1172 + def forward(

1173 + self,

1174 + input_ids: torch.Tensor | None,

1175 + positions: torch.Tensor,

1176 + intermediate_tensors: IntermediateTensors | None = None,

1177 + inputs_embeds: torch.Tensor | None = None,

1178 + ) -> torch.Tensor:

1179 + hidden_states = self.model(

1180 + input_ids=input_ids,

1181 + positions=positions,

1182 + intermediate_tensors=intermediate_tensors,

1183 + inputs_embeds=inputs_embeds,

1184 + )

1185 + return hidden_states

1186 +

1187 + def compute_logits(

1188 + self,

1189 + hidden_states: torch.Tensor,

1190 + ) -> torch.Tensor:

1191 + return self.logits_processor(self.lm_head, hidden_states)

1192 +

1193 + def make_empty_intermediate_tensors(

1194 + self, batch_size: int, dtype: torch.dtype, device: torch.device

1195 + ) -> IntermediateTensors:

1196 + return IntermediateTensors(

1197 + {

1198 + "hidden_states": torch.zeros(

1199 + (batch_size, self.config.hidden_size), dtype=dtype, device=device

1200 + ),

1201 + "residual": torch.zeros(

1202 + (batch_size, self.config.hidden_size), dtype=dtype, device=device

1203 + ),

1204 + }

1205 + )

1206 +

1207 + @classmethod

1208 + def get_mamba_state_shape_from_config(

1209 + cls,

1210 + vllm_config: VllmConfig,

1211 + ) -> tuple[tuple[int, ...], ...]:

1212 + """Calculate shape for linear attention cache."""

1213 + config = vllm_config.model_config.hf_config

1214 + tp_size = vllm_config.parallel_config.tensor_parallel_size

1215 +

1216 + head_dim = getattr(

1217 + config, "head_dim", config.hidden_size // config.num_attention_heads

1218 + )

1219 +

1220 + # Return base state shape from linear attention (no padding)

1221 + return MambaStateShapeCalculator.linear_attention_state_shape(

1222 + num_heads=config.num_attention_heads,

1223 + tp_size=tp_size,

1224 + head_dim=head_dim,

1225 + )

1226 +

1227 + @classmethod

1228 + def get_mamba_state_dtype_from_config(

1229 + cls,

1230 + vllm_config: VllmConfig,

1231 + ) -> tuple[torch.dtype, ...]:

1232 + return MambaStateDtypeCalculator.linear_attention_state_dtype(

1233 + vllm_config.model_config.dtype,

1234 + vllm_config.cache_config.mamba_cache_dtype,

1235 + )

1236 +

1237 + @classmethod

1238 + def get_mamba_state_copy_func(cls) -> tuple:

1239 + return MambaStateCopyFuncCalculator.linear_attention_state_copy_func()

1240 +

1241 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

1242 + loader = AutoWeightsLoader(self)

1243 + return loader.load_weights(weights)

1244 +

1245 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

1246 + return self.model.get_expert_mapping()

vllm/model_executor/models/registry.py

@@ -81,6 +81,7 @@ _TEXT_GENERATION_MODELS = {

81 "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"),

82 "BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"),

83 "BailingMoeV2ForCausalLM": ("bailing_moe", "BailingMoeV2ForCausalLM"),

84 + "BailingMoeV2_5ForCausalLM": ("bailing_moe_linear", "BailingMoeV25ForCausalLM"),

85 "BambaForCausalLM": ("bamba", "BambaForCausalLM"),

86 "BloomForCausalLM": ("bloom", "BloomForCausalLM"),

87 "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"),

vllm/transformers_utils/model_arch_config_convertor.py

@@ -245,6 +245,7 @@ class ModelArchConfigConvertorBase:

245 "longcat_flash",

246 "pangu_ultra_moe",

247 "pangu_ultra_moe_mtp",

248 + "bailing_hybrid",

249 ):

250 return self.hf_text_config.kv_lora_rank is not None

251 elif self.hf_text_config.model_type == "eagle":