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

[NPU][Feature] eliminate dsv3 redundant rotary embed calculation (#19842)

Browse Source
This commit is contained in:
liupeng374
2026-03-06 09:02:14 +08:00
committed by GitHub
parent b912d7ae19
commit 5471e4a492
1 changed files with 12 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py
+10
View File
@@ -266,7 +266,12 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
): ):
bsz, _ = hidden_states.view(-1, hidden_states.shape[-1]).shape bsz, _ = hidden_states.view(-1, hidden_states.shape[-1]).shape
self.dtype = hidden_states.dtype self.dtype = hidden_states.dtype
if self.layer_id == 0:
self.cos, self.sin = self.get_sin_cos(positions) self.cos, self.sin = self.get_sin_cos(positions)
self.rotary_emb.cos_cached, self.rotary_emb.sin_cache = self.cos, self.sin
else:
self.cos, self.sin = self.rotary_emb.cos_cached, self.rotary_emb.sin_cache
self.kvCache, self.kvCacheRope, self.slotmapping = ( self.kvCache, self.kvCacheRope, self.slotmapping = (
self.get_kv_cache_and_cache_idx(forward_batch) self.get_kv_cache_and_cache_idx(forward_batch)
) )
@@ -340,7 +345,12 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
self.has_preprocess_weights = True self.has_preprocess_weights = True
self.dtype = hidden_states.dtype self.dtype = hidden_states.dtype
if self.layer_id == 0:
cos, sin = self.get_sin_cos(positions) cos, sin = self.get_sin_cos(positions)
self.rotary_emb.cos_cached, self.rotary_emb.sin_cache = cos, sin
else:
cos, sin = self.rotary_emb.cos_cached, self.rotary_emb.sin_cache
k_cache, v_cache, slot_mapping = self.get_kv_cache_and_cache_idx(forward_batch) k_cache, v_cache, slot_mapping = self.get_kv_cache_and_cache_idx(forward_batch)
q_nope_out = torch.empty( q_nope_out = torch.empty(