From 1b99de7459ea6d4ecd9062ae0982a046ccf802ba Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Fri, 5 Jun 2026 23:48:08 +0800 Subject: [PATCH] Expose gated evidence for CP shared-KV bs>1 debugging The bs>1 prefill path has multiple coupled stages: scheduler admission, page-aligned batch planning, tensor splitting, direct cache writes, index top-k, MLA reuse, and disaggregated KV handoff. Add a default-off, rate-limited debug channel so production ETE runs can identify where batching or metadata semantics diverge without permanently increasing hot-path log volume. Constraint: Logs must be default-off and rate-limited because these paths execute per-rank and often per-layer. Rejected: Always-on INFO logs | would flood logs and add CPU overhead during normal prefill. Rejected: Only scheduler-side logging | insufficient to distinguish planner, index, MLA, and transfer handoff failures. Confidence: medium Scope-risk: moderate Directive: Keep bs>1 debug evidence env-gated; do not add unconditional per-layer or per-token logs in these paths. Tested: Local py_compile for touched files Tested: git diff --check Tested: Remote py_compile and targeted NSA CP utility tests: 5 passed Not-tested: Full ETE correctness with debug disabled --- python/sglang/srt/disaggregation/prefill.py | 62 +++++++++++ python/sglang/srt/environ.py | 2 + .../srt/layers/attention/nsa/nsa_indexer.py | 38 +++++++ .../sglang/srt/layers/attention/nsa/utils.py | 102 +++++++++++++++++- .../srt/layers/attention/nsa_backend.py | 48 +++++++++ python/sglang/srt/managers/scheduler.py | 46 ++++++++ .../unit/layers/test_nsa_cp_utils.py | 33 ++++++ 7 files changed, 326 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index b150436e5..d9c01ac05 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -68,6 +68,24 @@ def _cp_draft_shared_kv_debug(message: str, *args) -> None: logger.info("[CP_DRAFT_SHARED_KV] " + message, *args) +_CP_SHARED_KV_BS_GT1_PREFILL_DEBUG_COUNTS = {} + + +def _cp_shared_kv_bs_gt1_prefill_debug( + key: str, + message: str, + *args, +) -> None: + if not envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get(): + return + limit = int(envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG_LIMIT.get()) + count = _CP_SHARED_KV_BS_GT1_PREFILL_DEBUG_COUNTS.get(key, 0) + if limit > 0 and count >= limit: + return + _CP_SHARED_KV_BS_GT1_PREFILL_DEBUG_COUNTS[key] = count + 1 + logger.info("[CP_SHARED_KV_BS_GT1_DEBUG] event=%s " + message, key, *args) + + def _seq_summary(values) -> str: if values is None: return "None" @@ -622,6 +640,28 @@ class SchedulerDisaggregationPrefillMixin: logprob_pt = 0 # Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue next_token_ids = result.next_token_ids.tolist() + if envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get(): + spec_info = getattr(batch, "spec_info", None) + spec_hidden = getattr(spec_info, "hidden_states", None) + spec_topk = getattr(spec_info, "topk_index", None) + _cp_shared_kv_bs_gt1_prefill_debug( + "prefill_result_handoff", + "bs=%s rids=%s extend_lens=%s prefix_lens=%s next_token_ids=%s " + "has_spec=%s hidden_shape=%s topk_shape=%s out_cache_tokens=%s " + "inflight_before=%s", + len(batch.reqs), + [req.rid for req in batch.reqs[:8]], + list(getattr(batch, "extend_lens", []) or []), + list(getattr(batch, "prefix_lens", []) or []), + next_token_ids[:8], + spec_info is not None, + tuple(spec_hidden.shape) if spec_hidden is not None else None, + tuple(spec_topk.shape) if spec_topk is not None else None, + int(batch.out_cache_loc.numel()) + if getattr(batch, "out_cache_loc", None) is not None + else None, + len(self.disagg_prefill_inflight_queue), + ) if batch.return_logprob: if logits_output.next_token_logprobs is not None: logits_output.next_token_logprobs = ( @@ -989,6 +1029,28 @@ class SchedulerDisaggregationPrefillMixin: getattr(req, "already_computed", None), draft_prefix_overlap, ) + if envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get(): + _cp_shared_kv_bs_gt1_prefill_debug( + "send_kv_chunk", + "rid=%s room=%s start_idx=%s end_idx=%s last_chunk=%s " + "page_size=%s pages=%s state_pages=%s prefix_len=%s " + "host_hit_length=%s extend_input_len=%s fill_len=%s " + "origin_input_len=%s has_draft_pool=%s", + req.rid, + req.bootstrap_room, + start_idx, + end_idx, + last_chunk, + page_size, + _seq_summary(page_indices), + _seq_summary(state_indices), + prefix_len, + host_hit_length, + getattr(req, "extend_input_len", None), + len(req.fill_ids), + len(req.origin_input_ids), + has_draft_pool, + ) if has_draft_pool and draft_prefix_overlap > 0: _cp_draft_shared_kv_debug( "prefill_send_cachehit_draft_prefix rid=%s room=%s " diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 6a2280404..7e0ac848d 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -203,6 +203,8 @@ class Envs: SGLANG_FORCE_SHUTDOWN = EnvBool(False) SGLANG_DEBUG_MEMORY_POOL = EnvBool(False) SGLANG_DEBUG_CP_SHARED_KV = EnvBool(False) + SGLANG_CP_SHARED_KV_BS_GT1_DEBUG = EnvBool(False) + SGLANG_CP_SHARED_KV_BS_GT1_DEBUG_LIMIT = EnvInt(128) SGLANG_DEBUG_SORT_NVTX = EnvBool(False) SGLANG_DEBUG_MOE_SORT_NVTX = EnvBool(False) SGLANG_CP_SHARED_KV_CURRENT_REUSE = 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 486efc92a..2b0eb52d6 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -70,11 +70,13 @@ from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.attention.nsa.utils import ( cp_all_gather_rerange_output, cp_split_and_rebuild_data, + cp_shared_kv_bs_gt1_debug_enabled, get_cp_shared_kv_batch_plan, get_cp_shared_kv_local_out_cache_loc, get_cp_shared_kv_local_physical_out_cache_loc, is_nsa_enable_prefill_cp, is_nsa_prefill_cp_in_seq_split, + log_cp_shared_kv_bs_gt1_debug, nsa_use_prefill_cp, raise_cp_shared_kv_direct_write_error, select_cp_local_valid_rows_for_cache_write, @@ -2035,6 +2037,30 @@ class Indexer(MultiPlatformOp): and current_only_batch ) page_table_1 = None if current_only else metadata.get_page_table_1() + if cp_shared_kv_bs_gt1_debug_enabled(): + log_cp_shared_kv_bs_gt1_debug( + "index_topk_batch", + "layer=%s batch_size=%s q_tokens=%s weights_tokens=%s " + "uses_compute_q=%s current_only=%s has_current_index=%s " + "has_shared_index_buffer=%s page_table_shape=%s " + "kv_prev=%s kv_next=%s q_prev=%s q_next=%s " + "valid_q_prev=%s valid_q_next=%s", + layer_id, + batch_size, + int(q_fp8.shape[0]), + int(weights.shape[0]), + query_lengths.uses_compute_query_rows, + current_only, + current_index_kv is not None, + shared_index_buffer is not None, + tuple(page_table_1.shape) if page_table_1 is not None else None, + request_kv_len_prev, + request_kv_len_next, + request_actual_seq_q_prev, + request_actual_seq_q_next, + request_valid_seq_q_prev, + request_valid_seq_q_next, + ) def collect_segment( *, @@ -2344,6 +2370,18 @@ class Indexer(MultiPlatformOp): physical_out_loc = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch) if physical_out_loc is None: return False + if cp_shared_kv_bs_gt1_debug_enabled(): + log_cp_shared_kv_bs_gt1_debug( + "index_direct_write", + "layer=%s local_key_shape=%s local_out_tokens=%s " + "physical_tokens=%s pool=%s dtype=%s", + layer_id, + tuple(local_key.shape), + int(local_out_loc.numel()), + int(physical_out_loc.numel()), + forward_batch.token_to_kv_pool.__class__.__name__, + local_key.dtype, + ) log_cp_draft_shared_kv_debug( "index_write", "index_write layer=%s tokens=%s physical_tokens=%s pool=%s key_shape=%s", diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index cf6eea7f0..d7363afa0 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -31,6 +31,29 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) _CP_DRAFT_SHARED_KV_DEBUG_COUNTS = {} +_CP_SHARED_KV_BS_GT1_DEBUG_COUNTS = {} + + +def cp_shared_kv_bs_gt1_debug_enabled() -> bool: + return envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get() + + +def log_cp_shared_kv_bs_gt1_debug( + key: str, + message: str, + *args, + limit: Optional[int] = None, +) -> None: + if not cp_shared_kv_bs_gt1_debug_enabled(): + return + if limit is None: + limit = envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG_LIMIT.get() + limit = int(limit) + count = _CP_SHARED_KV_BS_GT1_DEBUG_COUNTS.get(key, 0) + if limit > 0 and count >= limit: + return + _CP_SHARED_KV_BS_GT1_DEBUG_COUNTS[key] = count + 1 + logger.info("[CP_SHARED_KV_BS_GT1_DEBUG] event=%s " + message, key, *args) def log_cp_draft_shared_kv_debug( @@ -697,7 +720,7 @@ def build_batch_page_aligned_in_seq_split_plan( flat_segment_request_ids.extend([req_id] * cp_segment_num) flat_segment_offsets.extend(split_prefix_list) - return CPSharedKVBatchPlan( + plan = CPSharedKVBatchPlan( batch_size=len(request_extend_lens), page_size=page_size, cp_size=cp_size, @@ -765,6 +788,29 @@ def build_batch_page_aligned_in_seq_split_plan( request_compute_seq_q_prev=request_compute_seq_q_prev, request_compute_seq_q_next=request_compute_seq_q_next, ) + if cp_shared_kv_bs_gt1_debug_enabled(): + log_cp_shared_kv_bs_gt1_debug( + "batch_plan", + "cp_rank=%s cp_size=%s bs=%s page_size=%s extend_lens=%s " + "prefix_lens=%s valid_pages=%s compute_pages=%s " + "valid_local_tokens=%s compute_local_tokens=%s " + "compute_padding=%s padding_tokens=%s last_owner=%s last_offset=%s", + cp_rank, + cp_size, + plan.batch_size, + page_size, + request_extend_lens, + request_prefix_lens, + request_valid_padded_pages, + request_compute_padded_pages, + request_valid_rank_local_tokens, + request_compute_rank_local_tokens, + plan.compute_padding_enabled, + request_compute_padding_tokens, + request_last_token_owner, + request_last_token_local_offset, + ) + return plan def get_cp_shared_kv_batch_plan(forward_batch: "ForwardBatch"): @@ -910,8 +956,25 @@ def split_tensor_by_cp_batch_plan( local_chunks.extend(req_segments[int(index)] for index in zigzag_index) if not local_chunks: - return tensor.new_empty((0, *tensor.shape[1:])) - return torch.cat(local_chunks, dim=0).view(-1, *tensor.shape[1:]) + result = tensor.new_empty((0, *tensor.shape[1:])) + else: + result = torch.cat(local_chunks, dim=0).view(-1, *tensor.shape[1:]) + if cp_shared_kv_bs_gt1_debug_enabled(): + log_cp_shared_kv_bs_gt1_debug( + f"split_tensor:{mode}:{split_kind}", + "mode=%s split_kind=%s bs=%s input_tokens=%s expected_tokens=%s " + "target_lens=%s local_rows=%s compute_padding=%s static_padded=%s", + mode, + split_kind, + batch_size, + input_tokens, + expected_tokens, + [int(x) for x in request_target_lens], + int(result.shape[0]), + compute_padding_enabled, + static_padded_tokens, + ) + return result def _get_cp_local_valid_row_indices_cache(forward_batch, plan, device: torch.device): @@ -1010,6 +1073,17 @@ def select_cp_local_valid_rows_for_cache_write( f"local_rows={local_rows} expected_compute_rows={expected_compute_rows} " f"valid_rows={int(indices.numel())}" ) + if cp_shared_kv_bs_gt1_debug_enabled(): + log_cp_shared_kv_bs_gt1_debug( + "valid_rows_select", + "local_rows=%s expected_compute_rows=%s valid_rows=%s " + "tensor_shape=%s dtype=%s", + local_rows, + expected_compute_rows, + int(indices.numel()), + tuple(local_tensor.shape), + local_tensor.dtype, + ) if indices.numel() == local_rows: return local_tensor if indices.numel() == 0: @@ -1740,11 +1814,30 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"): forward_batch, out_cache_loc.contiguous(), ) + valid_locs = local_out_cache_loc[local_out_cache_loc > 0] + if cp_shared_kv_bs_gt1_debug_enabled(): + log_cp_shared_kv_bs_gt1_debug( + "local_out_cache_loc", + "cp_rank=%s cp_size=%s batch_plan=%s compute_padding=%s " + "split_tokens=%s out_cache_tokens=%s local_tokens=%s valid_local_tokens=%s " + "page_size=%s forward_mode=%s", + layout.cp_rank, + layout.cp_size, + batch_plan is not None, + bool(getattr(batch_plan, "compute_padding_enabled", False)) + if batch_plan is not None + else False, + split_tokens, + out_cache_tokens, + int(local_out_cache_loc.numel()), + int(valid_locs.numel()), + layout.page_size, + getattr(forward_batch, "forward_mode", None), + ) if local_out_cache_loc.numel() == 0: forward_batch.cp_local_out_cache_loc = local_out_cache_loc return local_out_cache_loc - valid_locs = local_out_cache_loc[local_out_cache_loc > 0] if valid_locs.numel() > 0 and not torch.all(layout.owned_by_this_rank(valid_locs)): raise_cp_shared_kv_direct_write_error( "local_loc_owner_mismatch", @@ -1753,7 +1846,6 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"): layout.cp_size, layout.page_size, ) - forward_batch.cp_local_out_cache_loc = local_out_cache_loc return local_out_cache_loc diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index db3910fdf..d55b83f99 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -53,9 +53,11 @@ from sglang.srt.layers.attention.nsa.utils import ( can_nsa_prefill_cp_round_robin_split, cp_split_and_rebuild_data, compute_nsa_seqlens, + cp_shared_kv_bs_gt1_debug_enabled, get_cp_shared_kv_batch_plan, get_cp_shared_kv_local_out_cache_loc, is_nsa_enable_prefill_cp, + log_cp_shared_kv_bs_gt1_debug, nsa_cp_round_robin_split_data, nsa_cp_round_robin_split_q_seqs, nsa_use_prefill_cp, @@ -630,6 +632,20 @@ class NativeSparseAttnBackend( cache_loc, layout, ) + if cp_shared_kv_bs_gt1_debug_enabled(): + log_cp_shared_kv_bs_gt1_debug( + "mla_direct_write_filter", + "cp_rank=%s forward_mode=%s cache_tokens=%s owned_tokens=%s " + "k_shape=%s k_rope_shape=%s physical_tokens=%s dtype=%s", + layout.cp_rank, + forward_batch.forward_mode, + int(cache_loc.numel()), + int(physical_cache_loc.numel()), + tuple(k.shape), + tuple(k_rope.shape), + int(physical_cache_loc.numel()), + k.dtype, + ) if cp_shared_kv_debug_enabled(): valid_mask = cache_loc > 0 owned_count = int(owned_mask.sum().item()) @@ -1880,6 +1896,38 @@ class NativeSparseAttnBackend( if page_table_1 is not None else None, ) + if cp_shared_kv_bs_gt1_debug_enabled(): + log_cp_shared_kv_bs_gt1_debug( + "mla_forward_path", + "cp_rank=%s layer=%s is_draft=%s topk_transform=%s " + "prefix_lens=%s extend_lens=%s q_rows=%s k_rows=%s " + "k_rope_rows=%s kv_cache_shape=%s kv_cache_dtype=%s " + "page_table_shape=%s batch_plan=%s compute_padding=%s " + "has_prefetcher=%s can_current_reuse=%s current_rows=%s", + forward_batch.cp_shared_kv_layout.cp_rank, + layer.layer_id, + is_draft_mla_input, + getattr(topk_transform_method, "name", topk_transform_method), + [ + int(x) + for x in getattr(forward_batch, "extend_prefix_lens_cpu", []) + ], + [ + int(x) + for x in getattr(forward_batch, "extend_seq_lens_cpu", []) + ], + int(q_nope.shape[0]), + int(k.shape[0]) if k is not None else None, + int(k_rope.shape[0]) if k_rope is not None else None, + tuple(kv_cache.shape) if kv_cache is not None else None, + kv_cache.dtype if kv_cache is not None else None, + tuple(page_table_1.shape) if page_table_1 is not None else None, + batch_plan is not None, + compute_padding_current, + mla_prefetcher is not None, + can_reuse_current_kv, + current_kv_rows_for_reuse, + ) if can_reuse_current_kv: assert k is not None and k_rope is not None assert current_kv_rows_for_reuse is not None diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c64a0ecfa..9d2e66b3d 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -258,6 +258,24 @@ def _cp_draft_shared_kv_debug(message: str, *args) -> None: logger.info("[CP_DRAFT_SHARED_KV] " + message, *args) +_CP_SHARED_KV_BS_GT1_SCHED_DEBUG_COUNTS = {} + + +def _cp_shared_kv_bs_gt1_scheduler_debug( + key: str, + message: str, + *args, +) -> None: + if not envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get(): + return + limit = int(envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG_LIMIT.get()) + count = _CP_SHARED_KV_BS_GT1_SCHED_DEBUG_COUNTS.get(key, 0) + if limit > 0 and count >= limit: + return + _CP_SHARED_KV_BS_GT1_SCHED_DEBUG_COUNTS[key] = count + 1 + logger.info("[CP_SHARED_KV_BS_GT1_DEBUG] event=%s " + message, key, *args) + + _is_npu = is_npu() @@ -2490,6 +2508,34 @@ class Scheduler( exc, ) return None + if ( + getattr(self.server_args, "enable_nsa_prefill_cp_shared_kv", False) + and envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get() + ): + _cp_shared_kv_bs_gt1_scheduler_debug( + "scheduler_prefill_batch", + "bs=%s extend_lens=%s prefix_lens=%s seq_lens=%s " + "out_cache_tokens=%s chunked_req=%s enable_bs_gt1=%s " + "max_batch_reqs=%s max_total_extend=%s", + len(can_run_list), + list(getattr(new_batch, "extend_lens", []) or []), + list(getattr(new_batch, "prefix_lens", []) or []), + [ + int(x) + for x in getattr(new_batch, "seq_lens_cpu", torch.tensor([])).tolist() + ], + int(new_batch.out_cache_loc.numel()) + if getattr(new_batch, "out_cache_loc", None) is not None + else None, + self.chunked_req is not None, + getattr(self.server_args, "enable_cp_shared_kv_prefill_bs_gt1", None), + getattr(self.server_args, "cp_shared_kv_prefill_max_batch_requests", None), + getattr( + self.server_args, + "cp_shared_kv_prefill_max_total_extend_tokens", + None, + ), + ) # Record prefill stats for logging after forward new_batch.prefill_stats = PrefillStats.from_adder( diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index fe32030ed..aac6fbad5 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -6,6 +6,8 @@ import sys from types import SimpleNamespace from unittest.mock import patch +from sglang.srt.environ import envs +from sglang.srt.layers.attention.nsa import utils as nsa_utils from sglang.srt.layers.attention.nsa.utils import ( NSAContextParallelMetadata, PageAlignedCacheExtent, @@ -565,6 +567,37 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(plan.request_compute_seq_q_prev, [64]) self.assertEqual(plan.request_compute_seq_q_next, [0]) + def test_bs_gt1_debug_log_is_env_gated_and_limited(self): + nsa_utils._CP_SHARED_KV_BS_GT1_DEBUG_COUNTS.clear() + + with envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.override(False): + with patch.object(nsa_utils.logger, "info") as info: + nsa_utils.log_cp_shared_kv_bs_gt1_debug( + "unit_test", + "bs=%s", + 2, + ) + self.assertFalse(info.called) + + with envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.override(True): + with envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG_LIMIT.override(1): + with patch.object(nsa_utils.logger, "info") as info: + nsa_utils.log_cp_shared_kv_bs_gt1_debug( + "unit_test", + "bs=%s", + 2, + ) + nsa_utils.log_cp_shared_kv_bs_gt1_debug( + "unit_test", + "bs=%s", + 3, + ) + self.assertEqual(info.call_count, 1) + self.assertIn( + "[CP_SHARED_KV_BS_GT1_DEBUG]", + info.call_args.args[0], + ) + def test_index_topk_batch_lengths_follow_actual_q_rows_not_compute_alias(self): import torch