h3-blackwell-runtime/research/vortex_exact_attention/DESIGN.md

# 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:

```text
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 `2*H*N^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

1. Phase 2A register-allocation, shared-memory, handoff, sanitizer, and dual-pipe
   probes are complete; see `PHASE2A_CAPABILITY_REPORT.md`.
2. Retain CUDA block-scope mbarriers; do not use the rejected inline named
   barrier primitive.
3. Implement VEA-B only for one aligned short shape and exact D=128 arithmetic.
4. 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.