diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index dafeac92e..3fbcccf79 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -225,6 +225,9 @@ class Envs: 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) + # Optional hard cap for one fp32 MQA-logits chunk, in decimal GB. 0 = use + # row cap or memory budget. Mutually exclusive with *_CHUNK_MAX_ROWS. + SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB = EnvFloat(0.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 10f0cfd69..5ea47fa06 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1183,10 +1183,21 @@ class Indexer(MultiPlatformOp): 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()) + explicit_max_gb = float(envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB.get()) + if explicit_max_rows > 0 and explicit_max_gb > 0: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][mqa_logits_chunk] " + "SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS and " + "SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB are mutually exclusive" + ) 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 + if explicit_max_gb > 0: + explicit_max_bytes = int(explicit_max_gb * 1_000_000_000) + max_rows = max(1, explicit_max_bytes // max(bytes_per_row, 1)) + return min(max_rows, max(1, num_q)) max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1))) return min(max_rows, max(1, num_q)) 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 93228262a..e2ea9bfb5 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 @@ -33,6 +33,9 @@ class CPSharedKVPrefillBufferEstimatorContext: logprob_chunk_size: int bs_gt1_l1_prefetch_enabled: bool = False mqa_logits_chunk_max_rows: int = 0 + mqa_logits_chunk_max_bytes: int = 0 + mqa_logits_budget_bytes: int = 0 + cp_size: int = 1 @dataclass(frozen=True) @@ -112,6 +115,120 @@ def _vocab_shard_size(model_config: object | None, tp_size: int) -> int: return -(-vocab_size // max(int(tp_size), 1)) +def _cp_segment_valid_lengths( + *, + extend_len: int, + page_size: int, + cp_size: int, +) -> list[int]: + full_pages = int(extend_len) // int(page_size) + tail_tokens = int(extend_len) % int(page_size) + padded_pages = ceil_pages(extend_len, page_size) + cp_segment_num = int(cp_size) * 2 + if cp_segment_num <= 0: + return [] + + base_units = padded_pages // cp_segment_num + remainder_units = padded_pages % cp_segment_num + unit_counts = [ + base_units + (1 if idx < remainder_units else 0) + for idx in range(cp_segment_num) + ] + + split_list: list[int] = [] + page_cursor = 0 + for unit_count in unit_counts: + token_count = 0 + for _ in range(unit_count): + token_count += page_size if page_cursor < full_pages else tail_tokens + page_cursor += 1 + split_list.append(token_count) + return split_list + + +def _mqa_logits_row_cap( + *, + q_rows: int, + k_rows: int, + explicit_chunk_rows: int, + explicit_chunk_bytes: int, + budget_bytes: int, +) -> int: + if q_rows <= 0 or k_rows <= 0: + return 0 + if explicit_chunk_rows > 0 and explicit_chunk_bytes > 0: + raise ValueError( + "mqa_logits_chunk_max_rows and mqa_logits_chunk_max_bytes are " + "mutually exclusive" + ) + if explicit_chunk_rows > 0: + return min(q_rows, max(1, explicit_chunk_rows)) + if explicit_chunk_bytes > 0: + bytes_per_row = k_rows * 4 + return min(q_rows, max(1, int(explicit_chunk_bytes) // max(bytes_per_row, 1))) + if budget_bytes > 0: + bytes_per_row = k_rows * 4 + return min(q_rows, max(1, int(budget_bytes) // max(bytes_per_row, 1))) + return q_rows + + +def _estimate_cp_split_mqa_logits_peak_bytes( + *, + prefix_lens: Sequence[int], + extend_lens: Sequence[int], + page_size: int, + cp_size: int, + explicit_chunk_rows: int, + explicit_chunk_bytes: int, + budget_bytes: int, +) -> int: + if cp_size <= 1: + q_rows = sum(ceil_paged_tokens(tokens, page_size) for tokens in extend_lens) + k_rows = sum( + ceil_paged_tokens(prefix, page_size) + ceil_paged_tokens(extend, page_size) + for prefix, extend in zip(prefix_lens, extend_lens) + ) + row_cap = _mqa_logits_row_cap( + q_rows=q_rows, + k_rows=k_rows, + explicit_chunk_rows=explicit_chunk_rows, + explicit_chunk_bytes=explicit_chunk_bytes, + budget_bytes=budget_bytes, + ) + return row_cap * k_rows * 4 + + cp_segment_num = cp_size * 2 + max_peak_bytes = 0 + for cp_rank in range(cp_size): + q_rows = 0 + k_rows = 0 + mirror_idx = cp_segment_num - cp_rank - 1 + for prefix_len, extend_len in zip(prefix_lens, extend_lens): + split_list = _cp_segment_valid_lengths( + extend_len=int(extend_len), + page_size=page_size, + cp_size=cp_size, + ) + segment_end = 0 + for segment_idx, segment_len in enumerate(split_list): + segment_end += int(segment_len) + if segment_idx not in (cp_rank, mirror_idx) or segment_len <= 0: + continue + q_rows += int(segment_len) + # Match the fused CP MQA path: each owned segment materializes a + # compact K range up to this segment's logical end. + k_rows += ceil_paged_tokens(prefix_len, page_size) + segment_end + row_cap = _mqa_logits_row_cap( + q_rows=q_rows, + k_rows=k_rows, + explicit_chunk_rows=explicit_chunk_rows, + explicit_chunk_bytes=explicit_chunk_bytes, + budget_bytes=budget_bytes, + ) + max_peak_bytes = max(max_peak_bytes, row_cap * k_rows * 4) + return max_peak_bytes + + def estimate_cp_shared_kv_prefill_buffer_bytes( *, page_size: int, @@ -140,6 +257,9 @@ def estimate_cp_shared_kv_prefill_buffer_bytes( logprob_chunk_size=2048, bs_gt1_l1_prefetch_enabled=False, mqa_logits_chunk_max_rows=0, + mqa_logits_chunk_max_bytes=0, + mqa_logits_budget_bytes=0, + cp_size=1, ) if page_size <= 0: @@ -178,15 +298,19 @@ 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_logits_peak_bytes = _estimate_cp_split_mqa_logits_peak_bytes( + prefix_lens=prefix_lens, + extend_lens=extend_lens, + page_size=page_size, + cp_size=int(getattr(context, "cp_size", 1) or 1), + explicit_chunk_rows=int( + getattr(context, "mqa_logits_chunk_max_rows", 0) or 0 + ), + explicit_chunk_bytes=int( + getattr(context, "mqa_logits_chunk_max_bytes", 0) or 0 + ), + budget_bytes=int(getattr(context, "mqa_logits_budget_bytes", 0) or 0), ) - 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) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 49d42348b..1cf95e96c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -524,6 +524,48 @@ class Scheduler( def make_cp_shared_kv_prefill_buffer_estimator_context( self, ) -> CPSharedKVPrefillBufferEstimatorContext: + mqa_logits_chunk_max_rows = int( + envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS.get() + ) + mqa_logits_chunk_max_gb = float( + envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB.get() + ) + if mqa_logits_chunk_max_rows > 0 and mqa_logits_chunk_max_gb > 0: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][mqa_logits_chunk] " + "SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS and " + "SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB are mutually exclusive" + ) + mqa_logits_chunk_max_bytes = ( + int(mqa_logits_chunk_max_gb * 1_000_000_000) + if mqa_logits_chunk_max_gb > 0 + else 0 + ) + mqa_logits_budget_bytes = 0 + if str(self.device).startswith("cuda"): + try: + total_mem = torch.cuda.get_device_properties(self.device).total_memory + total_mem_budget = int(total_mem * 0.3) + if self.server_args.mem_fraction_static is None: + static_budget = total_mem_budget + else: + static_free_mem = int( + total_mem + * max(0.0, 1.0 - self.server_args.mem_fraction_static) + ) + static_budget = min( + int( + static_free_mem + * envs.SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION.get() + ), + total_mem_budget, + ) + mqa_logits_budget_bytes = max(1, static_budget) + except Exception: + # Keep scheduler admission CUDA-sync-free and non-fatal. A zero + # budget makes the estimator fall back to explicit chunk rows or + # the conservative unchunked peak. + mqa_logits_budget_bytes = 0 return CPSharedKVPrefillBufferEstimatorContext( kvcache=self.token_to_kv_pool_allocator.get_kvcache(), model_config=self.model_config, @@ -535,9 +577,10 @@ 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() - ), + mqa_logits_chunk_max_rows=mqa_logits_chunk_max_rows, + mqa_logits_chunk_max_bytes=mqa_logits_chunk_max_bytes, + mqa_logits_budget_bytes=mqa_logits_budget_bytes, + cp_size=self.attn_cp_size, ) 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 e631867d3..4d1dbd90e 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 @@ -116,6 +116,101 @@ def test_estimator_counts_mqa_logits_peak_from_extend_and_context_rows(): assert chunked.mqa_logits_peak_bytes == 1024 * k_rows * 4 +def test_estimator_counts_mqa_logits_peak_after_cp_in_seq_split(): + 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, + cp_size=8, + ) + estimate = estimate_cp_shared_kv_prefill_buffer_bytes( + page_size=64, + batch_size=1, + prefix_lens=[160_000], + extend_lens=[65_536], + context=context, + ) + + # 65,536 tokens = 1024 pages. With CP=8 the in-seq split has 16 + # page-aligned segments, 4096 valid rows each. A rank owns two zigzag + # segments, so local q rows are 8192, not the full 65,536. + q_rows = 4096 + 4096 + # The fused CP MQA path materializes K per owned segment. For rank 0 this + # is prefix+segment0_end plus prefix+segment15_end. + k_rows = (160_000 + 4096) + (160_000 + 65_536) + assert estimate.mqa_logits_peak_bytes == q_rows * k_rows * 4 + + chunked = estimate_cp_shared_kv_prefill_buffer_bytes( + page_size=64, + batch_size=1, + prefix_lens=[160_000], + extend_lens=[65_536], + context=replace(context, mqa_logits_chunk_max_rows=4096), + ) + assert chunked.mqa_logits_peak_bytes == 4096 * k_rows * 4 + + +def test_estimator_derives_mqa_chunk_rows_from_static_budget(): + prefix_len = 160_000 + extend_len = 65_536 + q_rows = 4096 + 4096 + k_rows = (prefix_len + 4096) + (prefix_len + extend_len) + 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, + cp_size=8, + mqa_logits_budget_bytes=1024 * k_rows * 4, + ) + + estimate = estimate_cp_shared_kv_prefill_buffer_bytes( + page_size=64, + batch_size=1, + prefix_lens=[prefix_len], + extend_lens=[extend_len], + context=context, + ) + + assert q_rows > 1024 + assert estimate.mqa_logits_peak_bytes == 1024 * k_rows * 4 + + +def test_estimator_caps_mqa_logits_with_explicit_gb_budget(): + prefix_len = 160_000 + extend_len = 65_536 + k_rows = (prefix_len + 4096) + (prefix_len + extend_len) + 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, + cp_size=8, + mqa_logits_chunk_max_bytes=2_000_000_000, + ) + + estimate = estimate_cp_shared_kv_prefill_buffer_bytes( + page_size=64, + batch_size=1, + prefix_lens=[prefix_len], + extend_lens=[extend_len], + context=context, + ) + + expected_rows = 2_000_000_000 // (k_rows * 4) + assert estimate.mqa_logits_peak_bytes == expected_rows * k_rows * 4 + + def test_smoke_check_allocates_and_releases_probe_with_device_module(monkeypatch): events = [] diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index 2805ae53b..937a21c4c 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -5854,6 +5854,38 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase): 128, ) + def test_nsa_mqa_logits_chunk_max_gb_caps_rows_by_bytes(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( + 0 + ), envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB.override(1.0): + self.assertEqual( + indexer._mqa_logits_chunk_max_rows( + num_q=8192, + num_k=1_000_000, + logits_budget_bytes=1, + ), + 250, + ) + + def test_nsa_mqa_logits_chunk_rows_and_gb_conflict_fail_fast(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 + ), envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB.override(1.0): + with self.assertRaisesRegex(RuntimeError, "mqa_logits_chunk"): + indexer._mqa_logits_chunk_max_rows( + num_q=8192, + num_k=1_000_000, + logits_budget_bytes=1, + ) + if __name__ == "__main__": unittest.main()