Skip to main content

/

DocsCode

Mojo struct

KBuffer

struct KBuffer[mut: Bool, dtype: DType, layout: Layout, address_space: AddressSpace, alignment: Int, origin: Origin[mut], masked: Bool, //, mma_shape: IndexList[3], k_group_size: Int, swizzle: OptionalReg[Swizzle], BN: Int, WN: Int, BK: Int, num_threads: Int]

Fields​

  • ​load_tile (LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)]):
  • ​mma_tile (LayoutTensor[dtype, Layout.row_major(ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1)), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)]):
  • ​smem_iter (LayoutTensorIter[dtype, blocked_product(Layout.row_major(BN, simd_width_of[dtype]()), Layout.row_major(1, (BK // simd_width_of[dtype]())), True), MutableAnyOrigin, address_space=AddressSpace(3), circular=True]):
  • ​bounds (Int):
  • ​global_iterator (LayoutTensorIter[dtype, LayoutTensor._compute_tile_layout[mut, dtype, layout, origin, address_space, Layout.__init__(IntTuple[__origin_of()](1), IntTuple[__origin_of()](1)), _get_layout_type(layout, address_space), _get_index_type(layout, address_space), masked, alignment, BN, BK]()[0], origin, address_space=address_space, axis=OptionalReg[Int]({:@stdlib::@builtin::@int::@Int {1}, 0}), layout_int_type=_get_layout_type(layout, address_space), linear_idx_type=_get_index_type(layout, address_space), masked=masked if masked else _tile_is_masked[layout, BN, BK]()]):

Implemented traits​

AnyType, UnknownDestructibility

Aliases​

__del__is_trivial​

alias __del__is_trivial = LayoutTensorIter[dtype, LayoutTensor._compute_tile_layout[mut, dtype, layout, origin, address_space, Layout.__init__(IntTuple[__origin_of()](1), IntTuple[__origin_of()](1)), _get_layout_type(layout, address_space), _get_index_type(layout, address_space), masked, alignment, BN, BK]()[0], origin, address_space=address_space, axis=OptionalReg[Int]({:@stdlib::@builtin::@int::@Int {1}, 0}), layout_int_type=_get_layout_type(layout, address_space), linear_idx_type=_get_index_type(layout, address_space), masked=masked if masked else _tile_is_masked[layout, BN, BK]()].__del__is_trivial if Int.__del__is_trivial if LayoutTensorIter[dtype, blocked_product(Layout.row_major(BN, simd_width_of[dtype]()), Layout.row_major(1, (BK // simd_width_of[dtype]())), True), MutableAnyOrigin, address_space=AddressSpace(3), circular=True].__del__is_trivial if LayoutTensor[dtype, Layout.row_major(ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1)), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial if LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[dtype, Layout.row_major(ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1)), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial if LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensorIter[dtype, blocked_product(Layout.row_major(BN, simd_width_of[dtype]()), Layout.row_major(1, (BK // simd_width_of[dtype]())), True), MutableAnyOrigin, address_space=AddressSpace(3), circular=True].__del__is_trivial if LayoutTensor[dtype, Layout.row_major(ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1)), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial if LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[dtype, Layout.row_major(ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1)), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial if LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else Int.__del__is_trivial if LayoutTensorIter[dtype, blocked_product(Layout.row_major(BN, simd_width_of[dtype]()), Layout.row_major(1, (BK // simd_width_of[dtype]())), True), MutableAnyOrigin, address_space=AddressSpace(3), circular=True].__del__is_trivial if LayoutTensor[dtype, Layout.row_major(ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1)), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial if LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[dtype, Layout.row_major(ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1)), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial if LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensorIter[dtype, blocked_product(Layout.row_major(BN, simd_width_of[dtype]()), Layout.row_major(1, (BK // simd_width_of[dtype]())), True), MutableAnyOrigin, address_space=AddressSpace(3), circular=True].__del__is_trivial if LayoutTensor[dtype, Layout.row_major(ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1)), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial if LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[dtype, Layout.row_major(ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1)), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial if LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial else LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)].__del__is_trivial

base_layout​

alias base_layout = Layout.row_major(BN, simd_width_of[dtype]())

GlobalTensorType​

alias GlobalTensorType = LayoutTensor[dtype, layout, origin, address_space=address_space, masked=masked, alignment=alignment]

GlobalTiledIteratorType​

alias GlobalTiledIteratorType = LayoutTensorIter[dtype, LayoutTensor._compute_tile_layout[mut, dtype, layout, origin, address_space, Layout.__init__(IntTuple[__origin_of()](1), IntTuple[__origin_of()](1)), _get_layout_type(layout, address_space), _get_index_type(layout, address_space), masked, alignment, BN, BK]()[0], origin, address_space=address_space, axis=OptionalReg[Int]({:@stdlib::@builtin::@int::@Int {1}, 0}), layout_int_type=_get_layout_type(layout, address_space), linear_idx_type=_get_index_type(layout, address_space), masked=masked if masked else _tile_is_masked[layout, BN, BK]()]

LoadTileType​

alias LoadTileType = LayoutTensor[dtype, Layout.row_major((ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size)) * ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)]

MMA_K​

alias MMA_K = mma_shape.__getitem__[3, DType.int64, Int](2)

MMA_N​

alias MMA_N = mma_shape.__getitem__[3, DType.int64, Int](1)

MMATileType​

alias MMATileType = LayoutTensor[dtype, Layout.row_major(ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1)), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)]

num_k_tiles​

alias num_k_tiles = ceildiv(BK, (mma_shape.__getitem__[3, DType.int64, Int](2) * k_group_size))

num_mmas​

alias num_mmas = ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1))

num_repeats​

alias num_repeats = (BK // simd_width_of[dtype]())

SharedIterType​

alias SharedIterType = LayoutTensorIter[dtype, blocked_product(Layout.row_major(BN, simd_width_of[dtype]()), Layout.row_major(1, (BK // simd_width_of[dtype]())), True), MutableAnyOrigin, address_space=AddressSpace(3), circular=True]

SharedTileType​

alias SharedTileType = LayoutTensor[dtype, blocked_product(Layout.row_major(BN, simd_width_of[dtype]()), Layout.row_major(1, (BK // simd_width_of[dtype]())), True), MutableAnyOrigin, address_space=AddressSpace(3), layout_int_type=_get_index_type(AddressSpace(3)), linear_idx_type=_get_index_type(AddressSpace(3))]

SharedWarpTileType​

alias SharedWarpTileType = LayoutTensor[dtype, LayoutTensor._compute_tile_layout[True, dtype, blocked_product(Layout.row_major(BN, simd_width_of[dtype]()), Layout.row_major(1, (BK // simd_width_of[dtype]())), True), MutableAnyOrigin, AddressSpace(3), Layout.__init__(IntTuple[__origin_of()](1), IntTuple[__origin_of()](1)), _get_index_type(AddressSpace(3)), _get_index_type(AddressSpace(3)), False, align_of[dtype](), WN, BK]()[0], MutableAnyOrigin, address_space=AddressSpace(3), layout_int_type=_get_index_type(AddressSpace(3)), linear_idx_type=_get_index_type(AddressSpace(3)), masked=_tile_is_masked[blocked_product(Layout.row_major(BN, simd_width_of[dtype]()), Layout.row_major(1, (BK // simd_width_of[dtype]())), True), WN, BK]()]

simd_width​

alias simd_width = simd_width_of[dtype]()

smem_layout​

alias smem_layout = blocked_product(Layout.row_major(BN, simd_width_of[dtype]()), Layout.row_major(1, (BK // simd_width_of[dtype]())), True)

thread_layout​

alias thread_layout = Layout.row_major((num_threads // 4), 4)

tiler_layout​

alias tiler_layout = Layout.row_major(1, (BK // simd_width_of[dtype]()))

wtile_dim0​

alias wtile_dim0 = WN

wtile_dim1​

alias wtile_dim1 = BK

Methods​

__init__​

__init__(out self, global_tile: LayoutTensor[dtype, layout, origin, address_space=address_space, masked=masked, alignment=alignment], num_b_rows: OptionalReg[Int], shared_ptr: UnsafePointer[Scalar[dtype], address_space=AddressSpace(3), mut=mut, origin=origin])

load_from_dram​

load_from_dram(mut self)

get_mma_tile​

get_mma_tile(self) -> LayoutTensor[dtype, Layout.row_major(ceildiv(WN, mma_shape.__getitem__[3, DType.int64, Int](1)), simd_width_of[dtype]()), MutableAnyOrigin, address_space=AddressSpace(5)]

Returns:

LayoutTensor

copy_to_shared​

copy_to_shared(self)

load_from_shared​

load_from_shared[accum_type: DType, mma_input_type: DType, _mma_shape: IndexList[3], _k_group_size: Int, transpose_b: Bool, k_mma: Int](self)

View source

View source

Was this page helpful?

PRODUCT

Mammoth

MAX

Mojo

Editions

Install

QUICK START

Documentation

Tutorials

Model Repo

Recipes

GPU Puzzles

Quickstart

SOLUTIONS

AI Agents

AI Inference

Batch processing

Chatbots

Code Generation

RAG & CAG

Research

DEVELOPERS

Modular Changelog

Tech Talks

Mojo🔥 Playground

Hackathons

AI Coding Assistant

FAQ

CONNECT

Blog

Community

Community Forum

Community Highlights

Report Issue

COMPANY

About Us

Culture

Careers

Talk to Us

@ Copyright - Modular Inc - 2025

Terms,

Privacy

&

Acceptable Use