diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index 9b220b184..8d14f7498 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/python/sglang/srt/models/deepseek_v4_nextn.py @@ -7,9 +7,17 @@ from torch import nn from transformers import PretrainedConfig from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size +from sglang.srt.layers.attention.nsa.utils import ( + can_nsa_cp_split, + is_nsa_enable_prefill_cp, + is_nsa_prefill_cp_round_robin_split, + nsa_use_prefill_cp, +) from sglang.srt.layers.dp_attention import ( _DpGatheredBufferWrapper, dp_gather_partial, + get_attention_cp_rank, + get_attention_cp_size, get_attention_dp_size, is_dp_attention_enabled, ) @@ -18,6 +26,12 @@ from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe.utils import get_moe_a2a_backend from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.utils.cp_utils import ( + cp_all_gather_rerange_output, + cp_split_and_rebuild_data, + cp_split_and_rebuild_position, + prepare_context_parallel_metadata, +) from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -90,6 +104,12 @@ class DeepseekV4ModelNextN(nn.Module): compress_ratio_override=COMPRESS_RATIO_NEXTN_LAYER, ) + self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() + if self.nsa_enable_prefill_cp: + self.cp_size = get_attention_cp_size() + else: + self.cp_size = None + self.shared_head = nn.Module() self.shared_head.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -145,6 +165,10 @@ class DeepseekV4ModelNextN(nn.Module): else: input_ids_global = input_ids + if nsa_use_prefill_cp(forward_batch): + hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) + positions = cp_split_and_rebuild_position(forward_batch, positions) + hidden_states = self.decoder( positions=positions, hidden_states=hidden_states, @@ -153,6 +177,14 @@ class DeepseekV4ModelNextN(nn.Module): input_ids_global=input_ids_global, ) + if nsa_use_prefill_cp(forward_batch): + hidden_states = cp_all_gather_rerange_output( + hidden_states, + self.cp_size, + forward_batch, + torch.cuda.current_stream(), + ) + pre_hc_head = hidden_states.flatten(1) hidden_states = self.hc_head( @@ -177,6 +209,13 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM): self.pp_group = get_pp_group() self.quant_config = quant_config self.determine_num_fused_shared_experts() + self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() + if self.nsa_enable_prefill_cp: + self.cp_rank = get_attention_cp_rank() + self.cp_size = get_attention_cp_size() + else: + self.cp_rank = None + self.cp_size = None self.model = DeepseekV4ModelNextN( config, quant_config, prefix=add_prefix("model", prefix) @@ -197,6 +236,26 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM): positions: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: + if self.nsa_enable_prefill_cp: + if can_nsa_cp_split(len(input_ids), self.cp_size, True, forward_batch): + forward_batch.attn_cp_metadata = prepare_context_parallel_metadata( + len(input_ids), + self.cp_rank, + self.cp_size, + forward_batch.seq_lens_cpu.tolist(), + ) + if is_nsa_prefill_cp_round_robin_split(): + metadata = forward_batch.attn_backend.forward_metadata + core_meta = metadata.core_attn_metadata + core_meta.apply_cp_reindex() + core_meta.init_flashmla_related() + if metadata.indexer_metadata is not None: + metadata.indexer_metadata = ( + forward_batch.attn_backend.init_forward_metadata_indexer( + core_meta + ) + ) + hidden_states, pre_hc_head = self.model(input_ids, positions, forward_batch) return self.logits_processor( input_ids, diff --git a/test/registered/dsv4/test_deepseek_v4_flash_fp4_b200.py b/test/registered/dsv4/test_deepseek_v4_flash_fp4_b200.py index 212f2b6b2..c3500c1bd 100644 --- a/test/registered/dsv4/test_deepseek_v4_flash_fp4_b200.py +++ b/test/registered/dsv4/test_deepseek_v4_flash_fp4_b200.py @@ -130,5 +130,51 @@ class TestDSV4FlashFP4B200Balanced(ServerSanityMixin, CustomTestCase): _gsm8k_check(self) +class TestDSV4FlashFP4B200Balanced_CP(ServerSanityMixin, CustomTestCase): + """Balanced recipe: TP=4, DP=4, DeepEP, EAGLE (1-step spec).""" + + @classmethod + def setUpClass(cls): + cls.model = try_cached_model(MODEL) + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=[ + "--trust-remote-code", + "--tp", + "4", + "--attn-cp-size", + "4", + "--enable-dp-attention", + "--moe-a2a-backend", + "deepep", + "--speculative-algorithm", + "EAGLE", + "--speculative-num-steps", + "1", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "2", + "--enable-nsa-prefill-context-parallel", + "--nsa-prefill-cp-mode", + "round-robin-split", + "--deepep-config", + DEEPEP_CONFIG, + ], + env=_DEEPEP_ENV, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + _gsm8k_check(self) + + if __name__ == "__main__": unittest.main()