Define Vortex exact attention contract

This commit is contained in:
Daniel Maddern 2026-08-26 13:20:33 +07:00
parent a29b8960b0
commit c5028427d7
7 changed files with 354 additions and 0 deletions

View file

@ -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.

View file

@ -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.

View file

@ -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`.

View file

@ -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"
]
}

View file

@ -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.

View file

@ -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`.

View file

@ -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.