diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 1e097983d..6ebc78114 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -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] diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 662cc2519..f33d97950 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 0d88f541a..a4d00ae00 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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