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

Standardize CP shared-KV runtime fallback markers

Browse Source 
Current reuse, TAI materialize, and TAI fused MLA store are optimized CP shared-KV paths. Their fallback helpers already emitted warnings, but the messages were not consistently grep-able with the standard CP shared-KV fallback marker.

This moves the marker into the helper layer so individual fallback call sites do not have to remember to include it, and keeps the existing per-reason rate limiting.

Constraint: Runtime fallback logs must be visible without enabling debug logging.

Rejected: Prefix only selected call sites | helper-level prefixing avoids future unmarked fallback messages.

Confidence: high

Scope-risk: narrow

Directive: New CP shared-KV runtime fallback helpers should use [CP_SHARED_KV_FALLBACK] and a component name.

Tested: Local py_compile for cp_shared_kv_runtime.py and test_cp_shared_kv_runtime.py.

Tested: Local extracted assertLogs check for all three fallback helper prefixes.

Tested: Remote g0034 pytest exact runtime helper tests: 2 passed, 3 warnings.

Not-tested: Full cp_shared_kv_runtime.py suite in this slice.

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-05-29 06:02:13 +08:00
parent e09c8256d7
commit dd69c5970c
3 changed files with 66 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
View File

@@ -143,7 +143,11 @@ def _log_current_reuse_fallback(
if count >= limit:
return
_CURRENT_REUSE_FALLBACK_LOG_COUNTS[key] = count + 1
logger.warning(message, *args)
logger.warning(
"[CP_SHARED_KV_FALLBACK][current_reuse] reason=%s " + message,
key,
*args,
)
@dataclass(frozen=True)
@@ -337,7 +341,11 @@ def _log_tai_materialize_fallback(
if count >= limit:
return
_TAI_MATERIALIZE_FALLBACK_LOG_COUNTS[key] = count + 1
logger.warning(message, *args)
logger.warning(
"[CP_SHARED_KV_FALLBACK][tai_materialize] reason=%s " + message,
key,
*args,
)
def _log_tai_fused_mla_store_fallback(
@@ -350,7 +358,11 @@ def _log_tai_fused_mla_store_fallback(
if count >= limit:
return
_TAI_FUSED_MLA_STORE_FALLBACK_LOG_COUNTS[key] = count + 1
logger.warning(message, *args)
logger.warning(
"[CP_SHARED_KV_FALLBACK][tai_fused_mla_store] reason=%s " + message,
key,
*args,
)
def _contiguous_for_tai(tensor: torch.Tensor) -> torch.Tensor: