diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 16f11a73f..f1f6603e1 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/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