diff --git a/python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py b/python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py index 5a3f92427..93228262a 100644 --- a/python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py +++ b/python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py @@ -32,6 +32,7 @@ class CPSharedKVPrefillBufferEstimatorContext: logprob_chunk_enabled: bool logprob_chunk_size: int bs_gt1_l1_prefetch_enabled: bool = False + mqa_logits_chunk_max_rows: int = 0 @dataclass(frozen=True) @@ -43,6 +44,7 @@ class CPSharedKVPrefillBufferEstimate: materialize_peak_bytes: int prefetch_peak_bytes: int logits_peak_bytes: int + mqa_logits_peak_bytes: int remap_peak_bytes: int transfer_descriptor_peak_bytes: int backup_descriptor_peak_bytes: int @@ -137,6 +139,7 @@ def estimate_cp_shared_kv_prefill_buffer_bytes( logprob_chunk_enabled=False, logprob_chunk_size=2048, bs_gt1_l1_prefetch_enabled=False, + mqa_logits_chunk_max_rows=0, ) if page_size <= 0: @@ -175,6 +178,15 @@ def estimate_cp_shared_kv_prefill_buffer_bytes( logits_rows = min(logits_rows, int(context.logprob_chunk_size)) vocab_shard = _vocab_shard_size(context.model_config, context.tp_size) logits_peak_bytes = logits_rows * vocab_shard * int(logits_dtype_bytes) + mqa_q_rows = sum(ceil_paged_tokens(tokens, page_size) for tokens in extend_lens) + mqa_k_rows = sum( + ceil_paged_tokens(prefix, page_size) + ceil_paged_tokens(extend, page_size) + for prefix, extend in zip(prefix_lens, extend_lens) + ) + mqa_chunk_rows = int(getattr(context, "mqa_logits_chunk_max_rows", 0) or 0) + if mqa_chunk_rows > 0: + mqa_q_rows = min(mqa_q_rows, mqa_chunk_rows) + mqa_logits_peak_bytes = mqa_q_rows * mqa_k_rows * 4 transfer_descriptor_peak_bytes = total_prefix_pages * int(descriptor_bytes) backup_descriptor_peak_bytes = total_extend_pages * int(descriptor_bytes) @@ -185,6 +197,7 @@ def estimate_cp_shared_kv_prefill_buffer_bytes( layer_forward_peak_bytes = ( materialize_peak_bytes + + mqa_logits_peak_bytes + remap_peak_bytes + prefetch_peak_bytes + backup_descriptor_peak_bytes @@ -211,6 +224,7 @@ def estimate_cp_shared_kv_prefill_buffer_bytes( materialize_peak_bytes=materialize_peak_bytes, prefetch_peak_bytes=prefetch_peak_bytes, logits_peak_bytes=logits_peak_bytes, + mqa_logits_peak_bytes=mqa_logits_peak_bytes, remap_peak_bytes=remap_peak_bytes, transfer_descriptor_peak_bytes=transfer_descriptor_peak_bytes, backup_descriptor_peak_bytes=backup_descriptor_peak_bytes, diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 99a696b92..2752c6b0e 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -698,7 +698,7 @@ class PrefillAdder: "admission_max_buffer_size", "stop rid=%s projected=%s limit=%s batch_size=%s " "layer_forward=%s logits_window=%s load_back_window=%s " - "materialize=%s prefetch=%s logits=%s remap=%s " + "materialize=%s prefetch=%s logits=%s mqa_logits=%s remap=%s " "transfer_desc=%s backup_desc=%s", rid, estimate.total_peak_bytes, @@ -710,6 +710,7 @@ class PrefillAdder: estimate.materialize_peak_bytes, estimate.prefetch_peak_bytes, estimate.logits_peak_bytes, + estimate.mqa_logits_peak_bytes, estimate.remap_peak_bytes, estimate.transfer_descriptor_peak_bytes, estimate.backup_descriptor_peak_bytes, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 29fa1dc4d..49d42348b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -535,6 +535,9 @@ class Scheduler( # gates off bs>1. Keep the estimate aligned with runtime until that # path is enabled. bs_gt1_l1_prefetch_enabled=False, + mqa_logits_chunk_max_rows=( + envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS.get() + ), ) def maybe_smoke_check_cp_shared_kv_prefill_buffer(self) -> None: diff --git a/test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py b/test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py index 6a09378a7..e631867d3 100644 --- a/test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py +++ b/test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py @@ -1,3 +1,4 @@ +from dataclasses import replace from types import SimpleNamespace import torch @@ -44,6 +45,7 @@ def test_estimator_uses_stream_aware_peak_instead_of_independent_max(): assert estimate.prefetch_peak_bytes > 0 assert estimate.layer_forward_peak_bytes == ( estimate.materialize_peak_bytes + + estimate.mqa_logits_peak_bytes + estimate.remap_peak_bytes + estimate.prefetch_peak_bytes + estimate.backup_descriptor_peak_bytes @@ -81,6 +83,39 @@ def test_estimator_keeps_bs_gt1_prefetch_zero_until_enabled(): assert estimate.total_peak_bytes >= estimate.materialize_peak_bytes +def test_estimator_counts_mqa_logits_peak_from_extend_and_context_rows(): + context = CPSharedKVPrefillBufferEstimatorContext( + kvcache=_fake_kvcache(), + model_config=SimpleNamespace(vocab_size=32), + tp_size=1, + page_size=64, + logprob_chunk_enabled=False, + logprob_chunk_size=2048, + bs_gt1_l1_prefetch_enabled=False, + ) + estimate = estimate_cp_shared_kv_prefill_buffer_bytes( + page_size=64, + batch_size=2, + prefix_lens=[40384, 8192], + extend_lens=[4096, 2048], + context=context, + ) + + q_rows = 4096 + 2048 + k_rows = 40384 + 4096 + 8192 + 2048 + assert estimate.mqa_logits_peak_bytes == q_rows * k_rows * 4 + assert estimate.total_peak_bytes >= estimate.mqa_logits_peak_bytes + + chunked = estimate_cp_shared_kv_prefill_buffer_bytes( + page_size=64, + batch_size=2, + prefix_lens=[40384, 8192], + extend_lens=[4096, 2048], + context=replace(context, mqa_logits_chunk_max_rows=1024), + ) + assert chunked.mqa_logits_peak_bytes == 1024 * k_rows * 4 + + def test_smoke_check_allocates_and_releases_probe_with_device_module(monkeypatch): events = []