Implementation:InternLM Lmdeploy Gemm ThreadGroupMap

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

Overview

Maps warp groups to MMA operation tiles within a CTA, providing the spatial decomposition of tiled MMA work across warps for the GEMM mainloop.

Description

This header defines two thread group mapping strategies:

• RakedThreadGroupMap: Maps groups of warps to sub-tiles of the CTA output in a raked (interleaved) pattern. Template parameters specify the total MMA dimensions (M, N, K), per-group tile sizes (TM, TN, TK), and group counts (GM, GN, GK). Each group's offset is computed as (m * footprintM, n * footprintN, k * footprintK).

• MMA_Map: A more flexible mapping that uses an ArrangementMN policy to define how groups are distributed in the M and N dimensions, with a separate K-dimension grouping factor. Supports both raked and blocked K-partitioning. Used to map warp groups to their starting MMA coordinates within a CTA tile.

Both provide get_offset(group_id) returning an int3 of (m, n, k) starting coordinates for MMA computation, and expose iteration counts and step sizes for the tiled loop structure.

Usage

Used by Tiled_MMA_v2 to determine each warp's MMA work assignment within the CTA.

Code Reference

Source Location

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

Signature

template<int M_, int N_, int K_, int TM, int TN, int TK, int GM, int GN, int GK>
struct RakedThreadGroupMap {
    static constexpr int kGroupCount = GM * GN * GK;
    __device__ static int3 get_offset(int group_id);
};

template<int M_, int N_, int K_, int tM_, int tN_, int tK_, class ArrangementMN, int gK, bool rK>
struct MMA_Map {
    static constexpr int kGroupCount = kGroupM * kGroupN * kGroupK;
    __device__ static int3 get_offset(int group_id);
};

Import

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

I/O Contract

Inputs

Name	Type	Required	Description
group_id	int	Yes	Index of the warp group within the CTA

Outputs

Name	Type	Description
offset	int3	(m, n, k) starting coordinates for this group's MMA tile

Usage Examples

using Map = RakedThreadGroupMap<128, 128, 32, 16, 16, 32, 4, 4, 1>;
int3 offset = Map::get_offset(group_id);  // returns MMA start coords

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime