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

perf(cp): add narrow output collection for NSA prefill CP

Browse Source 
Replace full hidden all-gather at prefill tail with per-request last-token
hidden collection, reducing communication from total_tokens x hidden_size
to cp_size x bs x hidden_size for both in-seq-split and round-robin modes.

- nsa/utils.py: add cp_collect_last_token_hidden() with mode-specific
  narrow collection helpers that only gather the last token hidden
- deepseek_v2.py: add _should_use_narrow_output_path() gate on
  DeepseekV2Model, fallback to full gather for EAGLE/return_logprob/
  capture_hidden batches
- logits_processor.py: add _is_compact_hidden_states() to bypass
  _get_pruned_states() when hidden is already compact
This commit is contained in:
laoyao0822
2026-04-22 04:55:06 +08:00
parent 7403f14511
commit bc9a9f128b
3 changed files with 141 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -137,7 +137,6 @@ def pad_nsa_cache_seqlens(forward_batch: "ForwardBatch", nsa_cache_seqlens):
@dataclass
class NSAContextParallelMetadata:
split_list: List[int] = None
max_rank_len: List[int] = None
zigzag_index: List[int] = None
@@ -158,9 +157,9 @@ class NSAContextParallelMetadata:
def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
if is_nsa_prefill_cp_round_robin_split():
cur_cp_seq_len = seq_len // cp_size
assert (
seq_len % cp_size == 0
), f"seq_len {seq_len} is not divisible by cp_size {cp_size} when nsa_prefill_cp_mode is round-robin-split"
assert seq_len % cp_size == 0, (
f"seq_len {seq_len} is not divisible by cp_size {cp_size} when nsa_prefill_cp_mode is round-robin-split"
)
else:
# TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
# Note: (self.cp_size * 2) To achieve load balancing for seq computation,
@@ -182,9 +181,9 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
if is_nsa_prefill_cp_round_robin_split():
cp_size = get_attention_cp_size()
assert (
input_.shape[0] % cp_size == 0
), f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}"
assert input_.shape[0] % cp_size == 0, (
f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}"
)
return nsa_cp_round_robin_split_data(input_)
input_list = list(
@@ -577,3 +576,54 @@ def prepare_input_dp_with_cp_dsa(
total_seq_lens=kv_len_origin,
)
return nsa_cp_metadata
def cp_collect_last_token_hidden(
hidden_states: torch.Tensor,
forward_batch: "ForwardBatch",
cp_size: int,
) -> torch.Tensor:
if is_nsa_prefill_cp_round_robin_split():
return _round_robin_collect_last_token(hidden_states, forward_batch, cp_size)
return _in_seq_collect_last_token(hidden_states, forward_batch, cp_size)
def _round_robin_collect_last_token(
hidden_states: torch.Tensor,
forward_batch: "ForwardBatch",
cp_size: int,
) -> torch.Tensor:
total_tokens = sum(forward_batch.extend_seq_lens_cpu)
owner = (total_tokens - 1) % cp_size
cp_rank = get_attention_cp_rank()
bs = len(forward_batch.extend_seq_lens_cpu)
cp_group = get_attention_cp_group()
if cp_rank == owner and hidden_states.shape[0] > 0:
local_last = hidden_states[-bs:].contiguous()
else:
local_last = hidden_states.new_zeros((bs, hidden_states.shape[1]))
gathered = hidden_states.new_empty((cp_size * bs, hidden_states.shape[1]))
attn_cp_all_gather_into_tensor(gathered, local_last)
return gathered[owner * bs : owner * bs + bs]
def _in_seq_collect_last_token(
hidden_states: torch.Tensor,
forward_batch: "ForwardBatch",
cp_size: int,
) -> torch.Tensor:
cp_rank = get_attention_cp_rank()
bs = len(forward_batch.extend_seq_lens_cpu)
owner = 0
cp_group = get_attention_cp_group()
if cp_rank == owner and hidden_states.shape[0] > 0:
local_last = hidden_states[-bs:].contiguous()
else:
local_last = hidden_states.new_zeros((bs, hidden_states.shape[1]))
gathered = hidden_states.new_empty((cp_size * bs, hidden_states.shape[1]))
attn_cp_all_gather_into_tensor(gathered, local_last)
return gathered[owner * bs : owner * bs + bs]

31
python/sglang/srt/layers/logits_processor.py
View File

@@ -204,7 +204,6 @@ class LogitsMetadata:
)
def compute_dp_attention_metadata(self):
cumtokens = torch.cumsum(self.global_num_tokens_for_logprob_gpu, dim=0)
dp_rank = get_attention_dp_rank()
if dp_rank == 0:
@@ -309,7 +308,16 @@ class LogitsProcessor(nn.Module):
if logits_metadata.forward_mode.is_dllm_extend():
return self._get_dllm_logits(hidden_states, lm_head, logits_metadata)
# Get the last hidden states and last logits for the next token prediction
is_compact = self._is_compact_hidden_states(hidden_states, logits_metadata)
if is_compact:
logits = self._get_logits(hidden_states, lm_head, logits_metadata)
return LogitsProcessorOutput(
next_token_logits=logits,
hidden_states=None,
mm_input_embeds=logits_metadata.mm_input_embeds,
)
(
pruned_states,
pruned_states_before_norm,
@@ -337,13 +345,11 @@ class LogitsProcessor(nn.Module):
del hidden_states
if not logits_metadata.extend_return_logprob:
# Compute logits for both input and sampled tokens.
logits = self._get_logits(pruned_states, lm_head, logits_metadata)
sampled_logits = (
logits[sample_indices] if sample_indices is not None else logits
)
# Decode mode or extend mode without return_logprob.
return LogitsProcessorOutput(
next_token_logits=sampled_logits,
hidden_states=hidden_states_to_store,
@@ -396,6 +402,17 @@ class LogitsProcessor(nn.Module):
mm_input_embeds=logits_metadata.mm_input_embeds,
)
def _is_compact_hidden_states(
self, hidden_states: torch.Tensor, logits_metadata: LogitsMetadata
) -> bool:
if not logits_metadata.forward_mode.is_extend():
return False
if logits_metadata.extend_return_logprob:
return False
if logits_metadata.extend_seq_lens is None:
return False
return hidden_states.shape[0] == len(logits_metadata.extend_seq_lens)
def _get_pruned_states(
self,
hidden_states: torch.Tensor,
@@ -800,9 +817,9 @@ class LogitsProcessor(nn.Module):
# Restore the full-pruned lm_head batch_info after chunk iteration.
if hasattr(lm_head, "reset_lm_head_pass"):
assert hasattr(
lm_head, "set_lm_head_pass"
), "lm_head must have set_lm_head_pass method and reset_lm_head_pass method at the same time"
assert hasattr(lm_head, "set_lm_head_pass"), (
"lm_head must have set_lm_head_pass method and reset_lm_head_pass method at the same time"
)
lm_head.reset_lm_head_pass()
# Concatenate the results

94
python/sglang/srt/models/deepseek_v2.py
View File

@@ -57,6 +57,7 @@ from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
from sglang.srt.layers.attention.nsa.utils import (
can_cp_split,
cp_all_gather_rerange_output,
cp_collect_last_token_hidden,
cp_split_and_rebuild_data,
cp_split_and_rebuild_position,
is_nsa_enable_prefill_cp,
@@ -114,7 +115,11 @@ from sglang.srt.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
PPProxyTensors,
)
from sglang.srt.models.deepseek_common.attention_backend_handler import (
AttentionBackendRegistry,
)
@@ -222,8 +227,7 @@ class DeepseekV2MLP(nn.Module):
self.down_proj.weight = self.down_proj.weight_packed
if hidden_act != "silu":
raise ValueError(
f"Unsupported activation: {hidden_act}. "
"Only silu is supported for now."
f"Unsupported activation: {hidden_act}. Only silu is supported for now."
)
self.act_fn = SiluAndMul()
@@ -321,7 +325,6 @@ class MoEGate(nn.Module):
and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
and _device_sm >= 90
):
# router gemm output float32
logits = dsv3_router_gemm(
hidden_states, self.weight, out_dtype=torch.float32
@@ -341,7 +344,6 @@ class MoEGate(nn.Module):
class DeepseekV2MoE(nn.Module):
def __init__(
self,
config: PretrainedConfig,
@@ -458,13 +460,15 @@ class DeepseekV2MoE(nn.Module):
else {}
),
)
is_packed_weight = hasattr(
self.shared_experts.gate_up_proj.quant_method, "quant_config"
) and self.shared_experts.gate_up_proj.quant_method.quant_config.get_name() in {
"awq",
"awq_marlin",
"moe_wna16",
}
is_packed_weight = (
hasattr(self.shared_experts.gate_up_proj.quant_method, "quant_config")
and self.shared_experts.gate_up_proj.quant_method.quant_config.get_name()
in {
"awq",
"awq_marlin",
"moe_wna16",
}
)
self.shared_experts_is_int8 = (
not is_packed_weight
and self.shared_experts.gate_up_proj.weight.dtype == torch.int8
@@ -486,9 +490,7 @@ class DeepseekV2MoE(nn.Module):
self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size
== self.shared_experts.down_proj.quant_method.quant_config.weight_block_size
)
self.shared_experts_weight_block_size = (
self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size
)
self.shared_experts_weight_block_size = self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size
self.top_k = config.num_experts_per_tok
@@ -647,7 +649,6 @@ class DeepseekV2MoE(nn.Module):
def _pre_combine_hook(
dispatcher: BaseDispatcher, combine_input: CombineInput
):
nonlocal shared_output
self.alt_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(self.alt_stream):
@@ -837,7 +838,6 @@ class DeepseekV2MoE(nn.Module):
def _post_dispatch_hook(
dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
):
combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
compute_overlap_args(dispatch_output, self.alt_stream)
)
@@ -855,7 +855,6 @@ class DeepseekV2MoE(nn.Module):
def _pre_combine_hook(
dispatcher: BaseDispatcher, combine_input: CombineInput
):
nonlocal shared_output
if (
@@ -893,7 +892,6 @@ class DeepseekV2MoE(nn.Module):
def _post_dispatch_hook(
dispatcher: BaseDispatcher, dispatch_output: DispatchOutput
):
combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
compute_overlap_args(dispatch_output, self.alt_stream)
)
@@ -1076,7 +1074,6 @@ class DeepseekV2AttentionMLA(
DeepseekMLARocmForwardMixin,
DeepseekMLACpuForwardMixin,
):
def __init__(
self,
config: PretrainedConfig,
@@ -1357,18 +1354,18 @@ class DeepseekV2AttentionMLA(
not get_attn_tp_context().input_scattered
and hidden_states[0].shape[0] == 0
):
assert (
not self.o_proj.reduce_results
), "short-circuiting allreduce will lead to hangs"
assert not self.o_proj.reduce_results, (
"short-circuiting allreduce will lead to hangs"
)
return hidden_states[0]
else:
if (
not get_attn_tp_context().input_scattered
and hidden_states.shape[0] == 0
):
assert (
not self.o_proj.reduce_results
), "short-circuiting allreduce will lead to hangs"
assert not self.o_proj.reduce_results, (
"short-circuiting allreduce will lead to hangs"
)
return hidden_states, None, forward_batch, None
attn_forward_method = self.dispatch_attn_forward_method(forward_batch)
@@ -1499,7 +1496,6 @@ class DeepseekV2AttentionMLA(
class DeepseekV2DecoderLayer(nn.Module):
def __init__(
self,
config: PretrainedConfig,
@@ -2044,17 +2040,34 @@ class DeepseekV2Model(nn.Module):
hidden_states, _ = self.norm(hidden_states, residual)
if self.pp_group.is_last_rank and nsa_use_prefill_cp(forward_batch):
# allgather + rerrange
hidden_states = cp_all_gather_rerange_output(
hidden_states,
self.cp_size,
forward_batch,
torch.cuda.current_stream(),
)
if self._should_use_narrow_output_path(forward_batch):
hidden_states = cp_collect_last_token_hidden(
hidden_states, forward_batch, self.cp_size
)
else:
hidden_states = cp_all_gather_rerange_output(
hidden_states,
self.cp_size,
forward_batch,
torch.cuda.current_stream(),
)
if len(aux_hidden_states) == 0:
return hidden_states
return hidden_states, aux_hidden_states
def _should_use_narrow_output_path(self, forward_batch):
if not nsa_use_prefill_cp(forward_batch):
return False
if not self.pp_group.is_last_rank:
return False
if not forward_batch.forward_mode.is_extend():
return False
if forward_batch.return_logprob:
return False
if forward_batch.capture_hidden_mode != CaptureHiddenMode.NULL:
return False
return True
class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
# for quark model load
@@ -2208,6 +2221,19 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
else:
return hidden_states
def _should_use_narrow_output_path(self, forward_batch: ForwardBatch) -> bool:
if not nsa_use_prefill_cp(forward_batch):
return False
if not self.pp_group.is_last_rank:
return False
if not forward_batch.forward_mode.is_extend():
return False
if forward_batch.return_logprob:
return False
if forward_batch.capture_hidden_mode != CaptureHiddenMode.NULL:
return False
return True
@property
def start_layer(self):
return self.model.start_layer