Reduce CP shared-KV prepare overhead without diagnostic log noise

The CP shared-KV path now has a gated tai-kernel replacement for NSA index
K/scale plus MQA range preparation, and Phase8 prefetch can skip tiny prefixes
that do not cover all CP lanes. The Phase9 plan documents the next scheduler
work for overlapping CP communication with peer-request attention windows.

Temporary diagnostic logs added while validating prefetch ownership and fused
index prepare routing were removed before committing so the runtime path does
not add log-only synchronization, log counters, or shape-reporting overhead.

Constraint: Production profiling showed small per-request CPU/GPU overhead from diagnostic logging and sync-prone debug counters.
Rejected: Keep fused-index prepare fallback/used logs behind a new env var | it leaves another runtime branch and logging surface for a path that should be benchmarked with profiler evidence instead.
Rejected: Keep owned page-count prefetch logs | they require sync-prone tensor reductions and were only useful for one-off diagnosis.
Confidence: medium
Scope-risk: moderate
Directive: Reintroduce CP shared-KV diagnostics only behind explicit debug paths, and avoid .item()/shape-heavy logging in hot prefill paths.
Tested: git diff --check for staged sglang-dev changes.
Tested: AST parse for environ.py, cp_shared_kv_prefetch.py, cp_shared_kv_runtime.py, nsa_indexer.py, and test_cp_shared_kv_runtime.py.
Not-tested: Full unit test suite.
Not-tested: Multi-node GLM5 prefill/decode/router runtime after this exact commit.

python/sglang/srt/environ.py

@@ -208,8 +208,10 @@ class Envs:
     SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False)
     SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
     SGLANG_CP_SHARED_KV_FUSED_MLA_STORE = EnvBool(False)
+    SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE = EnvBool(False)
     SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False)
     SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH = EnvBool(False)
+    SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1)
     SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)
     SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)
     SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1)

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

@@ -12,6 +12,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     cp_shared_kv_debug_enabled,
     cp_shared_kv_mla_prefetch_enabled,
     cp_shared_kv_mla_prefetch_log,
+    cp_shared_kv_mla_prefetch_min_prefix_pages,
     cp_shared_kv_mla_prefetch_should_log_layer,
     filter_locs_mappable_to_physical_pool,
     filter_pages_mappable_to_physical_pool,
@@ -181,6 +182,9 @@ class CpSharedKVMlaPrefetcher:
                 int(real_page_table.numel()),
             )
             return None
+        min_prefix_pages = cp_shared_kv_mla_prefetch_min_prefix_pages(layout.cp_size)
+        if prefix_pages < min_prefix_pages:
+            return None
 
         cp_group = get_attention_cp_group()
         if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1:
@@ -595,6 +599,9 @@ class CpSharedKVIndexPrefetcher:
                 int(real_page_table.numel()),
             )
             return None
+        min_prefix_pages = cp_shared_kv_mla_prefetch_min_prefix_pages(layout.cp_size)
+        if prefix_pages < min_prefix_pages:
+            return None
 
         cp_group = get_attention_cp_group()
         if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1:

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

@@ -41,6 +41,10 @@ def cp_shared_kv_tai_fused_mla_store_enabled() -> bool:
     return envs.SGLANG_CP_SHARED_KV_FUSED_MLA_STORE.get()
 
 
+def cp_shared_kv_tai_index_mqa_prepare_enabled() -> bool:
+    return envs.SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE.get()
+
+
 def cp_shared_kv_mla_prefetch_enabled() -> bool:
     return envs.SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH.get()
 
@@ -49,6 +53,20 @@ 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:
+    """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.
+    """
+
+    configured = envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.get()
+    if configured < 0:
+        return max(int(cp_size), 0)
+    return max(int(configured), 0)
+
+
 def cp_shared_kv_mla_prefetch_log(message: str, *args) -> None:
     if cp_shared_kv_mla_prefetch_log_enabled():
         logger.info("[CP_SHARED_KV_MLA_PREFETCH] " + message, *args)
@@ -219,6 +237,16 @@ def _load_tai_fused_mla_store_kernel():
         return None
 
 
+@lru_cache(maxsize=1)
+def _load_tai_index_mqa_prepare_kernel():
+    try:
+        from tai_kernel.nsa_prefill import prepare_cp_mqa_kv_and_range
+
+        return prepare_cp_mqa_kv_and_range
+    except Exception:
+        return None
+
+
 def _log_tai_materialize_fallback(
     key: str,
     message: str,
@@ -249,6 +277,51 @@ def _contiguous_for_tai(tensor: torch.Tensor) -> torch.Tensor:
     return tensor if tensor.is_contiguous() else tensor.contiguous()
 
 
+def try_tai_prepare_cp_mqa_index(
+    *,
+    index_buffer: torch.Tensor,
+    page_indices: torch.Tensor,
+    kv_len: int,
+    valid_q_count: int,
+    ke_start: int,
+    page_size: int,
+    index_head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None:
+    """Try the TAI fused GetK/GetS + MQA range prepare path.
+
+    The fallback path in SGLang launches separate GetK/GetS kernels and then
+    creates `ks`/`ke` tensors with torch elementwise ops.  The TAI kernel fuses
+    those preparation steps into one launch for the single-sequence CP pair path.
+    """
+
+    if not cp_shared_kv_tai_index_mqa_prepare_enabled():
+        return None
+
+    kernel = _load_tai_index_mqa_prepare_kernel()
+    if kernel is None:
+        return None
+
+    if index_buffer.dtype != torch.uint8:
+        return None
+    if not index_buffer.is_contiguous():
+        return None
+    if index_head_dim != 128 or page_size != 64:
+        return None
+
+    try:
+        return kernel(
+            index_buffer,
+            _contiguous_for_tai(page_indices),
+            kv_len=int(kv_len),
+            valid_q_count=int(valid_q_count),
+            ke_start=int(ke_start),
+            page_size=int(page_size),
+            index_head_dim=int(index_head_dim),
+        )
+    except Exception:
+        return None
+
+
 def try_tai_fused_mla_store(
     *,
     token_to_kv_pool,

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

@@ -24,6 +24,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     materialize_shared_paged_buffer,
     tensor_debug_checksum,
     tensor_debug_summary,
+    try_tai_prepare_cp_mqa_index,
 )
 from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor
 from sglang.srt.layers.layernorm import LayerNorm
@@ -1097,39 +1098,59 @@ class Indexer(MultiPlatformOp):
                 )
 
             ke_start = cp_kv_end - actual_seq_q + 1
-            ke_offset = torch.arange(
-                ke_start,
-                ke_start + valid_q_count,
-                dtype=torch.int32,
-                device=q_fp8.device,
-            )
 
             q_fp8 = q_fp8[:valid_q_count]
             weights = weights[:valid_q_count]
             kv_len = min(cp_kv_end, logical_kv_limit)
             if current_index_kv is None:
                 assert index_buffer is not None
-                k_fp8 = index_buf_accessor.GetK.execute(
-                    forward_batch.token_to_kv_pool,
-                    index_buffer,
-                    seq_len=kv_len,
-                    page_indices=block_tables[0],
-                )
-                k_scale = index_buf_accessor.GetS.execute(
-                    forward_batch.token_to_kv_pool,
-                    index_buffer,
-                    seq_len=kv_len,
+                tai_prepared = try_tai_prepare_cp_mqa_index(
+                    index_buffer=index_buffer,
                     page_indices=block_tables[0],
+                    kv_len=kv_len,
+                    valid_q_count=valid_q_count,
+                    ke_start=ke_start,
+                    page_size=page_size,
+                    index_head_dim=forward_batch.token_to_kv_pool.index_head_dim,
                 )
+                if tai_prepared is not None:
+                    k_fp8_u8, k_scale, ks, ke_offset = tai_prepared
+                    k_fp8 = k_fp8_u8.view(torch.float8_e4m3fn)
+                else:
+                    ke_offset = torch.arange(
+                        ke_start,
+                        ke_start + valid_q_count,
+                        dtype=torch.int32,
+                        device=q_fp8.device,
+                    )
+                    k_fp8 = index_buf_accessor.GetK.execute(
+                        forward_batch.token_to_kv_pool,
+                        index_buffer,
+                        seq_len=kv_len,
+                        page_indices=block_tables[0],
+                    )
+                    k_scale = index_buf_accessor.GetS.execute(
+                        forward_batch.token_to_kv_pool,
+                        index_buffer,
+                        seq_len=kv_len,
+                        page_indices=block_tables[0],
+                    )
 
-                k_fp8 = k_fp8.view(torch.float8_e4m3fn)
-                k_scale = k_scale.view(torch.float32).squeeze(-1)
+                    k_fp8 = k_fp8.view(torch.float8_e4m3fn)
+                    k_scale = k_scale.view(torch.float32).squeeze(-1)
+                    ks = torch.zeros_like(ke_offset)
             else:
+                ke_offset = torch.arange(
+                    ke_start,
+                    ke_start + valid_q_count,
+                    dtype=torch.int32,
+                    device=q_fp8.device,
+                )
                 k_fp8, k_scale = current_index_kv
                 k_fp8 = k_fp8[:kv_len].contiguous()
                 k_scale = k_scale[:kv_len].view(torch.float32).squeeze(-1).contiguous()
+                ks = torch.zeros_like(ke_offset)
             kv_fp8 = (k_fp8, k_scale)
-            ks = torch.zeros_like(ke_offset)
             ke = ke_offset
 
             with self._with_real_sm_count():