Implementation:InternLM Lmdeploy Gemm Sm90Utils

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

Overview

Utility header for SM90 GEMM kernels providing shared memory descriptor creation, GMMA descriptor iteration, warp-group fence operations, WGMMA invocation wrappers, and cluster geometry helpers.

Description

This header consolidates SM90-specific utilities shared across multiple GEMM kernel variants:

• make_smem_desc: Constructs a cute::GmmaDescriptor from a shared memory pointer and layout type, setting the 1024-byte stride offset required by GMMA operations

• SmemDescIterV2: A lightweight iterator over GMMA shared memory descriptors. Stores the descriptor as a uint64_t union, with Advance (increments by step, wraps at stage count), Reset (jumps to a specific stage), and +=/-= operators for fine-grained navigation

• wgmma / wgmma_impl: Template wrappers that expand a float[N] fragment array into the variadic register arguments expected by MMA_Atom::fma, with scale-out mode (Zero for clear, One for accumulate)

• warpgroup_fence_operand: Compiler fence (asm volatile("")) over accumulator register arrays to prevent reordering across GMMA boundaries

• arch::Cluster: Template helper for SM90 cluster geometry, computing CTA-to-cluster mappings, multicast masks for A and B operands, and coordinate transformations

Usage

Included by all SM90 GEMM kernel variants (v1-v5) for descriptor management and GMMA invocation.

Code Reference

Source Location

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

Signature

__device__ cute::GmmaDescriptor make_smem_desc(void* smem_ptr, int layout_type);

template<int Stages, int Step>
struct SmemDescIterV2 {
    __device__ void Advance(int stage);
    __device__ void Reset(int stage);
    __device__ operator uint64_t();
};

template<class MMA_Atom, int N>
__device__ void wgmma(uint64_t desc_a, uint64_t desc_b, float (&frag_C)[N], bool clear);

__device__ void warpgroup_fence_operand(float& reg);

template<int M_, int N_, Order order>
struct arch::Cluster { /* cluster geometry helpers */ };

Import

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

I/O Contract

Inputs

Name	Type	Required	Description
smem_ptr	void*	Yes	Shared memory pointer for descriptor creation
desc_a, desc_b	uint64_t	Yes	GMMA shared memory descriptors for A and B
frag_C	float[N]	Yes	Accumulator fragment array
clear	bool	Yes	Whether to zero-initialize or accumulate

Outputs

Name	Type	Description
GmmaDescriptor	uint64_t	Encoded GMMA shared memory descriptor
frag_C	float[N]	Updated accumulator fragments after GMMA

Usage Examples

auto desc_A = make_smem_desc(smem_A_ptr, 1);
SmemDescIterV2<4, step> iter_A{desc_A};
wgmma<MMA_Atom>(iter_A, iter_B, frag_C, /*clear=*/true);
iter_A.Advance(stage);

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime