Implementation:InternLM Lmdeploy Gemm MainloopSm80V2

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

Overview

Multi-stage software-pipelined GEMM mainloop for Ampere (SM80) that uses cp.async for overlapping global-to-shared memory transfers with MMA computation across N pipeline stages.

Description

The SM80 v2 mainloop leverages the cp.async instruction to decouple memory loading from computation, enabling deeper pipelines than the SM70 double-buffering approach. Key features:

• N-stage pipeline: Uses SmemIter with configurable stage count (typically 3-6 stages) to rotate shared memory buffers
• Async copy: GmemIteratorSm80::Prefetch issues cp.async instructions, committed with __pipeline_commit()
• Pipeline wait: __pipeline_wait_prior<N-2>() ensures the target stage is ready before MMA reads it
• Quantization group tracking: GroupIter modulo counter triggers scale operand (U/V) advancement at the correct group boundaries
• Prologue/body/epilogue: The pipeline is primed with an initial prologue filling stages before entering the steady-state body loop

The mainloop interleaves MMA operations with prefetch issue for maximum throughput.

Usage

Used by GemmUniversal when targeting SM80 (Ampere A100/A30) GPUs.

Code Reference

Source Location

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

Signature

template<class Impl_>
struct MainloopSm80_v2 {
    using Impl = Impl_;
    using FragC = typename Impl::FragC;

    __device__ void operator()(GmemIterA& gmem_A, GmemIterB& gmem_B,
                                GmemIterU& gmem_U, GmemIterV& gmem_V,
                                FragC& frag_C, int tile_iter,
                                SharedStorage& storage);
};

Import

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

I/O Contract

Inputs

Name	Type	Required	Description
gmem_A, gmem_B	GmemIteratorSm80	Yes	Async copy iterators for A and B
gmem_U, gmem_V	GmemIteratorSm80	Yes	Async copy iterators for scale factors
tile_iter	int	Yes	Number of K-tile iterations
storage	SharedStorage	Yes	N-stage shared memory buffer

Outputs

Name	Type	Description
frag_C	FragC&	Accumulated MMA result fragments

Usage Examples

// Primes the pipeline then runs steady-state loop
MainloopSm80_v2<Impl> mainloop{};
mainloop(gmem_A, gmem_B, gmem_U, gmem_V, frag_C, tile_iter, storage.mainloop);
// frag_C now contains accumulated result

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime