[Feature] support `gather` instead of `all_gather` when gathering the logits
#3,365 opened on Feb 7, 2025
Repository metrics
- Stars
- (28,442 stars)
- PR merge metrics
- (Avg merge 2d 1h) (1,000 merged PRs in 30d)
Description
Checklist
- 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sglang/discussions/new/choose Otherwise, it will be closed.
- 2. Please use English, otherwise it will be closed.
Motivation
We noticed that in the _get_logits function of vllm, gather instead of all_gather will be used under certain conditions (the main condition is that for non-tpu devices):
Code link:
The change from using all_gather to gather is initially added in this PR for your reference: https://github.com/vllm-project/vllm/pull/2221.
While in SGLang, we see currently all_gather is always used:
https://github.com/sgl-project/sglang/blob/e868d0b60eb2d435c5599165f787bca06bdc9c3d/python/sglang/srt/layers/logits_processor.py#L246
Does SGLang have the plan to add gather instead of only all_gather when gathering the logits? Per the practice in vllm, using gather seems to have better performance than all_gather on devices which have gather support.
Related resources
No response