Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

packages/realizations/cooperative/nvidia-sm86/src/Realization/Nvidia/SM86/StreamingAttentionSM86.alpha

3,230 lines493 declarations193.3 KiBSHA-256 cc0366b3d857

def · lines 1429–1429

sqrElement

Full file
dQ = s (the sum), turned back by the inverse half-split rotation and stored as the QKV pullback products' bfloat16 operand: row p, column c of head h at (p, 64 h + c) of the token-major [seq][width] projection gradient (argument 7), cos and sin the [seq][32] tables at arguments 12 and 13. The words the binary32 store (sqEpilogue) followed by the Q merge (AttentionGradientMergeSM86, attentionGradientMergeRotate with groups 1) stored: the same scaling FMUL, then the merge's instructions on the scaled words -- each sine negated, then FMUL, FNEG, FFMA; FMUL, FFMA -- and its packs (the odd component high). A thread holds columns c and c + 32 (n-tiles j and j + 4) for its rows. Each output pair is packed into its own accumulator register and stored once from it. The table loads take SB4: the pass's products free SB2 and SB3 (sm86TensorProductsFixed drops every wait on them). Registers: the K fragments' (dead once the products are done), from R100.
1429def sqrElement : Nat = 100

The compiler supplied declaration spans and resolved links from this source snapshot. This page does not assert that this file belongs to a checked closure.