diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 47b6c8f6e..caf158924 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -220,6 +220,11 @@ class Envs: # large bs) but coarser overlap. 1 = per-layer. SGLANG_CP_SHARED_KV_PER_LAYER_GROUP = EnvInt(8) SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False) + # NSA paged MQA-logits chunking equivalence test: when >0, force the paged + # topk path to chunk at this many query rows AND assert the chunked topk_result + # is byte-identical to the unchunked single-call result. For validation only + # (run a small batch so the unchunked reference fits); 0 = off (production). + SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS = EnvInt(0) SGLANG_CP_SHARED_KV_FUSED_MLA_STORE = EnvBool(False) SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE = EnvBool(False) SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 537c3b4cc..67c959642 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1073,6 +1073,7 @@ class Indexer(MultiPlatformOp): # When attn_tp_size > 1 or in the MAX_LEN padding mode, padding may exist in the hidden states, # and it is necessary to extract the actual q length. q_offset = sum(metadata.get_nsa_extend_len_cpu()) + topk_result = None if _is_hip: from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits @@ -1098,19 +1099,100 @@ class Indexer(MultiPlatformOp): WavePerEU=5, ) else: - logits = deep_gemm.fp8_paged_mqa_logits( - q_fp8[:q_offset], - kv_cache_fp8, - weights[:q_offset], - seqlens_32_2d, - block_tables, - schedule_metadata, - max_seq_len, - clean_logits=False, + device_index = q_fp8.device.index + assert device_index is not None + # The kernel allocates logits of width align(max_seq_len, 256) (DeepGEMM + # attention.hpp), so budget/chunk against the aligned width. + aligned_ctx = ((max_seq_len + 255) // 256) * 256 + force_rows = int(envs.SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS.get()) + need_chunk, logits_budget_bytes = self._should_chunk_mqa_logits( + q_offset, aligned_ctx, device_index, forward_batch=forward_batch ) + if force_rows > 0: + need_chunk = True + if not need_chunk: + logits = deep_gemm.fp8_paged_mqa_logits( + q_fp8[:q_offset], + kv_cache_fp8, + weights[:q_offset], + seqlens_32_2d, + block_tables, + schedule_metadata, + max_seq_len, + clean_logits=False, + ) + else: + # Bound the q_offset x align(max_seq_len,256) f32 logits buffer by + # chunking over query rows (each paged q-row is its own length-1 entry, + # so any row split is valid). Recompute the SM schedule per chunk (it + # encodes the work split for this chunk's context_lens). Run the topk + # transform per chunk with the per-chunk paged args (ke_offset / + # batch_idx_list / cu_seqlens_q override) so we never materialize the + # full logits buffer. Mirrors the ragged chunk loop. + bytes_per_row = aligned_ctx * self._MQA_LOGITS_BYTES_PER_ELEM + if force_rows > 0: + max_rows = force_rows + else: + max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1))) + max_rows = min(max(1, max_rows), q_offset) + seqlens_expanded_full = metadata.get_seqlens_expanded() + start = 0 + while start < q_offset: + end = min(start + max_rows, q_offset) + sched_chunk = deep_gemm.get_paged_mqa_logits_metadata( + seqlens_32_2d[start:end], blocksize, self.sm_count + ) + logits_chunk = deep_gemm.fp8_paged_mqa_logits( + q_fp8[start:end], + kv_cache_fp8, + weights[start:end], + seqlens_32_2d[start:end], + block_tables[start:end], + sched_chunk, + max_seq_len, + clean_logits=False, + ) + cu_chunk = torch.arange( + 0, + (end - start) + 1, + dtype=torch.int32, + device=logits_chunk.device, + ) + topk_chunk = metadata.topk_transform( + logits_chunk, + self.index_topk, + ke_offset=seqlens_expanded_full[start:end], + batch_idx_list=list(range(start, end)), + cu_seqlens_q_topk_override=cu_chunk, + ) + if topk_result is None: + topk_result = topk_chunk.new_full( + (q_offset, topk_chunk.shape[1]), -1 + ) + topk_result[start:end] = topk_chunk + start = end + if force_rows > 0: + # Equivalence gate: chunked topk_result must be byte-identical to + # the unchunked single-call path (run a small batch so this fits). + ref_logits = deep_gemm.fp8_paged_mqa_logits( + q_fp8[:q_offset], + kv_cache_fp8, + weights[:q_offset], + seqlens_32_2d, + block_tables, + schedule_metadata, + max_seq_len, + clean_logits=False, + ) + ref_topk = metadata.topk_transform(ref_logits, self.index_topk) + assert torch.equal(topk_result, ref_topk), ( + "[MQA_LOGITS_CHUNK_VERIFY] paged chunked topk_result != " + f"unchunked (q_offset={q_offset}, max_rows={max_rows})" + ) - # NOTE(dark): logits should be cleaned in topk_transform - topk_result = metadata.topk_transform(logits, self.index_topk) + if topk_result is None: + # NOTE(dark): logits should be cleaned in topk_transform + topk_result = metadata.topk_transform(logits, self.index_topk) # Restore possible padding exist in the hidden states. if not _is_hip and q_offset < q_fp8.shape[0]: pad_len = q_fp8.shape[0] - q_offset @@ -1158,8 +1240,35 @@ class Indexer(MultiPlatformOp): self._mqa_logits_budget_bytes[device_index] = budget_bytes return budget_bytes + def _current_free_mem_logits_budget( + self, device_index: int, forward_batch=None + ) -> int: + """0.5 x CURRENT free-memory budget for the MQA-logits buffer, queried at + most ONCE per forward. + + ``torch.cuda.mem_get_info`` host-syncs; the indexer runs once per layer, so + querying it per call serializes the host ~num_layers x per forward and + starves the GPU between batches (commit 40a0389a9c introduced the per-call + query for OOM safety -- this caches it without losing that safety). Free + memory is ~constant across layers within a forward: eager mode frees each + layer's activations and the KV pool is pre-reserved, so the driver-level + free-mem high-water-mark is set early and stays flat. We snapshot it on the + ``forward_batch`` (recreated per forward) and reuse it for later layers; a new + forward gets a fresh snapshot. Falls back to a direct query when no + ``forward_batch`` is threaded (keeps the call correct, just uncached). + """ + if forward_batch is not None: + cached = getattr(forward_batch, "_nsa_mqa_free_budget", None) + if cached is not None and cached[0] == device_index: + return cached[1] + free_mem, _ = torch.cuda.mem_get_info(device_index) + budget_bytes = max(1, int(free_mem * self._MQA_LOGITS_FREE_MEM_FRACTION)) + if forward_batch is not None: + forward_batch._nsa_mqa_free_budget = (device_index, budget_bytes) + return budget_bytes + def _should_chunk_mqa_logits( - self, num_q: int, num_k: int, device_index: int + self, num_q: int, num_k: int, device_index: int, forward_batch=None ) -> Tuple[bool, int]: """ Detect whether we need to chunk the MQA logits computation to avoid OOM @@ -1170,11 +1279,122 @@ class Indexer(MultiPlatformOp): return False, 0 logits_bytes = num_q * num_k * self._MQA_LOGITS_BYTES_PER_ELEM - logits_budget_bytes = self._get_mqa_logits_budget_bytes(device_index) + # Budget against CURRENT free memory, not the cached first-prefill / static + # estimate. The static headroom (1 - mem_fraction_static) is SHARED with the + # rest of the forward's activations, so a cached estimate over-counts what the + # logits buffer alone may use and OOMs at large batch (observed: 15.1 GiB + # logits "fit" a 24 GiB static budget but only 14.67 GiB was actually free). + # Reached only for large logits (post static-skip). The per-forward snapshot + # (see _current_free_mem_logits_budget) keeps that current-free-mem safety + # while collapsing the per-layer host-sync to once per forward. + # Keep the static guard during CUDA-graph capture (mem_get_info unreliable). + if get_is_capture_mode(): + logits_budget_bytes = self._get_mqa_logits_budget_bytes(device_index) + else: + logits_budget_bytes = self._current_free_mem_logits_budget( + device_index, forward_batch + ) need_chunk = logits_bytes > logits_budget_bytes return need_chunk, logits_budget_bytes + def _mqa_logits_topk_ragged_chunked( + self, + metadata, + q_fp8, + kv_fp8, + weights, + ks, + ke, + *, + actual_seq_q, + ke_offset, + batch_idx_list, + topk_indices_offset_override, + forward_batch=None, + ): + """RAGGED fp8_mqa_logits + topk_transform, byte-budget-chunked over query rows. + + Mirrors the `_get_topk_ragged` chunk loop so the unbounded `q_offset x kv_len` + f32 logits buffer can't OOM when the CP prefill batch grows. Per-row inputs + (q/weights/ks/ke/ke_offset/topk_indices_offset) are sliced; the shared `kv_fp8` + stays whole. The RAGGED transform keys off per-row `ks` + `ke_offset` + + `topk_indices_offset_override`, so per-chunk results are byte-identical + (`cu_seqlens_q`/`batch_idx_list` are unused once the override is set -- + nsa_backend.py:569). SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS>0 forces chunking + and asserts equivalence vs the single-call path. + """ + device_index = q_fp8.device.index + assert device_index is not None + q_offset = int(q_fp8.shape[0]) + k_offset = int(kv_fp8[0].shape[0]) + force_rows = int(envs.SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS.get()) + need_chunk, logits_budget_bytes = self._should_chunk_mqa_logits( + q_offset, k_offset, device_index, forward_batch=forward_batch + ) + if force_rows > 0: + need_chunk = True + + def _single(): + with self._with_real_sm_count(): + logits = deep_gemm.fp8_mqa_logits( + q_fp8, kv_fp8, weights, ks, ke, clean_logits=False + ) + return metadata.topk_transform( + logits, + self.index_topk, + ks=ks, + cu_seqlens_q=actual_seq_q, + ke_offset=ke_offset, + batch_idx_list=batch_idx_list, + topk_indices_offset_override=topk_indices_offset_override, + ) + + if not need_chunk: + return _single() + + bytes_per_row = k_offset * self._MQA_LOGITS_BYTES_PER_ELEM + if force_rows > 0: + max_rows = force_rows + else: + max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1))) + max_rows = min(max(1, max_rows), q_offset) + + topk_result = None + start = 0 + while start < q_offset: + end = min(start + max_rows, q_offset) + with self._with_real_sm_count(): + logits_chunk = deep_gemm.fp8_mqa_logits( + q_fp8[start:end], + kv_fp8, + weights[start:end], + ks[start:end], + ke[start:end], + clean_logits=False, + ) + topk_chunk = metadata.topk_transform( + logits_chunk, + self.index_topk, + ks=ks[start:end], + ke_offset=ke_offset[start:end], + topk_indices_offset_override=topk_indices_offset_override[start:end], + ) + if topk_result is None: + topk_result = topk_chunk.new_full( + (q_offset, topk_chunk.shape[1]), -1 + ) + topk_result[start:end] = topk_chunk + start = end + + if force_rows > 0: + ref = _single() + assert torch.equal(topk_result, ref), ( + "[MQA_LOGITS_CHUNK_VERIFY] cp-ragged chunked topk_result != " + f"unchunked (q_offset={q_offset}, max_rows={max_rows})" + ) + return topk_result + def _get_topk_ragged( self, enable_dual_stream: bool, @@ -1256,7 +1476,7 @@ class Indexer(MultiPlatformOp): 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 + q_offset, k_offset, device_index, forward_batch=forward_batch ) if not need_chunk: @@ -1588,23 +1808,18 @@ class Indexer(MultiPlatformOp): q_lens_list, dtype=torch.int32, device=q_fp8.device ) ke = ks + ke_offset - with self._with_real_sm_count(): - logits = deep_gemm.fp8_mqa_logits( - q_fp8, - kv_fp8, - weights, - ks, - ke, - clean_logits=False, - ) - topk_result = metadata.topk_transform( - logits, - self.index_topk, - ks=ks, - cu_seqlens_q=actual_seq_q, + topk_result = self._mqa_logits_topk_ragged_chunked( + metadata, + q_fp8, + kv_fp8, + weights, + ks, + ke, + actual_seq_q=actual_seq_q, ke_offset=ke_offset, batch_idx_list=batch_idx_list, topk_indices_offset_override=topk_indices_offset_override, + forward_batch=forward_batch, ) return topk_result else: @@ -1653,23 +1868,18 @@ class Indexer(MultiPlatformOp): q_lens_list, dtype=torch.int32, device=q_fp8.device ) ke = ks + ke_offset - with self._with_real_sm_count(): - logits = deep_gemm.fp8_mqa_logits( - q_fp8, - kv_fp8, - weights, - ks, - ke, - clean_logits=False, - ) - topk_result = metadata.topk_transform( - logits, - self.index_topk, - ks=ks, - cu_seqlens_q=actual_seq_q, + topk_result = self._mqa_logits_topk_ragged_chunked( + metadata, + q_fp8, + kv_fp8, + weights, + ks, + ke, + actual_seq_q=actual_seq_q, ke_offset=ke_offset, batch_idx_list=batch_idx_list, topk_indices_offset_override=topk_indices_offset_override, + forward_batch=forward_batch, ) return topk_result @@ -1748,23 +1958,18 @@ class Indexer(MultiPlatformOp): ke_offset = torch.cat(ke_offset_list, dim=0) ke = ks + ke_offset actual_seq_q = torch.cat(actual_seq_q_list, dim=0) - with self._with_real_sm_count(): - logits = deep_gemm.fp8_mqa_logits( - q_fp8, - kv_fp8, - weights, - ks, - ke, - clean_logits=False, - ) - topk_result = metadata.topk_transform( - logits, - self.index_topk, - ks=ks, - cu_seqlens_q=actual_seq_q, + topk_result = self._mqa_logits_topk_ragged_chunked( + metadata, + q_fp8, + kv_fp8, + weights, + ks, + ke, + actual_seq_q=actual_seq_q, ke_offset=ke_offset, batch_idx_list=batch_idx_list, topk_indices_offset_override=topk_indices_offset_override, + forward_batch=forward_batch, ) else: seq_len = int(forward_batch.seq_lens_cpu[batch_idx].item())