For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /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

• ​cta_group (Int): Cooperating CTAs (1 or 2).
• ​max_cols (Int): TMEM columns (512 for SM100).

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.

Args:

• ​smem_addr (SMemArray[UInt32, Int(1)]): Shared memory slot that receives the allocated TMEM address for cross-warp sharing with the epilogue.

from_shared

static def from_shared(smem_addr: SMemArray[UInt32, Int(1)]) -> Self

Get handle from existing allocation (epilogue warp).

Args:

• ​smem_addr (SMemArray[UInt32, Int(1)]): Shared memory slot holding the TMEM address previously written by allocate.

release_lock

def release_lock(self)

Release allocation lock before waiting for epilogue.

deallocate

def deallocate(self)

Free TMEM after epilogue completion.