From 384733639a917c74b28d2abcc0ddf412fa2a47d8 Mon Sep 17 00:00:00 2001 From: Trevor Morris Date: Mon, 13 Oct 2025 18:38:39 -0700 Subject: [PATCH] [DSv32] Use torch.compile for _get_logits_head_gate (#11565) --- python/sglang/srt/layers/attention/nsa/nsa_indexer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 13ab6b172..8921e1214 100644 --- 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