From e0ea8a485ccc82264447071f50e2d5259c06782c Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Thu, 11 Jun 2026 03:04:03 +0800 Subject: [PATCH] Revert NSA MQA logits chunking while performance is unresolved Temporarily revert the syh MQA-logits row-chunking port because the observed runtime performance is worse on the current bs>1 CP prefill workload. Keep the history explicit so the optimization can be revisited after profiling identifies where the extra overhead comes from. This reverts commit 4e49751406 (Bound NSA MQA logits peak memory). Constraint: Current priority is restoring the faster known path for remote ETE runs Rejected: Keep chunking behind the force/debug env only | the production auto-chunk path still changes runtime heuristics and should not stay until profiled Confidence: high Scope-risk: narrow Directive: Reintroduce logits chunking only with ETE performance evidence and forced-chunk equivalence coverage Tested: Local py_compile for environ.py and nsa_indexer.py Not-tested: Remote ETE performance after revert Co-authored-by: OmX --- python/sglang/srt/environ.py | 5 - .../srt/layers/attention/nsa/nsa_indexer.py | 317 ++++-------------- 2 files changed, 56 insertions(+), 266 deletions(-) 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())