daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Select Vortex exact attention architecture

Browse source
This commit is contained in:
Daniel Maddern 2026-08-26 13:39:20 +07:00
parent c5028427d7
commit b842badfc1
9 changed files with 593 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

168
research/vortex_exact_attention/DESIGN.md Normal file
View file

@ -0,0 +1,168 @@
# State-Ownership Architecture Decision
## Decision
Prototype VEA-B, the one-CTA Q128 paired-owner pipeline, after its capability
probes pass. VEA-A is the fallback if 320-thread role allocation fails. VEA-C is
the higher-risk second prototype only after DSM placement and handoff latency
are measured. These are projected designs, not implemented or achieved results.
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).
Four independent producer/consumer named-barrier pairs avoid a CTA-wide barrier
inside the loop. K/V TMA barriers are separate. Projected synchronization cost
is 6-15 ms.
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. Run register-allocation, shared-memory, named-barrier, and dual-pipe overlap
probes without model inference.
2. Reject VEA-B if it spills, cannot allocate 50-52 KiB predictably, or cannot
demonstrate disjoint-role overlap.
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.

17
research/vortex_exact_attention/EXPERIMENT_LOG.md
View file

@ -46,3 +46,20 @@ P1/P2 source is partial and mixed in
`research/shared_cuda_experimental_worktree/`. P3 source and launcher are
complete under `research/sage2_temporal_pair/`. Exact P0 variant patches are
unrecoverable; historical commands and environments must not be invented.
## 2026-08-26: Phase 1 State-Ownership Decision
Status: `research_retained`
Evaluated three architectures that permanently separate QK/softmax state from
PV/output state: VEA-A Q64 split-D intra-CTA, VEA-B Q128 paired owners, and
VEA-C two-CTA DSM ownership split. All retain strict tile order, exact FP8 score
bytes, per-tile RO rescaling, two-level PV accumulation, and final conversion.
Decision: recommend VEA-B after register/shared/barrier/dual-pipe capability
probes. Its projected `180-207 ms` mainloop screening range has a credible sub-190 path
without depending on unverified DSM. VEA-A is the fallback; VEA-C is conditional
on cluster residency and DSM latency. These are projections only.
No kernel was implemented, no inference was run, and production dispatch and
services were not modified.

83
research/vortex_exact_attention/HARDWARE_CAPABILITIES.md Normal file
View file

