diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 8d70695a3..b03deeb65 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/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) - half = c.shape[-1] // 2 - self._indexer_freqs_cis = torch.complex( - c[:, :half].contiguous(), c[:, half:].contiguous() - ) + self._indexer_freqs_cis = _shared_indexer_freqs_cis(self.rotary_emb) @contextlib.contextmanager def _with_real_sm_count(self):