chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Prevent stale CP shared-KV contracts from corrupting prefill

Browse Source 
CP shared-KV now uses CP-local current rows consistently across MLA/index current reuse, passes fp8 current-index K through the tai-kernel uint8 ABI, and clears the transient EAGLE CP-local hidden marker after draft capture. The disaggregation bootstrap also fingerprints the runtime source contract so prefill/decode mismatches fail fast instead of silently exchanging incompatible KV metadata.

Constraint: CP shared-KV batch paths flatten current K/V rows in CP-rank-local valid order, not global request order.

Constraint: tai-kernel current-index prepare validates current_index_k as uint8 bytes for fp8 payloads.

Rejected: Keep using global extend offsets for bs>1 current-index reuse | corrupts request-local bases once current_index_kv is CP-local.

Rejected: Infer CP-local EAGLE hidden semantics from tensor shape | static padding and bs>1 can make shape-based inference unsafe.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce forward_batch.out_cache_loc slicing in CP shared-KV current reuse without verifying CP-local owner-lane layout.

Tested: Remote container py_compile for touched runtime/test files.

Tested: Remote PYTHONPATH=python pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/disaggregation/test_common_conn_runtime_fingerprint.py (198 passed, 2 subtests passed).

Not-tested: Full remote ETE traffic after this commit; accept length and garbage-output recovery still require a fresh prefill/decode run.

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-04 20:22:29 +08:00
parent 3d6007246b
commit f50e2b1e00
9 changed files with 1463 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

93
test/registered/unit/disaggregation/test_common_conn_runtime_fingerprint.py Normal file
View File

@@ -0,0 +1,93 @@
import types
import unittest
from unittest.mock import Mock, patch
from sglang.srt.disaggregation.common import conn as common_conn
class TestDisaggregationRuntimeFingerprint(unittest.TestCase):
def _manager(self):
mgr = common_conn.CommonKVManager.__new__(common_conn.CommonKVManager)
mgr.prefill_info_table = {}
mgr.kv_args = types.SimpleNamespace(page_size=64, engine_rank=0)
mgr.server_args = types.SimpleNamespace(kv_cache_dtype="fp8_e4m3")
mgr._resolve_rank_mapping = Mock()
return mgr
def _response(self, payload, status_code=200):
response = Mock()
response.status_code = status_code
response.json.return_value = payload
response.text = ""
return response
def _payload(self, **overrides):
payload = {
"attn_tp_size": 8,
"attn_cp_size": 8,
"dp_size": 1,
"pp_size": 1,
"page_size": 64,
"kv_cache_dtype": "fp8_e4m3",
"follow_bootstrap_room": True,
}
payload.update(overrides)
return payload
def test_try_ensure_parallel_info_rejects_runtime_source_fingerprint_mismatch(self):
mgr = self._manager()
with patch.object(
common_conn.requests,
"get",
return_value=self._response(
self._payload(
runtime_source_fingerprint="prefill-source-hash",
runtime_source_root="/sgl-workspace/sglang-tai",
)
),
), patch.object(
common_conn,
"get_runtime_source_fingerprint",
return_value="decode-source-hash",
), patch.object(
common_conn,
"get_runtime_source_root",
return_value="/sgl-workspace/sglang",
):
with self.assertRaisesRegex(
RuntimeError,
"Runtime source fingerprint mismatch",
):
mgr.try_ensure_parallel_info("g0034:18992")
mgr._resolve_rank_mapping.assert_not_called()
self.assertEqual(mgr.prefill_info_table, {})
def test_try_ensure_parallel_info_accepts_matching_runtime_source_fingerprint(self):
mgr = self._manager()
with patch.object(
common_conn.requests,
"get",
return_value=self._response(
self._payload(
runtime_source_fingerprint="same-source-hash",
runtime_source_root="/sgl-workspace/sglang-tai",
)
),
), patch.object(
common_conn,
"get_runtime_source_fingerprint",
return_value="same-source-hash",
), patch.object(
common_conn,
"get_runtime_source_root",
return_value="/sgl-workspace/sglang-tai",
):
self.assertTrue(mgr.try_ensure_parallel_info("g0034:18992"))
mgr._resolve_rank_mapping.assert_called_once()
self.assertIn("g0034:18992", mgr.prefill_info_table)
if __name__ == "__main__":
unittest.main()

138
test/registered/unit/layers/test_nsa_cp_utils.py
View File

