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
OutputRegisterBuffer
struct OutputRegisterBuffer[dtype: DType, num_m_mmas: Int, num_n_mmas: Int, output_frag_size: Int]
Holds the attention output accumulator tile in register memory.
Stores the num_m_mmas * num_n_mmas MMA output fragments as a row-major
register TileTensor and applies the softmax row-sum denominator via
apply_softmax_denominator. Initialized zero-filled via __init__ and
resettable via zero.
Parameters
- dtype (
DType): Element dtype of the output accumulator tile stored in the register TileTensor. - num_m_mmas (
Int): Number of MMA tiles along the M (query) dimension of the output accumulator. - num_n_mmas (
Int): Number of MMA tiles along the N (key/value) dimension of the output accumulator. - output_frag_size (
Int): Per-lane element width of one MMA output fragment; the register tile is laid out as[num_n_mmas * num_m_mmas, output_frag_size].
Fields
- reg_tile (
OutputRegisterBuffer[dtype, num_m_mmas, num_n_mmas, output_frag_size].RegType):
Implemented traits
comptime members
reg_dtype
comptime reg_dtype = dtype
reg_layout
comptime reg_layout = row_major[(num_n_mmas * num_m_mmas), output_frag_size]()
RegType
comptime RegType = TileTensor[dtype, Layout[*?, *?], MutUntrackedOrigin, address_space=AddressSpace.LOCAL]
Methods
__init__
def __init__(out self)
apply_softmax_denominator
def apply_softmax_denominator[layout_type: TensorLayout, //](self, rowsum: TileTensor[dtype, layout_type, Storage=rowsum.Storage, address_space=rowsum.address_space, linear_idx_type=rowsum.linear_idx_type])
zero
def zero(self)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!