@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 452e8fd968
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 15, 18:30 UTC

[MODEL] Add support for Zamba2 models (#13185)

Signed-off-by: Yury Tokpanov <yury@zyphra.com> Signed-off-by: Quentin Anthony <qganthony@yahoo.com> Co-authored-by: Quentin Anthony <qganthony@yahoo.com> Co-authored-by: Tyler Michael Smith <tysmith@redhat.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

yury-tokpanov committed Mar 18, 2025 at 08:56 UTC 452e8fd9684283d1e0bed63a0c8dc1e3141f788e
9 files changed +1082 -27
docs/source/models/supported_models.md
+5
@@ -477,6 +477,11 @@ See [this page](#generative-models) for more information on how to use generativ
477 * `xverse/XVERSE-7B-Chat`, `xverse/XVERSE-13B-Chat`, `xverse/XVERSE-65B-Chat`, etc.
478 * ✅︎
479 * ✅︎
480 +- * `Zamba2ForCausalLM`
481 + * Zamba2
482 + * `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc.
483 + *
484 + *
485 :::
486
487 :::{note}
tests/models/decoder_only/language/test_hybrid.py
+29 -22
@@ -9,7 +9,7 @@ from vllm.sampling_params import SamplingParams
9 from ...utils import check_outputs_equal
10
11 # This test is for the hybrid models
12 -MODELS = ["ai21labs/Jamba-tiny-dev"]
12 +MODELS = ["ai21labs/Jamba-tiny-dev", "Zyphra/Zamba2-1.2B-instruct"]
13 # Bamba at Fp32 is too big for the CI (L4 GPU).
14 # MODELS = ["ai21labs/Jamba-tiny-dev", "ibm-ai-platform/Bamba-9B"]
15
@@ -27,17 +27,19 @@ def test_models(
27 ) -> None:
28
29 # numeric error produces different generation
30 - if 'Bamba' in model:
30 + if "Bamba" in model:
31 example_prompts.pop(3)
32
33 - with hf_runner(
34 - model,
35 - dtype=dtype,
36 - model_kwargs={
37 - "use_mamba_kernels":
38 - False, # mamba kernels are not installed so HF
39 - # don't use them
40 - }) as hf_model:
33 + model_kwargs = {
34 + "use_mamba_kernels": False, # mamba kernels are not installed so HF
35 + # don't use them
36 + }
37 + if "Zamba2" in model:
38 + # Zamba2 HF implementation automatically checks if mamba kernels are
39 + # installed
40 + model_kwargs = {}
41 +
42 + with hf_runner(model, dtype=dtype, model_kwargs=model_kwargs) as hf_model:
43 hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
44
45 with vllm_runner(model, dtype=dtype) as vllm_model:
@@ -112,26 +114,31 @@ def test_mamba_prefill_chunking_with_parallel_sampling(
114 def test_mamba_prefill_chunking(hf_runner, vllm_runner, example_prompts,
115 model: str, dtype: str,
116 max_tokens: int) -> None:
115 - # numeric error during prefill chucking produces different generation
117 + # numeric error during prefill chunking produces different generation
118 # compared to w/o prefill chunking for those examples, removed them for now
117 - if 'Jamba' in model:
119 + if "Jamba" in model:
120 example_prompts.pop(7)
121 example_prompts.pop(2)
122 example_prompts.pop(1)
121 - elif 'Bamba' in model:
123 + elif "Bamba" in model:
124 example_prompts.pop(6)
125 example_prompts.pop(3)
126 example_prompts.pop(2)
127 dtype = "half" # use a different dtype for Bamba
126 -
127 - with hf_runner(
128 - model,
129 - dtype=dtype,
130 - model_kwargs={
131 - "use_mamba_kernels":
132 - False, # mamba kernels are not installed so HF
133 - # don't use them
134 - }) as hf_model:
128 + elif "Zamba2" in model:
129 + example_prompts.pop(7)
130 + dtype = "half"
131 +
132 + model_kwargs = {
133 + "use_mamba_kernels": False, # mamba kernels are not installed so HF
134 + # don't use them
135 + }
136 + if "Zamba2" in model:
137 + # Zamba2 HF implementation automatically checks if mamba kernels are
138 + # installed
139 + model_kwargs = {}
140 +
141 + with hf_runner(model, dtype=dtype, model_kwargs=model_kwargs) as hf_model:
142 non_chunked = hf_model.generate_greedy(example_prompts, max_tokens)
143
144 with vllm_runner(model,
tests/models/registry.py
+2
@@ -195,6 +195,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
195 "XverseForCausalLM": _HfExamplesInfo("xverse/XVERSE-7B-Chat",
196 is_available_online=False,
197 trust_remote_code=True),
198 + "Zamba2ForCausalLM": _HfExamplesInfo("Zyphra/Zamba2-7B-instruct",
199 + min_transformers_version="4.49"),
200 # [Encoder-decoder]
201 "BartModel": _HfExamplesInfo("facebook/bart-base"),
202 "BartForConditionalGeneration": _HfExamplesInfo("facebook/bart-large-cnn"),
vllm/config.py
+14
@@ -821,6 +821,11 @@ class ModelConfig:
821 if qk_rope_head_dim and qk_nope_head_dim:
822 return qk_rope_head_dim + qk_nope_head_dim
823
824 + if hasattr(self.hf_text_config,
825 + "model_type") and (self.hf_text_config.model_type
826 + == "zamba2"):
827 + return self.hf_text_config.attention_head_dim
828 +
829 if self.is_attention_free:
830 return 0
831
@@ -944,6 +949,15 @@ class ModelConfig:
949 "cannot determine the num of "
950 f"{block_type.value} layers")
951
952 + if hasattr(self.hf_text_config,
953 + "model_type") and (self.hf_text_config.model_type
954 + == "zamba2"):
955 + if attn_block_type:
956 + return sum(t == "hybrid"
957 + for t in layers_block_type_value[start:end])
958 + else:
959 + return self.get_num_layers(parallel_config)
960 +
961 return sum(t == block_type.value
962 for t in layers_block_type_value[start:end])
963
vllm/model_executor/layers/mamba/mamba_mixer2.py
-1
@@ -245,7 +245,6 @@ class MambaMixer2(CustomOp):
245 assert num_heads % self.tp_size == 0, \
246 "Tensor parallel world size must divide num heads."
247
248 -
248 assert (n_groups % self.tp_size) == 0 or n_groups == 1, \
249 (
250 "If tensor parallel world size does not divide num_heads, "
vllm/model_executor/models/bamba.py
-2
@@ -38,8 +38,6 @@ from .utils import (is_pp_missing_parameter,
38 make_empty_intermediate_tensors_factory, make_layers,
39 maybe_prefix)
40
41 -KVCache = Tuple[torch.Tensor, torch.Tensor]
42 -
41
42 class BambaMLP(nn.Module):
43
vllm/model_executor/models/jamba.py
-2
@@ -36,8 +36,6 @@ from .utils import (is_pp_missing_parameter,
36 make_empty_intermediate_tensors_factory, make_layers,
37 maybe_prefix)
38
39 -KVCache = Tuple[torch.Tensor, torch.Tensor]
40 -
39
40 class JambaMoE(nn.Module):
41
vllm/model_executor/models/registry.py
+1
@@ -105,6 +105,7 @@ _TEXT_GENERATION_MODELS = {
105 "SolarForCausalLM": ("solar", "SolarForCausalLM"),
106 "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
107 "XverseForCausalLM": ("llama", "LlamaForCausalLM"),
108 + "Zamba2ForCausalLM": ("zamba2", "Zamba2ForCausalLM"),
109 # [Encoder-decoder]
110 "BartModel": ("bart", "BartForConditionalGeneration"),
111 "BartForConditionalGeneration": ("bart", "BartForConditionalGeneration"),
vllm/model_executor/models/zamba2.py new
+1031
@@ -0,0 +1,1031 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""PyTorch Zamba2 model implementation for vLLM.
3 +
4 +This module implements the Zamba2 architecture from
5 +https://arxiv.org/abs/2411.15242, which combines Mamba and Transformer
6 +architectures in a hybrid model optimized for efficient sequence modeling. The
7 +model alternates between state space model layers and attention-based layers.
8 +"""
9 +from itertools import cycle
10 +from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
11 +
12 +import torch
13 +from torch import nn
14 +from transformers import Zamba2Config
15 +
16 +from vllm.attention.layer import Attention
17 +from vllm.config import CacheConfig, VllmConfig
18 +from vllm.distributed import divide, get_tensor_model_parallel_world_size
19 +from vllm.forward_context import get_forward_context
20 +from vllm.model_executor.layers.activation import GeluAndMul
21 +from vllm.model_executor.layers.layernorm import RMSNorm
22 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,
23 + MergedColumnParallelLinear,
24 + QKVParallelLinear,
25 + ReplicatedLinear,
26 + RowParallelLinear)
27 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
28 +from vllm.model_executor.layers.mamba.mamba_mixer2 import (
29 + MambaMixer2, extra_groups_for_head_shards)
30 +from vllm.model_executor.layers.quantization import QuantizationConfig
31 +from vllm.model_executor.layers.rotary_embedding import get_rope
32 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
33 +from vllm.model_executor.layers.vocab_parallel_embedding import (
34 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
35 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
36 +from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
37 + MambaCacheParams)
38 +from vllm.model_executor.sampling_metadata import SamplingMetadata
39 +from vllm.sequence import IntermediateTensors
40 +
41 +from .interfaces import HasInnerState, IsHybrid, SupportsV0Only
42 +from .utils import maybe_prefix
43 +
44 +
45 +class Zamba2LoRA(nn.Module):
46 + """LoRA layer for the Zamba2 model.
47 +
48 + Implements a LoRA layer that is used in shared attention and gated MLP
49 + blocks.
50 + """
51 +
52 + def __init__(
53 + self,
54 + input_dim: int,
55 + rank: int,
56 + output_dim: Union[int, List[int]],
57 + quant_config: Optional[QuantizationConfig] = None,
58 + ):
59 + """Initialize the attention layer.
60 +
61 + Args:
62 + input_dim: input dimension
63 + rank: LoRA rank
64 + output_dim: output dimension
65 + quant_config: Configuration for model quantization
66 + """
67 + super().__init__()
68 +
69 + self.A = ColumnParallelLinear(input_dim,
70 + rank,
71 + bias=False,
72 + quant_config=quant_config,
73 + gather_output=True)
74 +
75 + if isinstance(output_dim, list):
76 + B_class = MergedColumnParallelLinear
77 + else:
78 + B_class = ColumnParallelLinear
79 + self.B = B_class(rank,
80 + output_dim,
81 + bias=False,
82 + quant_config=quant_config)
83 +
84 + def forward(
85 + self,
86 + hidden_states: torch.Tensor,
87 + ):
88 + lora_output, _ = self.A(hidden_states)
89 + lora_output, _ = self.B(lora_output)
90 + return lora_output
91 +
92 +
93 +class Zamba2Attention(nn.Module):
94 + """Multi-head attention mechanism for the Zamba2 model.
95 +
96 + Implements attention with parallel computation, QKV projections, optional
97 + adapters and rotary position embeddings. The attention is computed across
98 + distributed blocks for efficient processing.
99 + """
100 +
101 + def __init__(
102 + self,
103 + config: Zamba2Config,
104 + bare_block_idx: int,
105 + num_hybrid_layers: int,
106 + cache_config: Optional[CacheConfig] = None,
107 + quant_config: Optional[QuantizationConfig] = None,
108 + prefix: str = "",
109 + ) -> None:
110 + """Initialize the attention layer.
111 +
112 + Args:
113 + config: The Zamba2 model configuration
114 + bare_block_idx: Index of the bare attention block
115 + num_hybrid_layers: Total number of hybrid layers
116 + cache_config: Configuration for key-value caching
117 + quant_config: Configuration for model quantization
118 + prefix: Optional prefix for parameter names
119 + """
120 + super().__init__()
121 + tp_size = get_tensor_model_parallel_world_size()
122 + self.config = config
123 + self.num_hybrid_layers = num_hybrid_layers
124 + self.rope_theta = config.rope_theta
125 +
126 + self.attention_hidden_size = config.attention_hidden_size
127 + self.total_num_attention_heads = config.num_attention_heads
128 + assert self.total_num_attention_heads % tp_size == 0
129 + self.num_attention_heads = config.num_attention_heads // tp_size
130 + self.attention_head_dim = config.attention_head_dim
131 + self.qkv_size = self.attention_hidden_size // tp_size
132 + self.scale = (self.attention_head_dim / 2)**-0.5
133 +
134 + if (self.attention_head_dim *
135 + self.total_num_attention_heads) != self.attention_hidden_size:
136 + raise ValueError(
137 + f"attention_hidden_size must be divisible by"
138 + f" num_attention_heads"
139 + f" (got `attention_hidden_size`: {self.attention_hidden_size}"
140 + f" and `num_heads`: {self.num_attention_heads}).")
141 +
142 + self.qkv_proj = QKVParallelLinear(
143 + self.attention_hidden_size,
144 + self.attention_head_dim,
145 + self.total_num_attention_heads,
146 + bias=False,
147 + quant_config=quant_config,
148 + )
149 + self.o_proj = RowParallelLinear(self.attention_hidden_size,
150 + config.hidden_size,
151 + bias=False,
152 + quant_config=quant_config)
153 +
154 + # Even though in Zamba2 weights are shared between attention layers, KV
155 + # cache is unique for every attention layer. Hence, we need to define
156 + # separate Attention objects, because in recent vLLM KV cache tensors
157 + # are tied to specific Attention objects.
158 +
159 + # Initialize attention blocks with proper indexing
160 + self.dpa_list = nn.ModuleList([])
161 + j = bare_block_idx * (self.num_hybrid_layers + config.num_mem_blocks -
162 + 1) // config.num_mem_blocks
163 + for block_idx in range(self.num_hybrid_layers):
164 + if block_idx % config.num_mem_blocks == bare_block_idx:
165 + dpa = Attention(
166 + self.num_attention_heads,
167 + self.attention_head_dim,
168 + self.scale,
169 + cache_config=cache_config,
170 + prefix=f"{prefix}.attn.{j}",
171 + )
172 + j += 1
173 + else:
174 + dpa = nn.Identity()
175 + self.dpa_list.append(dpa)
176 +
177 + # Initialize adapter layers if enabled
178 + if config.use_shared_attention_adapter:
179 + self.linear_q_adapter_list = nn.ModuleList([])
180 + self.linear_k_adapter_list = nn.ModuleList([])
181 + self.linear_v_adapter_list = nn.ModuleList([])
182 +
183 + for block_idx in range(self.num_hybrid_layers):
184 + if block_idx % config.num_mem_blocks == bare_block_idx:
185 + linear_q_adapter = Zamba2LoRA(
186 + self.attention_hidden_size,
187 + config.adapter_rank,
188 + self.attention_hidden_size,
189 + quant_config=quant_config,
190 + )
191 + linear_k_adapter = Zamba2LoRA(
192 + self.attention_hidden_size,
193 + config.adapter_rank,
194 + self.attention_hidden_size,
195 + quant_config=quant_config,
196 + )
197 + linear_v_adapter = Zamba2LoRA(
198 + self.attention_hidden_size,
199 + config.adapter_rank,
200 + self.attention_hidden_size,
201 + quant_config=quant_config,
202 + )
203 + else:
204 + linear_q_adapter = nn.Identity()
205 + linear_k_adapter = nn.Identity()
206 + linear_v_adapter = nn.Identity()
207 +
208 + self.linear_q_adapter_list.append(linear_q_adapter)
209 + self.linear_k_adapter_list.append(linear_k_adapter)
210 + self.linear_v_adapter_list.append(linear_v_adapter)
211 +
212 + if config.use_mem_rope:
213 + self.rotary_emb = get_rope(
214 + head_size=self.attention_head_dim,
215 + rotary_dim=self.attention_head_dim,
216 + max_position=config.max_position_embeddings,
217 + base=self.rope_theta,
218 + rope_scaling=None,
219 + is_neox_style=True,
220 + )
221 +
222 + def forward(
223 + self,
224 + hidden_states: torch.Tensor,
225 + block_idx: int,
226 + position_ids: torch.Tensor,
227 + ) -> torch.Tensor:
228 + """Forward pass through the attention layer.
229 +
230 + Args:
231 + hidden_states: Input tensor [batch_size, seq_len, hidden_size]
232 + position_ids: Position IDs for positional embeddings
233 + block_idx: Current shared transformer block index
234 +
235 + Returns:
236 + Output tensor [batch_size, seq_len, hidden_size]
237 + """
238 + qkv, _ = self.qkv_proj(hidden_states)
239 + query_states, key_states, value_states = qkv.split([self.qkv_size] * 3,
240 + dim=-1)
241 +
242 + if self.config.use_shared_attention_adapter:
243 + # Apply adapter transformations to Q, K, V if enabled
244 + q_adapter = self.linear_q_adapter_list[block_idx]
245 + assert not isinstance(q_adapter, nn.Identity)
246 + q_lora_output = q_adapter(hidden_states)
247 + query_states = query_states + q_lora_output
248 +
249 + k_adapter = self.linear_k_adapter_list[block_idx]
250 + assert not isinstance(k_adapter, nn.Identity)
251 + k_lora_output = k_adapter(hidden_states)
252 + key_states = key_states + k_lora_output
253 +
254 + v_adapter = self.linear_v_adapter_list[block_idx]
255 + assert not isinstance(v_adapter, nn.Identity)
256 + v_lora_output = v_adapter(hidden_states)
257 + value_states = value_states + v_lora_output
258 +
259 + if self.config.use_mem_rope:
260 + query_states, key_states = self.rotary_emb(position_ids,
261 + query_states,
262 + key_states)
263 +
264 + y = self.dpa_list[block_idx](query_states, key_states, value_states)
265 + y, _ = self.o_proj(y)
266 + return y
267 +
268 +
269 +class Zamba2MLP(nn.Module):
270 + """Feed-forward MLP layer for the Zamba2 model.
271 +
272 + Implements a gated feed-forward network that projects inputs to a larger
273 + intermediate size, applies GELU activation with gating, then projects back
274 + to the original size. Includes optional adapter layers for model adaptation.
275 + """
276 +
277 + def __init__(
278 + self,
279 + config: Zamba2Config,
280 + bare_block_idx: int,
281 + num_hybrid_layers: Dict[int, int],
282 + quant_config: Optional[QuantizationConfig] = None,
283 + ) -> None:
284 + """Initialize the MLP layer.
285 +
286 + Args:
287 + config: The Zamba2 model configuration
288 + bare_block_idx: Index of the bare block in the model
289 + num_hybrid_layers: Total number of hybrid layers
290 + quant_config: Configuration for model quantization
291 + """
292 + super().__init__()
293 + self.config = config
294 + self.tp_size = get_tensor_model_parallel_world_size()
295 + self.num_hybrid_layers = num_hybrid_layers
296 + self.hidden_size = config.hidden_size
297 + self.intermediate_size = config.intermediate_size
298 +
299 + # Main projection layers with gating
300 + self.gate_up_proj = MergedColumnParallelLinear(
301 + self.hidden_size,
302 + 2 * [self.intermediate_size], # 2x for gate and input projections
303 + bias=self.config.add_bias_linear,
304 + quant_config=quant_config)
305 +
306 + self.down_proj = RowParallelLinear(self.intermediate_size,
307 + self.hidden_size,
308 + bias=self.config.add_bias_linear,
309 + quant_config=quant_config)
310 +
311 + # Only allow GELU activations
312 + if config.hidden_act != "gelu":
313 + raise ValueError(f"Only GELU activation is supported "
314 + f"(got `hidden_act`: {config.hidden_act})")
315 + self.act_fn = GeluAndMul()
316 +
317 + # Initialize adapter layers
318 + self.gate_up_proj_adapter_list = nn.ModuleList([])
319 + for block_idx in range(self.num_hybrid_layers):
320 + if block_idx % config.num_mem_blocks == bare_block_idx:
321 + gate_up_proj_adapter = Zamba2LoRA(
322 + config.hidden_size,
323 + config.adapter_rank,
324 + 2 * [self.intermediate_size],
325 + quant_config,
326 + )
327 + else:
328 + gate_up_proj_adapter = nn.Identity()
329 + self.gate_up_proj_adapter_list.append(gate_up_proj_adapter)
330 +
331 + def forward(self, hidden_states: torch.Tensor,
332 + block_idx: int) -> torch.Tensor:
333 + """Forward pass through the MLP layer.
334 +
335 + Args:
336 + hidden_states: Input tensor [batch_size, seq_len, hidden_size]
337 + block_idx: Current shared transformer block index
338 +
339 + Returns:
340 + Output tensor [batch_size, seq_len, hidden_size] after applying
341 + gated feed-forward transformation
342 + """
343 + # Project input to intermediate size with gating
344 + gate_up_states, _ = self.gate_up_proj(hidden_states)
345 +
346 + # Apply adapter transformation if present
347 + adapter = self.gate_up_proj_adapter_list[block_idx]
348 + assert not isinstance(adapter, nn.Identity)
349 + lora_output = adapter(hidden_states)
350 + gate_up_states = gate_up_states + lora_output
351 +
352 + # Apply GELU activation with gating
353 + hidden_states = self.act_fn(gate_up_states)
354 +
355 + # Project back to hidden size
356 + output, _ = self.down_proj(hidden_states)
357 + return output
358 +
359 +
360 +class Zamba2AttentionDecoderLayer(nn.Module):
361 + """Single decoder layer combining attention and feed-forward networks.
362 +
363 + This layer implements a standard transformer block with:
364 + - Input layer normalization
365 + - Multi-head self-attention
366 + - Pre-feed-forward layer normalization
367 + - Feed-forward network (MLP)
368 + """
369 +
370 + def __init__(
371 + self,
372 + config: Zamba2Config,
373 + bare_block_idx: int,
374 + num_hybrid_layers: int,
375 + cache_config: Optional[CacheConfig] = None,
376 + quant_config: Optional[QuantizationConfig] = None,
377 + prefix: str = "",
378 + ) -> None:
379 + """Initialize the decoder layer.
380 +
381 + Args:
382 + config: The Zamba2 model configuration
383 + bare_block_idx: Index of the bare block
384 + num_hybrid_layers: Total number of hybrid layers
385 + cache_config: Configuration for key-value caching
386 + quant_config: Configuration for model quantization
387 + prefix: Optional prefix for parameter names
388 + """
389 + super().__init__()
390 +
391 + # Initialize attention sublayer
392 + self.self_attn = Zamba2Attention(
393 + config,
394 + bare_block_idx=bare_block_idx,
395 + num_hybrid_layers=num_hybrid_layers,
396 + cache_config=cache_config,
397 + quant_config=quant_config,
398 + prefix=prefix,
399 + )
400 +
401 + # Initialize feed-forward sublayer
402 + self.feed_forward = Zamba2MLP(
403 + config,
404 + bare_block_idx=bare_block_idx,
405 + num_hybrid_layers=num_hybrid_layers,
406 + quant_config=quant_config,
407 + )
408 +
409 + # Initialize layer normalizations
410 + # Input normalization operates on concatenated states
411 + self.input_layernorm = RMSNorm(2 * config.hidden_size,
412 + eps=config.rms_norm_eps)
413 + # Pre-FF normalization operates on attention output
414 + self.pre_ff_layernorm = RMSNorm(config.hidden_size,
415 + eps=config.rms_norm_eps)
416 +
417 + def forward(
418 + self,
419 + hidden_states: torch.Tensor,
420 + original_hidden_states: torch.Tensor,
421 + block_idx: int,
422 + positions: torch.Tensor,
423 + ) -> torch.Tensor:
424 + """Forward pass through the decoder layer.
425 +
426 + Args:
427 + hidden_states: Input tensor from previous layer
428 + original_hidden_states: Original input tensor for residual
429 + connection
430 + block_idx: Current shared transformer block index
431 + positions: IDs for positional embeddings
432 +
433 + Returns:
434 + Transformed hidden states after attention and feed-forward
435 + """
436 +
437 + # The argument original_hidden_states is concatenated with hidden_states
438 + # (which is the output of the previous (mamba) layer).
439 + # The concatenated tensor is then used as input of the pre-attention
440 + # RMSNorm (see fig. 2 in https://arxiv.org/pdf/2405.16712).
441 + hidden_states = torch.concatenate(
442 + [hidden_states, original_hidden_states], dim=-1)
443 +
444 + # Layer norm before attention
445 + hidden_states = self.input_layernorm(hidden_states)
446 +
447 + # Self attention
448 + hidden_states = self.self_attn(
449 + hidden_states,
450 + position_ids=positions,
451 + block_idx=block_idx,
452 + )
453 +
454 + # Layer norm before feed-forward
455 + hidden_states = self.pre_ff_layernorm(hidden_states)
456 +
457 + # Feed-forward network
458 + hidden_states = self.feed_forward(hidden_states, block_idx=block_idx)
459 +
460 + return hidden_states
461 +
462 +
463 +class Zamba2MambaDecoderLayer(nn.Module):
464 + """Single Mamba decoder layer with normalization.
465 +
466 + This implements a Mamba block. It includes input normalization
467 + and can process sequences using either chunked or full
468 + computation depending on configuration.
469 + """
470 +
471 + def __init__(
472 + self,
473 + config: Zamba2Config,
474 + quant_config: Optional[QuantizationConfig] = None,
475 + ) -> None:
476 + """Initialize the Mamba decoder layer.
477 +
478 + Args:
479 + config: The Zamba2 model configuration
480 + quant_config: Configuration for model quantization
481 + """
482 + super().__init__()
483 +
484 + # Initialize Mamba mixer with expanded intermediate size
485 + intermediate_size = config.mamba_expand * config.hidden_size
486 + self.mamba = MambaMixer2(
487 + hidden_size=config.hidden_size,
488 + ssm_state_size=config.mamba_d_state,
489 + conv_kernel_size=config.mamba_d_conv,
490 + intermediate_size=intermediate_size,
491 + use_conv_bias=config.use_conv_bias,
492 + use_bias=config.add_bias_linear,
493 + n_groups=config.mamba_ngroups,
494 + num_heads=config.n_mamba_heads,
495 + head_dim=intermediate_size // config.n_mamba_heads,
496 + rms_norm_eps=config.rms_norm_eps,
497 + activation="silu",
498 + chunk_size=config.chunk_size,
499 + quant_config=quant_config,
500 + )
501 +
502 + # Input normalization
503 + self.input_layernorm = RMSNorm(config.hidden_size,
504 + eps=config.rms_norm_eps)
505 +
506 + def forward(
507 + self,
508 + hidden_states: torch.Tensor,
509 + mamba_cache_params: MambaCacheParams,
510 + sequence_idx: Optional[torch.Tensor] = None,
511 + transformer_hidden_states: Optional[torch.Tensor] = None,
512 + positions: Optional[torch.Tensor] = None,
513 + original_hidden_states: Optional[torch.Tensor] = None,
514 + ) -> torch.Tensor:
515 + """Forward pass through the Mamba decoder layer.
516 +
517 + Args:
518 + hidden_states: Input tensor [batch_size, seq_len, hidden_size]
519 + mamba_cache_params: Parameters for Mamba's state caches
520 + (one for conv, one for ssm)
521 + sequence_idx: Index tensor for identifying sequences in batch
522 + Required for proper chunked processing in prefill
523 + transformer_hidden_states: Optional output from transformer path
524 + Added to input if provided (used in hybrid architecture)
525 + positions: Optional position IDs (unused in Mamba)
526 + original_hidden_states: Optional original inputs (unused in Mamba)
527 +
528 + Returns:
529 + Transformed hidden states with residual connection applied
530 + """
531 + # Store input for residual connection
532 + residual = hidden_states
533 +
534 + # `transformer_hidden_states` is the output from shared
535 + # transformer + linear layer (see fig. 2 in
536 + # https://arxiv.org/pdf/2405.16712).
537 + # `transformer_hidden_states` is then added to the input to the mamba
538 + # layer below (as described in eq. (6) of
539 + # https://arxiv.org/pdf/2405.16712).
540 + if transformer_hidden_states is not None:
541 + hidden_states = hidden_states + transformer_hidden_states
542 +
543 + # Apply input normalization
544 + hidden_states = self.input_layernorm(hidden_states)
545 +
546 + # Process through Mamba mixer
547 + hidden_states = self.mamba(
548 + hidden_states,
549 + mamba_cache_params=mamba_cache_params,
550 + sequence_idx=sequence_idx,
551 + )
552 +
553 + # residual connection after mamba
554 + hidden_states = residual + hidden_states
555 +
556 + return hidden_states
557 +
558 +
559 +class Zamba2HybridLayer(nn.Module):
560 + """Hybrid layer combining Transformer and Mamba architectures.
561 +
562 + This layer implements the hybrid architecture described in the Zamba paper,
563 + where a shared transformer pathway processes input in parallel with a Mamba
564 + pathway. The transformer output is projected and added to the Mamba input
565 + for enhanced representation learning.
566 + """
567 +
568 + def __init__(
569 + self,
570 + shared_transformer: Zamba2AttentionDecoderLayer,
571 + config: Zamba2Config,
572 + block_idx: int,
573 + quant_config: Optional[QuantizationConfig] = None,
574 + ) -> None:
575 + """Initialize the hybrid layer.
576 +
577 + Args:
578 + shared_transformer: Transformer decoder layer for attention pathway
579 + linear: Linear projection for transformer output before Mamba
580 + mamba: Mamba decoder layer for state space pathway
581 + """
582 + super().__init__()
583 + self.block_idx = block_idx
584 + self.shared_transformer = shared_transformer
585 + self.linear = ReplicatedLinear(config.hidden_size,
586 + config.hidden_size,
587 + bias=False,
588 + quant_config=quant_config)
589 + self.mamba_decoder = Zamba2MambaDecoderLayer(config,
590 + quant_config=quant_config)
591 +
592 + def forward(
593 + self,
594 + hidden_states: torch.Tensor,
595 + original_hidden_states: torch.Tensor,
596 + positions: torch.Tensor,
597 + mamba_cache_params: Optional[MambaCacheParams] = None,
598 + sequence_idx: Optional[torch.Tensor] = None,
599 + ) -> torch.Tensor:
600 + """Forward pass through the hybrid layer.
601 +
602 + Processes input through parallel transformer and Mamba paths:
603 + 1. Transformer path processes input with attention
604 + 2. Transformer output is projected to match hidden size
605 + 3. Projected output is added to Mamba path input
606 + 4. Final output combines both paths' representations
607 +
608 + Args:
609 + hidden_states: Input tensor [batch_size, seq_len, hidden_size]
610 + original_hidden_states: Original input for transformer residual
611 + connection
612 + positions: Position IDs for positional embeddings
613 + mamba_cache_params: Parameters for Mamba's state caches
614 + (one for conv, one for ssm)
615 + sequence_idx: Indices for identifying sequences in batch,
616 + required for proper chunked processing in prefill
617 +
618 + Returns:
619 + Output tensor combining transformer and Mamba representations
620 + """
621 + # Process through transformer pathway
622 + transformer_hidden_states = self.shared_transformer(
623 + hidden_states,
624 + original_hidden_states=original_hidden_states,
625 + block_idx=self.block_idx,
626 + positions=positions,
627 + )
628 +
629 + # Project transformer output
630 + transformer_hidden_states, _ = self.linear(transformer_hidden_states)
631 +
632 + # Process through Mamba pathway with transformer injection
633 + layer_outputs = self.mamba_decoder(
634 + hidden_states,
635 + transformer_hidden_states=transformer_hidden_states,
636 + mamba_cache_params=mamba_cache_params,
637 + sequence_idx=sequence_idx,
638 + )
639 +
640 + return layer_outputs
641 +
642 +
643 +class Zamba2Model(nn.Module):
644 + """Core Zamba2 model combining transformer and Mamba architectures.
645 +
646 + The model processes input through a sequence of hybrid and Mamba-only
647 + layers, using token embeddings and final layer normalization.
648 + """
649 +
650 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
651 + """Initialize the Zamba2 model.
652 +
653 + Args:
654 + vllm_config: Configuration object containing model, cache,
655 + quantization and LoRA settings
656 + prefix: Optional prefix for parameter names in state dict
657 + """
658 + super().__init__()
659 +
660 + config = vllm_config.model_config.hf_config
661 + cache_config = vllm_config.cache_config
662 + quant_config = vllm_config.quant_config
663 + lora_config = vllm_config.lora_config
664 + is_lora_enabled = bool(lora_config)
665 + assert not is_lora_enabled
666 +
667 + self.config = config
668 + lora_vocab = ((lora_config.lora_extra_vocab_size *
669 + (lora_config.max_loras or 1)) if lora_config else 0)
670 + self.vocab_size = config.vocab_size + lora_vocab
671 + self.org_vocab_size = config.vocab_size
672 +
673 + # Initialize token embeddings
674 + self.embed_tokens = VocabParallelEmbedding(
675 + self.vocab_size,
676 + config.hidden_size,
677 + org_num_embeddings=config.vocab_size,
678 + )
679 +
680 + # Map hybrid layer indices to block indices
681 + layer2block_map = {
682 + layer_idx: block_idx
683 + for block_idx, layer_idx in enumerate(config.hybrid_layer_ids)
684 + }
685 +
686 + # Create cyclic iterator of transformer blocks
687 + blocks = cycle([
688 + Zamba2AttentionDecoderLayer(config,
689 + bare_block_idx=idx,
690 + num_hybrid_layers=len(layer2block_map),
691 + cache_config=cache_config,
692 + quant_config=quant_config,
693 + prefix=f"{prefix}")
694 + for idx in range(config.num_mem_blocks)
695 + ])
696 +
697 + # Initialize layers according to block type configuration
698 + layers = []
699 + for layer_idx, layer_type in enumerate(config.layers_block_type):
700 + if layer_type == "hybrid":
701 + block = next(blocks)
702 + block_idx = layer2block_map[layer_idx]
703 + layers.append(
704 + Zamba2HybridLayer(block, config, block_idx, quant_config))
705 + else:
706 + layers.append(
707 + Zamba2MambaDecoderLayer(config, quant_config=quant_config))
708 + self.layers = nn.ModuleList(layers)
709 +
710 + # Final layer normalization
711 + self.final_layernorm = RMSNorm(config.hidden_size,
712 + eps=config.rms_norm_eps)
713 +
714 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
715 + """Convert input token IDs to embeddings.
716 +
717 + Args:
718 + input_ids: Tensor of input token IDs
719 +
720 + Returns:
721 + Embedded representation of the input tokens
722 + """
723 + return self.embed_tokens(input_ids)
724 +
725 + def forward(
726 + self,
727 + input_ids: torch.Tensor,
728 + positions: torch.Tensor,
729 + mamba_cache_params: MambaCacheParams,
730 + inputs_embeds: Optional[torch.Tensor] = None,
731 + ) -> Union[torch.Tensor, IntermediateTensors]:
732 + """Forward pass through the model.
733 +
734 + Args:
735 + input_ids: Input token IDs
736 + positions: Position IDs for embeddings
737 + mamba_cache_params: Parameters for Mamba's state caches
738 + (one for conv, one for ssm)
739 + inputs_embeds: Optional pre-computed input embeddings
740 +
741 + Returns:
742 + Either final hidden states or intermediate tensors for pipeline
743 + parallelism
744 + """
745 + # Handle pipeline parallelism for first rank
746 + if inputs_embeds is None:
747 + inputs_embeds = self.get_input_embeddings(input_ids)
748 + hidden_states = inputs_embeds
749 +
750 + # pass a sequence index tensor, that is required for
751 + # proper continuous batching computation including
752 + # chunked prefill
753 + seq_idx = None
754 + attn_metadata = get_forward_context().attn_metadata
755 + if attn_metadata.num_prefills > 0:
756 + seq_idx = torch.zeros_like(input_ids, dtype=torch.int32)
757 + for i, (srt, end) in enumerate(
758 + zip(
759 + attn_metadata.query_start_loc,
760 + attn_metadata.query_start_loc[1:],
761 + )):
762 + seq_idx[srt:end] = i
763 + seq_idx.unsqueeze_(0)
764 +
765 + # Process through layers
766 + original_hidden_states = torch.clone(hidden_states)
767 + for layer_idx, layer in enumerate(self.layers):
768 + layer_outputs = layer(
769 + hidden_states,
770 + original_hidden_states=original_hidden_states,
771 + positions=positions,
772 + mamba_cache_params=mamba_cache_params.at_layer_idx(layer_idx),
773 + sequence_idx=seq_idx,
774 + )
775 + hidden_states = layer_outputs
776 +
777 + hidden_states = self.final_layernorm(hidden_states)
778 + return hidden_states
779 +
780 +
781 +class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsV0Only):
782 + """Zamba2 model with causal language modeling head.
783 +
784 + This class wraps the core Zamba2 model and adds:
785 + - A language modeling head for next token prediction
786 + - Mamba state caching functionality
787 + - Support for model parallelism and quantization
788 + - Sampling capabilities for text generation
789 + """
790 +
791 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
792 + """Initialize the Zamba2 model for causal language modeling.
793 +
794 + Args:
795 + vllm_config: Configuration containing model, cache, quantization,
796 + LoRA and scheduler settings
797 + prefix: Optional prefix for parameter names
798 +
799 + Raises:
800 + AssertionError: If prefix caching is enabled (not supported by
801 + Mamba)
802 + """
803 + config = vllm_config.model_config.hf_config
804 + cache_config = vllm_config.cache_config
805 + lora_config = vllm_config.lora_config
806 + scheduler_config = vllm_config.scheduler_config
807 + assert not cache_config.enable_prefix_caching, \
808 + "Mamba does not support prefix caching"
809 +
810 + super().__init__()
811 + self.config = config
812 + self.vllm_config = vllm_config
813 + self.scheduler_config = scheduler_config
814 + self.model_config = vllm_config.model_config
815 + self.unpadded_vocab_size = config.vocab_size
816 + if lora_config:
817 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
818 +
819 + # Initialize core model
820 + self.model = Zamba2Model(vllm_config=vllm_config,
821 + prefix=maybe_prefix(prefix, "model"))
822 +
823 + # Initialize language modeling head
824 + self.lm_head = ParallelLMHead(
825 + self.unpadded_vocab_size,
826 + config.hidden_size,
827 + org_num_embeddings=config.vocab_size,
828 + padding_size=DEFAULT_VOCAB_PADDING_SIZE
829 + # We need bigger padding if using lora for kernel
830 + # compatibility
831 + if not lora_config else lora_config.lora_vocab_padding_size,
832 + )
833 + # Tie weights with input embeddings if using same dimensions
834 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
835 +
836 + # Used to track and store by the Mamba cache between steps.
837 + self.mamba_cache: Optional[MambaCacheManager] = None
838 +
839 + # Initialize logits processing and sampling
840 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
841 + config.vocab_size)
842 + self.sampler = get_sampler()
843 +
844 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
845 + """Convert input token IDs to embeddings.
846 + Args:
847 + input_ids: Tensor of input token IDs
848 + Returns:
849 + Embedded representation of the input tokens
850 + """
851 + return self.model.get_input_embeddings(input_ids)
852 +
853 + def forward(self,
854 + input_ids: torch.Tensor,
855 + positions: torch.Tensor,
856 + inputs_embeds: Optional[torch.Tensor] = None,
857 + **kwargs) -> torch.Tensor:
858 + """Forward pass through the model.
859 +
860 + Args:
861 + input_ids: Input token IDs
862 + positions: Position IDs for embeddings
863 + inputs_embeds: Optional pre-computed input embeddings
864 + **kwargs: Additional arguments passed to cache manager
865 +
866 + Returns:
867 + Output hidden states
868 + """
869 + # Initialize Mamba cache if needed
870 + if self.mamba_cache is None:
871 + num_mamba_layers = self.config.num_hidden_layers
872 + self.mamba_cache = MambaCacheManager(
873 + self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,
874 + *self._get_mamba_cache_shape())
875 +
876 + # Get cache parameters for current run
877 + mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
878 +
879 + # Forward pass through model
880 + hidden_states = self.model(
881 + input_ids,
882 + positions,
883 + mamba_cache_params,
884 + inputs_embeds,
885 + )
886 +
887 + return hidden_states
888 +
889 + def copy_inputs_before_cuda_graphs(self, input_buffers: Dict[str,
890 + torch.Tensor],
891 + **kwargs) -> Dict[str, torch.Tensor]:
892 + """Copy inputs before CUDA graph capture.
893 +
894 + Args:
895 + input_buffers: Dictionary of input tensors
896 + **kwargs: Additional arguments passed to cache manager
897 +
898 + Returns:
899 + Updated input buffers
900 + """
901 + return self.mamba_cache.copy_inputs_before_cuda_graphs(
902 + input_buffers, **kwargs)
903 +
904 + def get_seqlen_agnostic_capture_inputs(
905 + self, batch_size: int) -> Dict[str, torch.Tensor]:
906 + """Get inputs for sequence-length-agnostic graph capture.
907 +
908 + Args:
909 + batch_size: Size of batch to capture
910 + Returns:
911 + Dictionary of capture inputs
912 + """
913 + return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
914 +
915 + def _get_mamba_cache_shape(
916 + self) -> Tuple[Tuple[int, int], Tuple[int, int]]:
917 + """Calculate shapes for Mamba's convolutional and state caches.
918 +
919 + Returns:
920 + Tuple containing:
921 + - conv_state_shape: Shape for convolutional state cache
922 + - temporal_state_shape: Shape for state space model cache
923 + """
924 + world_size = get_tensor_model_parallel_world_size()
925 +
926 + intermediate_size = self.config.mamba_expand * self.config.hidden_size
927 +
928 + # Extend groups if needed to ensure all groups needed by a head
929 + # are sharded together
930 +
931 + # if n_groups is not divisible by world_size, need to extend the shards
932 + # to ensure all groups needed by a head is sharded along with it
933 + n_groups = (self.config.mamba_ngroups + extra_groups_for_head_shards(
934 + self.config.mamba_ngroups, world_size))
935 +
936 + # Calculate conv state shape (includes groups)
937 + # - heads and n_groups are TP-ed
938 + conv_dim = (intermediate_size +
939 + 2 * n_groups * self.config.mamba_d_state)
940 + conv_state_shape = (
941 + divide(conv_dim, world_size),
942 + self.config.mamba_d_conv - 1,
943 + )
944 +
945 + # Calculate temporal state shape (per-head states)
946 + # These are not TP-ed as they depend on A, dt_bias, D
947 + # - they are typically small
948 + # e.g., (h_heads, d_head, d_state) = (128, 64, 128)
949 + temporal_state_shape = (
950 + divide(divide(intermediate_size, self.config.mamba_headdim),
951 + world_size),
952 + self.config.mamba_headdim,
953 + self.config.mamba_d_state,
954 + )
955 +
956 + return conv_state_shape, temporal_state_shape
957 +
958 + def compute_logits(
959 + self,
960 + hidden_states: torch.Tensor,
961 + sampling_metadata: SamplingMetadata,
962 + ) -> Optional[torch.Tensor]:
963 + """Compute logits for next token prediction.
964 +
965 + Args:
966 + hidden_states: Hidden states from model forward pass
967 + sampling_metadata: Metadata for sampling process
968 +
969 + Returns:
970 + Logits for next token prediction
971 + """
972 + logits = self.logits_processor(self.lm_head, hidden_states,
973 + sampling_metadata)
974 + return logits
975 +
976 + def sample(
977 + self,
978 + logits: Optional[torch.Tensor],
979 + sampling_metadata: SamplingMetadata,
980 + ) -> Optional[SamplerOutput]:
981 + """Sample next tokens from computed logits.
982 +
983 + Args:
984 + logits: Computed logits for next token prediction
985 + sampling_metadata: Metadata for sampling process
986 +
987 + Returns:
988 + Sampled tokens and related sampling information
989 + """
990 + next_tokens = self.sampler(logits, sampling_metadata)
991 + return next_tokens
992 +
993 + def load_weights(self, weights: Iterable[Tuple[str,
994 + torch.Tensor]]) -> Set[str]:
995 + stacked_params_mapping = [
996 + # (param_name, shard_name, shard_id)
997 + ("qkv_proj", "q_proj", "q"),
998 + ("qkv_proj", "k_proj", "k"),
999 + ("qkv_proj", "v_proj", "v"),
1000 + ]
1001 +
1002 + weights_dict = {}
1003 + for key, loaded_weight in weights:
1004 + if "A_log" in key:
1005 + key = key.replace("A_log", "A")
1006 + elif "adapter_list" in key:
1007 + key = key.replace("0.weight", "A.weight")
1008 + key = key.replace("1.weight", "B.weight")
1009 + weights_dict[key] = loaded_weight
1010 +
1011 + params_dict = dict(self.named_parameters())
1012 + loaded_params: Set[str] = set()
1013 + for chkpt_weight_name, loaded_weight in weights_dict.items():
1014 + for param_name, weight_name, shard_id in stacked_params_mapping:
1015 + if weight_name not in chkpt_weight_name:
1016 + continue
1017 + chkpt_weight_name = chkpt_weight_name.replace(
1018 + weight_name, param_name)
1019 + param = params_dict[chkpt_weight_name]
1020 + weight_loader = param.weight_loader
1021 + weight_loader(param, loaded_weight, shard_id)
1022 + break
1023 + else:
1024 + if chkpt_weight_name not in params_dict:
1025 + continue
1026 + param = params_dict[chkpt_weight_name]
1027 + weight_loader = getattr(param, "weight_loader",
1028 + default_weight_loader)
1029 + weight_loader(param, loaded_weight)
1030 + loaded_params.add(chkpt_weight_name)
1031 + return loaded_params