Skip to main content

/

DocsCode
Log in

Mojo struct

KBuffer

struct KBuffer[dtype: DType, layout: Layout, address_space: AddressSpace, alignment: Int, origin: Origin[], 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, 0 if (simd_width_of[dtype]() == 0) else ((div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) + -1) if ((((rem_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) == 0) ^ True) & ((simd_width_of[dtype]() < 0) ^ (BK < 0))) else (div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value)), 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 = True

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 = 0 if (simd_width_of[dtype]() == 0) else ((div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) + -1) if ((((rem_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) == 0) ^ True) & ((simd_width_of[dtype]() < 0) ^ (BK < 0))) else (div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value)

SharedIterType​

alias SharedIterType = LayoutTensorIter[dtype, blocked_product(Layout.row_major(BN, simd_width_of[dtype]()), Layout.row_major(1, 0 if (simd_width_of[dtype]() == 0) else ((div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) + -1) if ((((rem_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) == 0) ^ True) & ((simd_width_of[dtype]() < 0) ^ (BK < 0))) else (div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value)), 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, 0 if (simd_width_of[dtype]() == 0) else ((div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) + -1) if ((((rem_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) == 0) ^ True) & ((simd_width_of[dtype]() < 0) ^ (BK < 0))) else (div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value)), 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, 0 if (simd_width_of[dtype]() == 0) else ((div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) + -1) if ((((rem_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) == 0) ^ True) & ((simd_width_of[dtype]() < 0) ^ (BK < 0))) else (div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value)), 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, 0 if (simd_width_of[dtype]() == 0) else ((div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) + -1) if ((((rem_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) == 0) ^ True) & ((simd_width_of[dtype]() < 0) ^ (BK < 0))) else (div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value)), 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, 0 if (simd_width_of[dtype]() == 0) else ((div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) + -1) if ((((rem_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) == 0) ^ True) & ((simd_width_of[dtype]() < 0) ^ (BK < 0))) else (div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value)), True)

thread_layout​

alias thread_layout = Layout.row_major(((div_s num_threads._mlir_value, 4) + -1) if ((num_threads < 0) & (((rem_s num_threads._mlir_value, 4) == 0) ^ True)) else (div_s num_threads._mlir_value, 4), 4)

tiler_layout​

alias tiler_layout = Layout.row_major(1, 0 if (simd_width_of[dtype]() == 0) else ((div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) + -1) if ((((rem_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value) == 0) ^ True) & ((simd_width_of[dtype]() < 0) ^ (BK < 0))) else (div_s BK._mlir_value, 1 if (simd_width_of[dtype]() == 0) else simd_width_of[dtype]()._mlir_value))

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), alignment=alignment, 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