IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/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. /max/get-started.md).

Mojo struct

MLASparseSharedMemory

struct MLASparseSharedMemory[config: MLASparseConfig[config.qkv_dtype, config.b_topk_, config.num_mbars_, config.q_smem_depth_, config.q_tmem_depth_, config.cta_group_]]

Fields​

  • ​qkvo_union (UnsafeUnion[Array[Scalar[MLASparseSharedMemory[config].qkv_dtype], Int((mul (config // cta_group_), config.qk_depth))], Array[Scalar[MLASparseSharedMemory[config].qkv_dtype], Int((add (mul ((config // Int(2)) // cta_group_), b_topk_, 2), (mul (b_topk_ // cta_group_), config.qk_depth), (mul (config // cta_group_), q_smem_depth_)))], Array[Scalar[MLASparseSharedMemory[config].qkv_dtype], Int((mul (config // cta_group_), config.v_depth))]]):
  • ​scores (Array[Scalar[MLASparseSharedMemory[config].qkv_dtype], Int((mul (config // cta_group_), b_topk_))]):
  • ​p (Array[Float32, Int((mul (config // cta_group_), b_topk_))]):
  • ​prologue_q (Array[SharedMemBarrier, Int(1)]):
  • ​prologue_q_cp (Array[SharedMemBarrier, Int(1)]):
  • ​qk_ss_done (Array[SharedMemBarrier, MLASparseSharedMemory[config].num_mbars]):
  • ​qk_ts_done (Array[SharedMemBarrier, MLASparseSharedMemory[config].num_mbars]):
  • ​sv_p0_done (Array[SharedMemBarrier, MLASparseSharedMemory[config].num_mbars]):
  • ​sv_p1_done (Array[SharedMemBarrier, MLASparseSharedMemory[config].num_mbars]):
  • ​k_p0_ready (Array[SharedMemBarrier, MLASparseSharedMemory[config].num_mbars]):
  • ​k_p1_ready (Array[SharedMemBarrier, MLASparseSharedMemory[config].num_mbars]):
  • ​v_p0_ready (Array[SharedMemBarrier, MLASparseSharedMemory[config].num_mbars]):
  • ​v_p1_ready (Array[SharedMemBarrier, MLASparseSharedMemory[config].num_mbars]):
  • ​p_free (Array[SharedMemBarrier, MLASparseSharedMemory[config].num_mbars]):
  • ​so_ready (Array[SharedMemBarrier, MLASparseSharedMemory[config].num_mbars]):
  • ​k_valid_ready (Array[SharedMemBarrier, MLASparseSharedMemory[config].num_mbars]):
  • ​k_valid_free (Array[SharedMemBarrier, MLASparseSharedMemory[config].num_mbars]):
  • ​is_k_valid (Array[UInt8, (MLASparseSharedMemory[config].num_mbars * MLASparseSharedMemory[config].MASK_BYTES_PER_BUF)]):
  • ​tmem_addr (Array[UInt32, Int(1)]):
  • ​rowwise_max (Array[Float32, _resolve_warpgroup_size()]):
  • ​rowwise_sum (Array[Float32, _resolve_warpgroup_size()]):

Implemented traits​

AnyType, ImplicitlyDeletable

comptime members​

BH​

comptime BH = (config // config.cta_group)

FULL_Q_SIZE​

comptime FULL_Q_SIZE = ((config // cta_group_) * config)

FULL_Q_TYPE​

comptime FULL_Q_TYPE = Array[Scalar[MLASparseSharedMemory[config].qkv_dtype], Int((mul (config // cta_group_), config.qk_depth))]

INDICES_PER_LANE​

comptime INDICES_PER_LANE = 8

K_SIZE​

comptime K_SIZE = (MLASparseSharedMemory[config].TOPK_PER_CTA * config)

MASK_BYTES_PER_BUF​

comptime MASK_BYTES_PER_BUF = (config.B_TOPK // Int(8))

NUM_KV_VALID_LANES​

comptime NUM_KV_VALID_LANES = MLASparseSharedMemory[config].MASK_BYTES_PER_BUF

num_mbars​

comptime num_mbars = config.num_mbars

num_q_heads​

comptime num_q_heads = config.num_q_heads

NUM_SV_ATOMS​

comptime NUM_SV_ATOMS = 2

O_ATOM_PHYS_COLS​

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

O_SIZE​

comptime O_SIZE = ((config // cta_group_) * config)

O_TYPE​

comptime O_TYPE = Array[Scalar[MLASparseSharedMemory[config].qkv_dtype], Int((mul (config // cta_group_), config.v_depth))]

qk_depth​

comptime qk_depth = config.qk_depth

qkv_dtype​

comptime qkv_dtype

S_SIZE​

comptime S_SIZE = ((config // cta_group_) * config.B_TOPK)

SHARED_Q_SIZE​

comptime SHARED_Q_SIZE = ((config // cta_group_) * config.q_smem_depth)

SHARED_QKV_SIZE​

comptime SHARED_QKV_SIZE = (Int((add (mul (b_topk_ // cta_group_), config.qk_depth), (mul (config // cta_group_), q_smem_depth_))) + Int((mul ((config // Int(2)) // cta_group_), b_topk_, 2)))

SHARED_QKV_TYPE​

comptime SHARED_QKV_TYPE = Array[Scalar[MLASparseSharedMemory[config].qkv_dtype], Int((add (mul ((config // Int(2)) // cta_group_), b_topk_, 2), (mul (b_topk_ // cta_group_), config.qk_depth), (mul (config // cta_group_), q_smem_depth_)))]

SV_ATOM_MMA_N​

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

TOPK_PER_CTA​

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

V_BMN_PER_ATOM​

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

V_DEPTH_PER_CTA​

comptime V_DEPTH_PER_CTA = MLASparseSharedMemory[config].SV_ATOM_MMA_N

V_SIZE​

comptime V_SIZE = (config.B_TOPK * Int((mul ((config // Int(2)) // cta_group_), 2)))

V_SMEM_COLS_PER_CTA​

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

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact