From be4a3ec376eac11803b56b39b4eb473e5fb751d3 Mon Sep 17 00:00:00 2001
From: Minglei Zhu <mingleizhu1122@gmail.com>
Date: Sat, 6 Dec 2025 17:57:44 -0800
Subject: [PATCH] support piecewise cuda graph for Olmo models (#14476)

---
 python/sglang/srt/models/olmo.py  | 1 +
 python/sglang/srt/models/olmo2.py | 1 +
 2 files changed, 2 insertions(+)

diff --git a/python/sglang/srt/models/olmo.py b/python/sglang/srt/models/olmo.py
index 08239374d..0c1ecf857 100644
--- a/python/sglang/srt/models/olmo.py
+++ b/python/sglang/srt/models/olmo.py
@@ -318,6 +318,7 @@ class OlmoForCausalLM(nn.Module):
             )
         self.logits_processor = LogitsProcessor(config)
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,
diff --git a/python/sglang/srt/models/olmo2.py b/python/sglang/srt/models/olmo2.py
index 97754f362..a4b062f27 100644
--- a/python/sglang/srt/models/olmo2.py
+++ b/python/sglang/srt/models/olmo2.py
@@ -379,6 +379,7 @@ class Olmo2ForCausalLM(nn.Module):
     def get_attention_sliding_window_size(self):
         return get_attention_sliding_window_size(self.config)
 
+    @torch.no_grad()
     def forward(
         self,
         input_ids: torch.Tensor,