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

DeepSeek V4 MTP Support CP (#24934)

Browse Source 
Co-authored-by: zengpai <zengpai@baidu.com>
This commit is contained in:
Paiiii
2026-05-19 16:51:31 -07:00
committed by GitHub
co-authored by zengpai
parent beaff00331
commit 425dffbde3
2 changed files with 105 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v4_nextn.py
+59
View File
@@ -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,