Skip to main content

Mojo struct

MLAKVLayouts

struct MLAKVLayouts[k_nope_dtype: DType, k_rope_dtype: DType, kv_scale_dtype: DType, config: MLAConfig[config.qkv_dtype, rope_gmem_dtype=config.rope_gmem_dtype, rope_mma_dtype=config.rope_mma_dtype, scale_dtype=config.scale_dtype]]

Comptime layout and size metadata for MLA K/V tiles.

Implemented traits​

AnyType, ImplicitlyDestructible

comptime members​

k_bytes​

comptime k_bytes = (MLAKVLayouts[k_nope_dtype, k_rope_dtype, kv_scale_dtype, config].k_nope_bytes + MLAKVLayouts[k_nope_dtype, k_rope_dtype, kv_scale_dtype, config].k_rope_bytes)

k_elements​

comptime k_elements = (MLAKVLayouts[k_nope_dtype, k_rope_dtype, kv_scale_dtype, config].k_nope_tma_layout + MLAKVLayouts[k_nope_dtype, k_rope_dtype, kv_scale_dtype, config].k_rope_tma_layout)

k_nope_bytes​

comptime k_nope_bytes = (MLAKVLayouts[k_nope_dtype, k_rope_dtype, kv_scale_dtype, config].k_nope_tma_layout * size_of[k_nope_dtype]())

k_nope_tma_layout​

comptime k_nope_tma_layout = Layout[*?, *?].static_product

k_rope_bytes​

comptime k_rope_bytes = (MLAKVLayouts[k_nope_dtype, k_rope_dtype, kv_scale_dtype, config].k_rope_tma_layout * size_of[k_rope_dtype]())

k_rope_tma_layout​

comptime k_rope_tma_layout = Layout[*?, *?].static_product

k_tma_layout​

comptime k_tma_layout = Layout[*?, *?].static_product

KPairType​

comptime KPairType = TMADestination[k_nope_dtype, MLAKVLayouts[k_nope_dtype, k_rope_dtype, kv_scale_dtype, config].k_tma_layout]

SMemType​

comptime SMemType = UnsafePointer[Scalar[k_nope_dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]

v_bytes​

comptime v_bytes = (MLAKVLayouts[k_nope_dtype, k_rope_dtype, kv_scale_dtype, config].v_elements * size_of[k_nope_dtype]())

v_elements​

comptime v_elements = MLAKVLayouts[k_nope_dtype, k_rope_dtype, kv_scale_dtype, config].v_tma_layout

v_tma_layout​

comptime v_tma_layout = Layout[*?, *?].static_product

VPairType​

comptime VPairType = TMADestination[k_nope_dtype, MLAKVLayouts[k_nope_dtype, k_rope_dtype, kv_scale_dtype, config].v_tma_layout]

Models

Cookbook

Puzzles

Blog

Community

Contact