IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

FA4MiscMBars

struct FA4MiscMBars[*, num_qk_stages: Int = Int(1), num_pv_stages: Int = Int(1), num_kv_stages: Int = Int(2), use_order_barriers: Bool = True, use_shared_kv: Bool = False, pair_cta: Bool = False, num_q: Int = Int(2), splitk_partitions: Int = Int(1), BM: Int = Int(128), use_ws: Bool = False]

Manages all mbarrier resources for FA4.

This struct consolidates all mbarrier management including:

  • S barriers (score MMA synchronization)
  • C barriers (correction synchronization)
  • Order barriers (softmax ordering)
  • Q1Sync barriers (Q tile synchronization)
  • K/V pipeline barriers (separate K and V)
  • O pipeline barriers

Memory layout (count=128 first, then count=1): [S0_cons] [S1_cons] [C0] [C1] [Order*] | [S0_prod] [S1_prod] [Q1Sync**] [K] [V] [O_prod] *Order barriers only present when use_order_barriers=True **Q1Sync barriers only present when num_q == 2

Parameters​

  • ​num_qk_stages (Int): Number of stages for Q@K' MMA (K loading can be staged).
  • ​num_pv_stages (Int): Number of stages for P@V MMA (P writing can be staged).
  • ​num_kv_stages (Int): Number of KV buffer stages for double/triple buffering.
  • ​use_order_barriers (Bool): When True, allocate order barriers to prevent softmax warp group overlap. When False, order barriers are omitted.
  • ​use_shared_kv (Bool): Whether the K and V share the same pipeline, or separate.
  • ​pair_cta (Bool): Whether to use 1-cta or 2-cta implementation.
  • ​num_q (Int): Number of Q tiles per CTA. When 1, the Q1Sync slot is collapsed and K_offset shifts down by num_qk_stages. Must be 2 for any caller of q1_wait_mbar().
  • ​splitk_partitions (Int): Number of split-K partitions (P). When num_q == 1 and this exceeds 1, a single publish barrier is added so the cross-CTA O combine writer observes all P partitions' staged partials. Otherwise no extra barrier is allocated, keeping a byte-identical mbar layout.
  • ​BM (Int): Block size (rows per CTA). For 1Q split-K this is the number of WG0 rows that each arrive on the publish barrier, so its count is BM * P (every row of every partition). Only used to size the publish barrier; defaults to 128 (== WARPGROUP_SIZE on the 1Q path) for non-split-K callers.
  • ​use_ws (Bool): Warp-specialized packed-TMEM (MMA_M=32) datapath. When True, num_kv_stages counts depth-split 256x64 sub-tile ring slots ("Convention B"), so K_barriers = 2 * num_kv_stages (the num_qk_stages depth factor is already folded into the slot count). When False (default), the non-WS full-depth-tile count applies and the layout is byte-identical.

Fields​

  • ​mbar_base (MBarType):

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members​

C0_offset​

comptime C0_offset = (Int(2) * num_pv_stages)

C1_offset​

comptime C1_offset = (FA4MiscMBars[num_qk_stages=num_qk_stages, num_pv_stages=num_pv_stages, num_kv_stages=num_kv_stages, use_order_barriers=use_order_barriers, use_shared_kv=use_shared_kv, pair_cta=pair_cta, num_q=num_q, splitk_partitions=splitk_partitions, BM=BM, use_ws=use_ws].C0_offset + Int(2))

K_barriers​

comptime K_barriers = (Int(2) * num_kv_stages) if use_ws else ((Int(2) * num_qk_stages) * num_kv_stages)

K_offset​

comptime K_offset = (Int((add (mul num_pv_stages, 2), Int(2) if use_order_barriers else Int(0), 6)) + num_qk_stages if (eq num_q, 2) else Int(0))

num_order_barriers​

comptime num_order_barriers = Int(2) if use_order_barriers else Int(0)

number_warpgroup_count​

comptime number_warpgroup_count = FA4MiscMBars[num_qk_stages=num_qk_stages, num_pv_stages=num_pv_stages, num_kv_stages=num_kv_stages, use_order_barriers=use_order_barriers, use_shared_kv=use_shared_kv, pair_cta=pair_cta, num_q=num_q, splitk_partitions=splitk_partitions, BM=BM, use_ws=use_ws].S0_producer_offset

