minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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
+1
View File
@@ -413,6 +413,7 @@ 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)
if _is_npu:
emb = torch.cat((freqs, freqs), dim=-1) emb = torch.cat((freqs, freqs), dim=-1)
self.cos_cached_total = torch.cos(emb) * self.mscale self.cos_cached_total = torch.cos(emb) * self.mscale
self.sin_cached_total = torch.sin(emb) * self.mscale self.sin_cached_total = torch.sin(emb) * self.mscale