Mojo struct

MLASparseConfig

struct MLASparseConfig[qkv_dtype: DType]

Fields

num_q_heads (Int):
num_kv_heads (Int):
qk_depth (Int):
v_depth (Int):
top_k (Int):
group (Int):

Implemented traits

AnyType, ImplicitlyDestructible

`comptime` members

`B_TOPK`

comptime B_TOPK = 128

`cta_group`

comptime cta_group = 2

`k_swizzle_mode`

comptime k_swizzle_mode = TensorMapSwizzle.SWIZZLE_128B

`num_threads`

comptime num_threads = 512

`output_swizzle_mode`

comptime output_swizzle_mode = TensorMapSwizzle.SWIZZLE_128B

`q_smem_depth`

comptime q_smem_depth = 192

`q_swizzle_mode`

comptime q_swizzle_mode = TensorMapSwizzle.SWIZZLE_128B

`q_tmem_depth`

comptime q_tmem_depth = 384

`qkv_dtype_size`

comptime qkv_dtype_size = size_of[qkv_dtype]()

`sm100_tmem_cols`

comptime sm100_tmem_cols = 512

Methods

`init`

__init__(out self, *, num_q_heads: Int, num_kv_heads: Int, qk_depth: Int, v_depth: Int, top_k: Int, group: Int)

Fields​

Implemented traits​

comptime members​

B_TOPK​

cta_group​

k_swizzle_mode​

num_threads​

output_swizzle_mode​

q_smem_depth​

q_swizzle_mode​

q_tmem_depth​

qkv_dtype_size​

sm100_tmem_cols​

Methods​

__init__​