IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

MLAPrefillSparseCommon

struct MLAPrefillSparseCommon[KVLUTType: MHAOperand, output_dtype: DType, config: MLASparseConfig[config.qkv_dtype, config.b_topk_, config.num_mbars_, config.q_smem_depth_, config.q_tmem_depth_, config.cta_group_]]

Implemented traits​

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members​

accum_dtype​

comptime accum_dtype = DType.float32

B_TOPK_PER_CTA​

comptime B_TOPK_PER_CTA = (config.B_TOPK // config.cta_group)

CTA_MASK​

comptime CTA_MASK = UInt16(3) if (config.cta_group == Int(2)) else UInt16(1)

FP8_K_SWIZZLE​

comptime FP8_K_SWIZZLE = TensorMapSwizzle.SWIZZLE_NONE

FP8_V_SWIZZLE​

comptime FP8_V_SWIZZLE = TensorMapSwizzle.SWIZZLE_NONE

FULL_Q_TYPE​

comptime FULL_Q_TYPE = MLAPrefillSparseCommon[KVLUTType, output_dtype, config].SMemType.FULL_Q_TYPE

k_desc_shape​

comptime k_desc_shape = Index[Int, Int](Int(1), _gather4_box_width[qkv_dtype, config.qk_depth, TensorMapSwizzle(Int32(3))]())

k_gather_box​

comptime k_gather_box = _gather4_box_width[MLAPrefillSparseCommon[KVLUTType, output_dtype, config].qkv_dtype, config.qk_depth, MLAPrefillSparseCommon[KVLUTType, output_dtype, config].k_swizzle_mode]()

k_swizzle_mode​

comptime k_swizzle_mode = config.k_swizzle_mode

k_tile_height​

comptime k_tile_height = MLAPrefillSparseCommon[KVLUTType, output_dtype, config].B_TOPK_PER_CTA

k_tile_shape​

comptime k_tile_shape = Index[Int, Int](MLAPrefillSparseCommon[KVLUTType, output_dtype, config].k_tile_height, _gather4_box_width[qkv_dtype, config.qk_depth, TensorMapSwizzle(Int32(3))]())

k_tile_width​

comptime k_tile_width = config.qk_depth

k_tma_desc_shape_fp8​

comptime k_tma_desc_shape_fp8 = Index[Int, Int](Int(1), _gather4_box_width[DType.int64, (config // Int(8)), TensorMapSwizzle(Int32(0))]())

k_tma_dtype_fp8​

comptime k_tma_dtype_fp8 = DType.int64

k_tma_gather_box_fp8​

comptime k_tma_gather_box_fp8 = _gather4_box_width[DType.int64, (config // Int(8)), MLAPrefillSparseCommon[KVLUTType, output_dtype, config].k_tma_swizzle_fp8]()

k_tma_swizzle_fp8​

comptime k_tma_swizzle_fp8 = MLAPrefillSparseCommon[KVLUTType, output_dtype, config].FP8_K_SWIZZLE

k_tma_tile_shape_fp8​

comptime k_tma_tile_shape_fp8 = Index[Int, Int](MLAPrefillSparseCommon[KVLUTType, output_dtype, config].k_tile_height, _gather4_box_width[DType.int64, (config // Int(8)), TensorMapSwizzle(Int32(0))]())

k_tma_tile_width_fp8​

comptime k_tma_tile_width_fp8 = (config // Int(8))

NUM_Q_HEADS_PER_CTA​

comptime NUM_Q_HEADS_PER_CTA = (config // config.cta_group)

NUM_SV_ATOMS​

comptime NUM_SV_ATOMS = 2

O_ATOM_PHYS_COLS​

comptime O_ATOM_PHYS_COLS = ((config // Int(2)) // Int(2))

o_desc_shape​

comptime o_desc_shape = Index[Int, Int]((config // cta_group_), Int(64))

o_tile_shape​

comptime o_tile_shape = Index[Int, Int]((config // cta_group_), config)

O_TMEM_ADDR​

comptime O_TMEM_ADDR = 0

O_TMEM_ADDR_ATOM2​

comptime O_TMEM_ADDR_ATOM2 = (Int(0) + ((config // Int(2)) // Int(2)))

O_TYPE​

comptime O_TYPE = MLAPrefillSparseCommon[KVLUTType, output_dtype, config].SMemType.O_TYPE

P_TMEM_ADDR​

comptime P_TMEM_ADDR = 256

PADDED_HEADS_PER_CTA​

comptime PADDED_HEADS_PER_CTA = (config // config.cta_group)

q_desc_shape​

comptime q_desc_shape = _default_desc_shape[Int(3), MLAPrefillSparseCommon[KVLUTType, output_dtype, config].qkv_dtype, MLAPrefillSparseCommon[KVLUTType, output_dtype, config].q_tile_shape, config.q_swizzle_mode]()

q_smem_depth​

comptime q_smem_depth = config.q_smem_depth

q_tile_shape​

comptime q_tile_shape = Index[Int, Int, Int](Int(1), (config // cta_group_), config)

Q_TMEM_ADDR​

comptime Q_TMEM_ADDR = (Int(512) - (MLAPrefillSparseCommon[KVLUTType, output_dtype, config].q_tmem_depth // Int(2)))

q_tmem_depth​

comptime q_tmem_depth = config.q_tmem_depth

QKMMAOpType​

comptime QKMMAOpType = QKMMAOp[MLAPrefillSparseCommon[KVLUTType, output_dtype, config].qkv_dtype, DType.float32, config]

qkv_dtype​

comptime qkv_dtype

qkv_dtype_size​

comptime qkv_dtype_size = size_of[MLAPrefillSparseCommon[KVLUTType, output_dtype, config].qkv_dtype]()

SHARED_QKV_TYPE​

comptime SHARED_QKV_TYPE = MLAPrefillSparseCommon[KVLUTType, output_dtype, config].SMemType.SHARED_QKV_TYPE

SMemType​

comptime SMemType = MLASparseSharedMemory[config]

SV_ATOM_MMA_N​

comptime SV_ATOM_MMA_N = (config // Int(2))

SVMMAType​

comptime SVMMAType = SVMMAType[MLAPrefillSparseCommon[KVLUTType, output_dtype, config].qkv_dtype, DType.float32, config]

V_BMN_PER_ATOM​

comptime V_BMN_PER_ATOM = ((config // Int(2)) // config.cta_group)

V_DEPTH_PER_CTA​

comptime V_DEPTH_PER_CTA = MLAPrefillSparseCommon[KVLUTType, output_dtype, config].SV_ATOM_MMA_N

v_desc_shape​

comptime v_desc_shape = Index[Int, Int](Int(1), _gather4_box_width[qkv_dtype, Int((mul ((config // Int(2)) // cta_group_), 2)), TensorMapSwizzle(Int32(3))]())

v_gather_box​

comptime v_gather_box = _gather4_box_width[MLAPrefillSparseCommon[KVLUTType, output_dtype, config].qkv_dtype, Int((mul ((config // Int(2)) // cta_group_), 2)), MLAPrefillSparseCommon[KVLUTType, output_dtype, config].v_swizzle_mode]()

V_SMEM_COLS_PER_CTA​

comptime V_SMEM_COLS_PER_CTA = (((config // Int(2)) // cta_group_) * Int(2))

v_swizzle_mode​

comptime v_swizzle_mode = TensorMapSwizzle.SWIZZLE_128B

v_tile_height​

comptime v_tile_height = (config.B_TOPK // Int(2))

v_tile_shape​

comptime v_tile_shape = Index[Int, Int](MLAPrefillSparseCommon[KVLUTType, output_dtype, config].v_tile_height, _gather4_box_width[qkv_dtype, Int((mul ((config // Int(2)) // cta_group_), 2)), TensorMapSwizzle(Int32(3))]())

v_tile_width​

comptime v_tile_width = MLAPrefillSparseCommon[KVLUTType, output_dtype, config].V_SMEM_COLS_PER_CTA

v_tma_desc_shape_fp8​

comptime v_tma_desc_shape_fp8 = Index[Int, Int](Int(1), _gather4_box_width[DType.int64, (((config // Int(2)) // cta_group_) // Int(8)), TensorMapSwizzle(Int32(0))]())

v_tma_dtype_fp8​

comptime v_tma_dtype_fp8 = DType.int64

v_tma_gather_box_fp8​

comptime v_tma_gather_box_fp8 = _gather4_box_width[DType.int64, (((config // Int(2)) // cta_group_) // Int(8)), MLAPrefillSparseCommon[KVLUTType, output_dtype, config].v_tma_swizzle_fp8]()

v_tma_swizzle_fp8​

comptime v_tma_swizzle_fp8 = MLAPrefillSparseCommon[KVLUTType, output_dtype, config].FP8_V_SWIZZLE

v_tma_tile_height_fp8​

comptime v_tma_tile_height_fp8 = config.B_TOPK

v_tma_tile_shape_fp8​

comptime v_tma_tile_shape_fp8 = Index[Int, Int](MLAPrefillSparseCommon[KVLUTType, output_dtype, config].v_tma_tile_height_fp8, _gather4_box_width[DType.int64, (((config // Int(2)) // cta_group_) // Int(8)), TensorMapSwizzle(Int32(0))]())

v_tma_tile_width_fp8​

comptime v_tma_tile_width_fp8 = (((config // Int(2)) // cta_group_) // Int(8))

Methods​

mma​

static def mma[fp8_active: Bool = False](q_smem_ptr: Pointer[Scalar[Self.qkv_dtype], address_space=AddressSpace.SHARED], k_smem_ptr: Pointer[Scalar[Self.qkv_dtype], address_space=AddressSpace.SHARED], s_smem_ptr: Pointer[Scalar[Self.qkv_dtype], address_space=AddressSpace.SHARED], v_smem_ptr: Pointer[Scalar[Self.qkv_dtype], address_space=AddressSpace.SHARED], k_p0_ready: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED], k_p1_ready: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED], v_p0_ready: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED], v_p1_ready: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED], sv_p0_done: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED], sv_p1_done: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED], so_ready: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED], p_free: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED], qk_ss_done: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED], qk_ts_done: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED], k: UInt32, num_k_blocks: UInt32)

cp_q_from_smem_to_tmem​

static def cp_q_from_smem_to_tmem(smem_desc: MMASmemDescriptorPair, tmem_addr: UInt32)

kv_valid_producer​

static def kv_valid_producer(indices: TileTensor[DType.uint32, Storage=indices.Storage, linear_idx_type=indices.linear_idx_type], is_k_valid_ptr: Pointer[UInt8, address_space=AddressSpace.SHARED], k_valid_ready_ptr: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED], k_valid_free_ptr: Pointer[SharedMemBarrier, address_space=AddressSpace.SHARED], lane_idx: UInt32, indices_base: UInt32, num_kv_rows: Int32, top_k_length: Int32, num_k_blocks: Int)

Models

Cookbook

Puzzles

Blog

Community

Contact