Mojo struct
MHAAttentionConfig
struct MHAAttentionConfig[dtype: DType, //, token_gen: Bool, config: MHAConfig[dtype], group: Int]
Implemented traits
AnyType,
AttentionConfig,
Copyable,
ImplicitlyCopyable,
UnknownDestructibility
Aliases
__copyinit__is_trivial
alias __copyinit__is_trivial = True
__del__is_trivial
alias __del__is_trivial = True
depth_padded
alias depth_padded = False if MHAAttentionConfig[token_gen, config, group].USE_EXPERIMENTAL_CDNA4_MHA_KERNEL else True
double_buffer
alias double_buffer = True if MHAAttentionConfig[token_gen, config, group].USE_EXPERIMENTAL_CDNA4_MHA_KERNEL else False
full_kv
alias full_kv = True if MHAAttentionConfig[token_gen, config, group].USE_EXPERIMENTAL_CDNA4_MHA_KERNEL else False
shared_kv
alias shared_kv = False if MHAAttentionConfig[token_gen, config, group].USE_EXPERIMENTAL_CDNA4_MHA_KERNEL else True
USE_EXPERIMENTAL_CDNA4_MHA_KERNEL
alias USE_EXPERIMENTAL_CDNA4_MHA_KERNEL = token_gen.__invert__() if env_get_bool["USE_EXPERIMENTAL_CDNA4_MHA_KERNEL", False]() if _cdna_4_or_newer() else _cdna_4_or_newer() else env_get_bool["USE_EXPERIMENTAL_CDNA4_MHA_KERNEL", False]() if _cdna_4_or_newer() else _cdna_4_or_newer()
Methods
q_head_idx
q_tile_idx
kv_head_idx
get_mma_shape
get_q_offset
get_output_offset
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!