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 1, 23:45 UTC
[Attention] Mamba attention module refactor - LINEAR (#43556)
Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
wangxiyuan committed
Jun 4, 2026 at 18:45 UTC
90619351e33f9834963e6be0f756d3eafc20992c
7 files changed
+505
-551
tests/v1/attention/test_attention_backends_selection.py
+10
-9
@@ -54,15 +54,14 @@ from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend
54
(
55
MiniMaxText01LinearAttention,
56
dict(
57
- hidden_size=128,
58
- hidden_inner_size=256,
59
- num_heads=8,
60
- head_dim=32,
61
- max_position=2048,
62
- block_size=64,
63
- num_hidden_layer=12,
64
- layer_idx=0,
65
- linear_layer_idx=0,
57
+ config=SimpleNamespace(
58
+ hidden_size=256,
59
+ num_attention_heads=8,
60
+ head_dim=32,
61
+ num_hidden_layers=12,
62
+ block=64,
63
+ ),
64
+ prefix="layers.0.self_attn",
65
),
66
LinearAttentionBackend,
67
MambaAttentionBackendEnum.LINEAR,
@@ -88,6 +87,8 @@ def test_mamba_layers_get_attn_backend(
87
expected_mamba_type,
88
):
89
"""Test that Mamba-like layers return the correct attention backend."""
90
+ if layer_class is MiniMaxText01LinearAttention:
91
+ init_kwargs["vllm_config"] = default_vllm_config
92
layer = layer_class(**init_kwargs)
93
94
backend_class = layer.get_attn_backend()
vllm/model_executor/layers/mamba/linear/__init__.py
vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py
new
+384
@@ -0,0 +1,384 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import copy
4
+
5
+import torch
6
+import torch.nn.functional as F
7
+from transformers.configuration_utils import PretrainedConfig
8
+
9
+from vllm.config import (
10
+ VllmConfig,
11
+ get_current_vllm_config,
12
+)
13
+from vllm.distributed import (
14
+ get_tensor_model_parallel_rank,
15
+ get_tensor_model_parallel_world_size,
16
+)
17
+from vllm.forward_context import get_forward_context
18
+from vllm.model_executor.custom_op import PluggableLayer
19
+from vllm.model_executor.layers.fla.ops.layernorm_guard import (
20
+ RMSNormGated,
21
+ layernorm_fn,
22
+)
23
+from vllm.model_executor.layers.layernorm import RMSNorm
24
+from vllm.model_executor.layers.linear import (
25
+ ColumnParallelLinear,
26
+ QKVParallelLinear,
27
+ RowParallelLinear,
28
+)
29
+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,
34
+ linear_attention_decode,
35
+ linear_attention_prefill_and_mix,
36
+)
37
+from vllm.model_executor.layers.rotary_embedding import get_rope
38
+from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
39
+
40
+
41
+def _build_rope_parameters(config: PretrainedConfig) -> dict | None:
42
+ rope_parameters = copy.deepcopy(getattr(config, "rope_parameters", None)) or {}
43
+ if "rope_theta" not in rope_parameters and hasattr(config, "rope_theta"):
44
+ rope_parameters["rope_theta"] = config.rope_theta
45
+ if "partial_rotary_factor" not in rope_parameters and hasattr(
46
+ config, "partial_rotary_factor"
47
+ ):
48
+ rope_parameters["partial_rotary_factor"] = config.partial_rotary_factor
49
+
50
+ rope_scaling = getattr(config, "rope_scaling", None)
51
+ if isinstance(rope_scaling, dict):
52
+ rope_scaling = copy.deepcopy(rope_scaling)
53
+ if "type" in rope_scaling and "rope_type" not in rope_scaling:
54
+ rope_scaling["rope_type"] = rope_scaling.pop("type")
55
+ rope_parameters.update(rope_scaling)
56
+
57
+ return rope_parameters or None
58
+
59
+
60
+class BailingGroupRMSNormGate(RMSNormGated):
61
+ def __init__(
62
+ self,
63
+ hidden_size,
64
+ eps=1e-5,
65
+ group_size=None,
66
+ norm_before_gate=True,
67
+ device=None,
68
+ dtype=None,
69
+ ):
70
+ super().__init__(
71
+ hidden_size,
72
+ eps=eps,
73
+ group_size=group_size,
74
+ norm_before_gate=norm_before_gate,
75
+ device=device,
76
+ dtype=dtype,
77
+ activation="sigmoid",
78
+ )
79
+ # Add custom weight loader for TP sharding
80
+ self.weight.weight_loader = self._weight_loader
81
+
82
+ @staticmethod
83
+ def _weight_loader(param: torch.nn.Parameter, loaded_weight: torch.Tensor) -> None:
84
+ """Load weight with TP sharding."""
85
+ tp_size = get_tensor_model_parallel_world_size()
86
+ tp_rank = get_tensor_model_parallel_rank()
87
+ shard_size = loaded_weight.shape[0] // tp_size
88
+ shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
89
+ param.data.copy_(loaded_weight[shard].contiguous())
90
+
91
+
92
+# --8<-- [start:bailing_moe_linear_attention]
93
+@PluggableLayer.register("bailing_moe_linear_attention")
94
+class BailingMoELinearAttention(LinearAttention):
95
+ """Pluggable Bailing MoE Linear Attention layer which allows OOT backends
96
+ to add custom implementations.
97
+
98
+ This implements the linear attention mechanism from sglang, adapted for
99
+ vLLM's v1 engine with MambaBase interface support.
100
+ """
101
+
102
+ # --8<-- [end:bailing_moe_linear_attention]
103
+ def __init__(
104
+ self,
105
+ config: PretrainedConfig,
106
+ vllm_config: VllmConfig,
107
+ prefix: str = "linear_attn",
108
+ ):
109
+ super().__init__(config, vllm_config, prefix)
110
+
111
+ self.scaling = self.head_dim**-0.5
112
+
113
+ self.tp_heads = self.num_heads // self.tp_size
114
+
115
+ self.max_position_embeddings = config.max_position_embeddings
116
+ self.rope_theta = getattr(config, "rope_theta", 600000)
117
+
118
+ self.tp_kv_heads = self.num_heads // self.tp_size
119
+ self.q_size_per_rank = self.head_dim * self.tp_heads
120
+ self.kv_size_per_rank = self.head_dim * self.tp_kv_heads
121
+
122
+ self.use_qk_norm = getattr(config, "use_qk_norm", False)
123
+ self.linear_backend = "minimax"
124
+ self.linear_scale = self.linear_backend == "minimax"
125
+ self.linear_rope = getattr(config, "linear_rope", True)
126
+ if hasattr(config, "use_linear_silu"):
127
+ self.linear_silu = config.use_linear_silu
128
+ elif hasattr(config, "linear_silu"):
129
+ self.linear_silu = config.linear_silu
130
+ else:
131
+ self.linear_silu = False
132
+
133
+ self.query_key_value = QKVParallelLinear(
134
+ self.hidden_size,
135
+ self.head_dim,
136
+ self.num_heads,
137
+ self.num_heads, # MHA: kv_heads = num_heads
138
+ bias=(config.use_bias or config.use_qkv_bias),
139
+ quant_config=self.quant_config,
140
+ prefix=f"{prefix}.query_key_value",
141
+ )
142
+
143
+ if self.use_qk_norm:
144
+ self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
145
+ self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
146
+
147
+ self.g_proj = ColumnParallelLinear(
148
+ self.hidden_size,
149
+ self.hidden_inner_size,
150
+ bias=False,
151
+ quant_config=self.quant_config,
152
+ prefix=f"{prefix}.g_proj",
153
+ )
154
+ self.dense = RowParallelLinear(
155
+ self.hidden_inner_size,
156
+ self.hidden_size,
157
+ bias=config.use_bias,
158
+ quant_config=self.quant_config,
159
+ prefix=f"{prefix}.dense",
160
+ reduce_results=True,
161
+ )
162
+
163
+ self.group_norm_size = getattr(config, "group_norm_size", 1)
164
+ self.rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-5))
165
+ assert self.tp_size <= self.group_norm_size, (
166
+ "tp_size must be <= group_norm_size for local rms norm"
167
+ )
168
+ assert self.group_norm_size % self.tp_size == 0, (
169
+ "group_norm_size must be divisible by tp_size"
170
+ )
171
+
172
+ # When group_norm_size == 1, group_size equals hidden_size // tp_size
173
+ self.g_norm = BailingGroupRMSNormGate(
174
+ hidden_size=self.hidden_inner_size // self.tp_size,
175
+ eps=self.rms_norm_eps,
176
+ group_size=(
177
+ self.hidden_inner_size // self.group_norm_size
178
+ if self.group_norm_size > 1
179
+ else self.hidden_inner_size // self.tp_size
180
+ ),
181
+ )
182
+
183
+ # use fp32 rotary embedding
184
+ rope_parameters = _build_rope_parameters(config)
185
+
186
+ self.rotary_emb = get_rope(
187
+ self.head_dim,
188
+ max_position=self.max_position_embeddings,
189
+ is_neox_style=True,
190
+ rope_parameters=rope_parameters or None,
191
+ )
192
+
193
+ # Build slope tensor for linear attention decay
194
+ slope_rate = MiniMaxText01LinearAttention._build_slope_tensor(self.num_heads)
195
+ if self.num_hidden_layers <= 1:
196
+ self.slope_rate = slope_rate * (1 + 1e-5)
197
+ else:
198
+ self.slope_rate = slope_rate * (
199
+ 1 - self.layer_idx / (self.num_hidden_layers - 1) + 1e-5
200
+ )
201
+ self.tp_slope = self.slope_rate[
202
+ self.tp_rank * self.tp_heads : (self.tp_rank + 1) * self.tp_heads
203
+ ].contiguous()
204
+
205
+ # Register for compilation
206
+ compilation_config = get_current_vllm_config().compilation_config
207
+ if prefix in compilation_config.static_forward_context:
208
+ raise ValueError(f"Duplicate layer name: {prefix}")
209
+ compilation_config.static_forward_context[prefix] = self
210
+
211
+ @staticmethod
212
+ def weight_direct_load(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
213
+ """Load weight for linear attention layers.
214
+
215
+ For FP8 quantized parameters, we need to use the weight_loader if available,
216
+ as it handles special cases like tensor parallelism sharding.
217
+ """
218
+ # Check if param has a weight_loader (for vLLM ModelWeightParameter)
219
+ weight_loader = getattr(param, "weight_loader", None)
220
+ if weight_loader is not None:
221
+ # Use the weight_loader which handles TP sharding and quantization
222
+ weight_loader(param, loaded_weight)
223
+ else:
224
+ # Fall back to direct copy for standard tensors
225
+ assert param.size() == loaded_weight.size(), (
226
+ f"Shape mismatch: {param.shape} vs {loaded_weight.shape}"
227
+ )
228
+ param.data.copy_(loaded_weight)
229
+
230
+ def forward(
231
+ self,
232
+ hidden_states: torch.Tensor,
233
+ output: torch.Tensor,
234
+ positions: torch.Tensor,
235
+ ) -> None:
236
+ """Forward method called by torch.ops.vllm.linear_attention"""
237
+ torch.ops.vllm.linear_attention(
238
+ hidden_states,
239
+ output,
240
+ positions,
241
+ self.prefix,
242
+ )
243
+
244
+ def _forward(
245
+ self,
246
+ hidden_states: torch.Tensor,
247
+ output: torch.Tensor,
248
+ positions: torch.Tensor,
249
+ ) -> None:
250
+ """Actual forward implementation."""
251
+ forward_context = get_forward_context()
252
+ attn_metadata = forward_context.attn_metadata
253
+ if attn_metadata is not None:
254
+ assert isinstance(attn_metadata, dict)
255
+ attn_metadata = attn_metadata[self.prefix] # type: ignore
256
+ assert isinstance(attn_metadata, LinearAttentionMetadata)
257
+ num_actual_tokens = (
258
+ attn_metadata.num_prefill_tokens + attn_metadata.num_decode_tokens
259
+ )
260
+ else:
261
+ num_actual_tokens = hidden_states.shape[0]
262
+
263
+ # QKV projection
264
+ qkv, _ = self.query_key_value(hidden_states[:num_actual_tokens])
265
+
266
+ # use rotary_emb support fp32
267
+ qkv = qkv.to(torch.float32)
268
+ if self.linear_silu:
269
+ qkv = F.silu(qkv)
270
+
271
+ # Split q, k, v
272
+ q, k, v = torch.split(
273
+ qkv,
274
+ [self.q_size_per_rank, self.kv_size_per_rank, self.kv_size_per_rank],
275
+ dim=-1,
276
+ )
277
+
278
+ # Apply QK norm if needed
279
+ if self.use_qk_norm:
280
+ q = q.reshape(-1, self.tp_heads, self.head_dim)
281
+ k = k.reshape(-1, self.tp_kv_heads, self.head_dim)
282
+ q = layernorm_fn(
283
+ q,
284
+ self.query_layernorm.weight.data,
285
+ bias=None,
286
+ eps=self.rms_norm_eps,
287
+ is_rms_norm=True,
288
+ )
289
+ k = layernorm_fn(
290
+ k,
291
+ self.key_layernorm.weight.data,
292
+ bias=None,
293
+ eps=self.rms_norm_eps,
294
+ is_rms_norm=True,
295
+ )
296
+ q = q.reshape(-1, self.q_size_per_rank)
297
+ k = k.reshape(-1, self.kv_size_per_rank)
298
+
299
+ # Apply rotary embeddings
300
+ if self.linear_rope:
301
+ q, k = self.rotary_emb(positions[:num_actual_tokens], q, k)
302
+
303
+ # Reshape to [batch, heads, seq_len, head_dim]
304
+ q = q.view((qkv.shape[0], self.tp_heads, self.head_dim))
305
+ k = k.view((qkv.shape[0], self.tp_kv_heads, self.head_dim))
306
+ v = v.view((qkv.shape[0], self.tp_kv_heads, self.head_dim))
307
+
308
+ # Apply scaling if using minimax backend
309
+ if self.linear_scale:
310
+ q = q * self.scaling
311
+
312
+ # Get KV cache and state indices
313
+ if attn_metadata is not None:
314
+ kv_cache = self.kv_cache[0]
315
+ state_indices_tensor = attn_metadata.state_indices_tensor
316
+ clear_linear_attention_cache_for_new_sequences(
317
+ kv_cache, state_indices_tensor, attn_metadata
318
+ )
319
+
320
+ # Compute attention
321
+ decode_only = getattr(attn_metadata, "num_prefills", 0) == 0
322
+ if attn_metadata is None:
323
+ hidden = torch.empty(
324
+ (q.shape[0], q.shape[1] * q.shape[2]), device=q.device, dtype=q.dtype
325
+ )
326
+ else:
327
+ if not decode_only:
328
+ hidden = self._prefill_and_mix_infer(
329
+ q, k, v, kv_cache, state_indices_tensor, attn_metadata
330
+ )
331
+ else:
332
+ hidden = self._decode_infer(
333
+ q, k, v, kv_cache, state_indices_tensor, attn_metadata
334
+ )
335
+
336
+ # Apply group norm and gate (matching SGLang behavior)
337
+ gate, _ = self.g_proj(hidden_states[:num_actual_tokens])
338
+
339
+ if self.group_norm_size > 1:
340
+ hidden = self.g_norm(hidden, gate)
341
+ else:
342
+ hidden = self.g_norm(hidden)
343
+ hidden = F.sigmoid(gate) * hidden
344
+
345
+ hidden = hidden.to(hidden_states.dtype)
346
+
347
+ # Output projection
348
+ dense_out, _ = self.dense(hidden)
349
+ output[:num_actual_tokens] = dense_out
350
+
351
+ def _prefill_and_mix_infer(
352
+ self, q, k, v, kv_cache, state_indices_tensor, attn_metadata
353
+ ):
354
+ """Handle prefill (mixed with decode if any)."""
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):
370
+ """Handle decode (single token per sequence)."""
371
+ hidden = linear_attention_decode(
372
+ q,
373
+ k,
374
+ v,
375
+ kv_cache,
376
+ self.tp_slope,
377
+ state_indices_tensor,
378
+ q_start=0,
379
+ q_end=attn_metadata.num_decode_tokens,
380
+ slot_start=0,
381
+ slot_end=attn_metadata.num_decodes,
382
+ block_size=32,
383
+ )
384
+ return hidden
vllm/model_executor/layers/mamba/linear/base.py
new
+66
@@ -0,0 +1,66 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+from transformers import PretrainedConfig
5
+
6
+from vllm.config import (
7
+ VllmConfig,
8
+)
9
+from vllm.distributed.parallel_state import (
10
+ get_tensor_model_parallel_rank,
11
+ get_tensor_model_parallel_world_size,
12
+)
13
+from vllm.model_executor.custom_op import PluggableLayer
14
+from vllm.model_executor.layers.mamba.abstract import MambaBase
15
+from vllm.model_executor.layers.mamba.mamba_utils import (
16
+ MambaStateDtypeCalculator,
17
+ MambaStateShapeCalculator,
18
+)
19
+from vllm.model_executor.models.utils import extract_layer_index
20
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
21
+
22
+
23
+class LinearAttention(PluggableLayer, MambaBase):
24
+ """Base class for Linear attention layer."""
25
+
26
+ def __init__(
27
+ self, config: PretrainedConfig, vllm_config: VllmConfig, prefix: str = ""
28
+ ):
29
+ super().__init__()
30
+ self.layer_idx = extract_layer_index(prefix)
31
+ self.prefix = prefix
32
+ self.model_config = vllm_config.model_config
33
+ self.cache_config = vllm_config.cache_config
34
+ self.quant_config = vllm_config.quant_config
35
+
36
+ self.BLOCK = getattr(config, "block", 256)
37
+ self.hidden_size = config.hidden_size
38
+ self.num_heads = config.num_attention_heads
39
+ self.num_hidden_layers = config.num_hidden_layers
40
+ self.head_dim = (
41
+ config.head_dim
42
+ if hasattr(config, "head_dim")
43
+ else config.hidden_size // self.num_heads
44
+ )
45
+ self.hidden_inner_size = self.head_dim * self.num_heads
46
+
47
+ self.tp_size = get_tensor_model_parallel_world_size()
48
+ self.tp_rank = get_tensor_model_parallel_rank()
49
+ assert self.num_heads % self.tp_size == 0
50
+
51
+ @property
52
+ def mamba_type(self) -> MambaAttentionBackendEnum:
53
+ return MambaAttentionBackendEnum.LINEAR
54
+
55
+ def get_state_dtype(self) -> tuple[torch.dtype]:
56
+ assert self.model_config is not None
57
+ assert self.cache_config is not None
58
+ return MambaStateDtypeCalculator.linear_attention_state_dtype(
59
+ self.model_config.dtype,
60
+ self.cache_config.mamba_cache_dtype,
61
+ )
62
+
63
+ def get_state_shape(self) -> tuple[tuple[int, int, int], ...]:
64
+ return MambaStateShapeCalculator.linear_attention_state_shape(
65
+ num_heads=self.num_heads, tp_size=self.tp_size, head_dim=self.head_dim
66
+ )
vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py
renamed
+18
-68
@@ -7,30 +7,20 @@ from collections.abc import Callable
7
import torch
8
import torch.nn.functional as F
9
from einops import rearrange
10
-from torch import nn
10
12
-from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
13
-from vllm.distributed.parallel_state import (
14
- get_tensor_model_parallel_rank,
15
- get_tensor_model_parallel_world_size,
16
-)
11
+from vllm.config import get_current_vllm_config
12
from vllm.forward_context import ForwardContext, get_forward_context
13
+from vllm.model_executor.custom_op import PluggableLayer
14
from vllm.model_executor.layers.lightning_attn import (
15
lightning_attention,
16
linear_decode_forward_triton,
17
)
18
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
23
-from vllm.model_executor.layers.mamba.abstract import MambaBase
24
-from vllm.model_executor.layers.mamba.mamba_utils import (
25
- MambaStateDtypeCalculator,
26
- MambaStateShapeCalculator,
27
-)
19
+from vllm.model_executor.layers.mamba.linear.base import LinearAttention
20
from vllm.model_executor.layers.minimax_rms_norm import MiniMaxText01RMSNormTP
29
-from vllm.model_executor.layers.quantization import QuantizationConfig
21
from vllm.utils.torch_utils import direct_register_custom_op
22
from vllm.v1.attention.backend import AttentionMetadata
23
from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
33
-from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
24
25
26
def clear_linear_attention_cache_for_new_sequences(
@@ -157,79 +147,39 @@ class MiniMaxText01LinearKernel:
147
return rearrange(output.squeeze(0), "h n d -> n (h d)")
148
149
160
-class MiniMaxText01LinearAttention(nn.Module, MambaBase):
161
- @property
162
- def mamba_type(self) -> MambaAttentionBackendEnum:
163
- return MambaAttentionBackendEnum.LINEAR
164
-
165
- def get_state_dtype(self) -> tuple[torch.dtype]:
166
- assert self.model_config is not None
167
- assert self.cache_config is not None
168
- return MambaStateDtypeCalculator.linear_attention_state_dtype(
169
- self.model_config.dtype,
170
- self.cache_config.mamba_cache_dtype,
171
- )
172
-
173
- def get_state_shape(self) -> tuple[tuple[int, int, int], ...]:
174
- return MambaStateShapeCalculator.linear_attention_state_shape(
175
- num_heads=self.num_heads, tp_size=self.tp_size, head_dim=self.head_dim
176
- )
177
-
150
+@PluggableLayer.register("minimax_text_01_attention")
151
+class MiniMaxText01LinearAttention(LinearAttention):
152
def __init__(
153
self,
180
- hidden_size: int,
181
- hidden_inner_size: int,
182
- num_heads: int,
183
- head_dim: int,
184
- max_position: int,
185
- block_size: int,
186
- num_hidden_layer: int,
187
- model_config: ModelConfig | None = None,
188
- cache_config: CacheConfig | None = None,
189
- quant_config: QuantizationConfig | None = None,
190
- layer_idx: int = 0,
191
- linear_layer_idx: int = 0,
154
+ config,
155
+ vllm_config,
156
prefix: str = "linear_attn",
157
) -> None:
194
- super().__init__()
195
-
196
- self.layer_idx = layer_idx
197
- self.BLOCK = block_size
198
- self.hidden_size = hidden_size
199
- self.num_heads = num_heads
200
- self.head_dim = head_dim
201
- self.total_num_heads = num_heads
202
- self.hidden_inner_size = hidden_inner_size
203
- self.tp_size = get_tensor_model_parallel_world_size()
204
- self.tp_rank = get_tensor_model_parallel_rank()
205
-
206
- assert self.total_num_heads % self.tp_size == 0
207
- self.tp_heads = self.total_num_heads // self.tp_size
158
+ super().__init__(config, vllm_config, prefix)
159
+
160
+ self.tp_heads = self.num_heads // self.tp_size
161
self.qkv_size = self.num_heads * self.head_dim
162
self.tp_hidden = self.head_dim * self.tp_heads
210
- self.model_config = model_config
211
- self.cache_config = cache_config
212
- self.prefix = prefix
163
164
self.qkv_proj = ColumnParallelLinear(
215
- hidden_size,
165
+ self.hidden_size,
166
self.hidden_inner_size * 3,
167
bias=False,
218
- quant_config=quant_config,
168
+ quant_config=self.quant_config,
169
prefix=f"{prefix}.qkv_proj",
170
)
171
self.output_gate = ColumnParallelLinear(
222
- hidden_size,
172
+ self.hidden_size,
173
self.hidden_inner_size,
174
bias=False,
225
- quant_config=quant_config,
175
+ quant_config=self.quant_config,
176
prefix=f"{prefix}.output_gate",
177
)
178
self.out_proj = RowParallelLinear(
179
self.hidden_inner_size,
230
- hidden_size,
180
+ self.hidden_size,
181
bias=False,
232
- quant_config=quant_config,
182
+ quant_config=self.quant_config,
183
prefix=f"{prefix}.out_proj",
184
)
185
self.norm = MiniMaxText01RMSNormTP(
@@ -238,11 +188,11 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
188
)
189
190
slope_rate = MiniMaxText01LinearAttention._build_slope_tensor(self.num_heads)
241
- if num_hidden_layer <= 1:
191
+ if self.num_hidden_layers <= 1:
192
self.slope_rate = slope_rate * (1 + 1e-5)
193
else:
194
self.slope_rate = slope_rate * (
245
- 1 - layer_idx / (num_hidden_layer - 1) + 1e-5
195
+ 1 - self.layer_idx / (self.num_hidden_layers - 1) + 1e-5
196
)
197
self.tp_slope = self.slope_rate[
198
self.tp_rank * self.tp_heads : (self.tp_rank + 1) * self.tp_heads
vllm/model_executor/models/bailing_moe_linear.py
+13
-439
@@ -9,7 +9,7 @@ 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
12
+from vllm.config import CacheConfig, VllmConfig
13
from vllm.distributed import (
14
get_pp_group,
15
get_tensor_model_parallel_rank,
@@ -17,11 +17,6 @@ from vllm.distributed import (
17
)
18
from vllm.forward_context import get_forward_context
19
from vllm.logger import init_logger
20
-from vllm.model_executor.custom_op import PluggableLayer
21
-from vllm.model_executor.layers.fla.ops.layernorm_guard import (
22
- RMSNormGated,
23
- layernorm_fn,
24
-)
20
from vllm.model_executor.layers.fused_moe import (
21
FusedMoE,
22
fused_moe_make_expert_params_mapping,
@@ -30,25 +25,19 @@ from vllm.model_executor.layers.layernorm import RMSNorm
25
from vllm.model_executor.layers.linear import (
26
ColumnParallelLinear,
27
MergedColumnParallelLinear,
33
- QKVParallelLinear,
28
ReplicatedLinear,
29
RowParallelLinear,
30
)
31
from vllm.model_executor.layers.logits_processor import LogitsProcessor
38
-from vllm.model_executor.layers.mamba.abstract import MambaBase
39
-from vllm.model_executor.layers.mamba.linear_attn import (
40
- MiniMaxText01LinearAttention,
41
- MiniMaxText01LinearKernel,
42
- clear_linear_attention_cache_for_new_sequences,
43
- linear_attention_decode,
44
- linear_attention_prefill_and_mix,
32
+from vllm.model_executor.layers.mamba.linear.bailing_linear_attn import (
33
+ BailingMoELinearAttention,
34
+ _build_rope_parameters,
35
)
36
from vllm.model_executor.layers.mamba.mamba_utils import (
37
MambaStateCopyFuncCalculator,
38
MambaStateDtypeCalculator,
39
MambaStateShapeCalculator,
40
)
51
-from vllm.model_executor.layers.minimax_rms_norm import MiniMaxText01RMSNormTP
41
from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
42
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
43
from vllm.model_executor.layers.rotary_embedding import get_rope
@@ -63,8 +52,6 @@ from vllm.model_executor.model_loader.weight_utils import (
52
from vllm.model_executor.models.bailing_moe import BailingMLP
53
from vllm.sequence import IntermediateTensors
54
from vllm.v1.attention.backend import AttentionMetadata
66
-from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
67
-from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
55
56
from .interfaces import HasInnerState, IsHybrid, SupportsPP
57
from .utils import (
@@ -87,25 +74,6 @@ def is_linear_layer(layer_idx, layer_group_size):
74
return False
75
76
90
-def _build_rope_parameters(config: PretrainedConfig) -> dict | None:
91
- rope_parameters = copy.deepcopy(getattr(config, "rope_parameters", None)) or {}
92
- if "rope_theta" not in rope_parameters and hasattr(config, "rope_theta"):
93
- rope_parameters["rope_theta"] = config.rope_theta
94
- if "partial_rotary_factor" not in rope_parameters and hasattr(
95
- config, "partial_rotary_factor"
96
- ):
97
- rope_parameters["partial_rotary_factor"] = config.partial_rotary_factor
98
-
99
- rope_scaling = getattr(config, "rope_scaling", None)
100
- if isinstance(rope_scaling, dict):
101
- rope_scaling = copy.deepcopy(rope_scaling)
102
- if "type" in rope_scaling and "rope_type" not in rope_scaling:
103
- rope_scaling["rope_type"] = rope_scaling.pop("type")
104
- rope_parameters.update(rope_scaling)
105
-
106
- return rope_parameters or None
107
-
108
-
77
class BailingMoeV25MLAAttention(nn.Module):
78
"""
79
MLA Attention for BailingMoeV2.5 full attention layers.
@@ -397,400 +365,15 @@ class BailingMoeV25(nn.Module):
365
return final_hidden_states.view(num_tokens, hidden_size)
366
367
400
-BailingRMSNormTP = MiniMaxText01RMSNormTP
401
-
402
-
403
-class BailingGroupRMSNormGate(RMSNormGated):
404
- def __init__(
405
- self,
406
- hidden_size,
407
- eps=1e-5,
408
- group_size=None,
409
- norm_before_gate=True,
410
- device=None,
411
- dtype=None,
412
- ):
413
- super().__init__(
414
- hidden_size,
415
- eps=eps,
416
- group_size=group_size,
417
- norm_before_gate=norm_before_gate,
418
- device=device,
419
- dtype=dtype,
420
- activation="sigmoid",
421
- )
422
- # Add custom weight loader for TP sharding
423
- self.weight.weight_loader = self._weight_loader
424
-
425
- @staticmethod
426
- def _weight_loader(param: torch.nn.Parameter, loaded_weight: torch.Tensor) -> None:
427
- """Load weight with TP sharding."""
428
- tp_size = get_tensor_model_parallel_world_size()
429
- tp_rank = get_tensor_model_parallel_rank()
430
- shard_size = loaded_weight.shape[0] // tp_size
431
- shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
432
- param.data.copy_(loaded_weight[shard].contiguous())
433
-
434
-
435
-# --8<-- [start:bailing_moe_linear_attention]
436
-@PluggableLayer.register("bailing_moe_linear_attention")
437
-class BailingMoELinearAttention(PluggableLayer, MambaBase):
438
- """Pluggable Bailing MoE Linear Attention layer which allows OOT backends
439
- to add custom implementations.
440
-
441
- This implements the linear attention mechanism from sglang, adapted for
442
- vLLM's v1 engine with MambaBase interface support.
443
- """
444
-
445
- # --8<-- [end:bailing_moe_linear_attention]
446
-
447
- @property
448
- def mamba_type(self) -> MambaAttentionBackendEnum:
449
- return MambaAttentionBackendEnum.LINEAR
450
-
451
- def get_state_shape(self) -> tuple[tuple[int, ...], ...]:
452
- """Return state shape for linear attention cache.
453
-
454
- Must match the calculation in get_mamba_state_shape_from_config.
455
- """
456
- return MambaStateShapeCalculator.linear_attention_state_shape(
457
- num_heads=self.total_num_heads,
458
- tp_size=self.tp_size,
459
- head_dim=self.head_dim,
460
- )
461
-
462
- def get_state_dtype(self) -> tuple[torch.dtype, ...]:
463
- """Return state dtype for linear attention cache.
464
-
465
- Must match the calculation in get_mamba_state_dtype_from_config.
466
- """
467
- return MambaStateDtypeCalculator.linear_attention_state_dtype(
468
- self.model_config.dtype,
469
- self.cache_config.mamba_cache_dtype,
470
- )
471
-
472
- def __init__(
473
- self,
474
- config: PretrainedConfig,
475
- quant_config: QuantizationConfig | None = None,
476
- layer_id: int = 0,
477
- prefix: str = "linear_attn",
478
- model_config: ModelConfig | None = None,
479
- cache_config: CacheConfig | None = None,
480
- ):
481
- super().__init__()
482
-
483
- self.layer_id = layer_id
484
- self.hidden_size = config.hidden_size
485
- self.total_num_heads = config.num_attention_heads
486
- self.total_kv_heads = config.num_attention_heads # MHA
487
- self.tp_size = get_tensor_model_parallel_world_size()
488
- self.tp_rank = get_tensor_model_parallel_rank()
489
- self.model_config = model_config
490
- self.cache_config = cache_config
491
- self.prefix = prefix
492
-
493
- self.head_dim = (
494
- config.head_dim
495
- if hasattr(config, "head_dim")
496
- else config.hidden_size // self.total_num_heads
497
- )
498
-
499
- self.hidden_inner_size = self.head_dim * self.total_num_heads
500
- self.scaling = self.head_dim**-0.5
501
-
502
- assert self.total_num_heads % self.tp_size == 0
503
- self.tp_heads = self.total_num_heads // self.tp_size
504
-
505
- self.max_position_embeddings = config.max_position_embeddings
506
- self.rope_theta = getattr(config, "rope_theta", 600000)
507
-
508
- self.tp_kv_heads = self.total_kv_heads // self.tp_size
509
- self.q_size_per_rank = self.head_dim * self.tp_heads
510
- self.kv_size_per_rank = self.head_dim * self.tp_kv_heads
511
-
512
- self.use_qk_norm = getattr(config, "use_qk_norm", False)
513
- self.linear_backend = "minimax"
514
- self.linear_scale = self.linear_backend == "minimax"
515
- self.linear_rope = getattr(config, "linear_rope", True)
516
- if hasattr(config, "use_linear_silu"):
517
- self.linear_silu = config.use_linear_silu
518
- elif hasattr(config, "linear_silu"):
519
- self.linear_silu = config.linear_silu
520
- else:
521
- self.linear_silu = False
522
-
523
- # Block size for lightning attention
524
- self.BLOCK = getattr(config, "block", 256)
525
-
526
- self.query_key_value = QKVParallelLinear(
527
- self.hidden_size,
528
- self.head_dim,
529
- self.total_num_heads,
530
- self.total_num_heads, # MHA: kv_heads = num_heads
531
- bias=(config.use_bias or config.use_qkv_bias),
532
- quant_config=quant_config,
533
- prefix=f"{prefix}.query_key_value",
534
- )
535
-
536
- if self.use_qk_norm:
537
- self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
538
- self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
539
-
540
- self.g_proj = ColumnParallelLinear(
541
- self.hidden_size,
542
- self.hidden_inner_size,
543
- bias=False,
544
- quant_config=quant_config,
545
- prefix=f"{prefix}.g_proj",
546
- )
547
- self.dense = RowParallelLinear(
548
- self.hidden_inner_size,
549
- self.hidden_size,
550
- bias=config.use_bias,
551
- quant_config=quant_config,
552
- prefix=f"{prefix}.dense",
553
- reduce_results=True,
554
- )
555
-
556
- self.group_norm_size = getattr(config, "group_norm_size", 1)
557
- self.rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-5))
558
- assert self.tp_size <= self.group_norm_size, (
559
- "tp_size must be <= group_norm_size for local rms norm"
560
- )
561
- assert self.group_norm_size % self.tp_size == 0, (
562
- "group_norm_size must be divisible by tp_size"
563
- )
564
-
565
- # When group_norm_size == 1, group_size equals hidden_size // tp_size
566
- self.g_norm = BailingGroupRMSNormGate(
567
- hidden_size=self.hidden_inner_size // self.tp_size,
568
- eps=self.rms_norm_eps,
569
- group_size=(
570
- self.hidden_inner_size // self.group_norm_size
571
- if self.group_norm_size > 1
572
- else self.hidden_inner_size // self.tp_size
573
- ),
574
- )
575
-
576
- # use fp32 rotary embedding
577
- rope_parameters = _build_rope_parameters(config)
578
-
579
- self.rotary_emb = get_rope(
580
- self.head_dim,
581
- max_position=self.max_position_embeddings,
582
- is_neox_style=True,
583
- rope_parameters=rope_parameters or None,
584
- )
585
-
586
- # Build slope tensor for linear attention decay
587
- num_hidden_layers = config.num_hidden_layers
588
- slope_rate = MiniMaxText01LinearAttention._build_slope_tensor(
589
- self.total_num_heads
590
- )
591
- if num_hidden_layers <= 1:
592
- self.slope_rate = slope_rate * (1 + 1e-5)
593
- else:
594
- self.slope_rate = slope_rate * (
595
- 1 - layer_id / (num_hidden_layers - 1) + 1e-5
596
- )
597
- self.tp_slope = self.slope_rate[
598
- self.tp_rank * self.tp_heads : (self.tp_rank + 1) * self.tp_heads
599
- ].contiguous()
600
-
601
- # Register for compilation
602
- compilation_config = get_current_vllm_config().compilation_config
603
- if prefix in compilation_config.static_forward_context:
604
- raise ValueError(f"Duplicate layer name: {prefix}")
605
- compilation_config.static_forward_context[prefix] = self
606
-
607
- @staticmethod
608
- def weight_direct_load(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
609
- """Load weight for linear attention layers.
610
-
611
- For FP8 quantized parameters, we need to use the weight_loader if available,
612
- as it handles special cases like tensor parallelism sharding.
613
- """
614
- # Check if param has a weight_loader (for vLLM ModelWeightParameter)
615
- weight_loader = getattr(param, "weight_loader", None)
616
- if weight_loader is not None:
617
- # Use the weight_loader which handles TP sharding and quantization
618
- weight_loader(param, loaded_weight)
619
- else:
620
- # Fall back to direct copy for standard tensors
621
- assert param.size() == loaded_weight.size(), (
622
- f"Shape mismatch: {param.shape} vs {loaded_weight.shape}"
623
- )
624
- param.data.copy_(loaded_weight)
625
-
626
- def forward(
627
- self,
628
- hidden_states: torch.Tensor,
629
- output: torch.Tensor,
630
- positions: torch.Tensor,
631
- ) -> None:
632
- """Forward method called by torch.ops.vllm.linear_attention"""
633
- torch.ops.vllm.linear_attention(
634
- hidden_states,
635
- output,
636
- positions,
637
- self.prefix,
638
- )
639
-
640
- def _forward(
641
- self,
642
- hidden_states: torch.Tensor,
643
- output: torch.Tensor,
644
- positions: torch.Tensor,
645
- ) -> None:
646
- """Actual forward implementation."""
647
- forward_context = get_forward_context()
648
- attn_metadata: AttentionMetadata = forward_context.attn_metadata
649
- if attn_metadata is not None:
650
- assert isinstance(attn_metadata, dict)
651
- attn_metadata = attn_metadata[self.prefix]
652
- assert isinstance(attn_metadata, LinearAttentionMetadata)
653
- num_actual_tokens = (
654
- attn_metadata.num_prefill_tokens + attn_metadata.num_decode_tokens
655
- )
656
- else:
657
- num_actual_tokens = hidden_states.shape[0]
658
-
659
- # QKV projection
660
- qkv, _ = self.query_key_value(hidden_states[:num_actual_tokens])
661
-
662
- # use rotary_emb support fp32
663
- qkv = qkv.to(torch.float32)
664
- if self.linear_silu:
665
- qkv = F.silu(qkv)
666
-
667
- # Split q, k, v
668
- q, k, v = torch.split(
669
- qkv,
670
- [self.q_size_per_rank, self.kv_size_per_rank, self.kv_size_per_rank],
671
- dim=-1,
672
- )
673
-
674
- # Apply QK norm if needed
675
- if self.use_qk_norm:
676
- q = q.reshape(-1, self.tp_heads, self.head_dim)
677
- k = k.reshape(-1, self.tp_kv_heads, self.head_dim)
678
- q = layernorm_fn(
679
- q,
680
- self.query_layernorm.weight.data,
681
- bias=None,
682
- eps=self.rms_norm_eps,
683
- is_rms_norm=True,
684
- )
685
- k = layernorm_fn(
686
- k,
687
- self.key_layernorm.weight.data,
688
- bias=None,
689
- eps=self.rms_norm_eps,
690
- is_rms_norm=True,
691
- )
692
- q = q.reshape(-1, self.q_size_per_rank)
693
- k = k.reshape(-1, self.kv_size_per_rank)
694
-
695
- # Apply rotary embeddings
696
- if self.linear_rope:
697
- q, k = self.rotary_emb(positions[:num_actual_tokens], q, k)
698
-
699
- # Reshape to [batch, heads, seq_len, head_dim]
700
- q = q.view((qkv.shape[0], self.tp_heads, self.head_dim))
701
- k = k.view((qkv.shape[0], self.tp_kv_heads, self.head_dim))
702
- v = v.view((qkv.shape[0], self.tp_kv_heads, self.head_dim))
703
-
704
- # Apply scaling if using minimax backend
705
- if self.linear_scale:
706
- q = q * self.scaling
707
-
708
- # Get KV cache and state indices
709
- if attn_metadata is not None:
710
- kv_cache = self.kv_cache[0]
711
- state_indices_tensor = attn_metadata.state_indices_tensor
712
- clear_linear_attention_cache_for_new_sequences(
713
- kv_cache, state_indices_tensor, attn_metadata
714
- )
715
-
716
- # Compute attention
717
- decode_only = getattr(attn_metadata, "num_prefills", 0) == 0
718
- if attn_metadata is None:
719
- hidden = torch.empty(
720
- (q.shape[0], q.shape[1] * q.shape[2]), device=q.device, dtype=q.dtype
721
- )
722
- else:
723
- if not decode_only:
724
- hidden = self._prefill_and_mix_infer(
725
- q, k, v, kv_cache, state_indices_tensor, attn_metadata
726
- )
727
- else:
728
- hidden = self._decode_infer(
729
- q, k, v, kv_cache, state_indices_tensor, attn_metadata
730
- )
731
-
732
- # Apply group norm and gate (matching SGLang behavior)
733
- gate, _ = self.g_proj(hidden_states[:num_actual_tokens])
734
-
735
- if self.group_norm_size > 1:
736
- hidden = self.g_norm(hidden, gate)
737
- else:
738
- hidden = self.g_norm(hidden)
739
- hidden = F.sigmoid(gate) * hidden
740
-
741
- hidden = hidden.to(hidden_states.dtype)
742
-
743
- # Output projection
744
- dense_out, _ = self.dense(hidden)
745
- output[:num_actual_tokens] = dense_out
746
-
747
- def _prefill_and_mix_infer(
748
- self, q, k, v, kv_cache, state_indices_tensor, attn_metadata
749
- ):
750
- """Handle prefill (mixed with decode if any)."""
751
- return linear_attention_prefill_and_mix(
752
- q=q,
753
- k=k,
754
- v=v,
755
- kv_cache=kv_cache,
756
- state_indices_tensor=state_indices_tensor,
757
- attn_metadata=attn_metadata,
758
- slope_rate=self.tp_slope,
759
- block_size=self.BLOCK,
760
- decode_fn=self._decode_infer,
761
- prefix_fn=MiniMaxText01LinearKernel.jit_linear_forward_prefix,
762
- layer_idx=self.layer_id,
763
- )
764
-
765
- def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor, attn_metadata):
766
- """Handle decode (single token per sequence)."""
767
- hidden = linear_attention_decode(
768
- q,
769
- k,
770
- v,
771
- kv_cache,
772
- self.tp_slope,
773
- state_indices_tensor,
774
- q_start=0,
775
- q_end=attn_metadata.num_decode_tokens,
776
- slot_start=0,
777
- slot_end=attn_metadata.num_decodes,
778
- block_size=32,
779
- )
780
- return hidden
781
-
782
-
368
class BailingMoeV25DecoderLayer(nn.Module):
369
"""Decoder layer supporting both linear and full attention."""
370
371
def __init__(
372
self,
373
config: PretrainedConfig,
789
- quant_config: QuantizationConfig | None = None,
790
- layer_id: int = 0,
374
+ vllm_config: VllmConfig,
375
prefix: str = "layer",
792
- model_config: ModelConfig | None = None,
793
- cache_config: CacheConfig | None = None,
376
+ layer_id: int = 0,
377
) -> None:
378
super().__init__()
379
self.layer_id = layer_id
@@ -802,19 +385,16 @@ class BailingMoeV25DecoderLayer(nn.Module):
385
if self.attention_type == 0: # Linear attention
386
self.self_attn = BailingMoELinearAttention(
387
config,
805
- quant_config=quant_config,
806
- layer_id=layer_id,
388
+ vllm_config,
389
prefix=f"{prefix}.self_attn",
808
- model_config=model_config,
809
- cache_config=cache_config,
390
)
391
else: # Full attention
392
self.self_attn = BailingMoeV25MLAAttention(
393
config,
814
- quant_config=quant_config,
394
+ quant_config=vllm_config.quant_config,
395
layer_id=layer_id,
396
prefix=f"{prefix}.self_attn",
817
- cache_config=cache_config,
397
+ cache_config=vllm_config.cache_config,
398
)
399
400
# MLP/MoE
@@ -825,7 +405,7 @@ class BailingMoeV25DecoderLayer(nn.Module):
405
if is_moe_layer:
406
self.mlp = BailingMoeV25(
407
config,
828
- quant_config=quant_config,
408
+ quant_config=vllm_config.quant_config,
409
layer_id=layer_id,
410
prefix=f"{prefix}.mlp",
411
)
@@ -833,7 +413,7 @@ class BailingMoeV25DecoderLayer(nn.Module):
413
self.mlp = BailingMLP(
414
intermediate_size=config.intermediate_size,
415
config=config,
836
- quant_config=quant_config,
416
+ quant_config=vllm_config.quant_config,
417
reduce_results=True,
418
prefix=f"{prefix}.mlp",
419
)
@@ -896,10 +476,6 @@ class BailingMoeV25Model(nn.Module):
476
):
477
super().__init__()
478
config = vllm_config.model_config.hf_config
899
- model_config = vllm_config.model_config
900
- quant_config = vllm_config.quant_config
901
- cache_config = vllm_config.cache_config
902
-
479
self.config = config
480
self.vocab_size = config.vocab_size
481
self.embed_dim = config.hidden_size
@@ -934,11 +510,9 @@ class BailingMoeV25Model(nn.Module):
510
511
return BailingMoeV25DecoderLayer(
512
config=layer_config,
937
- quant_config=quant_config,
938
- layer_id=layer_idx,
513
+ vllm_config=vllm_config,
514
prefix=prefix,
940
- model_config=model_config,
941
- cache_config=cache_config,
515
+ layer_id=layer_idx,
516
)
517
518
self.start_layer, self.end_layer, self.layers = make_layers(
vllm/model_executor/models/minimax_text_01.py
+14
-35
@@ -15,7 +15,7 @@ from torch import nn
15
from transformers import MiniMaxConfig
16
17
from vllm.compilation.decorators import support_torch_compile
18
-from vllm.config import CacheConfig, ModelConfig, VllmConfig
18
+from vllm.config import CacheConfig, VllmConfig
19
from vllm.distributed.parallel_state import (
20
get_pp_group,
21
get_tensor_model_parallel_rank,
@@ -35,7 +35,9 @@ from vllm.model_executor.layers.linear import (
35
RowParallelLinear,
36
)
37
from vllm.model_executor.layers.logits_processor import LogitsProcessor
38
-from vllm.model_executor.layers.mamba.linear_attn import MiniMaxText01LinearAttention
38
+from vllm.model_executor.layers.mamba.linear.minimax_linear_attn import (
39
+ MiniMaxText01LinearAttention,
40
+)
41
from vllm.model_executor.layers.mamba.mamba_utils import (
42
MambaStateCopyFunc,
43
MambaStateCopyFuncCalculator,
@@ -277,9 +279,7 @@ class MiniMaxText01DecoderLayer(nn.Module):
279
def __init__(
280
self,
281
config: MiniMaxConfig,
280
- model_config: ModelConfig | None = None,
281
- cache_config: CacheConfig | None = None,
282
- quant_config: QuantizationConfig | None = None,
282
+ vllm_config: VllmConfig,
283
expert_num: int = 1,
284
layer_id: int = None,
285
linear_layer_id: int | None = None,
@@ -303,25 +303,9 @@ class MiniMaxText01DecoderLayer(nn.Module):
303
config.max_position_embeddings, config.max_model_len
304
)
305
if config.attention_type == 0:
306
- use_headxdim = True
307
- hidden_inner = (
308
- head_dim * config.num_attention_heads
309
- if use_headxdim
310
- else config.hidden_size
311
- )
306
self.self_attn = MiniMaxText01LinearAttention(
313
- hidden_size=self.hidden_size,
314
- hidden_inner_size=hidden_inner,
315
- num_heads=config.num_attention_heads,
316
- head_dim=head_dim,
317
- max_position=max_position_embeddings,
318
- block_size=config.block if hasattr(config, "block") else 256,
319
- num_hidden_layer=config.num_hidden_layers,
320
- model_config=model_config,
321
- cache_config=cache_config,
322
- quant_config=quant_config,
323
- layer_idx=self._ilayer,
324
- linear_layer_idx=linear_layer_id,
307
+ config,
308
+ vllm_config,
309
prefix=prefix,
310
)
311
elif config.attention_type == 1:
@@ -333,9 +317,9 @@ class MiniMaxText01DecoderLayer(nn.Module):
317
max_position=max_position_embeddings,
318
rope_parameters=config.rope_parameters,
319
sliding_window=config.sliding_window,
336
- quant_config=quant_config,
320
+ quant_config=vllm_config.quant_config,
321
layer_idx=self._ilayer,
338
- cache_config=cache_config,
322
+ cache_config=vllm_config.cache_config,
323
prefix=prefix,
324
)
325
else:
@@ -348,7 +332,7 @@ class MiniMaxText01DecoderLayer(nn.Module):
332
self.mlp = MiniMaxText01MLP(
333
hidden_size=self.hidden_size,
334
intermediate_size=config.intermediate_size,
351
- quant_config=quant_config,
335
+ quant_config=vllm_config.quant_config,
336
layer_idx=self._ilayer,
337
prefix=prefix,
338
)
@@ -359,7 +343,7 @@ class MiniMaxText01DecoderLayer(nn.Module):
343
hidden_size=config.hidden_size,
344
intermediate_size=config.intermediate_size,
345
layer_idx=self._ilayer,
362
- quant_config=quant_config,
346
+ quant_config=vllm_config.quant_config,
347
prefix=prefix,
348
)
349
@@ -410,7 +394,7 @@ class MiniMaxText01DecoderLayer(nn.Module):
394
self.shared_mlp = MiniMaxText01MLP(
395
hidden_size=self.hidden_size,
396
intermediate_size=shared_intermediate,
413
- quant_config=quant_config,
397
+ quant_config=vllm_config.quant_config,
398
layer_idx=self._ilayer,
399
prefix=prefix,
400
)
@@ -418,7 +402,7 @@ class MiniMaxText01DecoderLayer(nn.Module):
402
self.hidden_size,
403
1,
404
bias=False,
421
- quant_config=quant_config,
405
+ quant_config=vllm_config.quant_config,
406
params_dtype=torch.float32,
407
)
408
self.coefficient.weight.weight_loader = self.shared_moe_coefficient_loader
@@ -496,9 +480,6 @@ class MiniMaxText01Model(nn.Module):
480
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
481
super().__init__()
482
config: MiniMaxConfig = vllm_config.model_config.hf_config
499
- model_config = vllm_config.model_config
500
- quant_config = vllm_config.quant_config
501
- cache_config = vllm_config.cache_config
483
scheduler_config = vllm_config.scheduler_config
484
self.config = config
485
self.CONCAT_FFN = True
@@ -541,10 +522,8 @@ class MiniMaxText01Model(nn.Module):
522
layer_config.layer_idx = layer_idx
523
524
decoder_kwargs = {
544
- "quant_config": quant_config,
525
"layer_id": layer_idx,
546
- "model_config": model_config,
547
- "cache_config": cache_config,
526
+ "vllm_config": vllm_config,
527
}
528
529
if layer_config.attention_type == 0: