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 11, 05:00 UTC
adds jais 2 support (#30188)
Signed-off-by: sarathc-cerebras <sarath.chandran@cerebras.net> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
sarathc-cerebras committed
Dec 18, 2025 at 21:16 UTC
28d15ab56bd9d3fd17010bc4abaeec06988f7887
4 files changed
+534
docs/models/supported_models.md
+1
@@ -406,6 +406,7 @@ th {
406
| `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
407
| `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |
408
| `JAISLMHeadModel` | Jais | `inceptionai/jais-13b`, `inceptionai/jais-13b-chat`, `inceptionai/jais-30b-v3`, `inceptionai/jais-30b-chat-v3`, etc. | | ✅︎ |
409
+| `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | | ✅︎ |
410
| `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ |
411
| `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B-Base, Kimi-Linear-48B-A3B-Instruct | `moonshotai/Kimi-Linear-48B-A3B-Base`, `moonshotai/Kimi-Linear-48B-A3B-Instruct` | | ✅︎ |
412
| `Lfm2ForCausalLM` | LFM2 | `LiquidAI/LFM2-1.2B`, `LiquidAI/LFM2-700M`, `LiquidAI/LFM2-350M`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+3
@@ -295,6 +295,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
295
"internlm/internlm3-8b-instruct", trust_remote_code=True
296
),
297
"JAISLMHeadModel": _HfExamplesInfo("inceptionai/jais-13b-chat"),
298
+ "Jais2ForCausalLM": _HfExamplesInfo(
299
+ "inceptionai/Jais-2-8B-Chat", min_transformers_version="4.58"
300
+ ),
301
"JambaForCausalLM": _HfExamplesInfo(
302
"ai21labs/AI21-Jamba-1.5-Mini",
303
extras={
vllm/model_executor/models/jais2.py
new
+529
@@ -0,0 +1,529 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# Adapted from
5
+# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
6
+# Copyright 2023 The vLLM team.
7
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
8
+#
9
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
10
+# and OPT implementations in this library. It has been modified from its
11
+# original forms to accommodate minor architectural differences compared
12
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13
+#
14
+# Licensed under the Apache License, Version 2.0 (the "License");
15
+# you may not use this file except in compliance with the License.
16
+# You may obtain a copy of the License at
17
+#
18
+# http://www.apache.org/licenses/LICENSE-2.0
19
+#
20
+# Unless required by applicable law or agreed to in writing, software
21
+# distributed under the License is distributed on an "AS IS" BASIS,
22
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23
+# See the License for the specific language governing permissions and
24
+# limitations under the License.
25
+
26
+"""Inference-only Jais2 model compatible with HuggingFace weights."""
27
+
28
+from collections.abc import Iterable
29
+
30
+import torch
31
+from torch import nn
32
+from transformers import Jais2Config
33
+
34
+from vllm.attention.layer import Attention
35
+from vllm.compilation.decorators import support_torch_compile
36
+from vllm.config import CacheConfig, VllmConfig
37
+from vllm.distributed import (
38
+ get_pp_group,
39
+ get_tensor_model_parallel_world_size,
40
+)
41
+from vllm.model_executor.layers.activation import ReLUSquaredActivation
42
+from vllm.model_executor.layers.linear import (
43
+ ColumnParallelLinear,
44
+ QKVParallelLinear,
45
+ RowParallelLinear,
46
+)
47
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
48
+from vllm.model_executor.layers.quantization import QuantizationConfig
49
+from vllm.model_executor.layers.rotary_embedding import get_rope
50
+from vllm.model_executor.layers.vocab_parallel_embedding import (
51
+ DEFAULT_VOCAB_PADDING_SIZE,
52
+ ParallelLMHead,
53
+ VocabParallelEmbedding,
54
+)
55
+from vllm.model_executor.model_loader.weight_utils import (
56
+ default_weight_loader,
57
+ maybe_remap_kv_scale_name,
58
+)
59
+from vllm.sequence import IntermediateTensors
60
+
61
+from .interfaces import SupportsLoRA, SupportsPP
62
+from .utils import (
63
+ AutoWeightsLoader,
64
+ PPMissingLayer,
65
+ extract_layer_index,
66
+ is_pp_missing_parameter,
67
+ make_empty_intermediate_tensors_factory,
68
+ make_layers,
69
+ maybe_prefix,
70
+)
71
+
72
+
73
+class Jais2MLP(nn.Module):
74
+ def __init__(
75
+ self,
76
+ hidden_size: int,
77
+ intermediate_size: int,
78
+ hidden_act: str,
79
+ quant_config: QuantizationConfig | None = None,
80
+ bias: bool = False,
81
+ prefix: str = "",
82
+ ) -> None:
83
+ super().__init__()
84
+ self.up_proj = ColumnParallelLinear(
85
+ input_size=hidden_size,
86
+ output_size=intermediate_size,
87
+ bias=bias,
88
+ quant_config=quant_config,
89
+ prefix=f"{prefix}.up_proj",
90
+ )
91
+ self.down_proj = RowParallelLinear(
92
+ input_size=intermediate_size,
93
+ output_size=hidden_size,
94
+ bias=bias,
95
+ quant_config=quant_config,
96
+ prefix=f"{prefix}.down_proj",
97
+ )
98
+ self.act_fn = ReLUSquaredActivation()
99
+
100
+ def forward(self, x):
101
+ x, _ = self.up_proj(x)
102
+ x = self.act_fn(x)
103
+ x, _ = self.down_proj(x)
104
+ return x
105
+
106
+
107
+class Jais2Attention(nn.Module):
108
+ def __init__(
109
+ self,
110
+ config: Jais2Config,
111
+ hidden_size: int,
112
+ num_heads: int,
113
+ num_kv_heads: int,
114
+ max_position_embeddings: int = 8192,
115
+ quant_config: QuantizationConfig | None = None,
116
+ bias: bool = False,
117
+ cache_config: CacheConfig | None = None,
118
+ prefix: str = "",
119
+ ) -> None:
120
+ super().__init__()
121
+ layer_idx = extract_layer_index(prefix)
122
+ self.hidden_size = hidden_size
123
+ tp_size = get_tensor_model_parallel_world_size()
124
+ self.total_num_heads = num_heads
125
+ assert self.total_num_heads % tp_size == 0
126
+ self.num_heads = self.total_num_heads // tp_size
127
+ self.total_num_kv_heads = num_kv_heads
128
+ if self.total_num_kv_heads >= tp_size:
129
+ # Number of KV heads is greater than TP size, so we partition
130
+ # the KV heads across multiple tensor parallel GPUs.
131
+ assert self.total_num_kv_heads % tp_size == 0
132
+ else:
133
+ # Number of KV heads is less than TP size, so we replicate
134
+ # the KV heads across multiple tensor parallel GPUs.
135
+ assert tp_size % self.total_num_kv_heads == 0
136
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
137
+ # MistralConfig has an optional head_dim introduced by Mistral-Nemo
138
+ self.head_dim = getattr(
139
+ config, "head_dim", self.hidden_size // self.total_num_heads
140
+ )
141
+ self.q_size = self.num_heads * self.head_dim
142
+ self.kv_size = self.num_kv_heads * self.head_dim
143
+ self.scaling = self.head_dim**-0.5
144
+ self.max_position_embeddings = max_position_embeddings
145
+
146
+ self.qkv_proj = QKVParallelLinear(
147
+ hidden_size=hidden_size,
148
+ head_size=self.head_dim,
149
+ total_num_heads=self.total_num_heads,
150
+ total_num_kv_heads=self.total_num_kv_heads,
151
+ bias=bias,
152
+ quant_config=quant_config,
153
+ prefix=f"{prefix}.qkv_proj",
154
+ )
155
+
156
+ self.o_proj = RowParallelLinear(
157
+ input_size=self.total_num_heads * self.head_dim,
158
+ output_size=hidden_size,
159
+ bias=bias,
160
+ quant_config=quant_config,
161
+ prefix=f"{prefix}.o_proj",
162
+ )
163
+
164
+ is_neox_style = True
165
+ if quant_config is not None and quant_config.get_name() == "gguf":
166
+ is_neox_style = False
167
+
168
+ self.rotary_emb = get_rope(
169
+ self.head_dim,
170
+ rotary_dim=self.head_dim,
171
+ max_position=max_position_embeddings,
172
+ rope_parameters=getattr(config, "rope_parameters", None),
173
+ is_neox_style=is_neox_style,
174
+ )
175
+
176
+ if hasattr(config, "interleaved_sliding_window"):
177
+ interleaved_sliding_window = config.interleaved_sliding_window
178
+ if isinstance(interleaved_sliding_window, int):
179
+ sliding_window = interleaved_sliding_window
180
+ elif isinstance(interleaved_sliding_window, list):
181
+ sw_idx = layer_idx % len(interleaved_sliding_window)
182
+ sliding_window = interleaved_sliding_window[sw_idx]
183
+ else:
184
+ raise ValueError(
185
+ f"{type(interleaved_sliding_window)} is not supported."
186
+ )
187
+ else:
188
+ sliding_window = None
189
+
190
+ self.attn = Attention(
191
+ self.num_heads,
192
+ self.head_dim,
193
+ self.scaling,
194
+ num_kv_heads=self.num_kv_heads,
195
+ cache_config=cache_config,
196
+ quant_config=quant_config,
197
+ per_layer_sliding_window=sliding_window,
198
+ prefix=f"{prefix}.attn",
199
+ )
200
+
201
+ def forward(
202
+ self,
203
+ positions: torch.Tensor,
204
+ hidden_states: torch.Tensor,
205
+ ) -> torch.Tensor:
206
+ qkv, _ = self.qkv_proj(hidden_states)
207
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
208
+ q, k = self.rotary_emb(positions, q, k)
209
+ attn_output = self.attn(q, k, v)
210
+ output, _ = self.o_proj(attn_output)
211
+ return output
212
+
213
+
214
+class Jais2DecoderLayer(nn.Module):
215
+ def __init__(
216
+ self,
217
+ vllm_config: VllmConfig,
218
+ config: Jais2Config,
219
+ prefix: str = "",
220
+ ) -> None:
221
+ super().__init__()
222
+
223
+ config = config or vllm_config.model_config.hf_config
224
+ cache_config = vllm_config.cache_config
225
+ quant_config = self.get_quant_config(vllm_config)
226
+
227
+ self.hidden_size = config.hidden_size
228
+ max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
229
+ # Support abacusai/Smaug-72B-v0.1 with attention_bias
230
+ # Support internlm/internlm-7b with bias
231
+ attention_bias = getattr(config, "attention_bias", False) or getattr(
232
+ config, "bias", False
233
+ )
234
+ self.self_attn = Jais2Attention(
235
+ config=config,
236
+ hidden_size=self.hidden_size,
237
+ num_heads=config.num_attention_heads,
238
+ num_kv_heads=getattr(
239
+ config, "num_key_value_heads", config.num_attention_heads
240
+ ),
241
+ max_position_embeddings=max_position_embeddings,
242
+ quant_config=quant_config,
243
+ bias=attention_bias,
244
+ cache_config=cache_config,
245
+ prefix=f"{prefix}.self_attn",
246
+ )
247
+ self.mlp = Jais2MLP(
248
+ hidden_size=self.hidden_size,
249
+ intermediate_size=config.intermediate_size,
250
+ hidden_act=config.hidden_act,
251
+ quant_config=quant_config,
252
+ bias=getattr(config, "mlp_bias", False),
253
+ prefix=f"{prefix}.mlp",
254
+ )
255
+ self.input_layernorm = nn.LayerNorm(
256
+ config.hidden_size, eps=config.layer_norm_eps
257
+ )
258
+ self.post_attention_layernorm = nn.LayerNorm(
259
+ config.hidden_size, eps=config.layer_norm_eps
260
+ )
261
+
262
+ def forward(
263
+ self,
264
+ positions: torch.Tensor,
265
+ hidden_states: torch.Tensor,
266
+ residual: torch.Tensor | None,
267
+ ) -> tuple[torch.Tensor, torch.Tensor]:
268
+ # Self Attention
269
+ if residual is None:
270
+ residual = hidden_states
271
+ hidden_states = self.input_layernorm(hidden_states)
272
+ else:
273
+ hidden_states, residual = (
274
+ self.input_layernorm(hidden_states + residual),
275
+ hidden_states + residual,
276
+ )
277
+ hidden_states = self.self_attn(
278
+ positions=positions,
279
+ hidden_states=hidden_states,
280
+ )
281
+
282
+ # Fully Connected
283
+ hidden_states, residual = (
284
+ self.post_attention_layernorm(hidden_states + residual),
285
+ hidden_states + residual,
286
+ )
287
+ hidden_states = self.mlp(hidden_states)
288
+ return hidden_states, residual
289
+
290
+ def get_quant_config(self, vllm_config: VllmConfig) -> QuantizationConfig | None:
291
+ """Get quantization config for this layer. Override in subclasses."""
292
+ return vllm_config.quant_config
293
+
294
+
295
+@support_torch_compile
296
+class Jais2Model(nn.Module):
297
+ def __init__(
298
+ self,
299
+ vllm_config: VllmConfig,
300
+ prefix: str = "",
301
+ layer_type: type[nn.Module] = Jais2DecoderLayer,
302
+ ):
303
+ super().__init__()
304
+
305
+ config = vllm_config.model_config.hf_config
306
+ quant_config = vllm_config.quant_config
307
+ lora_config = vllm_config.lora_config
308
+
309
+ self.config = config
310
+ self.quant_config = quant_config
311
+ self.padding_idx = config.pad_token_id
312
+ lora_vocab = (
313
+ (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))
314
+ if lora_config
315
+ else 0
316
+ )
317
+ self.vocab_size = config.vocab_size + lora_vocab
318
+ self.org_vocab_size = config.vocab_size
319
+ if get_pp_group().is_first_rank or (
320
+ config.tie_word_embeddings and get_pp_group().is_last_rank
321
+ ):
322
+ self.embed_tokens = VocabParallelEmbedding(
323
+ self.vocab_size,
324
+ config.hidden_size,
325
+ org_num_embeddings=config.vocab_size,
326
+ quant_config=quant_config,
327
+ )
328
+ else:
329
+ self.embed_tokens = PPMissingLayer()
330
+ self.start_layer, self.end_layer, self.layers = make_layers(
331
+ config.num_hidden_layers,
332
+ lambda prefix: layer_type(
333
+ config=config,
334
+ vllm_config=vllm_config,
335
+ prefix=prefix,
336
+ ),
337
+ prefix=f"{prefix}.layers",
338
+ )
339
+ if get_pp_group().is_last_rank:
340
+ self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
341
+ else:
342
+ self.norm = PPMissingLayer()
343
+
344
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
345
+ ["hidden_states", "residual"], config.hidden_size
346
+ )
347
+
348
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
349
+ return self.embed_tokens(input_ids)
350
+
351
+ def forward(
352
+ self,
353
+ input_ids: torch.Tensor | None,
354
+ positions: torch.Tensor,
355
+ intermediate_tensors: IntermediateTensors | None,
356
+ inputs_embeds: torch.Tensor | None = None,
357
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
358
+ if get_pp_group().is_first_rank:
359
+ if inputs_embeds is not None:
360
+ hidden_states = inputs_embeds
361
+ else:
362
+ hidden_states = self.embed_input_ids(input_ids)
363
+ residual = None
364
+ else:
365
+ assert intermediate_tensors is not None
366
+ hidden_states = intermediate_tensors["hidden_states"]
367
+ residual = intermediate_tensors["residual"]
368
+
369
+ for i in range(self.start_layer, self.end_layer):
370
+ layer = self.layers[i]
371
+ hidden_states, residual = layer(positions, hidden_states, residual)
372
+
373
+ if not get_pp_group().is_last_rank:
374
+ return IntermediateTensors(
375
+ {"hidden_states": hidden_states, "residual": residual}
376
+ )
377
+
378
+ hidden_states, _ = self.norm(hidden_states + residual), residual
379
+ return hidden_states
380
+
381
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
382
+ stacked_params_mapping = [
383
+ # (param_name, shard_name, shard_id)
384
+ (".qkv_proj", ".q_proj", "q"),
385
+ (".qkv_proj", ".k_proj", "k"),
386
+ (".qkv_proj", ".v_proj", "v"),
387
+ ]
388
+ params_dict = dict(self.named_parameters())
389
+ loaded_params: set[str] = set()
390
+ for name, loaded_weight in weights:
391
+ if "rotary_emb.inv_freq" in name:
392
+ continue
393
+ if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
394
+ # Models trained using ColossalAI may include these tensors in
395
+ # the checkpoint. Skip them.
396
+ continue
397
+ if self.quant_config is not None and (
398
+ scale_name := self.quant_config.get_cache_scale(name)
399
+ ):
400
+ # Loading kv cache scales for compressed-tensors quantization
401
+ param = params_dict[scale_name]
402
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
403
+ loaded_weight = loaded_weight[0]
404
+ weight_loader(param, loaded_weight)
405
+ loaded_params.add(scale_name)
406
+ continue
407
+ if "scale" in name:
408
+ name = maybe_remap_kv_scale_name(name, params_dict)
409
+ if name is None:
410
+ continue
411
+ for param_name, weight_name, shard_id in stacked_params_mapping:
412
+ if weight_name not in name:
413
+ continue
414
+ name = name.replace(weight_name, param_name)
415
+ # Skip loading extra bias for GPTQ models.
416
+ if name.endswith(".bias") and name not in params_dict:
417
+ continue
418
+
419
+ if is_pp_missing_parameter(name, self):
420
+ continue
421
+
422
+ param = params_dict[name]
423
+ weight_loader = param.weight_loader
424
+ weight_loader(param, loaded_weight, shard_id)
425
+ break
426
+ else:
427
+ # Skip loading extra bias for GPTQ models.
428
+ if name.endswith(".bias") and name not in params_dict:
429
+ continue
430
+ # Remapping the name of FP8 kv-scale.
431
+ name = maybe_remap_kv_scale_name(name, params_dict)
432
+ if name is None:
433
+ continue
434
+
435
+ if is_pp_missing_parameter(name, self):
436
+ continue
437
+
438
+ param = params_dict[name]
439
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
440
+ weight_loader(param, loaded_weight)
441
+ loaded_params.add(name)
442
+ return loaded_params
443
+
444
+
445
+class Jais2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
446
+ packed_modules_mapping = {
447
+ "qkv_proj": ["q_proj", "k_proj", "v_proj"],
448
+ }
449
+
450
+ embedding_modules = {
451
+ "embed_tokens": "input_embeddings",
452
+ "lm_head": "output_embeddings",
453
+ }
454
+
455
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
456
+ super().__init__()
457
+ config = vllm_config.model_config.hf_config
458
+ quant_config = vllm_config.quant_config
459
+ lora_config = vllm_config.lora_config
460
+ self.config = config
461
+ self.lora_config = lora_config
462
+
463
+ self.model = self._init_model(
464
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
465
+ )
466
+
467
+ if get_pp_group().is_last_rank:
468
+ self.unpadded_vocab_size = config.vocab_size
469
+ if lora_config:
470
+ self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
471
+ self.lm_head = ParallelLMHead(
472
+ self.unpadded_vocab_size,
473
+ config.hidden_size,
474
+ org_num_embeddings=config.vocab_size,
475
+ padding_size=(
476
+ DEFAULT_VOCAB_PADDING_SIZE
477
+ # We need bigger padding if using lora for kernel
478
+ # compatibility
479
+ if not lora_config
480
+ else lora_config.lora_vocab_padding_size
481
+ ),
482
+ quant_config=quant_config,
483
+ prefix=maybe_prefix(prefix, "lm_head"),
484
+ )
485
+ if config.tie_word_embeddings:
486
+ self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
487
+
488
+ logit_scale = getattr(config, "logit_scale", 1.0)
489
+ self.logits_processor = LogitsProcessor(
490
+ self.unpadded_vocab_size, config.vocab_size, logit_scale
491
+ )
492
+ else:
493
+ self.lm_head = PPMissingLayer()
494
+
495
+ self.make_empty_intermediate_tensors = (
496
+ self.model.make_empty_intermediate_tensors
497
+ )
498
+
499
+ def _init_model(self, vllm_config: VllmConfig, prefix: str = ""):
500
+ return Jais2Model(vllm_config=vllm_config, prefix=prefix)
501
+
502
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
503
+ return self.model.embed_input_ids(input_ids)
504
+
505
+ def forward(
506
+ self,
507
+ input_ids: torch.Tensor,
508
+ positions: torch.Tensor,
509
+ intermediate_tensors: IntermediateTensors | None = None,
510
+ inputs_embeds: torch.Tensor | None = None,
511
+ ) -> torch.Tensor | IntermediateTensors:
512
+ model_output = self.model(
513
+ input_ids, positions, intermediate_tensors, inputs_embeds
514
+ )
515
+ return model_output
516
+
517
+ def compute_logits(
518
+ self,
519
+ hidden_states: torch.Tensor,
520
+ ) -> torch.Tensor | None:
521
+ logits = self.logits_processor(self.lm_head, hidden_states)
522
+ return logits
523
+
524
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
525
+ loader = AutoWeightsLoader(
526
+ self,
527
+ skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
528
+ )
529
+ return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+1
@@ -127,6 +127,7 @@ _TEXT_GENERATION_MODELS = {
127
"InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"),
128
"InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),
129
"JAISLMHeadModel": ("jais", "JAISLMHeadModel"),
130
+ "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"),
131
"JambaForCausalLM": ("jamba", "JambaForCausalLM"),
132
"KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"), # noqa: E501
133
"Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"),