[Feature]: Enhance AllToAll performance

[taylor-yb-lee]

🚀 The feature, motivation and pitch

(Part of the optimization effort for DeepSeek-R1 #12784)

In AutoDeploy backend for DeepSeek-R1 model, mla kernel (SwapsMmaAb, 1 CTA kernel) is 2x slower than that is used for PT(KeepsMmaAb, 2 CTA kernel).
But we cannot use 2CTA kernel in AutoDeploy, because

• To use 2 CTA kernels, we need attention_dp as PT does (see more detail in https://gitlab-master.nvidia.com/taylor/cc_reports/-/blob/main/deepseek_r1/compare_dsr1_pt_ad_perf_0421.md?ref_type=heads#the-fundamental-architectural-constraint)
• To use attention_dp, the MoE layer should use AllToAll (because tokens are distributed across the GPUs)

However, here the AllToAll performance in AutoDeploy is extremely slow.
So we need to have a performant AllToAll as PT.

How to reproduce:
working branch : https://github.com/nv-auto-deploy/TensorRT-LLM/tree/taylor/deepgemm_eran_mla

Add and change examples/auto_deploy/model_registry/configs/deepseek-r1.yaml as below:

max_num_tokens: 8192
detect_sharding:
   enable_attention_dp: true

Rum cmd:

bench-sweep --model deepseek-ai/DeepSeek-R1-0528 --config-path deepseek-r1.yaml--server-type trtllm-autodeploy --isl 1000 --osl 1000 --concurrencies 256 --world-size 8

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

[greg-kwasniewski1]

#13723