[Model] Optimize Sarvam MLA routing and preserve FP32 router logits (#56034)
Signed-off-by: Harshit Kedia <harshit@sarvam.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Harshit Kedia committed
Sep 16, 2026 at 17:32 UTC
6a2fdf9ac6f929f937dbe4617f1993c068045f81
1 file changed
+21
-10
vllm/model_executor/models/sarvam.py
+21
-10
@@ -36,7 +36,11 @@ from vllm.distributed import (
36
get_tensor_model_parallel_world_size,
37
)
38
from vllm.model_executor.layers.activation import SiluAndMul
39
-from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner
39
+from vllm.model_executor.layers.fused_moe import (
40
+ FusedMoEFactory,
41
+ GateLinear,
42
+ MoERunner,
43
+)
44
from vllm.model_executor.layers.layernorm import RMSNorm
45
from vllm.model_executor.layers.linear import (
46
ColumnParallelLinear,
@@ -306,11 +310,11 @@ class SarvamMLAMoE(nn.Module):
310
else:
311
self.router_dtype = torch.bfloat16
312
309
- self.gate = nn.Linear(
313
+ self.gate = GateLinear(
314
self.hidden_size,
315
self.num_experts,
312
- bias=False,
313
- dtype=self.router_dtype,
316
+ out_dtype=self.router_dtype,
317
+ prefix=f"{prefix}.gate",
318
)
319
320
if getattr(config, "moe_router_enable_expert_bias", True):
@@ -365,12 +369,7 @@ class SarvamMLAMoE(nn.Module):
369
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
370
num_tokens, hidden_dim = hidden_states.shape
371
hidden_states = hidden_states.view(-1, hidden_dim)
368
- router_logits = self.gate(
369
- hidden_states.to(self.router_dtype)
370
- if self.router_dtype is not None
371
- else hidden_states
372
- )
373
- router_logits = router_logits.to(hidden_states.dtype)
372
+ router_logits, _ = self.gate(hidden_states)
373
final_hidden = self.experts(
374
hidden_states=hidden_states,
375
router_logits=router_logits,
@@ -643,9 +642,21 @@ class SarvamMLAForCausalLM(
642
"gate_up_proj": ["gate_proj", "up_proj"],
643
}
644
645
+ @staticmethod
646
+ def _remap_config(config) -> None:
647
+ """Default the routing keys the released checkpoints omit."""
648
+ defaults = {
649
+ "n_group": 1,
650
+ "topk_group": 1,
651
+ }
652
+ for attr, default in defaults.items():
653
+ if getattr(config, attr, None) is None:
654
+ setattr(config, attr, default)
655
+
656
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
657
super().__init__()
658
config = vllm_config.model_config.hf_config
659
+ self._remap_config(config)
660
quant_config = vllm_config.quant_config
661
self.config = config
662
self.quant_config = quant_config