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

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:

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.

# 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.

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.

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.

# 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:

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.
• 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}