diff --git a/docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md b/docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md index 7100da517..0032c0f14 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md +++ b/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： diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index ce9cdd6ef..0edd9de4b 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 3df144bcd..4fda36787 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/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, diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 9ecca8d07..b4e2ebbe7 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/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, diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index 46c2593e2..03feb4675 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/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