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
MLASparseSharedMemoryFP8
struct MLASparseSharedMemoryFP8[config: MLASparseConfig[config.qkv_dtype, config.b_topk_, config.num_mbars_, config.q_smem_depth_, config.q_tmem_depth_, config.cta_group_], scale_block_size: Int]
Fields
- base (
MLASparseSharedMemory[config]): - k_scales (
Array[Float32, Int((mul (b_topk_ // cta_group_), ceildiv(config.qk_depth, scale_block_size) if (xor identical((xor (lt scale_block_size, 1), True), False), True) else Int(1)))]): - v_scales (
Array[Float32, Int((mul ceildiv(config.v_depth, scale_block_size) if (xor identical((xor (lt scale_block_size, 1), True), False), True) else Int(1), b_topk_))]): - k_fp8_tma_done (
Array[SharedMemBarrier, Int(2)]): - v_fp8_tma_done (
Array[SharedMemBarrier, Int(2)]):
Implemented traits
comptime members
K_scales_per_token
comptime K_scales_per_token = ceildiv(config.qk_depth, scale_block_size) if (scale_block_size > Int(0)) else Int(1)
K_SCALES_SIZE
comptime K_SCALES_SIZE = (MLASparseSharedMemoryFP8[config, scale_block_size].TOPK_PER_CTA * ceildiv(config.qk_depth, scale_block_size) if (xor identical((xor (lt scale_block_size, 1), True), False), True) else Int(1))
num_mbars
comptime num_mbars = 2
TOPK_PER_CTA
comptime TOPK_PER_CTA = (config.B_TOPK // config.cta_group)
V_scales_per_token
comptime V_scales_per_token = ceildiv(config.v_depth, scale_block_size) if (scale_block_size > Int(0)) else Int(1)
V_SCALES_SIZE
comptime V_SCALES_SIZE = (config.B_TOPK * ceildiv(config.v_depth, scale_block_size) if (xor identical((xor (lt scale_block_size, 1), True), False), True) else Int(1))