For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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 theS = Q@K'andO += P@VMMAs. - 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,
ImplicitlyCopyable,
ImplicitlyDeletable,
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, _safe=False]) -> 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
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!