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

Bound CP MQA logits buffers with row chunking

Browse Source 
CP shared-KV bs>1 can build large fp32 MQA-logits temporaries from
DeepGEMM fp8_mqa_logits. The official SGLang path already chunks normal NSA
MQA logits by query rows behind a cached memory budget; carry the same budget
control into our NSA indexer and extend it to CP-ragged topk paths that use
row-wise topk_indices_offset_override.

This keeps the previous one-time cached memory-budget behavior rather than the
recent current-free-mem per-forward variant that regressed performance. A new
optional max-rows env provides an explicit hard cap for debugging or controlled
ETE runs without adding host syncs.

Constraint: DeepGEMM materializes fp32 [q, k] logits internally, so row chunking is the narrowest way to cap temporary memory
Rejected: Restore the reverted syh current-free-mem implementation | it changed hot-path heuristics and showed poor runtime performance
Rejected: Split by K/context dimension | would change topk semantics and require a different transform contract
Confidence: medium
Scope-risk: moderate
Directive: CP-ragged chunking relies on topk_indices_offset_override being row-addressed; do not route non-ragged CP paths through it without separate validation
Tested: Local py_compile for environ.py, nsa_indexer.py, and test_cp_shared_kv_runtime.py
Tested: Remote g0034 cjy-glm5-new py_compile for environ.py, nsa_indexer.py, and test_cp_shared_kv_runtime.py
Tested: Remote pytest TestCpSharedKVTaiMaterializeIntegration, 17 passed
Not-tested: CUDA ETE high-cache-hit bs>1 workload memory/performance after chunking
Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-11 03:15:52 +08:00
parent e0ea8a485c
commit ddc1233955
3 changed files with 195 additions and 68 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 MQA logits are materialized as fp32 [q, k] buffers inside DeepGEMM.
# Lower values split query rows more aggressively to cap peak temporary memory.
SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION = EnvFloat(0.2)
# Optional hard cap for rows per MQA-logits chunk. 0 = use memory budget.
SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_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)

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

