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
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

View File

@@ -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