@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2e2000f352
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

[Model] Add LFM2 architecture (#22845)

Signed-off-by: Paul Pak <paulpak58@gmail.com>

Paul Pak committed Aug 21, 2025 at 01:35 UTC 2e2000f352d861dd2b527a48c5f12d295a93c3dd
11 files changed +960 -8
docs/models/supported_models.md
+1
@@ -373,6 +373,7 @@ th {
373 | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |
374 | `JAISLMHeadModel` | Jais | `inceptionai/jais-13b`, `inceptionai/jais-13b-chat`, `inceptionai/jais-30b-v3`, `inceptionai/jais-30b-chat-v3`, etc. | | ✅︎ | ✅︎ |
375 | `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ | ✅︎ |
376 +| `Lfm2ForCausalLM` | LFM2 | `LiquidAI/LFM2-1.2B`, `LiquidAI/LFM2-700M`, `LiquidAI/LFM2-350M`, etc. | ✅︎ | ✅︎ | ✅︎ |
377 | `LlamaForCausalLM` | Llama 3.1, Llama 3, Llama 2, LLaMA, Yi | `meta-llama/Meta-Llama-3.1-405B-Instruct`, `meta-llama/Meta-Llama-3.1-70B`, `meta-llama/Meta-Llama-3-70B-Instruct`, `meta-llama/Llama-2-70b-hf`, `01-ai/Yi-34B`, etc. | ✅︎ | ✅︎ | ✅︎ |
378 | `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ | ✅︎ |
379 | `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ | ✅︎ |
tests/models/language/generation/test_hybrid.py
+25 -8
@@ -31,6 +31,7 @@ HYBRID_MODELS = [
31 "hmellor/tiny-random-BambaForCausalLM",
32 "ibm-granite/granite-4.0-tiny-preview",
33 "tiiuae/Falcon-H1-0.5B-Base",
34 + "LiquidAI/LFM2-1.2B",
35 ]
36
37 HF_UNSUPPORTED_MODELS = [
@@ -52,6 +53,7 @@ V1_SUPPORTED_MODELS = [
53 "hmellor/tiny-random-BambaForCausalLM",
54 "ibm-granite/granite-4.0-tiny-preview",
55 "tiiuae/Falcon-H1-0.5B-Base",
56 + "LiquidAI/LFM2-1.2B",
57 ]
58
59 FULL_CUDA_GRAPH_MODELS = [
@@ -59,6 +61,10 @@ FULL_CUDA_GRAPH_MODELS = [
61 "Zyphra/Zamba2-1.2B-instruct",
62 ]
63
64 +V0_UNSUPPORTED_MODELS = [
65 + "LiquidAI/LFM2-1.2B",
66 +]
67 +
68 # Avoid OOM
69 MAX_NUM_SEQS = 4
70
@@ -94,9 +100,12 @@ def test_models(
100 else:
101 hf_outputs = None
102
97 - with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
98 - vllm_v0_outputs = vllm_model.generate_greedy_logprobs(
99 - example_prompts, max_tokens, num_logprobs)
103 + if model not in V0_UNSUPPORTED_MODELS:
104 + with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
105 + vllm_v0_outputs = vllm_model.generate_greedy_logprobs(
106 + example_prompts, max_tokens, num_logprobs)
107 + else:
108 + vllm_v0_outputs = None
109
110 if model in V1_SUPPORTED_MODELS:
111 with monkeypatch.context() as m:
@@ -112,7 +121,7 @@ def test_models(
121 else:
122 vllm_v1_outputs = None
123
115 - if hf_outputs is not None:
124 + if hf_outputs is not None and vllm_v0_outputs is not None:
125 check_logprobs_close(
126 outputs_0_lst=hf_outputs,
127 outputs_1_lst=vllm_v0_outputs,
@@ -122,6 +131,7 @@ def test_models(
131
132 if model in V1_SUPPORTED_MODELS:
133 ref_outputs = hf_outputs if hf_outputs is not None else vllm_v0_outputs
134 + assert ref_outputs is not None
135 check_logprobs_close(
136 outputs_0_lst=ref_outputs,
137 outputs_1_lst=vllm_v1_outputs,
@@ -140,6 +150,9 @@ def test_batching(
150 max_tokens: int,
151 num_logprobs: int,
152 ) -> None:
153 + if model in V0_UNSUPPORTED_MODELS:
154 + pytest.skip(
155 + f"Unsupported V0 Engine. Skipping `test_batching` on {model}.")
156
157 try:
158 model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
@@ -392,9 +405,12 @@ def test_full_cuda_graph(
405 else:
406 hf_outputs = None
407
395 - with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
396 - vllm_v0_outputs = vllm_model.generate_greedy_logprobs(
397 - example_prompts, max_tokens, num_logprobs)
408 + if model not in V0_UNSUPPORTED_MODELS:
409 + with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
410 + vllm_v0_outputs = vllm_model.generate_greedy_logprobs(
411 + example_prompts, max_tokens, num_logprobs)
412 + else:
413 + vllm_v0_outputs = None
414
415 with monkeypatch.context() as m:
416 m.setenv("VLLM_USE_V1", "1")
@@ -408,7 +424,7 @@ def test_full_cuda_graph(
424 vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
425 example_prompts, max_tokens, num_logprobs)
426
411 - if hf_outputs is not None:
427 + if hf_outputs is not None and vllm_v0_outputs is not None:
428 check_logprobs_close(
429 outputs_0_lst=hf_outputs,
430 outputs_1_lst=vllm_v0_outputs,
@@ -417,6 +433,7 @@ def test_full_cuda_graph(
433 )
434
435 ref_outputs = hf_outputs if hf_outputs is not None else vllm_v0_outputs
436 + assert ref_outputs is not None
437 check_logprobs_close(
438 outputs_0_lst=ref_outputs,
439 outputs_1_lst=vllm_v1_outputs,
tests/models/registry.py
+2
@@ -230,6 +230,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
230 "tiny": "ai21labs/Jamba-tiny-dev",
231 "random": "ai21labs/Jamba-tiny-random", # noqa: E501
232 }),
233 + "Lfm2ForCausalLM": _HfExamplesInfo("LiquidAI/LFM2-1.2B",
234 + min_transformers_version="4.54"),
235 "LlamaForCausalLM": _HfExamplesInfo("meta-llama/Llama-3.2-1B-Instruct",
236 extras={"guard": "meta-llama/Llama-Guard-3-1B", # noqa: E501
237 "hermes": "NousResearch/Hermes-3-Llama-3.1-8B", # noqa: E501
tests/models/test_initialization.py
+2
@@ -95,6 +95,8 @@ def can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch,
95
96 @pytest.mark.parametrize("model_arch", HF_EXAMPLE_MODELS.get_supported_archs())
97 def test_can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch):
98 + if model_arch == "Lfm2ForCausalLM":
99 + pytest.skip("Skipping until test supports V1-only models")
100 can_initialize(model_arch, monkeypatch, HF_EXAMPLE_MODELS)
101
102
vllm/config/compilation.py
+1
@@ -337,6 +337,7 @@ class CompilationConfig:
337 "vllm.unified_attention_with_output",
338 "vllm.mamba_mixer2",
339 "vllm.mamba_mixer",
340 + "vllm.short_conv",
341 ]
342
343 def compute_hash(self) -> str:
vllm/model_executor/layers/mamba/mamba_utils.py
+24
@@ -54,6 +54,16 @@ class MambaStateDtypeCalculator:
54
55 return (conv_state_dtype, temporal_state_dtype)
56
57 + @classmethod
58 + def short_conv_state_dtype(
59 + cls,
60 + model_dtype: Union[ModelDType, torch.dtype],
61 + mamba_cache_dtype: MambaDType,
62 + ) -> tuple[torch.dtype, ...]:
63 + conv_state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype,
64 + model_dtype)
65 + return (conv_state_dtype, )
66 +
67
68 class MambaStateShapeCalculator:
69
@@ -122,6 +132,20 @@ class MambaStateShapeCalculator:
132 tp_world_size), head_dim, state_size)
133 return conv_state_shape, temporal_state_shape
134
135 + @classmethod
136 + def short_conv_state_shape(
137 + cls,
138 + tp_world_size: int,
139 + intermediate_size: int,
140 + conv_kernel: int,
141 + use_v1: bool = True,
142 + ) -> tuple[tuple[int, int]]:
143 + conv_dim = divide(intermediate_size, tp_world_size)
144 + conv_state_shape = (conv_kernel - 1, conv_dim)
145 + if not use_v1:
146 + conv_state_shape = conv_state_shape[1], conv_state_shape[0]
147 + return (conv_state_shape, )
148 +
149 @classmethod
150 def extra_groups_for_head_shards(cls, ngroups: int, tp_size: int):
151 """Compute the increase in group numbers to account for
vllm/model_executor/layers/mamba/short_conv.py new
+262
@@ -0,0 +1,262 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Optional
5 +
6 +import torch
7 +
8 +from vllm import envs
9 +from vllm.attention.backends.abstract import AttentionMetadata
10 +from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
11 +from vllm.distributed import get_tensor_model_parallel_world_size
12 +from vllm.forward_context import ForwardContext, get_forward_context
13 +from vllm.model_executor.custom_op import CustomOp
14 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,
15 + MergedColumnParallelLinear,
16 + RowParallelLinear)
17 +from vllm.model_executor.layers.mamba.abstract import MambaBase
18 +from vllm.model_executor.layers.mamba.mamba2_metadata import update_metadata
19 +from vllm.model_executor.layers.mamba.mamba_utils import (
20 + MambaStateDtypeCalculator, MambaStateShapeCalculator)
21 +from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
22 + causal_conv1d_fn, causal_conv1d_update)
23 +from vllm.platforms import current_platform
24 +from vllm.utils import direct_register_custom_op
25 +from vllm.v1.attention.backends.short_conv_attn import (
26 + ShortConvAttentionMetadata)
27 +
28 +
29 +@CustomOp.register("short_conv")
30 +class ShortConv(MambaBase, CustomOp):
31 +
32 + def __init__(self,
33 + config,
34 + dim: int,
35 + layer_idx: int,
36 + model_config: Optional[ModelConfig] = None,
37 + cache_config: Optional[CacheConfig] = None,
38 + prefix: str = ""):
39 + super().__init__()
40 + self.config = config
41 + self.layer_idx = layer_idx
42 + self.conv_dim = dim
43 + self.L_cache = config.conv_L_cache
44 + self.bias = config.conv_bias
45 +
46 + self.conv = ColumnParallelLinear(
47 + input_size=self.L_cache,
48 + output_size=dim,
49 + bias=self.bias,
50 + prefix=f"{prefix}.conv1d",
51 + )
52 + # unsqueeze to fit conv1d weights shape into the linear weights shape.
53 + # Can't do this in `weight_loader` since it already exists in
54 + # `ColumnParallelLinear` and `set_weight_attrs`
55 + # doesn't allow to override it
56 + self.conv.weight.data = self.conv.weight.data.unsqueeze(1)
57 +
58 + self.in_proj = MergedColumnParallelLinear(
59 + input_size=dim,
60 + output_sizes=[dim] * 3,
61 + bias=self.bias,
62 + prefix=f"{prefix}.in_proj",
63 + )
64 + self.out_proj = RowParallelLinear(
65 + input_size=dim,
66 + output_size=dim,
67 + bias=self.bias,
68 + prefix=f"{prefix}.out_proj",
69 + )
70 +
71 + assert envs.VLLM_USE_V1, ("ShortConv layers are only supported in V1")
72 + compilation_config = get_current_vllm_config().compilation_config
73 + if prefix in compilation_config.static_forward_context:
74 + raise ValueError(f"Duplicate layer name: {prefix}")
75 + compilation_config.static_forward_context[prefix] = self
76 + # The outer list is for v0 PP virtual engine. Though this code path
77 + # only runs for v1, we have to do this to unify with the interface
78 + # of Attention + v0 PP.
79 + self.kv_cache = [(torch.tensor([]), )]
80 +
81 + self.model_config = model_config
82 + self.cache_config = cache_config
83 + self.prefix = prefix
84 +
85 + def forward_native(
86 + self,
87 + hidden_states: torch.Tensor,
88 + output: torch.Tensor,
89 + conv_metadata: ShortConvAttentionMetadata,
90 + ):
91 + return
92 +
93 + def forward(
94 + self,
95 + hidden_states: torch.Tensor,
96 + output: torch.Tensor,
97 + conv_metadata: ShortConvAttentionMetadata,
98 + ):
99 + torch.ops.vllm.short_conv(
100 + hidden_states,
101 + output,
102 + self.prefix,
103 + )
104 +
105 + def forward_cuda(
106 + self,
107 + hidden_states: torch.Tensor,
108 + output: torch.Tensor,
109 + conv_metadata: ShortConvAttentionMetadata,
110 + ):
111 + forward_context = get_forward_context()
112 + # ShortConvAttentionMetadata contains metadata necessary for the
113 + # short_conv triton kernels to operate in continuous batching and in
114 + # chunked prefill modes; they are computed at top-level model forward
115 + # since they stay the same and reused for all mamba layers in the same
116 + # iteration.
117 + attn_metadata: AttentionMetadata = forward_context.attn_metadata
118 + if attn_metadata is not None:
119 + assert isinstance(attn_metadata, dict)
120 + attn_metadata = attn_metadata[self.prefix]
121 + conv_metadata = attn_metadata
122 + assert isinstance(attn_metadata, ShortConvAttentionMetadata)
123 + self_kv_cache = self.kv_cache[forward_context.virtual_engine]
124 + conv_state = self_kv_cache[0].transpose(-1, -2)
125 + state_indices_tensor = attn_metadata.state_indices_tensor
126 + has_initial_states_p = attn_metadata.has_initial_states
127 +
128 + BCx, _ = self.in_proj(hidden_states)
129 +
130 + B, C, x = BCx.chunk(3, dim=-1)
131 +
132 + conv_weights = self.conv.weight.view(self.conv.weight.size(0),
133 + self.conv.weight.size(2))
134 +
135 + if attn_metadata is None:
136 + # V1 profile run
137 + Bx = (B * x).contiguous()
138 + hidden_states = C * Bx
139 + contextualized_states, _ = self.out_proj(hidden_states)
140 + return contextualized_states
141 +
142 + num_prefills = attn_metadata.num_prefills # request count
143 + num_decodes = attn_metadata.num_decode_tokens # token count (=request)
144 + num_prefill_tokens = attn_metadata.num_prefill_tokens # token count
145 + has_prefill = num_prefills > 0
146 + has_decode = num_decodes > 0
147 + num_actual_tokens = num_decodes + num_prefill_tokens
148 +
149 + # NOTE: V1 puts decode before prefill
150 + # Separate prefill and decode by splitting varlen input
151 + # Split along token dimension
152 + B_d, B_p = torch.split(
153 + B[:num_actual_tokens],
154 + [num_decodes, num_prefill_tokens],
155 + dim=0,
156 + )
157 + C_d, C_p = torch.split(
158 + C[:num_actual_tokens],
159 + [num_decodes, num_prefill_tokens],
160 + dim=0,
161 + )
162 + x_d, x_p = torch.split(
163 + x[:num_actual_tokens],
164 + [num_decodes, num_prefill_tokens],
165 + dim=0,
166 + )
167 + # Split along batch dimension
168 + state_indices_tensor_d, state_indices_tensor_p = torch.split(
169 + state_indices_tensor,
170 + [num_decodes, num_prefills],
171 + dim=0,
172 + )
173 + query_start_loc_p = (
174 + attn_metadata.query_start_loc[-num_prefills - 1:] -
175 + num_decodes if has_prefill else None)
176 +
177 + conv_output_list = []
178 +
179 + if has_prefill:
180 + Bx_p = (B_p * x_p).transpose(0, 1)
181 + if conv_metadata.cu_seqlen is None:
182 + conv_metadata = update_metadata(Bx_p, query_start_loc_p,
183 + conv_metadata)
184 + Bx = causal_conv1d_fn(Bx_p,
185 + conv_weights,
186 + self.conv.bias,
187 + activation=None,
188 + conv_states=conv_state,
189 + has_initial_state=has_initial_states_p,
190 + cache_indices=state_indices_tensor_p,
191 + metadata=conv_metadata,
192 + query_start_loc=query_start_loc_p).transpose(
193 + 0, 1)[:num_prefill_tokens]
194 +
195 + y = C_p * Bx
196 + conv_output_list.append(y)
197 +
198 + if has_decode:
199 + Bx_d = (B_d * x_d).contiguous()
200 + Bx = causal_conv1d_update(
201 + Bx_d,
202 + conv_state,
203 + conv_weights,
204 + self.conv.bias,
205 + activation=None,
206 + conv_state_indices=state_indices_tensor_d)
207 + y = C_d * Bx
208 + conv_output_list.insert(0, y)
209 +
210 + # Merge prefill and decode outputs before passing to gated MLP
211 + hidden_states = torch.vstack(conv_output_list)
212 +
213 + # Final linear projection
214 + output[:num_actual_tokens], _ = self.out_proj(hidden_states)
215 +
216 + def get_state_dtype(self) -> tuple[torch.dtype, ...]:
217 + assert self.model_config is not None
218 + assert self.cache_config is not None
219 + return MambaStateDtypeCalculator.short_conv_state_dtype(
220 + self.model_config.dtype,
221 + self.cache_config.mamba_cache_dtype,
222 + )
223 +
224 + def get_state_shape(self) -> tuple[tuple[int, ...]]:
225 + return MambaStateShapeCalculator.short_conv_state_shape(
226 + tp_world_size=get_tensor_model_parallel_world_size(),
227 + intermediate_size=self.conv_dim,
228 + conv_kernel=self.L_cache,
229 + )
230 +
231 + @property
232 + def mamba_type(self) -> str:
233 + return "short_conv"
234 +
235 +
236 +def short_conv(
237 + hidden_states: torch.Tensor,
238 + output: torch.Tensor,
239 + layer_name: str,
240 +) -> None:
241 + forward_context: ForwardContext = get_forward_context()
242 + self = forward_context.no_compile_layers[layer_name]
243 + self.forward_cuda(hidden_states=hidden_states,
244 + output=output,
245 + conv_metadata=None)
246 +
247 +
248 +def short_conv_fake(
249 + hidden_states: torch.Tensor,
250 + output: torch.Tensor,
251 + layer_name: str,
252 +) -> None:
253 + return
254 +
255 +
256 +direct_register_custom_op(
257 + op_name="short_conv",
258 + op_func=short_conv,
259 + mutates_args=["output"],
260 + fake_impl=short_conv_fake,
261 + dispatch_key=current_platform.dispatch_key,
262 +)
vllm/model_executor/models/lfm2.py new
+557
@@ -0,0 +1,557 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Iterable
4 +from typing import Any, Optional
5 +
6 +import torch
7 +import torch.nn as nn
8 +from transformers import Lfm2Config
9 +
10 +from vllm import envs
11 +from vllm.attention import Attention
12 +from vllm.compilation.decorators import support_torch_compile
13 +from vllm.config import CacheConfig, ModelConfig, VllmConfig
14 +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
15 +from vllm.model_executor.layers.activation import SiluAndMul
16 +from vllm.model_executor.layers.layernorm import RMSNorm
17 +from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
18 + QKVParallelLinear,
19 + RowParallelLinear)
20 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
21 +from vllm.model_executor.layers.mamba.mamba_utils import (
22 + MambaStateDtypeCalculator, MambaStateShapeCalculator)
23 +from vllm.model_executor.layers.mamba.short_conv import ShortConv
24 +from vllm.model_executor.layers.quantization import QuantizationConfig
25 +from vllm.model_executor.layers.rotary_embedding import get_rope
26 +from vllm.model_executor.layers.vocab_parallel_embedding import (
27 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
28 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
29 +from vllm.model_executor.sampling_metadata import SamplingMetadata
30 +from vllm.sequence import IntermediateTensors
31 +
32 +from .interfaces import (HasInnerState, IsHybrid, SupportsLoRA, SupportsPP,
33 + SupportsQuant)
34 +from .utils import (AutoWeightsLoader, PPMissingLayer, extract_layer_index,
35 + is_pp_missing_parameter,
36 + make_empty_intermediate_tensors_factory, make_layers,
37 + maybe_prefix)
38 +
39 +
40 +class Lfm2MLP(nn.Module):
41 +
42 + def __init__(
43 + self,
44 + dim: int,
45 + ff_dim: int,
46 + multiple_of: int,
47 + auto_adjust_ff_dim: bool,
48 + ffn_dim_multiplier: Optional[float],
49 + quant_config: Optional[QuantizationConfig] = None,
50 + prefix: str = "",
51 + ):
52 + super().__init__()
53 + if auto_adjust_ff_dim:
54 + ff_dim = int(2 * ff_dim / 3)
55 + # custom dim factor multiplier
56 + if ffn_dim_multiplier is not None:
57 + ff_dim = int(ffn_dim_multiplier * ff_dim)
58 + ff_dim = multiple_of * ((ff_dim + multiple_of - 1) // multiple_of)
59 +
60 + self.w1 = MergedColumnParallelLinear(
61 + input_size=dim,
62 + output_sizes=[ff_dim] * 2,
63 + bias=False,
64 + quant_config=quant_config,
65 + prefix=f"{prefix}.gate_up_proj",
66 + )
67 + self.w2 = RowParallelLinear(
68 + input_size=ff_dim,
69 + output_size=dim,
70 + bias=False,
71 + quant_config=quant_config,
72 + prefix=f"{prefix}.down_proj",
73 + )
74 + self.act_fn = SiluAndMul()
75 +
76 + def forward(self, x: torch.Tensor) -> torch.Tensor:
77 + gate_up, _ = self.w1(x)
78 + x = self.act_fn(gate_up)
79 + x, _ = self.w2(x)
80 + return x
81 +
82 +
83 +class Lfm2Attention(nn.Module):
84 +
85 + def __init__(
86 + self,
87 + config: Lfm2Config,
88 + layer_idx: int,
89 + hidden_size: int,
90 + num_heads: int,
91 + num_kv_heads: int,
92 + rope_theta: float = 10000,
93 + rope_scaling: Optional[dict[str, Any]] = None,
94 + max_position_embeddings: int = 8192,
95 + cache_config: Optional[CacheConfig] = None,
96 + quant_config: Optional[QuantizationConfig] = None,
97 + prefix: str = "",
98 + ) -> None:
99 + super().__init__()
100 + self.layer_idx = layer_idx
101 + self.hidden_size = hidden_size
102 + self.num_kv_heads = num_kv_heads
103 + tp_size = get_tensor_model_parallel_world_size()
104 + self.total_num_heads = num_heads
105 + assert self.total_num_heads % tp_size == 0
106 + self.num_heads = self.total_num_heads // tp_size
107 + self.total_num_kv_heads = num_kv_heads
108 + if self.total_num_kv_heads >= tp_size:
109 + # Number of KV heads is greater than TP size, so we partition
110 + # the KV heads across multiple tensor parallel GPUs.
111 + assert self.total_num_kv_heads % tp_size == 0
112 + else:
113 + # Number of KV heads is less than TP size, so we replicate
114 + # the KV heads across multiple tensor parallel GPUs.
115 + assert tp_size % self.total_num_kv_heads == 0
116 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
117 + self.head_dim = self.hidden_size // self.total_num_heads
118 +
119 + self.q_size = self.num_heads * self.head_dim
120 + self.kv_size = self.num_kv_heads * self.head_dim
121 + self.scaling = self.head_dim**-0.5
122 + self.rope_theta = rope_theta
123 + self.max_position_embeddings = max_position_embeddings
124 +
125 + self.qkv_proj = QKVParallelLinear(
126 + hidden_size=self.hidden_size,
127 + head_size=self.head_dim,
128 + total_num_heads=self.total_num_heads,
129 + total_num_kv_heads=self.total_num_kv_heads,
130 + bias=False,
131 + quant_config=quant_config,
132 + prefix=f"{prefix}.qkv_proj",
133 + )
134 + self.out_proj = RowParallelLinear(
135 + input_size=self.total_num_heads * self.head_dim,
136 + output_size=self.hidden_size,
137 + bias=False,
138 + quant_config=quant_config,
139 + prefix=f"{prefix}.out_proj",
140 + )
141 + self.rotary_emb = get_rope(
142 + self.head_dim,
143 + rotary_dim=self.head_dim,
144 + max_position=self.max_position_embeddings,
145 + base=self.rope_theta,
146 + rope_scaling=rope_scaling,
147 + is_neox_style=True,
148 + )
149 + self.attn = Attention(
150 + self.num_heads,
151 + self.head_dim,
152 + self.scaling,
153 + num_kv_heads=self.num_kv_heads,
154 + cache_config=cache_config,
155 + prefix=f"{prefix}.attn",
156 + )
157 + self.q_layernorm = RMSNorm(self.head_dim, eps=config.norm_eps)
158 + self.k_layernorm = RMSNorm(self.head_dim, eps=config.norm_eps)
159 +
160 + def forward(
161 + self,
162 + positions: torch.Tensor,
163 + hidden_states: torch.Tensor,
164 + ) -> torch.Tensor:
165 + n_tokens, _ = hidden_states.shape
166 + qkv, _ = self.qkv_proj(hidden_states)
167 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
168 + q = q.view(n_tokens, self.num_heads, self.head_dim).contiguous()
169 + k = k.view(n_tokens, self.num_kv_heads, self.head_dim).contiguous()
170 + q = self.q_layernorm(q)
171 + k = self.k_layernorm(k)
172 + q, k = self.rotary_emb(positions, q, k)
173 + q = q.view(n_tokens, self.num_heads * self.head_dim)
174 + k = k.view(n_tokens, self.num_kv_heads * self.head_dim)
175 + attn_output = self.attn(q, k, v)
176 + output, _ = self.out_proj(attn_output)
177 + return output
178 +
179 +
180 +class Lfm2AttentionDecoderLayer(nn.Module):
181 +
182 + def __init__(
183 + self,
184 + config: Lfm2Config,
185 + layer_idx: int,
186 + model_config: Optional[ModelConfig] = None,
187 + cache_config: Optional[CacheConfig] = None,
188 + quant_config: Optional[QuantizationConfig] = None,
189 + prefix: str = "",
190 + ) -> None:
191 + super().__init__()
192 + self.prefix = prefix
193 + self.config = config
194 + self.layer_idx = layer_idx
195 +
196 + rope_theta = getattr(config, "rope_theta", 10000)
197 + rope_scaling = getattr(config, "rope_scaling", None)
198 + if rope_scaling is not None and getattr(
199 + config, "original_max_position_embeddings", None):
200 + rope_scaling["original_max_position_embeddings"] = (
201 + config.original_max_position_embeddings)
202 + max_position_embeddings = getattr(config, "max_position_embeddings",
203 + 8192)
204 +
205 + self.self_attn = Lfm2Attention(
206 + config=config,
207 + layer_idx=layer_idx,
208 + hidden_size=config.hidden_size,
209 + num_heads=config.num_attention_heads,
210 + num_kv_heads=config.num_key_value_heads,
211 + rope_theta=rope_theta,
212 + rope_scaling=rope_scaling,
213 + max_position_embeddings=max_position_embeddings,
214 + cache_config=cache_config,
215 + quant_config=quant_config,
216 + prefix=f"{prefix}.self_attn",
217 + )
218 +
219 + self.feed_forward = Lfm2MLP(
220 + dim=config.block_dim,
221 + ff_dim=config.block_ff_dim,
222 + multiple_of=config.block_multiple_of,
223 + auto_adjust_ff_dim=config.block_auto_adjust_ff_dim,
224 + ffn_dim_multiplier=config.block_ffn_dim_multiplier,
225 + quant_config=quant_config,
226 + prefix=f"{prefix}.feed_forward",
227 + )
228 + self.operator_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
229 + self.ffn_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
230 +
231 + def forward(
232 + self,
233 + positions: torch.Tensor,
234 + hidden_states: torch.Tensor,
235 + residual: Optional[torch.Tensor],
236 + **kwargs,
237 + ) -> tuple[torch.Tensor, torch.Tensor]:
238 + if residual is None:
239 + residual = hidden_states
240 + hidden_states = self.operator_norm(hidden_states)
241 + else:
242 + hidden_states, residual = self.operator_norm(
243 + hidden_states, residual)
244 + hidden_states = self.self_attn(positions=positions,
245 + hidden_states=hidden_states)
246 + hidden_states, residual = self.ffn_norm(hidden_states, residual)
247 + return self.feed_forward(hidden_states), residual
248 +
249 +
250 +class Lfm2ShortConvDecoderLayer(nn.Module):
251 +
252 + def __init__(
253 + self,
254 + config: Lfm2Config,
255 + layer_idx: int,
256 + model_config: Optional[ModelConfig] = None,
257 + cache_config: Optional[CacheConfig] = None,
258 + quant_config: Optional[QuantizationConfig] = None,
259 + prefix: str = "",
260 + ) -> None:
261 + super().__init__()
262 + self.layer_idx = layer_idx
263 + self.conv = ShortConv(
264 + config=config,
265 + dim=config.conv_dim,
266 + layer_idx=layer_idx,
267 + model_config=model_config,
268 + cache_config=cache_config,
269 + prefix=f"{prefix}.conv",
270 + )
271 +
272 + self.feed_forward = Lfm2MLP(
273 + dim=config.block_dim,
274 + ff_dim=config.block_ff_dim,
275 + multiple_of=config.block_multiple_of,
276 + auto_adjust_ff_dim=config.block_auto_adjust_ff_dim,
277 + ffn_dim_multiplier=config.block_ffn_dim_multiplier,
278 + quant_config=quant_config,
279 + prefix=f"{prefix}.feed_forward",
280 + )
281 + self.operator_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
282 + self.ffn_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
283 +
284 + def forward(
285 + self,
286 + hidden_states: torch.Tensor,
287 + residual: Optional[torch.Tensor],
288 + **kwargs,
289 + ):
290 + if residual is None:
291 + residual = hidden_states
292 + hidden_states = self.operator_norm(hidden_states)
293 + else:
294 + hidden_states, residual = self.operator_norm(
295 + hidden_states, residual)
296 + output = torch.empty_like(hidden_states)
297 + self.conv(
298 + hidden_states,
299 + output,
300 + conv_metadata=None,
301 + )
302 + hidden_states, residual = self.ffn_norm(output, residual)
303 + hidden_states = self.feed_forward(hidden_states)
304 + return hidden_states, residual
305 +
306 +
307 +@support_torch_compile
308 +class Lfm2Model(nn.Module):
309 +
310 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
311 + super().__init__()
312 +
313 + config = vllm_config.model_config.hf_config
314 + model_config = vllm_config.model_config
315 + cache_config = vllm_config.cache_config
316 + quant_config = vllm_config.quant_config
317 + lora_config = vllm_config.lora_config
318 +
319 + self.config = config
320 + lora_vocab = ((lora_config.lora_extra_vocab_size *
321 + (lora_config.max_loras or 1)) if lora_config else 0)
322 + self.vocab_size = config.vocab_size + lora_vocab
323 + self.org_vocab_size = config.vocab_size
324 +
325 + self.embed_tokens = VocabParallelEmbedding(
326 + self.vocab_size,
327 + config.hidden_size,
328 + org_num_embeddings=config.vocab_size)
329 +
330 + def get_layer(prefix: str):
331 + layer_idx = extract_layer_index(prefix)
332 + is_attn = self.config.layer_types[layer_idx] == "full_attention"
333 + layer_class = (Lfm2AttentionDecoderLayer
334 + if is_attn else Lfm2ShortConvDecoderLayer)
335 + return layer_class(
336 + config,
337 + layer_idx,
338 + model_config,
339 + cache_config,
340 + quant_config=quant_config,
341 + prefix=prefix,
342 + )
343 +
344 + self.start_layer, self.end_layer, self.layers = make_layers(
345 + config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers")
346 + self.make_empty_intermediate_tensors = (
347 + make_empty_intermediate_tensors_factory(
348 + ["hidden_states", "residual"], config.hidden_size))
349 +
350 + if get_pp_group().is_last_rank:
351 + self.embedding_norm = RMSNorm(config.hidden_size,
352 + eps=config.norm_eps)
353 + else:
354 + self.embedding_norm = PPMissingLayer()
355 +
356 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
357 + return self.embed_tokens(input_ids)
358 +
359 + def forward(
360 + self,
361 + input_ids: torch.Tensor,
362 + positions: torch.Tensor,
363 + intermediate_tensors: Optional[IntermediateTensors] = None,
364 + inputs_embeds: Optional[torch.Tensor] = None,
365 + ) -> torch.Tensor:
366 + if get_pp_group().is_first_rank:
367 + if inputs_embeds is not None:
368 + hidden_states = inputs_embeds
369 + else:
370 + hidden_states = self.get_input_embeddings(input_ids)
371 + residual = None
372 + else:
373 + assert intermediate_tensors is not None
374 + hidden_states = intermediate_tensors["hidden_states"]
375 + residual = intermediate_tensors["residual"]
376 +
377 + for layer in self.layers[self.start_layer:self.end_layer]:
378 + hidden_states, residual = layer(
379 + positions=positions,
380 + hidden_states=hidden_states,
381 + residual=residual,
382 + )
383 + if not get_pp_group().is_last_rank:
384 + return IntermediateTensors({
385 + "hidden_states": hidden_states,
386 + "residual": residual
387 + })
388 + hidden_states, _ = self.embedding_norm(hidden_states, residual)
389 + return hidden_states
390 +
391 + def load_weights(self, weights: Iterable[tuple[str,
392 + torch.Tensor]]) -> set[str]:
393 + stacked_params_mapping = [
394 + (".qkv_proj", ".q_proj", "q"),
395 + (".qkv_proj", ".k_proj", "k"),
396 + (".qkv_proj", ".v_proj", "v"),
397 + (".w1", ".w1", 0),
398 + (".w1", ".w3", 1),
399 + ]
400 + params_dict = dict(self.named_parameters())
401 + loaded_params: set[str] = set()
402 + for name, loaded_weight in weights:
403 +
404 + for param_name, weight_name, shard_id in stacked_params_mapping:
405 + if weight_name not in name:
406 + continue
407 + name = name.replace(weight_name, param_name)
408 +
409 + if is_pp_missing_parameter(name, self):
410 + continue
411 + param = params_dict[name]
412 + weight_loader = param.weight_loader
413 + weight_loader(param, loaded_weight, shard_id)
414 + break
415 + else:
416 + if is_pp_missing_parameter(name, self):
417 + continue
418 + param = params_dict[name]
419 + weight_loader = getattr(param, "weight_loader",
420 + default_weight_loader)
421 + weight_loader(param, loaded_weight)
422 + loaded_params.add(name)
423 + return loaded_params
424 +
425 +
426 +class Lfm2ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
427 + IsHybrid, SupportsQuant):
428 + packed_modules_mapping = {
429 + "qkv_proj": [
430 + "q_proj",
431 + "k_proj",
432 + "v_proj",
433 + ],
434 + "w1": [
435 + "w1",
436 + "w3",
437 + ],
438 + }
439 +
440 + # LoRA specific attributes
441 + embedding_modules = {
442 + "embed_tokens": "input_embeddings",
443 + "lm_head": "output_embeddings",
444 + }
445 + embedding_padding_modules = ["lm_head"]
446 +
447 + @classmethod
448 + def get_mamba_state_dtype_from_config(
449 + cls,
450 + vllm_config: "VllmConfig",
451 + ) -> tuple[torch.dtype, ...]:
452 +
453 + return MambaStateDtypeCalculator.short_conv_state_dtype(
454 + vllm_config.model_config.dtype,
455 + vllm_config.cache_config.mamba_cache_dtype,
456 + )
457 +
458 + @classmethod
459 + def get_mamba_state_shape_from_config(
460 + cls,
461 + vllm_config: "VllmConfig",
462 + use_v1: bool = True,
463 + ) -> tuple[tuple[int, int]]:
464 + """ Calculate shapes for LFM2's convolutional cache.
465 +
466 + Args:
467 + vllm_config: vLLM config
468 + use_v1: Get shapes for V1 (or V0)
469 +
470 + Returns:
471 + Tuple containing:
472 + - conv_state_shape: Shape for convolutional state cache
473 + """
474 + parallel_config = vllm_config.parallel_config
475 + hf_config = vllm_config.model_config.hf_config
476 +
477 + return MambaStateShapeCalculator.short_conv_state_shape(
478 + tp_world_size=parallel_config.tensor_parallel_size,
479 + intermediate_size=hf_config.conv_dim,
480 + conv_kernel=hf_config.conv_L_cache,
481 + use_v1=use_v1,
482 + )
483 +
484 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
485 + config = vllm_config.model_config.hf_config
486 + quant_config = vllm_config.quant_config
487 + cache_config = vllm_config.cache_config
488 + lora_config = vllm_config.lora_config
489 + scheduler_config = vllm_config.scheduler_config
490 + assert (not cache_config.enable_prefix_caching
491 + ), "Lfm2 currently does not support prefix caching"
492 + assert envs.VLLM_USE_V1, (
493 + "Lfm2ForCausalLM doesn't support vLLM v0. Please enable v1")
494 +
495 + super().__init__()
496 + self.config = config
497 + self.vllm_config = vllm_config
498 + self.scheduler_config = scheduler_config
499 + self.model_config = vllm_config.model_config
500 +
501 + self.model = Lfm2Model(vllm_config=vllm_config,
502 + prefix=maybe_prefix(prefix, "model"))
503 +
504 + if get_pp_group().is_last_rank:
505 + self.unpadded_vocab_size = self.config.vocab_size
506 + if lora_config:
507 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
508 +
509 + self.lm_head = ParallelLMHead(
510 + self.unpadded_vocab_size,
511 + config.hidden_size,
512 + org_num_embeddings=config.vocab_size,
513 + padding_size=(
514 + DEFAULT_VOCAB_PADDING_SIZE
515 + # We need bigger padding if using lora for kernel
516 + # compatibility
517 + if not lora_config else
518 + lora_config.lora_vocab_padding_size),
519 + quant_config=quant_config,
520 + prefix=maybe_prefix(prefix, "lm_head"),
521 + )
522 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
523 + else:
524 + self.lm_head = PPMissingLayer()
525 +
526 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
527 + config.vocab_size)
528 +
529 + self.make_empty_intermediate_tensors = (
530 + self.model.make_empty_intermediate_tensors)
531 +
532 + def forward(
533 + self,
534 + input_ids: torch.Tensor,
535 + positions: torch.Tensor,
536 + intermediate_tensors: Optional[IntermediateTensors] = None,
537 + inputs_embeds: Optional[torch.Tensor] = None,
538 + **kwargs,
539 + ) -> torch.Tensor:
540 + hidden_states = self.model(input_ids, positions, intermediate_tensors,
541 + inputs_embeds)
542 + return hidden_states
543 +
544 + def compute_logits(self, hidden_states: torch.Tensor,
545 + sampling_metadata: SamplingMetadata) -> torch.Tensor:
546 + logits = self.logits_processor(self.lm_head, hidden_states,
547 + sampling_metadata)
548 + return logits
549 +
550 + def load_weights(self, weights: Iterable[tuple[str,
551 + torch.Tensor]]) -> set[str]:
552 + loader = AutoWeightsLoader(
553 + self,
554 + skip_prefixes=(["lm_head."]
555 + if self.config.tie_word_embeddings else None),
556 + )
557 + return loader.load_weights(weights)
\ No newline at end of file
vllm/model_executor/models/registry.py
+1
@@ -93,6 +93,7 @@ _TEXT_GENERATION_MODELS = {
93 "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),
94 "JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
95 "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
96 + "Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"),
97 "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
98 "Llama4ForCausalLM": ("llama4", "Llama4ForCausalLM"), # noqa: E501
99 # For decapoda-research/llama-*
vllm/v1/attention/backends/mamba_selectors.py
+4
@@ -4,6 +4,8 @@ from vllm.attention.backends.abstract import AttentionBackend
4 from vllm.v1.attention.backends.linear_attn import LinearAttentionBackend
5 from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionBackend
6 from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionBackend
7 +from vllm.v1.attention.backends.short_conv_attn import (
8 + ShortConvAttentionBackend)
9
10
11 def get_mamba_attn_backend(mamba_type: str) -> type[AttentionBackend]:
@@ -13,6 +15,8 @@ def get_mamba_attn_backend(mamba_type: str) -> type[AttentionBackend]:
15 return Mamba2AttentionBackend
16 if mamba_type == "linear_attention":
17 return LinearAttentionBackend
18 + if mamba_type == "short_conv":
19 + return ShortConvAttentionBackend
20
21 raise NotImplementedError(f"Mamba Attention type {mamba_type} is not "
22 "supported yet.")
vllm/v1/attention/backends/short_conv_attn.py new
+81
@@ -0,0 +1,81 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from dataclasses import dataclass
4 +from typing import ClassVar, Optional
5 +
6 +import torch
7 +
8 +from vllm.attention.backends.abstract import AttentionBackend
9 +from vllm.config import VllmConfig
10 +from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
11 + CommonAttentionMetadata,
12 + split_decodes_and_prefills)
13 +from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
14 +
15 +
16 +class ShortConvAttentionBackend(AttentionBackend):
17 +
18 + @staticmethod
19 + def get_builder_cls() -> type["ShortConvAttentionMetadataBuilder"]:
20 + return ShortConvAttentionMetadataBuilder
21 +
22 +
23 +@dataclass
24 +class ShortConvAttentionMetadata:
25 + num_prefills: int
26 + num_prefill_tokens: int
27 + num_decodes: int
28 + num_decode_tokens: int
29 +
30 + query_start_loc: torch.Tensor
31 + has_initial_states: torch.Tensor
32 + state_indices_tensor: torch.Tensor # shape: [batch,]
33 +
34 + # For causal_conv1d
35 + nums_dict: Optional[dict] = None
36 + cu_seqlen: Optional[int] = None
37 + batch_ptr: Optional[torch.tensor] = None
38 + token_chunk_offset_ptr: Optional[torch.tensor] = None
39 +
40 +
41 +class ShortConvAttentionMetadataBuilder(
42 + AttentionMetadataBuilder[ShortConvAttentionMetadata]):
43 +
44 + reorder_batch_threshold: ClassVar[int] = 1
45 +
46 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
47 + vllm_config: VllmConfig, device: torch.device):
48 + assert isinstance(kv_cache_spec, MambaSpec)
49 + self.kv_cache_spec = kv_cache_spec
50 +
51 + def build(self,
52 + common_prefix_len: int,
53 + common_attn_metadata: CommonAttentionMetadata,
54 + fast_build: bool = False) -> ShortConvAttentionMetadata:
55 + num_reqs = common_attn_metadata.num_reqs
56 + query_start_loc = common_attn_metadata.query_start_loc
57 +
58 + state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
59 +
60 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
61 + split_decodes_and_prefills(common_attn_metadata,
62 + decode_threshold=1))
63 + has_initial_states = None
64 + if num_prefills > 0:
65 + #[batch,]
66 + has_initial_states_cpu = (
67 + common_attn_metadata.
68 + num_computed_tokens_cpu[num_reqs - num_prefills:num_reqs] > 0)
69 + has_initial_states = has_initial_states_cpu.to(
70 + query_start_loc.device)
71 +
72 + attn_metadata = ShortConvAttentionMetadata(
73 + num_prefills=num_prefills,
74 + num_prefill_tokens=num_prefill_tokens,
75 + num_decodes=num_decodes,
76 + num_decode_tokens=num_decode_tokens,
77 + query_start_loc=query_start_loc,
78 + has_initial_states=has_initial_states,
79 + state_indices_tensor=state_indices_tensor,
80 + )
81 + return attn_metadata
\ No newline at end of file