[Bug] Performance regression on flashinfer.comm.trtllm_allreduce_fusion in 0.6.1

[Fridge003]

Ref: sgl-project/sglang#17237

In 0.5.3 it was correct, but after upgrading to 0.6.1 flashinfer.comm.trtllm_allreduce_fusion becomes really slow.

In sglang it's called here (with use_oneshot=None):
https://github.com/sgl-project/sglang/blob/17349168bb159a4e68ec7ac071aac4b8f67e5c60/python/sglang/srt/layers/flashinfer_comm_fusion.py#L194

It can be easily reproduced on small batch sizes (e.g. num_tokens=32, hidden_dim=7168) num_tokens=1 under checking

[noqcks]

hi! Im just learning CUDA and inference optimizations, so I apologize if my analysis is incorrect.

I believe there was a change in the function allreduce_sum between version 0.5.3 [ref] and 0.6.1 [ref]

there's a note about it

Optimization: process one element at a time to reduce register usage
Instead of storing acc_f32[VEC_SIZE] (32 bytes), process and convert immediately

so instead of processing all elements in parallel like in 0.5.3, we're computing one at a time and finishing it before moving onto the next. It reduces register usage from 32 bytes to 4 bytes, but also reduces parallelism.

[Fridge003]

Thanks @noqcks, does it mean the usage of this kernel should be changed? Like we need to pass in some different arguments?

[aleozlx]

@nv-yunzheq tried some repro and may have some questions as well.

[nv-yunzheq]

@Fridge003 I think we observe some regression when num_tokens is larger or equal to 32. But not when num_token euqals to 1. Do you know if the regression you observed only happens when num_tokens is small?
Which GPU are you testing? I was trying out on a B200 system.
Moreover, it seems the API will be deprecated. Could you switch to comm.allreduce_fusion API and enable API logging, so we could know the exact input parameters for this issue?
Thanks

[nv-yunzheq]

@Fridge003 I have a potentila fix in #2492. Are you able to take it and give it a try to see if it restore the performance?