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

MLAPrefillSparseQKVFP8

struct MLAPrefillSparseQKVFP8[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​

comptime B_TOPK = config.B_TOPK

B_TOPK_PER_CTA​

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

Common​

comptime Common = MLAPrefillSparseCommon[KVLUTType, output_dtype, config]

cta_group​

comptime cta_group = config.cta_group

CTA_MASK​

comptime CTA_MASK = UInt16(3) if MLAPrefillSparseQKVFP8[KVLUTType, output_dtype, config] else UInt16(1)

is_cg2​

comptime is_cg2 = (config.cta_group == Int(2))

kv_desc_shape​

comptime kv_desc_shape = Index[Int, Int](Int(1), _gather4_box_width[DType.float8_e4m3fn, config.qk_depth, TensorMapSwizzle(Int32(2))]())

kv_gather_box​

comptime kv_gather_box = _gather4_box_width[DType.float8_e4m3fn, config.qk_depth, SW64]()

kv_tile_shape​

comptime kv_tile_shape = Index[Int, Int](MLAPrefillSparseQKVFP8[KVLUTType, output_dtype, config].B_TOPK, _gather4_box_width[DType.float8_e4m3fn, config.qk_depth, TensorMapSwizzle(Int32(2))]())

NUM_PV_KHALVES​

comptime NUM_PV_KHALVES = (MLAPrefillSparseQKVFP8[KVLUTType, output_dtype, config].B_TOPK // Int(64))

NUM_Q_HEADS_PER_CTA​

comptime NUM_Q_HEADS_PER_CTA = (config // config.cta_group)

NUM_S_SLOTS​

comptime NUM_S_SLOTS = MLAPrefillSparseQKVFP8[KVLUTType, output_dtype, config].SMemType.NUM_S_SLOTS

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)))

P_FP8_BIAS​

comptime P_FP8_BIAS = SIMD(-2).__radd__(SIMD(8))

P_KHALF_STRIDE​

comptime P_KHALF_STRIDE = ((config // cta_group_) * Int(64))

PADDED_HEADS​

comptime PADDED_HEADS = config.padded_num_q_heads

PADDED_HEADS_PER_CTA​

comptime PADDED_HEADS_PER_CTA = (config // config.cta_group)

PV_BK​

comptime PV_BK = 64

PVAcc​

comptime PVAcc = SM100TensorAccumulator[DType.float8_e4m3fn, DType.float32, config.padded_num_q_heads, (config // Int(2)), Int(64), a_tmem=False, mma_kind=UMMAKind.KIND_F8F6F4, swizzle_a=SW64, swizzle_b=SW64, transpose_b=False, cta_group=config.cta_group]

q_desc_shape​

comptime q_desc_shape = _default_desc_shape[Int(3), DType.float8_e4m3fn, MLAPrefillSparseQKVFP8[KVLUTType, output_dtype, config].q_tile_shape, SW64]()

Q_SWIZZLE_COLS​

comptime Q_SWIZZLE_COLS = MLAPrefillSparseQKVFP8[KVLUTType, output_dtype, config].q_desc_shape[Int(2)]

q_tile_shape​

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

qk_depth​

comptime qk_depth = config.qk_depth

QKAcc​

comptime QKAcc = SM100TensorAccumulator[DType.float8_e4m3fn, DType.float32, config.padded_num_q_heads, MLAPrefillSparseQKVFP8[KVLUTType, output_dtype, config].B_TOPK, config.qk_depth, a_tmem=False, mma_kind=UMMAKind.KIND_F8F6F4, swizzle_a=SW64, swizzle_b=SW64, cta_group=config.cta_group]

RESCALE_THRESHOLD​

comptime RESCALE_THRESHOLD = SIMD(-2)

S_SLOT_STRIDE​

comptime S_SLOT_STRIDE = MLAPrefillSparseQKVFP8[KVLUTType, output_dtype, config].SMemType.S_SLOT_STRIDE

S_TMEM_BASE​

comptime S_TMEM_BASE = MLAPrefillSparseQKVFP8[KVLUTType, output_dtype, config].SMemType.S_TMEM_BASE

SMemType​

comptime SMemType = MLASparseSharedMemoryQKVFP8[config]

SV_ATOM_MMA_N​

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

V_ATOM_BYTE_OFFSET​

comptime V_ATOM_BYTE_OFFSET = (Int((mul ((config // Int(2)) // cta_group_), 64)) * size_of[DType.float8_e4m3fn]())

V_BMN_PER_ATOM​

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

v_depth​

comptime v_depth = config.v_depth

V_KHALF_STRIDE​

comptime V_KHALF_STRIDE = (Int((mul ((config // Int(2)) // cta_group_), 2)) * Int(64))

V_SMEM_COLS_PER_CTA​

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

Methods​

kernel​

static def kernel[TopKLengthLayout: TensorLayout, IndicesLayout: TensorLayout](q_tma_op: TMATensorTile[DType.float8_e4m3fn, Int(3), Self.q_tile_shape, Self.q_desc_shape], kv_tma_op: TMATensorTile[DType.float8_e4m3fn, Int(2), Self.kv_tile_shape, Self.kv_desc_shape], o_tma_op: TMATensorTile[output_dtype, Int(2), Self.o_tile_shape, Self.o_desc_shape], topk_lengths: TileTensor[DType.uint32, TopKLengthLayout, MutAnyOrigin], indices: TileTensor[DType.uint32, IndicesLayout, MutAnyOrigin], kv_lut: KVLUTType, scale: Float32, attn_sink_ptr: Optional[Pointer[Float32, ImmutAnyOrigin]], indices_stride: Int32) where (Int(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(IndicesLayout.__shape_types), [idx: __mlir_type.index] IndicesLayout.__shape_types[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(IndicesLayout.__shape_types), [idx: __mlir_type.index] IndicesLayout.__shape_types[idx]._ParamListType))[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(IndicesLayout.__shape_types), [idx: __mlir_type.index] IndicesLayout.__shape_types[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(IndicesLayout.__shape_types), [idx: __mlir_type.index] IndicesLayout.__shape_types[idx]._ParamListType))[idx]._ParamListType))[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(IndicesLayout.__shape_types), [idx: __mlir_type.index] IndicesLayout.__shape_types[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(IndicesLayout.__shape_types), [idx: __mlir_type.index] IndicesLayout.__shape_types[idx]._ParamListType))[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(IndicesLayout.__shape_types), [idx: __mlir_type.index] IndicesLayout.__shape_types[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(IndicesLayout.__shape_types), [idx: __mlir_type.index] IndicesLayout.__shape_types[idx]._ParamListType))[idx]._ParamListType))[idx]._ParamListType))[idx]._ParamListType)))) == Int(1)) if (Int(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))[idx]._ParamListType))[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))[idx]._ParamListType))[idx]._ParamListType))[idx]._ParamListType)))) == Int(1)) else (Int(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))[idx]._ParamListType))[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(#kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))), [idx: __mlir_type.index] #kgen.param_list.concat(#kgen.param_list.tabulate(len(TopKLengthLayout.__shape_types), [idx: __mlir_type.index] TopKLengthLayout.__shape_types[idx]._ParamListType))[idx]._ParamListType))[idx]._ParamListType))[idx]._ParamListType)))) == Int(1))

Models

Cookbook

Puzzles

Blog

Community

Contact