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]: