Implementation:InternLM Lmdeploy GemmUniversalSm90V3

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

Overview

SM90 GEMM v3 introduces the ScaledGmmaFP8_TN abstraction for batched/pipelined FP8 GMMA operations with per-block FP8 scaling, supporting grouped GEMM with runtime tensormap updates.

Description

Distinguishing features (v3): Refactors the GMMA computation to use the modular ScaledGmmaFP8_TN helper, which handles batched MMA operations and per-block scale application in a composable way.

Key characteristics:

• Tile size: 128x192x128 (narrower N than v1/v2)
• Warpgroups: 2 math (2M x 1N) + 1 scheduler = 3 (CTA_SIZE = 384)
• GMMA abstraction: Uses ScaledGmmaFP8_TN<WG_TILE_M, WG_TILE_N, TILE_K, 1, 1, 1, 1> for cleaner MMA invocation
• Per-block scaling: U-scales loaded via cp.async, V-scale predicate array tracks 128-element block boundaries
• Pipeline: 4-stage TMA pipeline
• Grouped GEMM: Supports dynamic tensormap replacement via tensormap_replace when switching between groups
• U-scale loading: Producer loads U-scales asynchronously using cp.async rather than TMA, with named barrier synchronization

Usage

Optimized for grouped FP8 GEMM workloads where per-block quantization scaling and heterogeneous problem sizes are required.

Code Reference

Source Location

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

Signature

template<Order raster_order, int multicast_a, int multicast_b, bool is_grouped_gemm_>
struct GemmUniversalSm90_v3 {
    static constexpr int TILE_M = 128, TILE_N = 192, TILE_K = 128;
    static constexpr int CTA_SIZE = 384;
    static constexpr int Stages = 4;

    using GMMA = ScaledGmmaFP8_TN<64, 192, 128, 1, 1, 1, 1>;
};

Import

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

I/O Contract

Inputs

Name	Type	Required	Description
tm_a, tm_b, tm_c, tm_u, tm_v	CUtensorMap	Yes	TMA descriptors for all operands
param_A..V, param_C	MatrixParam	Yes	Matrix layout parameters
sched	Scheduler	Yes	Tile scheduler with grouped GEMM support
tensormap_buf	CUtensorMap*	Yes	Writable TMA descriptor buffer for grouped GEMM

Outputs

Name	Type	Description
C matrix	global memory	BF16 output via TMA store

Usage Examples

// V3 is used when grouped GEMM with ScaledGmma is needed
GemmUniversalSm90_v3<kRowMajor, 1, 2, true> kernel;

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime