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

TMAProducerPipeline

struct TMAProducerPipeline[dtype: DType, config: FA4Config[config.qkv_dtype, rope_dtype_=config.rope_dtype_, scale_dtype_=config.scale_dtype_], is_k: Bool = True]

Unified producer pipeline for K and V TMA loading.

K loading (is_k=True): Can be staged (num_qk_stages chunks), uses k_major layout. V loading (is_k=False): Always complete (qk_stage=0), uses mn_major layout.

Parameters

• ​dtype (DType): Element dtype of the K or V tile being loaded via TMA.
• ​config (FA4Config[config.qkv_dtype, rope_dtype_=config.rope_dtype_, scale_dtype_=config.scale_dtype_]): FlashAttention-4 configuration providing tile sizes, swizzle mode, and stage counts.
• ​is_k (Bool): Whether this pipeline loads the K operand (True, k-major layout with staged qk_stages) or the V operand (False, mn-major layout with qk_stage=0); defaults to True.

Fields

• ​pipeline (StagedPipeline[config.num_kv_stages, config.num_qk_stages if is_k else Int(1)]):
• ​smem (TMAProducerPipeline[dtype, config, is_k].SMemType):

Implemented traits

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

bytes

comptime bytes = TMAProducerPipeline[dtype, config, is_k].tile_bytes

elements

comptime elements = TMAProducerPipeline[dtype, config, is_k].tile_layout.size()

elements_full

comptime elements_full = (tile_layout_k_major[dtype, config.k_rows_per_cta(), config.BK0, config.swizzle_mode]() if is_k else tile_layout_mn_major[dtype, config.v_cols_per_cta(), config.BK1, config.swizzle_mode]().size() * config) if is_k else TMAProducerPipeline[dtype, config, is_k].elements

KPairType

comptime KPairType = TMAProducerPipeline[dtype, config, is_k].PairType

num_qk_stages_effective

comptime num_qk_stages_effective = config.num_qk_stages if is_k else Int(1)

PairType

comptime PairType = TMADestination[dtype, tile_layout_k_major[dtype, config.k_rows_per_cta(), config.BK0, config.swizzle_mode]() if is_k else tile_layout_mn_major[dtype, config.v_cols_per_cta(), config.BK1, config.swizzle_mode]().size()]

SMemType

comptime SMemType = Pointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]

tile_bytes

comptime tile_bytes = (tile_layout_k_major[dtype, config.k_rows_per_cta(), config.BK0, config.swizzle_mode]() if is_k else tile_layout_mn_major[dtype, config.v_cols_per_cta(), config.BK1, config.swizzle_mode]().size() * size_of[dtype]())

tile_layout

comptime tile_layout = tile_layout_k_major[dtype, config.k_rows_per_cta(), config.BK0, config.swizzle_mode]() if is_k else tile_layout_mn_major[dtype, config.v_cols_per_cta(), config.BK1, config.swizzle_mode]()

Methods

__init__

def __init__(mbar: Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED], smem: Pointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]) -> Self

def __init__(pipeline: StagedPipeline[config.num_kv_stages, config.num_qk_stages if is_k else Int(1)], smem: Pointer[Scalar[dtype], MutAnyOrigin, address_space=AddressSpace.SHARED]) -> Self

get_smem

def get_smem[*, qk_stage: Int = Int(0)](self) -> Self.SMemType

Get smem pointer for current stage.

Parameters:

• ​qk_stage (Int): K-loading sub-stage whose smem offset to return (defaults to 0).

Returns:

Self.SMemType

get_tile

def get_tile[*, qk_stage: Int = Int(0)](self) -> Self.PairType

Get TMA destination for this stage.

Parameters:

• ​qk_stage (Int): K-loading sub-stage whose TMA destination to return (defaults to 0).

Returns:

Self.PairType

def get_tile[*, qk_stage: Int = Int(0)](self, e: Int32) -> Self.PairType

Get TMA destination with optional expect_bytes.

Parameters:

• ​qk_stage (Int): K-loading sub-stage whose TMA destination to return (defaults to 0).

Args:

• ​e (Int32): elect() result; when nonzero, issues expect_bytes on the producer mbarrier before returning.

Returns:

Self.PairType

acquire

def acquire[*, qk_stage: Int = Int(0)](self)

Wait for consumer to release the buffer.

Parameters:

• ​qk_stage (Int): K-loading sub-stage to acquire (defaults to 0).

commit_step

def commit_step(mut self)

Step the pipeline. Commit is handled by tma_op.async_copy.

get_k_smem

def get_k_smem[*, qk_stage: Int](self) -> Self.SMemType

Returns:

Self.SMemType

get_k

def get_k[*, qk_stage: Int](self) -> Self.PairType

Returns:

Self.PairType

def get_k[*, qk_stage: Int](self, e: Int32) -> Self.PairType

Returns:

Self.PairType

acquire_k

def acquire_k[*, qk_stage: Int](self)

get_v_smem

def get_v_smem(self) -> Self.SMemType

Returns:

Self.SMemType

get_v

def get_v(self, e: Int32) -> Self.PairType

Returns:

Self.PairType

acquire_v

def acquire_v(self)