@@ -1322,6 +1322,28 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(local[:, 0].tolist(), list(range(0, 16, 2)))
def test_cp_split_and_rebuild_data_preserves_non_token_dimensions(self):
import torch
forward_batch = SimpleNamespace(
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=2,
request_extend_lens=[4, 9],
request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]],
request_zigzag_indices=[[0, 3], [0, 3]],
)
)
tensor = torch.arange(13 * 2 * 3, dtype=torch.float32).view(13, 2, 3)
with patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
):
local = cp_split_and_rebuild_data(forward_batch, tensor)
self.assertEqual(tuple(local.shape), (8, 2, 3))
self.assertTrue(torch.equal(local[0], tensor[0]))
def test_cp_split_and_rebuild_data_uses_compute_padding_rows(self):
import torch
@@ -3309,6 +3331,122 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(deep_gemm_calls[0]["ks"], [0, 0, 3, 6, 10, 10, 10])
self.assertEqual(deep_gemm_calls[0]["ke"], [2, 3, 6, 10, 12, 13, 14])
def test_indexer_ragged_cp_index_current_batch_uses_cp_local_bases_and_uint8_k(
self,
):
import contextlib
import torch
from types import SimpleNamespace
from unittest.mock import patch
from sglang.srt.layers.attention.nsa import nsa_indexer
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
indexer = object.__new__(Indexer)
indexer.index_topk = 2
indexer._with_real_sm_count = lambda: contextlib.nullcontext()
prepare_calls = []
def fake_prepare(**kwargs):
prepare_calls.append(kwargs)
total_kv_len = int(kwargs["total_kv_len"])
return (
torch.zeros((total_kv_len, 1), dtype=torch.uint8),
torch.zeros((total_kv_len,), dtype=torch.float32),
torch.tensor([0, 0, 3, 3, 3], dtype=torch.int32),
torch.tensor([2, 3, 2, 3, 4], dtype=torch.int32),
)
class Metadata:
def get_page_table_64(self):
raise AssertionError("current cp_index path must not materialize index pages")
def topk_transform(self, logits, topk, **kwargs):
return (
torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32)
.view(-1, 1)
.repeat(1, topk)
)
forward_batch = SimpleNamespace(
token_to_kv_pool=SimpleNamespace(page_size=64, index_head_dim=1),
seq_lens_cpu=torch.tensor([5, 7], dtype=torch.int64),
extend_seq_lens_cpu=[5, 7],
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=2,
batch_plan=SimpleNamespace(
request_rank_local_offsets=[0, 2],
request_valid_rank_local_offsets=[0, 2],
),
),
)
current_index_k = (
torch.arange(6, dtype=torch.float32)
.to(torch.float8_e4m3fn)
.view(6, 1)
)
current_index_scale = torch.arange(6, dtype=torch.float32).view(6, 1)
current_index_kv = (current_index_k, current_index_scale)
with patch.object(
nsa_indexer,
"try_tai_prepare_cp_mqa_current_index_batch",
side_effect=fake_prepare,
create=True,
), patch.object(
nsa_indexer,
"deep_gemm",
SimpleNamespace(
fp8_mqa_logits=lambda q_fp8, kv_fp8, weights, ks, ke, clean_logits=False: torch.zeros(
(int(q_fp8.shape[0]), 8), dtype=torch.float32
)
),
):
result = Indexer._get_topk_ragged_with_cp(
indexer,
forward_batch,
layer_id=7,
q_fp8=torch.empty((5, 1), dtype=torch.float32),
weights=torch.empty((5, 1, 1), dtype=torch.float32),
metadata=Metadata(),
kv_len=0,
actual_seq_q=5,
cp_index=[(0, 1, 3), (1, 1, 4)],
current_index_kv=current_index_kv,
)
self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]])
self.assertEqual(len(prepare_calls), 1)
call = prepare_calls[0]
self.assertEqual(call["current_index_k"].dtype, torch.uint8)
self.assertEqual(call["current_bases"].tolist(), [0, 2])
self.assertEqual(call["kv_lens"].tolist(), [3, 4])
self.assertEqual(call["q_lens"].tolist(), [2, 3])
def test_eagle_capture_for_decode_clears_cp_local_hidden_marker(self):
import torch
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.speculative.eagle_worker import EAGLEWorker
worker = object.__new__(EAGLEWorker)
worker.topk = 1
draft_input = EagleDraftInput(
hidden_states=torch.full((4, 2), -1.0),
cp_local_hidden_states=True,
)
draft_output_hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
logits_output = LogitsProcessorOutput(
next_token_logits=torch.tensor([[0.1, 0.9], [0.7, 0.3]]),
hidden_states=draft_output_hidden,
)
worker.capture_for_decode(logits_output, draft_input)
self.assertIs(draft_input.hidden_states, draft_output_hidden)
self.assertFalse(draft_input.cp_local_hidden_states)
def test_indexer_in_seq_cp_pair_skips_materialize_when_current_index_reused(self):
import torch

104
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
View File

@@ -729,7 +729,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertIn("valid_current_rows=int(current_kv_rows_for_reuse)", body_source)
self.assertIn("pack_current_mla_kv_for_reuse", body_source)
self.assertIn("forward_batch.out_cache_loc[:valid_current_rows]", body_source)
self.assertIn("get_cp_shared_kv_local_out_cache_loc", body_source)
self.assertNotIn("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
@@ -755,12 +756,16 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
Path(__file__).resolve().parents[4]
/ "python/sglang/srt/layers/attention/nsa_backend.py"
).read_text()
method_start = source.index(" current_locs_for_reuse = None")
branch_start = source.index("eagle_draft_mla_branch = \"partial_current_sync\"")
branch_end = source.index("forward_partial_current_sync_compose", branch_start)
branch_source = source[branch_start:branch_end]
local_locs_source = source[method_start:branch_end]
self.assertIn("build_batch_prefix_slot_span", source)
self.assertIn("prefix_slot_span=", branch_source)
self.assertIn("get_cp_shared_kv_local_out_cache_loc", local_locs_source)
self.assertNotIn("current_locs = forward_batch.out_cache_loc", branch_source)
self.assertNotIn("or len(prefix_lens_cpu) != 1", branch_source)
def test_nsa_backend_topk_transform_uses_effective_forward_impl(self):
@@ -1504,12 +1509,18 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
Path(__file__).resolve().parents[4]
/ "python/sglang/srt/layers/attention/nsa/nsa_indexer.py"
).read_text()
branch_start = source.index(" if current_index_kv is not None:")
method_start = source.index(" def _maybe_materialize_shared_index_buffer")
branch_start = source.index(
" if current_index_kv is not None:", method_start
)
branch_end = source.index(" return materialized, dense_pages", branch_start)
branch_source = source[branch_start:branch_end]
branch_compact = "".join(branch_source.split())
self.assertIn("build_batch_prefix_slot_span", source)
self.assertIn("prefix_slot_span=", branch_source)
self.assertIn("get_cp_shared_kv_local_out_cache_loc", branch_source)
self.assertNotIn("current_locs=forward_batch.out_cache_loc", branch_compact)
self.assertNotIn("or len(prefix_lens_cpu) != 1", branch_source)
def test_ipc_page_descriptor_builder_maps_slots_to_owner_physical_pages(self):
@@ -2141,7 +2152,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
prepare_source = source[start:end]
self.assertIn("valid_current_rows", prepare_source)
self.assertIn("key[:valid_current_rows]", prepare_source)
self.assertIn("get_cp_shared_kv_local_out_cache_loc", prepare_source)
self.assertIn("cp_split_and_rebuild_data", prepare_source)
self.assertNotIn("forward_batch.out_cache_loc[:valid_current_rows]", prepare_source)
self.assertNotIn(
"key.shape[0] == forward_batch.out_cache_loc.numel()",
prepare_source,
@@ -2852,6 +2865,91 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
load_kernel.assert_not_called()
self.assertTrue(any("debug_enabled" in message for message in logs.output))
@unittest.skipIf(not torch.cuda.is_available(), "CUDA required")
def test_fp8_fused_mla_store_matches_sglang_fallback_for_cp_owned_tail_pages(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
from sglang.srt.layers.attention.nsa.quant_k_cache import (
quantize_k_cache_separate,
)
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.srt.mem_cache.utils import set_mla_kv_buffer_triton
try:
runtime._load_tai_fused_mla_store_kernel()
except Exception as exc:
self.skipTest(f"TAI fused MLA store kernel unavailable: {exc}")
page_size = 64
cp_size = 8
rank = 7
n_tokens = page_size * 3 + 17
layout = CpSharedKVLayout(
page_size=page_size,
cp_size=cp_size,
cp_rank=rank,
)
local_token_ids = torch.arange(n_tokens, device="cuda", dtype=torch.long)
local_pages = torch.div(local_token_ids, page_size, rounding_mode="floor")
logical_pages = 1 + rank + cp_size * local_pages
logical_locs = (
logical_pages * page_size + torch.remainder(local_token_ids, page_size)
).to(torch.int64)
physical_locs = layout.logical_locs_to_physical(logical_locs)
capacity_tokens = int(physical_locs.max().item()) + page_size + 1
torch.manual_seed(20260604)
k_nope = (
torch.randn((n_tokens, 1, 512), device="cuda", dtype=torch.bfloat16)
* 2.0
) + 0.25
latent_cache = torch.randn(
(n_tokens, 576), device="cuda", dtype=torch.bfloat16
)
k_rope = latent_cache[:, 512:].unsqueeze(1)
self.assertFalse(k_rope.is_contiguous())
expected = torch.zeros(
(capacity_tokens, 1, 656), dtype=torch.uint8, device="cuda"
)
nope_part, rope_part = quantize_k_cache_separate(k_nope, k_rope)
set_mla_kv_buffer_triton(expected, physical_locs, nope_part, rope_part)
class FakePool:
def __init__(self):
self.nsa_kv_cache_store_fp8 = True
self.page_size = page_size
self.start_layer = 0
self.kv_buffer = [
torch.zeros(
(capacity_tokens, 1, 656),
dtype=torch.uint8,
device="cuda",
)
]
pool = FakePool()
layer = SimpleNamespace(layer_id=0)
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):
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,
)
torch.cuda.synchronize()
self.assertTrue(used)
torch.testing.assert_close(
pool.kv_buffer[0],
expected,
atol=0,
rtol=0,
)
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