diff --git a/research/vortex_exact_attention/DESIGN.md b/research/vortex_exact_attention/DESIGN.md new file mode 100644 index 0000000..ad54b64 --- /dev/null +++ b/research/vortex_exact_attention/DESIGN.md @@ -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. diff --git a/research/vortex_exact_attention/EXPERIMENT_LOG.md b/research/vortex_exact_attention/EXPERIMENT_LOG.md index 224bfbf..772d3e4 100644 --- a/research/vortex_exact_attention/EXPERIMENT_LOG.md +++ b/research/vortex_exact_attention/EXPERIMENT_LOG.md @@ -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. diff --git a/research/vortex_exact_attention/HARDWARE_CAPABILITIES.md b/research/vortex_exact_attention/HARDWARE_CAPABILITIES.md new file mode 100644 index 0000000..d9afa02 --- /dev/null +++ b/research/vortex_exact_attention/HARDWARE_CAPABILITIES.md @@ -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. diff --git a/research/vortex_exact_attention/NUMERICAL_CONTRACT.md b/research/vortex_exact_attention/NUMERICAL_CONTRACT.md index d622348..e34f896 100644 --- a/research/vortex_exact_attention/NUMERICAL_CONTRACT.md +++ b/research/vortex_exact_attention/NUMERICAL_CONTRACT.md @@ -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 diff --git a/research/vortex_exact_attention/PERFORMANCE_MODEL.md b/research/vortex_exact_attention/PERFORMANCE_MODEL.md new file mode 100644 index 0000000..d54c475 --- /dev/null +++ b/research/vortex_exact_attention/PERFORMANCE_MODEL.md @@ -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. diff --git a/research/vortex_exact_attention/README.md b/research/vortex_exact_attention/README.md index 5c9eb70..a21ea47 100644 --- a/research/vortex_exact_attention/README.md +++ b/research/vortex_exact_attention/README.md @@ -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. diff --git a/research/vortex_exact_attention/benchmarks/architecture_model.json b/research/vortex_exact_attention/benchmarks/architecture_model.json new file mode 100644 index 0000000..dfe7b09 --- /dev/null +++ b/research/vortex_exact_attention/benchmarks/architecture_model.json @@ -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} +} diff --git a/research/vortex_exact_attention/tests/test_models.py b/research/vortex_exact_attention/tests/test_models.py new file mode 100644 index 0000000..9f06ec6 --- /dev/null +++ b/research/vortex_exact_attention/tests/test_models.py @@ -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() diff --git a/research/vortex_exact_attention/tools/validate_models.py b/research/vortex_exact_attention/tools/validate_models.py new file mode 100644 index 0000000..71e2745 --- /dev/null +++ b/research/vortex_exact_attention/tools/validate_models.py @@ -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")