IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

STMatrixOffsets

struct STMatrixOffsets[BM: Int, BN: Int, *, num_threads: Int, accum_dtype_size: Int, curr_repeat: Int, cumulative_repeat: Int, m_mma: Int]

Precomputed TMEM and local-fragment offsets for one st_matrix repeat column.

Parameters​

  • ​BM (Int): Number of rows in the BM x BN output tile (forwarded to STMatrixLayout).
  • ​BN (Int): Number of columns in the BM x BN output tile (forwarded to STMatrixLayout).
  • ​num_threads (Int): Number of threads participating in the st_matrix store (forwarded to STMatrixLayout).
  • ​accum_dtype_size (Int): Size in bytes of the accumulator element dtype (forwarded to STMatrixLayout).
  • ​curr_repeat (Int): Number of repeat columns in this power-of-two chunk of the st_matrix store.
  • ​cumulative_repeat (Int): Number of repeat columns already stored before this chunk, used as the TMEM column and local-fragment base offset.
  • ​m_mma (Int): M-tile index of the st_matrix store within the BM-row tile, selecting the 16-row TMEM quadrant.

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members​

b32_per_repeat​

comptime b32_per_repeat = (Int((mul accum_dtype_size, 4)) // Int(4))

local_frag_size_b32​

comptime local_frag_size_b32 = (curr_repeat * (Int((mul accum_dtype_size, 4)) // Int(4)))

ptr_offset​

comptime ptr_offset = ((Int((mul accum_dtype_size, 4)) // Int(4)) * Int((add (mul (BN // Int(8)), m_mma), cumulative_repeat)))

STLayout​

comptime STLayout = STMatrixLayout[BM, BN, num_threads=num_threads, accum_dtype_size=accum_dtype_size]

tmem_col_offset​

comptime tmem_col_offset = (Int((mul cumulative_repeat, 2)) * Int(4))

tmem_offset​

comptime tmem_offset = ((STMatrixOffsets[BM, BN, num_threads=num_threads, accum_dtype_size=accum_dtype_size, curr_repeat=curr_repeat, cumulative_repeat=cumulative_repeat, m_mma=m_mma].tmem_row_offset << Int(16)) + Int((mul cumulative_repeat, 8)))

tmem_row_offset​

comptime tmem_row_offset = (Int(16) * m_mma)

Methods​

__init__​

def __init__() -> Self

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact