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
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,
ImplicitlyCopyable,
ImplicitlyDeletable,
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, _safe=False]], 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))
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!