State-Ownership Architecture Decision

Decision

Prototype VEA-B, the one-CTA Q128 paired-owner pipeline. Phase 2A capability probes pass with CUDA block-scope mbarriers as the selected two-slot handoff. VEA-A remains the fallback if the complete 320-thread kernel fails exactness, resources, or latency. VEA-C is the higher-risk second prototype only after DSM placement and handoff latency are measured. No attention result is achieved.

Every design is side-by-side research under kernels/. The production Sage2 path remains the reference and fallback; no dispatch integration is authorized.

Non-Negotiable Dataflow

Keep the exact global K mean and V scale prepasses.
Keep per-32-row Q and per-64-row K INT8 quantization bytes/scales.
Traverse all 591 K64/V64 tiles in increasing order.
For each tile preserve QK, score scaling, max update, denominator update, FP8 probability conversion, two K32 FP16 PV MMAs, and FP32 RO addition order.
Never materialize an N-by-N score or probability tensor.
Never split KV traversal across independently normalized CTAs.
Emit contiguous token-major NHD BF16 output.

The exact handoff record for tile t is:

slot[t mod 2] = {
  epoch: t,
  p_fp8: exact post-softmax E4M3 bytes,
  ro_rescale: exact FP32 old_scale applied before tile-t PV,
  valid_rows: tail predicate
}

At completion the QK/softmax owner publishes the exact FP32 reciprocal of d. The output owner applies final normalization, V scale, and BF16 conversion in the reference order. Shared storage transfers bytes; it must not recompute or reassociate numerical state.

VEA-A: Q64 Split-D Intra-CTA Pipeline

Geometry: Q64 x K64 x D128, 256 threads, eight warps, 33,096 CTAs.

Role	Warps	State ownership	Registers/thread estimate
TMA/input producer	1	K/V slots, predicates, optional embedded P1/P2 conversion	32-48
QK/softmax	2	Q fragments, RS/RS_f8, m, d; one warp per Q32	96-128
PV/output	4	two D64 RO owners per Q32; instantaneous FP16 PV	88-120
control/store	1	epochs, barriers, final stores	32-64

State movement: each QK warp publishes one Q32xK64 FP8 probability tile and one FP32 ro_rescale per logical query row to a two-slot shared ring. Two PV warps consume disjoint D64 halves. m/d never move per tile; final reciprocal moves once. QK and PV fragments remain live simultaneously across different warps, never in the same warp.

Resource model: resident Q INT8 8 KiB, double K 16 KiB, double V 16 KiB, double score ring 8 KiB, and 3-4 KiB metadata, totaling 51-52 KiB. One CTA/SM is projected because two exceed SM121 shared memory. It retains the baseline eight active warps but separates dependency chains. Named barriers coordinate producer-ready and consumer-released epochs; projected synchronization cost is 8-18 ms.

Projected mainloop screening range: 192-222 ms. This can plausibly beat 220 ms, but a sub-190 result requires low handoff cost and real INT8/FP8 overlap.

Structural difference from P0-P3: P3 kept full RO, scores, m, and d in each compute warp. VEA-A permanently removes RO from QK warps and removes scores/m/d from PV warps, then splits RO by output columns.

VEA-B: Q128 Paired QK/PV Owners

Recommended geometry: Q128 x K64 x D128, 320 threads, ten warps, 16,576 CTAs.

Role	Warps	State ownership	Registers/thread estimate
K producer/sequencer	1	K slot, Q/K scales, epochs	32-48
V producer	1	V slot, V scales, optional embedded P2 conversion	40-56
QK/softmax producers	4	Q, RS/RS_f8, m, d; one per Q32	96-128
PV/output owners	4	full D128 RO and instantaneous FP16 PV; one per Q32	144-176

Each QK warp is paired with one PV warp. The producer publishes exact FP8 probability bytes and ro_rescale; the consumer first multiplies its persistent RO by that exact FP32 value, then executes the same ordered two-K32 FP16 PV operations and FP32 additions. m/d stay producer-owned. QK and PV fragments are concurrently live only in separate owner warps.

Before aliasing, shared memory is 66-68 KiB: Q 16, double K 16, double V 16, double score ring 16, metadata 2-4 KiB. After all persistent Q fragments load, the complete 16 KiB Q region aliases the complete two-slot score ring, yielding 50-52 KiB. Logical role-weighted live state is approximately 39K registers, but that is not a physical allocation claim. Conservative uniform allocation at the largest modeled 176-register role is 56,320 registers/CTA; compilation must remain at or below 200 registers/thread with zero material spills. One CTA gives ten active warps (20.83% of the 48-warp ceiling).

