Define Vortex exact attention contract

2026-08-26 13:20:33 +07:00 · 2026-08-26 13:20:33 +07:00 · c5028427d7
commit c5028427d7
parent a29b8960b0
7 changed files with 354 additions and 0 deletions
--- a/research/vortex_exact_attention/EXPERIMENT_LOG.md
+++ b/research/vortex_exact_attention/EXPERIMENT_LOG.md
@ -0,0 +1,48 @@
+# Experiment Log
+
+This is append-only. Projected ranges are design estimates, never achieved
+results.
+
+## 2026-08-26: Phase 0 Import And Verification
+
+Status: `research_retained`
+
+Hypothesis: the retained Sage inventory is sufficient to import the exact
+observable contract without repeating prior experiments.
+
+Actions:
+
+- Verified public SageAttention commit
+  `d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5` exists.
+- Rehashed ten retained contract, parity, timing, and NCU artifacts; all matched
+  `benchmarks/reference_inventory.json`.
+- Imported Q/K preparation, K smoothing, INT8 quantization, QK traversal,
+  online-softmax, FP8 V/P, two-level PV accumulation, final conversion, and
+  ragged-tail contracts into `NUMERICAL_CONTRACT.md`.
+- Recorded missing production binary/checkpoint hashes and missing standalone
+  canonical input tensors as blockers rather than fabricating provenance.
+- Did not execute inference, modify dispatch, or touch the production service.
+
+Decision: Phase 0 is verified as an imported specification. No kernel work is
+authorized until Phase 1 selects a structural state-ownership architecture.
+
+## Closed Prior Work
+
+| Work | Result | Preservation / reopening rule |
+| --- | --- | --- |
+| P0 register caps | Three-CTA threshold at 168 registers caused catastrophic spills; no gain | Timing/JSON/NCU retained; exact candidate source missing and must not be claimed reproducible |
+| P0 scope/lifetime narrowing | Fewer spills, approximately `0.06%` gain | Reopen only with structural ownership change |
+| P0 score aliasing | Approximately `0.01%`, noise | Do not repeat unchanged |
+| P0 early K prefetch | Approximately `0.30%` slower | Reopen only with a different producer/consumer pipeline |
+| P0 softmax interleaving | Neutral/slower | Do not repeat unchanged |
+| P0 padded V shared memory | Same excessive wavefronts, no gain | Reopen only if staging instructions/dataflow change |
+| P1 entry fusion | `20.9-23.4%` isolated; only `0.535-0.863%` complete-block | Reuse partial implementation only inside an owned attention pipeline |
+| P2 direct V preparation | `39.48%`, `4.17 ms`; only `0.909%` of block | Reuse partial implementation only inside an owned attention pipeline |
+| P3 temporal pairing | Exact/sanitizer-clean; approximately `0.10%`; private state remained live | Complete patch and launcher retained; new work must transfer/partition ownership |
+| Hopper WGMMA retarget | CUDA 13 rejects WGMMA for `sm_121a` | Reopen only with changed ISA/toolchain support |
+| SM121 UMMA INT8 QK | Public SM120/121 Blackwell MMA path does not expose a Sage-compatible UMMA INT8 operation | Retain supported INT8 `mma.sync` unless capability changes |
+
+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.
--- a/research/vortex_exact_attention/NUMERICAL_CONTRACT.md
+++ b/research/vortex_exact_attention/NUMERICAL_CONTRACT.md
@ -0,0 +1,162 @@
+# Exact Numerical Contract
+
+## Authority And Verification
+
+This document imports the retained Sage research inventory rather than
+reconstructing Phase 0. The reference is SageAttention 2.2.0 commit
+`d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5`, verified present in the public
+`thu-ml/SageAttention` repository on 2026-08-26. Ten retained local artifacts
+were rehashed; every SHA-256 matched `benchmarks/reference_inventory.json`.
+
+Observable equality means byte equality at every named boundary. Tolerance,
+ULP, cosine, visual, and statistical comparisons cannot accept `vortex_exact`.
+
+## Canonical Boundary
+
+- Batch 1, sequence 37,810, 56 heads, head dimension 128.
+- BF16 projection-strided NHD Q/K/V, non-causal self-attention.
+- Interleaved QKV row width 21,504; Q, K, and V offsets are 0, 7,168,
+  and 14,336 BF16 elements.
+- Logical Q/K/V shape `[1,N,56,128]`, with canonical element strides
+  `[813066240,21504,128,1]`.
+- Output is contiguous BF16 NHD `[1,N,56,128]`, directly reshapeable to
+  token-major `[N,7168]` without a copy.
+
+## Preparation
+
+1. Apply per-head Q/K RMSNorm with epsilon `1e-5` in the production reduction
+   order.
+2. Normalize in FP32 and round the normalized value to BF16 RN.
+3. Apply split-half RoPE to dimensions 0-95 using BF16 coefficients and FP32
+   arithmetic; round results to BF16 RN. Dimensions 96-127 remain normalized
+   but unrotated.
+4. Compute the exact global sequence mean for each K head/dimension and subtract
+   it before K quantization. Despite the local `smooth_k=False` call, the pinned
+   Sage dispatcher selects the FP8 path whose effective default performs K
+   smoothing. The observable mean bytes and scales are part of the contract.
+
+The production Comfy Kitchen RMSNorm/RoPE binary and its hash are not retained.
+P1 proved byte equality over retained randomized boundaries and blocks 0/24/49,
+but does not prove exceptional-value behavior for all possible inputs.
+
+## Q And K Quantization
+
+Q uses one INT8 scale for each 32-query-row warp group inside a 128-row CTA.
+K uses one INT8 scale for each 64-key-row block after mean subtraction. Each
+maximum spans all valid rows in the group and all 128 dimensions:
+
+```text
+amax = max(1e-7f, max(abs(x)))
+scale = amax / 127.0f
+q8 = cvt.rni.sat.s8.f32(x * (127.0f / amax))
+```
+
+Canonical scale shapes are Q `[1,56,1184]` and K `[1,56,591]`. Invalid tail
+rows contribute zero to maxima and are not written.
+
+## INT8 QK And Score Domain
+
+- Each warp owns 32 query rows and traverses 64 key rows at a time.
+- Head dimension 128 is accumulated as four ordered K=32 signed INT8
+  `mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32` operations.
+- INT32 scores convert to FP32 using the reference conversion instruction.
+- Attention scale is `1/sqrt(128)`, converted to base 2 with `log2(e)`.
+- The regular-tile and final-two-tile placement of dequantization multiplication
+  must remain exactly as in the pinned mainloop. Algebraic equivalence is not
+  sufficient if it changes rounding.
+
+## Online Softmax
+
+KV tiles are visited strictly in increasing order, tile 0 through tile 590.
+Per query row, initialize:
+
+```text
+RO = 0.0f
+m = -5000000.0f
+d = 1.0f
+```
+
+For each tile, preserve this sequence and lane reduction mapping:
+
+1. Compute scaled score fragments.
+2. Compute tile-local maximum, subtract FP8 offset `8.807`, and reduce with XOR
+   lane masks 1 then 2.
+3. Set `m_new = max(m_old, tile_max)`.
+4. Compute `old_scale = ex2.approx.ftz.f32(m_old - m_new)`.
+5. Rescale `d` and every persistent FP32 output accumulator by `old_scale`.
+6. Evaluate each probability with the reference fused multiply-add placement
+   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.
+
+Ownership may move between roles only if the transferred bytes cause the same
+instruction-level updates in the same order. Reassociating maxima, denominator
+sums, or output accumulation is forbidden.
+
+## V And PV
+
+V uses a per `(batch,head,dimension)` global sequence maximum:
+
+```text
+amax = max(abs(BF16 V))
+v_scale = amax / 2.25f
+V_fp8 = cvt.rn.satfinite.e4m3(BF16_V * (2.25f / amax))
+```
+
+The retained V layout is contiguous `[1,128,56,37824]` with 64-row padding and
+the per-16-row permutation `[0,1,8,9,2,3,10,11,4,5,12,13,6,7,14,15]`.
+
+PV has two exact accumulation levels:
+
+1. Within a 64-key tile, two ordered K=32
+   `mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16` operations accumulate
+   in an FP16 instantaneous fragment.
+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`.
+
+## Finalization And Tails
+
+After tile 590, reduce `d` with XOR masks 1 then 2, apply
+`rcp.approx.ftz.f32`, multiply `RO` by the reciprocal, multiply by the FP32 V
+scale, and convert FP32 pairs to BF16 RN. Store only valid query rows.
+
+At N=37,810 the final Q128 CTA and final K64 tile each have 50 valid rows.
+Invalid K score lanes receive `-5000000.0f` at the same point as the reference;
+invalid V rows are zero; invalid Q rows suppress stores. Tail cost is below
+1 ms and is not an optimization target.
+
+## Required Golden Gates
+
+| Gate | Existing oracle | Required future fixture |
+| --- | --- | --- |
+| Random short and ragged shapes | P1/P2/P3 retained JSON | Self-contained tensors plus exceptional values |
+| Prepared Q/K bytes and scales | P1 parity artifacts | Stored SHA-256 per tensor |
+| FP8 V bytes and scales | P2 randomized/repeat artifacts | Stored SHA-256 per tensor |
+| Block 24 attention | SHA `4c666c20f5f8f651158a2ced33ccff08f3bada07665c595b99008d171db30574` | Captured input Q/K/V and metadata |
+| Blocks 0/24/49 | Strided-NHD and P1 artifacts | Self-contained boundary tensors |
+| Two-step latents | Strided-NHD contract | Re-capture under frozen environment |
+| Canonical 12-step latents | video/audio SHA below | Re-capture inputs and full environment |
+
+Canonical final latent SHA-256 values:
+
+```text
+video c62d23a42972eab907ba42f93c50247ff17a9c454b4a53fe93d2e34f9fefe578
+audio 852005383770480a6503504e1ffec86dd1fb63a69c6400f92da18e39e0986de2
+```
+
+## Proof Limits And Blockers
+
+- No self-contained canonical Q/K/V tensor fixture is retained.
+- No checkpoint SHA-256 is retained.
+- The deployed Sage2 extension binary hash is not retained.
+- The exact production K-mean reduction tree and Comfy RMSNorm/RoPE binary are
+  not source-frozen in this repository.
+- NaN, infinity, denormal, signed-zero, and tie behavior is not exhaustively
+  proven by existing fixtures.
+
+These gaps do not invalidate retained decisions, but they block executable
+prototype acceptance. They must be captured before comparing a new kernel.
--- a/research/vortex_exact_attention/README.md
+++ b/research/vortex_exact_attention/README.md
@ -0,0 +1,82 @@
+# Vortex Exact Attention
+
+Vortex Exact Attention is a clean-sheet attention research project for MiniMax
+H3 on NVIDIA GB10 (`sm_121a`). Its claim target is the fastest reference-exact
+MiniMax H3 attention backend for long-context Blackwell inference. This is a
+research target, not an achieved performance claim.
+
+## Scope
+
+`vortex_exact` must preserve the observable behavior of SageAttention 2.2.0 at
+commit `d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5`. It may change scheduling,
+ownership, staging, and physical data movement, but not quantization, rounding,
+reduction, traversal, accumulation, or output-conversion semantics.
+
+Excluded from this backend:
+
+- sparsity;
+- cache reuse;
+- timestep or block skipping;
+- changed Q/K/V formats or scaling;
+- reordered reductions or approximate quality gates.
+
+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.
+- Kernel implementation: not started.
+- Production dispatch: unchanged; no `vortex_exact` route exists.
+
+## Authoritative Baseline
+
+| Item | Value |
+| --- | ---: |
+| Device | NVIDIA GB10, SM121 |
+| Workload | 1344x768, 124 frames, 37,810 tokens, 56 heads, D=128 |
+| Resident sampling median | `255.447 s` |
+| Sage2 share | `62.36%`, approximately `159.3 s` |
+| Uninstrumented mainloop | `237.089 ms` |
+| NCU grid / block | `(296,56,1)` / `(32,4,1)` |
+| KV traversal | 591 increasing 64-row tiles, two explicit stages |
+| Registers / dynamic shared memory | 255/thread / 32 KiB |
+| Achieved occupancy | `16.65-16.83%` |
+| No-eligible scheduler cycles | `63.53%` |
+| L2 hit / memory throughput | `98.85%` / approximately `31.74%` |
+
+The target is dependency latency and live-state ownership. Launch overhead,
+off-chip bandwidth, ragged tails, and ordinary shared-memory padding are closed
+as primary explanations.
+
+## Gates
+
+The first implementation may begin only after the Phase 1 design commit. No
+candidate can advance unless its model plausibly reduces simultaneous live
+state, avoids material score/probability tensors and catastrophic spilling,
+preserves the exact reduction sequence, and has a credible path below `220 ms`.
+A model below `190 ms` is preferred.
+
+Executable advancement later requires byte-exact short shapes, canonical
+attention, blocks 0/24/49, two-step latents, and 12-step latents, plus sanitizer,
+p95, block-level, and resident-service gates.
+
+## Layout
+
+- `NUMERICAL_CONTRACT.md`: imported exact Sage2 semantics and proof limits.
+- `DESIGN.md`: state-ownership architectures and selected prototype.
+- `PERFORMANCE_MODEL.md`: resource and latency model.
+- `HARDWARE_CAPABILITIES.md`: target-specific hardware facts and unknowns.
+- `EXPERIMENT_LOG.md`: append-only research decisions.
+- `benchmarks/`: machine-readable models and fixture manifests.
+- `kernels/`: future isolated implementation; currently contains no kernel.
+- `tests/`: future parity and sanitizer harnesses.
+- `tools/`: project-local verification/modeling utilities.
+
+## Provenance
+
+The reference is the public SageAttention repository at the pinned commit above.
+No upstream source is copied into this project. Before implementation, every
+borrowed algorithmic or source component must record its origin and license.
+P1/P2 local code may be reused only as provenance-preserving components inside
+an owned pipeline. See `NUMERICAL_CONTRACT.md` and `EXPERIMENT_LOG.md`.
--- a/research/vortex_exact_attention/benchmarks/reference_inventory.json
+++ b/research/vortex_exact_attention/benchmarks/reference_inventory.json
@ -0,0 +1,38 @@
+{
+  "schema": "vortex-exact-reference-inventory",
+  "version": 1,
+  "verified_at": "2026-08-26",
+  "reference": {
+    "name": "SageAttention 2.2.0",
+    "commit": "d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5",
+    "url": "https://github.com/thu-ml/SageAttention/tree/d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5"
+  },
+  "canonical": {
+    "tokens": 37810,
+    "heads": 56,
+    "head_dimension": 128,
+    "causal": false,
+    "attention_output_sha256": "4c666c20f5f8f651158a2ced33ccff08f3bada07665c595b99008d171db30574",
+    "video_latent_sha256": "c62d23a42972eab907ba42f93c50247ff17a9c454b4a53fe93d2e34f9fefe578",
+    "audio_latent_sha256": "852005383770480a6503504e1ffec86dd1fb63a69c6400f92da18e39e0986de2"
+  },
+  "artifacts": [
+    {"path": "benchmarks/gb10-sage-strided-nhd-contract-1344x768-124f-seed440420.json", "size_bytes": 6358, "sha256": "de934e74527121f808bbc3d8baf510c47dea09d8116616aabd2bd70d9bbf5e9f"},
+    {"path": "benchmarks/gb10-sage-strided-nhd-summary-1344x768-124f-seed440420.json", "size_bytes": 3045, "sha256": "6e4553bc01649ce196eb0c197b8b4baebed8a87a28e2d224c213684631849612"},
+    {"path": "benchmarks/gb10-sage2-block24-scheduler-baseline.json", "size_bytes": 8102, "sha256": "f7ef716f6541fe93b3db296ad9bd5b5aeae74ada457689ceaee6434a705424b6"},
+    {"path": "benchmarks/gb10-sage2-block24-mainloop-full.ncu-rep", "size_bytes": 6830626, "sha256": "5c34ba2b1a33f667fb519da0f54ac1dc68a6e2da894c67d1ce9dd02bae09b32a"},
+    {"path": "benchmarks/gb10-sage2-p1-real-parity.json", "size_bytes": 12023, "sha256": "378aac3a8754f2f6318c526798239ce58204caa5c186a416d2b71af7ddd6db19"},
+    {"path": "benchmarks/gb10-sage2-p2-vprep-analysis.json", "size_bytes": 1743, "sha256": "4aba487f54ae72b8468889c3551f825f4f3b8301c97a4061d21158b69f507320"},
+    {"path": "benchmarks/gb10-sage2-vprep-randomized.json", "size_bytes": 7299, "sha256": "0c05e04a46c2f36c98bd5d38fad1a7ed15227e5c441a3521a815b73e261e19e4"},
+    {"path": "benchmarks/gb10-sage2-p3-temporal-pair-analysis.json", "size_bytes": 2197, "sha256": "6cb2cd0606f07ab27dbaa55260c8f37d2d4e01cd526bdacc97a72ccbdeb863a9"},
+    {"path": "benchmarks/gb10-sage2-temporal-pair-short-20260825-p0.json", "size_bytes": 4348, "sha256": "e3c7b0f56125a358025184ca5c2b188d61999f8c3fdbd82fea41ee987b931f00"},
+    {"path": "benchmarks/gb10-sage2-temporal-pair-timing-20260825-p0.json", "size_bytes": 5487, "sha256": "74b69b08873f31e98a0d40c501186e1f5853176c653381a1d4e65ea5d9a28552"}
+  ],
+  "known_missing": [
+    "deployed SageAttention extension binary hash",
+    "checkpoint SHA-256",
+    "self-contained canonical Q/K/V input tensors",
+    "exact P0 candidate source patches",
+    "complete historical P0 commands and environments"
+  ]
+}
--- a/research/vortex_exact_attention/kernels/README.md
+++ b/research/vortex_exact_attention/kernels/README.md
@ -0,0 +1,5 @@
+# Kernels
+
+No kernel exists. Phase 0 and Phase 1 are specification and design work only.
+Future candidates remain isolated here until every exactness and performance
+gate passes; this directory must not be imported by production dispatch.
--- a/research/vortex_exact_attention/tests/README.md
+++ b/research/vortex_exact_attention/tests/README.md
@ -0,0 +1,14 @@
+# Tests
+
+Future test layers, in order:
+
+1. artifact and environment verification;
+2. randomized aligned/ragged prepared-byte parity;
+3. adversarial NaN, infinity, signed-zero, tie, and extreme-value behavior;
+4. compute-sanitizer and deterministic repetition;
+5. canonical attention and blocks 0/24/49;
+6. two-step and canonical 12-step video/audio latent parity;
+7. concurrent resident-service safety.
+
+Existing retained outputs are oracles, not a substitute for the missing
+self-contained canonical input fixtures recorded in `NUMERICAL_CONTRACT.md`.
--- a/research/vortex_exact_attention/tools/README.md
+++ b/research/vortex_exact_attention/tools/README.md
@ -0,0 +1,5 @@
+# Tools
+
+Project-local tools will verify reference artifacts, generate architecture
+models, and run isolated parity/profiling. No tool may alter production dispatch
+or start, stop, or profile the resident service without a later explicit gate.