Slightly improve the sampler to skip unnecessary steps

[merrymercy]

Allow to skip topk and topp if not needed

Co-authored-by: Sehoon Kim <[email protected]>

[gemini-code-assist]

Hello @merrymercy, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

Hello team, Gemini here with a summary of this pull request. The main goal of this PR, as indicated by the title, is to improve the efficiency of the sampler by skipping unnecessary computation steps. Specifically, it introduces logic to bypass top-k, top-p, and min-p filtering when none of these sampling methods are required for any request in the batch. This should lead to performance gains in scenarios where simpler sampling (like greedy or temperature sampling without filtering) is used.

In addition to the core optimization, the PR also refactors the custom logit processor logic, moving it from a class method to a module-level function, potentially to support more flexible use cases like speculative decoding. It also includes a minor fix related to grammar termination checks during vocabulary mask updates.

Highlights

• Sampler Optimization: Introduced a check to determine if top-k, top-p, or min-p sampling is needed for the current batch. If none are required, the sampler now directly samples from the probabilities using a new dedicated function, skipping the filtering steps.
• Logprob Calculation Adjustment: The calculation of log probabilities has been moved and updated. It now correctly uses the normalized probabilities after top-p filtering (if applied) to ensure accurate logprob values are returned.
• Custom Logit Processor Refactoring: The _apply_custom_logit_processor method within the Sampler class has been removed and its logic extracted into a new module-level function apply_custom_logit_processor. This new function also includes a parameter num_tokens_in_batch for potential future support of multi-token processing (e.g., speculative decoding).
• Sampling Batch Info Updates: The SamplingBatchInfo class now tracks whether any request in the batch specifically requires top-p (need_top_p_sampling) or top-k (need_top_k_sampling) sampling. These flags are used by the sampler to decide whether filtering is necessary.
• Minor Grammar Fix: Added a check (not grammar.is_terminated()) in the update_regex_vocab_mask method to avoid processing grammars that have already terminated.

Changelog

Click here to see the changelog

• python/sglang/srt/layers/sampler.py
  • Replaced get_tensor_model_parallel_group with get_tp_group (lines 8, 33).
  • Changed call from self._apply_custom_logit_processor to apply_custom_logit_processor (line 62).
  • Added logic to check if top-k, top-p, or min-p sampling is needed (can_sample_directly_from_probs) (lines 84-88).
  • Introduced a direct sampling path using sampling_from_probs_torch when no filtering is needed (lines 90-92).
  • Adjusted logprob calculation block to occur after sampling and use top_p_normalize_probs_torch if top-p was needed (lines 125-139).
  • Removed the _apply_custom_logit_processor method from the Sampler class (lines 163-194 removed).
  • Added a new module-level function sampling_from_probs_torch for direct sampling (lines 204-209).
  • Added a new module-level function apply_custom_logit_processor with num_tokens_in_batch parameter (lines 252-290).
• python/sglang/srt/sampling/sampling_batch_info.py
  • Imported TOP_K_ALL (line 12).
  • Added need_top_p_sampling and need_top_k_sampling boolean fields to SamplingBatchInfo (lines 32-36).
  • Populated need_top_p_sampling and need_top_k_sampling in from_schedule_batch (lines 144-145).
  • Added and not grammar.is_terminated() check in update_regex_vocab_mask (line 180).
  • Added merging logic for need_top_p_sampling and need_top_k_sampling in merge_batch (lines 321-322).

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a valuable optimization to the sampling process by skipping unnecessary top-k and top-p computations when they are not active. The changes are well-targeted and improve efficiency.

The refactoring of apply_custom_logit_processor into a standalone function enhances modularity and is a good improvement. Additionally, the updates to SamplingBatchInfo to support these optimizations are clear and correctly implemented.

One point for discussion is a subtle change in the NaN handling logic within the flashinfer backend path, which might alter when the process crashes due to NaNs. Please see the specific comment for details.

Overall, the PR is a positive step towards a more efficient sampler.

Summary of Findings

• Sampler Optimization: The core change effectively skips top-k/top-p sampling steps when need_top_p_sampling, need_top_k_sampling, and need_min_p_sampling are all false, which should improve performance in applicable scenarios.
• Logit Processor Refactoring: The apply_custom_logit_processor method was refactored into a standalone function, improving code modularity. It also now supports num_tokens_in_batch for speculative decoding.
• NaN Handling Behavior Change: There's a potential change in NaN handling behavior for the flashinfer backend. The condition for crashing on NaN is now solely dependent on self.use_nan_detection, whereas it previously also depended on crash_on_warnings(). This needs clarification.
• Improved Constrained Decoding Logic: In SamplingBatchInfo, the condition for applying regex vocab masks now also checks not grammar.is_terminated(), which is a good refinement for correctness.

Merge Readiness

The pull request introduces a good optimization and refactoring. However, there is one medium-severity issue regarding a potential change in NaN handling behavior that should be clarified or addressed before merging. Once this point is resolved, the PR should be in good shape for merging. As an AI, I am not authorized to approve pull requests; please ensure further review and approval from team members.