# How to Use K-Grouped GEMM for MoE Weight Backward Pass in DeepGEMM

> Learn how K-grouped GEMM speeds up MoE weight backward pass in DeepGEMM. This method efficiently calculates expert gradients with batched kernel launches for better performance.

- Repository: [DeepSeek/DeepGEMM](https://github.com/deepseek-ai/DeepGEMM)
- Tags: how-to-guide
- Published: 2026-04-19

---

**K-grouped GEMM enables efficient computation of Mixture-of-Experts (MoE) weight gradients by partitioning the K dimension into independent sub-ranges, allowing each expert's gradient to be calculated in a single batched kernel launch.**

The deepseek-ai/DeepGEMM library provides high-performance GPU kernels specifically designed for this pattern. When training MoE models, the weight backward pass requires computing gradients for each expert independently based on the tokens routed to them. K-grouped GEMM for MoE weight backward pass consolidates these independent operations into optimized kernels that handle varying K-sizes per expert without launching separate CUDA kernels.

## What Is K-Grouped GEMM?

**K-grouped GEMM** is a kernel family that groups the K dimension of a matrix multiplication into independent sub-ranges. Unlike standard GEMM where K is uniform across the batch, K-grouped kernels accept a list of K sizes—one per group—and internally schedule the computation so that each group processes its specific K-range.

In the context of MoE training, each expert acts as a group. The weight gradient computation for expert *e* only involves tokens routed to that expert, resulting in a unique K size per expert. The K-grouped abstraction maps naturally to this workload.

## Why Use K-Grouped GEMM for MoE Weight Gradients?

Traditional approaches launch separate GEMM kernels for each expert or pad tensors to a uniform K size. Both approaches suffer from kernel launch overhead or memory waste.

**K-grouped GEMM solves this by:**

- **Single kernel launch**: All experts are processed in one kernel, amortizing launch overhead across the expert count.
- **No padding**: The kernel respects the actual K-size of each expert, avoiding wasted computation on padded zeros.
- **Optimized memory access**: The kernel implementation in `deep_gemm` uses specialized Triton and CUDA kernels to maintain coalesced access patterns despite irregular K sizes.

## Prerequisites and Setup

Before implementing the backward pass, ensure you have the correct tensor layouts and precision formats.

**Required imports and utilities:**

```python
import torch
import deep_gemm
from deep_gemm.utils import per_token_cast_to_fp8, per_token_cast_to_fp4

```

**Precision requirements:**
- For FP8/FP4 weights: Use `per_token_cast_to_fp8` or `per_token_cast_to_fp4` with consistent quantization granularity (e.g., `gran_k=32`).
- For BF16 weights: Use the `k_grouped_bf16_gemm_tn_contiguous` kernel directly without casting.

## Implementing the MoE Weight Backward Pass

The implementation follows four distinct steps: preparing metadata, casting tensors, allocating output, and launching the kernel.

### Step 1: Prepare Grouping Metadata

First, calculate the K-size for each expert. For a given expert, the K dimension equals the number of tokens routed to that expert multiplied by the hidden dimension.

```python

# Example: 8 experts with varying token counts

num_experts = 8
tokens_per_expert = [1024, 980, 1010, 995, 1005, 1021, 987, 998]
hidden_size = 4096

# K-size per expert = tokens * hidden_size

ks = [t * hidden_size for t in tokens_per_expert]
ks_tensor = torch.tensor(ks, dtype=torch.int32, device='cuda')

```

The `ks` list and `ks_tensor` are required by the low-level kernel to determine the offset and size of each expert's K-range.

### Step 2: Cast Tensors to Supported Precision

Cast the forward activations (`a`) and upstream gradients (`b`) to FP8 (or FP4) using the provided utilities. Both must use the same quantization granularity.

```python
total_tokens = sum(tokens_per_expert)

# Activations: (M, K) where M = total tokens, K = hidden

a_bf16 = torch.randn((total_tokens, hidden_size), dtype=torch.bfloat16, device='cuda')
a_fp8, a_scale = per_token_cast_to_fp8(
    a_bf16, 
    use_ue8m0=True, 
    gran_k=32, 
    use_packed_ue8m0=True
)

# Upstream gradient: (N, K) where N = output dimension

output_dim = 4096
b_bf16 = torch.randn((output_dim, hidden_size), dtype=torch.bfloat16, device='cuda')
b_fp8, b_scale = per_token_cast_to_fp8(
    b_bf16,
    use_ue8m0=True,
    gran_k=32,
    use_packed_ue8m0=True
)

```

### Step 3: Allocate the Output Buffer

Allocate the weight gradient buffer `d` with shape `(total_K, N)`, where `total_K = sum(ks)`. This buffer receives the computed gradients for all experts.

```python
total_K = sum(ks)
d = torch.empty((total_K, output_dim), dtype=torch.bfloat16, device='cuda')

```

### Step 4: Launch the K-Grouped Kernel

Invoke the appropriate K-grouped GEMM kernel. For FP8/FP4 weights, use `k_grouped_fp8_gemm_tn_contiguous` (or `nt` variant depending on layout). For BF16, use `k_grouped_bf16_gemm_tn_contiguous`.

```python

# FP8/FP4 kernel

deep_gemm.k_grouped_fp8_gemm_tn_contiguous(
    a_fp8,           # (M, K) - activations

    b_fp8,           # (N, K) - upstream gradient

    d,               # (total_K, N) - weight gradient output

    ks,              # Python list of K sizes

    ks_tensor,       # torch.int32 tensor of K sizes

    c=None,          # Optional accumulator

    recipe=(1, 1, 32) # Quantization recipe

)

# Alternative: BF16 kernel (if not using FP8/FP4)

# deep_gemm.k_grouped_bf16_gemm_tn_contiguous(a_fp8, b_fp8, d, ks, ks_tensor, c=None)

```

The kernel automatically iterates over the `ks` ranges, computing each expert's weight gradient and writing it to the corresponding slice of `d`.

## Complete Code Example

Here is the full implementation combining all steps:

```python
import torch
import deep_gemm
from deep_gemm.utils import per_token_cast_to_fp8

def compute_moe_weight_grad(
    tokens_per_expert: list,
    hidden_size: int,
    output_dim: int,
    device: str = 'cuda'
):
    """
    Compute MoE weight gradients using K-grouped GEMM.
    
    Args:
        tokens_per_expert: List of token counts for each expert
        hidden_size: Input feature dimension (K)
        output_dim: Expert output dimension (N)
    """
    num_experts = len(tokens_per_expert)
    
    # 1. Prepare grouping metadata

    ks = [t * hidden_size for t in tokens_per_expert]
    ks_tensor = torch.tensor(ks, dtype=torch.int32, device=device)
    total_K = sum(ks)
    total_tokens = sum(tokens_per_expert)
    
    # 2. Prepare tensors (using random data for demonstration)

    # Activations: (total_tokens, hidden)

    a = torch.randn((total_tokens, hidden_size), 
                    dtype=torch.bfloat16, device=device)
    a_fp8, _ = per_token_cast_to_fp8(a, use_ue8m0=True, 
                                     gran_k=32, use_packed_ue8m0=True)
    
    # Upstream gradient: (output_dim, hidden)

    b = torch.randn((output_dim, hidden_size), 
                    dtype=torch.bfloat16, device=device)
    b_fp8, _ = per_token_cast_to_fp8(b, use_ue8m0=True, 
                                     gran_k=32, use_packed_ue8m0=True)
    
    # 3. Allocate output buffer

    d = torch.empty((total_K, output_dim), 
                    dtype=torch.bfloat16, device=device)
    
    # 4. Launch K-grouped GEMM

    deep_gemm.k_grouped_fp8_gemm_tn_contiguous(
        a_fp8, b_fp8, d, ks, ks_tensor, 
        c=None, recipe=(1, 1, 32)
    )
    
    return d

# Example usage

if __name__ == "__main__":
    tokens_per_expert = [1024, 980, 1010, 995, 1005, 1021, 987, 998]
    weight_grads = compute_moe_weight_grad(
        tokens_per_expert, 
        hidden_size=4096, 
        output_dim=4096
    )
    print(f"Computed weight gradients shape: {weight_grads.shape}")

```

## Key Implementation Details

**Kernel variants**: The `k_grouped_fp8_gemm_*` functions come in `tn` and `nt` variants referring to the transpose state of the input matrices. The `tn` variant expects `a` in row-major and `b` in column-major (or transposed) layout. Choose the variant that matches your tensor layout to avoid unnecessary transposition overhead.

**Accumulator parameter**: The `c` parameter allows on-device accumulation. If you need to accumulate gradients across multiple micro-batches, pass a pre-allocated tensor to `c` and set the accumulation flag. For pure weight-gradient computation with no accumulation, pass `c=None`.

**Quantization consistency**: When using FP8/FP4 kernels, both `a` (activations) and `b` (upstream gradients) must use the same quantization granularity (`gran_k`). Mismatched granularity causes incorrect results. The `recipe` parameter controls the accumulation precision and tile sizes.

## Summary

- **K-grouped GEMM** segments the K dimension into independent ranges, making it ideal for MoE weight gradients where each expert processes a different token count.
- **Implementation requires** preparing a `ks` list and `ks_tensor` containing per-expert K sizes, casting inputs to FP8/BF16, and allocating a contiguous output buffer.
- **Key functions** are `k_grouped_fp8_gemm_tn_contiguous` and `k_grouped_bf16_gemm_tn_contiguous`, exposed in [`deep_gemm/__init__.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/__init__.py).
- **Performance gains** come from fusing all expert computations into a single kernel launch, eliminating padding overhead and reducing launch latency.

## Frequently Asked Questions

### What is the difference between K-grouped GEMM and standard batched GEMM?

Standard batched GEMM requires all matrices in the batch to have identical dimensions (M, N, K). **K-grouped GEMM** allows each batch element (expert) to have a different K size while sharing the same M and N dimensions. This is essential for MoE training where load balancing results in uneven token distribution across experts.

### Which kernel should I use for BF16 weights instead of FP8?

For BF16 weights, use `deep_gemm.k_grouped_bf16_gemm_tn_contiguous`. This kernel accepts BF16 inputs directly without requiring the `per_token_cast_to_fp8` preprocessing step. The function signature is identical to the FP8 variant but omits the `recipe` parameter.

### How do I handle the accumulator parameter for gradient accumulation?

Pass a pre-allocated tensor of shape `(total_K, N)` to the `c` parameter. The kernel will add the computed gradients to the existing values in `c`. If you pass `c=None`, the kernel writes directly to `d` without accumulation. This is useful when accumulating gradients across multiple micro-batches in gradient accumulation training.

### What causes the "granularity mismatch" error when calling the kernel?

This error occurs when the quantization granularity (`gran_k`) used to cast `a` (activations) differs from the granularity used to cast `b` (upstream gradients). Both tensors must use the same `gran_k` value (typically 32) and the same packed format settings when calling `per_token_cast_to_fp8`.