Skip to main content

/

DocsCode

Mojo struct

MLAConfig

@register_passable(trivial) struct MLAConfig

Fields​

  • ​fa4_config (FA4Config):
  • ​MMA_M (Int):
  • ​BM (Int):
  • ​BN (Int):
  • ​BK0 (Int):
  • ​BK1 (Int):
  • ​depth (Int):
  • ​k_rope_depth (Int):
  • ​kv_depth (Int):
  • ​cache_depth (Int):
  • ​padded_depth (Int):
  • ​group (Int):
  • ​num_q_heads (Int):
  • ​num_kv_heads (Int):
  • ​TMEM_S1 (Int):
  • ​TMEM_O0 (Int):
  • ​TMEM_O1 (Int):
  • ​TMEM_P0 (Int):
  • ​TMEM_P1 (Int):
  • ​TMEM_C0 (Int):
  • ​TMEM_C1 (Int):
  • ​tmem_used (Int):
  • ​num_kv_stages (Int):
  • ​num_qk_stages (Int):
  • ​num_pv_stages (Int):
  • ​smem_used (Int):
  • ​dtype_size (Int):
  • ​split_m (Bool):
  • ​qkv_swizzle_mode (TensorMapSwizzle):
  • ​k_rope_swizzle_mode (TensorMapSwizzle):
  • ​output_swizzle_mode (TensorMapSwizzle):

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, RegisterPassable, TrivialRegisterPassable

comptime members​

__copy_ctor_is_trivial​

comptime __copy_ctor_is_trivial = True

__del__is_trivial​

comptime __del__is_trivial = True

__move_ctor_is_trivial​

comptime __move_ctor_is_trivial = True

mbar_size​

comptime mbar_size = size_of[DType.int64]()

num_correction_cols​

comptime num_correction_cols = 1

num_threads​

comptime num_threads = 512

sm100_smem_carveout​

comptime sm100_smem_carveout = (B200 - 1024)

sm100_tmem_cols​

comptime sm100_tmem_cols = 512

TMEM_S0​

comptime TMEM_S0 = 0

Methods​

__init__​

__init__(*, num_q_heads: Int, group: Int, depth: Int, qkv_dtype_size: Int, k_rope_dtype_size: Int, output_dtype_size: Int, page_size: Int) -> Self

num_qo​

num_qo(self) -> Int

Returns:

Int

supported​

supported(self) -> Bool

Returns:

Bool

correction_smem_elements​

correction_smem_elements(self) -> Int

Returns:

Int

num_active_warps_per_group​

num_active_warps_per_group(self) -> Int

Returns:

Int

num_active_threads_per_group​

num_active_threads_per_group(self) -> Int

Returns:

Int

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact