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

MhaPrefillV2

struct MhaPrefillV2[config: MhaConfigV2]

8-warp MHA forward kernel parameterized by MhaConfigV2.

Each block runs config.num_warps wave64 warps that share K/V SMEM via cooperative DMA. Warp w owns Q rows [w * q_block_size, (w + 1) * q_block_size) of the block's stripe and carries its own register-resident attention state.

Parameters

config (MhaConfigV2): Shape configuration (MhaConfigV2).

Implemented traits

AnyType, Deinitable, Movable

`comptime` members

`BM`

comptime BM = (config * config)

`D_FRAG_PER_LANE`

comptime D_FRAG_PER_LANE = (Int((mul config.depth, config.q_block_size)) // Int(64))

`DEPTH`

comptime DEPTH = config.depth

`k_swizzle`

comptime k_swizzle = Optional(Swizzle(Int(1), Int(0), Int(4)))

`k_swizzle2`

comptime k_swizzle2 = Optional(Swizzle(Int(1), Int(1), Int(4)))

`KTileLoader`

comptime KTileLoader = SubTileLoaderLDS[config.dtype, MhaPrefillV2[config].k_swizzle, MhaPrefillV2[config].k_swizzle2]

`KV_BLOCK`

comptime KV_BLOCK = config.kv_block

`NUM_HEADS`

comptime NUM_HEADS = config.num_heads

`NUM_KV_HEADS`

comptime NUM_KV_HEADS = config.num_kv_heads

`NUM_THREADS`

comptime NUM_THREADS = (config * Int(64))

`NUM_WARPS`

comptime NUM_WARPS = config.num_warps

`prescale_q`

comptime prescale_q = not config.dtype.is_float8().__bool__()

`Q_BLOCK_SIZE`

comptime Q_BLOCK_SIZE = config.q_block_size

`RESCALE_THRESHOLD`

comptime RESCALE_THRESHOLD = config.rescale_threshold

`v_swizzle`

comptime v_swizzle = Optional(None)

`VTileLoader`

comptime VTileLoader = SubTileLoaderLDS_st_8x32[config.dtype, config.kv_block, config.depth, Int(64) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> #lit.struct.extract<:!lit.struct<@nn::@attention::@gpu::@amd_structured::@mha_mma_op::@MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> #lit.struct.extract<:!lit.struct<@nn::@attention::@gpu::@amd_structured::@mha_mma_op::@MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> #lit.struct.extract<:!lit.struct<@nn::@attention::@gpu::@amd_structured::@mha_mma_op::@MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> #lit.struct.extract<:!lit.struct<@nn::@attention::@gpu::@amd_structured::@mha_mma_op::@MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> #lit.struct.extract<:!lit.struct<@nn::@attention::@gpu::@amd_structured::@mha_mma_op::@MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> #lit.struct.extract<:!lit.struct<@nn::@attention::@gpu::@amd_structured::@mha_mma_op::@MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<@std::@builtin::@dtype::@DType> #lit.struct.extract<:!lit.struct<@nn::@attention::@gpu::@amd_structured::@mha_mma_op::@MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) else Int(32), Int((mul config.num_warps, 64))]

Methods

`load_q`

static def load_q[layout: TensorLayout](q_warp_2d: TileTensor[config.dtype, layout, Engine=q_warp_2d.Engine, address_space=q_warp_2d.address_space, linear_idx_type=q_warp_2d.linear_idx_type]) -> TileTensor[config.dtype, Layout[TypeList[ComptimeInt[(config // Int(16) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) and config.fp8_mma_k_128 else Int(32))], ComptimeInt[(config // Int(128) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) and config.fp8_mma_k_128 else Int(64) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) else Int(16))], ComptimeInt[(Int((mul Int(128) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) and config.fp8_mma_k_128 else Int(64) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) else Int(16), Int(16) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) and config.fp8_mma_k_128 else Int(32))) // Int(64))]](), TypeList[ComptimeInt[Int((mul (Int((mul Int(128) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) and config.fp8_mma_k_128 else Int(64) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) else Int(16), Int(16) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) and config.fp8_mma_k_128 else Int(32))) // Int(64)), (config // Int(128) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) and config.fp8_mma_k_128 else Int(64) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) else Int(16))))], ComptimeInt[(Int((mul Int(128) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) and config.fp8_mma_k_128 else Int(64) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) else Int(16), Int(16) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) if identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 73) else identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 74) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 75) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 76) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 77) or identical(#pop.dtype_to_ui8<#lit.struct.extract<:!lit.struct<_std::_builtin::_dtype::_DType> #lit.struct.extract<:!lit.struct<_nn::_attention::_gpu::_amd_structured::_mha_mma_op::_MhaConfigV2> config, "dtype">, "_mlir_value">>, 78) and config.fp8_mma_k_128 else Int(32))) // Int(64))], ComptimeInt[Int(1)]]()], MutUntrackedOrigin, address_space=AddressSpace.LOCAL]

Loads the warp's Q sub-tile from gmem into a row_l register tile via RegTileLoader.

BF16 (d=128, MMA_K=16): 8 K-tiles × 1 buffer_load_bf16x8 per lane per K-tile = 8 loads × 16 B each. Per-lane fragment = 8 BF16 = 16 B fits in one buffer_load.

FP8 (d=128, MMA_K=64): 2 K-tiles, but each base tile per lane is 32 FP8 = 32 B which exceeds the 16-B buffer_load_lds max. Splits each K-tile load into 2 × 16-elt halves (16 B each) targeting the first / second half of the destination cell.

Parameters:

layout (TensorLayout): Layout of q_warp_2d (inferred).

Args:

q_warp_2d (TileTensor[config.dtype, layout, Engine=q_warp_2d.Engine, address_space=q_warp_2d.address_space, linear_idx_type=q_warp_2d.linear_idx_type]): Per-warp Q gmem sub-tile of shape (Q_BLOCK_SIZE, DEPTH) sliced from the per-head 2D Q view.

Returns:

`run`

static def run[k_t: MHAOperand, v_t: MHAOperand, mask_t: MHAMask, q_dtype: DType, output_dtype: DType, q_layout: TensorLayout, o_layout: TensorLayout, q_engine: TensorEngine, o_engine: TensorEngine, ragged: Bool = False, sink: Bool = False](q: TileTensor[q_dtype, q_layout, ImmutAnyOrigin, Engine=q_engine], k: k_t, v: v_t, o: TileTensor[output_dtype, o_layout, MutAnyOrigin, Engine=o_engine], mask_functor: mask_t, scale: Float32, num_keys: Int32, start_pos: Int32, sink_weights_ptr: Pointer[Scalar[q_dtype], ImmutAnyOrigin])

Multi-block 8-warp MHA forward (inference-only).

Grid: (NUM_HEADS, ceildiv(seq_len, BM), batch). Each block owns one (batch, head, BM-tile) slice; the 8 warps within split the BM-tile's Q rows.

Expected layouts / shapes:

q, o: (batch, seq_len, NUM_HEADS, DEPTH) row-major TileTensor. o's dtype matches config.output_dtype: BF16 for the production dispatcher (which holds a BF16 output buffer) or FP32 if the caller wants the unnormalized accumulator.
k, v: any MHAOperand whose block_paged_tile[KV_BLOCK] returns (KV_BLOCK, DEPTH) tiles per (batch, t*KV_BLOCK, kv_head, 0). LayoutTensorMHAOperand for contiguous test / bench buffers; KVCacheMHAOperand for paged production caches (page_size >= KV_BLOCK = 64).

batch and seq_len / num_keys may be dynamic; NUM_HEADS, NUM_KV_HEADS, DEPTH must be static. NUM_HEADS must be a multiple of NUM_KV_HEADS (GROUP = NUM_HEADS // NUM_KV_HEADS).

Parameters:

k_t (MHAOperand): K operand type (inferred); any MHAOperand whose block_paged_tile returns (KV_BLOCK, DEPTH) tiles.
v_t (MHAOperand): V operand type (inferred); same tile contract as k_t.
mask_t (MHAMask): Mask functor type (inferred); selects the per-tile masking strategy (causal, sliding-window, null, etc.).
q_dtype (DType): Element dtype of q (inferred); must equal config.dtype.
output_dtype (DType): Element dtype of o (inferred); must equal config.output_dtype.
q_layout (TensorLayout): Layout of the q TileTensor (inferred).
o_layout (TensorLayout): Layout of the o TileTensor (inferred).
q_engine (TensorEngine): Engine of the q TileTensor (inferred).
o_engine (TensorEngine): Engine of the o TileTensor (inferred).
ragged (Bool): Whether q is a per-sequence slice in a packed ragged batch (defaults to False).
sink (Bool): Whether to seed the online softmax with attention-sink weights (defaults to False).

Args:

q (TileTensor[q_dtype, q_layout, ImmutAnyOrigin, Engine=q_engine]): Q tile tensor.
k (k_t): K operand (MHAOperand).
v (v_t): V operand (MHAOperand).
o (TileTensor[output_dtype, o_layout, MutAnyOrigin, Engine=o_engine]): Output tile tensor (config.output_dtype, same shape as q).
mask_functor (mask_t): Per-tile mask predicate (causal, sliding-window, etc.). Evaluated inside the QK→softmax cluster; identity for unmasked attention.
scale (Float32): Softmax scale (1/sqrt(DEPTH)).
num_keys (Int32): Runtime length of the K/V sequence.
start_pos (Int32): Position of the first Q row in the global sequence: non-zero for prefill chunks of a longer generation. Used by the mask functor to compute the causal cutoff.
sink_weights_ptr (Pointer[Scalar[q_dtype], ImmutAnyOrigin]): Per-q-head attention-sink scalar weights. Read only when the comptime sink parameter is True; the non-sink path comptime-elides the load, so callers may pass Pointer[...].unsafe_dangling() when sink=False. Indexed by head_idx once per block at init time, cast to FP32, multiplied by log2e to land in the kernel's log2-units rowmax, and seeded into max_vec / max_vec_prev / norm_vec so the hot loop stays sink-agnostic.

`ragged_kernel`

static def ragged_kernel[k_t: MHAOperand, v_t: MHAOperand, mask_t: MHAMask, qkv_dtype: DType, output_dtype: DType, cross_attention: Bool = False, sink: Bool = False](q_ptr: Pointer[Scalar[qkv_dtype], ImmutAnyOrigin], k: k_t, v: v_t, output_ptr: Pointer[Scalar[output_dtype], MutAnyOrigin], mask_functor: mask_t, scale: Float32, input_row_offsets_ptr: Pointer[UInt32, ImmutAnyOrigin], kv_input_row_offsets_ptr: Pointer[UInt32, ImmutAnyOrigin], sink_weights_ptr: Pointer[Scalar[qkv_dtype], ImmutAnyOrigin])

Ragged-batch GPU kernel entry: per-sequence setup + run.

The non-ragged equivalent is run itself (which takes already-sliced per-batch TileTensors). For ragged, this wrapper does the per-block ragged setup so the launcher can pass a single packed Q pointer + input_row_offsets.

cross_attention=False (default): self-attention, where K/V length equals Q length plus any cached prefix. num_keys derives from start_pos + seq_len. kv_input_row_offsets_ptr is unused (caller may pass any well-typed stub).

cross_attention=True: encoder-decoder style. K/V lengths come from kv_input_row_offsets_ptr, independent of the Q-side offsets. Mirrors the FA2 contract at mha.mojo:1755-1762.

Parameters:

k_t (MHAOperand): K operand type (inferred); any MHAOperand whose block_paged_tile returns (KV_BLOCK, DEPTH) tiles.
v_t (MHAOperand): V operand type (inferred); same tile contract as k_t.
mask_t (MHAMask): Mask functor type (inferred); selects the per-tile masking strategy.
qkv_dtype (DType): Element dtype of Q, K, and V (inferred); must equal config.dtype.
output_dtype (DType): Element dtype of the output buffer (inferred); must equal config.output_dtype.
cross_attention (Bool): Whether K/V length is independent of Q (encoder-decoder style) (defaults to False).
sink (Bool): Whether to seed the online softmax with attention-sink weights (defaults to False).

Args:

q_ptr (Pointer[Scalar[qkv_dtype], ImmutAnyOrigin]): Pointer to the packed Q buffer, pre-offset per sequence by input_row_offsets_ptr.
k (k_t): K operand (MHAOperand).
v (v_t): V operand (MHAOperand).
output_ptr (Pointer[Scalar[output_dtype], MutAnyOrigin]): Pointer to the output buffer, same packing as q_ptr.
mask_functor (mask_t): Per-tile mask predicate (causal, sliding-window, etc.).
scale (Float32): Softmax scale (1/sqrt(DEPTH)).
input_row_offsets_ptr (Pointer[UInt32, ImmutAnyOrigin]): Cumulative uint32 row offsets for Q sequences, length batch_size + 1.
kv_input_row_offsets_ptr (Pointer[UInt32, ImmutAnyOrigin]): Cumulative uint32 row offsets for K/V sequences; read only when cross_attention is True.
sink_weights_ptr (Pointer[Scalar[qkv_dtype], ImmutAnyOrigin]): Per-q-head attention-sink scalar weights; read only when sink is True.

Parameters​

Implemented traits​

comptime members​

BM​

D_FRAG_PER_LANE​

DEPTH​

k_swizzle​

k_swizzle2​

KTileLoader​

KV_BLOCK​

NUM_HEADS​

NUM_KV_HEADS​

NUM_THREADS​

NUM_WARPS​

prescale_q​

Q_BLOCK_SIZE​

RESCALE_THRESHOLD​

v_swizzle​

VTileLoader​

Methods​

load_q​

run​

ragged_kernel​

Parameters

Implemented traits

`comptime` members

`BM`

`D_FRAG_PER_LANE`

`DEPTH`

`k_swizzle`

`k_swizzle2`

`KTileLoader`

`KV_BLOCK`

`NUM_HEADS`

`NUM_KV_HEADS`

`NUM_THREADS`

`NUM_WARPS`

`prescale_q`

`Q_BLOCK_SIZE`

`RESCALE_THRESHOLD`

`v_swizzle`

`VTileLoader`

Methods

`load_q`

`run`

`ragged_kernel`