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

Stabilize CP shared-KV batch padding semantics

Browse Source 
CP shared-KV bs>1 exposed three distinct padding domains: valid cache rows, CP page-tail compute rows, and MLP-sync flattened static padding. The previous implementation mixed these domains in direct-write and index top-k paths, so real requests failed when q/out_cache_loc lengths matched valid rows while metadata aliases described compute rows.\n\nThis change makes compute split strip only proven flattened static padding, keeps valid cache writes strict except for extend_num_tokens-proven static tails, marks CP-local EAGLE draft hidden state explicitly, and selects NSA index top-k query metadata by the actual q/weight row count.\n\nConstraint: CP shared-KV cache writes must never persist dummy page-tail or MLP static padding rows.\nConstraint: EAGLE draft hidden state can be CP-local before full CP metadata is visible in prepare_mlp_sync_batch.\nRejected: Use compute_padding_enabled as direct-write truncation proof | it silently accepts unknown out_cache_loc tails.\nRejected: Always consume compute q metadata in index top-k | actual q/weights can be valid-only after CP split.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not collapse valid rows, CP compute padding, and MLP static padding into one length condition; use explicit provenance.\nTested: remote py_compile for touched NSA files\nTested: remote targeted CP shared-KV padding/top-k regressions\nTested: remote pytest test_nsa_cp_utils.py test_cp_shared_kv_layout.py test_cp_shared_kv_runtime.py -k 'not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel' => 228 passed, 1 deselected, 5 warnings, 2 subtests passed\nNot-tested: full ETE replay after the final index top-k fix\nNot-tested: TAI current-index fast path dtype fallback
This commit is contained in:
laoyao0822
2026-06-04 07:25:11 +08:00
parent 02af370e87
commit 3d6007246b
7 changed files with 1403 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/layers/test_nsa_cp_utils.py
+515 -2
View File
@@ -27,6 +27,7 @@ from sglang.srt.layers.attention.nsa.utils import (
get_cp_shared_kv_local_out_cache_loc,
get_cp_shared_kv_local_physical_out_cache_loc,
get_cp_local_embedding_padded_token_count,
nsa_use_prefill_cp,
pad_cp_local_input_ids_for_embedding,
prepare_input_dp_with_cp_dsa,
select_cp_local_valid_rows_for_cache_write,
@@ -34,7 +35,9 @@ from sglang.srt.layers.attention.nsa.utils import (
split_in_seq_cp_local_pair,
)
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.models.deepseek_nextn import DeepseekModelNextN
from sglang.srt.speculative.eagle_info import EagleDraftInput
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
@@ -326,6 +329,45 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
):
self.assertTrue(can_cp_split(64, 8, True, forward_batch))
def test_can_cp_split_enables_cp_draft_shared_kv_draft_extend(self):
class DraftMode:
def is_context_parallel_extend(self, include_draft_extend_v2=False):
return False
def is_draft_extend(self, include_v2=False):
return True
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
extend_seq_lens_cpu=[4],
extend_prefix_lens_cpu=[0],
token_to_kv_pool=SimpleNamespace(page_size=64),
forward_mode=DraftMode(),
)
with (
patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}),
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
),
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
return_value=True,
),
):
self.assertTrue(can_cp_split(8, 8, True, forward_batch))
def test_nsa_use_prefill_cp_enables_cp_draft_shared_kv_draft_extend(self):
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
forward_mode=ForwardMode.DRAFT_EXTEND,
nsa_cp_metadata=NSAContextParallelMetadata(batch_size=1),
)
with patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}):
self.assertTrue(nsa_use_prefill_cp(forward_batch, True))
def test_can_cp_split_uses_compute_padding_per_request_for_batched_tiny_suffix(
self,
):
@@ -512,8 +554,8 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(plan.request_last_token_local_offset, [0])
# Compatibility aliases for cache/page accounting stay valid-token
# based. Query-length metadata is split separately below: attention and
# top-k consume compute rows, cache/current paths consume valid rows.
# based. Query-length metadata exposes both valid and compute rows:
# consumers must choose the view that matches their actual q layout.
self.assertEqual(plan.request_split_lists, plan.request_valid_split_lists)
self.assertEqual(plan.request_padded_pages, plan.request_valid_padded_pages)
self.assertEqual(plan.request_actual_seq_q_prev, [64])
@@ -523,6 +565,67 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(plan.request_compute_seq_q_prev, [64])
self.assertEqual(plan.request_compute_seq_q_next, [0])
def test_index_topk_batch_lengths_follow_actual_q_rows_not_compute_alias(self):
from sglang.srt.layers.attention.nsa.nsa_indexer import (
_select_batch_topk_query_lengths,
)
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[40387],
prefix_lens=[0],
page_size=64,
cp_size=8,
cp_rank=0,
)
valid_local_rows = (
plan.request_valid_seq_q_prev[0] + plan.request_valid_seq_q_next[0]
)
compute_local_rows = (
plan.request_compute_seq_q_prev[0] + plan.request_compute_seq_q_next[0]
)
self.assertEqual(valid_local_rows, 4995)
self.assertEqual(compute_local_rows, 5056)
selected = _select_batch_topk_query_lengths(
cp_metadata=NSAContextParallelMetadata(batch_size=1, batch_plan=plan),
batch_plan=plan,
batch_size=1,
q_tokens=valid_local_rows,
weights_tokens=valid_local_rows,
)
self.assertFalse(selected.uses_compute_query_rows)
self.assertEqual(selected.request_seq_q_prev, plan.request_valid_seq_q_prev)
self.assertEqual(selected.request_seq_q_next, plan.request_valid_seq_q_next)
self.assertEqual(
selected.request_valid_seq_q_prev, plan.request_valid_seq_q_prev
)
self.assertEqual(
selected.request_valid_seq_q_next, plan.request_valid_seq_q_next
)
selected_compute = _select_batch_topk_query_lengths(
cp_metadata=NSAContextParallelMetadata(batch_size=1, batch_plan=plan),
batch_plan=plan,
batch_size=1,
q_tokens=compute_local_rows,
weights_tokens=compute_local_rows,
)
self.assertTrue(selected_compute.uses_compute_query_rows)
self.assertEqual(
selected_compute.request_seq_q_prev, plan.request_compute_seq_q_prev
)
self.assertEqual(
selected_compute.request_seq_q_next, plan.request_compute_seq_q_next
)
self.assertEqual(
selected_compute.request_valid_seq_q_prev, plan.request_valid_seq_q_prev
)
self.assertEqual(
selected_compute.request_valid_seq_q_next, plan.request_valid_seq_q_next
)
def test_batch_plan_compute_padding_is_per_request_not_batch_total(self):
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[65, 1024],
@@ -752,6 +855,166 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
local_num_tokens=8,
)
def test_cp_draft_padding_keeps_local_hidden_when_static_tokens_are_shorter(self):
import torch
class FakeAttnBackend:
def get_cuda_graph_seq_len_fill_value(self):
return 0
model_runner = SimpleNamespace(attn_backend=FakeAttnBackend())
spec_info = EagleDraftInput(
hidden_states=torch.ones((64, 2), dtype=torch.float32),
verified_id=torch.tensor([1], dtype=torch.int64),
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
cp_local_hidden_states=True,
)
forward_batch = ForwardBatch(
forward_mode=ForwardMode.DRAFT_EXTEND,
batch_size=1,
input_ids=torch.arange(8, dtype=torch.int64),
req_pool_indices=torch.tensor([0], dtype=torch.int64),
seq_lens=torch.tensor([4], dtype=torch.int32),
out_cache_loc=torch.arange(8, dtype=torch.int64),
seq_lens_sum=4,
positions=torch.arange(8, dtype=torch.int64),
lora_ids=[None],
spec_info=spec_info,
uses_cp_shared_kv=True,
)
with patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}):
forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1)
self.assertEqual(tuple(forward_batch.spec_info.hidden_states.shape), (64, 2))
self.assertTrue(
torch.equal(
forward_batch.hidden_states_backup,
torch.ones((64, 2), dtype=torch.float32),
)
)
def test_cp_draft_padding_keeps_marked_cp_local_hidden_before_cp_flags_are_visible(
self,
):
import torch
class FakeAttnBackend:
def get_cuda_graph_seq_len_fill_value(self):
return 0
model_runner = SimpleNamespace(attn_backend=FakeAttnBackend())
spec_info = EagleDraftInput(
hidden_states=torch.ones((64, 2), dtype=torch.float32),
verified_id=torch.tensor([1], dtype=torch.int64),
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
cp_local_hidden_states=True,
)
forward_batch = ForwardBatch(
forward_mode=ForwardMode.DRAFT_EXTEND,
batch_size=1,
input_ids=torch.arange(8, dtype=torch.int64),
req_pool_indices=torch.tensor([0], dtype=torch.int64),
seq_lens=torch.tensor([4], dtype=torch.int32),
out_cache_loc=torch.arange(8, dtype=torch.int64),
seq_lens_sum=4,
positions=torch.arange(8, dtype=torch.int64),
lora_ids=[None],
spec_info=spec_info,
uses_cp_shared_kv=False,
)
forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1)
self.assertEqual(tuple(forward_batch.spec_info.hidden_states.shape), (64, 2))
self.assertTrue(
torch.equal(
forward_batch.hidden_states_backup,
torch.ones((64, 2), dtype=torch.float32),
)
)
def test_cp_draft_padding_keeps_marked_cp_local_hidden_after_forward_mode_rewrite(
self,
):
import torch
class FakeAttnBackend:
def get_cuda_graph_seq_len_fill_value(self):
return 0
model_runner = SimpleNamespace(attn_backend=FakeAttnBackend())
spec_info = EagleDraftInput(
hidden_states=torch.ones((64, 2), dtype=torch.float32),
verified_id=torch.tensor([1], dtype=torch.int64),
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
cp_local_hidden_states=True,
)
forward_batch = ForwardBatch(
# prepare_mlp_sync_batch can temporarily rewrite draft extend to
# EXTEND while static DP padding is being prepared. The draft
# side-channel contract must therefore be carried by spec_info, not
# inferred from forward_mode.
forward_mode=ForwardMode.EXTEND,
batch_size=1,
input_ids=torch.arange(8, dtype=torch.int64),
req_pool_indices=torch.tensor([0], dtype=torch.int64),
seq_lens=torch.tensor([4], dtype=torch.int32),
out_cache_loc=torch.arange(8, dtype=torch.int64),
seq_lens_sum=4,
positions=torch.arange(8, dtype=torch.int64),
lora_ids=[None],
spec_info=spec_info,
uses_cp_shared_kv=True,
)
forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1)
self.assertEqual(tuple(forward_batch.spec_info.hidden_states.shape), (64, 2))
self.assertTrue(
torch.equal(
forward_batch.hidden_states_backup,
torch.ones((64, 2), dtype=torch.float32),
)
)
def test_cp_draft_padding_rejects_unmarked_oversized_hidden(self):
import torch
class FakeAttnBackend:
def get_cuda_graph_seq_len_fill_value(self):
return 0
model_runner = SimpleNamespace(attn_backend=FakeAttnBackend())
spec_info = EagleDraftInput(
hidden_states=torch.ones((64, 2), dtype=torch.float32),
verified_id=torch.tensor([1], dtype=torch.int64),
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
)
forward_batch = ForwardBatch(
forward_mode=ForwardMode.DRAFT_EXTEND,
batch_size=1,
input_ids=torch.arange(8, dtype=torch.int64),
req_pool_indices=torch.tensor([0], dtype=torch.int64),
seq_lens=torch.tensor([4], dtype=torch.int32),
out_cache_loc=torch.arange(8, dtype=torch.int64),
seq_lens_sum=4,
positions=torch.arange(8, dtype=torch.int64),
lora_ids=[None],
spec_info=spec_info,
uses_cp_shared_kv=False,
)
with self.assertRaisesRegex(
RuntimeError,
r"\[CP_SHARED_KV_FAIL_FAST\]\[draft_hidden_static_padding_mismatch\]",
):
forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1)
def test_full_rerange_fails_fast_for_batch_metadata(self):
import torch
@@ -1087,6 +1350,90 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(local[0].tolist(), [128.0, 129.0])
self.assertTrue(torch.equal(local[1:], torch.zeros((63, 2))))
def test_cp_split_and_rebuild_data_ignores_trailing_static_padding_rows(self):
import torch
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[4],
prefix_lens=[0],
page_size=4,
cp_size=2,
cp_rank=1,
)
forward_batch = SimpleNamespace(
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=1,
batch_plan=plan,
)
)
tensor = torch.arange(8 * 2, dtype=torch.float32).view(8, 2)
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(local.shape, (4, 2))
self.assertTrue(torch.equal(local, torch.zeros((4, 2))))
def test_cp_split_and_rebuild_data_ignores_mlp_sync_static_padding_without_compute_padding(
self,
):
import torch
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[7],
prefix_lens=[0],
page_size=4,
cp_size=2,
cp_rank=1,
)
self.assertFalse(plan.compute_padding_enabled)
forward_batch = SimpleNamespace(
extend_num_tokens=8,
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=1,
batch_plan=plan,
),
)
tensor = torch.arange(8 * 2, dtype=torch.float32).view(8, 2)
expected = split_tensor_by_cp_batch_plan(
tensor[:7],
plan,
mode="data",
)
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.assertTrue(torch.equal(local, expected))
def test_cp_split_valid_kind_rejects_trailing_padding_rows(self):
import torch
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[4],
prefix_lens=[0],
page_size=4,
cp_size=2,
cp_rank=1,
)
with self.assertRaisesRegex(
RuntimeError,
"batch_gt1_split_input_len_mismatch",
):
split_tensor_by_cp_batch_plan(
torch.arange(8),
plan,
mode="1d",
split_kind="valid",
)
def test_cp_split_and_rebuild_1d_keeps_batch_request_boundaries(self):
import torch
@@ -1186,6 +1533,41 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(local.shape, (64,))
self.assertEqual(local.tolist(), list(range(40384, 40448)))
def test_cp_split_and_rebuild_position_ignores_mlp_sync_static_padding_without_compute_padding(
self,
):
import torch
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[7],
prefix_lens=[0],
page_size=4,
cp_size=2,
cp_rank=1,
)
self.assertFalse(plan.compute_padding_enabled)
forward_batch = SimpleNamespace(
extend_num_tokens=8,
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=1,
batch_plan=plan,
),
)
positions = torch.arange(8, dtype=torch.int32)
expected = split_tensor_by_cp_batch_plan(
positions[:7],
plan,
mode="position",
)
with patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
):
local = cp_split_and_rebuild_position(forward_batch, positions)
self.assertTrue(torch.equal(local, expected))
def test_cp_local_embedding_pad_len_uses_metadata_max_rank_len(self):
from types import SimpleNamespace
@@ -1299,6 +1681,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
)
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
extend_num_tokens=8,
cp_shared_kv_layout=CpSharedKVLayout(
page_size=page_size,
cp_size=2,
@@ -1368,6 +1751,136 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(local_locs.tolist(), [2 * page_size])
def test_local_out_cache_loc_ignores_trailing_static_padding_locs(self):
import torch
page_size = 4
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[4],
prefix_lens=[0],
page_size=page_size,
cp_size=2,
cp_rank=0,
)
valid_locs = torch.arange(1 * page_size, 2 * page_size, dtype=torch.int64)
static_padding_locs = torch.arange(
99 * page_size, 100 * page_size, dtype=torch.int64
)
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
extend_num_tokens=8,
cp_shared_kv_layout=CpSharedKVLayout(
page_size=page_size,
cp_size=2,
cp_rank=0,
),
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=1,
batch_plan=plan,
page_aligned=True,
page_size=page_size,
extend_prefix_len=0,
),
out_cache_loc=torch.cat((valid_locs, static_padding_locs)),
)
with patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
):
local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
self.assertEqual(local_locs.tolist(), valid_locs.tolist())
def test_local_out_cache_loc_ignores_mlp_sync_static_padding_without_compute_padding(
self,
):
import torch
page_size = 4
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[7],
prefix_lens=[0],
page_size=page_size,
cp_size=2,
cp_rank=1,
)
self.assertFalse(plan.compute_padding_enabled)
valid_locs = torch.cat(
(
torch.arange(1 * page_size, 2 * page_size, dtype=torch.int64),
torch.arange(2 * page_size, 2 * page_size + 3, dtype=torch.int64),
)
)
static_padding_locs = torch.tensor([99 * page_size], dtype=torch.int64)
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
extend_num_tokens=8,
cp_shared_kv_layout=CpSharedKVLayout(
page_size=page_size,
cp_size=2,
cp_rank=1,
),
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=1,
batch_plan=plan,
page_aligned=True,
page_size=page_size,
extend_prefix_len=0,
),
out_cache_loc=torch.cat((valid_locs, static_padding_locs)),
)
with patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
):
local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
self.assertEqual(
local_locs.tolist(),
[2 * page_size, 2 * page_size + 1, 2 * page_size + 2],
)
def test_local_out_cache_loc_rejects_unproven_trailing_padding_even_with_compute_padding(
self,
):
import torch
page_size = 4
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[4],
prefix_lens=[0],
page_size=page_size,
cp_size=2,
cp_rank=0,
)
self.assertTrue(plan.compute_padding_enabled)
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
cp_shared_kv_layout=CpSharedKVLayout(
page_size=page_size,
cp_size=2,
cp_rank=0,
),
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=1,
batch_plan=plan,
page_aligned=True,
page_size=page_size,
extend_prefix_len=0,
),
out_cache_loc=torch.cat(
(
torch.arange(1 * page_size, 2 * page_size, dtype=torch.int64),
torch.arange(99 * page_size, 100 * page_size, dtype=torch.int64),
)
),
)
with self.assertRaisesRegex(RuntimeError, "static_padded=None"):
get_cp_shared_kv_local_out_cache_loc(forward_batch)
def test_batch_local_physical_out_cache_loc_reuses_layer_invariant_plan(self):
import torch
from types import SimpleNamespace