[Model] Enable torch.compile for Sarvam MLA (#55817)
Signed-off-by: mohit-sarvam <mohit@sarvam.ai> Co-authored-by: Codex <noreply@openai.com>
Mohit Singla committed
Sep 8, 2026 at 19:31 UTC
613ab20f7f9082244240f4e9b2e7a4ee6769c38f
1 file changed
+11
vllm/model_executor/models/sarvam.py
+11
@@ -28,6 +28,7 @@ from itertools import islice
28
import torch
29
from torch import nn
30
31
+from vllm.compilation.decorators import support_torch_compile
32
from vllm.config import CacheConfig, ParallelConfig, VllmConfig
33
from vllm.distributed import (
34
get_pp_group,
@@ -448,6 +449,16 @@ class SarvamMLABlock(nn.Module):
449
return hidden_states, residual
450
451
452
+# PEP 563 stringifies annotations in this module, so the decorator cannot infer
453
+# these; they are the set DeepSeek-V2 infers for the same forward signature.
454
+@support_torch_compile(
455
+ dynamic_arg_dims={
456
+ "input_ids": 0,
457
+ "positions": 0,
458
+ "intermediate_tensors": 0,
459
+ "inputs_embeds": 0,
460
+ }
461
+)
462
class SarvamMLAModel(nn.Module, EagleModelMixin):
463
"""Sarvam MLA backbone with stage-local EAGLE3 auxiliary capture."""
464