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

Avoid full-prompt embedding in CP MTP prefill

Browse Source 
The CP draft shared-KV path only needs this rank's local draft tokens, but the previous compatibility path embedded the full prompt before CP-splitting. For long MTP/EAGLE prefill this recreates the large hidden activation that CP shared KV is trying to avoid.\n\nThis pads local draft input ids to the per-rank max token count recorded in NSA CP metadata, embeds the padded local tensor, then trims back to the true local length. That keeps rank shapes compatible while avoiding full-prompt embedding on every rank. Missing or stale metadata keeps the existing full-embedding fallback.\n\nConstraint: CP ranks can own uneven token counts, so the local embedding path needs a rank-uniform padded shape.\nRejected: Pad local ids to the full prompt length | this preserves compatibility but loses the intended memory reduction.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not remove the full-embedding fallback unless all CP draft metadata producers guarantee max_rank_len for every prefill path.\nTested: g0034 container py_compile for utils.py and deepseek_nextn.py; g0034 container pytest -q test/registered/unit/layers/test_nsa_cp_utils.py => 25 passed, 5 warnings.\nNot-tested: Full distributed E2E with HiCache cache-hit MTP accept-length recovery.
This commit is contained in:
laoyao0822
2026-05-23 15:48:36 +08:00
parent bacad1d498
commit ec7e9fbc57
3 changed files with 133 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
python/sglang/srt/layers/attention/nsa/utils.py
View File

