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

MLASmemStorage

struct MLASmemStorage[qkv_dtype: DType, rope_dtype: DType, num_mbars: Int, config: MLAConfig[config.qkv_dtype, rope_gmem_dtype=config.rope_gmem_dtype, rope_mma_dtype=config.rope_mma_dtype, scale_dtype_=config.scale_dtype_]]

Shared memory storage layout for the per-token-scale MLA prefill kernel.

Holds the Q (nope + rope), KV (nope + rope), q_scale, k_scale, correction, and barrier buffers sized from config for one CTA's worth of tile data.

Parameters

• ​qkv_dtype (DType): DType of the Q, K_nope, and V shared-memory buffers.
• ​rope_dtype (DType): DType of the Q_rope and K_rope shared-memory buffers.
• ​num_mbars (Int): Number of shared-memory barriers in the mbar_base array.
• ​config (MLAConfig[config.qkv_dtype, rope_gmem_dtype=config.rope_gmem_dtype, rope_mma_dtype=config.rope_mma_dtype, scale_dtype_=config.scale_dtype_]): MLAConfig supplying tile sizes, depths, and pipeline stage counts used to size each buffer.

Fields

• ​q_smem (Array[UInt8, Int((add (mul size_of[qkv_dtype](), config.BM, config.nope_depth), (mul size_of[rope_dtype](), config.BM, config.rope_depth)))]):
• ​kv_smem (Array[UInt8, Int((add (mul config.fa4_config.shared_kv_cols(), size_of[qkv_dtype](), config.BN, config.num_kv_stages, config.num_qk_stages), (mul size_of[rope_dtype](), config.BN, config.num_kv_stages, config.num_qk_stages, config.rope_depth)))]):
• ​q_scale_smem (Array[UInt8, Int((mul size_of[DType.float32](), config.BM))]):
• ​k_scale_smem (Array[UInt8, Int((mul size_of[DType.float32](), config.BN))]):
• ​correction_smem (Array[Float32, config.correction_smem_elements()]):
• ​mbar_base (Array[SharedMemBarrier, num_mbars]):
• ​tmem_addr (Array[UInt32, Int(1)]):

Implemented traits

AnyType, ImplicitlyDeletable

comptime members

correction_smem_size

comptime correction_smem_size = config.correction_smem_elements()

k_scale_bytes

comptime k_scale_bytes = (config * size_of[DType.float32]())

kv_bytes

comptime kv_bytes = (Int((mul config.fa4_config.shared_kv_cols(), size_of[qkv_dtype](), config.BN, config.num_kv_stages, config.num_qk_stages)) + Int((mul size_of[rope_dtype](), config.BN, config.num_kv_stages, config.num_qk_stages, config.rope_depth)))

kv_nope_bytes

comptime kv_nope_bytes = (Int((mul config.fa4_config.shared_kv_cols(), size_of[qkv_dtype](), config.BN)) * Int((mul config.num_kv_stages, config.num_qk_stages)))

kv_rope_bytes

comptime kv_rope_bytes = (Int((mul size_of[rope_dtype](), config.BN, config.rope_depth)) * Int((mul config.num_kv_stages, config.num_qk_stages)))

num_kv_stages

comptime num_kv_stages = (config * config)

q_bytes

comptime q_bytes = (Int((mul size_of[qkv_dtype](), config.BM, config.nope_depth)) + Int((mul size_of[rope_dtype](), config.BM, config.rope_depth)))

q_nope_bytes

comptime q_nope_bytes = (Int((mul config.BM, config.nope_depth)) * size_of[qkv_dtype]())

q_rope_bytes

comptime q_rope_bytes = (Int((mul config.BM, config.rope_depth)) * size_of[rope_dtype]())

q_scale_bytes

comptime q_scale_bytes = (config * size_of[DType.float32]())

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit