perf(nsa): cache MQA-logits budget + drop indexer rope clones

Two upstream NSA-indexer perf ports (verified against our diverged tree), cutting
CPU launches / host syncs on the prefill critical path for GLM-5.1-FP8 on B300:

- #25299 (beaff00331): cache the MQA-logits chunking budget per device so prefill
  stops calling torch.cuda.mem_get_info (a host sync) on every indexer pass. The
  budget is computed once and capped by the mem_fraction_static serving headroom;
  a static guard is used (uncached) during cuda-graph capture, and the first real
  prefill caches the free-memory budget.
- #22232 (671fe73961): replace `dst = src.clone()` slice write-backs of the RoPE
  output with a data_ptr-guarded `dst.copy_(src)`. q_rope/k_rope are torch.split
  views of query/key, so when RoPE runs in place src/dst alias and the write-back
  is a redundant no-op (guard skips it); otherwise one copy instead of clone+copy.
  Saves an alloc+copy per q/k per indexer call. (Skipped the PR's AMD-only
  @torch.compile cleanup.)

Both verified to import on the B300 / torch-2.11 image. Deliberately NOT taken:
#21332 (un-force MHA one-shot on Blackwell — risky, per user), #23856 (torch.mm
indexer GEMM — per user).

Refs: WI-2026-06-07-001

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 6bc23945e4ff5c68b0618173856c7daf4653fa10)

python/sglang/srt/layers/attention/nsa/nsa_indexer.py

@@ -413,6 +413,14 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor:
 
 
 class Indexer(MultiPlatformOp):
+    # MQA-logits chunking budget — cached per device so we avoid a per-call
+    # torch.cuda.mem_get_info host sync on the prefill critical path. (upstream PR #25299)
+    _MQA_LOGITS_BYTES_PER_ELEM = 4
+    _MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000
+    _MQA_LOGITS_FREE_MEM_FRACTION = 0.5
+    _MQA_LOGITS_TOTAL_MEM_FRACTION = 0.3
+    _mqa_logits_budget_bytes: Dict[int, int] = {}
+
     def __init__(
         self,
         hidden_size: int,
@@ -880,6 +888,15 @@ class Indexer(MultiPlatformOp):
         weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
         return weights
 
+    @staticmethod
+    def _update_rope_guarded(dst: torch.Tensor, src: torch.Tensor) -> None:
+        # Write RoPE output into the destination slice with a single copy (no clone).
+        # If the RoPE kernel wrote in place, src and dst alias the same memory and the
+        # write-back is a redundant no-op — skip it. (upstream PR #22232)
+        if src.data_ptr() == dst.data_ptr():
+            return
+        dst.copy_(src)
+
     def _get_q_k_bf16(
         self,
         q_lora: torch.Tensor,
@@ -928,8 +945,8 @@ class Indexer(MultiPlatformOp):
 
         q_rope, k_rope = self.rotary_emb(positions, q_rope, k_rope)
 
-        query[..., : self.rope_head_dim] = q_rope.clone()
-        key[..., : self.rope_head_dim] = k_rope.clone()
+        self._update_rope_guarded(query[..., : self.rope_head_dim], q_rope)
+        self._update_rope_guarded(key[..., : self.rope_head_dim], k_rope)
 
         if enable_dual_stream:
             current_stream = torch.cuda.current_stream()
@@ -968,7 +985,7 @@ class Indexer(MultiPlatformOp):
         )
 
         _, k_rope = self.rotary_emb(positions, k_rope, k_rope)
-        key[..., : self.rope_head_dim] = k_rope.clone()
+        self._update_rope_guarded(key[..., : self.rope_head_dim], k_rope)
         key = rotate_activation(key)
 
         return key
@@ -1096,24 +1113,57 @@ class Indexer(MultiPlatformOp):
             topk_result = torch.cat([topk_result, padding], dim=0)
         return topk_result
 
+    def _get_mqa_logits_budget_bytes(self, device_index: int) -> int:
+        # Cache the MQA-logits byte budget per device. torch.cuda.mem_get_info
+        # host-syncs, so query free memory at most once (after the first real
+        # prefill) and cap it by the workload-independent serving-memory headroom
+        # derived from mem_fraction_static. (upstream PR #25299)
+        cached_budget = self._mqa_logits_budget_bytes.get(device_index)
+        if cached_budget is not None:
+            return cached_budget
+
+        total_mem = torch.cuda.get_device_properties(device_index).total_memory
+        total_mem_budget = int(total_mem * self._MQA_LOGITS_TOTAL_MEM_FRACTION)
+        mem_fraction_static = get_global_server_args().mem_fraction_static
+        if mem_fraction_static is None:
+            static_budget = total_mem_budget
+        else:
+            static_free_mem = int(total_mem * max(0.0, 1.0 - mem_fraction_static))
+            static_budget = min(
+                int(static_free_mem * self._MQA_LOGITS_FREE_MEM_FRACTION),
+                total_mem_budget,
+            )
+        static_budget = max(1, static_budget)
+
+        # During CUDA graph capture keep the static guard but do NOT cache it; the
+        # first non-capture prefill caches the real free-memory budget below.
+        if get_is_capture_mode():
+            return static_budget
+
+        free_mem, _ = torch.cuda.mem_get_info(device_index)
+        budget_bytes = min(
+            int(free_mem * self._MQA_LOGITS_FREE_MEM_FRACTION), static_budget
+        )
+        budget_bytes = max(1, budget_bytes)
+        self._mqa_logits_budget_bytes[device_index] = budget_bytes
+        return budget_bytes
+
     def _should_chunk_mqa_logits(
-        self, num_q: int, num_k: int, device: torch.device
+        self, num_q: int, num_k: int, device_index: int
     ) -> Tuple[bool, int]:
         """
         Detect whether we need to chunk the MQA logits computation to avoid OOM
-        Return: (need_chunk, free_mem)
+        Return: (need_chunk, logits_budget_bytes)
         """
         # Quick static check for normal batches
-        if num_q * num_k < 8_000_000:  # 8M elements ≈ 32MB logits
+        if num_q * num_k < self._MQA_LOGITS_STATIC_SKIP_ELEMS:
             return False, 0
 
-        free_mem, total_mem = torch.cuda.mem_get_info(device)
-        bytes_per_elem = 4  # float32
-        logits_bytes = num_q * num_k * bytes_per_elem
+        logits_bytes = num_q * num_k * self._MQA_LOGITS_BYTES_PER_ELEM
+        logits_budget_bytes = self._get_mqa_logits_budget_bytes(device_index)
 
-        # Logits should not exceed 50% of free memory or 30% of total memory
-        need_chunk = (logits_bytes * 2 > free_mem) or (logits_bytes > total_mem * 0.3)
-        return need_chunk, free_mem
+        need_chunk = logits_bytes > logits_budget_bytes
+        return need_chunk, logits_budget_bytes
 
     def _get_topk_ragged(
         self,
@@ -1193,7 +1243,11 @@ class Indexer(MultiPlatformOp):
         token_to_batch_idx = metadata.get_token_to_batch_idx()
         q_offset = ks.shape[0]
         k_offset = k_fp8.shape[0]
-        need_chunk, free_mem = self._should_chunk_mqa_logits(q_offset, k_offset, device)
+        device_index = device.index
+        assert device_index is not None, "q_fp8 must be on an indexed CUDA device"
+        need_chunk, logits_budget_bytes = self._should_chunk_mqa_logits(
+            q_offset, k_offset, device_index
+        )
 
         if not need_chunk:
             assert q_fp8[:q_offset].shape[0] != 0
@@ -1222,10 +1276,8 @@ class Indexer(MultiPlatformOp):
             return topk_result
 
         # Chunk path
-        bytes_per_elem = 4  # float32
-        bytes_per_row = k_offset * bytes_per_elem
-        # Reserve 50% of free memory for logits
-        max_rows = max(1, int((free_mem * 0.5) // max(bytes_per_row, 1)))
+        bytes_per_row = k_offset * self._MQA_LOGITS_BYTES_PER_ELEM
+        max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1)))
         max_rows = min(max_rows, q_offset)
 
         global_topk_offset = metadata.attn_metadata.topk_indices_offset