172 lines
7.5 KiB
Markdown
172 lines
7.5 KiB
Markdown
# 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 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
|
|
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 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
|
|
|
|
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 and reload-verified Q/K/V/output plus 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
|
|
|
|
- Canonical Q/K/V/output tensors are retained in durable Phase 2A storage; all
|
|
reload exactly and the output matches the locked Sage2 hash.
|
|
- The checkpoint SHA-256 is
|
|
`72fa9269ce551fb63ff42a32d9b46d0c122e84b4b2c511e22fa698287b088f70`.
|
|
- The deployed Sage2 mainloop and fused-preparation binary SHA-256 values are
|
|
`ea084d6a252524fb40ab5ede15583a22e87ec783bb9ec25119b87283141c6397`
|
|
and `8139fa8ed9900c76fde4a76bb2d9f5b487b7196129df152ea7d091138a96318a`.
|
|
- 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.
|
|
|
|
The remaining exceptional-value and source-freezing gaps do not invalidate the
|
|
Phase 2A capability decision. They remain gates for executable attention parity,
|
|
not for the isolated resource and handoff probes.
|