diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index fd42d48bc..d496fc85d 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -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. diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 21ded56d7..6f96807fe 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -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]: diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 6d5bdddee..79756b170 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -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