det-attn

Bitwise-deterministic attention, invariant to KV sharding, loadable through kernels. With det-train's GEMM, the attention softmax combine is one of the two operations that make transformer training irreproducible under sequence and context parallelism; det-attn removes it. The reference contrast is the standard online-softmax combine, whose output bits depend on shard order.

Ring and context-parallel attention split the keys and values across devices and combine partial softmax states with an exp-rescale that is not associative in floating point, so the same attention over the same tokens returns different bits at different world sizes. This kernel makes the combine exact: composing KV shards is an integer addition, so the output is identical whether the sequence lives on one device or thirty-two, in any order.

Four shard rotations of the same attention. The float32 online-softmax output hashes to four different values (outputs differ by up to 1.5e-7); the det-attn output is torch.equal at every rotation, one hash, the same bytes.

Usage

import torch
from kernels import get_kernel

da = get_kernel("phanerozoic/det-attn", version=1, trust_remote_code=True)

# Q,K,V are [B, H, S, D]; GQA supported (K,V may have fewer heads).
out = da.det_attn(Q, K, V, causal=True)        # fused, deterministic

# context-parallel (ring) composition, bit-identical to the fused result:
out = da.det_attn_sharded(Q, K, V, kv_splits=[128, 128, 128, 128], causal=True)
assert torch.equal(out, da.det_attn(Q, K, V, causal=True))

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Symbol	Purpose
`det_attn(Q, K, V, causal, scale)`	fused deterministic attention, f32 out
`rowmax(Q, K, ..., kv_lo, kv_hi)`	per-query rowmax over a key range (combine by max)
`accumulate(Q, K, V, rowmax, ...)`	long-accumulator NUM/DEN digits for a shard (combine by add)
`finalize(NUM, DEN, shape)`	compose digits to the attention output
`det_attn_sharded(Q, K, V, kv_splits, ...)`	reference KV-shard composition

Method

Two exact, order-invariant reductions replace the online-softmax rescale:

a global row maximum rowmax_i = max_j (scale * Q_i . K_j), a max-reduce, order-invariant by construction; then
exact Kulisch long-accumulator sums of the exp-weighted values num_id = sum_j exp(s_ij - rowmax_i) V_jd and the denominator den_i = sum_j exp(s_ij - rowmax_i).

Because the row maximum is global, every weight and every product is a fixed float, so the pass-2 sums are exact integer accumulations into fixed bins. Composing KV shards is an int64 add of digit arrays, bit-identical to the single-device result and identical across ring step counts, block sizes, and world sizes. Scores are per-(i, j), so nothing depends on how keys are sharded.

Correctness

Shard invariance: the sharded composition and any permutation of KV blocks are torch.equal to the fused single-call result (the property the hero measures live; the float32 online-softmax combine over the same shards produces order-dependent bits at the 1e-7 level).
Causal and non-causal, MHA and GQA, f32 and bf16 inputs, D <= 256.
Repeated runs are bitwise identical.

Requirements and limits

NVIDIA GPU with compute capability 8.0+.
Two-pass exact accumulation on fp64 scalar arithmetic, not a fused tensor-core flash kernel; built for reproducibility rather than peak throughput. The regime is context-parallel training reproducibility, debugging, and audit.
D <= 256; f32 output.

References

The exact order-invariant softmax reduction over a Kulisch long accumulator; ring attention and context parallelism as the setting; phanerozoic/det-train for the GEMM half.