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

KVPipelineGeneric

struct KVPipelineGeneric[num_kv_stages: Int, num_qk_stages: Int, num_producer: Int, num_consumer: Int]

KVPipeline has num_kv_stages * num_qk_stages stages. num_kv_stages refers to how many K and V tiles we pipeline for performing the S = Q@K' and O += P@V MMAs. Each of these MMAs is broken up into num_qk_stages pipelined MMAs. We set step=False for all but the last MMA that completes the operation. An alternative implementation would separate the two, and potentially allow for more overall stages at the cost of slightly more bookkeeping.

Parameters

• ​num_kv_stages (Int): Number of KV tiles pipelined for the S = Q@K' and O += P@V MMAs.
• ​num_qk_stages (Int): Number of pipelined sub-MMAs each QK or PV MMA is broken into.
• ​num_producer (Int): Number of producer threads arriving on each producer mbarrier.
• ​num_consumer (Int): Number of consumer threads arriving on each consumer mbarrier.

Fields

• ​mbar (MBarType):
• ​state (PipelineState[num_kv_stages]):

Implemented traits

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

comptime members

num_stages

comptime num_stages = (num_kv_stages * num_qk_stages)

Methods

__init__

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

init

def init(self)

producer_mbar

def producer_mbar[qk_stage: Int](self) -> MBarType

Returns:

MBarType

consumer_mbar

def consumer_mbar[qk_stage: Int](self, idx: UInt32) -> MBarType

Returns:

MBarType

def consumer_mbar[qk_stage: Int](self) -> MBarType

Returns:

MBarType

producer_acquire

def producer_acquire[qk_stage: Int = (num_qk_stages - Int(1))](self)

Returns the dynamic pipe idx.

Parameters:

• ​qk_stage (Int): QK sub-stage index whose consumer mbarrier to wait on (defaults to the last QK stage).

consumer_wait

def consumer_wait[qk_stage: Int = (num_qk_stages - Int(1))](self)

consumer_release

def consumer_release[qk_stage: Int = (num_qk_stages - Int(1))](mut self, e: Int32)

num_mbars

static def num_mbars() -> UInt32

Returns:

UInt32