feat(deepseek_v3): Add grouped GEMM kernel for faster MoE computation
Author: bzantium
Pull request: no
Labels: name: Feature request

### Feature request

This is a proposal to integrate an optimized Mixture-of-Experts (MoE) forward pass for the `DeepseekV3` model using the `grouped_gemm` library. This new implementation significantly accelerates MoE computations by replacing the iterative execution over experts with a single, batched kernel call.

The key changes include:
* A new, optional `grouped_forward` path in the `DeepseekV3MoE` module that leverages fused expert weights.
* `GroupedLinear` and `GroupedDeepseekV3MLP` modules to store the weights of all experts in a single, contiguous tensor, enabling efficient processing.
* Public-facing `fuse_experts()` and `unfuse_experts()` methods on the `DeepseekV3ForCausalLM` model to allow users to dynamically switch between the standard PyTorch implementation and the optimized `grouped_gemm` version.

The implementation relies on the `grouped_gemm` library, which can be found here: [https://github.com/fanshiqing/grouped_gemm](https://github.com/fanshiqing/grouped_gemm).

### Motivation

The current MoE implementation in `DeepseekV3MoE` processes experts by iterating through them in a Python `for` loop. For models with a large number of experts, such as DeepseekV3, this approach becomes a major performance bottleneck due to the high overhead of launching multiple small GPU kernels. I'm always frustrated when a model's architecture can't be fully utilized due to inefficient, non-batched operations.

By leveraging a grouped GEMM kernel, we can perform the matrix multiplications for all activated experts in a single, batched operation. This minimizes kernel launch overhead and significantly improves GPU utilization, leading to substantial speedups for both training and inference. This optimization makes running and fine-tuning `DeepseekV3` more efficient and accessible.

### Your contribution

I will implement the full feature as detailed in the proposal.

Comments:
how do you got this idea?

Hi, @1himan! Thanks for the question.

I've been working on LLM pre-training and post-training, and recently I've been focusing on models with Mixture-of-Experts (MoE) architectures like DeepseekV3.

During training, I noticed that the standard MoE implementation, which processes each expert in a sequential for loop, was a major performance bottleneck. This naive approach leads to significant GPU kernel launch overhead and low GPU utilization, which was drastically slowing down our training cycles.

To address this, I looked for a way to batch the computations for all activated experts into a single operation. I realized that a grouped GEMM kernel was the perfect solution, as it's designed to handle exactly this kind of scenario. By fusing the expert weights and using a single kernel launch, we can minimize overhead and maximize GPU efficiency. That's what led to this proposal and PR.

hello, when using your code implementation, I encountered the following issue: According to the operator definition
```
Tensor unpermuted_output = torch::empty({num_tokens, num_cols}, torch::dtype(_st).device(torch::kCUDA).requires_grad(false));
```
the result generated by the `permute` operator is placed on the default device cuda:0. However, when executing with Transformers’ FSDP (Fully Sharded Data Parallel), different layers are distributed across different devices (e.g., cuda:1, cuda:2, etc.). After the `permute` operator is executed, the result ends up on cuda:0, which causes subsequent calculations to fail with errors like "tensors are not on the same device."

How should this issue be resolved?

Could you share the code to know about the context? I think the important thing is when `fuse_experts` is called. I think `fuse_experts` (on cpu) then FSDP model can solve the problem.
to: @plant310 

My code is:
```
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, device_map="auto")  
model.fuse_experts()  
```
Whether I perform fuse_experts() offline or online, I encounter the error:
`Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!`
However, if I set `device_map="cuda:0"`, the code executes successfully. But how should I handle this when running larger models?

My recommendation is to remove `device_map="auto"` when loading model. Load the model on cpu and run `fuse_experts()` then move the model to gpu devices.

is the main branch still using for loop?

Since https://github.com/huggingface/transformers/pull/42697 is merged, I will close this issue.