Implementation:InternLM Lmdeploy Gemm TiledMma

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

Overview

Wraps MMA atom operations and thread group maps into a unified Tiled_MMA_v2 abstraction that tiles MMA operations across a CTA's output matrix dimensions.

Description

Tiled_MMA_v2 composes an MMA_Atom (the hardware MMA instruction abstraction) with an MMA_Map (the warp-to-tile mapping) to define the complete tiled MMA computation:

• Atom: The primitive MMA instruction (e.g., 16x8x16 for HMMA, 8x8x4 for Volta, or SIMT)
• Map: How warp groups are assigned to sub-tiles of the CTA output
• Iteration: kMmaIterM x kMmaIterN x kMmaIterK total MMA iterations per warp
• Thread count: kGroupCount * Atom::kThreadCount threads total

The mma_k_iter method performs one K-iteration of MMA across all M and N tiles, with an optimization for column-major order that alternates M traversal direction on even/odd N iterations to improve register reuse.

Also provides Tiled_MMA_v3 and related helpers for shared memory copy integration.

Usage

Used by the GEMM mainloop implementations to execute the core tiled MMA computation.

Code Reference

Source Location

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

Signature

template<class MMA_Atom_, class MMA_Map_, Order order_ = kColMajor>
struct Tiled_MMA_v2 {
    using Atom = MMA_Atom_;
    using Map  = MMA_Map_;

    static constexpr int kMmaIterM, kMmaIterN, kMmaIterK;
    static constexpr int kThreadCount;

    __device__ static int3 get_offset(int thread_idx);

    template<class FragD, class FragA, class FragB, class FragC>
    __device__ static void mma_k_iter(FragD& frag_D, const FragA& frag_A,
                                       const FragB& frag_B, const FragC& frag_C);
};

Import

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

I/O Contract

Inputs

Name	Type	Required	Description
frag_A	register array	Yes	A operand fragments loaded from shared memory
frag_B	register array	Yes	B operand fragments loaded from shared memory
frag_C	register array	Yes	Accumulator input (for FMA: D = A*B + C)

Outputs

Name	Type	Description
frag_D	register array	MMA result accumulator

Usage Examples

using MMA = Tiled_MMA_v2<MMA_Atom_16816, MMA_Map, kColMajor>;
MMA::mma_k_iter(frag_D, frag_A, frag_B, frag_C);

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime