diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index caf158924..47b6c8f6e 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -220,11 +220,6 @@ 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 67c959642..537c3b4cc 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1073,7 +1073,6 @@ 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 @@ -1099,100 +1098,19 @@ class Indexer(MultiPlatformOp): WavePerEU=5, ) else: - 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 + 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, ) - 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})" - ) - if topk_result is None: - # NOTE(dark): logits should be cleaned in topk_transform - topk_result = metadata.topk_transform(logits, self.index_topk) + # 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 @@ -1240,35 +1158,8 @@ 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, forward_batch=None + 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 @@ -1279,122 +1170,11 @@ class Indexer(MultiPlatformOp): return False, 0 logits_bytes = num_q * num_k * self._MQA_LOGITS_BYTES_PER_ELEM - # 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 - ) + logits_budget_bytes = self._get_mqa_logits_budget_bytes(device_index) 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, @@ -1476,7 +1256,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, forward_batch=forward_batch + q_offset, k_offset, device_index ) if not need_chunk: @@ -1808,18 +1588,23 @@ class Indexer(MultiPlatformOp): q_lens_list, dtype=torch.int32, device=q_fp8.device ) ke = ks + ke_offset - topk_result = self._mqa_logits_topk_ragged_chunked( - metadata, - q_fp8, - kv_fp8, - weights, - ks, - ke, - actual_seq_q=actual_seq_q, + 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, 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: @@ -1868,18 +1653,23 @@ class Indexer(MultiPlatformOp): q_lens_list, dtype=torch.int32, device=q_fp8.device ) ke = ks + ke_offset - topk_result = self._mqa_logits_topk_ragged_chunked( - metadata, - q_fp8, - kv_fp8, - weights, - ks, - ke, - actual_seq_q=actual_seq_q, + 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, 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 @@ -1958,18 +1748,23 @@ 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) - topk_result = self._mqa_logits_topk_ragged_chunked( - metadata, - q_fp8, - kv_fp8, - weights, - ks, - ke, - actual_seq_q=actual_seq_q, + 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, 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())