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

Bound NSA MQA logits peak memory

Browse Source 
Paged and CP-ragged NSA indexer paths could materialize q x context fp32
MQA-logits buffers large enough to OOM high-cache-hit bs>1 prefill batches.
Port the syh branch chunking logic so paged and ragged paths split logits by
query rows when the estimated logits buffer exceeds the current free-memory
budget.

The free-memory query is cached on forward_batch so the OOM guard uses current
free memory without adding a torch.cuda.mem_get_info host sync on every layer.
The only new env kept from the syh commits is
SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS, which forces chunking for equivalence
validation.

Constraint: DeepGEMM fp8_mqa_logits still materializes fp32 logits internally, so limiting q rows is the least invasive way to cap peak memory
Rejected: Carry unrelated syh envs for page trace/source-fingerprint strictness | not part of the logits peak-memory fix
Rejected: Static mem_fraction-only budget | overestimates logits headroom shared with other forward activations
Confidence: medium
Scope-risk: moderate
Directive: Keep chunking row-split only; changing K/context partitioning needs topk_transform equivalence validation
Related: 40a0389a9c feat(nsa): chunk paged + CP-ragged MQA-logits by current-free-mem budget
Related: 108fa1f538 perf(nsa): cache MQA-logits free-mem budget per-forward
Tested: Local py_compile for environ.py and nsa_indexer.py
Tested: Remote g0034 cjy-glm5-new py_compile for environ.py and nsa_indexer.py
Not-tested: CUDA ETE run with forced SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS equivalence check
Not-tested: Full high-cache-hit bs>1 prefill OOM regression workload
Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-11 03:03:01 +08:00
parent cc908dd556
commit 4e49751406
2 changed files with 266 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

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