From b9ebf0ed631b1c3d81c358f8298566d7eea95038 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 20 Dec 2025 09:38:46 +0800 Subject: [PATCH] Clean `hidden_states_before_norm` (#15485) --- python/sglang/srt/layers/logits_processor.py | 44 +++++++++---------- python/sglang/srt/managers/schedule_batch.py | 4 -- python/sglang/srt/models/mimo_v2_flash.py | 7 ++- .../sglang/srt/models/mimo_v2_flash_nextn.py | 6 ++- 4 files changed, 29 insertions(+), 32 deletions(-) diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 48e735e6d..9426c108f 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -144,8 +144,6 @@ class LogitsMetadata: # Whether this batch is prefill-only (no token generation needed) is_prefill_only: bool = False - return_hidden_states_before_norm: bool = False - @classmethod def from_forward_batch(cls, forward_batch: ForwardBatch): if ( @@ -196,7 +194,6 @@ class LogitsMetadata: global_num_tokens_for_logprob_cpu=forward_batch.global_num_tokens_for_logprob_cpu, global_num_tokens_for_logprob_gpu=forward_batch.global_num_tokens_for_logprob_gpu, dp_padding_mode=DpPaddingMode.SUM_LEN, - return_hidden_states_before_norm=forward_batch.return_hidden_states_before_norm, ) def compute_dp_attention_metadata(self): @@ -405,6 +402,7 @@ class LogitsProcessor(nn.Module): ) # Get the last hidden states and last logits for the next token prediction + pruned_states_before_norm: Optional[torch.Tensor] = None if ( logits_metadata.forward_mode.is_decode_or_idle() or logits_metadata.forward_mode.is_target_verify() @@ -416,6 +414,7 @@ class LogitsProcessor(nn.Module): aux_pruned_states = [hidden for hidden in aux_hidden_states] sample_indices = None input_logprob_indices = None + elif ( logits_metadata.forward_mode.is_extend() and not logits_metadata.extend_return_logprob @@ -437,11 +436,8 @@ class LogitsProcessor(nn.Module): - 1 ) pruned_states = hidden_states[last_index] - pruned_states_before_norm = ( - hidden_states_before_norm[last_index] - if hidden_states_before_norm is not None - else None - ) + if hidden_states_before_norm is not None: + pruned_states_before_norm = hidden_states_before_norm[last_index] if aux_hidden_states is not None: aux_pruned_states = [hidden[last_index] for hidden in aux_hidden_states] sample_indices = None @@ -474,7 +470,7 @@ class LogitsProcessor(nn.Module): sample_indices = [] input_logprob_indices_pt = 0 input_logprob_indices = [] - pt, pruned_states, pruned_states_before_norm = 0, [], [] + pt, pruned_states_list, pruned_states_before_norm_list = 0, [], [] token_to_seq_idx = [] for idx, (extend_logprob_start_len, extend_len) in enumerate( @@ -493,9 +489,11 @@ class LogitsProcessor(nn.Module): # We always need at least 1 token to sample because that's required # by a caller. assert extend_len > start_len - pruned_states.append(hidden_states[pt + start_len : pt + extend_len]) + pruned_states_list.append( + hidden_states[pt + start_len : pt + extend_len] + ) if hidden_states_before_norm is not None: - pruned_states_before_norm.append( + pruned_states_before_norm_list.append( hidden_states_before_norm[pt + start_len : pt + extend_len] ) # Map each token to its sequence index, for chunked computation @@ -514,11 +512,9 @@ class LogitsProcessor(nn.Module): # Set the last token of the last sequence token_to_seq_idx.append(len(logits_metadata.extend_seq_lens_cpu) - 1) - pruned_states = torch.cat(pruned_states) + pruned_states = torch.cat(pruned_states_list) if hidden_states_before_norm is not None: - pruned_states_before_norm = torch.cat(pruned_states_before_norm) - else: - pruned_states_before_norm = None + pruned_states_before_norm = torch.cat(pruned_states_before_norm_list) sample_indices = torch.tensor( sample_indices, device=pruned_states.device, dtype=torch.int64 ) @@ -558,20 +554,20 @@ class LogitsProcessor(nn.Module): if sample_indices is not None else pruned_states ) - hidden_states_to_store_before_norm = ( - pruned_states_before_norm[sample_indices] - if sample_indices is not None - else pruned_states_before_norm - ) + if hidden_states_before_norm is not None: + hidden_states_to_store_before_norm = ( + pruned_states_before_norm[sample_indices] + if sample_indices is not None + else pruned_states_before_norm + ) else: assert False, "Should never reach" del hidden_states - if ( - logits_metadata.return_hidden_states_before_norm - and hidden_states_to_store_before_norm is not None - ): + if hidden_states_to_store_before_norm is not None: + # NOTE: when hidden_states_before_norm is provided, we always + # prefer to return it. hidden_states_to_store = hidden_states_to_store_before_norm if not logits_metadata.extend_return_logprob: diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 0140d62d1..748f11ff1 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1217,9 +1217,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # Diffusion LLM dllm_config: Optional[DllmConfig] = None - # For hidden states before normal - return_hidden_states_before_norm: bool = False - @classmethod def init_new( cls, @@ -2104,7 +2101,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): dllm_config=self.dllm_config, reqs=self.reqs, has_grammar=self.has_grammar, - return_hidden_states_before_norm=self.return_hidden_states_before_norm, mamba_track_indices=self.mamba_track_indices, mamba_track_mask=self.mamba_track_mask, mamba_track_seqlens=self.mamba_track_seqlens, diff --git a/python/sglang/srt/models/mimo_v2_flash.py b/python/sglang/srt/models/mimo_v2_flash.py index d2563ddc1..da352fd47 100644 --- a/python/sglang/srt/models/mimo_v2_flash.py +++ b/python/sglang/srt/models/mimo_v2_flash.py @@ -662,12 +662,15 @@ class MiMoV2Model(nn.Module): ) else: if hidden_states.shape[0] > 0: + if forward_batch.return_hidden_states_before_norm: + hidden_states_before_norm = ( + hidden_states if residual is None else hidden_states + residual + ) if residual is None: - hidden_states_before_norm = hidden_states hidden_states = self.norm(hidden_states) else: - hidden_states_before_norm = hidden_states + residual hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states, hidden_states_before_norm # If this function is called, it should always initialize KV cache scale diff --git a/python/sglang/srt/models/mimo_v2_flash_nextn.py b/python/sglang/srt/models/mimo_v2_flash_nextn.py index 1ce9ce47e..6440c3ef1 100644 --- a/python/sglang/srt/models/mimo_v2_flash_nextn.py +++ b/python/sglang/srt/models/mimo_v2_flash_nextn.py @@ -211,11 +211,13 @@ class MiMoV2ModelNextN(nn.Module): ) hidden_states_before_norm = None if not forward_batch.forward_mode.is_idle(): + if forward_batch.return_hidden_states_before_norm: + hidden_states_before_norm = ( + hidden_states if residual is None else hidden_states + residual + ) if residual is not None: - hidden_states_before_norm = hidden_states + residual hidden_states, _ = self.final_layernorm(hidden_states, residual) else: - hidden_states_before_norm = hidden_states hidden_states = self.final_layernorm(hidden_states) return hidden_states, hidden_states_before_norm