From 96b5b0ee4b728a057d760cd4f934a4082cbc3867 Mon Sep 17 00:00:00 2001 From: Vivaan Sharma <99265792+Vivaansharma2003@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:02:39 +0530 Subject: [PATCH] [Model] Enable torch.compile and CUDA graphs for EmbeddingGemma2 (#61572) Signed-off-by: Vivaan Sharma <99265792+Vivaansharma2003@users.noreply.github.com> Co-authored-by: Kevin H. Luu --- vllm/model_executor/models/embedding_gemma2.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/model_executor/models/embedding_gemma2.py b/vllm/model_executor/models/embedding_gemma2.py index 0f92aeaa37..efde2acd1e 100644 --- a/vllm/model_executor/models/embedding_gemma2.py +++ b/vllm/model_executor/models/embedding_gemma2.py @@ -18,6 +18,7 @@ from transformers.models.embedding_gemma2.processing_embedding_gemma2 import ( ) from transformers.video_utils import VideoMetadata +from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import get_tensor_model_parallel_world_size from vllm.logger import init_logger @@ -334,6 +335,7 @@ class EmbeddingGemma2TextPLE(nn.Module): return self.per_layer_projection_norm(proj) +@support_torch_compile class EmbeddingGemma2TextModel(nn.Module): """Text embedding backbone for EmbeddingGemma2 with bidirectional attention."""