How to Use DeepGEMM Weight Gradient Kernels for Training Backward Passes

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

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:

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	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	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	Python package re-exports exposing k_grouped_fp8_gemm_* functions to user code.	Module-level exports
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. 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) 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.

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