chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Revert NSA MQA logits chunking while performance is unresolved

Browse Source 
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 <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-11 03:04:03 +08:00
parent 4e49751406
commit e0ea8a485c
2 changed files with 56 additions and 266 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/srt/environ.py
View File

@@ -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)

317
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -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())