Route CP shared MLA store through TAI fused kernels without runtime spam

The shared-KV prefill path now optionally calls tai_kernel.nsa_prefill.fused_store_mla_kv before falling back to logical_locs_to_physical plus set_mla_kv_buffer. The fast path supports packed FP8 and BF16/FP16 direct KV buffers, while debug mode and kernel failures still preserve the existing fallback behavior. Success logging was removed after path verification because per-layer/per-rank logs are too noisy in normal server runs.

Constraint: Runtime must remain safe when tai-kernel is absent or debug checks are enabled
Rejected: Keep success logs permanently | floods prefill logs once every rank/layer starts using the fast path
Confidence: high
Scope-risk: moderate
Directive: Keep fallback warnings; do not re-add per-layer success logs outside explicit debug instrumentation
Tested: g0034 container python -m py_compile python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
Tested: g0034 container PYTHONPATH=python pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -q (40 passed)
Not-tested: Full multi-node PD server throughput after log removal

python/sglang/srt/environ.py

@@ -207,6 +207,7 @@ class Envs:
     SGLANG_DEBUG_MOE_SORT_NVTX = EnvBool(False)
     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_ENABLE_MLA_PREFETCH = EnvBool(False)
     SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH = EnvBool(False)
     SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)

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

@@ -14,6 +14,7 @@ logger = logging.getLogger(__name__)
 
 _DEBUG_LOG_COUNTS: dict[str, int] = {}
 _TAI_MATERIALIZE_FALLBACK_LOG_COUNTS: dict[str, int] = {}
+_TAI_FUSED_MLA_STORE_FALLBACK_LOG_COUNTS: dict[str, int] = {}
 _MLA_PREFETCH_LOG_PROBE_LAYER = 2
 _SORT_NVTX_ENABLED = envs.SGLANG_DEBUG_SORT_NVTX.get()
 
@@ -34,6 +35,10 @@ def cp_shared_kv_tai_materialize_enabled() -> bool:
     return envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.get()
 
 
+def cp_shared_kv_tai_fused_mla_store_enabled() -> bool:
+    return envs.SGLANG_CP_SHARED_KV_FUSED_MLA_STORE.get()
+
+
 def cp_shared_kv_mla_prefetch_enabled() -> bool:
     return envs.SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH.get()
 
@@ -85,6 +90,37 @@ def _tai_materialize_runtime_enabled() -> bool:
     return cp_shared_kv_tai_materialize_enabled() and not cp_shared_kv_debug_enabled()
 
 
+@lru_cache(maxsize=1)
+def _load_tai_fused_mla_store_kernel():
+    try:
+        from tai_kernel.nsa_prefill import fused_store_mla_kv
+
+        return fused_store_mla_kv
+    except ImportError:
+        try:
+            from tai_kernel.nsa_prefill import fused_quant_store_mla_kv
+
+            return fused_quant_store_mla_kv
+        except Exception as exc:
+            _log_tai_fused_mla_store_fallback(
+                "import_failed",
+                "CP shared KV tai fused MLA store kernel is unavailable; "
+                "falling back to torch MLA store. error=%s",
+                exc,
+                limit=1,
+            )
+            return None
+    except Exception as exc:
+        _log_tai_fused_mla_store_fallback(
+            "import_failed",
+            "CP shared KV tai fused MLA store kernel is unavailable; "
+            "falling back to torch MLA store. error=%s",
+            exc,
+            limit=1,
+        )
+        return None
+
+
 def _log_tai_materialize_fallback(
     key: str,
     message: str,
@@ -98,10 +134,91 @@ def _log_tai_materialize_fallback(
     logger.warning(message, *args)
 
 
+def _log_tai_fused_mla_store_fallback(
+    key: str,
+    message: str,
+    *args,
+    limit: int = 8,
+) -> None:
+    count = _TAI_FUSED_MLA_STORE_FALLBACK_LOG_COUNTS.get(key, 0)
+    if count >= limit:
+        return
+    _TAI_FUSED_MLA_STORE_FALLBACK_LOG_COUNTS[key] = count + 1
+    logger.warning(message, *args)
+
+
 def _contiguous_for_tai(tensor: torch.Tensor) -> torch.Tensor:
     return tensor if tensor.is_contiguous() else tensor.contiguous()
 
 
+def try_tai_fused_mla_store(
+    *,
+    token_to_kv_pool,
+    layer,
+    layout: CpSharedKVLayout,
+    logical_locs: torch.Tensor,
+    k_nope: torch.Tensor,
+    k_rope: torch.Tensor,
+) -> bool:
+    """Try the TAI fused MLA persistent-store path.
+
+    The fallback path first remaps shared-KV logical token ids to this rank's
+    compact physical rows, then runs quantize/store kernels or a torch direct
+    store.  The TAI kernel accepts the shared logical locs directly and writes
+    either the packed FP8 MLA KV bytes or the bf16/fp16 direct MLA rows into the
+    raw backing buffer, avoiding the torch remap and separate store launches.
+    """
+
+    if not cp_shared_kv_tai_fused_mla_store_enabled():
+        return False
+    if cp_shared_kv_debug_enabled():
+        _log_tai_fused_mla_store_fallback(
+            "debug_enabled",
+            "CP shared KV tai fused MLA store fallback (debug_enabled): "
+            "SGLANG_DEBUG_CP_SHARED_KV is enabled. layer=%s cp_rank=%s tokens=%s",
+            getattr(layer, "layer_id", None),
+            layout.cp_rank,
+            logical_locs.numel(),
+            limit=8,
+        )
+        return False
+
+    kernel = _load_tai_fused_mla_store_kernel()
+    if kernel is None:
+        return False
+
+    kv_buffer = getattr(token_to_kv_pool, "kv_buffer", None)
+    if kv_buffer is None:
+        _log_tai_fused_mla_store_fallback(
+            "missing_raw_buffer",
+            "CP shared KV tai fused MLA store requires token pool kv_buffer; "
+            "falling back to torch MLA store. pool=%s",
+            type(token_to_kv_pool).__name__,
+        )
+        return False
+
+    try:
+        layer_idx = int(layer.layer_id) - int(token_to_kv_pool.start_layer)
+        raw_layer_buffer = kv_buffer[layer_idx]
+        kernel(
+            k_nope,
+            k_rope,
+            raw_layer_buffer,
+            _contiguous_for_tai(logical_locs),
+            page_size=layout.page_size,
+            cp_size=layout.cp_size,
+        )
+        return True
+    except Exception as exc:
+        _log_tai_fused_mla_store_fallback(
+            "kernel_failed",
+            "CP shared KV tai fused MLA store failed; falling back to torch "
+            "MLA store. error=%s",
+            exc,
+        )
+        return False
+
+
 def _try_tai_materialize_shared_pages(
     page_buffer: torch.Tensor,
     logical_pages: torch.Tensor,

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -11,6 +11,9 @@ from sglang.srt.layers.attention.nsa.utils import (
     log_cp_shared_kv_direct_write_fallback,
     nsa_use_prefill_cp,
 )
+from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
+    try_tai_fused_mla_store,
+)
 from sglang.srt.layers.communicator import get_attn_tp_context
 from sglang.srt.layers.quantization.fp8_kernel import (
     fp8_dtype,
@@ -555,10 +558,19 @@ class DeepseekMLAForwardMixin:
             return True
 
         assert forward_batch.cp_shared_kv_layout is not None
+        layout = forward_batch.cp_shared_kv_layout
+        if try_tai_fused_mla_store(
+            token_to_kv_pool=forward_batch.token_to_kv_pool,
+            layer=self.attn_mqa,
+            layout=layout,
+            logical_locs=local_out_cache_loc,
+            k_nope=k_nope,
+            k_rope=k_pe,
+        ):
+            return True
+
         physical_out_cache_loc = (
-            forward_batch.cp_shared_kv_layout.logical_locs_to_physical(
-                local_out_cache_loc
-            ).contiguous()
+            layout.logical_locs_to_physical(local_out_cache_loc).contiguous()
         )
         forward_batch.token_to_kv_pool.set_mla_kv_buffer(
             self.attn_mqa,