adapt custom allreduce for tensorrt llm

[yizhang2077]

Motivation

adapt for tensorrt llm custom allreduce, currently still use vllm distributed.
After this pr is merged and sgl-kernel is stable, we only need replace vllm.distribued to sglang.srt.distributed, and add a monkey patch, then we can remove vllm distributed

Modifications

Checklist

• Format your code according to the Contributor Guide.
• Add unit tests as outlined in the Contributor Guide.
• Update documentation as needed, including docstrings or example tutorials.

[zhyncs]

@yizhang2077 Could you paste the unit test result

[yizhang2077]

@yizhang2077 Could you paste the unit test result

there are only correctness test here, do we need compare with vllm?

[zhyncs]

^ gentle ping cc @merrymercy

[zhyncs]

@yizhang2077 Could you paste the unit test result

there are only correctness test here, do we need compare with vllm?

ref #2481 (comment)

[zhyncs]

The most important case we care about is TP=8 and bs in [1, 1024]. The size is about 0 - 32MB. Can we do a more comprehensive test?

[merrymercy]

I think we can merge this one as it has correctness test. We can benchmark the performance part in future PRs.

Condition for switching to this

• Faster than or equal to vllm's custom allreduce on all cases (TP=2,4,8) x (bs=1,2,4,8, .. 128, 1024)
• Does not break AMD support

[yizhang2077]

performance part ref #2904, ci may fail temperarily until new sgl-kernel version release