Implementation:InternLM Lmdeploy GemmUniversal

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

Overview

The universal GEMM kernel template that composes a mainloop, epilogue, and scheduler into a complete GPU kernel for SM70/SM75/SM80 architectures, supporting quantized operands, split-K, and grouped GEMM.

Description

GemmUniversal is a composite template parameterized by Arch, Mainloop, Epilogue, and Scheduler. It defines the complete GEMM kernel structure:

1. The scheduler (GemmScheduler or DynamicScheduler) determines which tile and K-range this CTA processes.
2. Operand matrix data pointers are resolved via resolve_op, supporting flat, blocked, and indexed striding modes.
3. Global memory iterators for operands A, B, U (A-scales), and V (B-scales) are constructed with appropriate offsets.
4. The mainloop (SM70 or SM80) executes the tiled MMA accumulation loop.
5. The epilogue rearranges, scales, combines, and stores the result.

The SharedStorage union overlaps mainloop and epilogue shared memory to minimize footprint. The gemm_kernel global function wraps the kernel with architecture compatibility checking.

Usage

Instantiated by KernelImpl with specific template parameters for each supported configuration (data types, tile sizes, pipeline stages).

Code Reference

Source Location

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

Signature

template<class Arch_, class Mainloop, class Epilogue_, class Scheduler_>
struct GemmUniversal {
    using Impl = Mainloop;
    __device__ void operator()(const GemmParam& param, const EpilogueParam& epi_param,
                                Scheduler& sched, char* smem_buf);
};

template<class Kernel, class Param, class EpilogueParam, class Scheduler>
__global__ void gemm_kernel(Param param, EpilogueParam epi_param, Scheduler sched);

Import

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

I/O Contract

Inputs

Name	Type	Required	Description
param	GemmParam	Yes	Matrix pointers and strides for A, B, U, V
epi_param	EpilogueParam	Yes	Output matrix, scaling, activation configuration
sched	Scheduler	Yes	Tile scheduler (static or dynamic)
smem_buf	char*	Yes	Dynamic shared memory buffer

Outputs

Name	Type	Description
D matrix	global memory	GEMM output D = alpha * (A @ B) + beta * C with optional quantization and activation

Usage Examples

// Launched by KernelImpl::Launch
gemm_kernel<Gemm><<<grid, block, dynamic_smem_size, stream>>>(param, epilogue, sched);

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime