[Refactor] Unify input logprob processing on a single chunked path (#31655)

python/sglang/srt/layers/logits_processor.py

@@ -37,8 +37,9 @@ from sglang.srt.layers.dp_attention import (
 )
 from sglang.srt.layers.logprob_processor import (
     InputLogprobProcessor,
-    get_token_ids_logprobs_prefill,
+    LogprobStage,
-    get_top_logprobs_prefill,
+    get_token_ids_logprobs_raw,
+    get_top_logprobs_raw,
 )
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.model_executor.forward_batch_info import (
@@ -586,8 +587,6 @@ class LogitsProcessor(nn.Module):
                 )
                 input_logprob_indices_pt += extend_len - start_len
 
-            # 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_list)
             if hidden_states_before_norm is not None:
                 pruned_states_before_norm = torch.cat(pruned_states_before_norm_list)
@@ -921,8 +920,12 @@ class LogitsProcessor(nn.Module):
             (
                 input_token_ids_logprobs_val,
                 input_token_ids_logprobs_idx,
-            ) = get_token_ids_logprobs_prefill(
+            ) = get_token_ids_logprobs_raw(
-                sliced_logprobs, logits_metadata, no_copy_to_cpu=True
+                sliced_logprobs,
+                logits_metadata.token_ids_logprobs,
+                stage=LogprobStage.PREFILL,
+                extend_logprob_pruned_lens_cpu=logits_metadata.extend_logprob_pruned_lens_cpu,
+                no_copy_to_cpu=True,
             )
 
         # Get the logprob of top-k tokens
@@ -930,7 +933,12 @@ class LogitsProcessor(nn.Module):
             (
                 input_top_logprobs_val,
                 input_top_logprobs_idx,
-            ) = get_top_logprobs_prefill(sliced_logprobs, logits_metadata)
+            ) = get_top_logprobs_raw(
+                sliced_logprobs,
+                logits_metadata.top_logprobs_nums,
+                stage=LogprobStage.PREFILL,
+                extend_logprob_pruned_lens_cpu=logits_metadata.extend_logprob_pruned_lens_cpu,
+            )
 
         # MIS scores come from input_token_ids_logprobs_val (label-token logprobs),
         # not from per-position input_token_logprobs. However, the shared logprob

python/sglang/srt/layers/logprob_processor.py