O_producer_offset​

comptime O_producer_offset = (Int((add (mul num_pv_stages, 2), num_qk_stages if (eq num_q, 2) else Int(0), Int(2) if use_order_barriers else Int(0), Int((mul num_kv_stages, 2)) if use_ws else Int((mul num_kv_stages, num_qk_stages, 2)), 6)) + Int(0) if use_shared_kv else Int((mul num_kv_stages, 2)))

order_offset​

comptime order_offset = (FA4MiscMBars[num_qk_stages=num_qk_stages, num_pv_stages=num_pv_stages, num_kv_stages=num_kv_stages, use_order_barriers=use_order_barriers, use_shared_kv=use_shared_kv, pair_cta=pair_cta, num_q=num_q, splitk_partitions=splitk_partitions, BM=BM, use_ws=use_ws].C1_offset + Int(2))

Publish_count​

comptime Publish_count = Int(1) if (splitk_partitions > Int(1)) if (num_q == Int(1)) else (num_q == Int(1)) else Int(0)

Publish_offset​

comptime Publish_offset = (Int((add (mul num_pv_stages, 2), num_qk_stages if (eq num_q, 2) else Int(0), Int(2) if use_order_barriers else Int(0), Int(0) if use_shared_kv else Int((mul num_kv_stages, 2)), Int((mul num_kv_stages, 2)) if use_ws else Int((mul num_kv_stages, num_qk_stages, 2)), 6)) + Int(2))

Q1Sync_count​

comptime Q1Sync_count = num_qk_stages if (num_q == Int(2)) else Int(0)

Q1SyncIdx​

comptime Q1SyncIdx = (Int((add (mul num_pv_stages, 2), Int(2) if use_order_barriers else Int(0), 5)) + Int(1))

S0_consumer_offset​

comptime S0_consumer_offset = 0

S0_producer_offset​

comptime S0_producer_offset = (FA4MiscMBars[num_qk_stages=num_qk_stages, num_pv_stages=num_pv_stages, num_kv_stages=num_kv_stages, use_order_barriers=use_order_barriers, use_shared_kv=use_shared_kv, pair_cta=pair_cta, num_q=num_q, splitk_partitions=splitk_partitions, BM=BM, use_ws=use_ws].order_offset + Int(2) if use_order_barriers else Int(0))

S1_consumer_offset​

comptime S1_consumer_offset = num_pv_stages

S1_producer_offset​

comptime S1_producer_offset = (Int((add (mul num_pv_stages, 2), Int(2) if use_order_barriers else Int(0), 4)) + Int(1))

size​

comptime size = (Int((add (mul num_pv_stages, 2), num_qk_stages if (eq num_q, 2) else Int(0), Int(2) if use_order_barriers else Int(0), Int(0) if use_shared_kv else Int((mul num_kv_stages, 2)), Int((mul num_kv_stages, 2)) if use_ws else Int((mul num_kv_stages, num_qk_stages, 2)), 8)) + Int(1) if (xor (eq (xor (lt splitk_partitions, 2), True), False), True) if (eq num_q, 1) else (eq num_q, 1) else Int(0))

SPipelineConsumer​

comptime SPipelineConsumer = RolePipeline[Int(1), False, consumer_sub_stages=num_pv_stages]

SPipelineProducer​

comptime SPipelineProducer = RolePipeline[Int(1), consumer_sub_stages=num_pv_stages]

V_barriers​

comptime V_barriers = Int(0) if use_shared_kv else (Int(2) * num_kv_stages)

V_offset​

comptime V_offset = (Int((add (mul num_pv_stages, 2), num_qk_stages if (eq num_q, 2) else Int(0), Int(2) if use_order_barriers else Int(0), 6)) + Int((mul num_kv_stages, 2)) if use_ws else Int((mul num_kv_stages, num_qk_stages, 2)))

Methods​

__init__​

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

init​

def init(self, *, lane_idx: Int32)

producer_s0​

def producer_s0(self) -> Self.SPipelineProducer

Get S producer for warp group 0.

