h3-blackwell-runtime/research/vortex_exact_attention/NUMERICAL_CONTRACT.md
2026-08-26 14:44:28 +07:00

7.5 KiB

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:

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:

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:

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:

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.