@@ -66,17 +66,6 @@ def get_top_logprobs_raw(
     return top_logprobs_val, top_logprobs_idx
 
 
-def get_top_logprobs_prefill(
-    all_logprobs: torch.Tensor, logits_metadata: LogitsMetadata
-):
-    return get_top_logprobs_raw(
-        all_logprobs,
-        logits_metadata.top_logprobs_nums,
-        stage=LogprobStage.PREFILL,
-        extend_logprob_pruned_lens_cpu=logits_metadata.extend_logprob_pruned_lens_cpu,
-    )
-
-
 def get_top_logprobs(
     logprobs: torch.Tensor,
     top_logprobs_nums: List[int],
@@ -135,18 +124,6 @@ def get_token_ids_logprobs_raw(
     return vals, idxs
 
 
-def get_token_ids_logprobs_prefill(
-    all_logprobs, logits_metadata: LogitsMetadata, no_copy_to_cpu=False
-):
-    return get_token_ids_logprobs_raw(
-        all_logprobs,
-        logits_metadata.token_ids_logprobs,
-        stage=LogprobStage.PREFILL,
-        extend_logprob_pruned_lens_cpu=logits_metadata.extend_logprob_pruned_lens_cpu,
-        no_copy_to_cpu=no_copy_to_cpu,
-    )
-
-
 def get_token_ids_logprobs(logprobs, token_ids_logprobs, no_copy_to_cpu=False):
     return get_token_ids_logprobs_raw(
         logprobs,
@@ -387,30 +364,18 @@ class InputLogprobProcessor:
         logits_metadata: LogitsMetadata,
         skip_chunking_for_dp_attn: bool = False,
     ) -> Tuple[InputLogprobsResult, torch.Tensor]:
-        # Start to process input logprobs
+        # Non-chunked = one chunk covering every row. DP-attention must stay
-        # Determine whether to use chunked or non-chunked logits processing.
+        # single-chunk: the collective schedule cannot depend on per-rank rows.
-        # Skip chunking if:
+        if (
-        # 1. Chunking is disabled
-        # 2. Total count is below chunk size threshold
-        # 3. DP attention all-gather is enabled (can use "enable_dp_lm_head" to enable chunking)
-        should_skip_chunking = (
             not self.enable_logprobs_chunk
             or pruned_states.shape[0] <= self.logprobs_chunk_size
             or skip_chunking_for_dp_attn
-        )
+        ):
-
+            chunk_size = max(pruned_states.shape[0], 1)
-        if should_skip_chunking:
-            # Compute logits for both input and sampled tokens.
-            logits = get_logits_fn(pruned_states, lm_head, logits_metadata)
-            sampled_logits = (
-                logits[sample_indices] if sample_indices is not None else logits
-            )
-            input_logits = logits[input_logprob_indices]
-            del logits
-
-            logprobs_result = self.process_input_logprobs(input_logits, logits_metadata)
         else:
-            logprobs_result, sampled_logits = self.process_input_logprobs_by_chunk(
+            chunk_size = self.logprobs_chunk_size
+
+        return self._forward_by_chunk(
             pruned_states,
             sample_indices,
             input_logprob_indices,
@@ -418,45 +383,10 @@ class InputLogprobProcessor:
             lm_head,
             get_logits_fn,
             logits_metadata,
+            chunk_size,
         )
 
-        return logprobs_result, sampled_logits
+    def _forward_by_chunk(
-
-    def process_input_logprobs(self, input_logits, logits_metadata: LogitsMetadata):
-        input_logprobs = torch.nn.functional.log_softmax(input_logits, dim=-1)
-
-        # Get the logprob of top-k tokens
-        if logits_metadata.extend_return_top_logprob:
-            (
-                input_top_logprobs_val,
-                input_top_logprobs_idx,
-            ) = get_top_logprobs_prefill(input_logprobs, logits_metadata)
-        else:
-            input_top_logprobs_val = input_top_logprobs_idx = None
-
-        # Get the logprob of given token id
-        if logits_metadata.extend_token_ids_logprob:
-            (
-                input_token_ids_logprobs_val,
-                input_token_ids_logprobs_idx,
-            ) = get_token_ids_logprobs_prefill(input_logprobs, logits_metadata)
-        else:
-            input_token_ids_logprobs_val = input_token_ids_logprobs_idx = None
-
-        input_token_logprobs = input_logprobs[
-            torch.arange(input_logprobs.shape[0], device=input_logprobs.device),
-            logits_metadata.extend_input_logprob_token_ids_gpu,
-        ]
-
-        return InputLogprobsResult(
-            input_token_logprobs=input_token_logprobs,
-            input_top_logprobs_val=input_top_logprobs_val,
-            input_top_logprobs_idx=input_top_logprobs_idx,
-            input_token_ids_logprobs_val=input_token_ids_logprobs_val,
-            input_token_ids_logprobs_idx=input_token_ids_logprobs_idx,
-        )
-
-    def process_input_logprobs_by_chunk(
         self,
         pruned_states: torch.Tensor,
         sample_indices: torch.Tensor,
@@ -465,19 +395,9 @@ class InputLogprobProcessor:
         lm_head: VocabParallelEmbedding,
         get_logits_fn: Callable,
         logits_metadata: LogitsMetadata,
+        chunk_size: int,
     ) -> Tuple[InputLogprobsResult, torch.Tensor]:
-        """
+        """Compute input logprobs chunk by chunk to cap peak memory."""
-        compute logprobs for the output token from the hidden states.
-        To avoid using too much memory, we split pruned_states into chunks of
-        rows to compute input_logprobs separately, then concatenate the results.
-
-        Returns:
-            InputLogprobsResult: logprobs result
-            torch.Tensor: sampled logits
-        """
-
-        # The peak memory usage is proportional to the chunk size.
-        chunk_size = self.logprobs_chunk_size
         total_size = pruned_states.shape[0]
         num_chunks = (total_size + chunk_size - 1) // chunk_size
 
@@ -507,7 +427,7 @@ class InputLogprobProcessor:
             # Notify lm_head LoRA about the current chunk so it can swap
             # to the precomputed per-chunk batch_info.  This is a no-op
             # for non-LoRA lm_head modules.
-            if hasattr(lm_head, "set_lm_head_pass"):
+            if num_chunks > 1 and hasattr(lm_head, "set_lm_head_pass"):
                 lm_head.set_lm_head_pass(i)
 
             # Get indices for this chunk
@@ -525,7 +445,10 @@ class InputLogprobProcessor:
             # chunks whose shape happens to match the buffer.
             chunk_states = pruned_states[start_idx:end_idx]
             chunk_logits = get_logits_fn(
-                chunk_states, lm_head, logits_metadata, use_logits_buffer=False
+                chunk_states,
+                lm_head,
+                logits_metadata,
+                use_logits_buffer=num_chunks == 1,
             )
 
             # Initialize sampled_logits on first chunk
@@ -546,8 +469,11 @@ class InputLogprobProcessor:
                 sampled_logits[chunk_sample_mask] = chunk_logits[chunk_sample_indices]
 
             # Zero-logprob-row chunks still need the per-sequence bookkeeping below.
-            # Compute the logprobs of the chunk
+            # Compute the logprobs of the chunk. Free the raw logits before the
+            # out-of-place log_softmax: keeping all three alive is a 3x peak,
+            # which OOMs when the single chunk covers a large batch.
             chunk_input_logprobs = chunk_logits[chunk_indices]
+            del chunk_logits
             chunk_input_logprobs = torch.nn.functional.log_softmax(
                 chunk_input_logprobs, dim=-1
             )
@@ -597,9 +523,11 @@ class InputLogprobProcessor:
                 logits_metadata.extend_input_logprob_token_ids_gpu[mask_indices],
             ]
             input_token_logprobs.append(chunk_input_token_logprobs)
+            # Free before the next chunk's logits (bf16 + fp32) materialize.
+            del chunk_input_logprobs
 
         # Restore the full-pruned lm_head batch_info after chunk iteration.
-        if hasattr(lm_head, "reset_lm_head_pass"):
+        if num_chunks > 1 and 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"

python/sglang/srt/lora/layers.py

@@ -390,7 +390,7 @@ class ParallelLMHeadWithLoRA(BaseLayerWithLoRA):
     def set_lm_head_pass(self, pass_idx: int):
         """Set the active lm_head pass index before a logprobs chunk.
 
-        Called by InputLogprobProcessor.process_input_logprobs_by_chunk() before
+        Called by InputLogprobProcessor._forward_by_chunk() before
         each chunk's _get_logits call.  _get_lm_head_batch_info() will
         resolve to lm_head_pass_batch_infos[pass_idx].
         """

python/sglang/srt/lora/utils.py

@@ -545,7 +545,7 @@ def build_lm_head_pass_segments(
     Precompute per-pass segment info for lm_head LoRA logprobs processing.
 
     When InputLogprobProcessor uses chunked logprobs processing
-    (process_input_logprobs_by_chunk), pruned hidden states are split into
+    (_forward_by_chunk), pruned hidden states are split into
     fixed-size passes.  Each pass needs its own segmentation
     (weight_indices, seg_lens) so that lm_head LoRA operates on the
     correct adapter assignments per pass.

test/registered/unit/layers/test_logprob_chunk_stitching.py

@@ -46,7 +46,6 @@ def _build_batch(seq_specs, with_token_ids):
         input_logprob_indices.extend([lp_pt + i for i in range(n_lp)])
         lp_pt += rows
         pruned_lens.append(n_lp)
-    token_to_seq_idx.append(len(seq_specs) - 1)
     metadata = SimpleNamespace(
         extend_return_top_logprob=True,
         extend_token_ids_logprob=with_token_ids,