Implementation:InternLM Lmdeploy Gemm ScaledGmmaFp8Sm90

Knowledge Sources	InternLM_Lmdeploy
Domains	GPU_Kernels, GEMM
Last Updated	2026-02-07 15:00 GMT

Overview

Implements the ScaledGmmaFP8_TN template that wraps SM90 GMMA (warp-group MMA) operations for FP8 E4M3 matrix multiplication with batched, pipelined per-block scaling.

Description

ScaledGmmaFP8_TN is a composable building block used by SM90 GEMM variants v3-v5. It selects the appropriate GMMA operation atom based on tile dimensions (64x64 through 64x256), then structures the MMA computation into batched and pipelined iterations:

• Operation selection: Automatically picks the widest GMMA atom that divides the tile size (e.g., MMA_64x192x32_F32E4M3E4M3_SS_TN for 192-wide N)
• Iteration hierarchy: ITER_M x ITER_N outer iterations, PIPE_M x PIPE_N pipeline stages, BATCH_M x BATCH_N batched operations per pipeline stage
• Fragment types: FragC (accumulator registers), FragU (A-scale factors with 2-wide interleaving), FragV (B-scale factors)
• Scaled accumulation: After each GMMA batch, applies U * V scaling to the raw MMA output before adding to the running accumulator
• Descriptor stepping: Computes shared memory descriptor offsets (kStepMA, kStepNB, kStepKA, kStepKB) for navigating the swizzled smem layout

Usage

Used as the GMMA computation engine in GemmUniversalSm90_v3, v4, and v5.

Code Reference

Source Location

• Repository: InternLM_Lmdeploy
• File: src/turbomind/kernels/gemm/scaled_gmma_fp8_sm90.h

Signature

template<int TILE_M, int TILE_N, int TILE_K, int BATCH_M, int BATCH_N, int PIPE_M, int PIPE_N>
struct ScaledGmmaFP8_TN {
    using Operation = /* auto-selected GMMA atom */;
    using FragC = typename Operation::CRegisters[PIPE_M][PIPE_N][BATCH_M][BATCH_N];
    using AccumC = FragC[ITER_M][ITER_N];
    using FragU = float[ITER_M][PIPE_M][BATCH_M][2];
    using FragV = float[2];

    static constexpr int OP_M, OP_N, OP_K;
    static constexpr int ITER_M, ITER_N;
};

Import

#include "src/turbomind/kernels/gemm/scaled_gmma_fp8_sm90.h"

I/O Contract

Inputs

Name	Type	Required	Description
TILE_M, TILE_N, TILE_K	int (template)	Yes	Tile dimensions
BATCH_M, BATCH_N	int (template)	Yes	Batching factors
PIPE_M, PIPE_N	int (template)	Yes	Pipeline stage factors

Outputs

Name	Type	Description
AccumC	register array	Scaled FP32 accumulator fragments
Operation	GMMA type	Selected GMMA operation atom

Usage Examples

using GMMA = ScaledGmmaFP8_TN<64, 192, 128, 1, 1, 1, 1>;
GMMA::AccumC accum{};
// Issue GMMA with scale application

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime