Preserve FP8 CP shared-KV page contracts
NSA FP8 CP shared-KV reuse must operate on packed page-slot rows, not bf16 compact rows. The change keeps current-only and partial-current reuse inside the page-aligned materialization contract, fails fast for non-page-aligned CP split inputs, and prevents FP8 FlashMLA-KV prefill from reaching incompatible in-seq CP metadata. Constraint: NSA FP8 persistent MLA KV rows are packed 656-byte records and CP shared KV cache management is page-granular.\nConstraint: FlashMLA-KV prefill metadata is not CP-local after NSA in-seq splitting.\nRejected: Silently splice bf16 current rows into FP8 materialized cache | corrupts the packed cache layout.\nRejected: Keep FP8 FlashMLA-KV auto prefill under NSA CP | reaches num_splits shape errors after q-row splitting.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not re-enable FP8 FlashMLA-KV prefill for NSA in-seq CP until metadata is rebuilt after CP splitting or made CP-local.\nTested: Local git diff --check and py_compile for touched SGLang files.\nTested: Remote g0034 related unit sweep recorded in docs: test_nsa_cp_utils.py, test_cp_shared_kv_layout.py, test_cp_shared_kv_runtime.py, test_cp_hicache_metadata.py passed.\nNot-tested: Full FP8 ETE startup and performance run after this commit.
This commit is contained in:
laoyao0822
@@ -113,6 +113,7 @@ from sglang.srt.mem_cache.hiradix_cache import (
|
||||
PreparedCpHiCacheBackup,
|
||||
_compute_shared_hicache_token_capacities,
|
||||
)
|
||||
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
|
||||
from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode, _key_match_paged
|
||||
from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
@@ -191,6 +192,45 @@ class TestCpHiCacheImports(CustomTestCase):
|
||||
|
||||
|
||||
class TestHiRadixCacheCPDraftHostPool(CustomTestCase):
|
||||
def test_fp8_nsa_hicache_size_estimate_uses_packed_row_width(self):
|
||||
pool = object.__new__(NSATokenToKVPool)
|
||||
pool.store_dtype = torch.uint8
|
||||
pool.kv_cache_dim = 656
|
||||
pool.layer_num = 78
|
||||
pool.index_head_dim = 128
|
||||
pool.quant_block_size = 128
|
||||
|
||||
size_per_token = hiradix_cache._estimate_hicache_size_per_token(pool)
|
||||
|
||||
# FP8 NSA stores packed MLA rows as 656 uint8 bytes/token/layer plus
|
||||
# the paged index buffer: 128 uint8 K bytes + 4 uint8 scale bytes.
|
||||
self.assertEqual(size_per_token, (656 + 128 + 4) * 78)
|
||||
|
||||
def test_fp8_shared_budget_matches_target_and_one_layer_draft_capacity(self):
|
||||
bytes_per_layer = 656 + 128 + 4
|
||||
target_size_per_token = bytes_per_layer * 78
|
||||
draft_size_per_token = bytes_per_layer
|
||||
|
||||
target_tokens, draft_tokens = _compute_shared_hicache_token_capacities(
|
||||
total_host_bytes=int(150 * 1e9),
|
||||
target_size_per_token=target_size_per_token,
|
||||
draft_size_per_token=draft_size_per_token,
|
||||
page_size=64,
|
||||
)
|
||||
|
||||
self.assertGreaterEqual(draft_tokens, target_tokens)
|
||||
self.assertLessEqual(
|
||||
target_tokens * target_size_per_token
|
||||
+ draft_tokens * draft_size_per_token,
|
||||
int(150 * 1e9),
|
||||
)
|
||||
# The draft pool is much smaller in bytes because GLM-5 EAGLE draft has
|
||||
# one executable layer, but it still has at least target token capacity.
|
||||
self.assertLess(
|
||||
draft_tokens * draft_size_per_token,
|
||||
target_tokens * target_size_per_token // 32,
|
||||
)
|
||||
|
||||
def test_shared_budget_keeps_draft_at_least_target_capacity(self):
|
||||
target_tokens, draft_tokens = _compute_shared_hicache_token_capacities(
|
||||
total_host_bytes=1000,
|
||||
|
||||
@@ -672,10 +672,214 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
body_source = "".join(source[body_start:body_end].split())
|
||||
|
||||
self.assertIn("valid_current_rows=int(current_kv_rows_for_reuse)", body_source)
|
||||
self.assertIn("k[:valid_current_rows]", body_source)
|
||||
self.assertIn("k_rope[:valid_current_rows]", body_source)
|
||||
self.assertIn("pack_current_mla_kv_for_reuse", body_source)
|
||||
self.assertIn("forward_batch.out_cache_loc[:valid_current_rows]", body_source)
|
||||
|
||||
def test_flashmla_kv_current_only_reuse_keeps_page_slot_layout(self):
|
||||
from pathlib import Path
|
||||
|
||||
source = (
|
||||
Path(__file__).resolve().parents[4]
|
||||
/ "python/sglang/srt/layers/attention/nsa_backend.py"
|
||||
).read_text()
|
||||
branch_start = source.index(" if is_current_only_extend_batch")
|
||||
branch_end = source.index(" else:", branch_start)
|
||||
branch_source = source[branch_start:branch_end]
|
||||
|
||||
self.assertIn(
|
||||
"materialize_prefix_and_reuse_current_kv_page_slots", branch_source
|
||||
)
|
||||
self.assertIn("prefix_pages=0", branch_source)
|
||||
self.assertNotIn("kv_cache = current_kv_cache", branch_source)
|
||||
|
||||
def test_nsa_backend_topk_transform_uses_effective_forward_impl(self):
|
||||
from sglang.srt.layers.attention.nsa_backend import (
|
||||
NativeSparseAttnBackend,
|
||||
TopkTransformMethod,
|
||||
)
|
||||
|
||||
class Mode:
|
||||
def __init__(self, *, target=False, draft=False, decode=False):
|
||||
self.target = target
|
||||
self.draft = draft
|
||||
self.decode = decode
|
||||
|
||||
def is_decode_or_idle(self):
|
||||
return self.decode
|
||||
|
||||
def is_target_verify(self):
|
||||
return self.target
|
||||
|
||||
def is_draft_extend(self, include_v2=False):
|
||||
return self.draft
|
||||
|
||||
backend = NativeSparseAttnBackend.__new__(NativeSparseAttnBackend)
|
||||
backend.nsa_kv_cache_store_fp8 = True
|
||||
backend.nsa_prefill_impl = "flashmla_sparse"
|
||||
backend.nsa_decode_impl = "flashmla_kv"
|
||||
|
||||
# Normal extend follows the prefill implementation and therefore needs
|
||||
# ragged top-k for FP8 sparse prefill.
|
||||
self.assertEqual(
|
||||
backend.get_topk_transform_method(
|
||||
SimpleNamespace(forward_mode=Mode())
|
||||
),
|
||||
TopkTransformMethod.RAGGED,
|
||||
)
|
||||
# Target verify / draft / decode follow the decode implementation. When
|
||||
# decode is flashmla_kv, those paths need paged top-k even if prefill is
|
||||
# currently flashmla_sparse.
|
||||
for mode in (
|
||||
Mode(target=True),
|
||||
Mode(draft=True),
|
||||
Mode(decode=True),
|
||||
):
|
||||
self.assertEqual(
|
||||
backend.get_topk_transform_method(
|
||||
SimpleNamespace(forward_mode=mode)
|
||||
),
|
||||
TopkTransformMethod.PAGED,
|
||||
)
|
||||
|
||||
def test_flashmla_metadata_creation_uses_effective_forward_impl(self):
|
||||
from pathlib import Path
|
||||
|
||||
source = (
|
||||
Path(__file__).resolve().parents[4]
|
||||
/ "python/sglang/srt/layers/attention/nsa_backend.py"
|
||||
).read_text()
|
||||
metadata_start = source.index(" flashmla_metadata=(")
|
||||
metadata_end = source.index(" paged_mqa_schedule_metadata=", metadata_start)
|
||||
metadata_source = source[metadata_start:metadata_end]
|
||||
|
||||
self.assertIn(
|
||||
"_effective_nsa_impl_for_forward_batch(forward_batch)",
|
||||
metadata_source,
|
||||
)
|
||||
self.assertNotIn(
|
||||
'if self.nsa_decode_impl == "flashmla_kv"',
|
||||
metadata_source,
|
||||
)
|
||||
|
||||
def test_fp8_auto_prefill_cp_uses_sparse_not_flashmla_kv(self):
|
||||
from sglang.srt.layers.attention.nsa_backend import NativeSparseAttnBackend
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardMode
|
||||
|
||||
def make_backend():
|
||||
backend = NativeSparseAttnBackend.__new__(NativeSparseAttnBackend)
|
||||
backend.nsa_kv_cache_store_fp8 = True
|
||||
backend.enable_auto_select_prefill_impl = True
|
||||
backend.nsa_prefill_impl = "flashmla_auto"
|
||||
return backend
|
||||
|
||||
forward_batch = SimpleNamespace(
|
||||
forward_mode=ForwardMode.EXTEND,
|
||||
seq_lens_cpu=torch.tensor([40392], dtype=torch.int32),
|
||||
seq_lens_sum=40392,
|
||||
extend_num_tokens=40392,
|
||||
token_to_kv_pool=SimpleNamespace(dtype=torch.float8_e4m3fn),
|
||||
hisparse_coordinator=None,
|
||||
get_max_chunk_capacity=lambda: 65536,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"sglang.srt.utils.get_device_sm", return_value=90
|
||||
), patch(
|
||||
"sglang.srt.utils.is_blackwell", return_value=False
|
||||
), patch(
|
||||
"sglang.srt.layers.attention.nsa_backend.is_nsa_enable_prefill_cp",
|
||||
return_value=True,
|
||||
):
|
||||
backend = make_backend()
|
||||
backend.set_nsa_prefill_impl(forward_batch)
|
||||
self.assertEqual(backend.nsa_prefill_impl, "flashmla_sparse")
|
||||
|
||||
with patch(
|
||||
"sglang.srt.utils.get_device_sm", return_value=90
|
||||
), patch(
|
||||
"sglang.srt.utils.is_blackwell", return_value=False
|
||||
), patch(
|
||||
"sglang.srt.layers.attention.nsa_backend.is_nsa_enable_prefill_cp",
|
||||
return_value=False,
|
||||
):
|
||||
backend = make_backend()
|
||||
backend.set_nsa_prefill_impl(forward_batch)
|
||||
self.assertEqual(backend.nsa_prefill_impl, "flashmla_kv")
|
||||
|
||||
def test_explicit_fp8_flashmla_kv_cp_prefill_fails_before_kernel(self):
|
||||
from pathlib import Path
|
||||
|
||||
source = (
|
||||
Path(__file__).resolve().parents[4]
|
||||
/ "python/sglang/srt/layers/attention/nsa_backend.py"
|
||||
).read_text()
|
||||
branch_start = source.index(' elif nsa_impl == "flashmla_kv":')
|
||||
branch_end = source.index(' elif nsa_impl == "fa3":', branch_start)
|
||||
branch_source = source[branch_start:branch_end]
|
||||
|
||||
self.assertIn("nsa_use_prefill_cp(forward_batch)", branch_source)
|
||||
self.assertIn("NSA_FP8_FLASHMLA_KV_CP_PREFILL_UNSUPPORTED", branch_source)
|
||||
self.assertLess(
|
||||
branch_source.index("NSA_FP8_FLASHMLA_KV_CP_PREFILL_UNSUPPORTED"),
|
||||
branch_source.index("self._forward_flashmla_kv("),
|
||||
)
|
||||
|
||||
def test_pack_current_mla_kv_for_reuse_uses_tai_pack_for_fp8_cache(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
|
||||
class FakeTaiPack:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def __call__(self, k_nope, k_rope):
|
||||
self.calls.append((k_nope, k_rope))
|
||||
return torch.ones((k_nope.shape[0], 1, 656), dtype=torch.uint8)
|
||||
|
||||
fake_kernel = FakeTaiPack()
|
||||
k_nope = torch.zeros((3, 1, 512), dtype=torch.bfloat16)
|
||||
k_rope = torch.zeros((3, 1, 64), dtype=torch.bfloat16)
|
||||
fp8_kv_cache = torch.empty((128, 1, 656), dtype=torch.float8_e4m3fn)
|
||||
|
||||
with patch.object(
|
||||
runtime, "_load_tai_pack_quant_mla_kv_kernel", return_value=fake_kernel
|
||||
):
|
||||
packed = runtime.pack_current_mla_kv_for_reuse(
|
||||
k_nope,
|
||||
k_rope,
|
||||
kv_cache=fp8_kv_cache,
|
||||
)
|
||||
|
||||
self.assertEqual(len(fake_kernel.calls), 1)
|
||||
self.assertIs(fake_kernel.calls[0][0], k_nope)
|
||||
self.assertIs(fake_kernel.calls[0][1], k_rope)
|
||||
self.assertEqual(tuple(packed.shape), (3, 1, 656))
|
||||
self.assertEqual(packed.dtype, torch.float8_e4m3fn)
|
||||
self.assertTrue(
|
||||
torch.equal(
|
||||
packed.view(torch.uint8),
|
||||
torch.ones_like(packed.view(torch.uint8)),
|
||||
)
|
||||
)
|
||||
|
||||
def test_pack_current_mla_kv_for_reuse_fails_fast_when_fp8_pack_kernel_missing(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
|
||||
k_nope = torch.zeros((1, 1, 512), dtype=torch.bfloat16)
|
||||
k_rope = torch.zeros((1, 1, 64), dtype=torch.bfloat16)
|
||||
fp8_kv_cache = torch.empty((128, 1, 656), dtype=torch.float8_e4m3fn)
|
||||
|
||||
with patch.object(
|
||||
runtime, "_load_tai_pack_quant_mla_kv_kernel", return_value=None
|
||||
), self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
r"\[CP_SHARED_KV_FAIL_FAST\]\[fp8_current_pack\]",
|
||||
):
|
||||
runtime.pack_current_mla_kv_for_reuse(
|
||||
k_nope,
|
||||
k_rope,
|
||||
kv_cache=fp8_kv_cache,
|
||||
)
|
||||
|
||||
def test_runtime_fallback_helpers_use_standard_warning_marker(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
|
||||
|
||||
Reference in New Issue
Block a user