Implementation:InternLM Lmdeploy Gemm IteratorSm80
| Knowledge Sources | |
|---|---|
| Domains | GPU_Kernels, GEMM |
| Last Updated | 2026-02-07 15:00 GMT |
Overview
Global memory iterator for Ampere (SM80) architecture that uses cp.async instructions to asynchronously copy matrix tiles directly from global memory to shared memory.
Description
GmemIteratorSm80 is the SM80 (Ampere) counterpart to GmemIteratorSm70. Instead of loading data into registers and then storing to shared memory, it uses the cp.async hardware instruction to perform direct global-to-shared-memory transfers, enabling deeper software pipelining.
The iterator precomputes shared memory phase offsets for the swizzled layout (stored in phases_ array), source data pointers for each strided iteration (stored in src_data_vec_), and an optional L2 cache policy descriptor. The Prefetch method issues cp.async operations for all iteration points, while Advance moves the K-dimension offset for the next stage.
Supports batched prefetching where a single stage's loads are split across multiple MMA iterations for finer-grained overlap with computation.
Usage
Used by MainloopSm80_v2 to prefetch tiles of operands A, B, U, V asynchronously while MMA computation proceeds.
Code Reference
Source Location
- Repository: InternLM_Lmdeploy
- File: src/turbomind/kernels/gemm/iterator_sm80.h
Signature
template<class T, class Map, class SmemLayout, Pack kPack, Order kOrder,
bool AlignedC, bool AlignedS, Striding mode, class Policy_>
struct GmemIteratorSm80 {
using ThreadMap = Map;
using AccessType = Array<T, Map::kAccessC>;
static constexpr Striding kMode = mode;
// Prefetch, Advance, ClearSmem methods
};
Import
#include "src/turbomind/kernels/gemm/iterator_sm80.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| mat | MatrixData | Yes | Resolved matrix pointer and stride |
| offset | int2 | Yes | (M, K) starting offset |
| extent | int2 | Yes | (M, K) valid extent for predicates |
Outputs
| Name | Type | Description |
|---|---|---|
| shared memory | smem | Tile data asynchronously copied from global memory |
Usage Examples
GmemIteratorSm80<half, ThreadMap, SmemLayout, pack, kColMajor, true, true, Striding::kFlat, Policy>
gmem_A{mat_A, {offset_m, offset_k}, {extent_m, CTA_K}};
gmem_A.Prefetch(mask); // issues cp.async
__pipeline_commit(); // commit async group
gmem_A.Advance(); // move to next K tile