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

docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md

@@ -474,6 +474,50 @@ SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 fast path
 - fallback reason 可解释；
 - profiler 中 `materialize.token.local_copy` / `materialize.paged.local_copy` 下降。
 
+### P0 fused MLA persistent store 接入
+
+在 materialize 之前，shared KV direct-write 的 MLA persistent KV 写入还有一段
+独立开销：
+
+```text
+logical_locs_to_physical(...)
+quantize_k_cache_separate(...)
+set_mla_kv_buffer_triton(...)
+```
+
+对应的 tai-kernel fast path 已通过 SGLang 实验开关接入：
+
+```text
+SGLANG_CP_SHARED_KV_FUSED_MLA_STORE=0/1
+```
+
+默认关闭。开启后只影响 prefill CP shared-KV 的 MLA local persistent write：
+
+```text
+forward_mla.py
+  _maybe_write_cp_shared_local_mla_kv(...)
+    -> try_tai_fused_mla_store(...)
+       -> tai_kernel.nsa_prefill.fused_store_mla_kv(
+              k_nope, k_rope, raw_kv_buffer, logical_locs,
+              page_size, cp_size)
+    -> fallback logical_locs_to_physical + token_to_kv_pool.set_mla_kv_buffer
+```
+
+接入约束：
+
+1. `fused_store_mla_kv` 根据 raw KV buffer dtype 分发：
+   - `uint8 [capacity, 1, 656]`：packed NSA FP8 路径，
+     fused logical->physical + FP8 quant + scale/rope store；
+   - `bf16/fp16 [capacity, 1, 576]`：non-FP8 direct-store 路径，
+     fused logical->physical + MLA KV direct write。
+2. tai-kernel import 失败、shape/dtype 不支持、kernel runtime 失败时自动 fallback。
+3. `SGLANG_DEBUG_CP_SHARED_KV=1` 时强制保留原 PyTorch path，便于继续暴露
+   shared KV correctness 问题。
+4. fused kernel 消费 shared logical token locs；调用前仍要求 direct-write gate
+   已确认 `local_out_cache_loc` 属于当前 CP owner。
+5. 该路径不替代后续 attention 前的 shared KV materialize，只减少 persistent
+   write 阶段的 remap/quant/store kernel 碎片。
+
 ## Benchmark plan
 
 新增 benchmark 只测本地 materialize，不包含 CP all-reduce：

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(
+            layout.logical_locs_to_physical(local_out_cache_loc).contiguous()
-                local_out_cache_loc
-            ).contiguous()
         )
         forward_batch.token_to_kv_pool.set_mla_kv_buffer(
             self.attn_mqa,

test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py

@@ -268,6 +268,193 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         with envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.override(True):
             self.assertTrue(cp_shared_kv_mla_prefetch_log_enabled())
 
+    def test_fused_mla_store_uses_tai_kernel_when_enabled(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+        from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+
+        class FakeTaiFusedStore:
+            def __init__(self):
+                self.calls = []
+
+            def __call__(
+                self,
+                k_nope,
+                k_rope,
+                kv_buffer,
+                logical_locs,
+                *,
+                page_size,
+                cp_size,
+            ):
+                self.calls.append(
+                    (k_nope, k_rope, kv_buffer, logical_locs, page_size, cp_size)
+                )
+
+        class FakePool:
+            nsa_kv_cache_store_fp8 = True
+            page_size = 64
+            start_layer = 0
+
+            def __init__(self):
+                self.kv_buffer = [torch.zeros((128, 1, 656), dtype=torch.uint8)]
+
+        fake_kernel = FakeTaiFusedStore()
+        pool = FakePool()
+        layer = SimpleNamespace(layer_id=0)
+        layout = CpSharedKVLayout(page_size=64, cp_size=8, cp_rank=3)
+        logical_locs = torch.tensor([192, 193], dtype=torch.int64)
+        k_nope = torch.zeros((2, 1, 512), dtype=torch.bfloat16)
+        k_rope = torch.zeros((2, 1, 64), dtype=torch.bfloat16)
+
+        with patch.object(
+            runtime, "cp_shared_kv_tai_fused_mla_store_enabled", return_value=True
+        ), patch.object(
+            runtime, "cp_shared_kv_debug_enabled", return_value=False
+        ), patch.object(
+            runtime, "_load_tai_fused_mla_store_kernel", return_value=fake_kernel
+        ):
+            used = runtime.try_tai_fused_mla_store(
+                token_to_kv_pool=pool,
+                layer=layer,
+                layout=layout,
+                logical_locs=logical_locs,
+                k_nope=k_nope,
+                k_rope=k_rope,
+            )
+
+        self.assertTrue(used)
+        self.assertEqual(len(fake_kernel.calls), 1)
+        call = fake_kernel.calls[0]
+        self.assertIs(call[0], k_nope)
+        self.assertIs(call[1], k_rope)
+        self.assertIs(call[2], pool.kv_buffer[0])
+        self.assertIs(call[3], logical_locs)
+        self.assertEqual(call[4], 64)
+        self.assertEqual(call[5], 8)
+
+    def test_fused_mla_store_uses_tai_kernel_for_non_fp8_pool_when_enabled(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+        from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+
+        class FakeTaiFusedStore:
+            def __init__(self):
+                self.calls = []
+
+            def __call__(
+                self,
+                k_nope,
+                k_rope,
+                kv_buffer,
+                logical_locs,
+                *,
+                page_size,
+                cp_size,
+            ):
+                self.calls.append(
+                    (k_nope, k_rope, kv_buffer, logical_locs, page_size, cp_size)
+                )
+
+        class FakePool:
+            nsa_kv_cache_store_fp8 = False
+            page_size = 64
+            start_layer = 0
+
+            def __init__(self):
+                self.kv_buffer = [torch.zeros((128, 1, 576), dtype=torch.bfloat16)]
+
+        fake_kernel = FakeTaiFusedStore()
+        pool = FakePool()
+        layer = SimpleNamespace(layer_id=0)
+        layout = CpSharedKVLayout(page_size=64, cp_size=8, cp_rank=3)
+        logical_locs = torch.tensor([192, 193], dtype=torch.int64)
+        k_nope = torch.zeros((2, 1, 512), dtype=torch.bfloat16)
+        k_rope = torch.zeros((2, 1, 64), dtype=torch.bfloat16)
+
+        with patch.object(
+            runtime, "cp_shared_kv_tai_fused_mla_store_enabled", return_value=True
+        ), patch.object(
+            runtime, "cp_shared_kv_debug_enabled", return_value=False
+        ), patch.object(
+            runtime, "_load_tai_fused_mla_store_kernel", return_value=fake_kernel
+        ):
+            used = runtime.try_tai_fused_mla_store(
+                token_to_kv_pool=pool,
+                layer=layer,
+                layout=layout,
+                logical_locs=logical_locs,
+                k_nope=k_nope,
+                k_rope=k_rope,
+            )
+
+        self.assertTrue(used)
+        self.assertEqual(len(fake_kernel.calls), 1)
+        call = fake_kernel.calls[0]
+        self.assertIs(call[0], k_nope)
+        self.assertIs(call[1], k_rope)
+        self.assertIs(call[2], pool.kv_buffer[0])
+        self.assertIs(call[3], logical_locs)
+        self.assertEqual(call[4], 64)
+        self.assertEqual(call[5], 8)
+
+    def test_fused_mla_store_stays_off_by_default(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+        from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+
+        class FakePool:
+            nsa_kv_cache_store_fp8 = True
+            page_size = 64
+            start_layer = 0
+            kv_buffer = [torch.zeros((128, 1, 656), dtype=torch.uint8)]
+
+        with patch.object(
+            runtime, "cp_shared_kv_tai_fused_mla_store_enabled", return_value=False
+        ), patch.object(
+            runtime, "_load_tai_fused_mla_store_kernel"
+        ) as load_kernel:
+            used = runtime.try_tai_fused_mla_store(
+                token_to_kv_pool=FakePool(),
+                layer=SimpleNamespace(layer_id=0),
+                layout=CpSharedKVLayout(page_size=64, cp_size=8, cp_rank=0),
+                logical_locs=torch.tensor([64], dtype=torch.int64),
+                k_nope=torch.zeros((1, 1, 512), dtype=torch.bfloat16),
+                k_rope=torch.zeros((1, 1, 64), dtype=torch.bfloat16),
+            )
+
+        self.assertFalse(used)
+        load_kernel.assert_not_called()
+
+    def test_fused_mla_store_logs_debug_fallback_when_env_enabled(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+        from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+
+        class FakePool:
+            nsa_kv_cache_store_fp8 = True
+            page_size = 64
+            start_layer = 0
+            kv_buffer = [torch.zeros((128, 1, 656), dtype=torch.uint8)]
+
+        with patch.object(
+            runtime, "cp_shared_kv_tai_fused_mla_store_enabled", return_value=True
+        ), patch.object(
+            runtime, "cp_shared_kv_debug_enabled", return_value=True
+        ), patch.object(
+            runtime, "_load_tai_fused_mla_store_kernel"
+        ) as load_kernel, self.assertLogs(
+            runtime.logger.name, level="WARNING"
+        ) as logs:
+            used = runtime.try_tai_fused_mla_store(
+                token_to_kv_pool=FakePool(),
+                layer=SimpleNamespace(layer_id=0),
+                layout=CpSharedKVLayout(page_size=64, cp_size=8, cp_rank=0),
+                logical_locs=torch.tensor([64], dtype=torch.int64),
+                k_nope=torch.zeros((1, 1, 512), dtype=torch.bfloat16),
+                k_rope=torch.zeros((1, 1, 64), dtype=torch.bfloat16),
+            )
+
+        self.assertFalse(used)
+        load_kernel.assert_not_called()
+        self.assertTrue(any("debug_enabled" in message for message in logs.output))
+
     def test_token_range_materialize_uses_tai_kernel_when_enabled(self):
         from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
         from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout