# How to Use DeepGEMM Weight Gradient Kernels for Training Backward Passes

> Learn how to use DeepGEMM weight gradient kernels for efficient backward passes in MoE training. Accelerate your deep learning with optimized GEMM operations.

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

---

**Use DeepGEMM's K-grouped GEMM kernels via `k_grouped_fp8_gemm_nt_contiguous` (SM90) or `k_grouped_fp8_gemm_tn_contiguous` (SM100) to compute variable-K weight gradients for Mixture-of-Experts training backward passes.**

DeepGEMM provides specialized **K-grouped GEMM kernels** designed specifically for computing weight gradients during Mixture-of-Experts (MoE) training. These kernels handle the backward pass pattern where each expert processes a different number of tokens, keeping the **M** (batch) and **N** (output) dimensions fixed while allowing the **K** dimension (input features) to vary per expert.

## Understanding K-Grouped Weight Gradient Kernels

The weight gradient computation in MoE training requires aggregating gradients across variable numbers of tokens per expert. DeepGEMM implements this via K-grouped kernels that map directly to the CUDA architecture capabilities of SM90 (Hopper) and SM100 (Blackwell).

### API Entry Points

DeepGEMM exposes two primary Python functions for weight gradient computation in [`csrc/apis/gemm.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/apis/gemm.hpp) (lines 271-284):

- **`k_grouped_fp8_gemm_nt_contiguous`** – For SM90 architectures, computes non-transposed A by transposed B GEMM
- **`k_grouped_fp8_gemm_tn_contiguous`** – For SM100 architectures, computes transposed A by non-transposed B GEMM

For BF16 accumulation, the equivalent `k_grouped_bf16_gemm_tn_contiguous` is available at lines 527-543 in the same file.

### Dispatch and Tensor-Map Handling

The kernel selection logic (lines 306-314) automatically dispatches to the appropriate architecture-specific implementation:

- **SM90**: Uses `sm90_k_grouped_fp8_gemm_1d1d` with pre-allocated Tensor-Map buffers (lines 349-356) to stream variable-size K slices without memory copies
- **SM100**: Uses `sm100_k_grouped_fp8_gemm_1d1d` with optimized TN layout for Blackwell tensor cores

Shape validation (lines 286-302) ensures the sum of per-expert K dimensions matches the total input size, while early-exit logic (lines 298-301) skips zero-size groups for inactive experts.

## Implementing Weight Gradient Computation in Python

DeepGEMM's Python bindings expose the C++ kernels directly via `pybind11` (lines 447-456). The weight gradient workflow requires preparing per-expert token counts and contiguous memory layouts.

### FP8 Implementation Example

The following example computes weight gradients using FP8 precision on an SM90 GPU:

```python
import torch
import deep_gemm
from deep_gemm.testing import get_arch_major

# 1. Define per-expert token counts (K dimension slices)

ks = [32, 45, 27]  # 3 experts with variable tokens

ks_tensor = torch.tensor(ks, dtype=torch.int, device='cuda')

# 2. Prepare input dimensions

M, N = 64, 128  # Weight matrix shape M×N

total_k = sum(ks)

# 3. Create FP8 input tensors

# A: activations [sum(K_i), M]

# B: upstream gradients [sum(K_i), N]

A = torch.randn(total_k, M, dtype=torch.float8_e4m3fn, device='cuda')
B = torch.randn(total_k, N, dtype=torch.float8_e4m3fn, device='cuda')

# 4. Allocate output weight gradient

D = torch.empty(M, N, dtype=torch.bfloat16, device='cuda')

# 5. Dispatch to appropriate kernel

if get_arch_major() == 9:  # SM90 (Hopper)

    deep_gemm.k_grouped_fp8_gemm_nt_contiguous(
        (A, None),  # (tensor, scale_factor)

        (B, None),
        D,
        ks,
        ks_tensor,
        None,  # optional accumulator

        recipe=(1, 1, 128),  # (M-tile, N-tile, K-granularity)

    )
else:  # SM100 (Blackwell)

    deep_gemm.k_grouped_fp8_gemm_tn_contiguous(
        (A, None),
        (B, None),
        D,
        ks,
        ks_tensor,
        None,
        recipe=(1, 1, 128),
    )

```

The `recipe` parameter `(1, 1, 128)` specifies the K-granularity for the 1-D1-D kernel implementation. As enforced by `DG_HOST_ASSERT(gran_k == 32 or gran_k == 128)` in the source, only granularities of 32 or 128 are supported for optimal tensor core utilization.

### BF16 Alternative

For training scenarios requiring higher precision accumulation, replace the FP8 calls with:

```python
deep_gemm.k_grouped_bf16_gemm_tn_contiguous(
    (A_bf16, None),
    (B_bf16, None),
    D,
    ks,
    ks_tensor,
    None,
    recipe=(1, 1, 128),
)

```

This variant uses BF16 input tensors and accumulates into the output without FP8 quantization steps.

## Integration Best Practices for MoE Training

| Scenario | Implementation Strategy |
|----------|------------------------|
| **Mixed-precision training** | Keep forward passes in FP8 for speed, but accumulate weight gradients in BF16 using `k_grouped_bf16_gemm_tn_contiguous` for numerical stability. |
| **Empty expert groups** | The kernel automatically skips zero-size K slices (see early-exit logic at lines 298-301). Include zeros in `ks` for inactive experts without overhead. |
| **Tensor-Map overhead** (SM90 only) | The kernel handles temporary Tensor-Map buffer allocation internally (lines 349-356). No manual management required for variable K streaming. |
| **Performance tuning** | Use `recipe=(1, 1, 128)` for maximum throughput on SM90/SM100. The 1-D1-D kernel requires K-granularity of 32 or 128; 128 yields optimal tensor core occupancy. |

## Key Source Files and Implementation Details

| File | Purpose | Critical Lines |
|------|---------|----------------|
| [`csrc/apis/gemm.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/apis/gemm.hpp) | Core API definitions, dispatch logic, shape validation, and Python bindings for all K-grouped weight-gradient kernels. | 271-284 (API), 286-302 (validation), 306-314 (dispatch), 349-356 (Tensor-Map), 447-456 (Python), 527-543 (BF16) |
| [`tests/test_fp8_fp4.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_fp8_fp4.py) | End-to-end validation tests for K-grouped kernels, including per-expert token distribution and architecture detection. | 179-182 (architecture check), 185-190 (empty group handling) |
| [`deep_gemm/__init__.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/__init__.py) | Python package re-exports exposing `k_grouped_fp8_gemm_*` functions to user code. | Module-level exports |
| [`README.md`](https://github.com/deepseek-ai/DeepGEMM/blob/main/README.md) | Project documentation announcing weight-gradient kernel availability (2025-05-14) and pointing to K-grouped API usage. | News section |

## Summary

- **DeepGEMM weight gradient kernels** use K-grouped GEMM primitives to handle variable token counts per expert during MoE backward passes.
- **Two primary APIs** exist: `k_grouped_fp8_gemm_nt_contiguous` for SM90 and `k_grouped_fp8_gemm_tn_contiguous` for SM100, with BF16 variants available for stable accumulation.
- **Implementation requires** per-expert token counts (`ks`), contiguous input tensors with shape `[sum(K_i), M]` and `[sum(K_i), N]`, and a recipe tuple specifying K-granularity (32 or 128).
- **Automatic optimizations** include architecture-specific dispatch, internal Tensor-Map management for SM90, and early-exit handling for empty expert groups.

## Frequently Asked Questions

### What is the difference between the NT and TN variants of DeepGEMM weight gradient kernels?

The **NT** (non-transposed A, transposed B) variant `k_grouped_fp8_gemm_nt_contiguous` is optimized for SM90 (Hopper) GPUs, while the **TN** (transposed A, non-transposed B) variant `k_grouped_fp8_gemm_tn_contiguous` targets SM100 (Blackwell) tensor cores. The distinction reflects architectural differences in memory coalescing and tensor core instruction sets between the two GPU generations, with the API automatically selecting the correct implementation based on `device_runtime->get_arch_major()`.

### How do I handle experts with zero tokens in the backward pass?

DeepGEMM weight gradient kernels automatically handle empty expert groups through early-exit logic at lines 298-301 of [`csrc/apis/gemm.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/apis/gemm.hpp). Simply include zero values in your `ks` tensor for inactive experts; the kernel will skip these groups without launching unnecessary CUDA operations or causing synchronization overhead. This design supports dynamic expert routing where token distribution varies across training steps.

### Can I use BF16 instead of FP8 for weight gradient accumulation?

Yes, DeepGEMM provides `k_grouped_bf16_gemm_tn_contiguous` (lines 527-543 in [`csrc/apis/gemm.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/apis/gemm.hpp)) specifically for BF16 weight gradient accumulation. This is recommended for mixed-precision training scenarios where forward passes use FP8 for speed but backward weight updates require higher numerical precision to maintain training stability. The API signature matches the FP8 variants, accepting BF16 input tensors and accumulating into a BF16 output matrix.