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)
- 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
-
TmemAddress: Simple TMEM address wrapper for load/store operations. -
TmemAllocation: Handle to allocated Tensor Memory. -
TmemStage: A pipeline stage within TMEM for accumulator buffering.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!