From 613ab20f7f9082244240f4e9b2e7a4ee6769c38f Mon Sep 17 00:00:00 2001 From: Mohit Singla <140096022+mohit-sarvam@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:31:30 +0530 Subject: [PATCH] [Model] Enable torch.compile for Sarvam MLA (#55817) Signed-off-by: mohit-sarvam Co-authored-by: Codex --- vllm/model_executor/models/sarvam.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py index 3ce94b3f32..e790a27c31 100644 --- a/vllm/model_executor/models/sarvam.py +++ b/vllm/model_executor/models/sarvam.py @@ -28,6 +28,7 @@ from itertools import islice import torch from torch import nn +from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ParallelConfig, VllmConfig from vllm.distributed import ( get_pp_group, @@ -448,6 +449,16 @@ class SarvamMLABlock(nn.Module): return hidden_states, residual +# PEP 563 stringifies annotations in this module, so the decorator cannot infer +# these; they are the set DeepSeek-V2 infers for the same forward signature. +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": 0, + "intermediate_tensors": 0, + "inputs_embeds": 0, + } +) class SarvamMLAModel(nn.Module, EagleModelMixin): """Sarvam MLA backbone with stage-local EAGLE3 auxiliary capture."""