Fix DSA indexer fusion bug causing excessive memory consumption. (#29576)

Co-authored-by: Brayden Zhong <brayden@radixark.ai>

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -327,6 +327,16 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor:
     return hadamard_transform(x, scale=hidden_size**-0.5)
 
 
+def _shared_indexer_freqs_cis(rotary_emb: torch.nn.Module) -> torch.Tensor:
+    cached = getattr(rotary_emb, "_dsa_indexer_freqs_cis", None)
+    if cached is None:
+        c = rotary_emb.cos_sin_cache.to(torch.float32)
+        half = c.shape[-1] // 2
+        cached = torch.complex(c[:, :half].contiguous(), c[:, half:].contiguous())
+        rotary_emb._dsa_indexer_freqs_cis = cached
+    return cached
+
+
 class Indexer(MultiPlatformOp):
     _MQA_LOGITS_BYTES_PER_ELEM = 4
     _MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000
@@ -427,14 +437,9 @@ class Indexer(MultiPlatformOp):
         self.scale_fmt = scale_fmt
         self.softmax_scale = self.head_dim**-0.5
 
-        # freqs_cis is built from the fp32 cos/sin cache before any forward casts it to bf16.
         self._indexer_freqs_cis: Optional[torch.Tensor] = None
         if _use_dsa_indexer_fusion:
-            c = self.rotary_emb.cos_sin_cache.to(torch.float32)
+            self._indexer_freqs_cis = _shared_indexer_freqs_cis(self.rotary_emb)
-            half = c.shape[-1] // 2
-            self._indexer_freqs_cis = torch.complex(
-                c[:, :half].contiguous(), c[:, half:].contiguous()
-            )
 
     @contextlib.contextmanager
     def _with_real_sm_count(self):