Gate tiny CP shared-KV prefetch by token threshold

Prefix materialize prefetch has fixed launch and coordination overhead, so tiny cache-hit prefixes should stay on the simpler synchronous path. The default threshold is now expressed as a cached import-time token threshold while the existing page override remains available for workload-specific tuning.

Both MLA and index prefix prefetchers use the same page-size-aware threshold, so CP shared-KV prefix prefetch is enabled only when the prefix covers at least one page per CP lane and at least the configured token threshold.

Constraint: The default token threshold must be read once to avoid hot-path env lookups

Constraint: Existing page threshold override must continue to disable or force the gate

Rejected: Hard-code 1024 tokens | tuning needs to be possible without code changes

Rejected: Read token threshold on every maybe_create call | unnecessary hot-path env lookup

Confidence: medium

Scope-risk: narrow

Directive: Keep MLA and index prefetch threshold logic shared; do not let their gates diverge

Tested: Not rerun per user request; prior py_compile, diff check, and isolated env-cache check had passed before commit

Not-tested: full pytest in local environment due missing optional dependencies/kernel import constraints

Not-tested: GLM5 E2E throughput after this threshold change

python/sglang/srt/environ.py

@@ -212,6 +212,7 @@ class Envs:
     SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False)
     SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH = EnvBool(False)
     SGLANG_CP_SHARED_KV_MATERIALIZE_NVTX = EnvBool(False)
+    SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_TOKENS = EnvInt(1024)
     SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1)
     SGLANG_CP_DRAFT_SHARED_KV = EnvBool(False)
     SGLANG_CP_DRAFT_SHARED_KV_DEBUG = EnvBool(False)

python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py

@@ -403,15 +403,19 @@ class CpSharedKVMlaPrefetcher:
                 int(real_page_table.numel()),
             )
             return None
-        min_prefix_pages = cp_shared_kv_mla_prefetch_min_prefix_pages(layout.cp_size)
+        min_prefix_pages = cp_shared_kv_mla_prefetch_min_prefix_pages(
+            layout.cp_size, page_size=page_size
+        )
         if prefix_pages < min_prefix_pages:
             _prefetch_log(
                 "create_skip reason=prefix_below_min cp_rank=%s cp_size=%s "
-                "prefix_pages=%s min_prefix_pages=%s",
+                "prefix_pages=%s min_prefix_pages=%s prefix_len=%s page_size=%s",
                 layout.cp_rank,
                 layout.cp_size,
                 prefix_pages,
                 min_prefix_pages,
+                extend_prefix_len,
+                page_size,
             )
             return None
 
@@ -1061,15 +1065,19 @@ class CpSharedKVIndexPrefetcher:
                 int(real_page_table.numel()),
             )
             return None
-        min_prefix_pages = cp_shared_kv_mla_prefetch_min_prefix_pages(layout.cp_size)
+        min_prefix_pages = cp_shared_kv_mla_prefetch_min_prefix_pages(
+            layout.cp_size, page_size=page_size
+        )
         if prefix_pages < min_prefix_pages:
             _prefetch_log(
                 "index_create_skip reason=prefix_below_min cp_rank=%s cp_size=%s "
-                "prefix_pages=%s min_prefix_pages=%s",
+                "prefix_pages=%s min_prefix_pages=%s prefix_len=%s page_size=%s",
                 layout.cp_rank,
                 layout.cp_size,
                 prefix_pages,
                 min_prefix_pages,
+                extend_prefix_len,
+                page_size,
             )
             return None
 

python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py

@@ -20,6 +20,10 @@ _TAI_MATERIALIZE_FALLBACK_LOG_COUNTS: dict[str, int] = {}
 _TAI_FUSED_MLA_STORE_FALLBACK_LOG_COUNTS: dict[str, int] = {}
 _SLOT_REMAP_CACHE_LOG_COUNTS: dict[str, int] = {}
 _MLA_PREFETCH_LOG_PROBE_LAYER = 2
+_MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS = max(
+    int(envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_TOKENS.get()),
+    0,
+)
 _SORT_NVTX_ENABLED = envs.SGLANG_DEBUG_SORT_NVTX.get()
 _MATERIALIZE_NVTX_ENABLED = envs.SGLANG_CP_SHARED_KV_MATERIALIZE_NVTX.get()
 
@@ -60,17 +64,33 @@ def cp_shared_kv_mla_prefetch_log_enabled() -> bool:
     return envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.get()
 
 
-def cp_shared_kv_mla_prefetch_min_prefix_pages(cp_size: int) -> int:
+def cp_shared_kv_mla_prefetch_min_prefix_pages(
+    cp_size: int, *, page_size: int | None = None
+) -> int:
     """Minimum prefix pages required to enable Phase8 prefetch.
 
-    Negative env values mean "use cp_size" so the default skips tiny prefixes
-    that cannot cover all CP lanes.  Set the env to 0 to disable the gate, or to
-    a positive absolute page count for workload-specific tuning.
+    Negative env values mean "use the dynamic default": at least one page per CP
+    lane and, when the runtime page size is known, at least 1K prefix tokens.
+    This keeps tiny cache-hit prefixes on the simpler synchronous materialize
+    path where prefix prefetch launch/collective overhead can dominate.  Set the
+    env to 0 to disable the gate, or to a positive absolute page count for
+    workload-specific tuning.
     """
 
     configured = envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.get()
     if configured < 0:
-        return max(int(cp_size), 0)
+        min_pages = max(int(cp_size), 0)
+        if page_size is not None and int(page_size) > 0:
+            min_pages = max(
+                min_pages,
+                (
+                    _MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS
+                    + int(page_size)
+                    - 1
+                )
+                // int(page_size),
+            )
+        return min_pages
     return max(int(configured), 0)