@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4071c76cf3
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 14, 05:30 UTC

[V1] [Hybrid] Move MiniMaxLinearAttention into layers/mamba (#23831)

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

Thomas Parnell committed Aug 30, 2025 at 09:16 UTC 4071c76cf3cff46e14630f1e66cbf006b6eb51d3
2 files changed +448 -410
vllm/model_executor/layers/mamba/linear_attn.py new
+442
@@ -0,0 +1,442 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import math
5 +from typing import TYPE_CHECKING, Optional, Union
6 +
7 +if TYPE_CHECKING:
8 + from vllm.attention.backends.abstract import AttentionBackend
9 +
10 +from typing import TYPE_CHECKING
11 +
12 +import torch
13 +import torch.distributed
14 +import torch.nn.functional as F
15 +from einops import rearrange
16 +from torch import nn
17 +
18 +from vllm import envs
19 +from vllm.attention import AttentionMetadata
20 +from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
21 +from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
22 +from vllm.distributed.parallel_state import (
23 + get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)
24 +from vllm.forward_context import ForwardContext, get_forward_context
25 +from vllm.model_executor.custom_op import CustomOp
26 +from vllm.model_executor.layers.lightning_attn import (
27 + lightning_attention, linear_decode_forward_triton)
28 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,
29 + RowParallelLinear)
30 +from vllm.model_executor.layers.mamba.abstract import MambaBase
31 +from vllm.model_executor.layers.mamba.mamba_utils import (
32 + MambaStateDtypeCalculator, MambaStateShapeCalculator)
33 +from vllm.model_executor.layers.quantization.base_config import (
34 + QuantizationConfig)
35 +from vllm.platforms import current_platform
36 +from vllm.utils import direct_register_custom_op
37 +from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
38 +
39 +if TYPE_CHECKING:
40 + from vllm.attention.backends.abstract import AttentionBackend
41 +
42 +import torch
43 +import torch.distributed
44 +
45 +from vllm.model_executor.models.minimax_cache import MinimaxCacheParams
46 +
47 +
48 +class MiniMaxText01RMSNormTP(CustomOp):
49 + name = "MiniMaxText01RMSNormTP"
50 +
51 + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
52 + super().__init__()
53 + self.tp_world = get_tensor_model_parallel_world_size()
54 + self.tp_rank = get_tensor_model_parallel_rank()
55 + self.weight = nn.Parameter(torch.ones(int(hidden_size /
56 + self.tp_world)))
57 +
58 + self.weight.weight_loader = self.weight_loader
59 + self.variance_epsilon = eps
60 + return
61 +
62 + @staticmethod
63 + def weight_loader(
64 + param: nn.Parameter,
65 + loaded_weight: torch.Tensor,
66 + ) -> None:
67 + tp_world = get_tensor_model_parallel_world_size()
68 + tp_rank = get_tensor_model_parallel_rank()
69 +
70 + shard_size = loaded_weight.shape[0] // tp_world
71 + shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
72 + param.data.copy_(loaded_weight[shard])
73 + return
74 +
75 + def _forward(
76 + self,
77 + x: torch.Tensor,
78 + ) -> torch.Tensor:
79 + orig_dtype = x.dtype
80 + x = x.to(torch.float32)
81 + variance = x.pow(2).mean(dim=-1, keepdim=True, dtype=torch.float32)
82 + if self.tp_world > 1:
83 + variance = tensor_model_parallel_all_reduce(
84 + variance) / self.tp_world
85 + x = x * torch.rsqrt(variance + self.variance_epsilon)
86 +
87 + weight = self.weight
88 + if x.size(-1) != self.weight.size(0):
89 + if self.weight.size(0) < x.size(-1):
90 + repeat_count = (x.size(-1) + self.weight.size(0)) // x.size(-1)
91 + full_weight = self.weight.repeat(repeat_count)
92 + weight = full_weight[:x.size(-1)]
93 + else:
94 + weight = self.weight[:x.size(-1)]
95 +
96 + x = x.to(orig_dtype) * weight
97 + return x
98 +
99 + def forward(
100 + self,
101 + x: torch.Tensor,
102 + residual: Optional[torch.Tensor] = None,
103 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
104 + assert residual is None, "RMSNorm does not support residual connection."
105 + return self._forward(x)
106 +
107 +
108 +class MiniMaxText01LinearKernel:
109 +
110 + @staticmethod
111 + def jit_linear_forward_prefix(q: torch.Tensor,
112 + k: torch.Tensor,
113 + v: torch.Tensor,
114 + kv_caches: torch.Tensor,
115 + slope_rate: torch.Tensor,
116 + block_size: int,
117 + layer_idx: Optional[int] = None,
118 + **kwargs) -> torch.Tensor:
119 +
120 + slope_rate = slope_rate.to(torch.float32)
121 + should_pad_dim = q.dim() == 3
122 + if should_pad_dim:
123 + q = q.unsqueeze(0)
124 + k = k.unsqueeze(0)
125 + v = v.unsqueeze(0)
126 + b, h, n, d = q.shape
127 + e = d
128 + kv_history = kv_caches.reshape(1, h, d, e).contiguous()
129 + output, kv_history = lightning_attention(q,
130 + k,
131 + v,
132 + slope_rate,
133 + block_size=block_size,
134 + kv_history=kv_history)
135 + kv_caches.copy_(kv_history[:, :, -1, :, :].reshape(h, d, e))
136 + assert output.shape[0] == 1, "batch size must be 1"
137 + return rearrange(output.squeeze(0), "h n d -> n (h d)")
138 +
139 +
140 +class MiniMaxText01LinearAttention(nn.Module, MambaBase):
141 +
142 + @property
143 + def mamba_type(self) -> str:
144 + return "linear_attention"
145 +
146 + def get_attn_backend(self) -> type["AttentionBackend"]:
147 + from vllm.v1.attention.backends.linear_attn import (
148 + LinearAttentionBackend)
149 + return LinearAttentionBackend
150 +
151 + def get_state_dtype(self) -> tuple[torch.dtype]:
152 + assert self.model_config is not None
153 + assert self.cache_config is not None
154 + return MambaStateDtypeCalculator.linear_attention_state_dtype(
155 + self.model_config.dtype,
156 + self.cache_config.mamba_cache_dtype,
157 + )
158 +
159 + def get_state_shape(self) -> tuple[tuple[int, int, int], ...]:
160 + return MambaStateShapeCalculator.linear_attention_state_shape(
161 + num_heads=self.num_heads,
162 + tp_size=self.tp_size,
163 + head_dim=self.head_dim)
164 +
165 + def __init__(
166 + self,
167 + hidden_size: int,
168 + hidden_inner_size: int,
169 + num_heads: int,
170 + head_dim: int,
171 + max_position: int,
172 + block_size: int,
173 + num_hidden_layer: int,
174 + model_config: Optional[ModelConfig] = None,
175 + cache_config: Optional[CacheConfig] = None,
176 + quant_config: Optional[QuantizationConfig] = None,
177 + layer_idx: int = 0,
178 + linear_layer_idx: int = 0,
179 + prefix: str = "linear_attn",
180 + ) -> None:
181 + super().__init__()
182 +
183 + self.layer_idx = layer_idx
184 + self.BLOCK = block_size
185 + self.hidden_size = hidden_size
186 + self.num_heads = num_heads
187 + self.head_dim = head_dim
188 + self.total_num_heads = num_heads
189 + self.hidden_inner_size = hidden_inner_size
190 + self.tp_size = get_tensor_model_parallel_world_size()
191 + self.tp_rank = get_tensor_model_parallel_rank()
192 +
193 + assert self.total_num_heads % self.tp_size == 0
194 + self.tp_heads = self.total_num_heads // self.tp_size
195 + self.qkv_size = self.num_heads * self.head_dim
196 + self.tp_hidden = self.head_dim * self.tp_heads
197 + self.model_config = model_config
198 + self.cache_config = cache_config
199 + self.prefix = prefix
200 +
201 + self.qkv_proj = ColumnParallelLinear(
202 + hidden_size,
203 + self.hidden_inner_size * 3,
204 + bias=False,
205 + quant_config=quant_config,
206 + prefix=f"{prefix}.qkv_proj",
207 + )
208 + self.output_gate = ColumnParallelLinear(
209 + hidden_size,
210 + self.hidden_inner_size,
211 + bias=False,
212 + quant_config=quant_config,
213 + prefix=f"{prefix}.output_gate",
214 + )
215 + self.out_proj = RowParallelLinear(
216 + self.hidden_inner_size,
217 + hidden_size,
218 + bias=False,
219 + quant_config=quant_config,
220 + prefix=f"{prefix}.out_proj",
221 + )
222 + self.norm = MiniMaxText01RMSNormTP(
223 + self.hidden_inner_size,
224 + eps=1e-5,
225 + )
226 +
227 + slope_rate = MiniMaxText01LinearAttention._build_slope_tensor(
228 + self.num_heads)
229 + if num_hidden_layer <= 1:
230 + self.slope_rate = slope_rate * (1 + 1e-5)
231 + else:
232 + self.slope_rate = slope_rate * (1 - layer_idx /
233 + (num_hidden_layer - 1) + 1e-5)
234 + self.tp_slope = self.slope_rate[self.tp_rank *
235 + self.tp_heads:(self.tp_rank + 1) *
236 + self.tp_heads].contiguous()
237 +
238 + if envs.VLLM_USE_V1:
239 + compilation_config = get_current_vllm_config().compilation_config
240 + if prefix in compilation_config.static_forward_context:
241 + raise ValueError(f"Duplicate layer name: {prefix}")
242 + compilation_config.static_forward_context[prefix] = self
243 +
244 + @staticmethod
245 + def weight_direct_load(param: torch.Tensor,
246 + loaded_weight: torch.Tensor) -> None:
247 + assert param.size() == loaded_weight.size()
248 + param.data.copy_(loaded_weight)
249 + return
250 +
251 + @staticmethod
252 + def _build_slope_tensor(n_attention_heads: int):
253 +
254 + def get_slopes(n):
255 +
256 + def get_slopes_power_of_2(n):
257 + start = 2**(-(2**-(math.log2(n) - 3)))
258 + ratio = start
259 + return [start * ratio**i for i in range(n)]
260 +
261 + if math.log2(n).is_integer():
262 + return get_slopes_power_of_2(n)
263 + else:
264 + closest_power_of_2 = 2**math.floor(math.log2(n))
265 + return (get_slopes_power_of_2(closest_power_of_2) + get_slopes(
266 + 2 * closest_power_of_2)[0::2][:n - closest_power_of_2])
267 +
268 + slopes = torch.tensor(get_slopes(n_attention_heads),
269 + dtype=torch.float32).reshape(
270 + n_attention_heads, 1, 1)
271 + return slopes
272 +
273 + def _prefill_and_mix_infer(self, q, k, v, kv_cache, state_indices_tensor,
274 + attn_metadata):
275 + hidden = []
276 + for _prefill_idx in range(getattr(attn_metadata, "num_prefills", 0)):
277 + if _prefill_idx >= len(attn_metadata.query_start_loc):
278 + break
279 + if _prefill_idx >= len(state_indices_tensor):
280 + break
281 + # prefills are packed at end of batch in V1
282 + offset = attn_metadata.num_decode_tokens if envs.VLLM_USE_V1 else 0
283 + _start = attn_metadata.query_start_loc[offset + _prefill_idx]
284 + _end = attn_metadata.query_start_loc[offset + _prefill_idx + 1]
285 + slot_id = state_indices_tensor[offset + _prefill_idx]
286 + qs = q[_start:_end].transpose(0, 1).contiguous()
287 + ks = k[_start:_end].transpose(0, 1).contiguous()
288 + vs = v[_start:_end].transpose(0, 1).contiguous()
289 + slice_layer_cache = kv_cache[slot_id, ...]
290 +
291 + out_slice = MiniMaxText01LinearKernel.jit_linear_forward_prefix(
292 + qs,
293 + ks,
294 + vs,
295 + slice_layer_cache,
296 + self.tp_slope,
297 + self.BLOCK,
298 + layer_idx=self.layer_idx)
299 + hidden.append(out_slice.contiguous())
300 + if attn_metadata.num_decode_tokens > 0:
301 + hidden_decode = self._decode_infer(q, k, v, kv_cache,
302 + state_indices_tensor,
303 + attn_metadata)
304 + if envs.VLLM_USE_V1:
305 + hidden.insert(0, hidden_decode)
306 + else:
307 + hidden.append(hidden_decode)
308 +
309 + if not hidden:
310 + return torch.empty((0, q.size(-1)), device=q.device, dtype=q.dtype)
311 +
312 + hidden = torch.concat(hidden, dim=0).contiguous()
313 + return hidden
314 +
315 + def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor,
316 + attn_metadata):
317 + if not envs.VLLM_USE_V1:
318 + q = q[attn_metadata.num_prefill_tokens:].unsqueeze(2).contiguous()
319 + k = k[attn_metadata.num_prefill_tokens:].unsqueeze(2).contiguous()
320 + v = v[attn_metadata.num_prefill_tokens:].unsqueeze(2).contiguous()
321 + num_prefills = getattr(attn_metadata, "num_prefills", 0)
322 + slot_id = state_indices_tensor[num_prefills:]
323 + else:
324 + q = q[:attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()
325 + k = k[:attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()
326 + v = v[:attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()
327 + slot_id = state_indices_tensor[:attn_metadata.num_decodes]
328 + hidden = linear_decode_forward_triton(q, k, v, kv_cache, self.tp_slope,
329 + slot_id, 32)
330 + return hidden
331 +
332 + def forward(self, hidden_states: torch.Tensor, output: torch.Tensor,
333 + positions: torch.Tensor,
334 + kv_caches: MinimaxCacheParams) -> None:
335 + if not envs.VLLM_USE_V1:
336 + self._forward(hidden_states, output, positions, kv_caches)
337 + else:
338 + torch.ops.vllm.linear_attention(
339 + hidden_states,
340 + output,
341 + positions,
342 + self.prefix,
343 + )
344 +
345 + def _forward(self, hidden_states: torch.Tensor, output: torch.Tensor,
346 + positions: torch.Tensor,
347 + kv_caches: Optional[MinimaxCacheParams]) -> None:
348 + forward_context = get_forward_context()
349 + attn_metadata: AttentionMetadata = forward_context.attn_metadata
350 + if envs.VLLM_USE_V1 and attn_metadata is not None:
351 + assert isinstance(attn_metadata, dict)
352 + attn_metadata = attn_metadata[self.prefix]
353 + assert isinstance(attn_metadata, LinearAttentionMetadata)
354 + num_actual_tokens = attn_metadata.num_prefill_tokens + \
355 + attn_metadata.num_decode_tokens
356 + else:
357 + num_actual_tokens = hidden_states.shape[0]
358 +
359 + qkv, _ = self.qkv_proj(hidden_states[:num_actual_tokens])
360 + qkv32 = qkv.to(torch.float32)
361 + qkvact = torch.nn.functional.silu(qkv32)
362 + qkvact = qkvact.view((qkv.shape[0], self.tp_heads, -1))
363 + q, k, v = torch.split(qkvact, [self.head_dim] * 3, dim=-1)
364 + if envs.VLLM_USE_V1:
365 + if attn_metadata is not None:
366 + kv_cache = self.kv_cache[forward_context.virtual_engine][0]
367 + state_indices_tensor = attn_metadata.state_indices_tensor
368 +
369 + num_prefills = getattr(attn_metadata, "num_prefills", 0)
370 + if num_prefills > 0:
371 + num_decode_tokens = getattr(attn_metadata,
372 + "num_decode_tokens", 0)
373 + for prefill_idx in range(num_prefills):
374 + q_start = attn_metadata.query_start_loc[
375 + num_decode_tokens + prefill_idx]
376 + q_end = attn_metadata.query_start_loc[num_decode_tokens
377 + + prefill_idx +
378 + 1]
379 + query_len = q_end - q_start
380 + context_len = attn_metadata.seq_lens[
381 + num_decode_tokens + prefill_idx] - query_len
382 + if context_len == 0:
383 + block_to_clear = state_indices_tensor[
384 + num_decode_tokens + prefill_idx]
385 + kv_cache[block_to_clear, ...] = 0
386 + else:
387 + assert kv_caches is not None
388 + kv_cache = kv_caches.minimax_cache
389 + state_indices_tensor = kv_caches.state_indices_tensor
390 +
391 + decode_only = getattr(attn_metadata, "num_prefills", 0) == 0
392 + if attn_metadata is None:
393 + hidden = torch.empty((q.shape[0], q.shape[1] * q.shape[2]),
394 + device=q.device,
395 + dtype=q.dtype)
396 + else:
397 + if not decode_only:
398 + hidden = self._prefill_and_mix_infer(q, k, v, kv_cache,
399 + state_indices_tensor,
400 + attn_metadata)
401 + else:
402 + hidden = self._decode_infer(q, k, v, kv_cache,
403 + state_indices_tensor,
404 + attn_metadata)
405 + hidden = self.norm._forward(hidden)
406 + gate, _ = self.output_gate(hidden_states[:num_actual_tokens])
407 + hidden = F.sigmoid(gate) * hidden
408 + hidden = hidden.to(hidden_states.dtype)
409 +
410 + output[:num_actual_tokens], _ = self.out_proj(hidden)
411 +
412 +
413 +def linear_attention(
414 + hidden_states: torch.Tensor,
415 + output: torch.Tensor,
416 + positions: torch.Tensor,
417 + layer_name: str,
418 +) -> None:
419 + forward_context: ForwardContext = get_forward_context()
420 + self = forward_context.no_compile_layers[layer_name]
421 + self._forward(hidden_states=hidden_states,
422 + output=output,
423 + positions=positions,
424 + kv_caches=None)
425 +
426 +
427 +def linear_attention_fake(
428 + hidden_states: torch.Tensor,
429 + output: torch.Tensor,
430 + positions: torch.Tensor,
431 + layer_name: str,
432 +) -> None:
433 + return
434 +
435 +
436 +direct_register_custom_op(
437 + op_name="linear_attention",
438 + op_func=linear_attention,
439 + mutates_args=["output"],
440 + fake_impl=linear_attention_fake,
441 + dispatch_key=current_platform.dispatch_key,
442 +)
vllm/model_executor/models/minimax_text_01.py
+6 -410
@@ -1,45 +1,37 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Inference-only MiniMaxText01 model."""
4 -import math
4 from collections.abc import Iterable
5 from itertools import islice
6 from typing import TYPE_CHECKING, Optional, Union
7
8 if TYPE_CHECKING:
10 - from vllm.attention.backends.abstract import AttentionBackend
9 + pass
10
11 import regex as re
12 import torch
13 import torch.distributed
15 -import torch.nn.functional as F
16 -from einops import rearrange
14 from torch import nn
15 from transformers import MiniMaxConfig
16
17 from vllm import envs
18 from vllm.attention import Attention, AttentionMetadata
19 from vllm.compilation.decorators import support_torch_compile
23 -from vllm.config import (CacheConfig, ModelConfig, VllmConfig,
24 - get_current_vllm_config)
25 -from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
20 +from vllm.config import CacheConfig, ModelConfig, VllmConfig
21 from vllm.distributed.parallel_state import (
22 get_pp_group, get_tensor_model_parallel_rank,
23 get_tensor_model_parallel_world_size)
29 -from vllm.forward_context import ForwardContext, get_forward_context
30 -from vllm.model_executor.custom_op import CustomOp
24 +from vllm.forward_context import get_forward_context
25 from vllm.model_executor.layers.activation import SiluAndMul
26 from vllm.model_executor.layers.fused_moe import FusedMoE
27 from vllm.model_executor.layers.layernorm import RMSNorm
34 -from vllm.model_executor.layers.lightning_attn import (
35 - lightning_attention, linear_decode_forward_triton)
36 -from vllm.model_executor.layers.linear import (ColumnParallelLinear,
37 - MergedColumnParallelLinear,
28 +from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
29 QKVParallelLinear,
30 ReplicatedLinear,
31 RowParallelLinear)
32 from vllm.model_executor.layers.logits_processor import LogitsProcessor
42 -from vllm.model_executor.layers.mamba.abstract import MambaBase
33 +from vllm.model_executor.layers.mamba.linear_attn import (
34 + MiniMaxText01LinearAttention)
35 from vllm.model_executor.layers.mamba.mamba_utils import (
36 MambaStateDtypeCalculator, MambaStateShapeCalculator)
37 from vllm.model_executor.layers.quantization.base_config import (
@@ -50,10 +42,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
42 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
43 from vllm.model_executor.models.utils import maybe_prefix
44 from vllm.model_executor.sampling_metadata import SamplingMetadata
53 -from vllm.platforms import current_platform
45 from vllm.sequence import IntermediateTensors
55 -from vllm.utils import direct_register_custom_op
56 -from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
46
47 from .interfaces import HasInnerState, IsHybrid
48 from .minimax_cache import MinimaxCacheManager, MinimaxCacheParams
@@ -87,66 +76,6 @@ def weight_loader_with_alias(alias: str):
76 return wrapper
77
78
90 -class MiniMaxText01RMSNormTP(CustomOp):
91 - name = "MiniMaxText01RMSNormTP"
92 -
93 - def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
94 - super().__init__()
95 - self.tp_world = get_tensor_model_parallel_world_size()
96 - self.tp_rank = get_tensor_model_parallel_rank()
97 - self.weight = nn.Parameter(torch.ones(int(hidden_size /
98 - self.tp_world)))
99 -
100 - self.weight.weight_loader = self.weight_loader
101 - self.variance_epsilon = eps
102 - return
103 -
104 - @staticmethod
105 - def weight_loader(
106 - param: nn.Parameter,
107 - loaded_weight: torch.Tensor,
108 - ) -> None:
109 - tp_world = get_tensor_model_parallel_world_size()
110 - tp_rank = get_tensor_model_parallel_rank()
111 -
112 - shard_size = loaded_weight.shape[0] // tp_world
113 - shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
114 - param.data.copy_(loaded_weight[shard])
115 - return
116 -
117 - def _forward(
118 - self,
119 - x: torch.Tensor,
120 - ) -> torch.Tensor:
121 - orig_dtype = x.dtype
122 - x = x.to(torch.float32)
123 - variance = x.pow(2).mean(dim=-1, keepdim=True, dtype=torch.float32)
124 - if self.tp_world > 1:
125 - variance = tensor_model_parallel_all_reduce(
126 - variance) / self.tp_world
127 - x = x * torch.rsqrt(variance + self.variance_epsilon)
128 -
129 - weight = self.weight
130 - if x.size(-1) != self.weight.size(0):
131 - if self.weight.size(0) < x.size(-1):
132 - repeat_count = (x.size(-1) + self.weight.size(0)) // x.size(-1)
133 - full_weight = self.weight.repeat(repeat_count)
134 - weight = full_weight[:x.size(-1)]
135 - else:
136 - weight = self.weight[:x.size(-1)]
137 -
138 - x = x.to(orig_dtype) * weight
139 - return x
140 -
141 - def forward(
142 - self,
143 - x: torch.Tensor,
144 - residual: Optional[torch.Tensor] = None,
145 - ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
146 - assert residual is None, "RMSNorm does not support residual connection."
147 - return self._forward(x)
148 -
149 -
79 class MiniMaxText01MLP(nn.Module):
80
81 def __init__(
@@ -253,307 +182,6 @@ class MiniMaxText01MoE(nn.Module):
182 return final_hidden
183
184
256 -class MiniMaxText01LinearKernel:
257 -
258 - @staticmethod
259 - def jit_linear_forward_prefix(q: torch.Tensor,
260 - k: torch.Tensor,
261 - v: torch.Tensor,
262 - kv_caches: torch.Tensor,
263 - slope_rate: torch.Tensor,
264 - block_size: int,
265 - layer_idx: int = None,
266 - **kwargs) -> torch.Tensor:
267 -
268 - slope_rate = slope_rate.to(torch.float32)
269 - should_pad_dim = q.dim() == 3
270 - if should_pad_dim:
271 - q = q.unsqueeze(0)
272 - k = k.unsqueeze(0)
273 - v = v.unsqueeze(0)
274 - b, h, n, d = q.shape
275 - e = d
276 - kv_history = kv_caches.reshape(1, h, d, e).contiguous()
277 - output, kv_history = lightning_attention(q,
278 - k,
279 - v,
280 - slope_rate,
281 - block_size=block_size,
282 - kv_history=kv_history)
283 - kv_caches.copy_(kv_history[:, :, -1, :, :].reshape(h, d, e))
284 - assert output.shape[0] == 1, "batch size must be 1"
285 - return rearrange(output.squeeze(0), "h n d -> n (h d)")
286 -
287 -
288 -class MiniMaxText01LinearAttention(nn.Module, MambaBase):
289 -
290 - @property
291 - def mamba_type(self) -> str:
292 - return "linear_attention"
293 -
294 - def get_attn_backend(self) -> type["AttentionBackend"]:
295 - from vllm.v1.attention.backends.linear_attn import (
296 - LinearAttentionBackend)
297 - return LinearAttentionBackend
298 -
299 - def get_state_dtype(self) -> tuple[torch.dtype]:
300 - return MambaStateDtypeCalculator.linear_attention_state_dtype(
301 - self.model_config.dtype,
302 - self.cache_config.mamba_cache_dtype,
303 - )
304 -
305 - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
306 - return MambaStateShapeCalculator.linear_attention_state_shape(
307 - num_heads=self.num_heads,
308 - tp_size=self.tp_size,
309 - head_dim=self.head_dim)
310 -
311 - def __init__(
312 - self,
313 - hidden_size: int,
314 - hidden_inner_size: int,
315 - num_heads: int,
316 - head_dim: int,
317 - max_position: int,
318 - block_size: int,
319 - num_hidden_layer: int,
320 - model_config: Optional[ModelConfig] = None,
321 - cache_config: Optional[CacheConfig] = None,
322 - quant_config: Optional[QuantizationConfig] = None,
323 - layer_idx: int = 0,
324 - linear_layer_idx: int = 0,
325 - prefix: str = "linear_attn",
326 - ) -> None:
327 - super().__init__()
328 -
329 - self.layer_idx = layer_idx
330 - self.BLOCK = block_size
331 - self.hidden_size = hidden_size
332 - self.num_heads = num_heads
333 - self.head_dim = head_dim
334 - self.total_num_heads = num_heads
335 - self.hidden_inner_size = hidden_inner_size
336 - self.tp_size = get_tensor_model_parallel_world_size()
337 - self.tp_rank = get_tensor_model_parallel_rank()
338 -
339 - assert self.total_num_heads % self.tp_size == 0
340 - self.tp_heads = self.total_num_heads // self.tp_size
341 - self.qkv_size = self.num_heads * self.head_dim
342 - self.tp_hidden = self.head_dim * self.tp_heads
343 - self.model_config = model_config
344 - self.cache_config = cache_config
345 - self.prefix = prefix
346 -
347 - self.qkv_proj = ColumnParallelLinear(
348 - hidden_size,
349 - self.hidden_inner_size * 3,
350 - bias=False,
351 - quant_config=quant_config,
352 - prefix=f"{prefix}.qkv_proj",
353 - )
354 - self.output_gate = ColumnParallelLinear(
355 - hidden_size,
356 - self.hidden_inner_size,
357 - bias=False,
358 - quant_config=quant_config,
359 - prefix=f"{prefix}.output_gate",
360 - )
361 - self.out_proj = RowParallelLinear(
362 - self.hidden_inner_size,
363 - hidden_size,
364 - bias=False,
365 - quant_config=quant_config,
366 - prefix=f"{prefix}.out_proj",
367 - )
368 - self.norm = MiniMaxText01RMSNormTP(
369 - self.hidden_inner_size,
370 - eps=1e-5,
371 - )
372 -
373 - slope_rate = MiniMaxText01LinearAttention._build_slope_tensor(
374 - self.num_heads)
375 - if num_hidden_layer <= 1:
376 - self.slope_rate = slope_rate * (1 + 1e-5)
377 - else:
378 - self.slope_rate = slope_rate * (1 - layer_idx /
379 - (num_hidden_layer - 1) + 1e-5)
380 - self.tp_slope = self.slope_rate[self.tp_rank *
381 - self.tp_heads:(self.tp_rank + 1) *
382 - self.tp_heads].contiguous()
383 -
384 - if envs.VLLM_USE_V1:
385 - compilation_config = get_current_vllm_config().compilation_config
386 - if prefix in compilation_config.static_forward_context:
387 - raise ValueError(f"Duplicate layer name: {prefix}")
388 - compilation_config.static_forward_context[prefix] = self
389 -
390 - @staticmethod
391 - def weight_direct_load(param: torch.Tensor,
392 - loaded_weight: torch.Tensor) -> None:
393 - assert param.size() == loaded_weight.size()
394 - param.data.copy_(loaded_weight)
395 - return
396 -
397 - @staticmethod
398 - def _build_slope_tensor(n_attention_heads: int):
399 -
400 - def get_slopes(n):
401 -
402 - def get_slopes_power_of_2(n):
403 - start = 2**(-(2**-(math.log2(n) - 3)))
404 - ratio = start
405 - return [start * ratio**i for i in range(n)]
406 -
407 - if math.log2(n).is_integer():
408 - return get_slopes_power_of_2(n)
409 - else:
410 - closest_power_of_2 = 2**math.floor(math.log2(n))
411 - return (get_slopes_power_of_2(closest_power_of_2) + get_slopes(
412 - 2 * closest_power_of_2)[0::2][:n - closest_power_of_2])
413 -
414 - slopes = torch.tensor(get_slopes(n_attention_heads),
415 - dtype=torch.float32).reshape(
416 - n_attention_heads, 1, 1)
417 - return slopes
418 -
419 - def _prefill_and_mix_infer(self, q, k, v, kv_cache, state_indices_tensor,
420 - attn_metadata):
421 - hidden = []
422 - for _prefill_idx in range(getattr(attn_metadata, "num_prefills", 0)):
423 - if _prefill_idx >= len(attn_metadata.query_start_loc):
424 - break
425 - if _prefill_idx >= len(state_indices_tensor):
426 - break
427 - # prefills are packed at end of batch in V1
428 - offset = attn_metadata.num_decode_tokens if envs.VLLM_USE_V1 else 0
429 - _start = attn_metadata.query_start_loc[offset + _prefill_idx]
430 - _end = attn_metadata.query_start_loc[offset + _prefill_idx + 1]
431 - slot_id = state_indices_tensor[offset + _prefill_idx]
432 - qs = q[_start:_end].transpose(0, 1).contiguous()
433 - ks = k[_start:_end].transpose(0, 1).contiguous()
434 - vs = v[_start:_end].transpose(0, 1).contiguous()
435 - slice_layer_cache = kv_cache[slot_id, ...]
436 -
437 - out_slice = MiniMaxText01LinearKernel.jit_linear_forward_prefix(
438 - qs,
439 - ks,
440 - vs,
441 - slice_layer_cache,
442 - self.tp_slope,
443 - self.BLOCK,
444 - layer_idx=self.layer_idx)
445 - hidden.append(out_slice.contiguous())
446 - if attn_metadata.num_decode_tokens > 0:
447 - hidden_decode = self._decode_infer(q, k, v, kv_cache,
448 - state_indices_tensor,
449 - attn_metadata)
450 - if envs.VLLM_USE_V1:
451 - hidden.insert(0, hidden_decode)
452 - else:
453 - hidden.append(hidden_decode)
454 -
455 - if not hidden:
456 - return torch.empty((0, q.size(-1)), device=q.device, dtype=q.dtype)
457 -
458 - hidden = torch.concat(hidden, dim=0).contiguous()
459 - return hidden
460 -
461 - def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor,
462 - attn_metadata):
463 - if not envs.VLLM_USE_V1:
464 - q = q[attn_metadata.num_prefill_tokens:].unsqueeze(2).contiguous()
465 - k = k[attn_metadata.num_prefill_tokens:].unsqueeze(2).contiguous()
466 - v = v[attn_metadata.num_prefill_tokens:].unsqueeze(2).contiguous()
467 - num_prefills = getattr(attn_metadata, "num_prefills", 0)
468 - slot_id = state_indices_tensor[num_prefills:]
469 - else:
470 - q = q[:attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()
471 - k = k[:attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()
472 - v = v[:attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()
473 - slot_id = state_indices_tensor[:attn_metadata.num_decodes]
474 - hidden = linear_decode_forward_triton(q, k, v, kv_cache, self.tp_slope,
475 - slot_id, 32)
476 - return hidden
477 -
478 - def forward(self, hidden_states: torch.Tensor, output: torch.Tensor,
479 - positions: torch.Tensor,
480 - kv_caches: MinimaxCacheParams) -> None:
481 - if not envs.VLLM_USE_V1:
482 - self._forward(hidden_states, output, positions, kv_caches)
483 - else:
484 - torch.ops.vllm.linear_attention(
485 - hidden_states,
486 - output,
487 - positions,
488 - self.prefix,
489 - )
490 -
491 - def _forward(self, hidden_states: torch.Tensor, output: torch.Tensor,
492 - positions: torch.Tensor,
493 - kv_caches: Optional[MinimaxCacheParams]) -> None:
494 - forward_context = get_forward_context()
495 - attn_metadata: AttentionMetadata = forward_context.attn_metadata
496 - if envs.VLLM_USE_V1 and attn_metadata is not None:
497 - assert isinstance(attn_metadata, dict)
498 - attn_metadata = attn_metadata[self.prefix]
499 - assert isinstance(attn_metadata, LinearAttentionMetadata)
500 - num_actual_tokens = attn_metadata.num_prefill_tokens + \
501 - attn_metadata.num_decode_tokens
502 - else:
503 - num_actual_tokens = hidden_states.shape[0]
504 -
505 - qkv, _ = self.qkv_proj(hidden_states[:num_actual_tokens])
506 - qkv32 = qkv.to(torch.float32)
507 - qkvact = torch.nn.functional.silu(qkv32)
508 - qkvact = qkvact.view((qkv.shape[0], self.tp_heads, -1))
509 - q, k, v = torch.split(qkvact, [self.head_dim] * 3, dim=-1)
510 - if envs.VLLM_USE_V1:
511 - if attn_metadata is not None:
512 - kv_cache = self.kv_cache[forward_context.virtual_engine][0]
513 - state_indices_tensor = attn_metadata.state_indices_tensor
514 -
515 - num_prefills = getattr(attn_metadata, "num_prefills", 0)
516 - if num_prefills > 0:
517 - num_decode_tokens = getattr(attn_metadata,
518 - "num_decode_tokens", 0)
519 - for prefill_idx in range(num_prefills):
520 - q_start = attn_metadata.query_start_loc[
521 - num_decode_tokens + prefill_idx]
522 - q_end = attn_metadata.query_start_loc[num_decode_tokens
523 - + prefill_idx +
524 - 1]
525 - query_len = q_end - q_start
526 - context_len = attn_metadata.seq_lens[
527 - num_decode_tokens + prefill_idx] - query_len
528 - if context_len == 0:
529 - block_to_clear = state_indices_tensor[
530 - num_decode_tokens + prefill_idx]
531 - kv_cache[block_to_clear, ...] = 0
532 - else:
533 - kv_cache = kv_caches.minimax_cache
534 - state_indices_tensor = kv_caches.state_indices_tensor
535 -
536 - decode_only = getattr(attn_metadata, "num_prefills", 0) == 0
537 - if attn_metadata is None:
538 - hidden = torch.empty((q.shape[0], q.shape[1] * q.shape[2]),
539 - device=q.device,
540 - dtype=q.dtype)
541 - else:
542 - if not decode_only:
543 - hidden = self._prefill_and_mix_infer(q, k, v, kv_cache,
544 - state_indices_tensor,
545 - attn_metadata)
546 - else:
547 - hidden = self._decode_infer(q, k, v, kv_cache,
548 - state_indices_tensor,
549 - attn_metadata)
550 - hidden = self.norm._forward(hidden)
551 - gate, _ = self.output_gate(hidden_states[:num_actual_tokens])
552 - hidden = F.sigmoid(gate) * hidden
553 - hidden = hidden.to(hidden_states.dtype)
554 - output[:num_actual_tokens], _ = self.out_proj(hidden)
555 -
556 -
185 class MiniMaxText01Attention(nn.Module):
186
187 def __init__(
@@ -1397,35 +1025,3 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):
1025 tp_size=parallel_config.tensor_parallel_size,
1026 head_dim=hf_config.head_dim,
1027 )
1400 -
1401 -
1402 -def linear_attention(
1403 - hidden_states: torch.Tensor,
1404 - output: torch.Tensor,
1405 - positions: torch.Tensor,
1406 - layer_name: str,
1407 -) -> None:
1408 - forward_context: ForwardContext = get_forward_context()
1409 - self = forward_context.no_compile_layers[layer_name]
1410 - self._forward(hidden_states=hidden_states,
1411 - output=output,
1412 - positions=positions,
1413 - kv_caches=None)
1414 -
1415 -
1416 -def linear_attention_fake(
1417 - hidden_states: torch.Tensor,
1418 - output: torch.Tensor,
1419 - positions: torch.Tensor,
1420 - layer_name: str,
1421 -) -> None:
1422 - return
1423 -
1424 -
1425 -direct_register_custom_op(
1426 - op_name="linear_attention",
1427 - op_func=linear_attention,
1428 - mutates_args=["output"],
1429 - fake_impl=linear_attention_fake,
1430 - dispatch_key=current_platform.dispatch_key,
1431 -)