minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

support piecewise cuda graph for Olmo models (#14476)

Browse Source
This commit is contained in:
Minglei Zhu
2025-12-06 17:57:44 -08:00
committed by GitHub
parent ff6e3ea934
commit be4a3ec376
2 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/olmo.py
+1
View File
@@ -318,6 +318,7 @@ class OlmoForCausalLM(nn.Module):
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
@torch.no_grad()
def forward( def forward(
self, self,
input_ids: torch.Tensor, input_ids: torch.Tensor,
python/sglang/srt/models/olmo2.py
+1
View File
@@ -379,6 +379,7 @@ class Olmo2ForCausalLM(nn.Module):
def get_attention_sliding_window_size(self): def get_attention_sliding_window_size(self):
return get_attention_sliding_window_size(self.config) return get_attention_sliding_window_size(self.config)
@torch.no_grad()
def forward( def forward(
self, self,
input_ids: torch.Tensor, input_ids: torch.Tensor,