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

2025-10-13 18:38:39 -07:00
parent aaf7af1b17
commit 384733639a
1 changed files with 1 additions and 0 deletions
--- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
@@ -205,6 +205,7 @@ class Indexer(CustomOp):

        return ans

+    @torch.compile(dynamic=True)
    def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor):
        weights, _ = self.weights_proj(x)
        weights = weights * self.n_heads**-0.5