Add fused top-K softmax kernel for MoE

[WoosukKwon]

This PR ports a fused topk-softmax kernel from TensorRT-LLM v0.7.1.

TODO:

1. Port more MoE-related kernels
2. Use CUTLASS-based grouped GEMM kernels with appropriate tuning (if they perform better than the current Triton kernel).

[WoosukKwon]

@Yard1 @cadedaniel @pcmoritz Can any one of you review the PR?

[pcmoritz]

Yes, happy to review! Thanks a lot for writing this :)

[WoosukKwon]

@pcmoritz Thanks!

[pcmoritz]

Btw, I did a little bit of benchmarking on this PR and without touching any of the system parameters in the PR I'm already seeing a 1.5% - 3.5% end-to-end latency improvement. It is higher in the low latency regime. Concretely I tested on TP2 on H100 with 1000 input and 50 output tokens on Mixtral. So it seems worth merging this even though the low-level kernel code is not easy to follow -- most people can probably just treat it as a black box so it shouldn't have a big impact on maintainability.

[pcmoritz]

I will spend some more time trying to understand the implementation in topk_softmax_kernels.cu but no need to block on that since that's mostly the upstream code from https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu and in any case we should probably keep it close to that and not change it :)

[WoosukKwon]

@pcmoritz Thanks again for you review! Yes, I think we don't have to worry too much about the implementation details, at least at the moment, as I only made a minor change to the kernel.

[pcmoritz]

Sounds good, the PR looks great :)