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

Expose gated evidence for CP shared-KV bs>1 debugging

Browse Source 
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
This commit is contained in:
laoyao0822
2026-06-05 23:48:08 +08:00
parent 6eea77e5e9
commit 1b99de7459
7 changed files with 326 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/prefill.py
+62
View File
@@ -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 "
python/sglang/srt/environ.py
+2
View File
@@ -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)
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+38
View File
@@ -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",
python/sglang/srt/layers/attention/nsa/utils.py
+97 -5
View File
@@ -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
python/sglang/srt/layers/attention/nsa_backend.py
+48
View File
@@ -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
python/sglang/srt/managers/scheduler.py
+46
View File
@@ -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(