Mojo struct
MLAKVProducerPipeline
@register_passable(trivial)
struct MLAKVProducerPipeline[dtype: DType, config: FA4Config]
Fields
- kv_pipeline (
KVPipeline[config.num_kv_stages, config.num_mma_stages]): - smem (
MLAKVProducerPipeline[dtype, config].SMemType):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
Movable,
UnknownDestructibility
comptime members
__copyinit__is_trivial
comptime __copyinit__is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__moveinit__is_trivial
comptime __moveinit__is_trivial = True
k_bytes
comptime k_bytes = (MLAKVProducerPipeline[dtype, config].k_elements * size_of[dtype]())
k_elements
comptime k_elements = tile_layout_k_major[dtype, config.BN, config.BK0, config.swizzle_mode]().size()
k_layout
comptime k_layout = tile_layout_k_major[dtype, config.BN, 128, config.swizzle_mode]()
k_rope_layout
comptime k_rope_layout = tile_layout_k_major[dtype, config.BN, 64, config.swizzle_mode]()
KPairType
comptime KPairType = TMADestination[dtype, tile_layout_k_major[dtype, config.BN, config.BK0, config.swizzle_mode]()]
KType
comptime KType = LayoutTensor[dtype, tile_layout_k_major[dtype, config.BN, config.BK0, config.swizzle_mode](), MutAnyOrigin, address_space=AddressSpace.SHARED, layout_int_type=DType.int32, linear_idx_type=DType.int32, alignment=128]
SMemType
comptime SMemType = LegacyUnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED]
v_bytes
comptime v_bytes = (MLAKVProducerPipeline[dtype, config].v_elements * size_of[dtype]())
v_elements
comptime v_elements = tile_layout_mn_major[dtype, 128, config.BK1, config.swizzle_mode]().size()
VPairType
comptime VPairType = TMADestination[dtype, tile_layout_mn_major[dtype, 128, config.BK1, config.swizzle_mode]()]
VType
comptime VType = LayoutTensor[dtype, tile_layout_mn_major[dtype, 128, config.BK1, config.swizzle_mode](), MutAnyOrigin, address_space=AddressSpace.SHARED, layout_int_type=DType.int32, linear_idx_type=DType.int32, alignment=128]
Methods
__init__
__init__(mbar: LegacyUnsafePointer[SharedMemBarrier, address_space=AddressSpace.SHARED], smem: LegacyUnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED]) -> Self
__init__(kv_pipeline: KVPipeline[config.num_kv_stages, config.num_mma_stages], smem: LegacyUnsafePointer[Scalar[dtype], address_space=AddressSpace.SHARED]) -> Self
init
init(self)
Only one of the producer or consumer should call init().
get_kv_smem
get_kv_smem[*, mma_stage: Int](self) -> MLAKVProducerPipeline[dtype, config].SMemType
Returns:
MLAKVProducerPipeline
get_k
get_k[*, mma_stage: Int, expect: Bool = True](self) -> MLAKVProducerPipeline[dtype, config].KPairType
Returns:
MLAKVProducerPipeline
get_v
get_v[*, mma_stage: Int](self) -> MLAKVProducerPipeline[dtype, config].VPairType
Returns:
MLAKVProducerPipeline
acquire_kv
acquire_kv[*, mma_stage: Int = (config - 1)](self)
commit_kv_step
commit_kv_step(mut self)
Step the kv pipeline. The does not perform the commit on the mbars; that should be handled by the tma_op.async_copy.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!