Duplicate from sarvamai/sarvam-30b

Co-authored-by: Rahul <rahular@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,296 @@
+---
+language:
+  - en
+  - hi
+  - bn
+  - ta
+  - te
+  - mr
+  - gu
+  - kn
+  - ml
+  - pa
+  - or
+  - as
+  - ur
+  - sa
+  - ne
+  - sd
+  - kok
+  - mai
+  - doi
+  - mni
+  - sat
+  - ks
+  - bo
+library_name: transformers
+license: apache-2.0
+pipeline_tag: text-generation
+---
+
+![image](https://cdn-uploads.huggingface.co/production/uploads/60270a7c32856987162c641a/SivoCJWJqex41oprnwyuK.png)
+
+Want a bigger model? Download [Sarvam-105B](https://huggingface.co/sarvamai/sarvam-105b)!
+
+## Index
+
+1. [Introduction](#introduction)  
+2. [Architecture](#architecture)  
+3. [Benchmarks](#benchmarks)  
+   - Knowledge & Coding  
+   - Reasoning & Math  
+   - Agentic  
+4. [Inference](#inference)  
+   - Hugging Face  
+   - [vLLM](https://github.com/vllm-project/vllm)
+   - [SGLang](https://github.com/sgl-project/sglang)
+5. [Footnote](#footnote)  
+6. [Citation](#citation)  
+
+## Introduction
+
+**Sarvam-30B** is an advanced Mixture-of-Experts (MoE) model with 2.4B non-embedding active parameters, designed primarily for practical deployment. It combines strong reasoning, reliable coding ability, and best-in-class conversational quality across Indian languages. Sarvam-30B is built to run reliably in resource-constrained environments and can handle multilingual voice calls while performing tool calls.
+
+A major focus during training was the Indian context and languages, resulting in **state-of-the-art performance across 22 Indian languages** for its model size.
+
+Sarvam-30B is open-sourced under the **Apache License**. For more details, see our [blog](https://www.sarvam.ai/blogs/sarvam-30b-105b).
+
+## Architecture
+
+The 30B MoE model is designed for throughput and memory efficiency. It uses 19 layers, a dense FFN `intermediate_size` of 8192, `moe_intermediate_size` of 1024, top-6 routing, grouped KV heads (`num_key_value_heads=4`), and an extremely high rope_theta (`8e6`) for long-context stability without RoPE scaling. It has 128 experts with a shared expert, a routed scaling factor of 2.5, and auxiliary-loss-free router balancing. The 30B model focuses on throughput and memory efficiency through fewer layers, grouped KV attention, and smaller experts.
+
+## Benchmarks
+
+<details>
+  <summary>Knowledge & Coding</summary>
+
+| Benchmark | Sarvam-30B | Gemma 27B It | Mistral-3.2-24B | OLMo 3.1 32B Think | Nemotron-3-Nano-30B-A3B | Qwen3-30B-Thinking-2507 | GLM 4.7 Flash | GPT-OSS-20B |
+|---|---|---|---|---|---|---|---|---|
+| Math500 | 97.0 | 87.4 | 69.4 | 96.2 | 98.0 | 97.6 | 97.0 | 94.2 |
+| HumanEval | 92.1 | 88.4 | 92.9 | 95.1 | 97.6 | 95.7 | 96.3 | 95.7 |
+| MBPP | 92.7 | 81.8 | 78.3 | 58.7 | 91.9 | 94.3 | 91.8 | 95.3 |
+| Live Code Bench v6 | 70.0 | 28.0 | 26.0 | 73.0 | 68.3 | 66.0 | 64.0 | 61.0 |
+| MMLU | 85.1 | 81.2 | 80.5 | 86.4 | 84.0 | 88.4 | 86.9 | 85.3 |
+| MMLU Pro | 80.0 | 68.1 | 69.1 | 72.0 | 78.3 | 80.9 | 73.6 | 75.0 |
+| MILU | 76.8 | 69.2 | 67.9 | 69.9 | 64.8 | 82.6 | 75.6 | 73.7 |
+| Arena Hard v2 | 49.0 | 50.1 | 43.1 | 42.0 | 67.7 | 72.1 | 58.1 | 62.9 |
+| Writing Bench | 78.7 | 71.4 | 70.3 | 75.7 | 83.7 | 85.0 | 79.2 | 79.1 |
+
+</details>
+
+<details>
+  <summary>Reasoning & Math</summary>
+
+| Benchmark | Sarvam-30B | OLMo 3.1 32B | Nemotron-3-Nano-30B | Qwen3-30B-Thinking-2507 | GLM 4.7 Flash | GPT-OSS-20B |
+|---|---|---|---|---|---|---|
+| GPQA Diamond | 66.5 | 57.5 | 73.0 | 73.4 | 75.2 | 71.5 |
+| AIME 25 (w/ Tools) | 88.3 (96.7) | 78.1 (81.7) | 89.1 (99.2) | 85.0 (-) | 91.6 (-) | 91.7 (98.7) |
+| HMMT (Feb 25) | 73.3 | 51.7 | 85.0 | 71.4 | 85.0 | 76.7 |
+| HMMT (Nov 25) | 74.2 | 58.3 | 75.0 | 73.3 | 81.7 | 68.3 |
+| Beyond AIME | 58.3 | 48.5 | 64.0 | 61.0 | 60.0 | 46.0 |
+
+</details>
+
+<details>
+  <summary>Agentic</summary>
+
+| Benchmark | Sarvam-30B | Nemotron-3-Nano-30B | Qwen3-30B-Thinking-2507 | GLM 4.7 Flash | GPT-OSS-20B |
+|---|---|---|---|---|---|
+| BrowseComp | 35.5 | 23.8 | 2.9 | 42.8 | 28.3 |
+| SWE Bench Verified | 34.0 | 38.8 | 22.0 | 59.2 | 34.0 |
+| τ² Bench (avg.) | 45.7 | 49.0 | 47.7 | 79.5 | 48.7 |
+
+> See footnote for evaluation details.
+
+</details>
+
+## Inference
+
+<details>
+  <summary>Huggingface</summary>
+  
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+
+model_name = "sarvamai/sarvam-30b"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, device_map="auto")
+
+def generate_text(
+    prompt: str,
+    max_new_tokens: int = 2048,
+    temperature: float = 0.8,
+    top_p: float = 0.95,
+    repetition_penalty: float = 1.0,
+) -> None:
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")
+
+    generation_config = GenerationConfig(
+        max_new_tokens=max_new_tokens,
+        repetition_penalty=repetition_penalty,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True,
+    )
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            generation_config=generation_config,
+        )
+    return tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+prompts = [
+    "What is the capital city of New Zealand?",
+]
+
+for prompt in prompts:
+    templated_prompt = tokenizer.apply_chat_template(
+      [{"role": "user", "content": prompt}],
+      tokenize=False,
+      add_generation_prompt=True,
+      enable_thinking=True
+    )
+    output = generate_text(templated_prompt, max_new_tokens=512)
+    print("Prompt: ", prompt)
+    print("Generated text: ", output)
+    print("=" * 100)
+```
+</details>
+
+<details>
+  <summary>SGLang</summary>
+
+**Install latest SGLang from source**
+
+```bash
+git clone https://github.com/sgl-project/sglang.git
+cd sglang
+pip install -e "python[all]"
+```
+
+**Instantiate model and Run**
+
+```python
+import sglang as sgl
+from transformers import AutoTokenizer
+
+model_path = "sarvamai/sarvam-30b"
+engine = sgl.Engine(
+    model_path=model_path,
+    tp_size=2,
+    mem_fraction_static=0.8,
+    trust_remote_code=True,
+    dtype="bfloat16",
+    prefill_attention_backend="fa3",
+    decode_attention_backend="fa3",
+)
+
+sampling_params = {
+    "temperature": 0.8,
+    "max_new_tokens": 2048,
+    "repetition_penalty": 1.0,
+}
+
+prompts = [
+    "Which treaty formally ended World War I and imposed heavy reparations on Germany?",
+]
+
+outputs = engine.generate([
+    tokenizer.apply_chat_template([
+        {"role": "user", "content": prompt}],
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=True)
+        for prompt in prompts], 
+    sampling_params)
+for p, o in zip(prompts, outputs):
+    print("Prompt: ", p)
+    print("Generated text: ", o['text'])
+    print("=" * 100)
+```
+</details>
+
+<details>
+  <summary>vLLM</summary>
+
+Note: currently a PR is open for native support for the Sarvam models in vLLM ([link](https://github.com/vllm-project/vllm/pull/33942)). Therefore, we have 2 options here.
+
+#### Option 1: install from source (hard)
+
+* Use the custom fork here: [link](https://github.com/rahul-sarvam/vllm)
+* Follow the instructions here to install from source: [link](https://docs.vllm.ai/en/latest/getting_started/installation/gpu/index.html#build-wheel-from-source)
+
+#### Option 2: hot-patch (easy)
+
+* Run [hotpatch_vllm.py](./hotpatch_vllm.py)
+* This will do the following:
+  * install vllm=0.15.0
+  * add 2 model entries to `registry.py`
+  * download the model executors for `sarvam-105b` and `sarvam-30b`
+
+Once this is done, you can run vLLM as usual
+
+```python
+from vllm import LLM, SamplingParams
+from transformers import AutoTokenizer
+
+model_path = "sarvamai/sarvam-30b"
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+llm = LLM(model=model_path, 
+            trust_remote_code=True, 
+            max_model_len=2048, 
+            tensor_parallel_size=8, 
+            max_num_seqs=16,
+        )
+sampling_params = SamplingParams(
+                    temperature=0.8, 
+                    max_tokens=2048, 
+                    repetition_penalty=1.0,
+                    spaces_between_special_tokens=True
+                )
+
+prompts = [
+    "Who wrote The Picture of Dorian Gray?",
+]
+
+outputs = llm.generate([
+    tokenizer.apply_chat_template([
+        {"role": "user", "content": prompt}],
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=True)
+        for prompt in prompts], 
+    sampling_params)
+for p, o in zip(prompts, outputs):
+    print("Prompt: ", p)
+    print("Generated text: ", o.outputs[0].text)
+    print("=" * 100)
+```
+</details>
+
+### Footnote
+
+* **General settings**: All benchmarks are evaluated with a maximum context length of 65,536 tokens.
+* **Reasoning & Math benchmarks** (Math500, MMLU, MMLU Pro, GPQA Diamond, AIME 25, Beyond AIME, HMMT, HumanEval, MBPP): Evaluated with `temperature=1.0, top_p=1.0, max_new_tokens=65536`.
+* **Coding & Knowledge benchmarks** (Live Code Bench v6, Arena Hard v2, IF Eval):
+  Evaluated with `temperature=1.0, top_p=1.0, max_new_tokens=65536`.
+* **Writing Bench**:
+  Responses generated using official Writing-Bench parameters:
+  `temperature=0.7, top_p=0.8, top_k=20, max_length=16000`.
+  Scoring performed using the official Writing-Bench critic model with:
+  `temperature=1.0, top_p=0.95, max_length=2048`.
+* **Agentic benchmarks** (BrowseComp, SWE Bench Verified, τ² Bench): Evaluated with `temperature=0.5, top_p=1.0, max_new_tokens=32768`.
+
+## Citation
+```
+@misc{sarvam_sovereign_models,
+  title        = {Introducing Sarvam's Sovereign Models},
+  author       = {{Sarvam Foundation Models Team}},
+  year         = {2026},
+  howpublished = {\url{https://www.sarvam.ai/blogs/sarvam-30b-105b}},
+  note         = {Accessed: 2026-03-03}
+}
+```

chat_template.jinja

@@ -0,0 +1,97 @@
+{{- '[@BOS@]\n' }}
+{%- if tools -%}
+<|start_of_turn|><|tool_declare|>
+<tools>
+{% for tool in tools %}
+{{ tool | tojson(ensure_ascii=False) }}
+{% endfor %}
+</tools>
+{{- '<|end_of_turn|>\n' }}{%- endif -%}
+{%- macro visible_text(content) -%}
+    {%- if content is string -%}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping -%}
+        {%- for item in content -%}
+            {%- if item is mapping and item.type == 'text' -%}
+                {{- item.text }}
+            {%- elif item is string -%}
+                {{- item }}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- elif content is none -%}
+        {{- '' }}
+    {%- else -%}
+        {{- content }}
+    {%- endif -%}
+{%- endmacro -%}
+{%- set ns = namespace(last_user_index=-1) %}
+{%- for m in messages %}
+    {%- if m.role == 'user' %}
+        {% set ns.last_user_index = loop.index0 -%}
+    {%- endif %}
+{%- endfor %}
+{% for m in messages %}
+{%- if m.role == 'user' -%}<|start_of_turn|><|user|>
+{{ visible_text(m.content) }}
+{{- '<|nothink|>' if (enable_thinking is defined and not enable_thinking and not visible_text(m.content).endswith("<|nothink|>")) else '' -}}
+{{- '<|end_of_turn|>\n' }}
+{%- elif m.role == 'assistant' -%}
+{{- '<|start_of_turn|><|assistant|>\n' }}
+{%- set reasoning_content = '' %}
+{%- set content = visible_text(m.content) %}
+{%- if m.reasoning_content is string %}
+    {%- set reasoning_content = m.reasoning_content %}
+{%- else %}
+    {%- if '</think>' in content %}
+        {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+        {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+    {%- endif %}
+{%- endif %}
+{%- if loop.index0 > ns.last_user_index and reasoning_content -%}
+{{ '<think>' + reasoning_content.strip() +  '</think>'}}
+{%- else -%}
+{{ '<think></think>' }}
+{%- endif -%}
+{%- if content.strip() -%}
+{{ '\n' + content.strip() }}
+{%- endif -%}
+{% if m.tool_calls %}
+{% for tc in m.tool_calls %}
+{%- if tc.function %}
+    {%- set tc = tc.function %}
+{%- endif %}
+{{ '\n<tool_call>' + tc.name }}
+{% set _args = tc.arguments %}
+{% for k, v in _args.items() %}
+<arg_key>{{ k }}</arg_key>
+<arg_value>{{ v | tojson(ensure_ascii=False) if v is not string else v }}</arg_value>
+{% endfor %}
+</tool_call>{% endfor %}
+{% endif %}
+{{- '<|end_of_turn|>\n' }}
+{%- elif m.role == 'tool' -%}
+{%- if m.content is string -%}
+{%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+    {{- '<|start_of_turn|><|observation|>' }}
+{%- endif %}
+{{- '\n<tool_response>\n' }}
+{{- m.content }}
+{{- '\n</tool_response>' }}
+{%- else -%}
+<|start_of_turn|><|observation|>{% for tr in m.content %}
+
+<tool_response>
+{{ tr.output if tr.output is defined else tr }}
+</tool_response>{% endfor -%}
+{% endif -%}
+{%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+{{- '<|end_of_turn|>\n' }}{%- endif -%}
+{%- elif m.role == 'system' -%}
+<|start_of_turn|><|system|>
+{{ visible_text(m.content) }}
+{{- '<|end_of_turn|>\n' }}
+{%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{- '<|start_of_turn|><|assistant|>\n' }}
+{%- endif -%}

config.json

@@ -0,0 +1,53 @@
+{
+  "architectures": [
+    "SarvamMoEForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "attn_implementation": null,
+  "auto_map": {
+    "AutoConfig": "configuration_sarvam_moe.SarvamMoEConfig",
+    "AutoModel": "modeling_sarvam_moe.SarvamMoEModel",
+    "AutoModelForCausalLM": "modeling_sarvam_moe.SarvamMoEForCausalLM"
+  },
+  "dtype": "float32",
+  "embedding_dropout": 0.0,
+  "eos_token_id": 1,
+  "first_k_dense_replace": 1,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.006,
+  "intermediate_size": 8192,
+  "max_position_embeddings": 131072,
+  "max_window_layers": 19,
+  "model_type": "sarvam_moe",
+  "moe_intermediate_size": 1024,
+  "moe_router_enable_expert_bias": true,
+  "moe_shared_expert_intermediate_size": 1024,
+  "n_group": 1,
+  "norm_topk_prob": true,
+  "num_attention_heads": 64,
+  "num_experts": 128,
+  "num_experts_per_tok": 6,
+  "num_hidden_layers": 19,
+  "num_key_value_heads": 4,
+  "num_shared_experts": 1,
+  "output_dropout": 0.0,
+  "output_router_logits": false,
+  "pad_token_id": 0,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 8000000,
+  "routed_scaling_factor": 2.5,
+  "router_dtype": "fp32",
+  "score_function": "sigmoid",
+  "tie_word_embeddings": false,
+  "topk_group": 1,
+  "transformers_version": "4.57.2",
+  "use_bias": false,
+  "use_cache": true,
+  "use_qk_norm": true,
+  "use_qkv_bias": false,
+  "use_rmsnorm": true,
+  "vocab_size": 262144
+}

configuration_sarvam_moe.py

@@ -0,0 +1,103 @@
+from transformers.configuration_utils import PretrainedConfig
+
+
+class SarvamMoEConfig(PretrainedConfig):
+    model_type = "sarvam_moe"
+    def __init__(
+        self,
+        vocab_size=262144,
+        hidden_size=4096,
+        intermediate_size=8192,
+        num_hidden_layers=19,
+        num_attention_heads=16,
+        num_key_value_heads=4,
+        hidden_act="silu",
+        use_qkv_bias=False,
+        use_bias=False,
+        rms_norm_eps=1e-06,
+        tie_word_embeddings=False,
+        embedding_dropout=0.0,
+        attention_dropout=0.0,
+        output_dropout=0.0,
+        initializer_range=0.006,
+        max_position_embeddings=4096,
+        rope_theta=10000.0,
+        use_cache=True,
+        max_window_layers=19,
+        rope_scaling=None,
+        pad_token_id=0,
+        eos_token_id=1,
+        num_experts=128,
+        num_shared_experts=1,
+        num_experts_per_tok=6,
+        n_group=1,
+        topk_group=1,
+        moe_intermediate_size=1024,
+        first_k_dense_replace=1,
+        head_dim=256,
+        output_router_logits=False,
+        use_qk_norm=True,
+        moe_router_enable_expert_bias=True,
+        routed_scaling_factor=2.5,
+        attn_implementation: str = "eager",
+        **kwargs,
+    ):
+        self.num_hidden_layers = num_hidden_layers
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.use_qkv_bias = use_qkv_bias
+        self.use_bias = use_bias
+        self.rms_norm_eps = rms_norm_eps
+        self.embedding_dropout = embedding_dropout
+        self.attention_dropout = attention_dropout
+        self.output_dropout = output_dropout
+        self.initializer_range = initializer_range
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.use_cache = use_cache
+        self.max_window_layers = max_window_layers
+        self.head_dim = head_dim or hidden_size // num_attention_heads
+        self.rope_scaling = rope_scaling
+        self.use_qk_norm = use_qk_norm
+        self.moe_router_enable_expert_bias = moe_router_enable_expert_bias
+        self.routed_scaling_factor = routed_scaling_factor
+        self.num_experts = num_experts
+        self.num_shared_experts = num_shared_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.moe_intermediate_size = moe_intermediate_size
+        self.first_k_dense_replace = first_k_dense_replace
+        self.output_router_logits = output_router_logits
+        self.attn_implementation = attn_implementation
+        self._attn_implementation = attn_implementation
+
+        self.base_model_tp_plan = {
+            "layers.*.attention.query_key_value": "colwise",
+            "layers.*.attention.dense": "rowwise",
+            "layers.*.mlp.gate_proj": "colwise",
+            "layers.*.mlp.up_proj": "colwise",
+            "layers.*.mlp.down_proj": "rowwise",
+            "layers.*.mlp.experts.*.gate_proj": "colwise",
+            "layers.*.mlp.experts.*.up_proj": "colwise",
+            "layers.*.mlp.experts.*.down_proj": "rowwise",
+            "layers.*.mlp.shared_experts.gate_proj": "colwise",
+            "layers.*.mlp.shared_experts.up_proj": "colwise",
+            "layers.*.mlp.shared_experts.down_proj": "rowwise",
+        }
+        self.base_model_pp_plan = {
+            "word_embeddings": (["input_ids"], ["inputs_embeds"]),
+            "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+            "norm": (["hidden_states"], ["hidden_states"]),
+        }
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 26,
+  "pad_token_id": 0,
+  "transformers_version": "4.57.2"
+}

hotpatch_vllm.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import sys
+import subprocess
+from pathlib import Path
+from urllib.request import urlopen, Request
+
+
+HF_BLOB_URL = "https://huggingface.co/sarvamai/sarvam-30b/blob/main/sarvam.py"
+
+NEW_LINES = [
+    '    "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"),\n',
+    '    "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"),\n',
+]
+
+
+def run(cmd: list[str]) -> None:
+    print(f"+ {' '.join(cmd)}")
+    subprocess.check_call(cmd)
+
+
+def pip_install_vllm() -> None:
+    run([sys.executable, "-m", "pip", "install", "vllm==0.15.0"])
+
+
+def find_vllm_dir() -> Path:
+    import vllm  # type: ignore
+
+    vllm_dir = Path(vllm.__file__).resolve().parent
+    print(f"Detected vLLM package dir: {vllm_dir}")
+    return vllm_dir
+
+
+def patch_text_generation_models(registry_path: Path) -> None:
+    if not registry_path.exists():
+        raise FileNotFoundError(f"registry.py not found at: {registry_path}")
+
+    text = registry_path.read_text(encoding="utf-8")
+    lines = text.splitlines(keepends=True)
+
+    # Idempotency: if both keys already present, do nothing
+    if (
+        any('"SarvamMoEForCausalLM"' in l for l in lines)
+        and any('"SarvamMLAForCausalLM"' in l for l in lines)
+    ):
+        print("registry.py already contains Sarvam entries. Skipping patch.")
+        return
+
+    # Find the start of the _TEXT_GENERATION_MODELS dict
+    start_idx = None
+    for i, line in enumerate(lines):
+        if line.strip() == "_TEXT_GENERATION_MODELS = {":
+            start_idx = i
+            break
+
+    if start_idx is None:
+        raise RuntimeError(
+            "Could not find '_TEXT_GENERATION_MODELS = {' in registry.py. "
+            "vLLM version/layout may differ."
+        )
+
+    # Find the matching closing brace for that dict using brace depth
+    depth = 0
+    end_idx = None
+    for j in range(start_idx, len(lines)):
+        depth += lines[j].count("{")
+        depth -= lines[j].count("}")
+        if j > start_idx and depth == 0:
+            end_idx = j
+            break
+
+    if end_idx is None:
+        raise RuntimeError("Failed to find end of _TEXT_GENERATION_MODELS dict (brace matching).")
+
+    # Insert new entries just before the closing brace line
+    insert_at = end_idx
+    lines[insert_at:insert_at] = NEW_LINES
+
+    registry_path.write_text("".join(lines), encoding="utf-8")
+    print(f"Patched _TEXT_GENERATION_MODELS in: {registry_path}")
+
+
+def download_sarvam_py(dst: Path) -> None:
+    # Use /raw/ to download file contents, not HTML
+    raw_url = HF_BLOB_URL.replace("/blob/", "/raw/")
+    print(f"Downloading sarvam.py from: {raw_url}")
+
+    req = Request(raw_url, headers={"User-Agent": "vllm-hotpatch-script"})
+    with urlopen(req) as resp:
+        data = resp.read()
+
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    dst.write_bytes(data)
+    print(f"Wrote: {dst}")
+
+
+def main() -> None:
+    pip_install_vllm()
+
+    vllm_dir = find_vllm_dir()
+    registry_path = vllm_dir / "model_executor" / "models" / "registry.py"
+    sarvam_path = vllm_dir / "model_executor" / "models" / "sarvam.py"
+
+    patch_text_generation_models(registry_path)
+    download_sarvam_py(sarvam_path)
+
+    print("\nDone.")
+    print(f"- Registry patched: {registry_path}")
+    print(f"- Sarvam module installed: {sarvam_path}")
+
+
+if __name__ == "__main__":
+    main()

model-00001-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:efeac3cd0e9eb5a276aea37056e3abfd7cd0583b3e4f83ec871bf6fb62ead7ed
+size 4999645904

model-00002-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8fd3a943ae19559d1e9e0f5c3e6289843404bcfd7a699caf2f70bcc97fda2ba2
+size 4999647000

model-00003-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ef5e541b4732507466cecfad0f32810f677127d4739cbd23b74c23acdf5bde4
+size 4993389224

model-00004-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e9928a30813e2de87e91b06b56a2ffe18f2e46049a434eb585b351a52e8c3dbb
+size 4993389224

model-00005-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:718de0adb892541452bd3f9abbe11f108789b58c41d5bc776dbc0963f715ee6d
+size 4993389248

model-00006-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a475d430246ce0f64689ed324f98283e27d6e469e35c8c3050a5b3a819b500c
+size 4999647008

model-00007-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a612e55a2ce14c6df0a369564099635883235e80d93783e965742e61b645eda4
+size 4993389224

model-00008-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:01f625b20f0716e7eaf2b2bb2bc5ffdcb590ffcd91e790dfdd5850a080c1174a
+size 4993389224

model-00009-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c900c904bf8e6c1c4269bcd5669039be7afc6b62ad55675a8baedce82657aad7
+size 4993389248

model-00010-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a17cf9a02bfab152109d0894d8299e247c690d59a5dac705415d9c148dab3e8a
+size 4999647016

model-00011-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69d7ef47df1b3ec6fcb6b0be099abdb8d832db13f3dad323ad49395b8fa38ccf
+size 4993389216

model-00012-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ea7b63b62612fe78d9dd2e4e144d4f5a31c84dc8882cb4ec63f9c082dd86cc9
+size 4993389224

model-00013-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c138ca20e37dc040af76b10232bc9dffe496da28aff969a42816044f1efc4d1f
+size 4993389280

model-00014-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a22c6fb81a8f357eaa545a99da1fed0df6452e91253013c4f7d24a974e6cc724
+size 4999647320

model-00015-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b36d14b6ec4bd108331d91ecd405bebbba3d857c4c21b87a6371972c3e828cd0
+size 4993389512

model-00016-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f35318a600e102dbcb06473c6c8d3f15ec880b56bd42c2026cd1f9227ccaeb1
+size 4993389528

model-00017-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:14515810ae889da143c4f17c273f992ce577e8f8545c426093c5d763f5ce4663
+size 4993389536

model-00018-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ae4825d8232c200053b0b571937d299f58a7563b43c4e927e44dbdf51ac4b3e9
+size 4999647328

model-00019-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06e433bb2c02ca76a1e6f702a57f89d4d25b3fab5b832a4dbdea4aa888409316
+size 4993389512

model-00020-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0c6c0d47e3005836553859a42e0699235cc4a856c6954fe6d7f07ad59a5b6083
+size 4993389520

model-00021-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b9a8df1a1b2f740927500ac34b30ccea5faa4158b99550f533a51ae7505fcb8a
+size 4993389536

model-00022-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b63c24103409fc251ac09b281d607a9ced0b7e25da5e61d8142e2204bb2e405
+size 4999647336

model-00023-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88dbc5536a7bfcc89bf1c7f060c58f7c072012129ba60f64a173f52f75aa6244
+size 4993389504

model-00024-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f65c99212588029b979610c0f4b94dd18b22788284e9cbfbcf9ec4b70681dbf6
+size 4993389528

model-00025-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e359c0b5b2096d800a6923279a3dbaa15389b6b118b8a0fb907f9c31dc34a6af
+size 4431365192

model-00026-of-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f39a62c958f34788e01d2c45715a75c2a1238286994ebfb57335775beb2b5bf
+size 4294967424

modeling_sarvam_moe.py

@@ -0,0 +1,1025 @@
+"""PyTorch Sarvam MoE model."""
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import (
+    AttentionMaskConverter,
+    _prepare_4d_attention_mask,
+    _prepare_4d_causal_attention_mask,
+    _prepare_4d_causal_attention_mask_for_sdpa,
+)
+from transformers.modeling_outputs import MoeModelOutputWithPast
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import PreTrainedModel
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
+from transformers.utils import (
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+    logging,
+)
+from transformers.generation.utils import GenerationMixin
+from dataclasses import dataclass
+from transformers.utils import ModelOutput
+
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
+
+from .configuration_sarvam_moe import SarvamMoEConfig
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "SarvamMoEConfig"
+
+
+@dataclass
+class SarvamMoECausalLMOutputWithPast(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+    z_loss: Optional[torch.FloatTensor] = None
+    aux_loss: Optional[torch.FloatTensor] = None
+    router_logits: Optional[tuple[torch.FloatTensor]] = None
+
+
+class SarvamMoEModelOutputWithPast(MoeModelOutputWithPast):
+    pass
+
+
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+    return indices, cu_seqlens, max_seqlen_in_batch
+
+
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
+    return _prepare_4d_attention_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
+
+
+def _make_causal_mask(
+    input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
+):
+    return AttentionMaskConverter._make_causal_mask(
+        input_ids_shape=input_ids_shape, dtype=dtype, device=device, past_key_values_length=past_key_values_length
+    )
+
+
+class SarvamMoERMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+ALL_LAYERNORM_LAYERS.append(SarvamMoERMSNorm)
+
+
+class SarvamMoERotaryEmbedding(nn.Module):
+    def __init__(self, config: SarvamMoEConfig, device=None):
+        super().__init__()
+        self.config = config
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+        rope_scaling = getattr(config, "rope_scaling", None)
+        if rope_scaling is None:
+            self.rope_type = "default"
+            inv_freq, self.attention_scaling = self.compute_default_rope_parameters(
+                config, device
+            )
+        else:
+            self.rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", "default"))
+            if self.rope_type == "default":
+                inv_freq, self.attention_scaling = self.compute_default_rope_parameters(
+                    config, device
+                )
+            else:
+                rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+                inv_freq, self.attention_scaling = rope_init_fn(config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: SarvamMoEConfig,
+        device: Optional[torch.device] = None,
+        seq_len: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, float]:
+        """
+        Default RoPE parameters (classic rotary embedding).
+
+        Mirrors HF's default implementation: use `rope_theta`, head_dim and
+        return (inv_freq, attention_scaling).
+        """
+        base = config.rope_theta
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        inv_freq = 1.0 / (
+            base
+            ** (
+                torch.arange(0, dim, 2, dtype=torch.int64, device=device)
+                .to(dtype=torch.float32)
+                / dim
+            )
+        )
+        attention_factor = 1.0
+        return inv_freq, attention_factor
+
+    @torch.no_grad()
+    @dynamic_rope_update
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    rotary_dim = cos.shape[-1]
+    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+    q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+    k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+    q_embed = torch.cat([q_embed, q_pass], dim=-1)
+    k_embed = torch.cat([k_embed, k_pass], dim=-1)
+    return q_embed, k_embed
+
+
+class SarvamMoEMLP(nn.Module):
+    def __init__(self, config: SarvamMoEConfig, intermediate_size: int):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class SarvamMoEGate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+        self.gating_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.empty((self.num_experts, self.gating_dim)))
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.score_function = config.score_function
+        # Ideally, we should register the expert_bias as a buffer, but vllm complains about it.
+        # self.register_buffer("expert_bias", torch.zeros((self.num_experts)))
+        self.expert_bias = nn.Parameter(
+            torch.zeros((self.num_experts)),
+            requires_grad=False,
+        )
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        import torch.nn.init as init
+
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+
+    def group_limited_topk(self, scores: torch.Tensor):
+        num_tokens, _ = scores.size()
+        group_scores = scores.view(num_tokens, self.n_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
+        group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+        group_mask = torch.zeros_like(group_scores)
+        group_mask.scatter_(1, group_idx, 1)
+        score_mask = (
+            group_mask.unsqueeze(-1)
+            .expand(num_tokens, self.n_group, self.num_experts // self.n_group)
+            .reshape(num_tokens, -1)
+        )
+        masked_scores = scores.masked_fill(~score_mask.bool(), float("-inf"))
+        probs, top_indices = torch.topk(masked_scores, k=self.top_k, dim=-1)
+        return probs, top_indices
+
+    def forward(self, hidden_states):
+        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+        logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32))
+        scores = torch.sigmoid(logits.float()).type_as(logits)
+        scores_for_routing = scores + self.expert_bias
+        _, topk_idx = self.group_limited_topk(scores_for_routing)
+        scores = torch.gather(scores, dim=1, index=topk_idx).type_as(logits)
+        topk_weight = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if self.top_k > 1 else scores
+        topk_weight = topk_weight * self.routed_scaling_factor
+        return topk_idx, topk_weight, logits
+
+
+class SarvamMoEExperts(nn.ModuleList):
+    def __init__(self, config: SarvamMoEConfig):
+        # one MLP per expert
+        experts = [
+            SarvamMoEMLP(config=config, intermediate_size=config.moe_intermediate_size)
+            for _ in range(config.num_experts)
+        ]
+        super().__init__(experts)
+        self.config = config
+        self.num_experts_per_tok = config.num_experts_per_tok
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        top_k_index: torch.LongTensor,
+        top_k_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        hidden_states: (tokens, hidden_size) or (batch * seq, hidden_size)
+        top_k_index:   (tokens, top_k)
+        top_k_weights: (tokens, top_k)
+        """
+        tokens, hidden_dim = hidden_states.shape
+        flat_topk_idx = top_k_index.view(-1)
+
+        if self.training:
+            # training path: same as your previous logic
+            x = hidden_states.repeat_interleave(self.num_experts_per_tok, dim=0)
+            y = torch.empty_like(x)
+            for i, expert in enumerate(self):
+                mask = flat_topk_idx == i
+                if mask.any():
+                    y[mask] = expert(x[mask])
+            y = (y.view(*top_k_weights.shape, -1) * top_k_weights.unsqueeze(-1)).sum(dim=1)
+            return y.to(hidden_states.dtype)
+
+        # inference path: previous moe_infer logic
+        num_experts = len(self)
+        cnts = top_k_index.new_zeros((tokens, num_experts))
+        cnts.scatter_(1, top_k_index, 1)
+        tokens_per_expert = cnts.sum(dim=0)
+
+        idxs = top_k_index.view(-1).argsort()
+        sorted_tokens = hidden_states[idxs // top_k_index.shape[1]]
+
+        tokens_per_expert = tokens_per_expert.cpu().numpy().tolist()
+        outputs = []
+        start_idx = 0
+        for i, num_tokens in enumerate(tokens_per_expert):
+            end_idx = start_idx + num_tokens
+            if num_tokens == 0:
+                continue
+            expert = self[i]
+            tokens_for_expert = sorted_tokens[start_idx:end_idx]
+            expert_out = expert(tokens_for_expert)
+            outputs.append(expert_out.to(hidden_states.device))
+            start_idx = end_idx
+
+        outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
+        new_x = torch.empty_like(outs)
+        new_x[idxs] = outs
+
+        final_out = (
+            new_x.view(*top_k_index.shape, -1)
+            .type(top_k_weights.dtype)
+            .mul_(top_k_weights.unsqueeze(dim=-1))
+            .sum(dim=1)
+            .type(new_x.dtype)
+        )
+        return final_out
+
+
+class SarvamMoESparseMoeBlock(nn.Module):
+    def __init__(self, config: SarvamMoEConfig):
+        super().__init__()
+        self.config = config
+        self.num_experts_per_tok = config.num_experts_per_tok
+
+        # use the new experts container
+        self.experts = SarvamMoEExperts(config)
+        self.gate = SarvamMoEGate(config)
+
+        if config.num_shared_experts is not None:
+            self.shared_experts = SarvamMoEMLP(
+                config=config,
+                intermediate_size=config.moe_intermediate_size * config.num_shared_experts,
+            )
+
+    # _setup_experts no longer needed
+
+    def forward(self, hidden_states):
+        identity = hidden_states
+        bsz, seq_len, h = hidden_states.shape
+
+        topk_idx, topk_weight, router_logits = self.gate(hidden_states)
+
+        # flatten batch+seq for experts
+        flat_hidden = hidden_states.view(-1, h)
+        flat_topk_idx = topk_idx.view(-1, topk_idx.shape[-1])
+        flat_topk_weight = topk_weight.view(-1, topk_weight.shape[-1])
+
+        y = self.experts(flat_hidden, flat_topk_idx, flat_topk_weight)
+        y = y.view(bsz, seq_len, h)
+
+        if self.config.num_shared_experts is not None:
+            y = y + self.shared_experts(identity)
+
+        # router logits shape: (bsz, seq_len, num_experts)
+        router_info = (
+            router_logits.view(bsz, seq_len, -1),
+            topk_idx.view(bsz, seq_len, -1),
+        )
+        return y, router_info
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class SarvamMoEAttention(nn.Module):
+    is_causal = True    # vLLM / Transformers backend critical flag
+    def __init__(self, config: SarvamMoEConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+        self.attention_dropout = config.attention_dropout
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.head_dim or self.hidden_size // self.num_heads
+        partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
+        self.rope_dim = int(self.head_dim * partial_rotary_factor)
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.query_key_value = nn.Linear(
+            self.hidden_size,
+            (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim,
+            bias=config.use_qkv_bias,
+        )
+        if self.config.use_qk_norm:
+            self.query_layernorm = SarvamMoERMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.key_layernorm = SarvamMoERMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.dense = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.use_bias)
+        self.scaling = self.head_dim**-0.5
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs,
+    ):
+        bsz, q_len, _ = hidden_states.size()
+        qkv = self.query_key_value(hidden_states)
+        qkv = qkv.view(
+            bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim
+        )
+        query_states, key_states, value_states = qkv.split(
+            [self.num_heads, self.num_key_value_heads, self.num_key_value_heads],
+            dim=-2,
+        )
+        query_states = query_states.transpose(1, 2).contiguous()
+        key_states = key_states.transpose(1, 2).contiguous()
+        value_states = value_states.transpose(1, 2).contiguous()
+        if self.config.use_qk_norm:
+            query_states = self.query_layernorm(query_states)
+            key_states = self.key_layernorm(key_states)
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+        if past_key_value is not None:
+            if self.layer_idx is None:
+                raise ValueError(
+                    "When using cache, SarvamMoEAttention must be initialized with layer_idx."
+                )
+            cache_kwargs = {"sin": sin, "cos": cos}
+            key_states, value_states = past_key_value.update(
+                key_states, value_states, self.layer_idx, cache_kwargs
+            )
+        # NOTE: vLLM will set config._attn_implementation = "vllm"
+        if self.config._attn_implementation == "vllm":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+            attn_output, attn_weights = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                scaling=self.scaling,
+                **kwargs,
+            )
+            # vLLM backend may return [B, L, hidden] or [B*L, hidden]
+            if attn_output.dim() == 4:
+                # [B, H, L, Dh] -> [B, L, hidden]
+                attn_output = attn_output.transpose(1, 2).contiguous()
+                attn_output = attn_output.view(bsz, q_len, -1)
+            elif attn_output.dim() == 3:
+                if attn_output.shape[0] != bsz or attn_output.shape[1] != q_len:
+                    raise ValueError(
+                        f"Unexpected vLLM attention output shape {attn_output.shape}, "
+                        f"expected (bsz={bsz}, q_len={q_len}, hidden=*)"
+                    )
+            elif attn_output.dim() == 2:
+                attn_output = attn_output.view(bsz, q_len, -1)
+            else:
+                raise ValueError(
+                    f"Unsupported vLLM attention output rank {attn_output.dim()} "
+                    f"with shape {attn_output.shape}"
+                )
+            attn_output = self.dense(attn_output)
+            if not output_attentions:
+                attn_weights = None
+            return attn_output, attn_weights, past_key_value
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+        kv_seq_len = key_states.shape[-2]
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+        attn_output = self.dense(attn_output)
+        if not output_attentions:
+            attn_weights = None
+        return attn_output, attn_weights, past_key_value
+
+
+class SarvamMoEFlashAttention2(SarvamMoEAttention):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        output_attentions = False
+        bsz, q_len, _ = hidden_states.size()
+        qkv = self.query_key_value(hidden_states)
+        qkv = qkv.view(bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim)
+        query_states, key_states, value_states = qkv.split(
+            [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2
+        )
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+        if self.config.use_qk_norm:
+            query_states = self.query_layernorm(query_states)
+            key_states = self.key_layernorm(key_states)
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+        dropout_rate = self.attention_dropout if self.training else 0.0
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            if hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            elif torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            else:
+                target_dtype = self.query_key_value.weight.dtype
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+        attn_output = self._flash_attention_forward(
+            query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
+        )
+        attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+        attn_output = self.dense(attn_output)
+        if not output_attentions:
+            attn_weights = None
+        return attn_output, attn_weights, past_key_value
+
+    def _flash_attention_forward(
+        self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
+    ):
+        if not self._flash_attn_uses_top_left_mask:
+            causal = self.is_causal
+        else:
+            causal = self.is_causal and query_length != 1
+        if attention_mask is not None:
+            batch_size = query_states.shape[0]
+            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+                query_states, key_states, value_states, attention_mask, query_length
+            )
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+            attn_output_unpad = flash_attn_varlen_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q,
+                max_seqlen_k=max_seqlen_in_batch_k,
+                dropout_p=dropout,
+                softmax_scale=softmax_scale,
+                causal=causal,
+            )
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+        else:
+            attn_output = flash_attn_func(
+                query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
+            )
+        return attn_output
+
+    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+        key_layer = index_first_axis(
+            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        value_layer = index_first_axis(
+            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
+            )
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_in_batch_q = max_seqlen_in_batch_k
+            indices_q = indices_k
+        elif query_length == 1:
+            max_seqlen_in_batch_q = 1
+            cu_seqlens_q = torch.arange(
+                batch_size + 1, dtype=torch.int32, device=query_layer.device
+            )
+            indices_q = cu_seqlens_q[:-1]
+            query_layer = query_layer.squeeze(1)
+        else:
+            attention_mask = attention_mask[:, -query_length:]
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
+        return (
+            query_layer,
+            key_layer,
+            value_layer,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        )
+
+
+class SarvamMoESdpaAttention(SarvamMoEAttention):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if output_attentions:
+            return super().forward(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                **kwargs,
+            )
+        bsz, q_len, _ = hidden_states.size()
+        qkv = self.query_key_value(hidden_states)
+        qkv = qkv.view(bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim)
+        query_states, key_states, value_states = qkv.split(
+            [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2
+        )
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+        if self.config.use_qk_norm:
+            query_states = self.query_layernorm(query_states)
+            key_states = self.key_layernorm(key_states)
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+        if attention_mask is not None:
+            kv_seq_len = key_states.shape[-2]
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+        if query_states.device.type == "cuda" and attention_mask is not None:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            is_causal=self.is_causal and attention_mask is None and q_len > 1,
+        )
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+        attn_output = self.dense(attn_output)
+        return attn_output, None, past_key_value
+
+
+ATTENTION_CLASSES = {
+    "eager": SarvamMoEAttention,
+    "flash_attention_2": SarvamMoEFlashAttention2,
+    "sdpa": SarvamMoESdpaAttention,
+    "vllm": SarvamMoEAttention,
+}
+
+
+class SarvamMoEDecoderLayer(nn.Module):
+    def __init__(self, config: SarvamMoEConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.attention = ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
+        self.mlp = (
+            SarvamMoESparseMoeBlock(config)
+            if (config.num_experts is not None and layer_idx >= config.first_k_dense_replace)
+            else SarvamMoEMLP(config=config, intermediate_size=config.intermediate_size)
+        )
+        self.input_layernorm = SarvamMoERMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = SarvamMoERMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        output_router_logits: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, self_attn_weights, present_key_value = self.attention(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            position_embeddings=position_embeddings,
+            use_cache=use_cache,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        if isinstance(hidden_states, tuple):
+            hidden_states, router_logits = hidden_states
+        else:
+            router_logits = None
+        hidden_states = residual + hidden_states.to(residual.device)
+        outputs = (hidden_states,)
+        if output_attentions:
+            outputs += (self_attn_weights,)
+        if use_cache:
+            outputs += (present_key_value,)
+        if output_router_logits:
+            outputs += (router_logits,)
+        return outputs
+
+class SarvamMoEPreTrainedModel(PreTrainedModel):
+    config_class = SarvamMoEConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["SarvamMoEDecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+
+class SarvamMoEModel(SarvamMoEPreTrainedModel):
+    _supports_attention_backend = True
+    def __init__(self, config: SarvamMoEConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = []
+        for layer_idx in range(config.num_hidden_layers):
+            self.layers.append(SarvamMoEDecoderLayer(config, layer_idx))
+        self.layers = nn.ModuleList(self.layers)
+        self._use_sdpa = config._attn_implementation == "sdpa"
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+        self.norm = SarvamMoERMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = SarvamMoERotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.word_embeddings = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_router_logits: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, SarvamMoEModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape[:2]
+        elif inputs_embeds is not None:
+            batch_size, seq_length = inputs_embeds.shape[:2]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`transformers."
+                )
+                use_cache = False
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+        if position_ids is None:
+            position_ids = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+            position_ids = position_ids.unsqueeze(0)
+        if self._use_flash_attention_2:
+            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+        elif self._use_sdpa and not output_attentions:
+            attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+                attention_mask,
+                (batch_size, seq_length),
+                inputs_embeds,
+                past_seen_tokens,
+            )
+        else:
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask, (batch_size, seq_length), inputs_embeds, past_seen_tokens
+            )
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        all_router_logits = () if output_router_logits else None
+        next_decoder_cache = None
+        layers = self.layers
+        for decoder_layer in layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    output_router_logits,
+                    use_cache,
+                    position_embeddings,
+                    **kwargs,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    output_router_logits=output_router_logits,
+                    use_cache=use_cache,
+                    position_embeddings=position_embeddings,
+                    **kwargs,
+                )
+            hidden_states = layer_outputs[0]
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+            if output_router_logits and layer_outputs[-1] is not None:
+                all_router_logits += (layer_outputs[-1],)
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+        next_cache = None
+        if use_cache:
+            next_cache = next_decoder_cache
+        if not return_dict:
+            return tuple(
+                v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits] if v is not None
+            )
+        return SarvamMoEModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+            router_logits=all_router_logits,
+        )
+
+
+class SarvamMoEForCausalLM(SarvamMoEPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: SarvamMoEConfig):
+        super().__init__(config)
+        self.model = SarvamMoEModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.model.word_embeddings = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_router_logits: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, SarvamMoEModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_router_logits=output_router_logits,
+            return_dict=return_dict,
+            **kwargs,
+        )
+        loss = None
+        aux_loss = None
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+        logits = logits.float()
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.config.vocab_size, **kwargs)
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            if output_router_logits:
+                output = (aux_loss,) + output
+            return (loss,) + output if loss is not None else output
+        return SarvamMoECausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            aux_loss=aux_loss,
+            router_logits=outputs.router_logits,
+        )

sarvam.py

@@ -0,0 +1,788 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Copyright 2026 Sarvam AI team. All rights reserved.
+#
+# This code is based on Llama, Deepseek, and Bailing MoE implementations
+# in this library. It has been modified from its original forms to
+# accommodate Sarvam's MoE architectures.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable, Iterator
+from itertools import islice
+
+import torch
+from torch import nn
+
+from vllm.config import CacheConfig, ParallelConfig, VllmConfig
+from vllm.distributed import (
+    get_pp_group,
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.sequence import IntermediateTensors
+
+from .bailing_moe import BailingMoeForCausalLM
+from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
+from .utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    is_pp_missing_parameter,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_prefix,
+)
+
+
+def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
+    if scale <= 1:
+        return 1.0
+    return 0.1 * mscale * math.log(scale) + 1.0
+
+
+def _is_gate_expert_bias_name(name: str) -> bool:
+    return name.endswith(".mlp.gate.e_score_correction_bias") or name.endswith(
+        ".gate.e_score_correction_bias"
+    )
+
+
+def _zero_mean_tensor(t: torch.Tensor) -> torch.Tensor:
+    if t.numel() == 0:
+        return t
+    return t - t.mean()
+
+
+def _normalized_weights(
+    weights: Iterable[tuple[str, torch.Tensor]],
+) -> Iterator[tuple[str, torch.Tensor]]:
+    for name, w in weights:
+        if _is_gate_expert_bias_name(name):
+            yield name, _zero_mean_tensor(w)
+        else:
+            yield name, w
+
+
+class SarvamMLAAttention(nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        config,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.qk_nope_head_dim = config.qk_nope_head_dim
+        self.qk_rope_head_dim = config.qk_rope_head_dim
+        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        self.v_head_dim = config.v_head_dim
+
+        self.q_lora_rank = getattr(config, "q_lora_rank", None)
+        self.kv_lora_rank = config.kv_lora_rank
+
+        self.total_num_heads = config.num_attention_heads
+        tp_size = get_tensor_model_parallel_world_size()
+        assert self.total_num_heads % tp_size == 0
+        self.num_local_heads = self.total_num_heads // tp_size
+
+        self.scaling = self.qk_head_dim**-0.5
+        self.max_position_embeddings = config.max_position_embeddings
+
+        if self.q_lora_rank is not None:
+            self.q_a_proj = ReplicatedLinear(
+                self.hidden_size,
+                self.q_lora_rank,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.q_a_proj",
+            )
+            self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
+            self.q_b_proj = ColumnParallelLinear(
+                self.q_lora_rank,
+                self.total_num_heads * self.qk_head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.q_b_proj",
+            )
+            self.q_proj = None  # type: ignore
+        else:
+            self.q_proj = ColumnParallelLinear(
+                self.hidden_size,
+                self.total_num_heads * self.qk_head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.q_proj",
+            )
+            self.q_a_proj = None  # type: ignore
+            self.q_a_layernorm = None  # type: ignore
+            self.q_b_proj = None  # type: ignore
+
+        # KV latent (MQA-style) A-proj
+        self.kv_a_proj_with_mqa = ReplicatedLinear(
+            self.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.kv_a_proj_with_mqa",
+        )
+        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
+
+        # KV B-proj produces per-head K_nope and V
+        self.kv_b_proj = ColumnParallelLinear(
+            self.kv_lora_rank,
+            self.total_num_heads * (self.qk_nope_head_dim + self.v_head_dim),
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.kv_b_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.rotary_emb = get_rope(
+            self.qk_rope_head_dim,
+            # rotary_dim=self.qk_rope_head_dim,
+            max_position=config.max_position_embeddings,
+            rope_parameters=config.rope_parameters,
+            is_neox_style=False,
+        )
+
+        if config.rope_parameters.get("rope_type", None) == "deepseek_yarn":
+            mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
+            scaling_factor = config.rope_parameters["factor"]
+            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
+            self.scaling = self.scaling * mscale * mscale
+
+        mla_modules = MLAModules(
+            kv_a_layernorm=self.kv_a_layernorm,
+            kv_b_proj=self.kv_b_proj,
+            rotary_emb=self.rotary_emb,
+            o_proj=self.o_proj,
+            fused_qkv_a_proj=None,
+            kv_a_proj_with_mqa=self.kv_a_proj_with_mqa,
+            q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None,
+            q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
+            q_proj=self.q_proj if self.q_lora_rank is None else None,
+            indexer=None,
+            indexer_rotary_emb=None,
+            is_sparse=False,
+            topk_indices_buffer=None,
+        )
+
+        self.mla_attn = MultiHeadLatentAttentionWrapper(
+            self.hidden_size,
+            self.num_local_heads,
+            self.scaling,
+            self.qk_nope_head_dim,
+            self.qk_rope_head_dim,
+            self.v_head_dim,
+            self.q_lora_rank,
+            self.kv_lora_rank,
+            mla_modules,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.mla_attn(positions, hidden_states, llama_4_scaling=None)
+
+
+class SarvamMLAMLP(nn.Module):
+    def __init__(
+        self,
+        intermediate_size: int,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        reduce_results: bool = True,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.gate_up_proj = MergedColumnParallelLinear(
+            config.hidden_size,
+            [intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            config.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            reduce_results=reduce_results,
+            prefix=f"{prefix}.down_proj",
+        )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class SarvamMLAMoE(nn.Module):
+    def __init__(
+        self,
+        config,
+        parallel_config: ParallelConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_rank = get_tensor_model_parallel_rank()
+        self.hidden_size = config.hidden_size
+
+        self.num_experts = config.num_experts
+        self.top_k = config.num_experts_per_tok
+        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 2.5)
+
+        self.n_group = getattr(config, "n_group", None)
+        self.topk_group = getattr(config, "topk_group", None)
+        self.use_grouped_topk = self.n_group is not None and self.topk_group is not None
+
+        self.norm_expert_prob = getattr(config, "norm_topk_prob", True)
+
+        router_dtype_cfg = getattr(config, "router_dtype", "fp32")
+        if router_dtype_cfg is None:
+            self.router_dtype = None
+        elif router_dtype_cfg == "fp32":
+            self.router_dtype = torch.float32
+        else:
+            self.router_dtype = torch.bfloat16
+
+        self.gate = nn.Linear(
+            self.hidden_size,
+            self.num_experts,
+            bias=False,
+            dtype=self.router_dtype,
+        )
+
+        if getattr(config, "moe_router_enable_expert_bias", True):
+            self.gate.e_score_correction_bias = nn.Parameter(
+                torch.empty(
+                    (self.num_experts,),
+                    dtype=torch.float32,
+                )
+            )
+        else:
+            self.gate.e_score_correction_bias = None
+
+        self.score_function = getattr(config, "score_function", "sigmoid")
+        self.num_shared_experts = getattr(config, "num_shared_experts", 1)
+        if self.num_shared_experts > 0:
+            if hasattr(config, "moe_shared_expert_intermediate_size"):
+                shared_int = config.moe_shared_expert_intermediate_size
+            else:
+                shared_int = config.moe_intermediate_size
+            shared_int *= self.num_shared_experts
+            self.shared_experts = SarvamMLAMLP(
+                intermediate_size=shared_int,
+                config=config,
+                quant_config=quant_config,
+                reduce_results=False,
+                prefix=f"{prefix}.shared_experts",
+            )
+        else:
+            self.shared_experts = None
+
+        self.experts = SharedFusedMoE(
+            shared_experts=self.shared_experts,
+            num_experts=self.num_experts,
+            top_k=self.top_k,
+            hidden_size=self.hidden_size,
+            intermediate_size=config.moe_intermediate_size,
+            reduce_results=False,
+            renormalize=self.norm_expert_prob,
+            quant_config=quant_config,
+            prefix=f"{prefix}.experts",
+            scoring_func=self.score_function,
+            e_score_correction_bias=self.gate.e_score_correction_bias,
+            num_expert_group=self.n_group,
+            topk_group=self.topk_group,
+            use_grouped_topk=self.use_grouped_topk,
+            routed_scaling_factor=self.routed_scaling_factor,
+        )
+
+    def maybe_get_fused_moe(self) -> SharedFusedMoE:
+        return self.experts
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        router_logits = self.gate(
+            hidden_states.to(self.router_dtype)
+            if self.router_dtype is not None
+            else hidden_states
+        )
+        router_logits = router_logits.to(hidden_states.dtype)
+        final_hidden = self.experts(
+            hidden_states=hidden_states,
+            router_logits=router_logits,
+        )
+
+        if self.shared_experts is not None:
+            shared_output, expert_output = final_hidden
+        else:
+            shared_output, expert_output = None, final_hidden
+
+        # expert_output *= self.routed_scaling_factor
+
+        if shared_output is not None:
+            expert_output = expert_output + shared_output
+
+        if self.tp_size > 1:
+            expert_output = self.experts.maybe_all_reduce_tensor_model_parallel(
+                expert_output
+            )
+
+        return expert_output.view(num_tokens, hidden_dim)
+
+
+class SarvamMLABlock(nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+        parallel_config = vllm_config.parallel_config
+        layer_idx = int(prefix.split(".")[-1])
+        hidden_size = config.hidden_size
+        dense_intermediate = getattr(config, "intermediate_size", 16384)
+
+        self.input_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
+        self.self_attn = SarvamMLAAttention(
+            vllm_config=vllm_config,
+            config=config,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.post_attention_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
+        use_moe = hasattr(config, "num_experts") and config.num_experts is not None
+        first_k_dense = getattr(config, "first_k_dense_replace", 1)
+        moe_layer_freq = getattr(config, "moe_layer_freq", 1)
+        if use_moe:
+            is_moe_layer = layer_idx >= first_k_dense and (
+                (layer_idx - first_k_dense) % moe_layer_freq == 0
+            )
+        else:
+            is_moe_layer = False
+
+        if is_moe_layer:
+            self.mlp = SarvamMLAMoE(
+                config=config,
+                parallel_config=parallel_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+        else:
+            self.mlp = SarvamMLAMLP(
+                intermediate_size=dense_intermediate,
+                config=config,
+                quant_config=quant_config,
+                reduce_results=True,
+                prefix=f"{prefix}.mlp",
+            )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        positions: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+class SarvamMLAModel(nn.Module):
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.embed_dim = config.hidden_size
+        self.tie_word_embeddings = getattr(config, "tie_word_embeddings", False)
+        if get_pp_group().is_first_rank or (
+            self.tie_word_embeddings and get_pp_group().is_last_rank
+        ):
+            self.embed_tokens = VocabParallelEmbedding(
+                self.vocab_size,
+                self.embed_dim,
+                quant_config=quant_config,
+                prefix=f"{prefix}.embed_tokens",
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+
+        self.embedding_dropout = torch.nn.Dropout(
+            getattr(config, "embedding_dropout", 0.0)
+        )
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: SarvamMLABlock(
+                vllm_config=vllm_config,
+                prefix=prefix,
+            ),
+            prefix=f"{prefix}.layers",
+        )
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], config.hidden_size
+        )
+        if get_pp_group().is_last_rank:
+            self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.embed_input_ids(input_ids)
+            hidden_states = self.embedding_dropout(hidden_states)
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+
+        for layer in islice(self.layers, self.start_layer, self.end_layer):
+            hidden_states, residual = layer(
+                hidden_states,
+                positions,
+                residual,
+            )
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors(
+                {"hidden_states": hidden_states, "residual": residual}
+            )
+        if residual is None:
+            hidden_states = self.norm(hidden_states)
+        else:
+            hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states
+
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        return SharedFusedMoE.make_expert_params_mapping(
+            self,
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=self.config.num_experts,
+        )
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        """Load weights with stacked gate+up and MoE expert remapping."""
+        weights = _normalized_weights(weights)
+        stacked_params_mapping = [
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        loaded_params: set[str] = set()
+        expert_params_mapping = self.get_expert_mapping()
+
+        for name, loaded_weight in weights:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                if "mlp.experts" in name:
+                    continue
+                new_name = name.replace(weight_name, param_name)
+                if new_name.endswith(".bias") and new_name not in params_dict:
+                    continue
+                if new_name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(new_name, self):
+                    continue
+
+                param = params_dict[new_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight, shard_id)
+                loaded_params.add(new_name)
+                break
+            else:
+                mapped = False
+                for (
+                    param_name,
+                    weight_name,
+                    expert_id,
+                    shard_id,
+                ) in expert_params_mapping:
+                    if weight_name not in name:
+                        continue
+
+                    new_name = name.replace(weight_name, param_name)
+                    if is_pp_missing_parameter(new_name, self):
+                        continue
+                    if new_name not in params_dict:
+                        continue
+
+                    param = params_dict[new_name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(
+                        param,
+                        loaded_weight,
+                        name,
+                        shard_id=shard_id,
+                        expert_id=expert_id,
+                    )
+                    loaded_params.add(new_name)
+                    mapped = True
+                    break
+
+                if mapped:
+                    continue
+
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                if name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(name, self):
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                loaded_params.add(name)
+
+        return loaded_params
+
+
+class SarvamMixtureOfExperts(MixtureOfExperts):
+    def extract_moe_parameters(self, example_moe: SarvamMLAMoE | None) -> None:
+        if example_moe is None:
+            raise RuntimeError("No SarvamMLAMoE layer found in model.layers.")
+
+        self.num_logical_experts = example_moe.num_experts
+        self.num_routed_experts = example_moe.num_experts  # routed pool size
+        self.num_shared_experts = getattr(example_moe.config, "num_shared_experts", 1)
+
+        self.num_physical_experts = self.num_logical_experts
+        self.num_local_physical_experts = self.num_logical_experts
+        self.num_redundant_experts = 0
+
+    def update_physical_experts_metadata(
+        self,
+        num_physical_experts: int,
+        num_local_physical_experts: int,
+    ) -> None:
+        self.num_physical_experts = num_physical_experts
+        self.num_local_physical_experts = num_local_physical_experts
+        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
+
+        for moe in self.moe_mlp_layers:
+            moe.n_physical_experts = num_physical_experts
+            moe.n_local_physical_experts = num_local_physical_experts
+            moe.n_redundant_experts = self.num_redundant_experts
+
+            fused = moe.experts
+            if hasattr(fused, "n_local_physical_experts"):
+                fused.n_local_physical_experts = num_local_physical_experts
+            if hasattr(fused, "n_physical_experts"):
+                fused.n_physical_experts = num_physical_experts
+            if hasattr(fused, "n_redundant_experts"):
+                fused.n_redundant_experts = self.num_redundant_experts
+            if hasattr(fused, "update_expert_map"):
+                fused.update_expert_map()
+
+    def set_eplb_state(self, eplb_state) -> None:
+        self.eplb_state = eplb_state
+        for moe in self.moe_layers:
+            if hasattr(moe, "set_eplb_state"):
+                moe.set_eplb_state(eplb_state)
+
+
+class SarvamMLAForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SarvamMixtureOfExperts):
+    packed_modules_mapping = {
+        "q_proj": ["q_proj"],
+        "q_a_proj": ["q_a_proj"],
+        "q_b_proj": ["q_b_proj"],
+        "kv_a_proj_with_mqa": ["kv_a_proj_with_mqa"],
+        "kv_b_proj": ["kv_b_proj"],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        self.config = config
+        self.quant_config = quant_config
+
+        self.model = SarvamMLAModel(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "model"),
+        )
+
+        self.tie_word_embeddings = getattr(config, "tie_word_embeddings", False)
+        if get_pp_group().is_last_rank:
+            if self.tie_word_embeddings:
+                self.lm_head = self.model.embed_tokens
+            else:
+                self.lm_head = ParallelLMHead(
+                    config.vocab_size,
+                    config.hidden_size,
+                    quant_config=quant_config,
+                    prefix=maybe_prefix(prefix, "lm_head"),
+                )
+            self.logits_processor = LogitsProcessor(config.vocab_size)
+        else:
+            self.lm_head = PPMissingLayer()
+            self.logits_processor = None  # type: ignore
+
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+
+        self.expert_weights = []
+        self.num_moe_layers = 0
+
+        self.moe_layers = []
+        self.moe_mlp_layers = []
+
+        example_moe = None
+        for layer in self.model.layers:
+            if isinstance(layer, PPMissingLayer):
+                continue
+            if isinstance(layer.mlp, SarvamMLAMoE):
+                example_moe = layer.mlp
+                self.moe_mlp_layers.append(layer.mlp)
+                self.moe_layers.append(layer.mlp.experts)
+                self.num_moe_layers += 1
+
+        self.extract_moe_parameters(example_moe)
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        return self.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        if not get_pp_group().is_last_rank:
+            return None
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        return logits
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
+        )
+        return loader.load_weights(weights)
+
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        return self.model.get_expert_mapping()
+
+
+class SarvamMoEForCausalLM(BailingMoeForCausalLM):
+    """Same as BailingMoeForCausalLM, but normalizes gate expert_bias pre-load."""
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        return super().load_weights(_normalized_weights(weights))

special_tokens_map.json

@@ -0,0 +1,33 @@
+{
+  "boi_token": "<|start_of_image|>",
+  "bos_token": {
+    "content": "[@BOS@]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eoi_token": "<|end_of_image|>",
+  "eos_token": {
+    "content": "<|end_of_turn|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "image_token": "<|image_soft_token|>",
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a574ceaaff7c7a8f091179c53fd17ae33567089c099d4ff37d4cb3bc1a87e80e
+size 33627251