[Model] Remove BambaForCausalLM (#45990)
Signed-off-by: Xianbao QIAN <xianbao.qian@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Tiezhen WANG committed
Jun 18, 2026 at 21:51 UTC
d682968aa9fcd7e7a78218b548c52fc198a87a6c
5 files changed
+6
-529
docs/models/supported_models.md
-1
@@ -374,7 +374,6 @@ th {
374
| `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | ✅︎ | ✅︎ |
375
| `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ |
376
| `BailingMoeV2_5ForCausalLM` | Ling | `inclusionAI/Ling-2.5-1T`, `inclusionAI/Ring-2.5-1T` | | ✅︎ |
377
-| `BambaForCausalLM` | Bamba | `ibm-ai-platform/Bamba-9B-fp8`, `ibm-ai-platform/Bamba-9B` | ✅︎ | ✅︎ |
377
| `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ |
378
| `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
379
| `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ |
tests/models/language/generation/test_hybrid.py
+5
-6
@@ -37,7 +37,6 @@ HYBRID_MODELS = [
37
"ai21labs/Jamba-tiny-dev",
38
"pfnet/plamo-2-1b",
39
"Zyphra/Zamba2-1.2B-instruct",
40
- "hmellor/tiny-random-BambaForCausalLM",
40
"ibm-granite/granite-4.0-tiny-preview",
41
"tiiuae/Falcon-H1-0.5B-Base",
42
"LiquidAI/LFM2-1.2B",
@@ -439,7 +438,7 @@ def _get_vLLM_output(
438
return outs, vllm_model
439
440
442
-@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
441
+@pytest.mark.parametrize("model", [HYBRID_MODELS[0]])
442
@pytest.mark.parametrize("max_tokens", [64])
443
@pytest.mark.parametrize("n_repetitions", [2])
444
# If num_logprobs is set to -1, then the stringent version
@@ -503,7 +502,7 @@ def test_apc_single_prompt(
502
)
503
504
506
-@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
505
+@pytest.mark.parametrize("model", [HYBRID_MODELS[0]])
506
@pytest.mark.parametrize("max_tokens", [64])
507
@pytest.mark.parametrize("n_repetitions", [2])
508
# If num_logprobs is set to -1, then the stringent version
@@ -584,7 +583,7 @@ def test_apc_single_prompt_block_align_alignment(
583
)
584
585
587
-@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
586
+@pytest.mark.parametrize("model", [HYBRID_MODELS[0]])
587
@pytest.mark.parametrize("max_tokens", [64])
588
@pytest.mark.parametrize("n_repetitions", [2])
589
# If num_logprobs is set to -1, then the stringent version
@@ -653,7 +652,7 @@ def test_apc_multiple_prompts_all_cached_outputs(
652
)
653
654
656
-@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
655
+@pytest.mark.parametrize("model", [HYBRID_MODELS[0]])
656
@pytest.mark.parametrize("max_tokens", [64])
657
@pytest.mark.parametrize("n_repetitions", [2])
658
# If num_logprobs is set to -1, then the stringent version
@@ -738,7 +737,7 @@ def test_apc_multiple_prompts_block_align_alignment(
737
)
738
739
741
-@pytest.mark.parametrize("model", [HYBRID_MODELS[0], HYBRID_MODELS[3]])
740
+@pytest.mark.parametrize("model", [HYBRID_MODELS[0]])
741
@pytest.mark.parametrize("max_tokens", [64])
742
@pytest.mark.parametrize("n_repetitions", [2])
743
# If num_logprobs is set to -1, then the stringent version
tests/models/registry.py
-4
@@ -223,10 +223,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
223
"BailingMoeV2_5ForCausalLM": _HfExamplesInfo(
224
"inclusionAI/Ring-2.5-1T", trust_remote_code=True
225
),
226
- "BambaForCausalLM": _HfExamplesInfo(
227
- "ibm-ai-platform/Bamba-9B-v1",
228
- extras={"tiny": "hmellor/tiny-random-BambaForCausalLM"},
229
- ),
226
"BloomForCausalLM": _HfExamplesInfo(
227
"bigscience/bloom-560m", {"1b": "bigscience/bloomz-1b1"}
228
),
vllm/model_executor/models/bamba.py
deleted
-517
@@ -1,517 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Inference-only Bamba model."""
4
-
5
-# Added by the IBM Team, 2024
6
-from collections.abc import Iterable
7
-
8
-import torch
9
-from torch import nn
10
-from transformers import BambaConfig
11
-
12
-from vllm.compilation.decorators import support_torch_compile
13
-from vllm.config import CacheConfig, ModelConfig, VllmConfig
14
-from vllm.distributed import get_tensor_model_parallel_world_size
15
-from vllm.distributed.parallel_state import get_pp_group
16
-from vllm.model_executor.layers.activation import SiluAndMul
17
-from vllm.model_executor.layers.attention import Attention
18
-from vllm.model_executor.layers.layernorm import RMSNorm
19
-from vllm.model_executor.layers.linear import (
20
- MergedColumnParallelLinear,
21
- QKVParallelLinear,
22
- RowParallelLinear,
23
-)
24
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
25
-from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
26
-from vllm.model_executor.layers.mamba.mamba_utils import (
27
- MambaStateCopyFunc,
28
- MambaStateCopyFuncCalculator,
29
- MambaStateDtypeCalculator,
30
- MambaStateShapeCalculator,
31
-)
32
-from vllm.model_executor.layers.quantization import QuantizationConfig
33
-from vllm.model_executor.layers.rotary_embedding import get_rope
34
-from vllm.model_executor.layers.vocab_parallel_embedding import (
35
- ParallelLMHead,
36
- VocabParallelEmbedding,
37
-)
38
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
39
-from vllm.sequence import IntermediateTensors
40
-
41
-from .interfaces import (
42
- HasInnerState,
43
- IsHybrid,
44
- SupportsLoRA,
45
- SupportsMambaPrefixCaching,
46
- SupportsPP,
47
- SupportsQuant,
48
-)
49
-from .utils import (
50
- AutoWeightsLoader,
51
- is_pp_missing_parameter,
52
- make_empty_intermediate_tensors_factory,
53
- make_layers,
54
- maybe_prefix,
55
-)
56
-
57
-
58
-class BambaMLP(nn.Module):
59
- def __init__(
60
- self,
61
- config: BambaConfig,
62
- quant_config: QuantizationConfig | None = None,
63
- bias: bool = False,
64
- prefix: str = "",
65
- ) -> None:
66
- super().__init__()
67
- self.gate_up_proj = MergedColumnParallelLinear(
68
- input_size=config.hidden_size,
69
- output_sizes=[config.intermediate_size] * 2,
70
- bias=bias,
71
- quant_config=quant_config,
72
- prefix=f"{prefix}.gate_up_proj",
73
- )
74
- self.down_proj = RowParallelLinear(
75
- input_size=config.intermediate_size,
76
- output_size=config.hidden_size,
77
- bias=bias,
78
- quant_config=quant_config,
79
- prefix=f"{prefix}.down_proj",
80
- )
81
- if config.hidden_act != "silu":
82
- raise ValueError(
83
- f"Unsupported activation: {config.hidden_act}. "
84
- "Only silu is supported for now."
85
- )
86
- self.act_fn = SiluAndMul()
87
-
88
- def forward(self, x):
89
- x, _ = self.gate_up_proj(x)
90
- x = self.act_fn(x)
91
- x, _ = self.down_proj(x)
92
- return x
93
-
94
-
95
-class BambaMixerDecoderLayer(nn.Module):
96
- def __init__(
97
- self,
98
- config: BambaConfig,
99
- layer_idx: int,
100
- model_config: ModelConfig | None = None,
101
- cache_config: CacheConfig | None = None,
102
- quant_config: QuantizationConfig | None = None,
103
- prefix: str = "",
104
- ) -> None:
105
- super().__init__()
106
- self.config = config
107
- self.mamba = MambaMixer2(
108
- hidden_size=config.hidden_size,
109
- ssm_state_size=config.mamba_d_state,
110
- conv_kernel_size=config.mamba_d_conv,
111
- intermediate_size=config.mamba_expand * config.hidden_size,
112
- use_conv_bias=config.mamba_conv_bias,
113
- use_bias=config.mamba_proj_bias,
114
- n_groups=config.mamba_n_groups,
115
- num_heads=config.mamba_n_heads,
116
- head_dim=config.mamba_d_head,
117
- rms_norm_eps=config.rms_norm_eps,
118
- activation=config.hidden_act,
119
- model_config=model_config,
120
- cache_config=cache_config,
121
- quant_config=quant_config,
122
- prefix=f"{prefix}.mixer",
123
- )
124
-
125
- self.feed_forward = BambaMLP(
126
- config, quant_config=quant_config, prefix=f"{prefix}.feed_forward"
127
- )
128
- self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
129
- self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
130
-
131
- def forward(
132
- self,
133
- hidden_states: torch.Tensor,
134
- residual: torch.Tensor | None,
135
- **kwargs,
136
- ):
137
- if residual is None:
138
- residual = hidden_states
139
- hidden_states = self.input_layernorm(hidden_states)
140
- else:
141
- hidden_states, residual = self.input_layernorm(hidden_states, residual)
142
-
143
- output = self.mamba(hidden_states)
144
- # Fully Connected
145
- hidden_states, residual = self.pre_ff_layernorm(output, residual)
146
- hidden_states = self.feed_forward(hidden_states)
147
- return hidden_states, residual
148
-
149
-
150
-class BambaAttentionDecoderLayer(nn.Module):
151
- def __init__(
152
- self,
153
- config: BambaConfig,
154
- layer_idx: int,
155
- model_config: ModelConfig | None = None,
156
- cache_config: CacheConfig | None = None,
157
- quant_config: QuantizationConfig | None = None,
158
- prefix: str = "",
159
- ) -> None:
160
- super().__init__()
161
- max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
162
- self.hidden_size = config.hidden_size
163
- tp_size = get_tensor_model_parallel_world_size()
164
- self.total_num_heads = config.num_attention_heads
165
- assert self.total_num_heads % tp_size == 0
166
- self.num_heads = self.total_num_heads // tp_size
167
- self.total_num_kv_heads = config.num_key_value_heads
168
- if self.total_num_kv_heads >= tp_size:
169
- # Number of KV heads is greater than TP size, so we partition
170
- # the KV heads across multiple tensor parallel GPUs.
171
- assert self.total_num_kv_heads % tp_size == 0
172
- else:
173
- # Number of KV heads is less than TP size, so we replicate
174
- # the KV heads across multiple tensor parallel GPUs.
175
- assert tp_size % self.total_num_kv_heads == 0
176
- self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
177
- self.head_dim = config.hidden_size // self.total_num_heads
178
- self.q_size = self.num_heads * self.head_dim
179
- self.kv_size = self.num_kv_heads * self.head_dim
180
- self.scaling = self.head_dim**-0.5
181
- self.max_position_embeddings = max_position_embeddings
182
-
183
- rotary_dim = getattr(config, "attn_rotary_emb", self.head_dim)
184
- config.rope_parameters["partial_rotary_factor"] = rotary_dim / self.head_dim
185
-
186
- self.rotary_emb = get_rope(
187
- head_size=self.head_dim,
188
- max_position=max_position_embeddings,
189
- rope_parameters=config.rope_parameters,
190
- is_neox_style=True,
191
- dtype=torch.get_default_dtype(), # see impl of get_rope
192
- )
193
-
194
- self.qkv_proj = QKVParallelLinear(
195
- config.hidden_size,
196
- self.head_dim,
197
- self.total_num_heads,
198
- self.total_num_kv_heads,
199
- bias=False,
200
- quant_config=quant_config,
201
- prefix=f"{prefix}.qkv_proj",
202
- )
203
- self.o_proj = RowParallelLinear(
204
- self.total_num_heads * self.head_dim,
205
- config.hidden_size,
206
- bias=False,
207
- quant_config=quant_config,
208
- prefix=f"{prefix}.o_proj",
209
- )
210
-
211
- self.attn = Attention(
212
- self.num_heads,
213
- self.head_dim,
214
- self.scaling,
215
- num_kv_heads=self.num_kv_heads,
216
- cache_config=cache_config,
217
- prefix=f"{prefix}.attn",
218
- )
219
-
220
- self.feed_forward = BambaMLP(
221
- config, quant_config=quant_config, prefix=f"{prefix}.feed_forward"
222
- )
223
- self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
224
- self.pre_ff_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
225
-
226
- def self_attention(
227
- self,
228
- positions: torch.Tensor,
229
- hidden_states: torch.Tensor,
230
- **kwargs,
231
- ) -> torch.Tensor:
232
- qkv, _ = self.qkv_proj(hidden_states)
233
- q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
234
-
235
- q, k = self.rotary_emb(positions, q, k)
236
- attn_output = self.attn(q, k, v)
237
- output, _ = self.o_proj(attn_output)
238
- return output
239
-
240
- def forward(
241
- self,
242
- positions: torch.Tensor,
243
- hidden_states: torch.Tensor,
244
- residual: torch.Tensor | None,
245
- **kwargs,
246
- ):
247
- if residual is None:
248
- residual = hidden_states
249
- hidden_states = self.input_layernorm(hidden_states)
250
- else:
251
- hidden_states, residual = self.input_layernorm(hidden_states, residual)
252
-
253
- hidden_states = self.self_attention(
254
- positions=positions,
255
- hidden_states=hidden_states,
256
- )
257
- # Fully Connected
258
- hidden_states, residual = self.pre_ff_layernorm(hidden_states, residual)
259
- hidden_states = self.feed_forward(hidden_states)
260
- return hidden_states, residual
261
-
262
-
263
-ALL_DECODER_LAYER_TYPES = {
264
- "attention": BambaAttentionDecoderLayer,
265
- "mamba": BambaMixerDecoderLayer,
266
-}
267
-
268
-
269
-@support_torch_compile
270
-class BambaModel(nn.Module):
271
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
272
- super().__init__()
273
-
274
- config: BambaConfig = vllm_config.model_config.hf_config
275
- model_config = vllm_config.model_config
276
- cache_config = vllm_config.cache_config
277
- quant_config = vllm_config.quant_config
278
-
279
- self.config = config
280
-
281
- self.vocab_size = config.vocab_size
282
-
283
- self.embed_tokens = VocabParallelEmbedding(
284
- self.vocab_size,
285
- config.hidden_size,
286
- )
287
-
288
- def get_layer(prefix: str):
289
- layer_idx = int(prefix.rsplit(".", 1)[1])
290
- layer_class = ALL_DECODER_LAYER_TYPES[config.layers_block_type[layer_idx]]
291
- return layer_class(
292
- config,
293
- layer_idx,
294
- model_config,
295
- cache_config,
296
- quant_config=quant_config,
297
- prefix=prefix,
298
- )
299
-
300
- self.start_layer, self.end_layer, self.layers = make_layers(
301
- config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
302
- )
303
- self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
304
- ["hidden_states", "residual"], config.hidden_size
305
- )
306
-
307
- self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
308
-
309
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
310
- return self.embed_tokens(input_ids)
311
-
312
- def forward(
313
- self,
314
- input_ids: torch.Tensor | None,
315
- positions: torch.Tensor,
316
- intermediate_tensors: IntermediateTensors | None = None,
317
- inputs_embeds: torch.Tensor | None = None,
318
- ) -> torch.Tensor:
319
- if get_pp_group().is_first_rank:
320
- if inputs_embeds is not None:
321
- hidden_states = inputs_embeds
322
- else:
323
- hidden_states = self.embed_input_ids(input_ids)
324
- residual = None
325
- else:
326
- assert intermediate_tensors is not None
327
- hidden_states = intermediate_tensors["hidden_states"]
328
- residual = intermediate_tensors["residual"]
329
-
330
- residual = None
331
- for i, layer in enumerate(self.layers):
332
- hidden_states, residual = layer(
333
- positions=positions,
334
- hidden_states=hidden_states,
335
- residual=residual,
336
- )
337
-
338
- if not get_pp_group().is_last_rank:
339
- return IntermediateTensors(
340
- {"hidden_states": hidden_states, "residual": residual}
341
- )
342
- hidden_states, _ = self.final_layernorm(hidden_states, residual)
343
- return hidden_states
344
-
345
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
346
- stacked_params_mapping = [
347
- # (param_name, shard_name, shard_id)
348
- ("qkv_proj", "q_proj", "q"),
349
- ("qkv_proj", "k_proj", "k"),
350
- ("qkv_proj", "v_proj", "v"),
351
- ("gate_up_proj", "gate_proj", 0),
352
- ("gate_up_proj", "up_proj", 1),
353
- ]
354
-
355
- params_dict = dict(self.named_parameters())
356
- loaded_params: set[str] = set()
357
- for name, loaded_weight in weights:
358
- if "rotary_emb.inv_freq" in name:
359
- continue
360
-
361
- if "A_log" in name:
362
- name = name.replace("A_log", "A")
363
-
364
- if ".self_attn." in name:
365
- name = name.replace(".self_attn", "")
366
-
367
- for param_name, weight_name, shard_id in stacked_params_mapping:
368
- if weight_name not in name:
369
- continue
370
-
371
- name = name.replace(weight_name, param_name)
372
- # Skip loading extra bias for GPTQ models.
373
- if name.endswith(".bias") and name not in params_dict:
374
- continue
375
- # Skip layers on other devices.
376
- if is_pp_missing_parameter(name, self):
377
- continue
378
- param = params_dict[name]
379
- weight_loader = param.weight_loader
380
- weight_loader(param, loaded_weight, shard_id)
381
- break
382
- else:
383
- # Skip loading extra bias for GPTQ models.
384
- if name.endswith(".bias") and name not in params_dict:
385
- continue
386
- if is_pp_missing_parameter(name, self):
387
- continue
388
-
389
- param = params_dict[name]
390
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
391
- weight_loader(param, loaded_weight)
392
- loaded_params.add(name)
393
- return loaded_params
394
-
395
-
396
-class BambaForCausalLM(
397
- nn.Module,
398
- HasInnerState,
399
- SupportsLoRA,
400
- SupportsPP,
401
- IsHybrid,
402
- SupportsQuant,
403
- SupportsMambaPrefixCaching,
404
-):
405
- packed_modules_mapping = {
406
- "qkv_proj": [
407
- "q_proj",
408
- "k_proj",
409
- "v_proj",
410
- ],
411
- "gate_up_proj": ["up_proj", "down_proj"],
412
- }
413
-
414
- # LoRA specific attributes
415
- embedding_modules = {
416
- "embed_tokens": "input_embeddings",
417
- "lm_head": "output_embeddings",
418
- }
419
-
420
- @classmethod
421
- def get_mamba_state_dtype_from_config(
422
- cls,
423
- vllm_config: "VllmConfig",
424
- ) -> tuple[torch.dtype, torch.dtype]:
425
- return MambaStateDtypeCalculator.mamba2_state_dtype(
426
- vllm_config.model_config.dtype,
427
- vllm_config.cache_config.mamba_cache_dtype,
428
- vllm_config.cache_config.mamba_ssm_cache_dtype,
429
- )
430
-
431
- @classmethod
432
- def get_mamba_state_shape_from_config(
433
- cls,
434
- vllm_config: "VllmConfig",
435
- ) -> tuple[tuple[int, int], tuple[int, int, int]]:
436
- """Calculate shapes for Mamba's convolutional and state caches.
437
-
438
- Args:
439
- vllm_config: vLLM config
440
-
441
- Returns:
442
- Tuple containing:
443
- - conv_state_shape: Shape for convolutional state cache
444
- - temporal_state_shape: Shape for state space model cache
445
- """
446
- parallel_config = vllm_config.parallel_config
447
- hf_config = vllm_config.model_config.hf_config
448
- intermediate_size = hf_config.mamba_expand * hf_config.hidden_size
449
-
450
- return MambaStateShapeCalculator.mamba2_state_shape(
451
- intermediate_size=intermediate_size,
452
- tp_world_size=parallel_config.tensor_parallel_size,
453
- n_groups=hf_config.mamba_n_groups,
454
- num_heads=hf_config.mamba_n_heads,
455
- head_dim=hf_config.mamba_d_head,
456
- state_size=hf_config.mamba_d_state,
457
- conv_kernel=hf_config.mamba_d_conv,
458
- )
459
-
460
- @classmethod
461
- def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
462
- return MambaStateCopyFuncCalculator.mamba2_state_copy_func()
463
-
464
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
465
- config = vllm_config.model_config.hf_config
466
- self.vllm_config = vllm_config
467
- self.model_config = vllm_config.model_config
468
-
469
- scheduler_config = vllm_config.scheduler_config
470
- self.quant_config = vllm_config.quant_config
471
-
472
- super().__init__()
473
- self.config = config
474
- self.scheduler_config = scheduler_config
475
- self.model = BambaModel(
476
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
477
- )
478
-
479
- self.lm_head = ParallelLMHead(
480
- config.vocab_size,
481
- config.hidden_size,
482
- prefix=maybe_prefix(prefix, "lm_head"),
483
- )
484
-
485
- self.logits_processor = LogitsProcessor(config.vocab_size)
486
-
487
- self.make_empty_intermediate_tensors = (
488
- self.model.make_empty_intermediate_tensors
489
- )
490
-
491
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
492
- return self.model.embed_input_ids(input_ids)
493
-
494
- def forward(
495
- self,
496
- input_ids: torch.Tensor | None,
497
- positions: torch.Tensor,
498
- intermediate_tensors: IntermediateTensors | None = None,
499
- inputs_embeds: torch.Tensor | None = None,
500
- **kwargs,
501
- ):
502
- hidden_states = self.model(
503
- input_ids, positions, intermediate_tensors, inputs_embeds
504
- )
505
-
506
- return hidden_states
507
-
508
- def compute_logits(
509
- self,
510
- hidden_states: torch.Tensor,
511
- ) -> torch.Tensor | None:
512
- logits = self.logits_processor(self.lm_head, hidden_states)
513
- return logits
514
-
515
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
516
- loader = AutoWeightsLoader(self)
517
- return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+1
-1
@@ -84,7 +84,6 @@ _TEXT_GENERATION_MODELS = {
84
"BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"),
85
"BailingMoeV2ForCausalLM": ("bailing_moe", "BailingMoeV2ForCausalLM"),
86
"BailingMoeV2_5ForCausalLM": ("bailing_moe_linear", "BailingMoeV25ForCausalLM"),
87
- "BambaForCausalLM": ("bamba", "BambaForCausalLM"),
87
"BloomForCausalLM": ("bloom", "BloomForCausalLM"),
88
"ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"),
89
"ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"),
@@ -733,6 +732,7 @@ _PREVIOUSLY_SUPPORTED_MODELS = {
732
"MllamaForConditionalGeneration": "0.10.2",
733
"XverseForCausalLM": "0.23.0",
734
"Dots1ForCausalLM": "0.23.0",
735
+ "BambaForCausalLM": "0.23.0",
736
}
737
738
_OOT_SUPPORTED_MODELS = {