diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 53cbbd1fd..a11100ed7 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/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, - 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 diff --git a/python/sglang/srt/layers/logprob_processor.py b/python/sglang/srt/layers/logprob_processor.py index d78651707..b23723e9a 100644 --- a/python/sglang/srt/layers/logprob_processor.py +++ b/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,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" diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index 613342a47..d1bb04ad1 100644 --- a/python/sglang/srt/lora/layers.py +++ b/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]. """ diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py index 619c65df2..eeccd2d1a 100644 --- a/python/sglang/srt/lora/utils.py +++ b/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. diff --git a/test/registered/unit/layers/test_logprob_chunk_stitching.py b/test/registered/unit/layers/test_logprob_chunk_stitching.py index 21af55cdc..48489fdc5 100644 --- a/test/registered/unit/layers/test_logprob_chunk_stitching.py +++ b/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,