Mojo module
tmem
Tensor Memory (TMEM) abstractions for SM100 Blackwell GPUs.
TMEM is dedicated memory for MMA accumulators, separate from registers and shared memory. This module provides type-safe abstractions:
- TmemAllocation: Manages TMEM lifecycle (alloc/dealloc)
- TmemTensor: Layout-parameterized typed view over TMEM accumulators
- TmemStage: Represents a pipeline stage for accumulator buffering
- TmemAddress: Simple address wrapper for TMEM load operations
comptime values
TMEM_LOWER_ROW_OFFSET
comptime TMEM_LOWER_ROW_OFFSET = 1048576
Structs
-
BlockScaledTmem: TMEM region for block-scaled matmul with typed tile accessors. -
TmemAddress: Simple TMEM address wrapper for load/store operations. -
TmemAllocation: Handle to allocated Tensor Memory. -
TmemArrayType: Array of tiles in Tensor Memory (TMEM). -
TmemFragments: Paired upper/lower accumulator fragments from TMEM. -
TmemStage: A pipeline stage within TMEM for accumulator buffering. -
TmemTensor: Typed tensor view over Tensor Memory (TMEM) for MMA accumulators.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!