minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Liangsheng Yin
2026-07-19 16:03:23 -07:00
committed by GitHub
parent 1a317839d7
commit b3570a4531
5 changed files with 49 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/logits_processor.py
+15 -7
View File
@@ -37,8 +37,9 @@ from sglang.srt.layers.dp_attention import (
)
from sglang.srt.layers.logprob_processor import (
InputLogprobProcessor,
get_token_ids_logprobs_prefill,
get_top_logprobs_prefill,
LogprobStage,
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(
sliced_logprobs, logits_metadata, no_copy_to_cpu=True
) = get_token_ids_logprobs_raw(
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
+32 -104
View File
@@ -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,76 +364,29 @@ class InputLogprobProcessor:
logits_metadata: LogitsMetadata,
skip_chunking_for_dp_attn: bool = False,
) -> Tuple[InputLogprobsResult, torch.Tensor]:
# Start to process input logprobs
# Determine whether to use chunked or non-chunked logits processing.
# Skip chunking 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 = (
# Non-chunked = one chunk covering every row. DP-attention must stay
# single-chunk: the collective schedule cannot depend on per-rank rows.
if (
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)
else:
chunk_size = self.logprobs_chunk_size
return self._forward_by_chunk(
pruned_states,
sample_indices,
input_logprob_indices,
token_to_seq_idx,
lm_head,
get_logits_fn,
logits_metadata,
chunk_size,
)
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(
pruned_states,
sample_indices,
input_logprob_indices,
token_to_seq_idx,
lm_head,
get_logits_fn,
logits_metadata,
)
return logprobs_result, sampled_logits
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(
def _forward_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 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
"""Compute input logprobs chunk by chunk to cap peak memory."""
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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
-1
View File
@@ -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,