Two ready and two free block-scope mbarriers provide role-asymmetric two-slot reuse: six producer warps arrive without waiting on ready, four consumer warps wait and consume, then the roles reverse on free. Completion is CTA-scoped per slot; it is not a four-pair barrier claim. K/V TMA barriers are separate. The corrected Phase 2A 48-block, 591-epoch probe measures 0.076800 ms p50 and 0.078880 ms p95 for barriers alone, far below the 11.85445 ms budget. The inline-PTX named-barrier alternative measures 0.056128 ms p50 but is rejected because racecheck reports five hazards; mbarrier reports none.

Projected mainloop screening range: 180-207 ms; conservative complete attention range with 25-28 ms non-overlapped exact preparation is 205-235 ms. This has a credible mainloop path below 190 ms without relying on DSM.

Structural difference from P0-P3: no QK warp owns RO, no PV warp owns RS/RS_f8/m/d, and load warps own neither. P1/P2 can be embedded to feed slots, but their standalone boundaries are not benchmark targets.

VEA-C: Two-CTA DSM Ownership Split

Geometry: one two-CTA cluster per (head,Q128) work item. The producer CTA owns QK/softmax; the consumer CTA owns PV/output. Both traverse the same 591 tiles.

CTA / role	Warps	State ownership	Registers/thread estimate
Producer load	1	Q/K slots and scales	32-48
Producer QK/softmax	4	Q, RS/RS_f8, m, d	96-128
Producer publication	1	DSM epochs and barriers	32-48
Consumer V load	1	V slots and scale	40-56
Consumer PV/output	4	full RO and instantaneous FP16 PV	144-176
Consumer final/store	1	reciprocal consumption and BF16 store	40-64

A two-slot DSM ring carries Q128xK64 FP8 probabilities (16 KiB total), exact rescale records, and final reciprocal. No arithmetic reduction crosses CTAs. The producer cannot overwrite an epoch until the consumer releases it; the consumer cannot observe tile t+1 before completing tile t.

Projected shared memory is approximately 50 KiB producer plus 18 KiB consumer. Projected pair register demand is 50-60K and 12 warps. Residency is not yet credible evidence: SM121 cluster placement, simultaneous residency, and DSM round-trip latency must be measured first. If a pair consumes two SMs with poor work-item throughput, reject it regardless of single-cluster latency.

Projected mainloop range if feasibility probes pass: 174-210 ms; projected synchronization/DSM cost is 12-30 ms. This offers the strongest ownership split and possible sub-190 path, but has the weakest current hardware evidence.

Structural difference from P0-P3: RS/RS_f8/m/d and RO are in separate CTAs, not private duplicates in alternating warps. DSM is an explicit state-transfer boundary rather than a temporal schedule rearrangement.

Excluded Architecture Classes

Class	Reason
Global score/probability ring	Adds approximately `2HN^2`, about 160 GB, of L2 score traffic and materially stores probabilities
Split-KV multi-CTA reduction	Changes online-softmax and RO reduction order
K/V broadcast-only cluster	Optimizes a 98.85%-hit data path without reducing RO ownership
Register cap or source lifetime retune	Closed P0; cannot cross residency cliff without spills
Temporal QK/PV pairing with private state	Closed P3; simultaneous state remains allocated
Standalone P1 or P2	Correct but below complete-block gate; reusable only as embedded producers
Hopper WGMMA or FP8/FP4 QK	Unsupported on SM121 or changes exact arithmetic

Prototype Order And Stop Conditions

Phase 2A register-allocation, shared-memory, handoff, sanitizer, and dual-pipe probes are complete; see PHASE2A_CAPABILITY_REPORT.md.
Retain CUDA block-scope mbarriers; do not use the rejected inline named barrier primitive.
Implement VEA-B only for one aligned short shape and exact D=128 arithmetic.
Consider VEA-C only after DSM feasibility; use VEA-A if B fails role/resource allocation.

Stop immediately on changed KV order, any material score tensor, nonzero output difference, forced local spilling, inability to beat the 220 ms model gate, or fallback to Sage2 for main computation.

Interface And Future Ulysses Layout

The isolated API consumes projection-strided BF16 NHD views directly. Its work descriptor must carry q_start, q_count, kv_count, head_start, head_count, source token/head strides, destination token/head strides, and a future rank destination. No HND transpose or contiguous input copy is allowed.

Output stores land directly in token-major [token, local_head*128+d] rows for the local output projection. Ragged local token counts and nonzero destination offsets are legal, but each rank must receive the globally ordered KV sequence before attention. Independently normalized KV partials must never be merged.

VEA-B is optimized first for the single-GPU Q128 shape. A future SM120/SM100 or Ulysses schedule may choose Q64 or a different producer geometry, but it must retain this numerical and descriptor contract and pass separately tuned gates.

9.6 KiB Raw Blame History