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
This commit is contained in:
laoyao0822
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user