For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo struct
TmemAllocation
struct TmemAllocation[cta_group: Int, max_cols: Int = Int(512)]
Handle to allocated Tensor Memory.
Lifecycle: allocate() → use → release_lock() → wait → deallocate()
Parameters
Fields
- addr (
UInt32):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
SmemAddrStorage
comptime SmemAddrStorage = SMemArray[UInt32, Int(1)]
Methods
__init__
def __init__(addr: UInt32) -> Self
allocate
static def allocate(smem_addr: SMemArray[UInt32, Int(1)]) -> Self
Allocate TMEM (MMA warp). Address stored in smem for epilogue.
from_shared
static def from_shared(smem_addr: SMemArray[UInt32, Int(1)]) -> Self
Get handle from existing allocation (epilogue warp).
release_lock
def release_lock(self)
Release allocation lock before waiting for epilogue.
deallocate
def deallocate(self)
Free TMEM after epilogue completion.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!