@ -0,0 +1,83 @@
# Hardware Capabilities
## Interpretation Rule
Blackwell-native means using target-specific scheduling, state ownership, or
data movement supported by the target. Compiling the SM89 algorithm for
`sm_121a` is not sufficient. SM100, SM120, and SM121 are separate schedules.
## Capability Matrix
| Capability | SM100 / B200 | SM120 / RTX PRO 6000 | SM121 / GB10 |
| --- | --- | --- | --- |
| Compute capability | 10.0 | 12.0 | 12.1 |
| TMA global/shared tensor copies | Documented | Documented | Documented |
| Native matrix programming model | `tcgen05.mma` with TMEM | warp-level `mma.sync` | SM120-family warp-level `mma.sync` |
| Programmer-visible TMEM | Yes | Not exposed | Not exposed |
| Native UMMA/tcgen05 INT8 | Exposed | Not exposed | Not exposed |
| Sage-compatible INT8 path | `tcgen05.kind::i8` can be investigated | supported classic INT8 `mma.sync` | supported classic INT8 `mma.sync` |
| FP8 PV path | `tcgen05` FP8 forms | warp-level FP8 `mma.sync` | warp-level FP8 `mma.sync` |
| Hopper WGMMA | Not a portable target path | Not a target path | CUDA 13 rejects for `sm_121a` |
| Registers per SM / max per thread | 64K / 255 | 64K / 255 | 64K / 255 |
| Shared memory per SM / block | 228 / 227 KiB | 100 / 99 KiB | 100 / 99 KiB |
| Maximum resident warps | 64 | 48 | 48 |
| Maximum resident threads | 2048 | 1536 | 1536 |
The public SM120/121 Blackwell MMA interface does not expose a Sage2-compatible
UMMA INT8 QK operation. Exact SM121 work therefore retains the supported INT8
`mma.sync` arithmetic while making TMA staging, role scheduling, ownership, and
handoff Blackwell-specific. FP8/FP4 QK would define `vortex_fast`, not
`vortex_exact`.
## SM121 Facts From Retained Evidence
- The current 128-thread kernel uses 255 registers/thread and 32 KiB dynamic
shared memory.
- Two CTAs provide eight resident warps, approximately `16.67%` of the 48-warp
ceiling.
- A 168-register cap reaches three CTAs but spills catastrophically; compiler
lifetime cleanup is not enough.
- CUDA 13 rejects `wgmma.fence`, `wgmma.mma_async`, `wgmma.commit_group`, and
`wgmma.wait_group` for `sm_121a`.
- CUTLASS 4.6 executes SM120-family block-scaled narrow MMA on GB10, but that is
not an exact replacement for Sage2 INT8 QK.
- TMA producer/consumer scheduling is available, but role-specific register
redistribution, named-barrier cost, and concurrent INT8/FP8 issue must be
measured rather than assumed.
## Required Capability Probes Before Prototype
1. Compile empty 256- and 320-thread role-specialized CTAs; record per-thread
registers, aggregate register allocation, occupancy, and spill traffic.
2. Test whether any supported register-allocation control can give producer and
consumer warps different practical budgets on SM121.
3. Measure one-CTA residency at 48, 52, 56, 64, and 68 KiB shared memory.
4. Measure TMA plus named-barrier two-slot handoff for 4 and 8 KiB payloads.
5. Verify simultaneous INT8 QK and FP8 PV warps produce actual overlap in pipe
counters, not merely alternating aggregate utilization.
6. Establish thread-block cluster/DSM support, placement, simultaneous
residency, and remote shared-memory latency on GB10 before accepting VEA-C.
7. Record sustained clocks and power behavior for 128-, 256-, and 320-thread
CTAs.
## Portability Policy
- SM121 is first and must compile specifically for `sm_121a`.
- SM120 may share API and contract code, but requires its own geometry,
resource model, profiler evidence, and acceptance gates.
- SM100 should use a separate `tcgen05`/TMEM design. Porting the SM121 register
accumulator schedule unchanged would discard the principal SM100 capability.
## Sources
- NVIDIA CUDA GPU Compute Capability: https://developer.nvidia.com/cuda-gpus
- CUDA Programming Guide: https://docs.nvidia.com/cuda/cuda-programming-guide/
- PTX ISA: https://docs.nvidia.com/cuda/parallel-thread-execution/
- CUTLASS Blackwell functionality:
https://github.com/NVIDIA/cutlass/blob/main/media/docs/cpp/blackwell_functionality.md
- CUTLASS `mma_sm100_umma.hpp`, `mma_sm120.hpp`, and architecture config.
- Local measurements: `SAGE2_BLACKWELL_DESIGN.md` and retained P0/post-FC2 NCU
reports.
Documented API support does not establish throughput, latency, queue depth, or
bit identity. Those remain target-specific measurements.

11
research/vortex_exact_attention/NUMERICAL_CONTRACT.md
View file

@ -88,7 +88,10 @@ For each tile, preserve this sequence and lane reduction mapping:
and `ex2.approx.ftz.f32`.
7. Add exponent values to `d` in the reference fragment order.
8. Convert probabilities FP32 to E4M3 RN, saturate-finite.
9. Execute PV for this tile before advancing.
9. Execute PV for this tile before the output owner numerically consumes the
next tile. A QK producer may prepare one bounded future tile only after it
publishes all exact tile state; the output owner must consume epochs in
strict order and apply every RO rescale/PV update before epoch `t+1`.
Ownership may move between roles only if the transferred bytes cause the same
instruction-level updates in the same order. Reassociating maxima, denominator
@ -115,8 +118,10 @@ PV has two exact accumulation levels:
2. The FP16 fragment is unpacked and added to persistent FP32 `RO` in the
reference order across tiles.
Changing either level, carrying a material probability tensor, or combining KV
partials across CTAs violates `vortex_exact`.
Changing either level, carrying a sequence-scale/global probability tensor, or
combining KV partials across CTAs violates `vortex_exact`. A bounded one- or
two-tile shared-memory handoff of the exact post-conversion FP8 bytes is allowed
only to transfer ownership; it is not an arithmetic or storage-format boundary.
## Finalization And Tails

92
research/vortex_exact_attention/PERFORMANCE_MODEL.md Normal file
View file

@ -0,0 +1,92 @@
# Performance Model
All candidate values are projections. No Vortex kernel has been implemented or
timed.
## Baseline
| Metric | Value |
| --- | ---: |
| Mainloop p50 | `237.089 ms` uninstrumented |
| Mainloop plus Sage quantization/preparation | approximately `257.7 ms` |
| Full owned boundary including Q/K RMSNorm and RoPE | approximately `269.9 ms` |
| Grid / CTA | 16,576 CTAs / 128 threads |
| KV iterations | 591 |
| Registers / shared memory | 255/thread / 32 KiB |
| Resident CTAs / warps | 2 / 8 |
| No-eligible cycles | `63.53%` |
| Combined INT8 QK + FP8 PV tensor activity | `75.54%` |
| L2 hit / memory throughput | `98.85%` / approximately `31.74%` |
The model targets shorter fixed-dependency chains and disjoint live-state
ownership. It assigns no speedup to launch removal, DRAM, tails, or padding.
## Resource Projection
| Candidate | Threads | Role-weighted registers/CTA | Shared memory | CTAs/SM projection | Active warps | QK owns RO? | PV owns m/d? |
| --- | ---: | ---: | ---: | ---: | ---: | --- | --- |
| VEA-A | 256 | 24-32K | 51-52 KiB | 1 | 8 | No | No |
| VEA-B | 320 | 56,320 at uniform 176 registers | 50-52 KiB aliased | 1 | 10 | No | No |
| VEA-C pair | 384 total | 50-60K pair | 50+18 KiB | unverified cluster | 12 | No | No |
Logical ownership shortens per-warp dependency chains but does not prove
role-weighted physical register allocation. VEA-B therefore uses the uniform
largest-role projection; compiled allocation above 200 registers/thread or any
material local spilling rejects it.
## Latency Budget
| Component | VEA-A | VEA-B | VEA-C | Basis |
| --- | ---: | ---: | ---: | --- |
| Recovered dependency-stall fraction assumption | 22-35% | 30-42% | 38-50% | screening assumption, not measured |
| Handoff and synchronization | 8-18 ms | 6-15 ms | 12-30 ms | unmeasured named-barrier/DSM budget |
| Projected mainloop | `192-222 ms` | `180-207 ms` | `174-210 ms` | heuristic screen, not achieved |
| Non-overlapped exact preparation | 25-28 ms | 25-28 ms | 25-28 ms | P1 entry 12.33 + K mean 2.37 + K quant 3.83 + P2 V 6.39 = 24.92 ms, plus margin |
| Projected complete attention | `217-250 ms` | `205-235 ms` | `199-238 ms` | not achieved |
The screening range starts from `237.089 ms`, assumes a candidate-specific
fraction of the measured `63.53%` no-eligible region becomes useful overlap,
then adds the handoff budget. Scheduler percentages do not translate linearly
to time, so this is a rejection model rather than a performance prediction.
Capability microbenchmarks must replace every assumption before implementation.
## Rejection Gates
- Compiled local spills materially above zero.
- Mainloop model lower bound cannot beat `220 ms`.
- Shared memory prevents projected resident work.
- QK and PV cannot issue concurrently across roles.
- Handoff requires a global N-by-N score/probability tensor.
- Handoff changes FP8 probability bytes, RO rescale order, tile order, or final
reciprocal/conversion.
- p95 has no bounded synchronization path.
VEA-B is selected because it is the only intra-CTA design with a credible sub-190 range
that does not depend on unverified DSM.
## End-To-End Scenario Math
With baseline sampling `255.447 s` and attention share approximately `159.3 s`:
```text
sampling(s) = 255.447 - 159.3 + 159.3 / attention_speedup
```
| Attention speedup scenario | Estimated sampling | Status |
| ---: | ---: | --- |
| 1.15x minimum viable | `234.67 s` | target only |
| 1.30x strong | `218.70 s` | target only |
| 1.50x research-grade | `202.35 s` | target only |
| 2.00x breakthrough | `175.80 s` | target only |
| 2.70x stretch | `155.15 s` | target only; approximately half the original `310.11 s` baseline |
These scenarios are not Vortex results and do not account for interactions,
clock changes, or non-attention regressions.
## Required Candidate Evidence
Executable candidates must use alternating AB/BA runs with one resident model
and report p10/p50/p90/p95, registers, occupancy/resident CTAs, eligible and
no-eligible cycles, INT8/FP8 utilization, instruction mix, shared traffic and
conflicts, L2 requests/hit rate, off-chip proxy, spills, synchronization stalls,
and kernel count. One-sample smoke timings cannot advance a candidate.

9
research/vortex_exact_attention/README.md
View file

@ -25,7 +25,7 @@ Such work belongs in a future `vortex_fast` backend with a separate contract.
## Status
- Phase 0: imported and verified retained Sage2 exactness specification.
- Phase 1: design decision and architecture model in progress.
- Phase 1: design decision complete; VEA-B is the recommended prototype.
- Kernel implementation: not started.
- Production dispatch: unchanged; no `vortex_exact` route exists.
@ -73,6 +73,13 @@ p95, block-level, and resident-service gates.
- `tests/`: future parity and sanitizer harnesses.
- `tools/`: project-local verification/modeling utilities.
Phase 1 selects a Q128, ten-warp paired-owner pipeline in which four QK/softmax
warps own `RS/RS_f8/m/d`, four distinct PV warps own `RO`, and two producer
warps own K/V staging. This differs structurally from P3 because complete state
is transferred or partitioned rather than retained in every compute warp. Its
`180-207 ms` mainloop range is a heuristic screening model, not a measured
result; physical register allocation and overlap remain unverified.
## Provenance
The reference is the public SageAttention repository at the pinned commit above.

71
research/vortex_exact_attention/benchmarks/architecture_model.json Normal file
View file

@ -0,0 +1,71 @@
{
"schema": "vortex-exact-attention-architecture-model",
"version": 1,
"status": "projected_not_measured",
"baseline": {
"device": "NVIDIA GB10", "architecture": "sm_121a", "tokens": 37810,
"heads": 56, "head_dimension": 128, "mainloop_p50_ms": 237.089,
"grid": [296, 56, 1], "block": [32, 4, 1], "kv_iterations": 591,
"stages": 2, "registers_per_thread": 255, "dynamic_shared_kib": 32,
"occupancy_percent": [16.65, 16.83], "no_eligible_cycles_percent": 63.53,
"l2_hit_percent": 98.85, "memory_throughput_percent": 31.74
},
"hard_gates": {
"mainloop_model_must_beat_ms": 220, "preferred_model_below_ms": 190,
"byte_exact": true, "material_score_tensor": false,
"material_local_spills": false,
"kv_tile_order": "strictly_increasing_0_through_590",
"pv_accumulation": "two_ordered_k32_fp16_mma_then_ordered_fp32_ro_add"
},
"model_basis": {
"kind": "heuristic_rejection_screen",
"formula": "baseline minus assumed recovered fraction of no-eligible region plus handoff budget",
"exact_preparation_components_ms": {"p1_entry": 12.33, "k_mean": 2.37, "k_quant": 3.83, "p2_v": 6.39},
"exact_preparation_ms": [25, 28],
"unresolved": ["compiled registers", "named barrier latency", "INT8_FP8 overlap", "DSM placement and latency"]
},
"candidates": [
{
"id": "VEA-A", "name": "Q64 split-D intra-CTA pipeline", "recommended_order": 3,
"geometry": {"q": 64, "k": 64, "d": 128, "threads": 256, "warps": 8},
"ownership": {"qk_fragments": "warps_1_2", "m_d": "warps_1_2", "output_accumulation": "warps_3_6_split_d64", "load_control": "warps_0_7"},
"handoff": {"mechanism": "two_slot_shared_ring_named_barriers", "payload": "fp8_probability_bytes_plus_fp32_ro_rescale", "bounded_tile_local": true, "consumer_order": "strict_epoch_order", "qk_and_pv_fragments_simultaneously_live": true, "same_warp": false},
"registers_per_thread": {"load": [32, 48], "qk_softmax": [96, 128], "pv_output": [88, 120], "control": [32, 64]},
"shared_kib": {"q": 8, "k_double": 16, "v_double": 16, "score_double": 8, "metadata": [3, 4], "total": [51, 52]},
"registers_per_cta_uniform_projection": 32768,
"resident_ctas_per_sm_projected": 1, "synchronization_ms_projected": [8, 18],
"recovered_no_eligible_fraction": [0.22, 0.35],
"mainloop_ms_projected": [192, 222], "complete_attention_ms_projected": [217, 250],
"capability_dependencies": ["named_barrier_cost", "INT8_FP8_overlap"],
"structural_difference": "RO is removed from QK warps and split by D; m/d and scores are removed from PV warps"
},
{
"id": "VEA-B", "name": "Q128 paired QK/PV owner pipeline", "recommended_order": 1,
"geometry": {"q": 128, "k": 64, "d": 128, "threads": 320, "warps": 10},
"ownership": {"qk_fragments": "warps_2_5", "m_d": "warps_2_5", "output_accumulation": "warps_6_9", "load_control": "warps_0_1"},
"handoff": {"mechanism": "four_paired_two_slot_shared_rings_named_barriers", "payload": "fp8_probability_bytes_plus_fp32_ro_rescale_and_final_reciprocal", "bounded_tile_local": true, "consumer_order": "strict_epoch_order", "qk_and_pv_fragments_simultaneously_live": true, "same_warp": false},
"registers_per_thread": {"k_producer": [32, 48], "v_producer": [40, 56], "qk_softmax": [96, 128], "pv_output": [144, 176]},
"logical_role_weighted_registers": 39000, "registers_per_cta_uniform_projection": 56320, "compiled_register_limit_per_thread": 200,
"shared_kib": {"unaliased": [66, 68], "alias_saving": 16, "total": [50, 52]},
"resident_ctas_per_sm_projected": 1, "synchronization_ms_projected": [6, 15],
"recovered_no_eligible_fraction": [0.30, 0.42],
"mainloop_ms_projected": [180, 207], "complete_attention_ms_projected": [205, 235],
"capability_dependencies": ["compiled_registers_at_or_below_200", "zero_material_spills", "named_barrier_cost", "INT8_FP8_overlap"],
"structural_difference": "QK/softmax and RO/PV have permanent disjoint owners; P1/P2 may feed slots but are not standalone boundaries"
},
{
"id": "VEA-C", "name": "two-CTA DSM ownership split", "recommended_order": 2,
"geometry": {"q": 128, "k": 64, "d": 128, "ctas": 2, "threads_total": 384, "warps_total": 12},
"ownership": {"qk_fragments": "producer_cta_warps_1_4", "m_d": "producer_cta_warps_1_4", "output_accumulation": "consumer_cta_warps_1_4", "load_control": "producer_and_consumer_warps_0_5"},
"handoff": {"mechanism": "two_slot_dsm_ring_cluster_barriers", "payload": "fp8_probability_bytes_plus_fp32_ro_rescale_and_final_reciprocal", "bounded_tile_local": true, "consumer_order": "strict_epoch_order", "qk_and_pv_fragments_simultaneously_live": true, "same_warp": false, "same_cta": false},
"registers_per_thread": {"producer_load_publish": [32, 48], "qk_softmax": [96, 128], "consumer_v": [40, 56], "pv_output": [144, 176], "consumer_store": [40, 64]},
"registers_per_cluster_projected": [50000, 60000], "shared_kib": {"producer": 50, "consumer": 18, "cluster_total": 68},
"resident_cluster_projection": "unverified", "synchronization_ms_projected": [12, 30],
"recovered_no_eligible_fraction": [0.38, 0.50],
"mainloop_ms_projected": [174, 210], "complete_attention_ms_projected": [199, 238],
"capability_dependencies": ["DSM_support", "cluster_simultaneous_residency", "DSM_round_trip_latency", "INT8_FP8_overlap"],
"structural_difference": "QK/softmax state and RO/PV state reside in separate CTAs and cross an explicit DSM byte handoff"
}
],
"decision": {"recommended_prototype": "VEA-B", "fallback": "VEA-A", "conditional_second": "VEA-C", "kernel_implementation_started": false, "production_dispatch_changed": false}
}

