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 7, 02:45 UTC
Adding support to Sarvam's MoE models (#33942)
Signed-off-by: rahul-sarvam <140298821+rahul-sarvam@users.noreply.github.com>
rahul-sarvam committed
Mar 8, 2026 at 01:16 UTC
85f50eb41fa43783b64e07d768ba3ac6d4ed7a5a
4 files changed
+802
docs/models/supported_models.md
+2
@@ -469,6 +469,8 @@ th {
469
| `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B`, etc. | ✅︎ | ✅︎ |
470
| `Qwen3NextForCausalLM` | Qwen3NextMoE | `Qwen/Qwen3-Next-80B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
471
| `RWForCausalLM` | Falcon RW | `tiiuae/falcon-40b`, etc. | | ✅︎ |
472
+| `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ |
473
+| `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ |
474
| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
475
| `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ |
476
| `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | |
tests/models/registry.py
+12
@@ -480,6 +480,18 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
480
min_transformers_version="4.56.3",
481
),
482
"RWForCausalLM": _HfExamplesInfo("tiiuae/falcon-40b"),
483
+ "SarvamMoEForCausalLM": _HfExamplesInfo(
484
+ "sarvamai/sarvam-30b",
485
+ trust_remote_code=True,
486
+ max_model_len=4096,
487
+ is_available_online=True,
488
+ ),
489
+ "SarvamMLAForCausalLM": _HfExamplesInfo(
490
+ "sarvamai/sarvam-105b",
491
+ trust_remote_code=True,
492
+ max_model_len=4096,
493
+ is_available_online=True,
494
+ ),
495
"SeedOssForCausalLM": _HfExamplesInfo(
496
"ByteDance-Seed/Seed-OSS-36B-Instruct",
497
trust_remote_code=True,
vllm/model_executor/models/registry.py
+2
@@ -191,6 +191,8 @@ _TEXT_GENERATION_MODELS = {
191
"Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
192
"Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
193
"RWForCausalLM": ("falcon", "FalconForCausalLM"),
194
+ "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"),
195
+ "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"),
196
"SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"),
197
"Step1ForCausalLM": ("step1", "Step1ForCausalLM"),
198
"Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"),
vllm/model_executor/models/sarvam.py
new
+786
@@ -0,0 +1,786 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+#
4
+# Copyright 2026 Sarvam AI team. All rights reserved.
5
+#
6
+# This code is based on Llama, Deepseek, and Bailing MoE implementations
7
+# in this library. It has been modified from its original forms to
8
+# accommodate Sarvam's MoE architectures.
9
+#
10
+# Licensed under the Apache License, Version 2.0 (the "License");
11
+# you may not use this file except in compliance with the License.
12
+# You may obtain a copy of the License at
13
+#
14
+# http://www.apache.org/licenses/LICENSE-2.0
15
+#
16
+# Unless required by applicable law or agreed to in writing, software
17
+# distributed under the License is distributed on an "AS IS" BASIS,
18
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19
+# See the License for the specific language governing permissions and
20
+# limitations under the License.
21
+
22
+from __future__ import annotations
23
+
24
+import math
25
+from collections.abc import Iterable, Iterator
26
+from itertools import islice
27
+
28
+import torch
29
+from torch import nn
30
+
31
+from vllm.config import CacheConfig, ParallelConfig, VllmConfig
32
+from vllm.distributed import (
33
+ get_pp_group,
34
+ get_tensor_model_parallel_rank,
35
+ get_tensor_model_parallel_world_size,
36
+)
37
+from vllm.model_executor.layers.activation import SiluAndMul
38
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
39
+from vllm.model_executor.layers.layernorm import RMSNorm
40
+from vllm.model_executor.layers.linear import (
41
+ ColumnParallelLinear,
42
+ MergedColumnParallelLinear,
43
+ ReplicatedLinear,
44
+ RowParallelLinear,
45
+)
46
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
47
+from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
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
+ ParallelLMHead,
52
+ VocabParallelEmbedding,
53
+)
54
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
55
+from vllm.sequence import IntermediateTensors
56
+
57
+from .bailing_moe import BailingMoeForCausalLM
58
+from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
59
+from .utils import (
60
+ AutoWeightsLoader,
61
+ PPMissingLayer,
62
+ is_pp_missing_parameter,
63
+ make_empty_intermediate_tensors_factory,
64
+ make_layers,
65
+ maybe_prefix,
66
+)
67
+
68
+
69
+def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
70
+ if scale <= 1:
71
+ return 1.0
72
+ return 0.1 * mscale * math.log(scale) + 1.0
73
+
74
+
75
+def _is_gate_expert_bias_name(name: str) -> bool:
76
+ return name.endswith(".mlp.gate.e_score_correction_bias") or name.endswith(
77
+ ".gate.e_score_correction_bias"
78
+ )
79
+
80
+
81
+def _zero_mean_tensor(t: torch.Tensor) -> torch.Tensor:
82
+ if t.numel() == 0:
83
+ return t
84
+ return t - t.mean()
85
+
86
+
87
+def _normalized_weights(
88
+ weights: Iterable[tuple[str, torch.Tensor]],
89
+) -> Iterator[tuple[str, torch.Tensor]]:
90
+ for name, w in weights:
91
+ if _is_gate_expert_bias_name(name):
92
+ yield name, _zero_mean_tensor(w)
93
+ else:
94
+ yield name, w
95
+
96
+
97
+class SarvamMLAAttention(nn.Module):
98
+ def __init__(
99
+ self,
100
+ vllm_config: VllmConfig,
101
+ config,
102
+ cache_config: CacheConfig | None = None,
103
+ quant_config: QuantizationConfig | None = None,
104
+ prefix: str = "",
105
+ ) -> None:
106
+ super().__init__()
107
+
108
+ self.config = config
109
+ self.hidden_size = config.hidden_size
110
+ self.qk_nope_head_dim = config.qk_nope_head_dim
111
+ self.qk_rope_head_dim = config.qk_rope_head_dim
112
+ self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
113
+ self.v_head_dim = config.v_head_dim
114
+
115
+ self.q_lora_rank = getattr(config, "q_lora_rank", None)
116
+ self.kv_lora_rank = config.kv_lora_rank
117
+
118
+ self.total_num_heads = config.num_attention_heads
119
+ tp_size = get_tensor_model_parallel_world_size()
120
+ assert self.total_num_heads % tp_size == 0
121
+ self.num_local_heads = self.total_num_heads // tp_size
122
+
123
+ self.scaling = self.qk_head_dim**-0.5
124
+ self.max_position_embeddings = config.max_position_embeddings
125
+
126
+ if self.q_lora_rank is not None:
127
+ self.q_a_proj = ReplicatedLinear(
128
+ self.hidden_size,
129
+ self.q_lora_rank,
130
+ bias=False,
131
+ quant_config=quant_config,
132
+ prefix=f"{prefix}.q_a_proj",
133
+ )
134
+ self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
135
+ self.q_b_proj = ColumnParallelLinear(
136
+ self.q_lora_rank,
137
+ self.total_num_heads * self.qk_head_dim,
138
+ bias=False,
139
+ quant_config=quant_config,
140
+ prefix=f"{prefix}.q_b_proj",
141
+ )
142
+ self.q_proj = None # type: ignore
143
+ else:
144
+ self.q_proj = ColumnParallelLinear(
145
+ self.hidden_size,
146
+ self.total_num_heads * self.qk_head_dim,
147
+ bias=False,
148
+ quant_config=quant_config,
149
+ prefix=f"{prefix}.q_proj",
150
+ )
151
+ self.q_a_proj = None # type: ignore
152
+ self.q_a_layernorm = None # type: ignore
153
+ self.q_b_proj = None # type: ignore
154
+
155
+ # KV latent (MQA-style) A-proj
156
+ self.kv_a_proj_with_mqa = ReplicatedLinear(
157
+ self.hidden_size,
158
+ self.kv_lora_rank + self.qk_rope_head_dim,
159
+ bias=False,
160
+ quant_config=quant_config,
161
+ prefix=f"{prefix}.kv_a_proj_with_mqa",
162
+ )
163
+ self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
164
+
165
+ # KV B-proj produces per-head K_nope and V
166
+ self.kv_b_proj = ColumnParallelLinear(
167
+ self.kv_lora_rank,
168
+ self.total_num_heads * (self.qk_nope_head_dim + self.v_head_dim),
169
+ bias=False,
170
+ quant_config=quant_config,
171
+ prefix=f"{prefix}.kv_b_proj",
172
+ )
173
+
174
+ self.o_proj = RowParallelLinear(
175
+ self.total_num_heads * self.v_head_dim,
176
+ self.hidden_size,
177
+ bias=False,
178
+ quant_config=quant_config,
179
+ prefix=f"{prefix}.o_proj",
180
+ )
181
+
182
+ self.rotary_emb = get_rope(
183
+ self.qk_rope_head_dim,
184
+ # rotary_dim=self.qk_rope_head_dim,
185
+ max_position=config.max_position_embeddings,
186
+ rope_parameters=config.rope_parameters,
187
+ is_neox_style=False,
188
+ )
189
+
190
+ if config.rope_parameters.get("rope_type", None) == "deepseek_yarn":
191
+ mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
192
+ scaling_factor = config.rope_parameters["factor"]
193
+ mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
194
+ self.scaling = self.scaling * mscale * mscale
195
+
196
+ mla_modules = MLAModules(
197
+ kv_a_layernorm=self.kv_a_layernorm,
198
+ kv_b_proj=self.kv_b_proj,
199
+ rotary_emb=self.rotary_emb,
200
+ o_proj=self.o_proj,
201
+ fused_qkv_a_proj=None,
202
+ kv_a_proj_with_mqa=self.kv_a_proj_with_mqa,
203
+ q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None,
204
+ q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
205
+ q_proj=self.q_proj if self.q_lora_rank is None else None,
206
+ indexer=None,
207
+ indexer_rotary_emb=None,
208
+ is_sparse=False,
209
+ topk_indices_buffer=None,
210
+ )
211
+
212
+ self.mla_attn = MultiHeadLatentAttentionWrapper(
213
+ self.hidden_size,
214
+ self.num_local_heads,
215
+ self.scaling,
216
+ self.qk_nope_head_dim,
217
+ self.qk_rope_head_dim,
218
+ self.v_head_dim,
219
+ self.q_lora_rank,
220
+ self.kv_lora_rank,
221
+ mla_modules,
222
+ cache_config=cache_config,
223
+ quant_config=quant_config,
224
+ prefix=prefix,
225
+ )
226
+
227
+ def forward(
228
+ self,
229
+ positions: torch.Tensor,
230
+ hidden_states: torch.Tensor,
231
+ ) -> torch.Tensor:
232
+ return self.mla_attn(positions, hidden_states, llama_4_scaling=None)
233
+
234
+
235
+class SarvamMLAMLP(nn.Module):
236
+ def __init__(
237
+ self,
238
+ intermediate_size: int,
239
+ config,
240
+ quant_config: QuantizationConfig | None = None,
241
+ reduce_results: bool = True,
242
+ prefix: str = "",
243
+ ) -> None:
244
+ super().__init__()
245
+
246
+ self.gate_up_proj = MergedColumnParallelLinear(
247
+ config.hidden_size,
248
+ [intermediate_size] * 2,
249
+ bias=False,
250
+ quant_config=quant_config,
251
+ prefix=f"{prefix}.gate_up_proj",
252
+ )
253
+ self.down_proj = RowParallelLinear(
254
+ intermediate_size,
255
+ config.hidden_size,
256
+ bias=False,
257
+ quant_config=quant_config,
258
+ reduce_results=reduce_results,
259
+ prefix=f"{prefix}.down_proj",
260
+ )
261
+ self.act_fn = SiluAndMul()
262
+
263
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
264
+ gate_up, _ = self.gate_up_proj(x)
265
+ x = self.act_fn(gate_up)
266
+ x, _ = self.down_proj(x)
267
+ return x
268
+
269
+
270
+class SarvamMLAMoE(nn.Module):
271
+ def __init__(
272
+ self,
273
+ config,
274
+ parallel_config: ParallelConfig,
275
+ quant_config: QuantizationConfig | None = None,
276
+ prefix: str = "",
277
+ ) -> None:
278
+ super().__init__()
279
+
280
+ self.config = config
281
+ self.tp_size = get_tensor_model_parallel_world_size()
282
+ self.tp_rank = get_tensor_model_parallel_rank()
283
+ self.hidden_size = config.hidden_size
284
+
285
+ self.num_experts = config.num_experts
286
+ self.top_k = config.num_experts_per_tok
287
+ self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 2.5)
288
+
289
+ self.n_group = getattr(config, "n_group", None)
290
+ self.topk_group = getattr(config, "topk_group", None)
291
+ self.use_grouped_topk = self.n_group is not None and self.topk_group is not None
292
+
293
+ self.norm_expert_prob = getattr(config, "norm_topk_prob", True)
294
+
295
+ router_dtype_cfg = getattr(config, "router_dtype", "fp32")
296
+ if router_dtype_cfg is None:
297
+ self.router_dtype = None
298
+ elif router_dtype_cfg == "fp32":
299
+ self.router_dtype = torch.float32
300
+ else:
301
+ self.router_dtype = torch.bfloat16
302
+
303
+ self.gate = nn.Linear(
304
+ self.hidden_size,
305
+ self.num_experts,
306
+ bias=False,
307
+ dtype=self.router_dtype,
308
+ )
309
+
310
+ if getattr(config, "moe_router_enable_expert_bias", True):
311
+ self.gate.e_score_correction_bias = nn.Parameter(
312
+ torch.empty(
313
+ (self.num_experts,),
314
+ dtype=torch.float32,
315
+ )
316
+ )
317
+ else:
318
+ self.gate.e_score_correction_bias = None
319
+
320
+ self.score_function = getattr(config, "score_function", "sigmoid")
321
+ self.num_shared_experts = getattr(config, "num_shared_experts", 1)
322
+ if self.num_shared_experts > 0:
323
+ if hasattr(config, "moe_shared_expert_intermediate_size"):
324
+ shared_int = config.moe_shared_expert_intermediate_size
325
+ else:
326
+ shared_int = config.moe_intermediate_size
327
+ shared_int *= self.num_shared_experts
328
+ self.shared_experts = SarvamMLAMLP(
329
+ intermediate_size=shared_int,
330
+ config=config,
331
+ quant_config=quant_config,
332
+ reduce_results=False,
333
+ prefix=f"{prefix}.shared_experts",
334
+ )
335
+ else:
336
+ self.shared_experts = None
337
+
338
+ self.experts = SharedFusedMoE(
339
+ shared_experts=self.shared_experts,
340
+ num_experts=self.num_experts,
341
+ top_k=self.top_k,
342
+ hidden_size=self.hidden_size,
343
+ intermediate_size=config.moe_intermediate_size,
344
+ reduce_results=False,
345
+ renormalize=self.norm_expert_prob,
346
+ quant_config=quant_config,
347
+ prefix=f"{prefix}.experts",
348
+ scoring_func=self.score_function,
349
+ e_score_correction_bias=self.gate.e_score_correction_bias,
350
+ num_expert_group=self.n_group,
351
+ topk_group=self.topk_group,
352
+ use_grouped_topk=self.use_grouped_topk,
353
+ routed_scaling_factor=self.routed_scaling_factor,
354
+ )
355
+
356
+ def maybe_get_fused_moe(self) -> SharedFusedMoE:
357
+ return self.experts
358
+
359
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
360
+ num_tokens, hidden_dim = hidden_states.shape
361
+ hidden_states = hidden_states.view(-1, hidden_dim)
362
+ router_logits = self.gate(
363
+ hidden_states.to(self.router_dtype)
364
+ if self.router_dtype is not None
365
+ else hidden_states
366
+ )
367
+ router_logits = router_logits.to(hidden_states.dtype)
368
+ final_hidden = self.experts(
369
+ hidden_states=hidden_states,
370
+ router_logits=router_logits,
371
+ )
372
+
373
+ if self.shared_experts is not None:
374
+ shared_output, expert_output = final_hidden
375
+ else:
376
+ shared_output, expert_output = None, final_hidden
377
+
378
+ if shared_output is not None:
379
+ expert_output = expert_output + shared_output
380
+
381
+ if self.tp_size > 1:
382
+ expert_output = self.experts.maybe_all_reduce_tensor_model_parallel(
383
+ expert_output
384
+ )
385
+
386
+ return expert_output.view(num_tokens, hidden_dim)
387
+
388
+
389
+class SarvamMLABlock(nn.Module):
390
+ def __init__(
391
+ self,
392
+ vllm_config: VllmConfig,
393
+ prefix: str = "",
394
+ ) -> None:
395
+ super().__init__()
396
+ config = vllm_config.model_config.hf_config
397
+ cache_config = vllm_config.cache_config
398
+ quant_config = vllm_config.quant_config
399
+ parallel_config = vllm_config.parallel_config
400
+ layer_idx = int(prefix.split(".")[-1])
401
+ hidden_size = config.hidden_size
402
+ dense_intermediate = getattr(config, "intermediate_size", 16384)
403
+
404
+ self.input_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
405
+ self.self_attn = SarvamMLAAttention(
406
+ vllm_config=vllm_config,
407
+ config=config,
408
+ cache_config=cache_config,
409
+ quant_config=quant_config,
410
+ prefix=f"{prefix}.self_attn",
411
+ )
412
+ self.post_attention_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
413
+ use_moe = hasattr(config, "num_experts") and config.num_experts is not None
414
+ first_k_dense = getattr(config, "first_k_dense_replace", 1)
415
+ moe_layer_freq = getattr(config, "moe_layer_freq", 1)
416
+ if use_moe:
417
+ is_moe_layer = layer_idx >= first_k_dense and (
418
+ (layer_idx - first_k_dense) % moe_layer_freq == 0
419
+ )
420
+ else:
421
+ is_moe_layer = False
422
+
423
+ if is_moe_layer:
424
+ self.mlp = SarvamMLAMoE(
425
+ config=config,
426
+ parallel_config=parallel_config,
427
+ quant_config=quant_config,
428
+ prefix=f"{prefix}.mlp",
429
+ )
430
+ else:
431
+ self.mlp = SarvamMLAMLP(
432
+ intermediate_size=dense_intermediate,
433
+ config=config,
434
+ quant_config=quant_config,
435
+ reduce_results=True,
436
+ prefix=f"{prefix}.mlp",
437
+ )
438
+
439
+ def forward(
440
+ self,
441
+ hidden_states: torch.Tensor,
442
+ positions: torch.Tensor,
443
+ residual: torch.Tensor | None,
444
+ ) -> tuple[torch.Tensor, torch.Tensor]:
445
+ if residual is None:
446
+ residual = hidden_states
447
+ hidden_states = self.input_layernorm(hidden_states)
448
+ else:
449
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
450
+
451
+ hidden_states = self.self_attn(
452
+ positions=positions,
453
+ hidden_states=hidden_states,
454
+ )
455
+ hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
456
+ hidden_states = self.mlp(hidden_states)
457
+ return hidden_states, residual
458
+
459
+
460
+class SarvamMLAModel(nn.Module):
461
+ def __init__(
462
+ self,
463
+ *,
464
+ vllm_config: VllmConfig,
465
+ prefix: str = "",
466
+ ) -> None:
467
+ super().__init__()
468
+
469
+ config = vllm_config.model_config.hf_config
470
+ quant_config = vllm_config.quant_config
471
+
472
+ self.config = config
473
+ self.vocab_size = config.vocab_size
474
+ self.embed_dim = config.hidden_size
475
+ self.tie_word_embeddings = getattr(config, "tie_word_embeddings", False)
476
+ if get_pp_group().is_first_rank or (
477
+ self.tie_word_embeddings and get_pp_group().is_last_rank
478
+ ):
479
+ self.embed_tokens = VocabParallelEmbedding(
480
+ self.vocab_size,
481
+ self.embed_dim,
482
+ quant_config=quant_config,
483
+ prefix=f"{prefix}.embed_tokens",
484
+ )
485
+ else:
486
+ self.embed_tokens = PPMissingLayer()
487
+
488
+ self.embedding_dropout = torch.nn.Dropout(
489
+ getattr(config, "embedding_dropout", 0.0)
490
+ )
491
+ self.start_layer, self.end_layer, self.layers = make_layers(
492
+ config.num_hidden_layers,
493
+ lambda prefix: SarvamMLABlock(
494
+ vllm_config=vllm_config,
495
+ prefix=prefix,
496
+ ),
497
+ prefix=f"{prefix}.layers",
498
+ )
499
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
500
+ ["hidden_states", "residual"], config.hidden_size
501
+ )
502
+ if get_pp_group().is_last_rank:
503
+ self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps)
504
+ else:
505
+ self.norm = PPMissingLayer()
506
+
507
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
508
+ return self.embed_tokens(input_ids)
509
+
510
+ def forward(
511
+ self,
512
+ input_ids: torch.Tensor,
513
+ positions: torch.Tensor,
514
+ intermediate_tensors: IntermediateTensors | None,
515
+ inputs_embeds: torch.Tensor | None = None,
516
+ ) -> torch.Tensor | IntermediateTensors:
517
+ if get_pp_group().is_first_rank:
518
+ if inputs_embeds is not None:
519
+ hidden_states = inputs_embeds
520
+ else:
521
+ hidden_states = self.embed_input_ids(input_ids)
522
+ hidden_states = self.embedding_dropout(hidden_states)
523
+ residual = None
524
+ else:
525
+ assert intermediate_tensors is not None
526
+ hidden_states = intermediate_tensors["hidden_states"]
527
+ residual = intermediate_tensors["residual"]
528
+
529
+ for layer in islice(self.layers, self.start_layer, self.end_layer):
530
+ hidden_states, residual = layer(
531
+ hidden_states,
532
+ positions,
533
+ residual,
534
+ )
535
+ if not get_pp_group().is_last_rank:
536
+ return IntermediateTensors(
537
+ {"hidden_states": hidden_states, "residual": residual}
538
+ )
539
+ if residual is None:
540
+ hidden_states = self.norm(hidden_states)
541
+ else:
542
+ hidden_states, _ = self.norm(hidden_states, residual)
543
+ return hidden_states
544
+
545
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
546
+ return SharedFusedMoE.make_expert_params_mapping(
547
+ self,
548
+ ckpt_gate_proj_name="gate_proj",
549
+ ckpt_down_proj_name="down_proj",
550
+ ckpt_up_proj_name="up_proj",
551
+ num_experts=self.config.num_experts,
552
+ )
553
+
554
+ def load_weights(
555
+ self,
556
+ weights: Iterable[tuple[str, torch.Tensor]],
557
+ ) -> set[str]:
558
+ """Load weights with stacked gate+up and MoE expert remapping."""
559
+ weights = _normalized_weights(weights)
560
+ stacked_params_mapping = [
561
+ ("gate_up_proj", "gate_proj", 0),
562
+ ("gate_up_proj", "up_proj", 1),
563
+ ]
564
+
565
+ params_dict = dict(self.named_parameters(remove_duplicate=False))
566
+ loaded_params: set[str] = set()
567
+ expert_params_mapping = self.get_expert_mapping()
568
+
569
+ for name, loaded_weight in weights:
570
+ for param_name, weight_name, shard_id in stacked_params_mapping:
571
+ if weight_name not in name:
572
+ continue
573
+ if "mlp.experts" in name:
574
+ continue
575
+ new_name = name.replace(weight_name, param_name)
576
+ if new_name.endswith(".bias") and new_name not in params_dict:
577
+ continue
578
+ if new_name not in params_dict:
579
+ continue
580
+ if is_pp_missing_parameter(new_name, self):
581
+ continue
582
+
583
+ param = params_dict[new_name]
584
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
585
+ weight_loader(param, loaded_weight, shard_id)
586
+ loaded_params.add(new_name)
587
+ break
588
+ else:
589
+ mapped = False
590
+ for (
591
+ param_name,
592
+ weight_name,
593
+ expert_id,
594
+ shard_id,
595
+ ) in expert_params_mapping:
596
+ if weight_name not in name:
597
+ continue
598
+
599
+ new_name = name.replace(weight_name, param_name)
600
+ if is_pp_missing_parameter(new_name, self):
601
+ continue
602
+ if new_name not in params_dict:
603
+ continue
604
+
605
+ param = params_dict[new_name]
606
+ weight_loader = getattr(
607
+ param, "weight_loader", default_weight_loader
608
+ )
609
+ weight_loader(
610
+ param,
611
+ loaded_weight,
612
+ name,
613
+ shard_id=shard_id,
614
+ expert_id=expert_id,
615
+ )
616
+ loaded_params.add(new_name)
617
+ mapped = True
618
+ break
619
+
620
+ if mapped:
621
+ continue
622
+
623
+ if name.endswith(".bias") and name not in params_dict:
624
+ continue
625
+ if name not in params_dict:
626
+ continue
627
+ if is_pp_missing_parameter(name, self):
628
+ continue
629
+
630
+ param = params_dict[name]
631
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
632
+ weight_loader(param, loaded_weight)
633
+ loaded_params.add(name)
634
+
635
+ return loaded_params
636
+
637
+
638
+class SarvamMixtureOfExperts(MixtureOfExperts):
639
+ def extract_moe_parameters(self, example_moe: SarvamMLAMoE | None) -> None:
640
+ if example_moe is None:
641
+ raise RuntimeError("No SarvamMLAMoE layer found in model.layers.")
642
+
643
+ self.num_logical_experts = example_moe.num_experts
644
+ self.num_routed_experts = example_moe.num_experts # routed pool size
645
+ self.num_shared_experts = getattr(example_moe.config, "num_shared_experts", 1)
646
+
647
+ self.num_physical_experts = self.num_logical_experts
648
+ self.num_local_physical_experts = self.num_logical_experts
649
+ self.num_redundant_experts = 0
650
+
651
+ def update_physical_experts_metadata(
652
+ self,
653
+ num_physical_experts: int,
654
+ num_local_physical_experts: int,
655
+ ) -> None:
656
+ self.num_physical_experts = num_physical_experts
657
+ self.num_local_physical_experts = num_local_physical_experts
658
+ self.num_redundant_experts = num_physical_experts - self.num_logical_experts
659
+
660
+ for moe in self.moe_mlp_layers:
661
+ moe.n_physical_experts = num_physical_experts
662
+ moe.n_local_physical_experts = num_local_physical_experts
663
+ moe.n_redundant_experts = self.num_redundant_experts
664
+
665
+ fused = moe.experts
666
+ if hasattr(fused, "n_local_physical_experts"):
667
+ fused.n_local_physical_experts = num_local_physical_experts
668
+ if hasattr(fused, "n_physical_experts"):
669
+ fused.n_physical_experts = num_physical_experts
670
+ if hasattr(fused, "n_redundant_experts"):
671
+ fused.n_redundant_experts = self.num_redundant_experts
672
+ if hasattr(fused, "update_expert_map"):
673
+ fused.update_expert_map()
674
+
675
+ def set_eplb_state(self, eplb_state) -> None:
676
+ self.eplb_state = eplb_state
677
+ for moe in self.moe_layers:
678
+ if hasattr(moe, "set_eplb_state"):
679
+ moe.set_eplb_state(eplb_state)
680
+
681
+
682
+class SarvamMLAForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SarvamMixtureOfExperts):
683
+ packed_modules_mapping = {
684
+ "q_proj": ["q_proj"],
685
+ "q_a_proj": ["q_a_proj"],
686
+ "q_b_proj": ["q_b_proj"],
687
+ "kv_a_proj_with_mqa": ["kv_a_proj_with_mqa"],
688
+ "kv_b_proj": ["kv_b_proj"],
689
+ "gate_up_proj": ["gate_proj", "up_proj"],
690
+ }
691
+
692
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
693
+ super().__init__()
694
+ config = vllm_config.model_config.hf_config
695
+ quant_config = vllm_config.quant_config
696
+ self.config = config
697
+ self.quant_config = quant_config
698
+
699
+ self.model = SarvamMLAModel(
700
+ vllm_config=vllm_config,
701
+ prefix=maybe_prefix(prefix, "model"),
702
+ )
703
+
704
+ self.tie_word_embeddings = getattr(config, "tie_word_embeddings", False)
705
+ if get_pp_group().is_last_rank:
706
+ if self.tie_word_embeddings:
707
+ self.lm_head = self.model.embed_tokens
708
+ else:
709
+ self.lm_head = ParallelLMHead(
710
+ config.vocab_size,
711
+ config.hidden_size,
712
+ quant_config=quant_config,
713
+ prefix=maybe_prefix(prefix, "lm_head"),
714
+ )
715
+ self.logits_processor = LogitsProcessor(config.vocab_size)
716
+ else:
717
+ self.lm_head = PPMissingLayer()
718
+ self.logits_processor = None # type: ignore
719
+
720
+ self.make_empty_intermediate_tensors = (
721
+ self.model.make_empty_intermediate_tensors
722
+ )
723
+
724
+ self.expert_weights = []
725
+ self.num_moe_layers = 0
726
+
727
+ self.moe_layers = []
728
+ self.moe_mlp_layers = []
729
+
730
+ example_moe = None
731
+ for layer in self.model.layers:
732
+ if isinstance(layer, PPMissingLayer):
733
+ continue
734
+ if isinstance(layer.mlp, SarvamMLAMoE):
735
+ example_moe = layer.mlp
736
+ self.moe_mlp_layers.append(layer.mlp)
737
+ self.moe_layers.append(layer.mlp.experts)
738
+ self.num_moe_layers += 1
739
+
740
+ self.extract_moe_parameters(example_moe)
741
+
742
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
743
+ return self.model.embed_input_ids(input_ids)
744
+
745
+ def forward(
746
+ self,
747
+ input_ids: torch.Tensor,
748
+ positions: torch.Tensor,
749
+ intermediate_tensors: IntermediateTensors | None = None,
750
+ inputs_embeds: torch.Tensor | None = None,
751
+ ) -> torch.Tensor | IntermediateTensors:
752
+ return self.model(
753
+ input_ids=input_ids,
754
+ positions=positions,
755
+ intermediate_tensors=intermediate_tensors,
756
+ inputs_embeds=inputs_embeds,
757
+ )
758
+
759
+ def compute_logits(
760
+ self,
761
+ hidden_states: torch.Tensor,
762
+ ) -> torch.Tensor | None:
763
+ if not get_pp_group().is_last_rank:
764
+ return None
765
+ logits = self.logits_processor(self.lm_head, hidden_states)
766
+ return logits
767
+
768
+ def load_weights(
769
+ self,
770
+ weights: Iterable[tuple[str, torch.Tensor]],
771
+ ) -> set[str]:
772
+ loader = AutoWeightsLoader(
773
+ self,
774
+ skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
775
+ )
776
+ return loader.load_weights(weights)
777
+
778
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
779
+ return self.model.get_expert_mapping()
780
+
781
+
782
+class SarvamMoEForCausalLM(BailingMoeForCausalLM):
783
+ """Same as BailingMoeForCausalLM, but normalizes gate expert_bias pre-load."""
784
+
785
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
786
+ return super().load_weights(_normalized_weights(weights))