@@ -418,10 +418,13 @@ class Indexer(MultiPlatformOp):
# torch.cuda.mem_get_info host sync on the prefill critical path. (upstream PR #25299)
_MQA_LOGITS_BYTES_PER_ELEM = 4
_MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000
_MQA_LOGITS_FREE_MEM_FRACTION = 0.5
_MQA_LOGITS_TOTAL_MEM_FRACTION = 0.3
_mqa_logits_budget_bytes: Dict[int, int] = {}
@staticmethod
def _mqa_logits_free_mem_fraction() -> float:
return envs.SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION.get()
def __init__(
self,
hidden_size: int,
@@ -1124,6 +1127,7 @@ class Indexer(MultiPlatformOp):
return topk_result
def _get_mqa_logits_budget_bytes(self, device_index: int) -> int:
free_mem_fraction = self._mqa_logits_free_mem_fraction()
# Cache the MQA-logits byte budget per device. torch.cuda.mem_get_info
# host-syncs, so query free memory at most once (after the first real
# prefill) and cap it by the workload-independent serving-memory headroom
@@ -1140,7 +1144,7 @@ class Indexer(MultiPlatformOp):
else:
static_free_mem = int(total_mem * max(0.0, 1.0 - mem_fraction_static))
static_budget = min(
int(static_free_mem * self._MQA_LOGITS_FREE_MEM_FRACTION),
int(static_free_mem * free_mem_fraction),
total_mem_budget,
)
static_budget = max(1, static_budget)
@@ -1152,7 +1156,7 @@ class Indexer(MultiPlatformOp):
free_mem, _ = torch.cuda.mem_get_info(device_index)
budget_bytes = min(
int(free_mem * self._MQA_LOGITS_FREE_MEM_FRACTION), static_budget
int(free_mem * free_mem_fraction), static_budget
)
budget_bytes = max(1, budget_bytes)
self._mqa_logits_budget_bytes[device_index] = budget_bytes
@@ -1175,6 +1179,98 @@ class Indexer(MultiPlatformOp):
need_chunk = logits_bytes > logits_budget_bytes
return need_chunk, logits_budget_bytes
def _mqa_logits_chunk_max_rows(
self, num_q: int, num_k: int, logits_budget_bytes: int
) -> int:
explicit_max_rows = int(envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS.get())
if explicit_max_rows > 0:
return min(max(1, explicit_max_rows), max(1, num_q))
bytes_per_row = num_k * self._MQA_LOGITS_BYTES_PER_ELEM
max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1)))
return min(max_rows, max(1, num_q))
def _mqa_logits_topk_ragged_chunked(
self,
metadata: BaseIndexerMetadata,
q_fp8: torch.Tensor,
kv_fp8: Tuple[torch.Tensor, torch.Tensor],
weights: torch.Tensor,
ks: torch.Tensor,
ke: torch.Tensor,
*,
ke_offset: torch.Tensor,
topk_indices_offset_override: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
) -> torch.Tensor:
q_offset = int(q_fp8.shape[0])
if q_offset == 0:
return torch.full(
(0, self.index_topk),
-1,
dtype=torch.int32,
device=q_fp8.device,
)
device_index = q_fp8.device.index
assert device_index is not None, "q_fp8 must be on an indexed CUDA device"
k_offset = int(kv_fp8[0].shape[0])
need_chunk, logits_budget_bytes = self._should_chunk_mqa_logits(
q_offset, k_offset, device_index
)
max_rows = self._mqa_logits_chunk_max_rows(
q_offset, k_offset, logits_budget_bytes
)
need_chunk = need_chunk or max_rows < q_offset
def _topk_for_slice(start: int, end: int) -> torch.Tensor:
with self._with_real_sm_count():
if _is_hip:
from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits
kv, scale = kv_fp8
logits = fp8_mqa_logits(
q_fp8[start:end],
kv,
scale,
weights[start:end],
ks[start:end],
ke[start:end],
)
else:
logits = deep_gemm.fp8_mqa_logits(
q_fp8[start:end],
kv_fp8,
weights[start:end],
ks[start:end],
ke[start:end],
clean_logits=False,
)
return metadata.topk_transform(
logits,
self.index_topk,
ks=ks[start:end],
ke_offset=ke_offset[start:end],
topk_indices_offset_override=topk_indices_offset_override[start:end],
)
if not need_chunk:
return _topk_for_slice(0, q_offset)
topk_result = None
start = 0
while start < q_offset:
end = min(start + max_rows, q_offset)
topk_chunk = _topk_for_slice(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
assert topk_result is not None
return topk_result
def _get_topk_ragged(
self,
enable_dual_stream: bool,
@@ -1286,9 +1382,9 @@ class Indexer(MultiPlatformOp):
return topk_result
# Chunk path
bytes_per_row = k_offset * self._MQA_LOGITS_BYTES_PER_ELEM
max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1)))
max_rows = min(max_rows, q_offset)
max_rows = self._mqa_logits_chunk_max_rows(
q_offset, k_offset, logits_budget_bytes
)
global_topk_offset = metadata.attn_metadata.topk_indices_offset
@@ -1588,23 +1684,16 @@ 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,
topk_result = self._mqa_logits_topk_ragged_chunked(
metadata,
q_fp8,
kv_fp8,
weights,
ks=ks,
cu_seqlens_q=actual_seq_q,
ke=ke,
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 +1742,16 @@ 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,
topk_result = self._mqa_logits_topk_ragged_chunked(
metadata,
q_fp8,
kv_fp8,
weights,
ks=ks,
cu_seqlens_q=actual_seq_q,
ke=ke,
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 +1830,16 @@ 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,
topk_result = self._mqa_logits_topk_ragged_chunked(
metadata,
q_fp8,
kv_fp8,
weights,
ks=ks,
cu_seqlens_q=actual_seq_q,
ke=ke,
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())
@@ -1863,16 +1938,6 @@ class Indexer(MultiPlatformOp):
k_scale = k_scale[:kv_len].view(torch.float32).squeeze(-1).contiguous()
kv_fp8 = (k_fp8, k_scale)
ke = 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_indices_offset_override = None
cu_seqlens_q_topk_override = None
if (
@@ -1884,7 +1949,17 @@ class Indexer(MultiPlatformOp):
# produced a zero offset per query. Reuse `ks` and avoid the
# post-MQA metadata kernels entirely.
topk_indices_offset_override = ks
actual_seq_q_tensor = None
valid_topk_result = self._mqa_logits_topk_ragged_chunked(
metadata,
q_fp8,
kv_fp8,
weights,
ks=ks,
ke=ke,
ke_offset=ke_offset,
topk_indices_offset_override=topk_indices_offset_override,
forward_batch=forward_batch,
)
elif valid_q_count == actual_seq_q and actual_seq_q_cu_tensor is not None:
cu_seqlens_q_topk_override = actual_seq_q_cu_tensor
elif actual_seq_q_tensor is None or valid_q_count != actual_seq_q:
@@ -1894,15 +1969,25 @@ class Indexer(MultiPlatformOp):
cu_seqlens_q_topk_override[1] = actual_seq_q_tensor.reshape(-1)[0]
elif actual_seq_q_tensor.ndim == 0:
actual_seq_q_tensor = actual_seq_q_tensor.reshape(1)
valid_topk_result = metadata.topk_transform(
logits,
self.index_topk,
ks=ks,
cu_seqlens_q=actual_seq_q_tensor,
ke_offset=ke_offset,
topk_indices_offset_override=topk_indices_offset_override,
cu_seqlens_q_topk_override=cu_seqlens_q_topk_override,
)
if topk_indices_offset_override is None:
with self._with_real_sm_count():
logits = deep_gemm.fp8_mqa_logits(
q_fp8,
kv_fp8,
weights,
ks,
ke,
clean_logits=False,
)
valid_topk_result = metadata.topk_transform(
logits,
self.index_topk,
ks=ks,
cu_seqlens_q=actual_seq_q_tensor,
ke_offset=ke_offset,
topk_indices_offset_override=topk_indices_offset_override,
cu_seqlens_q_topk_override=cu_seqlens_q_topk_override,
)
if valid_q_count == actual_seq_q:
topk_result = valid_topk_result
else:

37
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
View File

@@ -5817,6 +5817,43 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase):
self.assertIs(dense_pages, fallback_pages)
logger.warning.assert_not_called()
def test_nsa_mqa_logits_chunk_budget_uses_env_fraction(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import nsa_indexer
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
indexer = object.__new__(Indexer)
indexer._mqa_logits_budget_bytes = {}
with envs.SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION.override(0.25), patch(
"sglang.srt.layers.attention.nsa.nsa_indexer.get_is_capture_mode",
return_value=False,
), patch(
"sglang.srt.layers.attention.nsa.nsa_indexer.get_global_server_args",
return_value=SimpleNamespace(mem_fraction_static=0.5),
), patch.object(
nsa_indexer.torch.cuda, "get_device_properties"
) as props, patch.object(
nsa_indexer.torch.cuda, "mem_get_info", return_value=(80_000, 100_000)
):
props.return_value = SimpleNamespace(total_memory=100_000)
self.assertEqual(indexer._get_mqa_logits_budget_bytes(0), 12_500)
def test_nsa_mqa_logits_chunk_max_rows_overrides_budget_rows(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
indexer = object.__new__(Indexer)
with envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS.override(128):
self.assertEqual(
indexer._mqa_logits_chunk_max_rows(
num_q=1024,
num_k=4096,
logits_budget_bytes=4096 * 4 * 512,
),
128,
)
if __name__ == "__main__":
unittest.main()