EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[DSv32] Use torch.compile for _get_logits_head_gate (#11565)

Browse Source
This commit is contained in:
Trevor Morris
2025-10-13 18:38:39 -07:00
committed by GitHub
parent aaf7af1b17
commit 384733639a
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -205,6 +205,7 @@ class Indexer(CustomOp):
return ans return ans
@torch.compile(dynamic=True)
def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor): def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor):
weights, _ = self.weights_proj(x) weights, _ = self.weights_proj(x)
weights = weights * self.n_heads**-0.5 weights = weights * self.n_heads**-0.5