[Bugfix][Model] Enable torch.compile for StableLM (#54969)
Signed-off-by: Djordje Ramic <djoramic@amd.com>
djramic committed
Sep 3, 2026 at 12:21 UTC
0d3ede3e3bd62cae2fc5c465b9bc2e62028b1e71
1 file changed
+2
vllm/model_executor/models/stablelm.py
+2
@@ -29,6 +29,7 @@ import torch
29
from torch import nn
30
from transformers import StableLmConfig
31
32
+from vllm.compilation.decorators import support_torch_compile
33
from vllm.config import CacheConfig, VllmConfig
34
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
35
from vllm.model_executor.layers.activation import SiluAndMul
@@ -213,6 +214,7 @@ class StablelmDecoderLayer(nn.Module):
214
return hidden_states, residual
215
216
217
+@support_torch_compile
218
class StableLMEpochModel(nn.Module):
219
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
220
super().__init__()