Returns:

Self.SPipelineProducer

producer_s1​

def producer_s1(self) -> Self.SPipelineProducer

Get S producer for warp group 1.

Returns:

Self.SPipelineProducer

consumer_s​

def consumer_s(self, wg_idx: UInt32) -> Self.SPipelineConsumer

Get S consumer for given warp group.

Args:

  • ​wg_idx (UInt32): Warp group index (0 or 1) selecting the S consumer pipeline.

Returns:

Self.SPipelineConsumer

consumer_c0​

def consumer_c0(self) -> RolePipeline[Int(1), False]

Returns:

RolePipeline[Int(1), False]

consumer_c1​

def consumer_c1(self) -> RolePipeline[Int(1), False]

Returns:

RolePipeline[Int(1), False]

producer_c​

def producer_c(self, wg_idx: UInt32) -> RolePipeline[Int(1)]

Returns:

RolePipeline[Int(1)]

pipeline_order_wait​

def pipeline_order_wait(self, wg_idx: UInt32) -> MBarType

Returns:

MBarType

pipeline_order_arrive​

def pipeline_order_arrive(self, wg_idx: UInt32) -> MBarType

Returns:

MBarType

q1_wait_mbar​

def q1_wait_mbar(self) -> MBarType

Returns:

MBarType

get_k_mbars​

def get_k_mbars(self) -> MBarType

Returns base pointer for K pipeline barriers.

Returns:

MBarType

get_v_mbars​

def get_v_mbars(self) -> MBarType

Returns base pointer for V pipeline barriers. In shared mode, returns the same as get_k_mbars (shared pipeline).

Returns:

MBarType

combined_p_o_consumer​

def combined_p_o_consumer(self, wg_idx: UInt32) -> MBarType

Combined P+O consumer barrier for given warp group.

Arrived at by BOTH softmax (P ready) and correction (O rescaled). Returns S_consumer[0] for wg_idx=0 or wg_idx=1.

Args:

  • ​wg_idx (UInt32): Warp group index (0 or 1) selecting the consumer barrier slot.

Returns:

MBarType

consumer_o​

def consumer_o(self) -> RolePipeline[Int(2), False, consumer_sub_stages=num_pv_stages]

Get O consumer pipeline.

Wait side: O_producer barriers (stride 1, indexed by stage). Release side: combined S+O barriers (S_consumer[0] per wg, stride num_pv_stages).

Returns:

RolePipeline[Int(2), False, consumer_sub_stages=num_pv_stages]

consumer_o0​

def consumer_o0(self) -> RolePipeline[Int(1), False, consumer_sub_stages=num_pv_stages]

Single-O (1Q wide-V) O consumer: a ONE-stage pipeline on WG0's O-producer barrier only.

The standard consumer_o() is a 2-stage pipeline that alternates between the two per-WG O-producer barriers (O_producer_offset+0 for WG0, +1 for WG1). The single-O path runs a single warp group (WG0) that accumulates ALL K-tiles into the single (aliased) O0, so the correction warp must wait on ONLY O_producer_offset+0 with an incrementing phase, never the never-produced +1 (which would deadlock). Release side is WG0's combined P+O consumer barrier, as in producer_o0.

Returns:

RolePipeline[Int(1), False, consumer_sub_stages=num_pv_stages]

producer_o0​

def producer_o0(self) -> RolePipeline[Int(1)]

Get O producer for warp group 0.

Returns:

RolePipeline[Int(1)]

producer_o1​

def producer_o1(self) -> RolePipeline[Int(1)]

Get O producer for warp group 1.

Returns:

RolePipeline[Int(1)]

publish_mbar​

def publish_mbar(self) -> MBarType

Split-K cross-CTA O-combine publish barrier (count=BM * P).

Every WG0 row of every partition CTA arrive_clusters on every peer's copy (BM rows × P partitions arrivals per copy); the softmax threads wait on it before the writer's DSMEM peer reads. Per-row arrivals mean the publish sites need no CTA-local named_barrier to collect rows. Only present for 1Q split-K (Publish_count > 0).

Returns:

MBarType

num_mbars​

static def num_mbars() -> UInt32

Returns:

UInt32

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact