[Feature] support `gather` instead of `all_gather` when gathering the logits

[chunyuan-w]

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:

• logits = tensor_model_parallel_gather(logits)

• condition of whether using all_gather or gather

The change from using all_gather to gather is initially added in this PR for your reference: vllm-project/vllm#2221.

While in SGLang, we see currently all_gather is always used:

sglang/python/sglang/srt/layers/logits_processor.py

Line 246 in e868d0b

logits = tensor_model_parallel_all_gather(logits)

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

[zhaochenyang20]

Good idea. We should have someone working on this.

[voidxb]

I wanna own this issue @zhaochenyang20

[zhaochenyang20]

@voidxb Sure. Could have a try.

[voidxb]

https://github.com/vllm-project/vllm/pull/12071/files in this vllm pr, I see the use_all_gather var but firstly, I see parallel_config = get_current_vllm_config().parallel_config self.use_all_gather = current_platform.is_tpu() \ or envs.VLLM_USE_V1 \ or parallel_config.distributed_executor_backend == "external_launcher" # noqa so it seems that in new v1 arch, vllm uses all_gather firstly.

then from vllm/distributed/device_communicators/base_device_communicator.py and vllm/model_executor/layers/logits_processor.py, I found logits could be return None but in model_runner , logits just is delivered to sampler and sampler.forward does not allow logits could be None(assert). @chunyuan-w so do you really use this feature in production? in vllm and it really works well?

I do a version to implement it, but its output seems bad(without token broadcast, just like vllm), and in my opinion, I think it might should change more code to let workers(tp_rank >0) could handle None logits.

and as I could see, if we don't use all_gather in get_logits, tprank > 0 worker should get token from rank 0 worker as well, so we should broadcast token from rank0 to other ranks.In all_gather, other rank get token from logits(after sampling), but in gather method, we have to add a new path to get token.
so Is It really an opt?
the fix (use gather instead of all_gather ) in this pr is not the core fetaure and was merged in Jan this year, so I have some doubts.
Or could a logits gather & a token broadcast obviously decrease communication cost(compare to a logits all gather)?

thanks

[zhaochenyang20]

cc @FrankLeeeee

[chunyuan-w]

Hi @voidxb thanks for looking into this.

Yes we do use gather and it works well.
VLLM_USE_V1 is still turned off by default so we're not running into V1 currently.
When using gather, only the driver worker (rank 0 of tensor parallel group) will perform sampling, the other workers directly return [] here and won't call the self.model.sample.

[voidxb]

ok, I have see the difference
in vllm, before forward, it broadcasts model input to other ranks
and only rank0(driver worker) do sampling
so in that case, all_gather logits is unnecessary(token broadcast + logits all gather -> token broadcast + logits gather)
but in sglang, we have no token broadcast step, so all gather is necessary now
I will try to add token broadcast step

[zhaochenyang20]

@voidxb Thankss so muchc!