48
research/vortex_exact_attention/tests/test_models.py Normal file
View file

@ -0,0 +1,48 @@
import importlib.util
import copy
from pathlib import Path
import unittest
class ModelContracts(unittest.TestCase):
@classmethod
def setUpClass(cls):
path = Path(__file__).resolve().parents[1] / "tools/validate_models.py"
spec = importlib.util.spec_from_file_location("validate_vortex_models", path)
cls.module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(cls.module)
def test_retained_reference_and_architecture_models(self):
self.module.validate()
def test_rejects_same_warp_state_ownership(self):
inventory, model = self.module.load_models()
changed = copy.deepcopy(model)
changed["candidates"][0]["handoff"]["same_warp"] = True
with self.assertRaisesRegex(RuntimeError, "one warp"):
self.module.validate_documents(inventory, changed)
def test_rejects_inconsistent_shared_memory(self):
inventory, model = self.module.load_models()
changed = copy.deepcopy(model)
changed["candidates"][1]["shared_kib"]["total"] = [52, 56]
with self.assertRaisesRegex(RuntimeError, "VEA-B shared-memory"):
self.module.validate_documents(inventory, changed)
def test_rejects_inconsistent_screening_arithmetic(self):
inventory, model = self.module.load_models()
changed = copy.deepcopy(model)
changed["candidates"][0]["mainloop_ms_projected"] = [192, 220]
with self.assertRaisesRegex(RuntimeError, "screening arithmetic"):
self.module.validate_documents(inventory, changed)
def test_rejects_measured_result_label(self):
inventory, model = self.module.load_models()
changed = copy.deepcopy(model)
changed["status"] = "achieved"
with self.assertRaisesRegex(RuntimeError, "measurements"):
self.module.validate_documents(inventory, changed)
if __name__ == "__main__":
unittest.main()

98
research/vortex_exact_attention/tools/validate_models.py Normal file
View file

@ -0,0 +1,98 @@
"""Validate Vortex Exact Attention's retained reference and design models."""
from __future__ import annotations
import hashlib
import json
import math
from pathlib import Path
PROJECT = Path(__file__).resolve().parents[1]
ROOT = PROJECT.parents[1]
def load_models() -> tuple[dict, dict]:
inventory = json.loads((PROJECT / "benchmarks/reference_inventory.json").read_text())
model = json.loads((PROJECT / "benchmarks/architecture_model.json").read_text())
return inventory, model
def validate_documents(inventory: dict, model: dict, root: Path = ROOT) -> None:
if len(inventory["artifacts"]) != 10:
raise RuntimeError("the imported reference inventory must contain ten artifacts")
for artifact in inventory["artifacts"]:
path = root / artifact["path"]
if path.stat().st_size != artifact["size_bytes"]:
raise RuntimeError(f"size mismatch: {path}")
if hashlib.sha256(path.read_bytes()).hexdigest() != artifact["sha256"]:
raise RuntimeError(f"SHA-256 mismatch: {path}")
if model["status"] != "projected_not_measured":
raise RuntimeError("design estimates must not be labeled as measurements")
if len(model["candidates"]) < 3:
raise RuntimeError("at least three state-ownership designs are required")
preparation_components = model["model_basis"]["exact_preparation_components_ms"]
preparation_range = model["model_basis"]["exact_preparation_ms"]
if preparation_range[0] != math.ceil(sum(preparation_components.values())):
raise RuntimeError("exact-preparation lower bound is inconsistent")
candidate_ids = [candidate["id"] for candidate in model["candidates"]]
if set(candidate_ids) != {"VEA-A", "VEA-B", "VEA-C"} or len(candidate_ids) != 3:
raise RuntimeError("the three reviewed architecture IDs must be unique and retained")
for candidate in model["candidates"]:
ownership = candidate["ownership"]
if ownership["qk_fragments"] == ownership["output_accumulation"]:
raise RuntimeError(f"{candidate['id']} does not separate QK and RO ownership")
if ownership["m_d"] == ownership["output_accumulation"]:
raise RuntimeError(f"{candidate['id']} does not separate m/d and RO ownership")
if candidate["handoff"]["same_warp"]:
raise RuntimeError(f"{candidate['id']} retains QK and PV state in one warp")
if not candidate["handoff"]["bounded_tile_local"]:
raise RuntimeError(f"{candidate['id']} uses a material handoff")
if candidate["handoff"]["consumer_order"] != "strict_epoch_order":
raise RuntimeError(f"{candidate['id']} does not preserve tile consumption order")
if candidate["mainloop_ms_projected"][0] >= 220:
raise RuntimeError(f"{candidate['id']} cannot plausibly beat the model gate")
recovery = candidate["recovered_no_eligible_fraction"]
synchronization = candidate["synchronization_ms_projected"]
baseline = model["baseline"]["mainloop_p50_ms"]
no_eligible = model["baseline"]["no_eligible_cycles_percent"] / 100.0
expected_mainloop = [
round(baseline * (1.0 - no_eligible * recovery[1]) + synchronization[0]),
round(baseline * (1.0 - no_eligible * recovery[0]) + synchronization[1]),
]
if candidate["mainloop_ms_projected"] != expected_mainloop:
raise RuntimeError(f"{candidate['id']} screening arithmetic is inconsistent")
expected_complete = [value + margin for value, margin in zip(
candidate["mainloop_ms_projected"], preparation_range, strict=True,
)]
if candidate["complete_attention_ms_projected"] != expected_complete:
raise RuntimeError(f"{candidate['id']} complete-attention arithmetic is inconsistent")
if not candidate["capability_dependencies"]:
raise RuntimeError(f"{candidate['id']} omits unresolved capability dependencies")
by_id = {candidate["id"]: candidate for candidate in model["candidates"]}
a_shared = by_id["VEA-A"]["shared_kib"]
if a_shared["total"] != [
a_shared["q"] + a_shared["k_double"] + a_shared["v_double"]
+ a_shared["score_double"] + value for value in a_shared["metadata"]
]:
raise RuntimeError("VEA-A shared-memory arithmetic is inconsistent")
b_shared = by_id["VEA-B"]["shared_kib"]
if b_shared["total"] != [value - b_shared["alias_saving"] for value in b_shared["unaliased"]]:
raise RuntimeError("VEA-B shared-memory alias arithmetic is inconsistent")
c_shared = by_id["VEA-C"]["shared_kib"]
if c_shared["cluster_total"] != c_shared["producer"] + c_shared["consumer"]:
raise RuntimeError("VEA-C shared-memory arithmetic is inconsistent")
decision = model["decision"]
if decision["kernel_implementation_started"] or decision["production_dispatch_changed"]:
raise RuntimeError("Phase 1 must remain design-only and isolated")
def validate() -> None:
validate_documents(*load_models())
if __name__ == "__main__":
validate()
print("Vortex exact reference and architecture models are valid")