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

Add NPU condition for cosine and sine caching (#25399)

Browse Source
This commit is contained in:
Cheng Wan
2026-05-15 05:21:53 -07:00
committed by GitHub
parent fd9525436b
commit 3f7e538b2f
1 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/rotary_embedding/rope_variant.py
+4 -3
View File
@@ -413,9 +413,10 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
cos = freqs.cos() * self.mscale cos = freqs.cos() * self.mscale
sin = freqs.sin() * self.mscale sin = freqs.sin() * self.mscale
cache = torch.cat((cos, sin), dim=-1) cache = torch.cat((cos, sin), dim=-1)
emb = torch.cat((freqs, freqs), dim=-1) if _is_npu:
self.cos_cached_total = torch.cos(emb) * self.mscale emb = torch.cat((freqs, freqs), dim=-1)
self.sin_cached_total = torch.sin(emb) * self.mscale self.cos_cached_total = torch.cos(emb) * self.mscale
self.sin_cached_total = torch.sin(emb) * self.mscale
return cache return cache
def get_cos_cached_total(self): def get_cos_cached_total(self):