[Feature] support `gather` instead of `all_gather` when gathering the logits
#3365 aperta il 7 feb 2025
Metriche repository
- Star
- (28.442 star)
- Metriche merge PR
- (Merge medio 2g 1h) (1000 PR mergiate in 30 g)
Descrizione
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