Clean hidden_states_before_norm (#15485)

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 = (
+            if hidden_states_before_norm is not None:
-                hidden_states_before_norm[last_index]
+                pruned_states_before_norm = hidden_states_before_norm[last_index]
-                if hidden_states_before_norm is not None
-                else None
-            )
             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)
+                pruned_states_before_norm = torch.cat(pruned_states_before_norm_list)
-            else:
-                pruned_states_before_norm = None
             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 = (
+                    if hidden_states_before_norm is not None:
-                        pruned_states_before_norm[sample_indices]
+                        hidden_states_to_store_before_norm = (
-                        if sample_indices is not None
+                            pruned_states_before_norm[sample_indices]
-                        else pruned_states_before_norm
+                            if sample_indices is not None
-                    )
+                            else pruned_states_before_norm
+                        )
             else:
                 assert False, "Should never reach"
 
         del hidden_states
 
-        if (
+        if hidden_states_to_store_before_norm is not None:
-            logits_metadata.return_hidden_states_before_norm
+            # NOTE: when hidden_states_before_norm is provided, we always
-            and hidden_states_to_store_before_norm is not None
+            # prefer to return it.
-        ):
             hidden_states_to_store = hidden_states_to_store_before_norm
 
         if not logits_metadata.extend_return_logprob:

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,

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

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