@@ -461,6 +461,45 @@ def cp_split_and_rebuild_1d(forward_batch, input_: torch.Tensor):
).view(-1)
def get_cp_local_embedding_padded_token_count(forward_batch, local_num_tokens: int):
metadata = getattr(forward_batch, "nsa_cp_metadata", None)
max_rank_len = getattr(metadata, "max_rank_len", None)
if not max_rank_len:
return None
try:
padded_token_count = int(max_rank_len[0])
except (TypeError, ValueError, IndexError):
return None
if padded_token_count < int(local_num_tokens):
return None
return padded_token_count
def pad_cp_local_input_ids_for_embedding(
forward_batch,
local_input_ids: torch.Tensor,
*,
pad_token_id: int = 0,
):
local_num_tokens = local_input_ids.shape[0]
padded_token_count = get_cp_local_embedding_padded_token_count(
forward_batch, local_num_tokens
)
if padded_token_count is None:
return None
if padded_token_count == local_num_tokens:
return local_input_ids
pad_input_ids = local_input_ids.new_full(
(padded_token_count - local_num_tokens,), pad_token_id
)
return torch.cat((local_input_ids, pad_input_ids), dim=0)
def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
"""Return this CP rank's local logical out_cache_loc for direct writes.

48
python/sglang/srt/models/deepseek_nextn.py
View File

@@ -32,8 +32,10 @@ from sglang.srt.layers.attention.nsa.utils import (
cp_split_and_rebuild_1d,
cp_split_and_rebuild_data,
cp_split_and_rebuild_position,
get_cp_local_embedding_padded_token_count,
is_nsa_enable_prefill_cp,
nsa_use_prefill_cp,
pad_cp_local_input_ids_for_embedding,
prepare_input_dp_with_cp_dsa,
)
from sglang.srt.layers.dp_attention import (
@@ -161,6 +163,33 @@ class DeepseekModelNextN(nn.Module):
)
return None
def _embed_cp_local_input_ids(
self,
forward_batch: ForwardBatch,
local_input_ids: torch.Tensor,
*,
full_num_tokens: int,
) -> Optional[torch.Tensor]:
local_num_tokens = local_input_ids.shape[0]
padded_token_count = get_cp_local_embedding_padded_token_count(
forward_batch, local_num_tokens
)
if padded_token_count is None:
self._debug_cp_draft_shared_kv(
"fallback reason=missing_or_stale_embedding_pad_len "
f"full_tokens={full_num_tokens} local_tokens={local_num_tokens}"
)
return None
local_input_ids = pad_cp_local_input_ids_for_embedding(
forward_batch, local_input_ids
)
hidden_states = self.embed_tokens(local_input_ids)
if hidden_states.shape[0] != local_num_tokens:
hidden_states = hidden_states[:local_num_tokens]
return hidden_states
def forward(
self,
input_ids: torch.Tensor,
@@ -179,9 +208,8 @@ class DeepseekModelNextN(nn.Module):
use_cp = nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp)
use_cp_local_draft = use_cp and envs.SGLANG_CP_DRAFT_SHARED_KV.get()
if use_cp_local_draft:
local_num_tokens = cp_split_and_rebuild_1d(
forward_batch, input_ids
).shape[0]
local_input_ids = cp_split_and_rebuild_1d(forward_batch, input_ids)
local_num_tokens = local_input_ids.shape[0]
local_positions = cp_split_and_rebuild_position(forward_batch, positions)
spec_hidden_states = self._get_cp_local_spec_hidden_states(
forward_batch,
@@ -194,11 +222,17 @@ class DeepseekModelNextN(nn.Module):
else:
positions = local_positions
if input_embeds is None:
# Embed full input first so all ranks see the same tensor
# shape in the TP all-reduce, then CP-split the result.
hidden_states = cp_split_and_rebuild_data(
forward_batch, self.embed_tokens(input_ids)
hidden_states = self._embed_cp_local_input_ids(
forward_batch,
local_input_ids,
full_num_tokens=input_ids.shape[0],
)
if hidden_states is None:
# Conservative compatibility fallback: embed full input
# so all TP ranks all-reduce the same shape, then CP-split.
hidden_states = cp_split_and_rebuild_data(
forward_batch, self.embed_tokens(input_ids)
)
elif input_embeds.shape[0] == local_num_tokens:
hidden_states = input_embeds
elif input_embeds.shape[0] == input_ids.shape[0]:

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

@@ -12,6 +12,8 @@ from sglang.srt.layers.attention.nsa.utils import (
cp_split_and_rebuild_1d,
get_cp_shared_kv_local_out_cache_loc,
get_cp_shared_kv_local_physical_out_cache_loc,
get_cp_local_embedding_padded_token_count,
pad_cp_local_input_ids_for_embedding,
split_in_seq_cp_local_pair,
)
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -293,6 +295,57 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(local_locs.tolist(), [2, 3, 12, 13])
def test_cp_local_embedding_pad_len_uses_metadata_max_rank_len(self):
from types import SimpleNamespace
import torch
forward_batch = SimpleNamespace(
nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[4096] * 8)
)
self.assertEqual(
get_cp_local_embedding_padded_token_count(forward_batch, 4040), 4096
)
self.assertEqual(
get_cp_local_embedding_padded_token_count(forward_batch, 4096), 4096
)
self.assertEqual(
pad_cp_local_input_ids_for_embedding(
SimpleNamespace(
nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[6] * 8)
),
torch.tensor([11, 12, 13, 14]),
).tolist(),
[11, 12, 13, 14, 0, 0],
)
self.assertEqual(
pad_cp_local_input_ids_for_embedding(
SimpleNamespace(
nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[4] * 8)
),
torch.tensor([11, 12, 13, 14]),
).tolist(),
[11, 12, 13, 14],
)
missing_metadata = SimpleNamespace(nsa_cp_metadata=None)
self.assertIsNone(
get_cp_local_embedding_padded_token_count(missing_metadata, 4040)
)
self.assertIsNone(
pad_cp_local_input_ids_for_embedding(
missing_metadata, torch.tensor([11, 12, 13, 14])
)
)
stale_metadata = SimpleNamespace(
nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[4039] * 8)
)
self.assertIsNone(
get_cp_local_embedding_padded_token_count(stale_metadata, 4040)
)
def test_local_out_cache_loc_requires_compute_owner_pages(self):
import torch
from types import SimpleNamespace