Implementation:InternLM Lmdeploy Gemm ThreadMap
| Knowledge Sources | |
|---|---|
| Domains | GPU_Kernels, GEMM |
| Last Updated | 2026-02-07 15:00 GMT |
Overview
Defines the ThreadMap template that maps individual threads within warps to memory access positions in contiguous (C) and strided (S) dimensions for global and shared memory operations.
Description
ThreadMap computes how threads within a CTA access a 2D tile of data:
- Dimensions:
DimC(contiguous, e.g., row for row-major) andDimS(strided, e.g., column) - Access granularity:
AccessCelements per thread per access - Warp decomposition: Each warp is split into
kWarpThreadCthreads along C andkWarpThreadSthreads along S - Warp tiling: Warps are partitioned along S first (to reduce strided iterations), then C
- Iteration counts:
kIterCandkIterSare the number of iterations each thread must execute
The get_offset(warp_id, lane_id) method returns the (C, S) starting offset for a given thread. Constants kDeltaC and kDeltaS give the stride between iterations. kAlignedC and kAlignedS indicate whether boundary checking can be skipped.
Also provides ThreadMapV2 variants and epilogue-specific thread maps that support different access patterns and partial tile handling.
Usage
Used by memory iterators (GmemIteratorSm70, GmemIteratorSm80) and the epilogue to determine each thread's memory access pattern.
Code Reference
Source Location
- Repository: InternLM_Lmdeploy
- File: src/turbomind/kernels/gemm/thread_map.h
Signature
template<int DimC, int DimS, int AccessC, int WarpCount, int WarpThreadC = ...>
struct ThreadMap {
static constexpr int kIterC, kIterS;
static constexpr int kDeltaC, kDeltaS;
static constexpr int kAccessC;
static constexpr bool kAlignedC, kAlignedS;
__device__ static int2 get_offset(int warp_id, int lane_id);
};
Import
#include "src/turbomind/kernels/gemm/thread_map.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| DimC, DimS | int (template) | Yes | Tile dimensions in contiguous and strided axes |
| AccessC | int (template) | Yes | Elements accessed per thread per iteration |
| WarpCount | int (template) | Yes | Number of warps in the CTA |
| warp_id | int | Yes | Warp index within the CTA |
| lane_id | int | Yes | Thread lane within the warp (0-31) |
Outputs
| Name | Type | Description |
|---|---|---|
| offset | int2 | (C, S) starting position for this thread |
Usage Examples
using Map = ThreadMap<128, 64, 8, 4>;
int2 offset = Map::get_offset(warp_id, lane_id);
// Thread accesses Map::kAccessC elements starting at offset, iterating kIterC x kIterS times