support piecewise cuda graph for Olmo models (#14476)

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,

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,