Large source region · 3,230 lines
module Realization.Nvidia.SM86.StreamingAttentionSM86
import Accelerator.SM86.Control
import Accelerator.SM86.Immediate
import Accelerator.SM86.Instruction
import Accelerator.SM86.NumericSemantics
import Accelerator.SM86.Operands
import Accelerator.SM86.Program
import Accelerator.SM86.Types
import Realization.Nvidia.SM86.StallCompaction
import Std.Natural
import Std.Foundation
-- Causal attention as a streaming normalised sum: the
-- forward pass of one head's softmax(Q K^T scale) V without the T x T
-- scores ever leaving a thread's registers. A block of four warps takes 64
-- queries of one head (grid x the query tiles, grid y the compact K/V heads,
-- grid z each K/V head's query group); each warp 16 of them. For every
-- 64-key tile up to and including the diagonal one,
-- the warp forms its 16 x 64 scores with the tensor cores, folds them into
-- a running row maximum m, a running row sum l and a running weighted sum O
-- by the merge rule for normalised weighted sums -- a new maximum m' scales
-- the old l and O by 2^(m - m') -- and adds the tile's P V. At the end
-- O / l is the attention output and m + log2 l each row's log-sum-exp (in
-- base 2), which the backward pass recomputes the probabilities from.
--
-- Scores are in base-2 units: t = S c with c the binary32 nearest
-- scale / ln 2, given at run time; P = 2^(t - m). P is rounded to binary16
-- for its product with V (as the materialised path rounds its
-- probabilities); the sums stay binary32. The reduction order differs from
-- the materialised path's, a numerical change the plan's contract admits
-- (the gradient check and the loss gate decide it).
--
-- Operands come straight from global memory (cached in L1): Q as the A
-- fragments of HMMA.16816 (row-major, [head][T][64] binary16), K's rows as
-- the B fragments of S = Q K^T ([head][T][64]), V^T's rows as the B
-- fragments of P V ([head][64][T]). No shared memory and no barrier: each
-- warp works alone. The output is binary16, token-major across heads
-- ([T][heads x 64], the heads merged), and the log-sum-exp binary32
-- ([head][T]).
--
-- Parameters (constant bank 0 from the SM86 parameter base): pointers to
-- Q, K, V^T, O and L, then c.
def saR = (lambda unrestricted index : Nat . (sm86Register (nat-to-byte index)))
def saU = sm86Unsigned32FromNaturalTruncated
def saNone = (constructor SM86Barrier SM86BarrierNone)
def saSB0 = (constructor SM86Barrier SM86Barrier0)
def saSB1 = (constructor SM86Barrier SM86Barrier1)
def saSB2 = (constructor SM86Barrier SM86Barrier2)
def saSB3 = (constructor SM86Barrier SM86Barrier3)
def saSB4 = (constructor SM86Barrier SM86Barrier4)
def saSB5 = (constructor SM86Barrier SM86Barrier5)
-- the wait mask naming scoreboards: bit n waits SBn
def saWaitNone : Nat = 0
def saWait0 : Nat = 1
def saWait1 : Nat = 2
def saWait2 : Nat = 4
def saWait3 : Nat = 8
def saWait4 : Nat = 16
def saWait5 : Nat = 32
def saWaitAll : Nat = 63
-- Generated with the largest stall (15): Realization.Nvidia.SM86.
-- StallCompaction shrinks each to what the fixed latencies need, and
-- SM86.Scoreboard checks the result; variable-latency results are ordered
-- by the scoreboards named here.
def saCtl = (lambda unrestricted write : (family SM86Barrier) . (lambda unrestricted wait : Nat .
(constructor SM86Control SM86ControlValue (byte 15) (constructor SM86YieldMode SM86Continue)
write saNone (nat-to-byte wait) (byte 0))))
def saPlain = (saCtl saNone saWaitNone)
def saAfter = (lambda unrestricted wait : Nat . (saCtl saNone wait))
def saOp = (lambda unrestricted body : (family SM86InstructionBody) . (lambda unrestricted tail : (family SM86Program) .
(constructor SM86Program SM86ProgramNext (sm86Instruction body) tail)))
def saUnless = (lambda unrestricted predicate : (family SM86Predicate) . (lambda unrestricted body : (family SM86InstructionBody) .
(lambda unrestricted tail : (family SM86Program) .
(constructor SM86Program SM86ProgramNext (sm86NegatedPredicatedInstruction predicate body) tail))))
def saWhen = (lambda unrestricted predicate : (family SM86Predicate) . (lambda unrestricted body : (family SM86InstructionBody) .
(lambda unrestricted tail : (family SM86Program) .
(constructor SM86Program SM86ProgramNext (sm86PredicatedInstruction predicate body) tail))))
def saP0 = (constructor SM86Predicate SM86Predicate0)
def saP1 = (constructor SM86Predicate SM86Predicate1)
-- body 0, body 1, ..., body (count - 1), then the tail
def saFor = (lambda unrestricted count : Nat .
(lambda unrestricted body : (pi unrestricted index : Nat . (pi unrestricted tail : (family SM86Program) . (family SM86Program))) .
(lambda unrestricted tail : (family SM86Program) .
(app
(nat-eliminate (lambda unrestricted n : Nat . (pi unrestricted start : Nat . (family SM86Program)))
(lambda unrestricted start : Nat . tail)
(lambda unrestricted predecessor : Nat .
(lambda unrestricted rest : (pi unrestricted start : Nat . (family SM86Program)) .
(lambda unrestricted start : Nat . (body start (rest (succ start))))))
count)
0))))
-- ---- instructions ----
def saMovImm = (lambda unrestricted d : Nat . (lambda unrestricted value : Nat .
(saOp (constructor SM86InstructionBody SM86MoveImmediate (saR d) (saU value) saPlain))))
def saMovImmAfter = (lambda unrestricted d : Nat . (lambda unrestricted value : Nat . (lambda unrestricted wait : Nat .
(saOp (constructor SM86InstructionBody SM86MoveImmediate (saR d) (saU value) (saAfter wait))))))
def saMovConstAfter = (lambda unrestricted d : Nat . (lambda unrestricted offset : Nat .
(lambda unrestricted wait : Nat .
(saOp (constructor SM86InstructionBody SM86MoveConstant (saR d) (byte 0) (saU offset) (saAfter wait))))))
def saMovConst = (lambda unrestricted d : Nat . (lambda unrestricted offset : Nat .
(saMovConstAfter d offset saWaitNone)))
def saS2R = (lambda unrestricted d : Nat . (lambda unrestricted special : (family SM86SpecialRegister) .
(saOp (constructor SM86InstructionBody SM86SpecialToRegister (saR d) special (saCtl saSB5 saWaitNone)))))
-- d = a * immediate + c
def saImad = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted immediate : Nat . (lambda unrestricted c : Nat .
(saOp (constructor SM86InstructionBody SM86IntegerMultiplyAddImmediate (saR d) (saR a) (saU immediate) (saR c) saPlain))))))
-- d = a + immediate
def saAddImm = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted immediate : Nat .
(saOp (constructor SM86InstructionBody SM86IntegerAddThreeImmediate (saR d) (saR a) (saU immediate) saPlain)))))
def saShr = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted amount : Nat .
(saOp (constructor SM86InstructionBody SM86ShiftRightImmediate (saR d) (saR a) (nat-to-byte amount) saPlain)))))
-- d = a & b
def saAnd = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted b : Nat .
(saOp (constructor SM86InstructionBody SM86LogicThreeInputTruthTable (saR d) (saR a) (saR b) (byte 192) saPlain)))))
-- d (pair) = a * b + c[0][offset] (64-bit)
def saWide = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted b : Nat . (lambda unrestricted offset : Nat .
(saOp (constructor SM86InstructionBody SM86IntegerMultiplyAddWideConstant (saR d) (saR a) (saR b) (byte 0) (saU offset) saPlain))))))
def saLoad = (lambda unrestricted d : Nat . (lambda unrestricted address : Nat . (lambda unrestricted offset : Nat .
(lambda unrestricted write : (family SM86Barrier) . (lambda unrestricted wait : Nat .
(saOp (constructor SM86InstructionBody SM86LoadGlobal (saR d) (saR address) (saU offset) (saCtl write wait))))))))
def saStore = (lambda unrestricted address : Nat . (lambda unrestricted value : Nat . (lambda unrestricted offset : Nat .
(lambda unrestricted wait : Nat .
(constructor SM86InstructionBody SM86StoreGlobal (saR address) (saR value) (saU offset) (saAfter wait))))))
-- The accumulator source is independent of the destination: RZ begins a
-- product without clearing its four destination registers. Subsequent slices
-- read the destination so their accumulation order is unchanged.
def saHmmaFrom = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted b : Nat .
(lambda unrestricted c : Nat . (lambda unrestricted write : (family SM86Barrier) . (lambda unrestricted wait : Nat .
(saOp (constructor SM86InstructionBody SM86TensorCoreHalfMatrixMultiplyAccumulate16x8x16Float32
(saR d) (saR a) (saR b) (saR c) (saCtl write wait)))))))))
def saHmma = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted b : Nat .
(saHmmaFrom d a b d))))
def saFmul = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted b : Nat . (lambda unrestricted wait : Nat .
(saOp (constructor SM86InstructionBody SM86FloatMultiply (saR d) (saR a) (saR b) (saAfter wait)))))))
def saFadd = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted b : Nat . (lambda unrestricted wait : Nat .
(saOp (constructor SM86InstructionBody SM86FloatAdd (saR d) (saR a) (saR b) (saAfter wait)))))))
def saFfma = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted b : Nat . (lambda unrestricted c : Nat .
(saOp (constructor SM86InstructionBody SM86FloatFusedMultiplyAdd (saR d) (saR a) (saR b) (saR c) saPlain))))))
def saFmax = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted b : Nat . (lambda unrestricted wait : Nat .
(saOp (constructor SM86InstructionBody SM86FloatMinimumOrMaximum (saR d) (saR a) (saR b)
(constructor SM86FloatExtremum SM86FloatMaximum) (saAfter wait)))))))
def saFneg = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat .
(saOp (constructor SM86InstructionBody SM86FloatNegate (saR d) (saR a) saPlain))))
def saMufu = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted operation : (family SM86MultiFunction) .
(lambda unrestricted wait : Nat .
(saOp (constructor SM86InstructionBody SM86MultiFunctionUnitApproximation (saR d) (saR a) operation (saCtl saSB4 wait)))))))
def saEx2 = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted wait : Nat .
(saMufu d a (constructor SM86MultiFunction SM86ExponentialBase2) wait))))
-- butterfly across the lanes of a quad. SHFL may read its source after
-- issue, and every use here overwrites that source once the result is in;
-- SB5 is also its read barrier, so the result's wait retires the read too
-- (SM86.Scoreboard does not let a write barrier retire late reads).
def saShfl = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted lane : Nat . (lambda unrestricted wait : Nat .
(saOp (constructor SM86InstructionBody SM86WarpShuffle (saR d) (saR a) (nat-to-byte lane) (saU 31)
(constructor SM86ShuffleMode SM86ShuffleButterfly)
(constructor SM86Control SM86ControlValue (byte 15) (constructor SM86YieldMode SM86Continue)
saSB5 saSB5 (nat-to-byte wait) (byte 0))))))))
-- d = (binary16 high, binary16 low)
def saPack = (lambda unrestricted d : Nat . (lambda unrestricted high : Nat . (lambda unrestricted low : Nat .
(saOp (constructor SM86InstructionBody SM86FloatPairToPackedHalfPair (saR d) (saR high) (saR low) saPlain)))))
-- P = a > immediate (unsigned)
def saGreater = (lambda unrestricted predicate : (family SM86Predicate) . (lambda unrestricted a : Nat . (lambda unrestricted immediate : Nat .
(saOp (constructor SM86InstructionBody SM86PredicateGreaterThanImmediate predicate (saR a) (saU immediate) saPlain)))))
-- ---- the geometry ----
def saHeadWidth : Nat = 64
def saTile : Nat = 64
def saThreads : Nat = 128
def saHalfBytes : Nat = 2
-- binary32 -infinity
def saMinusInfinity : Nat = 0xff800000
def saParameterBase : Nat = 0x160
def saArgument = (lambda unrestricted index : Nat . (naturalAdd saParameterBase (naturalMultiply index 8)))
def saConstantBytes : Nat = (saArgument 6)
-- ---- registers ----
def saTid : Nat = 0
def saLane : Nat = 1
def saWarp : Nat = 2
def saQuad : Nat = 3
def saQuadLane : Nat = 4
def saOne : Nat = 5
def saKeyOffset : Nat = 6
def saValueOffset : Nat = 7
def saCount : Nat = 8
def saMask : Nat = 9
def saKeyPointer : Nat = 10
def saValuePointer : Nat = 12
def saPointer : Nat = 14
-- Output stores can still be in flight when the log-sum pointer is formed.
-- Keep their address pair live until the stores retire instead of reusing it.
def saForwardLogSumPointer : Nat = 194
-- A global store may still consume its source register after the next pack
-- issues. Give each of the sixteen output stores a distinct packed word;
-- physical row-zero readback otherwise varied despite identical Q/K/V and
-- log-sum inputs. These registers follow the two-word log-sum address pair.
def saForwardOutputPackBase : Nat =
(naturalAdd saForwardLogSumPointer 2)
def saForwardOutputPackRegister =
(lambda unrestricted index : Nat .
(naturalAdd saForwardOutputPackBase index))
def saMaskBase : Nat = 16
def saScratch : Nat = 17
def saScale : Nat = 18
def saTileIndex : Nat = 19
def saHead : Nat = 192
def saRow : Nat = 193
-- Query backward occupies R194..R243, so the grouped K/V selector lives
-- above that entire allocation rather than aliasing its transpose pointer.
def saKeyValueHead : Nat = 244
-- the A fragments of Q: k-slice kt's four registers
def saQ = (lambda unrestricted kt : Nat . (naturalAdd 20 (naturalMultiply kt 4)))
def saQReg = (lambda unrestricted kt : Nat . (lambda unrestricted k : Nat . (naturalAdd (saQ kt) k)))
-- the score accumulators: n-tile nt's four binary32 registers
def saS = (lambda unrestricted nt : Nat . (naturalAdd 36 (naturalMultiply nt 4)))
def saO = (lambda unrestricted nt : Nat . (naturalAdd 68 (naturalMultiply nt 4)))
-- K's (then V^T's) B fragments: slice k, n-tile nt, half h
def saB = (lambda unrestricted k : Nat . (lambda unrestricted nt : Nat .
(naturalAdd 100 (naturalMultiply (naturalAdd (naturalMultiply k 8) nt) 2))))
-- P's A fragments: k-slice kk
def saPA = (lambda unrestricted kk : Nat . (naturalAdd 164 (naturalMultiply kk 4)))
def saM = (lambda unrestricted h : Nat . (naturalAdd 180 h))
def saL = (lambda unrestricted h : Nat . (naturalAdd 182 h))
def saTmp = (lambda unrestricted n : Nat . (naturalAdd 184 n))
def saRegisters : Nat = 248
-- the element (row half h, column b) of an accumulator tile: c0 c1 the
-- first row, c2 c3 the row 8 below
def saElem = (lambda unrestricted tile : Nat . (lambda unrestricted h : Nat . (lambda unrestricted b : Nat .
(naturalAdd tile (naturalAdd (naturalMultiply 2 h) b)))))
-- ---- one key tile ----
-- the tile's K fragments (all 32 slices x tiles x halves), S zeroed
def saLoadKeys = (lambda unrestricted tail : (family SM86Program) .
(saFor 32 (lambda unrestricted i : Nat .
(let unrestricted k = (naturalDivideUnchecked i 8) in (let unrestricted nt = (naturalModuloUnchecked i 8) in
(lambda unrestricted rest : (family SM86Program) .
(saLoad (saB k nt) saKeyPointer (naturalAdd (naturalMultiply nt 1024) (naturalMultiply k 32)) saSB0 (naturalSelect (naturalIsZero i) saWait3 saWaitNone)
(saLoad (succ (saB k nt)) saKeyPointer (naturalAdd (naturalMultiply nt 1024) (naturalAdd (naturalMultiply k 32) 16)) saSB0 saWaitNone
rest))))))
(saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 36 i) 0))
tail)))
-- S = Q K^T: slice by slice, the eight n-tiles of a slice in flight together
def saScores = (lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted k : Nat .
(saFor 8 (lambda unrestricted nt : Nat .
(saHmma (saS nt) (saQ k) (saB k nt) saSB2
(naturalSelect (naturalIsZero nt) (naturalSelect (naturalIsZero k) (naturalAdd saWait0 saWait1) saWait2) saWaitNone)))))
tail))
-- the tile's V^T fragments into the registers K's had (the products that
-- read them have finished: the first load waits for them)
def saLoadValues = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 32 (lambda unrestricted i : Nat .
(let unrestricted k = (naturalDivideUnchecked i 8) in (let unrestricted nt = (naturalModuloUnchecked i 8) in
(let unrestricted rowOffset = (naturalMultiply nt (naturalMultiply 8 (naturalMultiply seq saHalfBytes))) in
(lambda unrestricted rest : (family SM86Program) .
(saLoad (saB k nt) saValuePointer (naturalAdd rowOffset (naturalMultiply k 32)) saSB0 (naturalSelect (naturalIsZero i) saWait2 saWaitNone)
(saLoad (succ (saB k nt)) saValuePointer (naturalAdd rowOffset (naturalAdd (naturalMultiply k 32) 16)) saSB0 saWaitNone
rest)))))))
tail)))
-- the diagonal tile's mask: element (row r + 8h, column 8 nt + 2 (lane % 4) + b)
-- is masked when its column exceeds its row: when 8 nt + b - 8 h exceeds
-- 16 warp + lane / 4 - 2 (lane % 4). saMask holds that plus 64, plus 4096
-- for every tile after this one, so no element of an earlier tile is masked
def saMaskTile = (lambda unrestricted tail : (family SM86Program) .
(saFor 32 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
(let unrestricted b = (naturalModuloUnchecked i 2) in
(let unrestricted threshold = (naturalSaturatingSubtract (naturalAdd (naturalAdd (naturalMultiply 8 nt) b) 63) (naturalMultiply 8 h)) in
(lambda unrestricted rest : (family SM86Program) .
(saGreater saP0 saMask threshold
(saUnless saP0 (constructor SM86InstructionBody SM86MoveImmediate (saR (saElem (saS nt) h b)) (saU saMinusInfinity) saPlain)
rest))))))))
tail))
-- one row half h: its sixteen scores, scaled to base 2 (the products have
-- finished: the first multiply waits for them)
def saScaleRow = (lambda unrestricted h : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted e = (saElem (saS (naturalDivideUnchecked i 2)) h (naturalModuloUnchecked i 2)) in
(saFmul e e saScale saWaitNone)))
tail)))
-- the row's maximum over the quad's 64 columns, in tmp h
def saRowMax = (lambda unrestricted h : Nat . (lambda unrestricted tail : (family SM86Program) .
(let unrestricted into = (saTmp h) in
(saFmax into (saElem (saS 0) h 0) (saElem (saS 0) h 1) saWaitNone
(saFor 14 (lambda unrestricted j : Nat .
(let unrestricted i = (naturalAdd j 2) in
(saFmax into into (saElem (saS (naturalDivideUnchecked i 2)) h (naturalModuloUnchecked i 2)) saWaitNone)))
(saShfl (saTmp 2) into 1 saWaitNone
(saFmax into into (saTmp 2) saWait5
(saShfl (saTmp 2) into 2 saWaitNone
(saFmax into into (saTmp 2) saWait5
tail)))))))))
-- m' = max(m, the row's maximum) in tmp h; alpha = 2^(m - m') in tmp 4 + h;
-- m = m'; -m' in tmp 6 + h
def saRescaleFactor = (lambda unrestricted h : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFmax (saTmp h) (saTmp h) (saM h) saWaitNone
(saFneg (saTmp (naturalAdd 6 h)) (saTmp h)
(saFadd (saTmp (naturalAdd 4 h)) (saM h) (saTmp (naturalAdd 6 h)) saWaitNone
(saEx2 (saTmp (naturalAdd 4 h)) (saTmp (naturalAdd 4 h)) saWaitNone
(saOp (constructor SM86InstructionBody SM86IntegerAddThreeImmediate (saR (saM h)) (saR (saTmp h)) (saU 0) saPlain)
tail)))))))
-- P = 2^(t - m') in place of t; the row's sum of P into l, rescaled
def saExponentials = (lambda unrestricted h : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted e = (saElem (saS (naturalDivideUnchecked i 2)) h (naturalModuloUnchecked i 2)) in
(lambda unrestricted rest : (family SM86Program) .
(saFadd e e (saTmp (naturalAdd 6 h)) saWaitNone
(saEx2 e e saWaitNone rest)))))
-- l = l alpha + (the sum of the sixteen P), once the exponentials are in
(saFmul (saL h) (saL h) (saTmp (naturalAdd 4 h)) saWait4
(saFor 16 (lambda unrestricted i : Nat .
(saFadd (saL h) (saL h) (saElem (saS (naturalDivideUnchecked i 2)) h (naturalModuloUnchecked i 2)) saWaitNone))
tail)))))
-- O's row half h rescaled by alpha (the previous tile's products have
-- finished: the first multiply waits for them)
def saRescaleOutput = (lambda unrestricted h : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted e = (saElem (saO (naturalDivideUnchecked i 2)) h (naturalModuloUnchecked i 2)) in
(saFmul e e (saTmp (naturalAdd 4 h)) (naturalSelect (naturalIsZero i) saWait3 saWaitNone))))
tail)))
-- P's A fragments for P V: k-slice kk is n-tiles 2 kk and 2 kk + 1
def saPackP = (lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted kk : Nat .
(let unrestricted left = (saS (naturalMultiply 2 kk)) in (let unrestricted right = (saS (succ (naturalMultiply 2 kk))) in
(lambda unrestricted rest : (family SM86Program) .
(saPack (saPA kk) (succ left) left
(saPack (succ (saPA kk)) (naturalAdd left 3) (naturalAdd left 2)
(saPack (naturalAdd (saPA kk) 2) (succ right) right
(saPack (naturalAdd (saPA kk) 3) (naturalAdd right 3) (naturalAdd right 2)
rest))))))))
tail))
-- O += P V: the first product waits for V^T's loads
def saValues = (lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted k : Nat .
(saFor 8 (lambda unrestricted nt : Nat .
(saHmma (saO nt) (saPA k) (saB k nt) saSB3
(naturalSelect (naturalIsZero nt) (naturalSelect (naturalIsZero k) saWait0 saWait3) saWaitNone)))))
tail))
def saTileBody = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
(saWide saKeyPointer saKeyOffset saOne (saArgument 1)
(saWide saValuePointer saValueOffset saOne (saArgument 2)
(saImad saMask saCount 4096 saMaskBase
(saLoadKeys
(saScores
(saLoadValues seq
(saMaskTile
(saScaleRow 0 (saScaleRow 1
(saRowMax 0 (saRowMax 1
(saRescaleFactor 0 (saRescaleFactor 1
(saExponentials 0 (saExponentials 1
(saRescaleOutput 0 (saRescaleOutput 1
(saPackP
(saValues
tail)))))))))))))))))))))
-- ---- the loop's tail: the next key tile, one fewer left ----
def saLoopTail = (lambda unrestricted seq : Nat . (lambda unrestricted bodyInstructions : Nat .
(saAddImm saKeyOffset saKeyOffset (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes))
(saAddImm saValueOffset saValueOffset (naturalMultiply saTile saHalfBytes)
(saAddImm saCount saCount 4294967295
(saGreater saP1 saCount 0
(saWhen saP1
(constructor SM86InstructionBody SM86Branch
(saU (naturalSaturatingSubtract 4294967296 (naturalMultiply 16 (naturalAdd bodyInstructions 5))))
(sm86Unsigned32 (byte 255) (byte 255) (byte 131) (byte 3))
saPlain)
sm86ProgramEmpty)))))))
-- ---- the epilogue: O / l to the merged plane, m + log2 l ----
-- 1/l in registers of its own: the loop's in-place EX2 of the rescale factor
-- reads saTmp 4 + h without a read barrier, so SM86.Scoreboard keeps that
-- source live and refuses a later write to it.
def saForwardReciprocal = (lambda unrestricted h : Nat . (naturalAdd 212 h))
-- likewise the output address: the prologue's Q loads name saPointer
def saForwardOutputPointer : Nat = 214
def saEpilogue = (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted tail : (family SM86Program) .
(let unrestricted rowBytes = (naturalMultiply heads (naturalMultiply saHeadWidth saHalfBytes)) in
(saFor 2 (lambda unrestricted h : Nat . (lambda unrestricted rest : (family SM86Program) .
(saShfl (saTmp 2) (saL h) 1 saWait3
(saFadd (saL h) (saL h) (saTmp 2) saWait5
(saShfl (saTmp 2) (saL h) 2 saWaitNone
(saFadd (saL h) (saL h) (saTmp 2) saWait5
(saMufu (saForwardReciprocal h) (saL h) (constructor SM86MultiFunction SM86Reciprocal) saWaitNone
(saMufu (saTmp (naturalAdd 6 h)) (saL h) (constructor SM86MultiFunction SM86LogarithmBase2) saWaitNone
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted e = (saElem (saO (naturalDivideUnchecked i 2)) h (naturalModuloUnchecked i 2)) in
(saFmul e e (saForwardReciprocal h) (naturalSelect (naturalIsZero i) saWait4 saWaitNone))))
(saFadd (saM h) (saM h) (saTmp (naturalAdd 6 h)) saWaitNone
rest))))))))))
-- the output: this thread's rows r and r + 8, columns 8 nt + 2 (lane % 4)
(saWide saForwardOutputPointer saScratch saOne (saArgument 3)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 2) in (let unrestricted h = (naturalModuloUnchecked i 2) in
(lambda unrestricted rest : (family SM86Program) .
(saPack (saForwardOutputPackRegister i) (succ (saElem (saO nt) h 0)) (saElem (saO nt) h 0)
(saOp (saStore saForwardOutputPointer (saForwardOutputPackRegister i) (naturalAdd (naturalMultiply nt 16) (naturalMultiply h (naturalMultiply 8 rowBytes))) saWaitNone)
rest))))))
-- the log-sum-exp, by the quad's first lane
(saWide saForwardLogSumPointer saRow saOne (saArgument 4)
(saGreater saP0 saQuadLane 0
(saUnless saP0 (saStore saForwardLogSumPointer (saM 0) 0 saWaitNone)
(saUnless saP0 (saStore saForwardLogSumPointer (saM 1) 32 saWaitNone)
tail)))))))))))
-- ---- block order ----
-- Order 0: grid x the query tile, y the K/V head, z the query group.
-- Order 1: grid x the K/V head, y the query group, z the tile. Blocks are
-- dispatched in linear order, x fastest, so order 1 issues one tile of every
-- head before the next tile; the forward and query passes also count their
-- tiles down from the diagonal one, so the longest causal blocks start first
-- and the short ones fill the tail. Each block computes exactly what it
-- computed before.
def saAxis = (lambda unrestricted order : Nat .
(lambda unrestricted first : (family SM86SpecialRegister) .
(lambda unrestricted second : (family SM86SpecialRegister) .
(nat-eliminate (lambda unrestricted mode : Nat . (family SM86SpecialRegister))
first
(lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86SpecialRegister) . second))
order))))
def saTileAxis = (lambda unrestricted order : Nat . (saAxis order (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdX) (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdZ)))
def saKeyValueAxis = (lambda unrestricted order : Nat . (saAxis order (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdY) (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdX)))
def saGroupAxis = (lambda unrestricted order : Nat . (saAxis order (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdZ) (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdY)))
-- order 1: tile = tiles - 1 - z, once the special registers are in; saCount
-- is free until the loop count is formed from the tile
def saCountDownTile = (lambda unrestricted order : Nat . (lambda unrestricted tiles : Nat .
(lambda unrestricted tail : (family SM86Program) .
(nat-eliminate (lambda unrestricted mode : Nat . (family SM86Program))
tail
(lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86Program) .
(saMovImm saCount (naturalSaturatingSubtract tiles 1)
(saImad saTileIndex saTileIndex 4294967295 saCount tail))))
order))))
-- ---- the program ----
-- seq a multiple of 64, heads the number of 64-wide heads
-- Grid Y selects the compact K/V head and grid Z selects one of its query
-- heads. This avoids a device integer divide and keeps K/V planes compact.
-- The equal-head case uses one group and a unit Z dimension.
def streamingAttentionGroupedSM86Admitted =
(lambda unrestricted seq : Nat .
(lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(naturalAnd (naturalNonzero keyValueHeads)
(naturalAnd (naturalNonzero heads)
(naturalAnd (naturalNonzero seq)
(naturalAnd (naturalEqual (naturalModuloUnchecked seq saTile) 0)
(naturalAnd
(naturalEqual
(naturalModuloUnchecked heads
(naturalSelect (naturalNonzero keyValueHeads) keyValueHeads 1)) 0)
(naturalAnd (naturalLessOrEqual heads 65535)
(naturalLessOrEqual
(naturalMultiply heads (naturalMultiply saHeadWidth (naturalMultiply seq 4)))
4294967295))))))))))
def streamingAttentionForwardGroupedOrderedUncheckedSM86 = (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
(let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
(let unrestricted rowBytes = (naturalMultiply heads (naturalMultiply saHeadWidth saHalfBytes)) in
(let unrestricted body = (saTileBody seq sm86ProgramEmpty) in
(saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
(saS2R saTileIndex (saTileAxis order)
(saS2R saKeyValueHead (saKeyValueAxis order)
(saS2R saHead (saGroupAxis order)
(saMovImm saOne 1
-- the special registers are in once this waits
(saMovImmAfter saScratch 31 saWait5
(saImad saHead saKeyValueHead groups saHead
(saCountDownTile order (naturalDivideUnchecked seq saTile)
(saAnd saLane saTid saScratch
(saShr saWarp saTid 5
(saShr saQuad saLane 2
(saMovImm saScratch 3
(saAnd saQuadLane saLane saScratch
-- the thread's query row: 64 tile + 16 warp + lane / 4
(saImad saRow saWarp 16 saQuad
(saImad saRow saTileIndex saTile saRow
-- Q: head h's plane, the row, the quad lane's pair of columns
(saMovImm saScratch 0
(saImad saScratch saRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
(saImad saScratch saQuadLane 4 saScratch
(saImad saScratch saHead headBytes saScratch
(saWide saPointer saScratch saOne (saArgument 0)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
(saLoad (saQReg kt j) saPointer
(naturalAdd (naturalMultiply kt 32)
(naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
(naturalMultiply (naturalDivideUnchecked j 2) 16)))
saSB1 saWaitNone))))
-- the output's offset in the merged plane (kept in saScratch): the row,
-- head h's 64 columns, the quad lane's pair
(saMovImm saScratch 0
(saImad saScratch saRow rowBytes saScratch
(saImad saScratch saHead (naturalMultiply saHeadWidth saHalfBytes) saScratch
(saImad saScratch saQuadLane 4 saScratch
-- the log-sum-exp's (kept in saRow): 4 (h seq + row)
(saImad saRow saHead seq saRow
(saMovImm saKeyOffset 0
(saImad saRow saRow 4 saKeyOffset
-- K from key 0: the quad's row lane / 4, its pair of columns
(saImad saKeyOffset saQuad (naturalMultiply saHeadWidth saHalfBytes) saKeyOffset
(saImad saKeyOffset saQuadLane 4 saKeyOffset
(saImad saKeyOffset saKeyValueHead headBytes saKeyOffset
-- V^T from key 0: row lane / 4 of head h's 64 x seq plane
(saMovImm saValueOffset 0
(saImad saValueOffset saQuad (naturalMultiply seq saHalfBytes) saValueOffset
(saImad saValueOffset saQuadLane 4 saValueOffset
(saImad saValueOffset saKeyValueHead headBytes saValueOffset
-- the tiles up to the diagonal one; the mask's base, 16 warp + lane / 4
-- - 2 (lane % 4) + 64 - 4096
(saAddImm saCount saTileIndex 1
(saImad saMaskBase saWarp 16 saQuad
(saImad saMaskBase saQuadLane 4294967294 saMaskBase
(saAddImm saMaskBase saMaskBase (naturalSaturatingSubtract 4294967296 (naturalSaturatingSubtract 4096 64))
(saMovConst saScale (saArgument 5)
(saMovImm (saM 0) saMinusInfinity (saMovImm (saM 1) saMinusInfinity
(saMovImm (saL 0) 0 (saMovImm (saL 1) 0
(saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
(sm86ProgramAppend body
(sm86ProgramAppend (saLoopTail seq (sm86ProgramCount body))
(saEpilogue seq heads
(saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty)))))))))))))))))))))))))))))))))))))))))))))))))))))))))
def streamingAttentionForwardGroupedUncheckedSM86 = (streamingAttentionForwardGroupedOrderedUncheckedSM86 0)
def streamingAttentionForwardGroupedSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionForwardGroupedUncheckedSM86 seq heads keyValueHeads)))))
def streamingAttentionForwardSM86 = (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(streamingAttentionForwardGroupedUncheckedSM86 seq heads heads)))
def streamingAttentionForwardSM86Registers : Nat = saRegisters
def streamingAttentionForwardSM86Threads : Nat = saThreads
def streamingAttentionForwardSM86ConstantBytes : Nat = saConstantBytes
def streamingAttentionForwardQueryArgument : Nat = 0
def streamingAttentionForwardKeyArgument : Nat = 1
def streamingAttentionForwardValueArgument : Nat = 2
def streamingAttentionForwardOutputArgument : Nat = 3
def streamingAttentionForwardLogSumArgument : Nat = 4
def streamingAttentionForwardScaleArgument : Nat = 5
-- ===========================================================================
-- The forward pass with each key tile's K and V^T staged in shared memory.
-- The register kernel above has every warp fetch its own B fragments from
-- global memory, then wait for them, before each tile's products: with two
-- warps per scheduler little hides that latency. Here the block's 128 threads
-- copy tile t + 1 (cp.async, 16 bytes a thread a step) into one shared buffer
-- while the warps compute tile t from the other, and LDSM.x4 fills exactly
-- the B-fragment registers the register kernel loads (same values, same
-- lanes), so every product and every result word is unchanged. A row of 64
-- halves takes 144 bytes (128 + 16 of padding), which keeps LDSM's eight row
-- reads and the copies' 16-byte stores on distinct banks. The loop body is
-- emitted once per buffer (LDSM's address register is never rewritten); an
-- odd tile count leaves after the first body by a forward branch.
def sfRowBytes : Nat = 144
def sfValueTile : Nat = 9216
def sfBufferBytes : Nat = 18432
def streamingAttentionForwardSharedBytes : Nat = 36864
-- registers the register kernel leaves free (R216..R243)
def sfKeyOffset = (lambda unrestricted i : Nat . (naturalAdd 216 i))
def sfValueOffset = (lambda unrestricted i : Nat . (naturalAdd 220 i))
def sfPair = (lambda unrestricted j : Nat . (naturalAdd 224 (naturalMultiply 2 j)))
def sfStore : Nat = 240
def sfLane : Nat = 241
def sfT0 : Nat = 242
def sfT1 : Nat = 243
-- d (pair) = a + c[0][offset] (64-bit), after `wait`
def sfWide = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted offset : Nat .
(lambda unrestricted wait : Nat .
(saOp (constructor SM86InstructionBody SM86IntegerMultiplyAddWideConstant (saR d) (saR a) (saR saOne)
(byte 0) (saU offset) (saAfter wait)))))))
-- LDGSTS of chunk j (its pointer pair) to shared sfStore + offset. SB5 is its
-- read barrier, one live token: each chunk waits for the previous one's.
def sfCopyBody = (lambda unrestricted offset : Nat . (lambda unrestricted j : Nat . (lambda unrestricted wait : Nat .
(constructor SM86InstructionBody SM86LoadGlobalToShared (saR sfStore) (saU offset) (saR (sfPair j)) (saU 0)
(constructor SM86Control SM86ControlValue (byte 15) (constructor SM86YieldMode SM86Continue)
saNone saSB5 (nat-to-byte wait) (byte 0))))))
def sfCommit = (saOp (constructor SM86InstructionBody SM86CommitAsyncGroup (saCtl saSB0 saWaitNone)))
-- DEPBAR.LE SB0, n after `wait`
def sfWaitGroups = (lambda unrestricted n : Nat . (lambda unrestricted wait : Nat .
(saOp (constructor SM86InstructionBody SM86WaitAsyncGroups (nat-to-byte n) (saAfter wait)))))
def sfBarrier = (saOp (constructor SM86InstructionBody SM86BarrierSynchronize saPlain))
-- A tile's eight chunk pointers re-formed (the first after the previous
-- copies have read theirs), its chunks copied into `buffer` -- K's rows r,
-- r + 16, r + 32, r + 48 and V^T's the same, chunk c, for thread 8 r + c --
-- predicated on P1 when `guarded`, and the group closed.
def sfCopyTile = (lambda unrestricted buffer : Nat . (lambda unrestricted guarded : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted i : Nat .
(sfWide (sfPair i) (sfKeyOffset i) (saArgument 1) (naturalSelect (naturalIsZero i) saWait5 saWaitNone)))
(saFor 4 (lambda unrestricted i : Nat .
(sfWide (sfPair (naturalAdd 4 i)) (sfValueOffset i) (saArgument 2) saWaitNone))
(saFor 8 (lambda unrestricted j : Nat .
(let unrestricted body =
(sfCopyBody
(naturalAdd buffer
(naturalAdd (naturalSelect (naturalLess j 4) 0 sfValueTile)
(naturalMultiply (naturalModuloUnchecked j 4) (naturalMultiply 16 sfRowBytes))))
j (naturalSelect (naturalIsZero j) saWaitNone saWait5)) in
(nat-eliminate (lambda unrestricted mode : Nat . (pi unrestricted rest : (family SM86Program) . (family SM86Program)))
(saOp body)
(lambda unrestricted predecessor : Nat .
(lambda unrestricted induction : (pi unrestricted rest : (family SM86Program) . (family SM86Program)) .
(saWhen saP1 body)))
guarded)))
(sfCommit tail)))))))
-- The tile's B fragments from shared `base`: LDSM.x4 (k, p) fills n-tiles
-- 2p and 2p + 1 of slice k, rows 16 p.. and halves 16 k.. of the tile. Lane
-- L names row 8 (L / 16) + L % 8 and half 8 ((L / 8) % 2) of its matrix.
def sfLoadFragments = (lambda unrestricted base : Nat . (lambda unrestricted wait : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted k = (naturalDivideUnchecked i 4) in (let unrestricted pair = (naturalModuloUnchecked i 4) in
(saOp (constructor SM86InstructionBody SM86LoadSharedMatrix (saR (saB k (naturalMultiply 2 pair))) (saR sfLane)
(saU (naturalAdd base (naturalAdd (naturalMultiply pair (naturalMultiply 16 sfRowBytes)) (naturalMultiply k 32))))
(constructor SM86SharedMatrixCount SM86SharedMatrix4)
(constructor SM86SharedMatrixTranspose SM86SharedMatrixNotTransposed)
(saCtl saSB1 (naturalSelect (naturalIsZero i) wait saWaitNone)))))))
tail))))
-- The same B fragments from a tile staged as rows of the other operand
-- (keys down, columns across): LDSM.T (k, p) reads keys 16 k.., columns
-- 16 p.. through `lane` (lane L names key 8 ((L / 8) % 2) + L % 8, column
-- 8 (L / 16)), filling the registers sfLoadFragments fills from the
-- transposed tile.
def sfLoadFragmentsTransposed = (lambda unrestricted lane : Nat . (lambda unrestricted base : Nat . (lambda unrestricted wait : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted k = (naturalDivideUnchecked i 4) in (let unrestricted pair = (naturalModuloUnchecked i 4) in
(saOp (constructor SM86InstructionBody SM86LoadSharedMatrix (saR (saB k (naturalMultiply 2 pair))) (saR lane)
(saU (naturalAdd base (naturalAdd (naturalMultiply k (naturalMultiply 16 sfRowBytes)) (naturalMultiply pair 32))))
(constructor SM86SharedMatrixCount SM86SharedMatrix4)
(constructor SM86SharedMatrixTranspose SM86SharedMatrixTransposed)
(saCtl saSB1 (naturalSelect (naturalIsZero i) wait saWaitNone)))))))
tail)))))
-- `yes` (condition nonzero) or `no`, each awaiting its tail
def sfChooseWith = (lambda unrestricted condition : Nat .
(lambda unrestricted yes : (pi unrestricted tail : (family SM86Program) . (family SM86Program)) .
(lambda unrestricted no : (pi unrestricted tail : (family SM86Program) . (family SM86Program)) .
(nat-eliminate (lambda unrestricted n : Nat . (pi unrestricted tail : (family SM86Program) . (family SM86Program))) no
(lambda unrestricted p : Nat . (lambda unrestricted ignored : (pi unrestricted tail : (family SM86Program) . (family SM86Program)) . yes))
condition))))
-- S = Q K^T and O += P V as the register kernel's, the first product waiting
-- for the fragments (SB1, which also ordered Q's loads)
-- The first slice reads RZ for its accumulator, so no tile zeroes the
-- scores first: +0 + the products is the same word either way.
def sfScores = (lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted k : Nat .
(saFor 8 (lambda unrestricted nt : Nat .
(saHmmaFrom (saS nt) (saQ k) (saB k nt)
(naturalSelect (naturalIsZero k) sm86RegisterZero (saS nt)) saSB2
(naturalSelect (naturalIsZero nt) (naturalSelect (naturalIsZero k) saWait1 saWait2) saWaitNone)))))
tail))
def sfValues = (lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted k : Nat .
(saFor 8 (lambda unrestricted nt : Nat .
(saHmma (saO nt) (saPA k) (saB k nt) saSB3
(naturalSelect (naturalIsZero nt) (naturalSelect (naturalIsZero k) saWait1 saWait3) saWaitNone)))))
tail))
-- One tile from `buffer` while the next is copied into `other`: all warps
-- are done with `other` (the barrier), the copies go out while another tile
-- follows (P1), this tile's group retires (and the copies' register reads),
-- every thread's copies are visible (the barrier), then the register
-- kernel's tile body from shared fragments.
-- `rows` 1: V is staged as rows ([K/V head][T][64], K's layout), the next
-- tile 64 rows on, and O += P V's B fragments are its tile read transposed
-- (LDSM.T through sfT1, as the query pass reads K^T): the same values in
-- the same lanes as V^T's tile read plain.
-- The causal mask can change only the diagonal tile, which every pass
-- walks last: while another tile follows (P1, set at the body's top), a
-- forward branch skips the mask's 32 compares and 32 predicated moves
-- (which would leave every score as it is).
-- The schedule checks follow a straight line, so each body is checked
-- twice: through the mask (the diagonal tile) and with the branch taken
-- (`sfMaskSkipped`: the branch, then what follows the mask).
def sfMaskBranch = (lambda unrestricted bytes : Nat . (lambda unrestricted tail : (family SM86Program) .
(saWhen saP1
(constructor SM86InstructionBody SM86Branch
(saU bytes)
(sm86Unsigned32 (byte 0) (byte 0) (byte 128) (byte 3))
saPlain)
tail)))
-- the bytes the branch skips: the mask's 64 instructions
def sfMaskBytes : Nat = (naturalMultiply 16 64)
def sfMaskDiagonal = (lambda unrestricted tail : (family SM86Program) .
(sfMaskBranch sfMaskBytes (saMaskTile tail)))
def sfMaskSkipped = (lambda unrestricted tail : (family SM86Program) .
(sfMaskBranch sfMaskBytes tail))
def sfTileBodyMasked = (lambda unrestricted mask : (pi unrestricted tail : (family SM86Program) . (family SM86Program)) .
(lambda unrestricted rows : Nat . (lambda unrestricted buffer : Nat . (lambda unrestricted other : Nat .
(lambda unrestricted tail : (family SM86Program) .
(sfBarrier
(saGreater saP1 saCount 1
(saFor 4 (lambda unrestricted i : Nat . (saAddImm (sfKeyOffset i) (sfKeyOffset i) (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes))))
(saFor 4 (lambda unrestricted i : Nat . (saAddImm (sfValueOffset i) (sfValueOffset i)
(naturalSelect rows (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes)) (naturalMultiply saTile saHalfBytes))))
(sfCopyTile other 1
(sfWaitGroups 1 saWait5
(sfBarrier
(saImad saMask saCount 4096 saMaskBase
(sfLoadFragments buffer saWaitNone
(sfScores
(sfChooseWith rows
(sfLoadFragmentsTransposed sfT1 (naturalAdd buffer sfValueTile) saWait2)
(sfLoadFragments (naturalAdd buffer sfValueTile) saWait2)
(mask
(saScaleRow 0 (saScaleRow 1
(saRowMax 0 (saRowMax 1
(saRescaleFactor 0 (saRescaleFactor 1
(saExponentials 0 (saExponentials 1
(saRescaleOutput 0 (saRescaleOutput 1
(saPackP
(sfValues
tail)))))))))))))))))))))))))))))
def sfTileBodyWith = (sfTileBodyMasked sfMaskDiagonal)
def sfTileBody = (sfTileBodyWith 0)
-- after the first body: one tile fewer; leave for the epilogue (past the
-- second body and its tail) when none is left
def sfFirstTail = (lambda unrestricted secondInstructions : Nat .
(saAddImm saCount saCount 4294967295
(saGreater saP1 saCount 0
(saUnless saP1
(constructor SM86InstructionBody SM86Branch
(saU (naturalMultiply 16 (naturalAdd secondInstructions 3)))
(sm86Unsigned32 (byte 0) (byte 0) (byte 128) (byte 3))
saPlain)
sm86ProgramEmpty))))
-- after the second: back to the first body while tiles are left
def sfSecondTail = (lambda unrestricted firstInstructions : Nat . (lambda unrestricted secondInstructions : Nat .
(saAddImm saCount saCount 4294967295
(saGreater saP1 saCount 0
(saWhen saP1
(constructor SM86InstructionBody SM86Branch
(saU (naturalSaturatingSubtract 4294967296
(naturalMultiply 16 (naturalAdd (naturalAdd firstInstructions 3) (naturalAdd secondInstructions 3)))))
(sm86Unsigned32 (byte 255) (byte 255) (byte 131) (byte 3))
saPlain)
sm86ProgramEmpty)))))
-- the LDSM lane address, the copy role (row r = tid / 8, chunk c = tid % 8:
-- shared r 144 + 16 c; K's global offset at key r, V^T's at row r, chunk c,
-- of the block's K/V head; three more rows 16 apart), then tile 0 into
-- buffer 0
def sfSetupWithCopy = (lambda unrestricted rows : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted headBytes : Nat .
(lambda unrestricted copy : (pi unrestricted tail : (family SM86Program) . (family SM86Program)) .
(lambda unrestricted tail : (family SM86Program) .
(saShr sfT0 saLane 4
(saMovImm sfT1 7
(saAnd sfT1 saLane sfT1
(saImad sfT1 sfT0 8 sfT1
(saMovImm sfLane 0
(saImad sfLane sfT1 sfRowBytes sfLane
(saShr sfT0 saLane 3
(saAnd sfT0 sfT0 saOne
(saImad sfLane sfT0 16 sfLane
(saShr sfT0 saTid 3
(saMovImm sfT1 7
(saAnd sfT1 saTid sfT1
(saMovImm sfStore 0
(saImad sfStore sfT1 16 sfStore
(saAddImm (sfKeyOffset 0) sfStore 0
(saAddImm (sfValueOffset 0) sfStore 0
(saImad sfStore sfT0 sfRowBytes sfStore
(saImad (sfKeyOffset 0) sfT0 (naturalMultiply saHeadWidth saHalfBytes) (sfKeyOffset 0)
(saImad (sfKeyOffset 0) saKeyValueHead headBytes (sfKeyOffset 0)
-- V's row stride: rows (`rows` 1, K's) or V^T's seq halves
(let unrestricted valueRow = (naturalSelect rows (naturalMultiply saHeadWidth saHalfBytes) (naturalMultiply seq saHalfBytes)) in
(saImad (sfValueOffset 0) sfT0 valueRow (sfValueOffset 0)
(saImad (sfValueOffset 0) saKeyValueHead headBytes (sfValueOffset 0)
(saFor 3 (lambda unrestricted i : Nat .
(saAddImm (sfKeyOffset (succ i)) (sfKeyOffset 0) (naturalMultiply (succ i) (naturalMultiply 16 (naturalMultiply saHeadWidth saHalfBytes)))))
(saFor 3 (lambda unrestricted i : Nat .
(saAddImm (sfValueOffset (succ i)) (sfValueOffset 0) (naturalMultiply (succ i) (naturalMultiply 16 valueRow))))
-- rows: the transposed LDSM lane address in sfT1 (free from here on):
-- key 8 ((L / 8) % 2) + L % 8, column chunk L / 16
(sfChooseWith rows
(lambda unrestricted next : (family SM86Program) .
(saShr sfT0 saLane 3
(saAnd sfT0 sfT0 saOne
(saMovImm sfT1 7
(saAnd sfT1 saLane sfT1
(saImad sfT1 sfT0 8 sfT1
(saImad sfT1 sfT1 sfRowBytes 255
(saShr sfT0 saLane 4
(saImad sfT1 sfT0 16 sfT1 next)))))))))
(lambda unrestricted next : (family SM86Program) . next)
(copy tail)))))))))))))))))))))))))))))))
def sfSetupWith = (lambda unrestricted rows : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted headBytes : Nat .
(sfSetupWithCopy rows seq headBytes (sfCopyTile 0 0)))))
def sfSetup = (sfSetupWith 0)
def sfPrologueWithCopy = (lambda unrestricted rows : Nat . (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(lambda unrestricted copy : (pi unrestricted tail : (family SM86Program) . (family SM86Program)) .
(lambda unrestricted rest : (family SM86Program) .
(let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
(let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
(let unrestricted rowBytes = (naturalMultiply heads (naturalMultiply saHeadWidth saHalfBytes)) in
(saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
(saS2R saTileIndex (saTileAxis order)
(saS2R saKeyValueHead (saKeyValueAxis order)
(saS2R saHead (saGroupAxis order)
(saMovImm saOne 1
-- the special registers are in once this waits
(saMovImmAfter saScratch 31 saWait5
(saImad saHead saKeyValueHead groups saHead
(saCountDownTile order (naturalDivideUnchecked seq saTile)
(saAnd saLane saTid saScratch
(saShr saWarp saTid 5
(saShr saQuad saLane 2
(saMovImm saScratch 3
(saAnd saQuadLane saLane saScratch
-- the thread's query row: 64 tile + 16 warp + lane / 4
(saImad saRow saWarp 16 saQuad
(saImad saRow saTileIndex saTile saRow
-- Q: head h's plane, the row, the quad lane's pair of columns
(saMovImm saScratch 0
(saImad saScratch saRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
(saImad saScratch saQuadLane 4 saScratch
(saImad saScratch saHead headBytes saScratch
(saWide saPointer saScratch saOne (saArgument 0)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
(saLoad (saQReg kt j) saPointer
(naturalAdd (naturalMultiply kt 32)
(naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
(naturalMultiply (naturalDivideUnchecked j 2) 16)))
saSB1 saWaitNone))))
-- the output's offset in the merged plane (kept in saScratch): the row,
-- head h's 64 columns, the quad lane's pair
(saMovImm saScratch 0
(saImad saScratch saRow rowBytes saScratch
(saImad saScratch saHead (naturalMultiply saHeadWidth saHalfBytes) saScratch
(saImad saScratch saQuadLane 4 saScratch
-- the log-sum-exp's (kept in saRow): 4 (h seq + row)
(saImad saRow saHead seq saRow
(saMovImm saKeyOffset 0
(saImad saRow saRow 4 saKeyOffset
-- K from key 0: the quad's row lane / 4, its pair of columns
(saImad saKeyOffset saQuad (naturalMultiply saHeadWidth saHalfBytes) saKeyOffset
(saImad saKeyOffset saQuadLane 4 saKeyOffset
(saImad saKeyOffset saKeyValueHead headBytes saKeyOffset
-- V^T from key 0: row lane / 4 of head h's 64 x seq plane
(saMovImm saValueOffset 0
(saImad saValueOffset saQuad (naturalMultiply seq saHalfBytes) saValueOffset
(saImad saValueOffset saQuadLane 4 saValueOffset
(saImad saValueOffset saKeyValueHead headBytes saValueOffset
-- the tiles up to the diagonal one; the mask's base, 16 warp + lane / 4
-- - 2 (lane % 4) + 64 - 4096
(saAddImm saCount saTileIndex 1
(saImad saMaskBase saWarp 16 saQuad
(saImad saMaskBase saQuadLane 4294967294 saMaskBase
(saAddImm saMaskBase saMaskBase (naturalSaturatingSubtract 4294967296 (naturalSaturatingSubtract 4096 64))
(saMovConst saScale (saArgument 5)
(saMovImm (saM 0) saMinusInfinity (saMovImm (saM 1) saMinusInfinity
(saMovImm (saL 0) 0 (saMovImm (saL 1) 0
(saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
(sfSetupWithCopy rows seq headBytes copy rest))))))))))))))))))))))))))))))))))))))))))))))))))))))))
def sfPrologueWith = (lambda unrestricted rows : Nat . (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat .
(lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(sfPrologueWithCopy rows order seq heads keyValueHeads (sfCopyTile 0 0)))))))
def sfPrologue = (sfPrologueWith 0)
-- the program: both bodies and their tails, then the register kernel's
-- epilogue (reached from either tail)
def sfEpilogue = (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(saEpilogue seq heads
(saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty))))
def sfBodyWithA = (lambda unrestricted rows : Nat . (sfTileBodyWith rows 0 sfBufferBytes sm86ProgramEmpty))
def sfBodyWithB = (lambda unrestricted rows : Nat . (sfTileBodyWith rows sfBufferBytes 0 sm86ProgramEmpty))
def sfBodyA : (family SM86Program) = (sfBodyWithA 0)
def sfBodyB : (family SM86Program) = (sfBodyWithB 0)
-- the bodies with the mask's branch taken (every tile but the diagonal)
def sfBodySkippedWithA = (lambda unrestricted rows : Nat . (sfTileBodyMasked sfMaskSkipped rows 0 sfBufferBytes sm86ProgramEmpty))
def sfBodySkippedWithB = (lambda unrestricted rows : Nat . (sfTileBodyMasked sfMaskSkipped rows sfBufferBytes 0 sm86ProgramEmpty))
def streamingAttentionForwardSharedOrderedWithSM86 =
(lambda unrestricted rows : Nat . (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(let unrestricted bodyA = (sfBodyWithA rows) in (let unrestricted bodyB = (sfBodyWithB rows) in
(sfPrologueWith rows order seq heads keyValueHeads
(sm86ProgramAppend bodyA
(sm86ProgramAppend (sfFirstTail (sm86ProgramCount bodyB))
(sm86ProgramAppend bodyB
(sm86ProgramAppend (sfSecondTail (sm86ProgramCount bodyA) (sm86ProgramCount bodyB))
(sfEpilogue seq heads)))))))))))))
def streamingAttentionForwardSharedOrderedUncheckedSM86 = (streamingAttentionForwardSharedOrderedWithSM86 0)
-- What SM86.Scoreboard can follow of it: the straight-line pass reads one
-- body after the other, and the softmax's in-place EX2 (no read barrier)
-- leaves each score register's read live, so the second body's first write
-- of the scores is refused there -- the register kernel has the same pair
-- across its back edge, which the pass never follows. Each path from the
-- prologue through one body to the epilogue is checked whole.
def streamingAttentionForwardSharedPathSM86 =
(lambda unrestricted body : (family SM86Program) . (lambda unrestricted seq : Nat .
(lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(sfPrologue 1 seq heads keyValueHeads
(sm86ProgramAppend body (sfEpilogue seq heads)))))))
def streamingAttentionForwardSharedHeavyFirstSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionForwardSharedOrderedUncheckedSM86 1 seq heads keyValueHeads)))))
-- V as rows ([K/V head][T][64] half, K's layout, the backward passes' V):
-- the forward needs no transposed copy of V
-- ===========================================================================
-- The shared forward with V as rows, software-pipelined: while a warp works
-- through tile j's softmax it issues tile j+1's score products from K's next
-- tile, so the tensor cores need not idle through the softmax. The scores
-- alternate between two accumulator blocks (S at R36, S' at R100), and K's
-- and V's B fragments are loaded a 16-key slice at a time into two slice
-- buffers (R132, R148) instead of a tile at once (R100..R163). K and V keep
-- separate double buffers in the same shared layout (K at 0 and 18432, V
-- 9216 on): tile j's body copies K's tile j + 2 into the K buffer tile j's
-- scores have finished with and V's tile j + 1 into the V buffer tile
-- j - 1's products have finished with, so every copy has a whole tile to
-- land. Every product, softmax step and rounding is the shared forward's,
-- in the same order; only the issue order interleaves two independent
-- streams.
def sfpScoresB = (lambda unrestricted nt : Nat . (naturalAdd 100 (naturalMultiply nt 4)))
def sfpSlice = (lambda unrestricted k : Nat . (naturalAdd 132 (naturalMultiply (naturalModuloUnchecked k 2) 16)))
def sfpB = (lambda unrestricted k : Nat . (lambda unrestricted nt : Nat . (naturalAdd (sfpSlice k) (naturalMultiply 2 nt))))
def sfpKeyBuffer = (lambda unrestricted b : Nat . (naturalMultiply b sfBufferBytes))
def sfpValueBuffer = (lambda unrestricted b : Nat . (naturalAdd (naturalMultiply b sfBufferBytes) sfValueTile))
-- an instruction unpredicated (0), on P1 (1) or on P0 (2)
def sfpOp = (lambda unrestricted guard : Nat . (lambda unrestricted body : (family SM86InstructionBody) .
(lambda unrestricted tail : (family SM86Program) .
(nat-eliminate (lambda unrestricted n : Nat . (family SM86Program))
(saOp body tail)
(lambda unrestricted p : Nat . (lambda unrestricted ignored : (family SM86Program) .
(nat-eliminate (lambda unrestricted m : Nat . (family SM86Program))
(saWhen saP1 body tail)
(lambda unrestricted q : Nat . (lambda unrestricted ignored2 : (family SM86Program) . (saWhen saP0 body tail)))
p)))
guard))))
def sfpLdsm = (lambda unrestricted d : Nat . (lambda unrestricted lane : Nat . (lambda unrestricted offset : Nat .
(lambda unrestricted transposed : Nat . (lambda unrestricted wait : Nat .
(constructor SM86InstructionBody SM86LoadSharedMatrix (saR d) (saR lane) (saU offset)
(constructor SM86SharedMatrixCount SM86SharedMatrix4)
(eliminate StdBool (lambda unrestricted c : (family StdBool) . (family SM86SharedMatrixTranspose)) (stdBoolFromNatural transposed)
(branch StdTrue . (constructor SM86SharedMatrixTranspose SM86SharedMatrixTransposed))
(branch StdFalse . (constructor SM86SharedMatrixTranspose SM86SharedMatrixNotTransposed)))
(saCtl saSB1 wait)))))))
def sfpHmma = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted b : Nat . (lambda unrestricted c : Nat .
(lambda unrestricted write : (family SM86Barrier) . (lambda unrestricted wait : Nat .
(constructor SM86InstructionBody SM86TensorCoreHalfMatrixMultiplyAccumulate16x8x16Float32
(saR d) (saR a) (saR b) (saR c) (saCtl write wait))))))))
-- K's slice k (halves 16 k.., keys 16 pair..) as sfLoadFragments's registers
-- of that slice, into its slice buffer
def sfpKeySlice = (lambda unrestricted guard : Nat . (lambda unrestricted base : Nat . (lambda unrestricted k : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted pair : Nat .
(sfpOp guard (sfpLdsm (sfpB k (naturalMultiply 2 pair)) sfLane
(naturalAdd base (naturalAdd (naturalMultiply pair (naturalMultiply 16 sfRowBytes)) (naturalMultiply k 32))) 0 saWaitNone)))
tail)))))
-- slice k's products into `scores`, as sfScores's (from RZ at k = 0): the
-- first after the slice's fragments (SB1) and, past k = 0, the previous
-- slice's products (SB2)
def sfpScoreSlice = (lambda unrestricted guard : Nat . (lambda unrestricted scores : (pi unrestricted nt : Nat . Nat) .
(lambda unrestricted k : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 8 (lambda unrestricted nt : Nat .
(sfpOp guard (sfpHmma (scores nt) (saQ k) (sfpB k nt)
(naturalSelect (naturalIsZero k) sm86RegisterZero (scores nt)) saSB2
(naturalSelect (naturalIsZero nt) (naturalSelect (naturalIsZero k) saWait1 (naturalAdd saWait1 saWait2)) saWaitNone))))
tail)))))
-- the four slices, two loaded ahead: each slice's products issue before the
-- slice two on overwrites its buffer
def sfpScores = (lambda unrestricted guard : Nat . (lambda unrestricted scores : (pi unrestricted nt : Nat . Nat) .
(lambda unrestricted base : Nat . (lambda unrestricted tail : (family SM86Program) .
(sfpKeySlice guard base 0 (sfpKeySlice guard base 1
(sfpScoreSlice guard scores 0 (sfpKeySlice guard base 2
(sfpScoreSlice guard scores 1 (sfpKeySlice guard base 3
(sfpScoreSlice guard scores 2
(sfpScoreSlice guard scores 3 tail))))))))))))
-- V's slice kk (keys 16 kk.., columns 16 pair..) read transposed through
-- sfT1, as sfLoadFragmentsTransposed's registers of that slice
def sfpValueSlice = (lambda unrestricted base : Nat . (lambda unrestricted kk : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted pair : Nat .
(saOp (sfpLdsm (sfpB kk (naturalMultiply 2 pair)) sfT1
(naturalAdd base (naturalAdd (naturalMultiply kk (naturalMultiply 16 sfRowBytes)) (naturalMultiply pair 32))) 1 saWaitNone)))
tail))))
-- O += P V's slice kk, as sfValues's
def sfpValueProducts = (lambda unrestricted kk : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 8 (lambda unrestricted nt : Nat .
(saOp (sfpHmma (saO nt) (saPA kk) (sfpB kk nt) (saO nt) saSB3
(naturalSelect (naturalIsZero nt) (naturalSelect (naturalIsZero kk) saWait1 (naturalAdd saWait1 saWait3)) saWaitNone))))
tail)))
def sfpValues = (lambda unrestricted base : Nat . (lambda unrestricted tail : (family SM86Program) .
(sfpValueSlice base 0 (sfpValueSlice base 1
(sfpValueProducts 0 (sfpValueSlice base 2
(sfpValueProducts 1 (sfpValueSlice base 3
(sfpValueProducts 2
(sfpValueProducts 3 tail))))))))))
-- the softmax steps on `scores` (saMaskTile, saScaleRow, saRowMax,
-- saExponentials, saPackP with the score block a parameter); the first
-- scaling waits for the scores' products (SB2)
def sfpMaskTile = (lambda unrestricted scores : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted tail : (family SM86Program) .
(saFor 32 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
(let unrestricted b = (naturalModuloUnchecked i 2) in
(let unrestricted threshold = (naturalSaturatingSubtract (naturalAdd (naturalAdd (naturalMultiply 8 nt) b) 63) (naturalMultiply 8 h)) in
(lambda unrestricted rest : (family SM86Program) .
(saGreater saP0 saMask threshold
(saUnless saP0 (constructor SM86InstructionBody SM86MoveImmediate (saR (saElem (scores nt) h b)) (saU saMinusInfinity) saPlain)
rest))))))))
tail)))
def sfpMaskDiagonal = (lambda unrestricted scores : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted tail : (family SM86Program) .
(sfMaskBranch sfMaskBytes (sfpMaskTile scores tail))))
def sfpMaskSkipped = (lambda unrestricted scores : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted tail : (family SM86Program) .
(sfMaskBranch sfMaskBytes tail)))
def sfpScaleRow = (lambda unrestricted scores : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted h : Nat . (lambda unrestricted wait : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted e = (saElem (scores (naturalDivideUnchecked i 2)) h (naturalModuloUnchecked i 2)) in
(saFmul e e saScale (naturalSelect (naturalIsZero i) wait saWaitNone))))
tail)))))
def sfpRowMax = (lambda unrestricted scores : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted h : Nat . (lambda unrestricted tail : (family SM86Program) .
(let unrestricted into = (saTmp h) in
(saFmax into (saElem (scores 0) h 0) (saElem (scores 0) h 1) saWaitNone
(saFor 14 (lambda unrestricted j : Nat .
(let unrestricted i = (naturalAdd j 2) in
(saFmax into into (saElem (scores (naturalDivideUnchecked i 2)) h (naturalModuloUnchecked i 2)) saWaitNone)))
(saShfl (saTmp 2) into 1 saWaitNone
(saFmax into into (saTmp 2) saWait5
(saShfl (saTmp 2) into 2 saWaitNone
(saFmax into into (saTmp 2) saWait5
tail))))))))))
def sfpExponentials = (lambda unrestricted scores : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted h : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted e = (saElem (scores (naturalDivideUnchecked i 2)) h (naturalModuloUnchecked i 2)) in
(lambda unrestricted rest : (family SM86Program) .
(saFadd e e (saTmp (naturalAdd 6 h)) saWaitNone
(saEx2 e e saWaitNone rest)))))
(saFmul (saL h) (saL h) (saTmp (naturalAdd 4 h)) saWait4
(saFor 16 (lambda unrestricted i : Nat .
(saFadd (saL h) (saL h) (saElem (scores (naturalDivideUnchecked i 2)) h (naturalModuloUnchecked i 2)) saWaitNone))
tail))))))
def sfpPackP = (lambda unrestricted scores : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted kk : Nat .
(let unrestricted left = (scores (naturalMultiply 2 kk)) in (let unrestricted right = (scores (succ (naturalMultiply 2 kk))) in
(lambda unrestricted rest : (family SM86Program) .
(saPack (saPA kk) (succ left) left
(saPack (succ (saPA kk)) (naturalAdd left 3) (naturalAdd left 2)
(saPack (naturalAdd (saPA kk) 2) (succ right) right
(saPack (naturalAdd (saPA kk) 3) (naturalAdd right 3) (naturalAdd right 2)
rest))))))))
tail)))
def sfpSoftmax = (lambda unrestricted scores : (pi unrestricted nt : Nat . Nat) .
(sfpScaleRow scores 0 saWait2 (sfpScaleRow scores 1 saWaitNone
(sfpRowMax scores 0 (sfpRowMax scores 1
(saRescaleFactor 0 (saRescaleFactor 1
(sfpExponentials scores 0 (sfpExponentials scores 1
(saRescaleOutput 0 (saRescaleOutput 1
(sfpPackP scores sm86ProgramEmpty))))))))))))
-- `a`'s instructions with `b`'s interleaved: one of `b` after every `every`
-- of `a` (the first after `every`), the rest of either after the other ends
def sfpInterleaveFrom = (lambda unrestricted every : Nat . (lambda unrestricted a : (family SM86Program) .
(eliminate SM86Program (lambda unrestricted current : (family SM86Program) . (pi unrestricted b : (family SM86Program) . (pi unrestricted n : Nat . (family SM86Program)))) a
(branch SM86ProgramEnd . (lambda unrestricted b : (family SM86Program) . (lambda unrestricted n : Nat . b)))
(branch SM86ProgramNext instruction rest induction .
(lambda unrestricted b : (family SM86Program) . (lambda unrestricted n : Nat .
(nat-eliminate (lambda unrestricted m : Nat . (family SM86Program))
(eliminate SM86Program (lambda unrestricted current : (family SM86Program) . (family SM86Program)) b
(branch SM86ProgramEnd . (constructor SM86Program SM86ProgramNext instruction (induction b every)))
(branch SM86ProgramNext bInstruction bRest bInduction .
(constructor SM86Program SM86ProgramNext bInstruction
(constructor SM86Program SM86ProgramNext instruction (induction bRest every)))))
(lambda unrestricted p : Nat . (lambda unrestricted ignored : (family SM86Program) .
(constructor SM86Program SM86ProgramNext instruction (induction b p))))
n)))))))
def sfpInterleave = (lambda unrestricted every : Nat . (lambda unrestricted a : (family SM86Program) .
(lambda unrestricted b : (family SM86Program) . (sfpInterleaveFrom every a b every))))
def sfpEvery : Nat = 4
-- K's (into `keyBuffer`, guard `keyGuard`) and V's (`valueBuffer`,
-- `valueGuard`) four chunk rows of their next tiles, as sfCopyTile's
def sfpCopy = (lambda unrestricted keyBuffer : Nat . (lambda unrestricted keyGuard : Nat .
(lambda unrestricted valueBuffer : Nat . (lambda unrestricted valueGuard : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted i : Nat .
(sfWide (sfPair i) (sfKeyOffset i) (saArgument 1) (naturalSelect (naturalIsZero i) saWait5 saWaitNone)))
(saFor 4 (lambda unrestricted i : Nat .
(sfWide (sfPair (naturalAdd 4 i)) (sfValueOffset i) (saArgument 2) saWaitNone))
(saFor 8 (lambda unrestricted j : Nat .
(sfpOp (naturalSelect (naturalLess j 4) keyGuard valueGuard)
(sfCopyBody
(naturalAdd (naturalSelect (naturalLess j 4) keyBuffer valueBuffer)
(naturalMultiply (naturalModuloUnchecked j 4) (naturalMultiply 16 sfRowBytes)))
j (naturalSelect (naturalIsZero j) saWaitNone saWait5))))
(sfCommit tail)))))))))
-- One tile, j, with its scores in `scores` and tile j + 1's to come in
-- `next`; `b` is j % 2. All warps are done with K's buffer b (tile j's
-- products ran in the previous body) and V's buffer 1 - b (tile j - 1's):
-- K's tile j + 2 (P0: two tiles follow) and V's tile j + 1 (P1: one does)
-- are copied there; the previous body's copies retire and are visible to
-- every thread (the barrier); the diagonal mask, then tile j's softmax with
-- tile j + 1's products interleaved (on P1), then O += P V from V's buffer b.
def sfpTileBodyMasked = (lambda unrestricted mask : (pi unrestricted scores : (pi unrestricted nt : Nat . Nat) . (pi unrestricted tail : (family SM86Program) . (family SM86Program))) .
(lambda unrestricted scores : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted next : (pi unrestricted nt : Nat . Nat) .
(lambda unrestricted b : Nat . (lambda unrestricted tail : (family SM86Program) .
(let unrestricted other = (naturalSaturatingSubtract 1 b) in
(sfBarrier
(saGreater saP1 saCount 1
(saGreater saP0 saCount 2
(saFor 4 (lambda unrestricted i : Nat . (saAddImm (sfKeyOffset i) (sfKeyOffset i) (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes))))
(saFor 4 (lambda unrestricted i : Nat . (saAddImm (sfValueOffset i) (sfValueOffset i) (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes))))
(sfpCopy (sfpKeyBuffer b) 2 (sfpValueBuffer other) 1
(sfWaitGroups 1 saWait5
(sfBarrier
(saImad saMask saCount 4096 saMaskBase
(mask scores
(sm86ProgramAppend (sfpInterleave sfpEvery (sfpSoftmax scores) (sfpScores 1 next (sfpKeyBuffer other) sm86ProgramEmpty))
(sfpValues (sfpValueBuffer b)
tail))))))))))))))))))
def sfpTileBody = (sfpTileBodyMasked sfpMaskDiagonal)
-- K's four chunk rows of its next tile alone
def sfpKeyCopy = (lambda unrestricted keyBuffer : Nat . (lambda unrestricted keyGuard : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted i : Nat .
(sfWide (sfPair i) (sfKeyOffset i) (saArgument 1) (naturalSelect (naturalIsZero i) saWait5 saWaitNone)))
(saFor 4 (lambda unrestricted j : Nat .
(sfpOp keyGuard
(sfCopyBody (naturalAdd keyBuffer (naturalMultiply j (naturalMultiply 16 sfRowBytes)))
j (naturalSelect (naturalIsZero j) saWaitNone saWait5))))
(sfCommit tail))))))
-- the prologue's copies: tile 0's K and V into buffers 0, K's tile 1 into
-- K's buffer 1 (P1: a second tile); then K's offsets at tile 1, V's at 0.
-- Tile 0's scores once its copies are visible.
def sfpFirstCopies = (lambda unrestricted tail : (family SM86Program) .
(sfCopyTile 0 0
(saFor 4 (lambda unrestricted i : Nat . (saAddImm (sfKeyOffset i) (sfKeyOffset i) (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes))))
(saGreater saP1 saCount 1
(sfpKeyCopy (sfpKeyBuffer 1) 1
(sfWaitGroups 0 saWait5
(sfBarrier
(sfpScores 0 saS (sfpKeyBuffer 0) tail))))))))
def sfpBodyA = (sfpTileBody saS sfpScoresB 0 sm86ProgramEmpty)
def sfpBodyB = (sfpTileBody sfpScoresB saS 1 sm86ProgramEmpty)
-- every tile but the diagonal: the mask's branch taken
def sfpBodySkippedA = (sfpTileBodyMasked sfpMaskSkipped saS sfpScoresB 0 sm86ProgramEmpty)
def sfpBodySkippedB = (sfpTileBodyMasked sfpMaskSkipped sfpScoresB saS 1 sm86ProgramEmpty)
def sfpPrologue = (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat . (sfPrologueWithCopy 1 order seq heads keyValueHeads sfpFirstCopies)))))
def streamingAttentionForwardPipelinedOrderedSM86 =
(lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(sfpPrologue order seq heads keyValueHeads
(sm86ProgramAppend sfpBodyA
(sm86ProgramAppend (sfFirstTail (sm86ProgramCount sfpBodyB))
(sm86ProgramAppend sfpBodyB
(sm86ProgramAppend (sfSecondTail (sm86ProgramCount sfpBodyA) (sm86ProgramCount sfpBodyB))
(sfEpilogue seq heads))))))))))
def streamingAttentionForwardPipelinedPathSM86 =
(lambda unrestricted body : (family SM86Program) . (lambda unrestricted seq : Nat .
(lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(sfpPrologue 1 seq heads keyValueHeads
(sm86ProgramAppend body (sfEpilogue seq heads)))))))
def streamingAttentionForwardPipelinedHeavyFirstSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionForwardPipelinedOrderedSM86 1 seq heads keyValueHeads)))))
def streamingAttentionForwardSharedRowsHeavyFirstSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionForwardSharedOrderedWithSM86 1 1 seq heads keyValueHeads)))))
def streamingAttentionForwardSharedRowsPathSM86 =
(lambda unrestricted body : (family SM86Program) . (lambda unrestricted seq : Nat .
(lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(sfPrologueWith 1 1 seq heads keyValueHeads
(sm86ProgramAppend body (sfEpilogue seq heads)))))))
-- ===========================================================================
-- The backward pass, deterministic and in two launches (no atomics). With
-- P = 2^(S c - L) recomputed from the forward's base-2 log-sum-exp L, the
-- probability gradient dP = dO V^T and the row term D = rowsum(dO * O)
-- (= rowsum(P * dP)), the score gradient is dS = P * (dP - D); then
-- dQ = s dS K, dK = s dS^T Q and dV = P^T dO, s the score scale.
--
-- The query launch (streamingAttentionQuerySM86): a block per 64 queries of
-- a query head, grid (T / 64, K/V heads, groups). It forms D for its rows from dO's
-- fragments and the forward's output O (and writes it for the key launch),
-- then walks the key tiles up to the diagonal as the forward does, keeping
-- dQ. Operands: Q, K, V, dO ([head][T][64], half), K^T ([head][64][T],
-- half), O ([T][heads x 64], half, the forward's), L ([head][T]); results
-- dQ ([head][T][64], binary32) and D ([head][T], binary32). Parameters:
-- Q, K, V, K^T, dO, O, L, dQ, D, then c and s.
def sbStore64 = (lambda unrestricted address : Nat . (lambda unrestricted value : Nat . (lambda unrestricted offset : Nat .
(lambda unrestricted wait : Nat .
(constructor SM86InstructionBody SM86StoreGlobal64 (saR address) (saR value) (saU offset) (saAfter wait))))))
def sbWiden = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted high : Nat . (lambda unrestricted wait : Nat .
(saOp (constructor SM86InstructionBody SM86HalfToFloat (saR d) (saR a)
(eliminate StdBool (lambda unrestricted current : (family StdBool) . (family SM86HalfSelector)) (stdBoolFromNatural high)
(branch StdTrue . (constructor SM86HalfSelector SM86HighHalf))
(branch StdFalse . (constructor SM86HalfSelector SM86LowHalf)))
(saAfter wait)))))))
def sbNegate = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted wait : Nat .
(saOp (constructor SM86InstructionBody SM86FloatNegate (saR d) (saR a) (saAfter wait))))))
-- registers beyond the forward's
def sqKTOffset : Nat = 194
def sqDOut = (lambda unrestricted kt : Nat . (naturalAdd 196 (naturalMultiply kt 4)))
def sqDP = (lambda unrestricted nt : Nat . (naturalAdd 212 (naturalMultiply nt 4)))
-- the row terms: -L (184 + h) and -D (186 + h) for the rows r and r + 8
def sqNegL = (lambda unrestricted h : Nat . (naturalAdd 184 h))
def sqNegD = (lambda unrestricted h : Nat . (naturalAdd 186 h))
def sqTmp = (lambda unrestricted n : Nat . (naturalAdd 188 n))
-- Q loads may still use R14/R15 when the output-gradient address is formed.
-- Give dO its own pair; otherwise some row terms read the wrong operand and
-- vary across identical physical runs. The dO loads finish before the row-
-- term stores, which can reuse this pair. Both finish before the dP loop
-- assigns these words, keeping register demand below 256.
def sqOutputGradientLoadPointer : Nat = 240
def sqRowTermStorePointer : Nat = sqOutputGradientLoadPointer
def sqRegisters : Nat = 248
-- the dQ accumulators are the forward's O's; dS's A fragments P's
-- K^T's rows as the B fragments of dQ += dS K: slice kk (keys 16 kk ..),
-- n-tile nt (columns 8 nt ..) of head h's 64 x seq plane
def sqLoadKeysT = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 32 (lambda unrestricted i : Nat .
(let unrestricted k = (naturalDivideUnchecked i 8) in (let unrestricted nt = (naturalModuloUnchecked i 8) in
(let unrestricted rowOffset = (naturalMultiply nt (naturalMultiply 8 (naturalMultiply seq saHalfBytes))) in
(lambda unrestricted rest : (family SM86Program) .
(saLoad (saB k nt) saValuePointer (naturalAdd rowOffset (naturalMultiply k 32)) saSB0 (naturalSelect (naturalIsZero i) saWait2 saWaitNone)
(saLoad (succ (saB k nt)) saValuePointer (naturalAdd rowOffset (naturalAdd (naturalMultiply k 32) 16)) saSB0 saWaitNone
rest)))))))
tail)))
-- a tile's rows (K's or V's) as B fragments, the first load waiting on `wait`
def sqLoadRows = (lambda unrestricted pointer : Nat . (lambda unrestricted wait : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 32 (lambda unrestricted i : Nat .
(let unrestricted k = (naturalDivideUnchecked i 8) in (let unrestricted nt = (naturalModuloUnchecked i 8) in
(lambda unrestricted rest : (family SM86Program) .
(saLoad (saB k nt) pointer (naturalAdd (naturalMultiply nt 1024) (naturalMultiply k 32)) saSB0 (naturalSelect (naturalIsZero i) wait saWaitNone)
(saLoad (succ (saB k nt)) pointer (naturalAdd (naturalMultiply nt 1024) (naturalAdd (naturalMultiply k 32) 16)) saSB0 saWaitNone
rest))))))
tail))))
-- accumulators = A (4-register fragments from `a`) x the B fragments, slice
-- by slice; the first product waits on `first`, each slice after on SB2
def sqProducts = (lambda unrestricted accumulator : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted a : Nat .
(lambda unrestricted first : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted k : Nat .
(saFor 8 (lambda unrestricted nt : Nat .
(saHmma (accumulator nt) (naturalAdd a (naturalMultiply k 4)) (saB k nt) saSB2
(naturalSelect (naturalIsZero nt) (naturalSelect (naturalIsZero k) first saWait2) saWaitNone)))))
tail)))))
-- P = 2^(t - L) and dS = P (dP - D), into S's registers (the products have
-- finished: the first operation waits for them)
-- sqProducts with the first slice reading RZ for its accumulator: a tile's
-- products start from +0 without the accumulators being zeroed first
def sqProductsFresh = (lambda unrestricted accumulator : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted a : Nat .
(lambda unrestricted first : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted k : Nat .
(saFor 8 (lambda unrestricted nt : Nat .
(saHmmaFrom (accumulator nt) (naturalAdd a (naturalMultiply k 4)) (saB k nt)
(naturalSelect (naturalIsZero k) sm86RegisterZero (accumulator nt)) saSB2
(naturalSelect (naturalIsZero nt) (naturalSelect (naturalIsZero k) first saWait2) saWaitNone)))))
tail)))))
def sqScoreGradient = (lambda unrestricted tail : (family SM86Program) .
(saFor 32 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
(let unrestricted b = (naturalModuloUnchecked i 2) in
(let unrestricted e = (saElem (saS nt) h b) in
(lambda unrestricted rest : (family SM86Program) .
(saFfma e e saScale (sqNegL h)
(saEx2 e e saWaitNone rest))))))))
(saFor 32 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
(let unrestricted b = (naturalModuloUnchecked i 2) in
(let unrestricted e = (saElem (saS nt) h b) in (let unrestricted g = (saElem (sqDP nt) h b) in
(lambda unrestricted rest : (family SM86Program) .
(saFadd g g (sqNegD h) saWaitNone
(saFmul e e g (naturalSelect (naturalIsZero i) saWait4 saWaitNone) rest)))))))))
tail)))
def sqTileBody = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
(saImad saMask saCount 4096 saMaskBase
-- S = Q K^T
(saWide saKeyPointer saKeyOffset saOne (saArgument 1)
(sqLoadRows saKeyPointer saWait3
(saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 36 i) 0))
(saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 212 i) 0))
(sqProducts saS (saQ 0) (naturalAdd saWait0 saWait1)
-- dP = dO V^T (V's rows at K's offset)
(saWide saKeyPointer saKeyOffset saOne (saArgument 2)
(sqLoadRows saKeyPointer saWait2
(sqProducts sqDP (sqDOut 0) saWait0
-- K^T's fragments for dQ, once dP's products are done
(saWide saValuePointer sqKTOffset saOne (saArgument 3)
(sqLoadKeysT seq
(saMaskTile
(sqScoreGradient
(saPackP
(saValues
tail)))))))))))))))))
def sqLoopTail = (lambda unrestricted bodyInstructions : Nat .
(saAddImm saKeyOffset saKeyOffset (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes))
(saAddImm sqKTOffset sqKTOffset (naturalMultiply saTile saHalfBytes)
(saAddImm saCount saCount 4294967295
(saGreater saP1 saCount 0
(saWhen saP1
(constructor SM86InstructionBody SM86Branch
(saU (naturalSaturatingSubtract 4294967296 (naturalMultiply 16 (naturalAdd bodyInstructions 5))))
(sm86Unsigned32 (byte 255) (byte 255) (byte 131) (byte 3))
saPlain)
sm86ProgramEmpty))))))
-- D for the thread's rows: its sixteen dO values of row r (and of r + 8)
-- against O's, summed over the quad; -D and -L kept, D written by the
-- quad's first lane
-- D for the thread's rows, stored at 4 (h seq + row) by the quad's first
-- lane through `storePointer`
def sqRowTermsWith = (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted storePointer : Nat . (lambda unrestricted tail : (family SM86Program) .
(let unrestricted rowBytes = (naturalMultiply heads (naturalMultiply saHeadWidth saHalfBytes)) in
-- O in dO's fragment pattern, into the B registers (free before the loop).
-- Q's loads name saPointer as their address and may read it after issue:
-- re-form it only once they (SB1) have retired, else an O address can
-- replace a Q address before its late read (the c28712e8 hazard; the
-- checker found it at ordinal 64).
(saOp (constructor SM86InstructionBody SM86IntegerMultiplyAddWideConstant (saR saPointer) (saR saScratch) (saR saOne)
(byte 0) (saU (saArgument 5)) (saAfter saWait1))
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
(saLoad (naturalAdd 100 i) saPointer
(naturalAdd (naturalMultiply kt 32)
(naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 rowBytes))
(naturalMultiply (naturalDivideUnchecked j 2) 16)))
saSB0 saWaitNone))))
(saMovImm (sqTmp 0) 0 (saMovImm (sqTmp 1) 0
-- register j of slice kt is row r (j even) or r + 8 (j odd); both halves
(saFor 32 (lambda unrestricted i : Nat .
(let unrestricted r = (naturalDivideUnchecked i 2) in (let unrestricted high = (naturalModuloUnchecked i 2) in
(let unrestricted h = (naturalModuloUnchecked (naturalModuloUnchecked r 4) 2) in
(lambda unrestricted rest : (family SM86Program) .
(sbWiden (sqTmp 2) (naturalAdd 100 r) high (naturalSelect (naturalIsZero i) (naturalAdd saWait0 saWait1) saWaitNone)
(sbWiden (sqTmp 3) (naturalAdd (sqDOut 0) r) high saWaitNone
(saFfma (sqTmp h) (sqTmp 2) (sqTmp 3) (sqTmp h)
rest))))))))
(saFor 2 (lambda unrestricted h : Nat . (lambda unrestricted rest : (family SM86Program) .
(saShfl (sqTmp 2) (sqTmp h) 1 saWaitNone
(saFadd (sqTmp h) (sqTmp h) (sqTmp 2) saWait5
(saShfl (sqTmp 2) (sqTmp h) 2 saWaitNone
(saFadd (sqTmp h) (sqTmp h) (sqTmp 2) saWait5
(saFneg (sqNegD h) (sqTmp h)
rest)))))))
(saWide storePointer saRow saOne (saArgument 8)
(saGreater saP0 saQuadLane 0
(saUnless saP0 (saStore storePointer (sqTmp 0) 0 saWaitNone)
(saUnless saP0 (saStore storePointer (sqTmp 1) 32 saWaitNone)
tail)))))))))))))))
-- the query pass's: D, then -L at 4 (h seq + row)
def sqRowTerms = (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted tail : (family SM86Program) .
(sqRowTermsWith seq heads sqRowTermStorePointer
(saWide saPointer saRow saOne (saArgument 6)
(saLoad (sqTmp 2) saPointer 0 saSB1 saWaitNone
(saLoad (sqTmp 3) saPointer 32 saSB1 saWaitNone
(sbNegate (sqNegL 0) (sqTmp 2) saWait1
(sbNegate (sqNegL 1) (sqTmp 3) saWaitNone
tail)))))))))
-- dQ = s (the sum): the thread's rows r and r + 8, pairs of binary32 columns
def sqEpilogue = (lambda unrestricted tail : (family SM86Program) .
(saMovConst (sqTmp 4) (saArgument 10)
(saFor 32 (lambda unrestricted i : Nat .
(saFmul (naturalAdd 68 i) (naturalAdd 68 i) (sqTmp 4) (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
(saWide saPointer saScratch saOne (saArgument 7)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 2) in (let unrestricted h = (naturalModuloUnchecked i 2) in
(saOp (sbStore64 saPointer (saElem (saO nt) h 0)
(naturalAdd (naturalMultiply nt 32) (naturalMultiply h (naturalMultiply 8 (naturalMultiply saHeadWidth 4))))
saWaitNone)))))
tail)))))
-- 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.
def sqrElement : Nat = 100
def sqrColumn : Nat = 101
def sqrRow : Nat = 102
def sqrPosition : Nat = 103
def sqrHead : Nat = 104
def sqrOutput : Nat = 105
def sqrPointer : Nat = 106
def sqrTable : Nat = 108
def sqrFour : Nat = 109
def sqrCosine : Nat = 110
def sqrSine : Nat = 112
-- row half h, cos (s 0) or sin (s 1), n-tile j < 4, element e
def sqrWord = (lambda unrestricted h : Nat . (lambda unrestricted s : Nat . (lambda unrestricted j : Nat .
(lambda unrestricted e : Nat .
(naturalAdd 114 (naturalAdd (naturalMultiply 16 h) (naturalAdd (naturalMultiply 8 s)
(naturalAdd (naturalMultiply 2 j) e))))))))
def sqrOut = (lambda unrestricted i : Nat . (naturalAdd 146 i))
def streamingAttentionQueryRotatedCosineArgument : Nat = 12
def streamingAttentionQueryRotatedSineArgument : Nat = 13
def streamingAttentionQueryFromScoresRotatedConstantBytes : Nat = (saArgument 14)
-- log2 of a power of two below 2^32
def sqrLog2 = (lambda unrestricted n : Nat .
(nat-eliminate (lambda unrestricted current : Nat . Nat) 0
(lambda unrestricted k : Nat . (lambda unrestricted found : Nat .
(naturalSelect (naturalLessOrEqual (naturalPowerOfTwo (succ k)) n) (succ k) found)))
31))
-- row half h, n-tiles j and j + 4: the inverse rotation, both pairs packed
-- into their first accumulators and stored
def sqrPairs = (lambda unrestricted width : Nat . (lambda unrestricted h : Nat . (lambda unrestricted j : Nat .
(lambda unrestricted tail : (family SM86Program) .
(let unrestricted first = (saElem (saO j) h 0) in
(let unrestricted second = (saElem (saO (naturalAdd j 4)) h 0) in
(saFor 2 (lambda unrestricted e : Nat . (lambda unrestricted rest : (family SM86Program) .
(let unrestricted x1 = (saElem (saO j) h e) in
(let unrestricted x2 = (saElem (saO (naturalAdd j 4)) h e) in
(let unrestricted cosine = (sqrWord h 0 j e) in
(let unrestricted sine = (sqrWord h 1 j e) in
(saFmul (sqrOut e) x2 sine saWaitNone
(saFneg (sqrOut e) (sqrOut e)
(saFfma (sqrOut e) x1 cosine (sqrOut e)
(saFmul (sqrOut (naturalAdd 2 e)) x1 sine saWaitNone
(saFfma (sqrOut (naturalAdd 2 e)) x2 cosine (sqrOut (naturalAdd 2 e))
rest)))))))))))
(saOp (constructor SM86InstructionBody SM86FloatPairToPackedBFloat16Pair
(saR first) (saR (sqrOut 1)) (saR (sqrOut 0)) saPlain)
(saOp (constructor SM86InstructionBody SM86FloatPairToPackedBFloat16Pair
(saR second) (saR (sqrOut 3)) (saR (sqrOut 2)) saPlain)
(saOp (saStore sqrPointer first
(naturalAdd (naturalMultiply h (naturalMultiply 8 (naturalMultiply width saHalfBytes)))
(naturalMultiply j 16)) saWaitNone)
(saOp (saStore sqrPointer second
(naturalAdd (naturalMultiply h (naturalMultiply 8 (naturalMultiply width saHalfBytes)))
(naturalMultiply (naturalAdd j 4) 16)) saWaitNone)
tail)))))))))))
def sqRotatedEpilogue = (lambda unrestricted seq : Nat . (lambda unrestricted width : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saMovConst (sqTmp 4) (saArgument 10)
(saFor 32 (lambda unrestricted i : Nat .
(saFmul (naturalAdd 68 i) (naturalAdd 68 i) (sqTmp 4) (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
-- the column, position and head from the binary32 store's offset (bytes
-- of [head][seq][64]): element e = offset / 4, column e % 64, row e / 64
-- = head seq + position
(saShr sqrElement saScratch 2
(saMovImm sqrColumn 63
(saAnd sqrColumn sqrElement sqrColumn
(saShr sqrRow sqrElement 6
(saMovImm sqrPosition (naturalSaturatingSubtract seq 1)
(saAnd sqrPosition sqrRow sqrPosition
(saShr sqrHead sqrRow (sqrLog2 seq)
-- the output: (position width + 64 head + column) 2 bytes
(saImad sqrOutput sqrPosition width sqrColumn
(saImad sqrOutput sqrHead saHeadWidth sqrOutput
(saImad sqrOutput sqrOutput saHalfBytes 255
(saWide sqrPointer sqrOutput saOne (saArgument 7)
-- the tables: (position 32 + column) 4 bytes
(saImad sqrTable sqrPosition 32 sqrColumn
(saMovImm sqrFour 4
(saWide sqrCosine sqrTable sqrFour (saArgument streamingAttentionQueryRotatedCosineArgument)
(saWide sqrSine sqrTable sqrFour (saArgument streamingAttentionQueryRotatedSineArgument)
-- every table word: rows r and r + 8, columns 8 j + c and + 1 (j < 4)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted h = (naturalDivideUnchecked i 8) in
(let unrestricted j = (naturalDivideUnchecked (naturalModuloUnchecked i 8) 2) in
(let unrestricted e = (naturalModuloUnchecked i 2) in
(let unrestricted offset = (naturalMultiply 4 (naturalAdd (naturalMultiply 256 h) (naturalAdd (naturalMultiply 8 j) e))) in
(lambda unrestricted after : (family SM86Program) .
(saLoad (sqrWord h 0 j e) sqrCosine offset saSB4 saWaitNone
(saLoad (sqrWord h 1 j e) sqrSine offset saSB4 saWaitNone after))))))))
-- the sines negated (the inverse rotation), once every word is in
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted h = (naturalDivideUnchecked i 8) in
(let unrestricted j = (naturalDivideUnchecked (naturalModuloUnchecked i 8) 2) in
(let unrestricted e = (naturalModuloUnchecked i 2) in
(sbNegate (sqrWord h 1 j e) (sqrWord h 1 j e) (naturalSelect i saWaitNone saWait4))))))
(saFor 2 (lambda unrestricted h : Nat . (lambda unrestricted rest : (family SM86Program) .
(saFor 4 (lambda unrestricted j : Nat . (sqrPairs width h j)) rest)))
tail)))))))))))))))))))))))
def streamingAttentionQueryGroupedOrderedUncheckedSM86 = (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
(let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
(let unrestricted rowBytes = (naturalMultiply heads (naturalMultiply saHeadWidth saHalfBytes)) in
(let unrestricted body = (sqTileBody seq sm86ProgramEmpty) in
(saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
(saS2R saTileIndex (saTileAxis order)
(saS2R saKeyValueHead (saKeyValueAxis order)
(saS2R saHead (saGroupAxis order)
(saMovImm saOne 1
(saMovImmAfter saScratch 31 saWait5
(saImad saHead saKeyValueHead groups saHead
(saCountDownTile order (naturalDivideUnchecked seq saTile)
(saAnd saLane saTid saScratch
(saShr saWarp saTid 5
(saShr saQuad saLane 2
(saMovImm saScratch 3
(saAnd saQuadLane saLane saScratch
(saImad saRow saWarp 16 saQuad
(saImad saRow saTileIndex saTile saRow
-- Q's and dO's A fragments: head h's plane, the row, the quad lane's pair
(saMovImm saScratch 0
(saImad saScratch saRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
(saImad saScratch saQuadLane 4 saScratch
(saImad saScratch saHead headBytes saScratch
(saWide saPointer saScratch saOne (saArgument 0)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
(saLoad (saQReg kt j) saPointer
(naturalAdd (naturalMultiply kt 32)
(naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
(naturalMultiply (naturalDivideUnchecked j 2) 16)))
saSB1 saWaitNone))))
(saWide sqOutputGradientLoadPointer saScratch saOne (saArgument 4)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
(saLoad (naturalAdd (sqDOut kt) j) sqOutputGradientLoadPointer
(naturalAdd (naturalMultiply kt 32)
(naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
(naturalMultiply (naturalDivideUnchecked j 2) 16)))
saSB1 saWaitNone))))
-- dQ's offset (binary32 rows), kept in saKeyPointer's partner until the
-- epilogue: saScratch becomes O's (merged plane) for the row terms
(saMovImm saKeyOffset 0
(saImad saKeyOffset saRow (naturalMultiply saHeadWidth 4) saKeyOffset
(saImad saKeyOffset saQuadLane 8 saKeyOffset
(saImad saKeyOffset saHead (naturalMultiply seq (naturalMultiply saHeadWidth 4)) saKeyOffset
(saMovImm saScratch 0
(saImad saScratch saRow rowBytes saScratch
(saImad saScratch saHead (naturalMultiply saHeadWidth saHalfBytes) saScratch
(saImad saScratch saQuadLane 4 saScratch
-- the row terms' offset: 4 (h seq + row)
(saImad saRow saHead seq saRow
(saMovImm saValueOffset 0
(saImad saRow saRow 4 saValueOffset
(sqRowTerms seq heads
-- dQ's offset into saScratch; K's and V's from key 0, K^T's
(saAddImm saScratch saKeyOffset 0
(saMovImm saKeyOffset 0
(saImad saKeyOffset saQuad (naturalMultiply saHeadWidth saHalfBytes) saKeyOffset
(saImad saKeyOffset saQuadLane 4 saKeyOffset
(saImad saKeyOffset saKeyValueHead headBytes saKeyOffset
(saMovImm sqKTOffset 0
(saImad sqKTOffset saQuad (naturalMultiply seq saHalfBytes) sqKTOffset
(saImad sqKTOffset saQuadLane 4 sqKTOffset
(saImad sqKTOffset saKeyValueHead headBytes sqKTOffset
(saAddImm saCount saTileIndex 1
(saImad saMaskBase saWarp 16 saQuad
(saImad saMaskBase saQuadLane 4294967294 saMaskBase
(saAddImm saMaskBase saMaskBase (naturalSaturatingSubtract 4294967296 (naturalSaturatingSubtract 4096 64))
(saMovConst saScale (saArgument 9)
(saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
(sm86ProgramAppend body
(sm86ProgramAppend (sqLoopTail (sm86ProgramCount body))
(sqEpilogue
(saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))
def streamingAttentionQueryGroupedUncheckedSM86 = (streamingAttentionQueryGroupedOrderedUncheckedSM86 0)
def streamingAttentionQueryGroupedSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionQueryGroupedUncheckedSM86 seq heads keyValueHeads)))))
def streamingAttentionQuerySM86 = (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(streamingAttentionQueryGroupedUncheckedSM86 seq heads heads)))
-- ===========================================================================
-- The query pass with K and V staged in shared memory, as the forward above.
-- dQ's B fragments (K^T's rows in the register kernel) are the shared K
-- tile read transposed (LDSM.T: the same keys and columns in the same
-- lanes), so a buffer holds K's and V's rows only and two fit two blocks.
-- Registers: the loop no longer names R6, R7, R10..R13, R180..R183 or R194.
def sqsOffset = (lambda unrestricted i : Nat . (naturalAdd 10 i))
def sqsPair = (lambda unrestricted j : Nat . (naturalAdd 180 (naturalMultiply 2 j)))
def sqsStore : Nat = 6
def sqsLane : Nat = 7
def sqsLaneTransposed : Nat = 194
def sqsT0 : Nat = 212
def sqsT1 : Nat = 213
-- chunk j: K's (j < 4) or V's row r + 16 (j % 4), chunk c, through pointer
-- pair j % 2; each copy waits for the previous one's register reads, which
-- also frees the pair the next pointer re-forms
def sqsCopyTile = (lambda unrestricted buffer : Nat . (lambda unrestricted guarded : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 8 (lambda unrestricted j : Nat .
(let unrestricted body =
(constructor SM86InstructionBody SM86LoadGlobalToShared (saR sqsStore)
(saU (naturalAdd buffer
(naturalAdd (naturalSelect (naturalLess j 4) 0 sfValueTile)
(naturalMultiply (naturalModuloUnchecked j 4) (naturalMultiply 16 sfRowBytes)))))
(saR (sqsPair (naturalModuloUnchecked j 2))) (saU 0)
(constructor SM86Control SM86ControlValue (byte 15) (constructor SM86YieldMode SM86Continue)
saNone saSB5 (nat-to-byte saWait5) (byte 0))) in
(lambda unrestricted rest : (family SM86Program) .
(saOp (constructor SM86InstructionBody SM86IntegerMultiplyAddWideConstant
(saR (sqsPair (naturalModuloUnchecked j 2))) (saR (sqsOffset (naturalModuloUnchecked j 4))) (saR saOne)
(byte 0) (saU (naturalSelect (naturalLess j 4) (saArgument 1) (saArgument 2))) saPlain)
(nat-eliminate (lambda unrestricted mode : Nat . (family SM86Program))
(saOp body rest)
(lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86Program) .
(saWhen saP1 body rest)))
guarded)))))
(sfCommit tail)))))
-- a tile's B fragments from shared `base` through lane address `lane`,
-- plain (rows as the forward's) or transposed (LDSM.T (kk, p): keys 16 kk..,
-- columns 16 p..; lane L names key 8 ((L / 8) % 2) + L % 8, column 8 (L / 16))
def sqsLoadFragments = (lambda unrestricted lane : Nat . (lambda unrestricted transposed : Nat .
(lambda unrestricted base : Nat . (lambda unrestricted wait : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted k = (naturalDivideUnchecked i 4) in (let unrestricted pair = (naturalModuloUnchecked i 4) in
(saOp (constructor SM86InstructionBody SM86LoadSharedMatrix (saR (saB k (naturalMultiply 2 pair))) (saR lane)
(saU (naturalAdd base
(naturalSelect transposed
(naturalAdd (naturalMultiply k (naturalMultiply 16 sfRowBytes)) (naturalMultiply pair 32))
(naturalAdd (naturalMultiply pair (naturalMultiply 16 sfRowBytes)) (naturalMultiply k 32)))))
(constructor SM86SharedMatrixCount SM86SharedMatrix4)
(eliminate StdBool (lambda unrestricted c : (family StdBool) . (family SM86SharedMatrixTranspose)) (stdBoolFromNatural transposed)
(branch StdTrue . (constructor SM86SharedMatrixTranspose SM86SharedMatrixTransposed))
(branch StdFalse . (constructor SM86SharedMatrixTranspose SM86SharedMatrixNotTransposed)))
(saCtl saSB1 (naturalSelect (naturalIsZero i) wait saWaitNone)))))))
tail))))))
def sqsTileBodyMasked = (lambda unrestricted mask : (pi unrestricted tail : (family SM86Program) . (family SM86Program)) .
(lambda unrestricted buffer : Nat . (lambda unrestricted other : Nat .
(lambda unrestricted tail : (family SM86Program) .
(sfBarrier
(saGreater saP1 saCount 1
(saFor 4 (lambda unrestricted i : Nat . (saAddImm (sqsOffset i) (sqsOffset i) (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes))))
(sqsCopyTile other 1
(sfWaitGroups 1 saWait5
(sfBarrier
(saImad saMask saCount 4096 saMaskBase
-- S = Q K^T
(sqsLoadFragments sqsLane 0 buffer saWait3
(sqProductsFresh saS (saQ 0) saWait1
-- dP = dO V^T
(sqsLoadFragments sqsLane 0 (naturalAdd buffer sfValueTile) saWait2
(sqProductsFresh sqDP (sqDOut 0) saWait1
-- K^T's fragments for dQ: K's tile transposed, once dP's products are done
(sqsLoadFragments sqsLaneTransposed 1 buffer saWait2
(mask
(sqScoreGradient
(saPackP
(sfValues
tail))))))))))))))))))))
def sqsTileBody = (sqsTileBodyMasked sfMaskDiagonal)
-- lane addresses, the copy role and K/V offsets (both planes share K's
-- layout), then tile 0 into buffer 0
def sqsSetup = (lambda unrestricted headBytes : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saShr sqsT0 saLane 4
(saMovImm sqsT1 7
(saAnd sqsT1 saLane sqsT1
(saImad sqsT1 sqsT0 8 sqsT1
(saMovImm sqsLane 0
(saImad sqsLane sqsT1 sfRowBytes sqsLane
(saShr sqsT1 saLane 3
(saAnd sqsT1 sqsT1 saOne
(saImad sqsLane sqsT1 16 sqsLane
-- transposed: row 8 ((L / 8) % 2) + L % 8, column chunk L / 16
(saMovImm sqsLaneTransposed 7
(saAnd sqsLaneTransposed saLane sqsLaneTransposed
(saImad sqsT1 sqsT1 8 sqsLaneTransposed
(saMovImm sqsLaneTransposed 0
(saImad sqsLaneTransposed sqsT1 sfRowBytes sqsLaneTransposed
(saImad sqsLaneTransposed sqsT0 16 sqsLaneTransposed
(saShr sqsT0 saTid 3
(saMovImm sqsT1 7
(saAnd sqsT1 saTid sqsT1
(saMovImm sqsStore 0
(saImad sqsStore sqsT1 16 sqsStore
(saAddImm (sqsOffset 0) sqsStore 0
(saImad sqsStore sqsT0 sfRowBytes sqsStore
(saImad (sqsOffset 0) sqsT0 (naturalMultiply saHeadWidth saHalfBytes) (sqsOffset 0)
(saImad (sqsOffset 0) saKeyValueHead headBytes (sqsOffset 0)
(saFor 3 (lambda unrestricted i : Nat .
(saAddImm (sqsOffset (succ i)) (sqsOffset 0) (naturalMultiply (succ i) (naturalMultiply 16 (naturalMultiply saHeadWidth saHalfBytes)))))
(sqsCopyTile 0 0 tail))))))))))))))))))))))))))))
def sqsPrologue = (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(lambda unrestricted rest : (family SM86Program) .
(let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
(let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
(let unrestricted rowBytes = (naturalMultiply heads (naturalMultiply saHeadWidth saHalfBytes)) in
(saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
(saS2R saTileIndex (saTileAxis order)
(saS2R saKeyValueHead (saKeyValueAxis order)
(saS2R saHead (saGroupAxis order)
(saMovImm saOne 1
(saMovImmAfter saScratch 31 saWait5
(saImad saHead saKeyValueHead groups saHead
(saCountDownTile order (naturalDivideUnchecked seq saTile)
(saAnd saLane saTid saScratch
(saShr saWarp saTid 5
(saShr saQuad saLane 2
(saMovImm saScratch 3
(saAnd saQuadLane saLane saScratch
(saImad saRow saWarp 16 saQuad
(saImad saRow saTileIndex saTile saRow
-- Q's and dO's A fragments: head h's plane, the row, the quad lane's pair
(saMovImm saScratch 0
(saImad saScratch saRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
(saImad saScratch saQuadLane 4 saScratch
(saImad saScratch saHead headBytes saScratch
(saWide saPointer saScratch saOne (saArgument 0)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
(saLoad (saQReg kt j) saPointer
(naturalAdd (naturalMultiply kt 32)
(naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
(naturalMultiply (naturalDivideUnchecked j 2) 16)))
saSB1 saWaitNone))))
(saWide sqOutputGradientLoadPointer saScratch saOne (saArgument 4)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
(saLoad (naturalAdd (sqDOut kt) j) sqOutputGradientLoadPointer
(naturalAdd (naturalMultiply kt 32)
(naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
(naturalMultiply (naturalDivideUnchecked j 2) 16)))
saSB1 saWaitNone))))
-- dQ's offset (binary32 rows), kept in saKeyPointer's partner until the
-- epilogue: saScratch becomes O's (merged plane) for the row terms
(saMovImm saKeyOffset 0
(saImad saKeyOffset saRow (naturalMultiply saHeadWidth 4) saKeyOffset
(saImad saKeyOffset saQuadLane 8 saKeyOffset
(saImad saKeyOffset saHead (naturalMultiply seq (naturalMultiply saHeadWidth 4)) saKeyOffset
(saMovImm saScratch 0
(saImad saScratch saRow rowBytes saScratch
(saImad saScratch saHead (naturalMultiply saHeadWidth saHalfBytes) saScratch
(saImad saScratch saQuadLane 4 saScratch
-- the row terms' offset: 4 (h seq + row)
(saImad saRow saHead seq saRow
(saMovImm saValueOffset 0
(saImad saRow saRow 4 saValueOffset
(sqRowTerms seq heads
-- dQ's offset into saScratch; K's and V's from key 0, K^T's
(saAddImm saScratch saKeyOffset 0
(saMovImm saKeyOffset 0
(saImad saKeyOffset saQuad (naturalMultiply saHeadWidth saHalfBytes) saKeyOffset
(saImad saKeyOffset saQuadLane 4 saKeyOffset
(saImad saKeyOffset saKeyValueHead headBytes saKeyOffset
(saMovImm sqKTOffset 0
(saImad sqKTOffset saQuad (naturalMultiply seq saHalfBytes) sqKTOffset
(saImad sqKTOffset saQuadLane 4 sqKTOffset
(saImad sqKTOffset saKeyValueHead headBytes sqKTOffset
(saAddImm saCount saTileIndex 1
(saImad saMaskBase saWarp 16 saQuad
(saImad saMaskBase saQuadLane 4294967294 saMaskBase
(saAddImm saMaskBase saMaskBase (naturalSaturatingSubtract 4294967296 (naturalSaturatingSubtract 4096 64))
(saMovConst saScale (saArgument 9)
(saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
(sqsSetup headBytes rest)))))))))))))))))))))))))))))))))))))))))))))))))))))))))))
def sqsEpilogue : (family SM86Program) =
(sqEpilogue (saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty))
def sqsBodyA : (family SM86Program) = (sqsTileBody 0 sfBufferBytes sm86ProgramEmpty)
def sqsBodyB : (family SM86Program) = (sqsTileBody sfBufferBytes 0 sm86ProgramEmpty)
def sqsBodySkippedA : (family SM86Program) = (sqsTileBodyMasked sfMaskSkipped 0 sfBufferBytes sm86ProgramEmpty)
def sqsBodySkippedB : (family SM86Program) = (sqsTileBodyMasked sfMaskSkipped sfBufferBytes 0 sm86ProgramEmpty)
def streamingAttentionQuerySharedOrderedUncheckedSM86 =
(lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(sqsPrologue order seq heads keyValueHeads
(sm86ProgramAppend sqsBodyA
(sm86ProgramAppend (sfFirstTail (sm86ProgramCount sqsBodyB))
(sm86ProgramAppend sqsBodyB
(sm86ProgramAppend (sfSecondTail (sm86ProgramCount sqsBodyA) (sm86ProgramCount sqsBodyB))
sqsEpilogue)))))))))
def streamingAttentionQuerySharedPathSM86 =
(lambda unrestricted body : (family SM86Program) . (lambda unrestricted seq : Nat .
(lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(sqsPrologue 1 seq heads keyValueHeads (sm86ProgramAppend body sqsEpilogue))))))
def streamingAttentionQuerySharedHeavyFirstSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionQuerySharedOrderedUncheckedSM86 1 seq heads keyValueHeads)))))
def streamingAttentionQuerySharedBytes : Nat = streamingAttentionForwardSharedBytes
-- dQ = s (the sum) stored straight into a fused projection's gradient,
-- token-major [seq][width] with head h's columns at 64 h: binary32 at
-- argument 7 and bfloat16 at argument 12 (the projection products'
-- operand), instead of the [head][seq][64] plane an inverse layout pass
-- then copied and converted. The same scaling FMUL as sqEpilogue; each
-- pair stored as binary32 from its accumulators, then packed by F2FP.BF16
-- (the even column low) into a register of its own and stored. The row,
-- position and head come from the binary32 offset as sqRotatedEpilogue
-- takes them; its registers (the K fragments', dead once the products are
-- done).
def sqpHalfPointer : Nat = 108
def streamingAttentionQueryProjectionHalfArgument : Nat = 12
def sqProjectionEpilogue = (lambda unrestricted seq : Nat . (lambda unrestricted width : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saMovConst (sqTmp 4) (saArgument 10)
(saFor 32 (lambda unrestricted i : Nat .
(saFmul (naturalAdd 68 i) (naturalAdd 68 i) (sqTmp 4) (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
(saShr sqrElement saScratch 2
(saMovImm sqrColumn 63
(saAnd sqrColumn sqrElement sqrColumn
(saShr sqrRow sqrElement 6
(saMovImm sqrPosition (naturalSaturatingSubtract seq 1)
(saAnd sqrPosition sqrRow sqrPosition
(saShr sqrHead sqrRow (sqrLog2 seq)
-- the element: position width + 64 head + column
(saImad sqrOutput sqrPosition width sqrColumn
(saImad sqrOutput sqrHead saHeadWidth sqrOutput
(saMovImm sqrFour 4
(saWide sqrPointer sqrOutput sqrFour (saArgument 7)
(saMovImm sqrFour saHalfBytes
(saWide sqpHalfPointer sqrOutput sqrFour (saArgument streamingAttentionQueryProjectionHalfArgument)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 2) in (let unrestricted h = (naturalModuloUnchecked i 2) in
(lambda unrestricted after : (family SM86Program) .
(saOp (sbStore64 sqrPointer (saElem (saO nt) h 0)
(naturalAdd (naturalMultiply nt 32) (naturalMultiply h (naturalMultiply 8 (naturalMultiply width 4))))
saWaitNone)
(saOp (constructor SM86InstructionBody SM86FloatPairToPackedHalfPair
(saR (sqrOut i)) (saR (saElem (saO nt) h 1)) (saR (saElem (saO nt) h 0)) saPlain)
(saOp (saStore sqpHalfPointer (sqrOut i)
(naturalAdd (naturalMultiply nt 16) (naturalMultiply h (naturalMultiply 8 (naturalMultiply width saHalfBytes))))
saWaitNone)
after)))))))
tail)))))))))))))))))))
def streamingAttentionQueryProjectionConstantBytes : Nat = (saArgument 13)
-- a power-of-two sequence, the query heads within the projection's width,
-- its binary32 byte offsets within 32 bits
def streamingAttentionQueryProjectionAdmitted =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat . (lambda unrestricted width : Nat .
(naturalAnd (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads)
(naturalAnd (naturalEqual (naturalPowerOfTwo (sqrLog2 seq)) seq)
(naturalAnd (naturalLessOrEqual (naturalMultiply heads saHeadWidth) width)
(naturalLess (naturalMultiply seq (naturalMultiply width 4)) (naturalPowerOfTwo 32)))))))))
def streamingAttentionQuerySharedProjectionSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat . (lambda unrestricted width : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionQueryProjectionAdmitted seq heads keyValueHeads width) 1) .
(sqsPrologue 1 seq heads keyValueHeads
(sm86ProgramAppend sqsBodyA
(sm86ProgramAppend (sfFirstTail (sm86ProgramCount sqsBodyB))
(sm86ProgramAppend sqsBodyB
(sm86ProgramAppend (sfSecondTail (sm86ProgramCount sqsBodyA) (sm86ProgramCount sqsBodyB))
(sqProjectionEpilogue seq width
(saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty))))))))))))
def streamingAttentionQuerySM86Registers : Nat = sqRegisters
def streamingAttentionQuerySM86Threads : Nat = saThreads
def streamingAttentionQuerySM86ConstantBytes : Nat = (saArgument 11)
def streamingAttentionQueryQueryArgument : Nat = 0
def streamingAttentionQueryKeyArgument : Nat = 1
def streamingAttentionQueryValueArgument : Nat = 2
def streamingAttentionQueryTransposedKeyArgument : Nat = 3
def streamingAttentionQueryOutputGradientArgument : Nat = 4
def streamingAttentionQueryAttentionOutputArgument : Nat = 5
def streamingAttentionQueryLogSumArgument : Nat = 6
def streamingAttentionQueryQueryGradientArgument : Nat = 7
def streamingAttentionQueryRowTermArgument : Nat = 8
def streamingAttentionQueryBaseTwoScaleArgument : Nat = 9
def streamingAttentionQueryNaturalScaleArgument : Nat = 10
-- ===========================================================================
-- The query pass split in two around the key pass, which stores the score
-- gradients it forms anyway (sksStoreScoreGradients):
-- the row terms (streamingAttentionRowTermsSM86), before the key pass: the
-- query pass's prologue without Q, forming D for its 64 rows by the
-- same instructions (sqRowTermsWith) from the same fragments;
-- dQ (streamingAttentionQueryFromScoresSM86), after it: for each key tile
-- up to the diagonal, in the query pass's order, dQ += dS K from dS^T's
-- tile read transposed (LDSM.T) into the A registers saPackP fills and K's
-- read as the query pass reads it -- the same HMMAs on the same words --
-- then the query pass's epilogue.
-- S, P and dP are formed once (in the key pass) instead of twice.
-- the row terms: grid and parameters as the query pass's (dO 4, O 5, L 6,
-- D 8), one block of four warps per 64 rows of a query head
def streamingAttentionRowTermsOrderedUncheckedSM86 = (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat .
(lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
(let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
(let unrestricted rowBytes = (naturalMultiply heads (naturalMultiply saHeadWidth saHalfBytes)) in
(saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
(saS2R saTileIndex (saTileAxis order)
(saS2R saKeyValueHead (saKeyValueAxis order)
(saS2R saHead (saGroupAxis order)
(saMovImm saOne 1
(saMovImmAfter saScratch 31 saWait5
(saImad saHead saKeyValueHead groups saHead
(saCountDownTile order (naturalDivideUnchecked seq saTile)
(saAnd saLane saTid saScratch
(saShr saWarp saTid 5
(saShr saQuad saLane 2
(saMovImm saScratch 3
(saAnd saQuadLane saLane saScratch
(saImad saRow saWarp 16 saQuad
(saImad saRow saTileIndex saTile saRow
-- dO's A fragments: head h's plane, the row, the quad lane's pair
(saMovImm saScratch 0
(saImad saScratch saRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
(saImad saScratch saQuadLane 4 saScratch
(saImad saScratch saHead headBytes saScratch
(saWide sqOutputGradientLoadPointer saScratch saOne (saArgument 4)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
(saLoad (naturalAdd (sqDOut kt) j) sqOutputGradientLoadPointer
(naturalAdd (naturalMultiply kt 32)
(naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
(naturalMultiply (naturalDivideUnchecked j 2) 16)))
saSB1 saWaitNone))))
-- O's offset (merged plane) and the row terms', as the query pass's
(saMovImm saScratch 0
(saImad saScratch saRow rowBytes saScratch
(saImad saScratch saHead (naturalMultiply saHeadWidth saHalfBytes) saScratch
(saImad saScratch saQuadLane 4 saScratch
(saImad saRow saHead seq saRow
(saMovImm saValueOffset 0
(saImad saRow saRow 4 saValueOffset
-- D through its own pair: the dO loads' address pair is never re-formed
(sqRowTermsWith seq heads saKeyPointer
(saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty)))))))))))))))))))))))))))))))))))))
def streamingAttentionRowTermsHeavyFirstSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionRowTermsOrderedUncheckedSM86 1 seq heads keyValueHeads)))))
def streamingAttentionRowTermsSM86Registers : Nat = sqRegisters
def streamingAttentionRowTermsSM86Threads : Nat = saThreads
def streamingAttentionRowTermsSM86ConstantBytes : Nat = streamingAttentionQuerySM86ConstantBytes
-- dS^T's layout: [head][32-query sub-tile][key][32 queries] half (the key
-- pass's sub-tile, skSubTile), so a key-pass warp's sub-tile is one
-- contiguous kilobyte (16 keys of 64 bytes)
def sksScoreGradientRowBytes : Nat = (naturalMultiply 32 saHalfBytes)
def sksScoreGradientSubTileBytes = (lambda unrestricted seq : Nat . (naturalMultiply seq sksScoreGradientRowBytes))
-- dQ from dS^T: parameters K 1, dQ 7, s 10 (the query pass's) and dS^T 11
-- (the key pass's argument 14); grid as the query pass's. Three shared
-- stages of K's tile and dS^T's (64 rows of 128 bytes each, unpadded: a
-- row's chunk c at c ^ (row % 8), so the eight rows an LDSM reads fall in
-- distinct banks), the copies two tiles ahead.
-- K's and dS^T's offsets of the thread's copy rows r + 16 m (LDGSTS's
-- global immediate is 11 bits: each row group keeps its own offset)
def sqdKeyOffset = (lambda unrestricted m : Nat . (naturalAdd 180 m))
def sqdScoreOffset = (lambda unrestricted m : Nat . (naturalAdd 184 m))
def sqdStore : Nat = 188
def sqdLaneScores : Nat = 189
def sqdT2 : Nat = 190
-- K^T's lanes, one a column pair p (the swizzle does not commute with +2p)
def sqdLaneKeys = (lambda unrestricted p : Nat . (naturalAdd 202 p))
def sqdT0 : Nat = 200
def sqdT1 : Nat = 201
def sqdPair = (lambda unrestricted j : Nat . (naturalAdd 196 (naturalMultiply 2 j)))
def streamingAttentionQueryFromScoresArgument : Nat = 11
def streamingAttentionQueryFromScoresConstantBytes : Nat = (saArgument 12)
def sqdRowBytes : Nat = 128
def sqdTileBytes : Nat = (naturalMultiply saTile sqdRowBytes)
def sqdStageBytes : Nat = (naturalMultiply 2 sqdTileBytes)
def sqdStages : Nat = 3
def sqdStage = (lambda unrestricted b : Nat . (naturalMultiply b sqdStageBytes))
-- d = a ^ b (LOP3 0x3c)
def sqdXor = (lambda unrestricted d : Nat . (lambda unrestricted a : Nat . (lambda unrestricted b : Nat .
(saOp (constructor SM86InstructionBody SM86LogicThreeInputTruthTable (saR d) (saR a) (saR b) (byte 60) saPlain)))))
-- A tile's copies into stage `buffer`: chunk j < 4 K's row r + 16 j, j >= 4
-- dS^T's row r + 16 (j - 4) (the tile's keys, this block's 64 queries),
-- chunk c, thread 8 r + c (rows 16 apart share the swizzle r % 8), each
-- through pointer pair j % 2 formed just before it, each copy after the
-- previous one's reads; predicated on P1 when `guarded`; then the next
-- tile's offsets, 64 keys on.
def sqdCopyTile = (lambda unrestricted seq : Nat . (lambda unrestricted buffer : Nat . (lambda unrestricted guarded : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 8 (lambda unrestricted j : Nat .
(let unrestricted m = (naturalModuloUnchecked j 4) in
(let unrestricted body =
(constructor SM86InstructionBody SM86LoadGlobalToShared (saR sqdStore)
(saU (naturalAdd buffer
(naturalAdd (naturalSelect (naturalLess j 4) 0 sqdTileBytes)
(naturalMultiply m (naturalMultiply 16 sqdRowBytes)))))
(saR (sqdPair (naturalModuloUnchecked j 2))) (saU 0)
(constructor SM86Control SM86ControlValue (byte 15) (constructor SM86YieldMode SM86Continue)
saNone saSB5 (nat-to-byte saWait5) (byte 0))) in
(lambda unrestricted rest : (family SM86Program) .
(saOp (constructor SM86InstructionBody SM86IntegerMultiplyAddWideConstant
(saR (sqdPair (naturalModuloUnchecked j 2)))
(saR (naturalSelect (naturalLess j 4) (sqdKeyOffset m) (sqdScoreOffset m))) (saR saOne)
(byte 0) (saU (naturalSelect (naturalLess j 4) (saArgument 1) (saArgument streamingAttentionQueryFromScoresArgument)))
saPlain)
(nat-eliminate (lambda unrestricted mode : Nat . (family SM86Program))
(saOp body rest)
(lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86Program) .
(saWhen saP1 body rest)))
guarded))))))
(saFor 4 (lambda unrestricted m : Nat .
(saAddImm (sqdKeyOffset m) (sqdKeyOffset m) (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes))))
(saFor 4 (lambda unrestricted m : Nat .
(saAddImm (sqdScoreOffset m) (sqdScoreOffset m) (naturalMultiply saTile sksScoreGradientRowBytes)))
(sfCommit tail))))))))
-- dS's A fragments for slice kk (keys 16 kk..) from dS^T's staged tile:
-- LDSM.T of its rows (keys) through the lane naming key 8 (L / 16) + L % 8,
-- query 16 warp + 8 ((L / 8) % 2): register j holds row g + 8 (j % 2),
-- keys 16 kk + 8 (j / 2) + 2 (L % 4).., the words saPackP packs
def sqdLoadScores = (lambda unrestricted base : Nat . (lambda unrestricted wait : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 4 (lambda unrestricted kk : Nat .
(saOp (constructor SM86InstructionBody SM86LoadSharedMatrix (saR (saPA kk)) (saR sqdLaneScores)
(saU (naturalAdd base (naturalMultiply kk (naturalMultiply 16 sqdRowBytes))))
(constructor SM86SharedMatrixCount SM86SharedMatrix4)
(constructor SM86SharedMatrixTranspose SM86SharedMatrixTransposed)
(saCtl saSB1 (naturalSelect (naturalIsZero kk) wait saWaitNone)))))
tail))))
-- K^T's B fragments of slice k (keys 16 k..), n-tile pair p (columns
-- 16 p..), LDSM.T through pair p's lane: the registers the query pass's
-- sqsLoadFragments fills from its padded tile
def sqdLoadKeys = (lambda unrestricted base : Nat . (lambda unrestricted wait : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted k = (naturalDivideUnchecked i 4) in (let unrestricted pair = (naturalModuloUnchecked i 4) in
(saOp (constructor SM86InstructionBody SM86LoadSharedMatrix (saR (saB k (naturalMultiply 2 pair))) (saR (sqdLaneKeys pair))
(saU (naturalAdd base (naturalMultiply k (naturalMultiply 16 sqdRowBytes))))
(constructor SM86SharedMatrixCount SM86SharedMatrix4)
(constructor SM86SharedMatrixTranspose SM86SharedMatrixTransposed)
(saCtl saSB1 (naturalSelect (naturalIsZero i) wait saWaitNone)))))))
tail))))
-- One key tile from stage `buffer`: all warps are done with the stage two
-- on (the barrier: it held the previous tile), the copies of the tile two
-- on go out while one follows (P1), this tile's group retires (at most two
-- newer pending), every copy is visible (the barrier), then once the last
-- tile's products have read their operands (SB3) dS's and K^T's fragments
-- and dQ += dS K.
def sqdTileBody = (lambda unrestricted seq : Nat . (lambda unrestricted buffer : Nat .
(lambda unrestricted tail : (family SM86Program) .
(sfBarrier
(saGreater saP1 saCount 2
(sqdCopyTile seq (sqdStage (naturalModuloUnchecked (naturalAdd (naturalDivideUnchecked buffer sqdStageBytes) 2) sqdStages)) 1
(sfWaitGroups 2 saWait5
(sfBarrier
(sqdLoadScores (naturalAdd buffer sqdTileBytes) saWait3
(sqdLoadKeys buffer saWaitNone
(sfValues
tail)))))))))))
-- the LDSM lanes, the copy role (shared r 128 + 16 (c ^ r % 8)), K's
-- offsets (K/V head, key r + 16 m, chunk c) and dS^T's (head, the block's
-- queries 64 i + 8 c, key r + 16 m), then tiles 0 and 1 (P1 when a second
-- follows)
def sqdSetupAhead = (lambda unrestricted ahead : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted headBytes : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saShr sqdT0 saLane 4
(saMovImm sqdT1 7
(saAnd sqdT1 saLane sqdT1
(saShr sqdT2 saLane 3
(saAnd sqdT2 sqdT2 saOne
-- K^T: key 8 ((L / 8) % 2) + L % 8, chunk L / 16 + 2 p, swizzled by L % 8
(saImad (sqdKeyOffset 0) sqdT2 8 sqdT1
(saMovImm (sqdKeyOffset 1) 0
(saImad (sqdKeyOffset 1) (sqdKeyOffset 0) sqdRowBytes (sqdKeyOffset 1)
(saFor 4 (lambda unrestricted p : Nat . (lambda unrestricted rest : (family SM86Program) .
(saAddImm (sqdKeyOffset 2) sqdT0 (naturalMultiply 2 p)
(sqdXor (sqdKeyOffset 2) (sqdKeyOffset 2) sqdT1
(saImad (sqdLaneKeys p) (sqdKeyOffset 2) 16 (sqdKeyOffset 1)
rest)))))
-- dS^T: key 8 (L / 16) + L % 8, chunk 2 warp + (L / 8) % 2, swizzled by L % 8
(saImad (sqdKeyOffset 0) sqdT0 8 sqdT1
(saImad (sqdKeyOffset 2) saWarp 2 sqdT2
(sqdXor (sqdKeyOffset 2) (sqdKeyOffset 2) sqdT1
(saMovImm sqdLaneScores 0
(saImad sqdLaneScores (sqdKeyOffset 0) sqdRowBytes sqdLaneScores
(saImad sqdLaneScores (sqdKeyOffset 2) 16 sqdLaneScores
-- the copy role: row r = tid / 8, chunk c = tid % 8
(saShr sqdT0 saTid 3
(saMovImm sqdT1 7
(saAnd sqdT1 saTid sqdT1
(saMovImm sqdT2 7
(saAnd sqdT2 sqdT0 sqdT2
(sqdXor sqdT2 sqdT2 sqdT1
(saMovImm sqdStore 0
(saImad sqdStore sqdT2 16 sqdStore
(saImad sqdStore sqdT0 sqdRowBytes sqdStore
(saMovImm (sqdKeyOffset 0) 0
(saImad (sqdKeyOffset 0) sqdT1 16 (sqdKeyOffset 0)
(saImad (sqdKeyOffset 0) sqdT0 (naturalMultiply saHeadWidth saHalfBytes) (sqdKeyOffset 0)
(saImad (sqdKeyOffset 0) saKeyValueHead headBytes (sqdKeyOffset 0)
-- dS^T's ([head][sub-tile][key][32 queries]): chunk c is sub-tile
-- 2 i + c / 4, queries 8 (c % 4).., key r
(saShr (sqdScoreOffset 1) sqdT1 2
(saMovImm (sqdScoreOffset 2) 3
(saAnd (sqdScoreOffset 2) sqdT1 (sqdScoreOffset 2)
(saMovImm (sqdScoreOffset 0) 0
(saImad (sqdScoreOffset 0) (sqdScoreOffset 2) 16 (sqdScoreOffset 0)
(saImad (sqdScoreOffset 0) sqdT0 sksScoreGradientRowBytes (sqdScoreOffset 0)
(saImad (sqdScoreOffset 0) (sqdScoreOffset 1) (sksScoreGradientSubTileBytes seq) (sqdScoreOffset 0)
(saImad (sqdScoreOffset 0) saTileIndex (naturalMultiply 2 (sksScoreGradientSubTileBytes seq)) (sqdScoreOffset 0)
(saImad (sqdScoreOffset 0) saHead (naturalMultiply seq (naturalMultiply seq saHalfBytes)) (sqdScoreOffset 0)
(saFor 3 (lambda unrestricted i : Nat .
(saAddImm (sqdKeyOffset (succ i)) (sqdKeyOffset 0) (naturalMultiply (succ i) (naturalMultiply 16 (naturalMultiply saHeadWidth saHalfBytes)))))
(saFor 3 (lambda unrestricted i : Nat .
(saAddImm (sqdScoreOffset (succ i)) (sqdScoreOffset 0) (naturalMultiply (succ i) (naturalMultiply 16 sksScoreGradientRowBytes))))
(sqdCopyTile seq (sqdStage 0) 0
(nat-eliminate (lambda unrestricted n : Nat . (family SM86Program))
tail
(lambda unrestricted p : Nat . (lambda unrestricted ignored : (family SM86Program) .
(saGreater saP1 saCount 1
(sqdCopyTile seq (sqdStage 1) 1 tail))))
(naturalSaturatingSubtract ahead 1))))))))))))))))))))))))))))))))))))))))))))))
def sqdSetup = (sqdSetupAhead 2)
def sqdPrologueAhead = (lambda unrestricted ahead : Nat . (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(lambda unrestricted rest : (family SM86Program) .
(let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
(let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
(saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
(saS2R saTileIndex (saTileAxis order)
(saS2R saKeyValueHead (saKeyValueAxis order)
(saS2R saHead (saGroupAxis order)
(saMovImm saOne 1
(saMovImmAfter saScratch 31 saWait5
(saImad saHead saKeyValueHead groups saHead
(saCountDownTile order (naturalDivideUnchecked seq saTile)
(saAnd saLane saTid saScratch
(saShr saWarp saTid 5
(saShr saQuad saLane 2
(saMovImm saScratch 3
(saAnd saQuadLane saLane saScratch
(saImad saRow saWarp 16 saQuad
(saImad saRow saTileIndex saTile saRow
-- dQ's offset (binary32 rows), as the query pass's, for its epilogue
(saMovImm saScratch 0
(saImad saScratch saRow (naturalMultiply saHeadWidth 4) saScratch
(saImad saScratch saQuadLane 8 saScratch
(saImad saScratch saHead (naturalMultiply seq (naturalMultiply saHeadWidth 4)) saScratch
-- key tiles 0 .. i, dQ from +0
(saAddImm saCount saTileIndex 1
(saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
(sqdSetupAhead ahead seq headBytes rest))))))))))))))))))))))))))))))
def sqdPrologue = (sqdPrologueAhead 2)
-- the three stages' bodies; the loop leaves after the first or second when
-- no tile is left and returns from the third to the first
def sqdBody = (lambda unrestricted seq : Nat . (lambda unrestricted b : Nat . (sqdTileBody seq (sqdStage b) sm86ProgramEmpty)))
def streamingAttentionQueryFromScoresOrderedWithSM86 =
(lambda unrestricted epilogue : (family SM86Program) .
(lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(let unrestricted length = (sm86ProgramCount (sqdBody seq 0)) in
(sqdPrologue order seq heads keyValueHeads
(sm86ProgramAppend (sqdBody seq 0)
(sm86ProgramAppend (sfFirstTail (naturalAdd length (naturalAdd 3 length)))
(sm86ProgramAppend (sqdBody seq 1)
(sm86ProgramAppend (sfFirstTail length)
(sm86ProgramAppend (sqdBody seq 2)
(sm86ProgramAppend (sfSecondTail (naturalAdd length (naturalAdd 3 length)) length)
epilogue)))))))))))))
def streamingAttentionQueryFromScoresOrderedUncheckedSM86 =
(streamingAttentionQueryFromScoresOrderedWithSM86 sqsEpilogue)
def streamingAttentionQueryFromScoresPathSM86 =
(lambda unrestricted body : (family SM86Program) . (lambda unrestricted seq : Nat .
(lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(sqdPrologue 1 seq heads keyValueHeads (sm86ProgramAppend body sqsEpilogue))))))
def streamingAttentionQueryFromScoresHeavyFirstSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionQueryFromScoresOrderedUncheckedSM86 1 seq heads keyValueHeads)))))
-- dQ stored as the projection gradient's inverse-rotated bfloat16 Q slice
-- (sqRotatedEpilogue): a power-of-two sequence, the query heads within the
-- projection's width, its byte offsets within 32 bits
def streamingAttentionQueryFromScoresRotatedAdmitted =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat . (lambda unrestricted width : Nat .
(naturalAnd (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads)
(naturalAnd (naturalEqual (naturalPowerOfTwo (sqrLog2 seq)) seq)
(naturalAnd (naturalLessOrEqual (naturalMultiply heads saHeadWidth) width)
(naturalLess (naturalMultiply seq (naturalMultiply width saHalfBytes)) (naturalPowerOfTwo 32)))))))))
def sqsRotatedEpilogue = (lambda unrestricted seq : Nat . (lambda unrestricted width : Nat .
(sqRotatedEpilogue seq width (saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty))))
def streamingAttentionQueryFromScoresRotatedPathSM86 =
(lambda unrestricted body : (family SM86Program) . (lambda unrestricted seq : Nat .
(lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat . (lambda unrestricted width : Nat .
(sqdPrologue 1 seq heads keyValueHeads (sm86ProgramAppend body (sqsRotatedEpilogue seq width))))))))
def streamingAttentionQueryFromScoresRotatedHeavyFirstSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat . (lambda unrestricted width : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionQueryFromScoresRotatedAdmitted seq heads keyValueHeads width) 1) .
(streamingAttentionQueryFromScoresOrderedWithSM86 (sqsRotatedEpilogue seq width) 1 seq heads keyValueHeads))))))
-- dQ from the stored dS^T stored into a fused projection's gradient
-- (sqProjectionEpilogue: binary32 at argument 7, bfloat16 at 12)
def streamingAttentionQueryFromScoresProjectionConstantBytes : Nat = (saArgument 13)
def streamingAttentionQueryFromScoresProjectionHeavyFirstSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat . (lambda unrestricted width : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionQueryProjectionAdmitted seq heads keyValueHeads width) 1) .
(streamingAttentionQueryFromScoresOrderedWithSM86
(sqProjectionEpilogue seq width (saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty))
1 seq heads keyValueHeads))))))
-- ---- the same in two shared stages ----
-- One key tile from `buffer` while the next is copied into `other` (the
-- forward pass's double buffer): 32 KB of shared stages, so with the
-- program's own registers (Realization.Nvidia.SM86.RegisterCompactionSM86:
-- 160) three blocks fit an SM. Every product, fragment and order is the
-- three-stage pass's.
def sqd2TileBody = (lambda unrestricted seq : Nat . (lambda unrestricted buffer : Nat . (lambda unrestricted other : Nat .
(lambda unrestricted tail : (family SM86Program) .
(sfBarrier
(saGreater saP1 saCount 1
(sqdCopyTile seq other 1
(sfWaitGroups 1 saWait5
(sfBarrier
(sqdLoadScores (naturalAdd buffer sqdTileBytes) saWait3
(sqdLoadKeys buffer saWaitNone
(sfValues
tail))))))))))))
def sqd2BodyA = (lambda unrestricted seq : Nat . (sqd2TileBody seq (sqdStage 0) (sqdStage 1) sm86ProgramEmpty))
def sqd2BodyB = (lambda unrestricted seq : Nat . (sqd2TileBody seq (sqdStage 1) (sqdStage 0) sm86ProgramEmpty))
def streamingAttentionQueryFromScoresDoubleProjectionHeavyFirstSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat . (lambda unrestricted width : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionQueryProjectionAdmitted seq heads keyValueHeads width) 1) .
(sqdPrologueAhead 1 1 seq heads keyValueHeads
(sm86ProgramAppend (sqd2BodyA seq)
(sm86ProgramAppend (sfFirstTail (sm86ProgramCount (sqd2BodyB seq)))
(sm86ProgramAppend (sqd2BodyB seq)
(sm86ProgramAppend (sfSecondTail (sm86ProgramCount (sqd2BodyA seq)) (sm86ProgramCount (sqd2BodyB seq)))
(sqProjectionEpilogue seq width (saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty))))))))))))
def streamingAttentionQueryFromScoresDoubleSharedBytes : Nat = (naturalMultiply 2 sqdStageBytes)
def streamingAttentionQueryFromScoresSM86Registers : Nat = sqRegisters
def streamingAttentionQueryFromScoresSM86Threads : Nat = saThreads
def streamingAttentionQueryFromScoresSharedBytes : Nat = (naturalMultiply sqdStages sqdStageBytes)
-- The key launch (streamingAttentionKeySM86): a block of two warps per 32
-- keys of a K/V head, grid (T / 32, K/V heads, groups). A warp holds its 16 keys' K and V
-- rows as A fragments and walks the 32-query sub-tiles from the last one
-- down to its diagonal (the last, masked): S^T = K Q^T, dP^T = V dO^T,
-- P^T = 2^(S^T c - L), dS^T = P^T (dP^T - D) with L and D the queries',
-- dV += P^T dO and dK += dS^T Q. Operands: Q, K, V, dO ([head][T][64],
-- half), Q^T, dO^T ([head][64][T], half), L and D ([head][T]); results dK
-- ([head][T][64]) and dV^T ([head][64][T]), binary32. Parameters: Q, K,
-- V, dO, Q^T, dO^T, L, D, dK, dV^T, then c and s.
def skK = (lambda unrestricted kt : Nat . (naturalAdd 20 (naturalMultiply kt 4)))
def skV = (lambda unrestricted kt : Nat . (naturalAdd 36 (naturalMultiply kt 4)))
def skS = (lambda unrestricted nt : Nat . (naturalAdd 52 (naturalMultiply nt 4)))
def skDP = (lambda unrestricted nt : Nat . (naturalAdd 68 (naturalMultiply nt 4)))
def skDV = (lambda unrestricted nt : Nat . (naturalAdd 84 (naturalMultiply nt 4)))
def skDK = (lambda unrestricted nt : Nat . (naturalAdd 116 (naturalMultiply nt 4)))
-- the B fragments of slice k, n-tile nt among `tiles` n-tiles
def skB = (lambda unrestricted tiles : Nat . (lambda unrestricted k : Nat . (lambda unrestricted nt : Nat .
(naturalAdd 148 (naturalMultiply (naturalAdd (naturalMultiply k tiles) nt) 2)))))
def skPA = (lambda unrestricted kk : Nat . (naturalAdd 180 (naturalMultiply kk 4)))
def skSA = (lambda unrestricted kk : Nat . (naturalAdd 188 (naturalMultiply kk 4)))
-- -L and -D of the thread's columns: n-tile nt, column b
def skNegL = (lambda unrestricted nt : Nat . (lambda unrestricted b : Nat . (naturalAdd 196 (naturalAdd (naturalMultiply nt 2) b))))
def skNegD = (lambda unrestricted nt : Nat . (lambda unrestricted b : Nat . (naturalAdd 204 (naturalAdd (naturalMultiply nt 2) b))))
def skTmp = (lambda unrestricted n : Nat . (naturalAdd 212 n))
def skHead : Nat = 220
def skRow : Nat = 221
def skLogOffset : Nat = 222
def skKeyValueHead : Nat = 223
-- dK stores remain in flight while the dV address is formed. Their pair
-- cannot share the address registers used by the later dV stores.
def skKeyGradientStorePointer : Nat = 224
def skRegisters : Nat = 232
def skSubTile : Nat = 32
def skThreads : Nat = 64
-- The wide grouped launch can leave the 3090's channel waiting with no
-- retired fence. A per-head launch uses the same body and output planes but
-- reads its K/V head and within-group query head from one CB0 slot. Keeping
-- this selection in the realization avoids copying the backward kernel into
-- a system; the launch schedule determines the head traversal.
def skLoadHeadIndicesOrdered = (lambda unrestricted order : Nat . (lambda unrestricted fromConstant : Nat .
(lambda unrestricted tail : (family SM86Program) .
(nat-eliminate (lambda unrestricted mode : Nat . (family SM86Program))
(saS2R skKeyValueHead (saKeyValueAxis order)
(saS2R skHead (saGroupAxis order) tail))
(lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86Program) .
(saMovConst skKeyValueHead (naturalAdd (saArgument 12) 4)
(saMovConst skHead (saArgument 12) tail))))
fromConstant))))
def skLoadHeadIndices = (skLoadHeadIndicesOrdered 0)
-- A key launch may own a contiguous subset of key tiles. The caller derives
-- the subset's base from its launch schedule; the ordinary unsplit kernel
-- keeps using CTA x directly. This keeps all query-loop and output-address
-- semantics in one implementation while bounding the work of one launch.
def skLoadTileBase = (lambda unrestricted fromConstant : Nat .
(lambda unrestricted tail : (family SM86Program) .
(nat-eliminate (lambda unrestricted mode : Nat . (family SM86Program))
tail
(lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86Program) .
-- Tile rebasing consumes CTA x before the later head-selection
-- wait. S2R is variable latency: the immediate-result machine model
-- cannot prove this ordering. Drain SB5 here, before its first use;
-- Proof.CheckedAttentionGridSelection checks the control dependency.
(saMovConstAfter saScratch (saArgument 13) saWait5
(saImad saTileIndex saScratch 1 saTileIndex tail))))
fromConstant)))
-- rows ([T][64] half) as the B fragments of a product over the 64 columns:
-- slices k 0..3, n-tiles nt 0..3 (the sub-tile's 32 rows)
def skLoadRows = (lambda unrestricted pointer : Nat . (lambda unrestricted wait : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted k = (naturalDivideUnchecked i 4) in (let unrestricted nt = (naturalModuloUnchecked i 4) in
(lambda unrestricted rest : (family SM86Program) .
(saLoad (skB 4 k nt) pointer (naturalAdd (naturalMultiply nt 1024) (naturalMultiply k 32)) saSB0 (naturalSelect (naturalIsZero i) wait saWaitNone)
(saLoad (succ (skB 4 k nt)) pointer (naturalAdd (naturalMultiply nt 1024) (naturalAdd (naturalMultiply k 32) 16)) saSB0 saWaitNone
rest))))))
tail))))
-- transposed rows ([64][T] half) as the B fragments of a product over the
-- sub-tile's 32 positions: slices kk 0..1, n-tiles nt 0..7 (the 64 columns)
def skLoadColumns = (lambda unrestricted seq : Nat . (lambda unrestricted pointer : Nat . (lambda unrestricted wait : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted k = (naturalDivideUnchecked i 8) in (let unrestricted nt = (naturalModuloUnchecked i 8) in
(let unrestricted rowOffset = (naturalMultiply nt (naturalMultiply 8 (naturalMultiply seq saHalfBytes))) in
(lambda unrestricted rest : (family SM86Program) .
(saLoad (skB 8 k nt) pointer (naturalAdd rowOffset (naturalMultiply k 32)) saSB0 (naturalSelect (naturalIsZero i) wait saWaitNone)
(saLoad (succ (skB 8 k nt)) pointer (naturalAdd rowOffset (naturalAdd (naturalMultiply k 32) 16)) saSB0 saWaitNone
rest)))))))
tail)))))
def skProducts = (lambda unrestricted initialize : Nat . (lambda unrestricted accumulator : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted a : (pi unrestricted k : Nat . Nat) .
(lambda unrestricted slices : Nat . (lambda unrestricted tiles : Nat . (lambda unrestricted barrier : (family SM86Barrier) .
(lambda unrestricted wait : Nat . (lambda unrestricted first : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor slices (lambda unrestricted k : Nat .
(saFor tiles (lambda unrestricted nt : Nat .
(saHmmaFrom (accumulator nt) (a k) (skB tiles k nt)
(naturalSelect (naturalIsZero k)
(naturalSelect (naturalIsZero initialize) (accumulator nt) sm86RegisterZero)
(accumulator nt)) barrier
(naturalSelect (naturalIsZero nt) (naturalSelect (naturalIsZero k) first wait) saWaitNone)))))
tail))))))))))
-- the diagonal sub-tile's mask: element (key r + 8h, query 8 nt + 2 (lane % 4)
-- + b) is kept when the query is not before the key: saMask holds
-- 2 (lane % 4) - 16 warp - lane / 4 + 128, plus 4096 for every sub-tile
-- before the diagonal one
def skMaskTile = (lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
(let unrestricted b = (naturalModuloUnchecked i 2) in
(let unrestricted threshold = (naturalSaturatingSubtract (naturalAdd 127 (naturalMultiply 8 h)) (naturalAdd (naturalMultiply 8 nt) b)) in
(lambda unrestricted rest : (family SM86Program) .
(saGreater saP0 saMask threshold
(saUnless saP0 (constructor SM86InstructionBody SM86MoveImmediate (saR (saElem (skS nt) h b)) (saU saMinusInfinity) saPlain)
rest))))))))
tail))
-- The key pass's mask likewise changes only its last (diagonal) sub-tile:
-- while another follows (P1), a forward branch skips its 16 compares and
-- 16 predicated moves.
def sksMaskBytes : Nat = (naturalMultiply 16 32)
def sksMaskDiagonal = (lambda unrestricted tail : (family SM86Program) .
(sfMaskBranch sksMaskBytes (skMaskTile tail)))
def sksMaskSkipped = (lambda unrestricted tail : (family SM86Program) .
(sfMaskBranch sksMaskBytes tail))
-- Pack a 16-element score plane into two eight-half A fragments. Source
-- tiles are visited before later packed destinations can overwrite them;
-- destination may therefore start at the source plane's first register.
-- The checked packing root covers every destination word, signed zero and
-- BF16 ties. This is a register transformation, not a change of precision.
def skPackRowFragments = (lambda unrestricted source : (pi unrestricted nt : Nat . Nat) .
(lambda unrestricted destination : (pi unrestricted kk : Nat . Nat) .
(lambda unrestricted wait : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 2 (lambda unrestricted kk : Nat .
(let unrestricted left = (source (naturalMultiply 2 kk)) in
(let unrestricted right = (source (succ (naturalMultiply 2 kk))) in
(lambda unrestricted rest : (family SM86Program) .
(saOp (constructor SM86InstructionBody SM86FloatPairToPackedHalfPair
(saR (destination kk)) (saR (succ left)) (saR left)
(saAfter (naturalSelect (naturalIsZero kk) wait saWaitNone)))
(saPack (succ (destination kk)) (naturalAdd left 3) (naturalAdd left 2)
(saPack (naturalAdd (destination kk) 2) (succ right) right
(saPack (naturalAdd (destination kk) 3) (naturalAdd right 3) (naturalAdd right 2)
rest)))))))) tail)))))
-- Multiply a probability by its centred cotangent. Keeping the destination
-- explicit permits either input plane to be reused after its last consumer;
-- the add then multiply and their binary32 rounding order stay unchanged.
def skWeightedDifference = (lambda unrestricted destination : Nat .
(lambda unrestricted probability : Nat . (lambda unrestricted cotangent : Nat .
(lambda unrestricted negativeRowTerm : Nat . (lambda unrestricted wait : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFadd cotangent cotangent negativeRowTerm wait
(saFmul destination probability cotangent saWaitNone tail))))))))
-- P^T = 2^(t - L), dS^T = P^T (dP^T - D), in S^T's registers
def skScoreGradient = (lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
(let unrestricted b = (naturalModuloUnchecked i 2) in
(let unrestricted e = (saElem (skS nt) h b) in
(lambda unrestricted rest : (family SM86Program) .
(saFfma e e saScale (skNegL nt b)
(saEx2 e e saWaitNone rest))))))))
-- P^T's fragments, before dS^T replaces it
(skPackRowFragments skS skPA saWait4
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
(let unrestricted b = (naturalModuloUnchecked i 2) in
(let unrestricted e = (saElem (skS nt) h b) in (let unrestricted g = (saElem (skDP nt) h b) in
(lambda unrestricted rest : (family SM86Program) .
(skWeightedDifference e e g (skNegD nt b) saWaitNone rest))))))))
(skPackRowFragments skS skSA saWaitNone
tail)))))
def skTileBody = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
(saImad saMask saCount 4096 saMaskBase
-- the sub-tile's L and D (their negatives)
(saWide saPointer skLogOffset saOne (saArgument 6)
(saFor 8 (lambda unrestricted i : Nat .
(saLoad (skTmp i) saPointer (naturalAdd (naturalMultiply (naturalDivideUnchecked i 2) 32) (naturalMultiply (naturalModuloUnchecked i 2) 4))
saSB1 (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
(saFor 8 (lambda unrestricted i : Nat .
(sbNegate (skNegL (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2)) (skTmp i) (naturalSelect (naturalIsZero i) saWait1 saWaitNone)))
(saWide saPointer skLogOffset saOne (saArgument 7)
(saFor 8 (lambda unrestricted i : Nat .
(saLoad (skTmp i) saPointer (naturalAdd (naturalMultiply (naturalDivideUnchecked i 2) 32) (naturalMultiply (naturalModuloUnchecked i 2) 4))
saSB1 saWaitNone))
(saFor 8 (lambda unrestricted i : Nat .
(sbNegate (skNegD (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2)) (skTmp i) (naturalSelect (naturalIsZero i) saWait1 saWaitNone)))
-- S^T = K Q^T
(saWide saKeyPointer saKeyOffset saOne (saArgument 0)
(skLoadRows saKeyPointer saWait3
-- Keep the qualified ordinary schedule. A first-product RZ candidate
-- removed these clears but gave overlapping 100-step timings and higher
-- measured energy; see the 2026-10-01 attention-product-origin fixture.
(saFor 16 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 52 i) 0))
(saFor 16 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
(skProducts 0 skS skK 4 4 saSB2 saWait2 (naturalAdd saWait0 saWait1)
-- dP^T = V dO^T
(saWide saKeyPointer saKeyOffset saOne (saArgument 3)
(skLoadRows saKeyPointer saWait2
(skProducts 0 skDP skV 4 4 saSB2 saWait2 saWait0
-- dO^T's fragments for dV, once dP^T's products are done
(saWide saValuePointer saValueOffset saOne (saArgument 5)
(skLoadColumns seq saValuePointer saWait2
(skMaskTile
(skScoreGradient
(skProducts 0 skDV skPA 2 8 saSB3 saWait3 saWait0
-- Q^T's fragments for dK, once dV's products have read dO^T's
(saWide saValuePointer saValueOffset saOne (saArgument 4)
(skLoadColumns seq saValuePointer saWait3
(skProducts 0 skDK skSA 2 8 saSB3 saWait3 saWait0
tail)))))))))))))))))))))))))
def skLoopTail = (lambda unrestricted bodyInstructions : Nat .
(saAddImm saKeyOffset saKeyOffset (naturalSaturatingSubtract 4294967296 (naturalMultiply skSubTile (naturalMultiply saHeadWidth saHalfBytes)))
(saAddImm saValueOffset saValueOffset (naturalSaturatingSubtract 4294967296 (naturalMultiply skSubTile saHalfBytes))
(saAddImm skLogOffset skLogOffset (naturalSaturatingSubtract 4294967296 (naturalMultiply skSubTile 4))
(saAddImm saCount saCount 4294967295
(saGreater saP1 saCount 0
(saWhen saP1
(constructor SM86InstructionBody SM86Branch
(saU (naturalSaturatingSubtract 4294967296 (naturalMultiply 16 (naturalAdd bodyInstructions 6))))
(sm86Unsigned32 (byte 255) (byte 255) (byte 131) (byte 3))
saPlain)
sm86ProgramEmpty)))))))
-- dK = s (the sum), rows; dV^T, columns: element (key r + 8h, column
-- 8 nt + 2 (lane % 4) + b) at (8 nt + b) seq + 8 h words past the thread's
def skEpilogue = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
(saMovConst (skTmp 0) (saArgument 11)
(saFor 32 (lambda unrestricted i : Nat .
(saFmul (naturalAdd 116 i) (naturalAdd 116 i) (skTmp 0) (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
(saWide skKeyGradientStorePointer saScratch saOne (saArgument 8)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 2) in (let unrestricted h = (naturalModuloUnchecked i 2) in
(saOp (sbStore64 skKeyGradientStorePointer (saElem (skDK nt) h 0)
(naturalAdd (naturalMultiply nt 32) (naturalMultiply h (naturalMultiply 8 (naturalMultiply saHeadWidth 4))))
saWaitNone)))))
-- dV^T's offset: head h's 64 x seq plane, the quad lane's first column,
-- the key
(saMovImm (skTmp 1) 0
(saImad (skTmp 1) saQuadLane (naturalMultiply 8 seq) (skTmp 1)
(saImad (skTmp 1) skHead (naturalMultiply saHeadWidth (naturalMultiply seq 4)) (skTmp 1)
(saImad (skTmp 1) skRow 4 (skTmp 1)
(saWide saPointer (skTmp 1) saOne (saArgument 9)
(saFor 32 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
(let unrestricted b = (naturalModuloUnchecked i 2) in
(saOp (saStore saPointer (saElem (skDV nt) h b)
(naturalAdd (naturalMultiply (naturalAdd (naturalMultiply 8 nt) b) (naturalMultiply seq 4)) (naturalMultiply h 32))
saWaitNone))))))
tail))))))))))))
-- dK = s (the sum) and dV stored straight into a fused projection's
-- gradient, token-major [seq][width] (the launch points the arguments at the
-- key and value columns): dK binary32 at argument 8 and bfloat16 at 15, dV
-- binary32 at 9 and bfloat16 at 16 (14 is the shared pass's dS^T), instead of the [head][seq][64] dK and
-- [head][64][seq] dV^T planes the inverse layout passes copied and
-- converted. dK takes skEpilogue's scaling FMUL; dV, as there, none. Each
-- pair is stored as binary32 from its accumulators, then packed by
-- F2FP.BF16 (the even column low) into a register of its own and stored.
-- Position, head and column come from dK's binary32 offset (as
-- sqProjectionEpilogue); dV's rows and columns are dK's. Registers: the K,
-- V and S fragments' (dead once the products are done).
def skpElement : Nat = 52
def skpColumn : Nat = 53
def skpRow : Nat = 54
def skpPosition : Nat = 55
def skpHead : Nat = 56
def skpOutput : Nat = 57
def skpBytes : Nat = 58
def skpPointer : Nat = 60
def skpHalfPointer : Nat = 62
def skpValuePointer : Nat = 64
def skpValueHalfPointer : Nat = 66
def skpPacked = (lambda unrestricted i : Nat . (naturalAdd 20 i))
def skpValuePacked = (lambda unrestricted i : Nat . (naturalAdd 36 i))
def streamingAttentionKeyProjectionHalfArgument : Nat = 15
def streamingAttentionKeyProjectionValueHalfArgument : Nat = 16
def streamingAttentionKeyProjectionConstantBytes : Nat = (saArgument 17)
-- the 16 pairs of `accumulator` (n-tile nt, row half h) from `pointer`
-- (binary32) and `halfPointer` (bfloat16), each packed into `packed` i
def skpStorePairs = (lambda unrestricted width : Nat .
(lambda unrestricted accumulator : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted packed : (pi unrestricted i : Nat . Nat) .
(lambda unrestricted pointer : Nat . (lambda unrestricted halfPointer : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 2) in (let unrestricted h = (naturalModuloUnchecked i 2) in
(lambda unrestricted after : (family SM86Program) .
(saOp (sbStore64 pointer (saElem (accumulator nt) h 0)
(naturalAdd (naturalMultiply nt 32) (naturalMultiply h (naturalMultiply 8 (naturalMultiply width 4))))
saWaitNone)
(saOp (constructor SM86InstructionBody SM86FloatPairToPackedHalfPair
(saR (packed i)) (saR (saElem (accumulator nt) h 1)) (saR (saElem (accumulator nt) h 0)) saPlain)
(saOp (saStore halfPointer (packed i)
(naturalAdd (naturalMultiply nt 16) (naturalMultiply h (naturalMultiply 8 (naturalMultiply width saHalfBytes))))
saWaitNone)
after)))))))
tail)))))))
def skProjectionEpilogue = (lambda unrestricted width : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
(saMovConst (skTmp 0) (saArgument 11)
(saFor 32 (lambda unrestricted i : Nat .
(saFmul (naturalAdd 116 i) (naturalAdd 116 i) (skTmp 0) (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
(saShr skpElement saScratch 2
(saMovImm skpColumn 63
(saAnd skpColumn skpElement skpColumn
(saShr skpRow skpElement 6
(saMovImm skpPosition (naturalSaturatingSubtract seq 1)
(saAnd skpPosition skpRow skpPosition
(saShr skpHead skpRow (sqrLog2 seq)
(saImad skpOutput skpPosition width skpColumn
(saImad skpOutput skpHead saHeadWidth skpOutput
(saMovImm skpBytes 4
(saWide skpPointer skpOutput skpBytes (saArgument 8)
(saWide skpValuePointer skpOutput skpBytes (saArgument 9)
(saMovImm skpBytes saHalfBytes
(saWide skpHalfPointer skpOutput skpBytes (saArgument streamingAttentionKeyProjectionHalfArgument)
(saWide skpValueHalfPointer skpOutput skpBytes (saArgument streamingAttentionKeyProjectionValueHalfArgument)
(skpStorePairs width skDK skpPacked skpPointer skpHalfPointer
(skpStorePairs width skDV skpValuePacked skpValuePointer skpValueHalfPointer
tail))))))))))))))))))))))
-- a power-of-two sequence, the heads within the width, its binary32 byte
-- offsets within 32 bits
def streamingAttentionKeyProjectionAdmitted =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat . (lambda unrestricted width : Nat .
(naturalAnd (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads)
(naturalAnd (naturalEqual (naturalPowerOfTwo (sqrLog2 seq)) seq)
(naturalAnd (naturalLessOrEqual (naturalMultiply heads saHeadWidth) width)
(naturalLess (naturalMultiply seq (naturalMultiply width 4)) (naturalPowerOfTwo 32)))))))))
-- Each query head writes a separate dK/dV plane. A later reduction sums
-- those planes into the compact K/V gradient; concurrent CTAs never race on
-- one compact gradient address.
def streamingAttentionKeyGroupedOrderedWithSM86 = (lambda unrestricted epilogue : (pi unrestricted seq : Nat . (pi unrestricted tail : (family SM86Program) . (family SM86Program))) . (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat . (lambda unrestricted fromConstant : Nat . (lambda unrestricted fromTileBase : Nat .
(let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
(let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
(let unrestricted subTiles = (naturalDivideUnchecked seq skSubTile) in
(let unrestricted body = (skTileBody seq sm86ProgramEmpty) in
(saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
(saS2R saTileIndex (saTileAxis order)
(skLoadTileBase fromTileBase
(skLoadHeadIndicesOrdered order fromConstant
(saMovImm saOne 1
(saMovImmAfter saScratch 31 saWait5
(saImad skHead skKeyValueHead groups skHead
(saAnd saLane saTid saScratch
(saShr saWarp saTid 5
(saShr saQuad saLane 2
(saMovImm saScratch 3
(saAnd saQuadLane saLane saScratch
-- the thread's key row: 32 tile + 16 warp + lane / 4
(saImad skRow saWarp 16 saQuad
(saImad skRow saTileIndex skSubTile skRow
-- K's and V's A fragments
(saMovImm saScratch 0
(saImad saScratch skRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
(saImad saScratch saQuadLane 4 saScratch
(saImad saScratch skKeyValueHead headBytes saScratch
(saWide saPointer saScratch saOne (saArgument 1)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
(saLoad (naturalAdd (skK kt) j) saPointer
(naturalAdd (naturalMultiply kt 32)
(naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
(naturalMultiply (naturalDivideUnchecked j 2) 16)))
saSB1 saWaitNone))))
(saWide saPointer saScratch saOne (saArgument 2)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
(saLoad (naturalAdd (skV kt) j) saPointer
(naturalAdd (naturalMultiply kt 32)
(naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
(naturalMultiply (naturalDivideUnchecked j 2) 16)))
saSB1 saWaitNone))))
-- dK's offset (binary32 rows) in saScratch
(saMovImm saScratch 0
(saImad saScratch skRow (naturalMultiply saHeadWidth 4) saScratch
(saImad saScratch saQuadLane 8 saScratch
(saImad saScratch skHead (naturalMultiply seq (naturalMultiply saHeadWidth 4)) saScratch
-- the queries start at the last sub-tile: Q's and dO's rows (row lane / 4,
-- the quad lane's pair), Q^T's and dO^T's columns (row lane / 4), L's and D's
(saMovImm saKeyOffset 0
(saImad saKeyOffset saQuad (naturalMultiply saHeadWidth saHalfBytes) saKeyOffset
(saImad saKeyOffset saQuadLane 4 saKeyOffset
(saImad saKeyOffset skHead headBytes saKeyOffset
(saAddImm saKeyOffset saKeyOffset (naturalMultiply (naturalSaturatingSubtract subTiles 1) (naturalMultiply skSubTile (naturalMultiply saHeadWidth saHalfBytes)))
(saMovImm saValueOffset 0
(saImad saValueOffset saQuad (naturalMultiply seq saHalfBytes) saValueOffset
(saImad saValueOffset saQuadLane 4 saValueOffset
(saImad saValueOffset skHead headBytes saValueOffset
(saAddImm saValueOffset saValueOffset (naturalMultiply (naturalSaturatingSubtract subTiles 1) (naturalMultiply skSubTile saHalfBytes))
(saMovImm skLogOffset 0
(saImad skLogOffset saQuadLane 8 skLogOffset
(saImad skLogOffset skHead (naturalMultiply seq 4) skLogOffset
(saAddImm skLogOffset skLogOffset (naturalMultiply (naturalSaturatingSubtract subTiles 1) (naturalMultiply skSubTile 4))
-- the sub-tiles from the last down to the diagonal; the mask's base,
-- 2 (lane % 4) - 16 warp - lane / 4 + 128 - 4096
(saMovImm saCount subTiles
(saImad saCount saTileIndex 4294967295 saCount
(saMovImm saMaskBase 0
(saImad saMaskBase saQuadLane 2 saMaskBase
(saImad saMaskBase saWarp 4294967280 saMaskBase
(saImad saMaskBase saQuad 4294967295 saMaskBase
(saAddImm saMaskBase saMaskBase (naturalSaturatingSubtract 4294967296 (naturalSaturatingSubtract 4096 128))
(saMovConst saScale (saArgument 10)
(saFor 64 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 84 i) 0))
(sm86ProgramAppend body
(sm86ProgramAppend (skLoopTail (sm86ProgramCount body))
(epilogue seq
(saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))
def streamingAttentionKeyGroupedOrderedUncheckedSM86 = (streamingAttentionKeyGroupedOrderedWithSM86 skEpilogue)
-- ===========================================================================
-- The key pass with each query sub-tile staged in shared memory, as the
-- forward and query passes. A step's operands are Q's and dO's 32 rows
-- (plain for S^T and dP^T, read transposed by LDSM.T for dK and dV, the same
-- values in the same lanes as the register kernel's Q^T and dO^T loads) and
-- the sub-tile's 32 L and D words (global loads, each pair its own pointer).
-- The 64 threads copy row r + 8 i, chunk c of each plane (thread 8 r + c).
-- The prologue's K and V fragment
-- loads keep their own address pairs, so no load address is ever re-formed
-- before its late read and the late-read guard is not needed.
def sksQueryTile : Nat = 4608
def sksBufferBytes : Nat = 9216
def streamingAttentionKeySharedBytes : Nat = 18432
def sksOffset = (lambda unrestricted i : Nat . (naturalAdd 226 i))
def sksPair = (lambda unrestricted j : Nat . (naturalAdd 232 (naturalMultiply 2 j)))
def sksStore : Nat = 236
def sksLane : Nat = 237
def sksLaneTransposed : Nat = 238
-- c 16 during the setup
def sksChunkBytes : Nat = 240
def sksT0 : Nat = 241
def sksT1 : Nat = 242
-- dS^T for the dQ product (streamingAttentionQueryFromScoresSM86): the
-- [head][key][query] half plane's byte offset of the thread's key row and
-- quad-lane column pair at the current sub-tile, and its address pair
def sksScoreGradientOffset : Nat = 239
def sksScoreGradientPointer : Nat = 230
def sksT2 : Nat = 243
-- head h's seq x seq plane, the key row, the quad lane's pair at the last
-- sub-tile (skHead, skRow and saQuadLane formed)
def sksScoreGradientOffsetAtLast = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
(saMovImm sksScoreGradientOffset 0
(saImad sksScoreGradientOffset saQuadLane 4 sksScoreGradientOffset
(saImad sksScoreGradientOffset skRow sksScoreGradientRowBytes sksScoreGradientOffset
(saImad sksScoreGradientOffset skHead (naturalMultiply seq (naturalMultiply seq saHalfBytes)) sksScoreGradientOffset
(saAddImm sksScoreGradientOffset sksScoreGradientOffset
(naturalMultiply (naturalSaturatingSubtract (naturalDivideUnchecked seq skSubTile) 1) (sksScoreGradientSubTileBytes seq))
tail)))))))
def streamingAttentionKeyScoreGradientArgument : Nat = 14
-- A packed dS^T word, its read barrier SB5: the next sub-tile's row-term
-- loads and copies wait on SB5 before anything rewrites the packed words
-- or the address pair, and each word after the first waits for the one
-- before it to have read the shared pair (one live SB5 token, as the
-- copies'). Word i: slice kk = i / 4, register j = i % 4 of skSA kk -- row
-- r + 8 (j % 2), n-tile 2 kk + j / 2, its column pair.
def sksScoreGradientWord = (lambda unrestricted seq : Nat . (lambda unrestricted value : (pi unrestricted i : Nat . Nat) .
(lambda unrestricted i : Nat .
(let unrestricted j = (naturalModuloUnchecked i 4) in
(let unrestricted nt = (naturalAdd (naturalMultiply 2 (naturalDivideUnchecked i 4)) (naturalDivideUnchecked j 2)) in
(constructor SM86InstructionBody SM86StoreGlobal (saR sksScoreGradientPointer) (saR (value i))
(saU (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 sksScoreGradientRowBytes))
(naturalMultiply nt 16)))
(constructor SM86Control SM86ControlValue (byte 15) (constructor SM86YieldMode SM86Continue)
saNone saSB5 (nat-to-byte (naturalSelect (naturalIsZero i) saWaitNone saWait5)) (byte 0))))))))
def sksPackedScoreGradient = (lambda unrestricted i : Nat .
(naturalAdd (skSA (naturalDivideUnchecked i 4)) (naturalModuloUnchecked i 4)))
-- the sub-tile's 32 x 32 words of dS^T (P^T (dP^T - D) packed as dK's A
-- fragments: the same words the query pass packs for dQ), then the next
-- sub-tile's offset, 32 queries down
def sksStoreScoreGradients = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
(sfWide sksScoreGradientPointer sksScoreGradientOffset (saArgument streamingAttentionKeyScoreGradientArgument) saWait5
(saFor 8 (lambda unrestricted i : Nat . (saOp (sksScoreGradientWord seq sksPackedScoreGradient i)))
(saAddImm sksScoreGradientOffset sksScoreGradientOffset
(naturalSaturatingSubtract 4294967296 (sksScoreGradientSubTileBytes seq))
tail)))))
-- An odd key tile's 32 keys against the 32 queries below its diagonal lie
-- inside the dQ product's diagonal 64 x 64 tile but are never walked here:
-- every one is masked (P = +0), so the query pass's dS there is a signed
-- zero, and a zero product leaves the dQ accumulators' words unchanged (they
-- start at +0, which no sum of zeros turns into -0). Write +0 there, through
-- the offset the loop leaves one sub-tile below the diagonal.
def sksZeroBelowDiagonal = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
(saAnd sksT2 saTileIndex saOne
(saGreater saP0 sksT2 0
(sfWide sksScoreGradientPointer sksScoreGradientOffset (saArgument streamingAttentionKeyScoreGradientArgument) saWait5
(saFor 8 (lambda unrestricted i : Nat .
(saWhen saP0 (sksScoreGradientWord seq (lambda unrestricted k : Nat . sm86RegisterZero) i)))
tail))))))
-- chunk j: Q's (j < 4) or dO's row r + 8 (j % 4), chunk c, through pointer
-- pair j % 2, each copy after the previous one's reads
def sksCopyTile = (lambda unrestricted buffer : Nat . (lambda unrestricted guarded : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 8 (lambda unrestricted j : Nat .
(let unrestricted body =
(constructor SM86InstructionBody SM86LoadGlobalToShared (saR sksStore)
(saU (naturalAdd buffer
(naturalAdd (naturalSelect (naturalLess j 4) 0 sksQueryTile)
(naturalMultiply (naturalModuloUnchecked j 4) (naturalMultiply 8 sfRowBytes)))))
(saR (sksPair (naturalModuloUnchecked j 2))) (saU 0)
(constructor SM86Control SM86ControlValue (byte 15) (constructor SM86YieldMode SM86Continue)
saNone saSB5 (nat-to-byte saWait5) (byte 0))) in
(lambda unrestricted rest : (family SM86Program) .
(saOp (constructor SM86InstructionBody SM86IntegerMultiplyAddWideConstant
(saR (sksPair (naturalModuloUnchecked j 2))) (saR (sksOffset (naturalModuloUnchecked j 4))) (saR saOne)
(byte 0) (saU (naturalSelect (naturalLess j 4) (saArgument 0) (saArgument 3))) saPlain)
(nat-eliminate (lambda unrestricted mode : Nat . (family SM86Program))
(saOp body rest)
(lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86Program) .
(saWhen saP1 body rest)))
guarded)))))
(sfCommit tail)))))
-- B fragments from shared `base`: plain, 4 slices x 4 n-tiles (32 rows) --
-- LDSM.x4 (k, p) n-tiles 2p, 2p + 1 of slice k -- or transposed, 2 slices x
-- 8 n-tiles (64 columns) -- LDSM.T (kk, p) rows 16 kk.., columns 16 p..
def sksLoadFragments = (lambda unrestricted transposed : Nat .
(lambda unrestricted base : Nat . (lambda unrestricted wait : Nat .
(lambda unrestricted tail : (family SM86Program) .
(saFor 8 (lambda unrestricted i : Nat .
(let unrestricted k = (naturalSelect transposed (naturalDivideUnchecked i 4) (naturalDivideUnchecked i 2)) in
(let unrestricted pair = (naturalSelect transposed (naturalModuloUnchecked i 4) (naturalModuloUnchecked i 2)) in
(saOp (constructor SM86InstructionBody SM86LoadSharedMatrix
(saR (skB (naturalSelect transposed 8 4) k (naturalMultiply 2 pair)))
(saR (naturalSelect transposed sksLaneTransposed sksLane))
(saU (naturalAdd base
(naturalSelect transposed
(naturalAdd (naturalMultiply k (naturalMultiply 16 sfRowBytes)) (naturalMultiply pair 32))
(naturalAdd (naturalMultiply pair (naturalMultiply 16 sfRowBytes)) (naturalMultiply k 32)))))
(constructor SM86SharedMatrixCount SM86SharedMatrix4)
(eliminate StdBool (lambda unrestricted c : (family StdBool) . (family SM86SharedMatrixTranspose)) (stdBoolFromNatural transposed)
(branch StdTrue . (constructor SM86SharedMatrixTranspose SM86SharedMatrixTransposed))
(branch StdFalse . (constructor SM86SharedMatrixTranspose SM86SharedMatrixNotTransposed)))
(saCtl saSB1 (naturalSelect (naturalIsZero i) wait saWaitNone)))))))
tail)))))
-- the thread's eight L (or D) words: column 8 nt + 2 (lane % 4) + b, through
-- `pointer`. Each load names SB5 as its read barrier and waits for the one
-- before it, so the pointer can be re-formed for the next sub-tile.
-- (Shared memory would serve these too, but SM86 LDS/STS are encoded with
-- the .X4 address scale: see the 2026-10-02 shared-attention fixture.)
def sksLoadRowTerms = (lambda unrestricted pointer : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 8 (lambda unrestricted i : Nat .
(saOp (constructor SM86InstructionBody SM86LoadGlobal (saR (skTmp i)) (saR pointer)
(saU (naturalAdd (naturalMultiply (naturalDivideUnchecked i 2) 32) (naturalMultiply (naturalModuloUnchecked i 2) 4)))
(constructor SM86Control SM86ControlValue (byte 15) (constructor SM86YieldMode SM86Continue)
saSB1 saSB5 (nat-to-byte saWait5) (byte 0)))))
tail)))
-- the D pointer pair, beside L's (saPointer); the loop never names R6/R7
-- The same loads into the row-term registers themselves (`target` nt b),
-- their data under SB4, so the sub-tile's products run while they are in
-- flight; the terms are negated in place once the products are issued.
-- SB4 is the score pass's EX2 barrier, which nothing between these loads
-- and that negation names.
def sksLoadRowTermsInto = (lambda unrestricted target : (pi unrestricted nt : Nat . (pi unrestricted b : Nat . Nat)) .
(lambda unrestricted pointer : Nat . (lambda unrestricted tail : (family SM86Program) .
(saFor 8 (lambda unrestricted i : Nat .
(saOp (constructor SM86InstructionBody SM86LoadGlobal
(saR (target (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2))) (saR pointer)
(saU (naturalAdd (naturalMultiply (naturalDivideUnchecked i 2) 32) (naturalMultiply (naturalModuloUnchecked i 2) 4)))
(constructor SM86Control SM86ControlValue (byte 15) (constructor SM86YieldMode SM86Continue)
saSB4 saSB5 (nat-to-byte saWait5) (byte 0)))))
tail))))
def sksRowTermPointer : Nat = 6
-- P^T and dS^T as skScoreGradient, with dS^T = P^T (dP^T - D) left in dP^T's
-- registers instead of P^T's: the in-place EX2 reads each score register
-- late (no read barrier), so nothing may write the scores again in the step.
-- The same additions, multiplications and roundings in the same order.
def sksScoreGradient = (lambda unrestricted tail : (family SM86Program) .
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
(let unrestricted b = (naturalModuloUnchecked i 2) in
(let unrestricted e = (saElem (skS nt) h b) in
(lambda unrestricted rest : (family SM86Program) .
(saFfma e e saScale (skNegL nt b)
(saEx2 e e saWaitNone rest))))))))
(skPackRowFragments skS skPA saWait4
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
(let unrestricted b = (naturalModuloUnchecked i 2) in
(let unrestricted e = (saElem (skS nt) h b) in (let unrestricted g = (saElem (skDP nt) h b) in
(lambda unrestricted rest : (family SM86Program) .
(skWeightedDifference g e g (skNegD nt b) saWaitNone rest))))))))
(skPackRowFragments skDP skSA saWaitNone
tail)))))
def sksTileBodyMasked = (lambda unrestricted mask : (pi unrestricted tail : (family SM86Program) . (family SM86Program)) .
(lambda unrestricted seq : Nat . (lambda unrestricted buffer : Nat . (lambda unrestricted other : Nat .
(lambda unrestricted tail : (family SM86Program) .
(sfBarrier
(saGreater saP1 saCount 1
(saFor 4 (lambda unrestricted i : Nat .
(saAddImm (sksOffset i) (sksOffset i) (naturalSaturatingSubtract 4294967296 (naturalMultiply skSubTile (naturalMultiply saHeadWidth saHalfBytes)))))
(sksCopyTile other 1
(sfWaitGroups 1 saWait5
(sfBarrier
(saImad saMask saCount 4096 saMaskBase
-- the sub-tile's L and D, in flight under the products below
(saWide saPointer skLogOffset saOne (saArgument 6)
(sksLoadRowTermsInto skNegL saPointer
(saWide sksRowTermPointer skLogOffset saOne (saArgument 7)
(sksLoadRowTermsInto skNegD sksRowTermPointer
(saAddImm skLogOffset skLogOffset (naturalSaturatingSubtract 4294967296 (naturalMultiply skSubTile 4))
-- S^T = K Q^T
(sksLoadFragments 0 buffer saWait3
(skProducts 1 skS skK 4 4 saSB2 saWait2 saWait1
-- dP^T = V dO^T
(sksLoadFragments 0 (naturalAdd buffer sksQueryTile) saWait2
(skProducts 1 skDP skV 4 4 saSB2 saWait2 saWait1
-- dO^T's fragments for dV, once dP^T's products are done
(sksLoadFragments 1 (naturalAdd buffer sksQueryTile) saWait2
(mask
-- their negatives, once they have arrived (exact, as before)
(saFor 8 (lambda unrestricted i : Nat .
(sbNegate (skNegL (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2))
(skNegL (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2))
(naturalSelect (naturalIsZero i) (naturalAdd saWait4 saWait5) saWaitNone)))
(saFor 8 (lambda unrestricted i : Nat .
(sbNegate (skNegD (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2))
(skNegD (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2)) saWaitNone))
(sksScoreGradient
(skProducts 0 skDV skPA 2 8 saSB3 saWait3 saWait1
-- Q^T's fragments for dK, once dV's products have read dO^T's
(sksLoadFragments 1 buffer saWait3
(skProducts 0 skDK skSA 2 8 saSB3 saWait3 saWait1
(sksStoreScoreGradients seq
tail))))))))))))))))))))))))))))))
def sksTileBody = (sksTileBodyMasked sksMaskDiagonal)
-- lane addresses, the copy role, Q/dO and L/D offsets at the last sub-tile,
-- then that sub-tile into buffer 0
def sksSetup = (lambda unrestricted seq : Nat . (lambda unrestricted headBytes : Nat .
(lambda unrestricted tail : (family SM86Program) .
(let unrestricted last = (naturalSaturatingSubtract (naturalDivideUnchecked seq skSubTile) 1) in
(saShr sksT0 saLane 4
(saMovImm sksT1 7
(saAnd sksT1 saLane sksT1
(saImad sksT1 sksT0 8 sksT1
(saMovImm sksLane 0
(saImad sksLane sksT1 sfRowBytes sksLane
(saShr sksT1 saLane 3
(saAnd sksT1 sksT1 saOne
(saImad sksLane sksT1 16 sksLane
(saMovImm sksLaneTransposed 7
(saAnd sksLaneTransposed saLane sksLaneTransposed
(saImad sksT1 sksT1 8 sksLaneTransposed
(saMovImm sksLaneTransposed 0
(saImad sksLaneTransposed sksT1 sfRowBytes sksLaneTransposed
(saImad sksLaneTransposed sksT0 16 sksLaneTransposed
(saShr sksT0 saTid 3
(saMovImm sksT1 7
(saAnd sksT1 saTid sksT1
(saMovImm sksChunkBytes 0
(saImad sksChunkBytes sksT1 16 sksChunkBytes
(saImad sksStore sksT0 sfRowBytes sksChunkBytes
(saImad (sksOffset 0) sksT0 (naturalMultiply saHeadWidth saHalfBytes) sksChunkBytes
(saImad (sksOffset 0) skHead headBytes (sksOffset 0)
(saAddImm (sksOffset 0) (sksOffset 0) (naturalMultiply last (naturalMultiply skSubTile (naturalMultiply saHeadWidth saHalfBytes)))
(saFor 3 (lambda unrestricted i : Nat .
(saAddImm (sksOffset (succ i)) (sksOffset 0) (naturalMultiply (succ i) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))))
(sksCopyTile 0 0 tail))))))))))))))))))))))))))))))
def sksPrologue = (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat . (lambda unrestricted fromConstant : Nat . (lambda unrestricted fromTileBase : Nat .
(lambda unrestricted rest : (family SM86Program) .
(let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
(let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
(let unrestricted subTiles = (naturalDivideUnchecked seq skSubTile) in
(saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
(saS2R saTileIndex (saTileAxis order)
(skLoadTileBase fromTileBase
(skLoadHeadIndicesOrdered order fromConstant
(saMovImm saOne 1
(saMovImmAfter saScratch 31 saWait5
(saImad skHead skKeyValueHead groups skHead
(saAnd saLane saTid saScratch
(saShr saWarp saTid 5
(saShr saQuad saLane 2
(saMovImm saScratch 3
(saAnd saQuadLane saLane saScratch
-- the thread's key row: 32 tile + 16 warp + lane / 4
(saImad skRow saWarp 16 saQuad
(saImad skRow saTileIndex skSubTile skRow
-- K's and V's A fragments
(saMovImm saScratch 0
(saImad saScratch skRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
(saImad saScratch saQuadLane 4 saScratch
(saImad saScratch skKeyValueHead headBytes saScratch
(saWide saKeyPointer saScratch saOne (saArgument 1)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
(saLoad (naturalAdd (skK kt) j) saKeyPointer
(naturalAdd (naturalMultiply kt 32)
(naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
(naturalMultiply (naturalDivideUnchecked j 2) 16)))
saSB1 saWaitNone))))
(saWide saValuePointer saScratch saOne (saArgument 2)
(saFor 16 (lambda unrestricted i : Nat .
(let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
(saLoad (naturalAdd (skV kt) j) saValuePointer
(naturalAdd (naturalMultiply kt 32)
(naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
(naturalMultiply (naturalDivideUnchecked j 2) 16)))
saSB1 saWaitNone))))
-- dK's offset (binary32 rows) in saScratch
(saMovImm saScratch 0
(saImad saScratch skRow (naturalMultiply saHeadWidth 4) saScratch
(saImad saScratch saQuadLane 8 saScratch
(saImad saScratch skHead (naturalMultiply seq (naturalMultiply saHeadWidth 4)) saScratch
-- the queries start at the last sub-tile: Q's and dO's rows (row lane / 4,
-- the quad lane's pair), Q^T's and dO^T's columns (row lane / 4), L's and D's
(saMovImm saKeyOffset 0
(saImad saKeyOffset saQuad (naturalMultiply saHeadWidth saHalfBytes) saKeyOffset
(saImad saKeyOffset saQuadLane 4 saKeyOffset
(saImad saKeyOffset skHead headBytes saKeyOffset
(saAddImm saKeyOffset saKeyOffset (naturalMultiply (naturalSaturatingSubtract subTiles 1) (naturalMultiply skSubTile (naturalMultiply saHeadWidth saHalfBytes)))
(saMovImm saValueOffset 0
(saImad saValueOffset saQuad (naturalMultiply seq saHalfBytes) saValueOffset
(saImad saValueOffset saQuadLane 4 saValueOffset
(saImad saValueOffset skHead headBytes saValueOffset
(saAddImm saValueOffset saValueOffset (naturalMultiply (naturalSaturatingSubtract subTiles 1) (naturalMultiply skSubTile saHalfBytes))
(saMovImm skLogOffset 0
(saImad skLogOffset saQuadLane 8 skLogOffset
(saImad skLogOffset skHead (naturalMultiply seq 4) skLogOffset
(saAddImm skLogOffset skLogOffset (naturalMultiply (naturalSaturatingSubtract subTiles 1) (naturalMultiply skSubTile 4))
(sksScoreGradientOffsetAtLast seq
-- the sub-tiles from the last down to the diagonal; the mask's base,
-- 2 (lane % 4) - 16 warp - lane / 4 + 128 - 4096
(saMovImm saCount subTiles
(saImad saCount saTileIndex 4294967295 saCount
(saMovImm saMaskBase 0
(saImad saMaskBase saQuadLane 2 saMaskBase
(saImad saMaskBase saWarp 4294967280 saMaskBase
(saImad saMaskBase saQuad 4294967295 saMaskBase
(saAddImm saMaskBase saMaskBase (naturalSaturatingSubtract 4294967296 (naturalSaturatingSubtract 4096 128))
(saMovConst saScale (saArgument 10)
(saFor 64 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 84 i) 0))
(sksSetup seq headBytes rest)))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))
def sksEpilogue = (lambda unrestricted seq : Nat .
(sksZeroBelowDiagonal seq
(skEpilogue seq (saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty))))
-- the bodies at sequence length `seq` (dS^T's row stride)
def sksBodyA = (lambda unrestricted seq : Nat . (sksTileBody seq 0 sksBufferBytes sm86ProgramEmpty))
def sksBodyB = (lambda unrestricted seq : Nat . (sksTileBody seq sksBufferBytes 0 sm86ProgramEmpty))
def sksBodySkippedA = (lambda unrestricted seq : Nat . (sksTileBodyMasked sksMaskSkipped seq 0 sksBufferBytes sm86ProgramEmpty))
def sksBodySkippedB = (lambda unrestricted seq : Nat . (sksTileBodyMasked sksMaskSkipped seq sksBufferBytes 0 sm86ProgramEmpty))
def streamingAttentionKeySharedOrderedUncheckedSM86 =
(lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat . (lambda unrestricted fromConstant : Nat . (lambda unrestricted fromTileBase : Nat .
(sksPrologue order seq heads keyValueHeads fromConstant fromTileBase
(sm86ProgramAppend (sksBodyA seq)
(sm86ProgramAppend (sfFirstTail (sm86ProgramCount (sksBodyB seq)))
(sm86ProgramAppend (sksBodyB seq)
(sm86ProgramAppend (sfSecondTail (sm86ProgramCount (sksBodyA seq)) (sm86ProgramCount (sksBodyB seq)))
(sksEpilogue seq))))))))))))
-- The same storing dK and dV into a fused projection's gradient
-- (skProjectionEpilogue) after the +0 block below the diagonal.
def streamingAttentionKeySharedProjectionSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat . (lambda unrestricted width : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionKeyProjectionAdmitted seq heads keyValueHeads width) 1) .
(sksPrologue 1 seq heads keyValueHeads 0 1
(sm86ProgramAppend (sksBodyA seq)
(sm86ProgramAppend (sfFirstTail (sm86ProgramCount (sksBodyB seq)))
(sm86ProgramAppend (sksBodyB seq)
(sm86ProgramAppend (sfSecondTail (sm86ProgramCount (sksBodyA seq)) (sm86ProgramCount (sksBodyB seq)))
(sksZeroBelowDiagonal seq
(skProjectionEpilogue width seq
(saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty)))))))))))))
def streamingAttentionKeySharedPathSM86 =
(lambda unrestricted body : (family SM86Program) . (lambda unrestricted seq : Nat .
(lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(sksPrologue 1 seq heads keyValueHeads 0 1 (sm86ProgramAppend body (sksEpilogue seq)))))))
-- order 1, heads from the grid, key tiles from the launch's constant
def streamingAttentionKeySharedHeavyFirstSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionKeySharedOrderedUncheckedSM86 1 seq heads keyValueHeads 0 1)))))
def streamingAttentionKeyGroupedWithHeadSourceUncheckedSM86 = (streamingAttentionKeyGroupedOrderedUncheckedSM86 0)
def streamingAttentionKeyGroupedUncheckedSM86 = (lambda unrestricted seq : Nat .
(lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(streamingAttentionKeyGroupedWithHeadSourceUncheckedSM86 seq heads keyValueHeads 0 0))))
def streamingAttentionKeyGroupedSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionKeyGroupedUncheckedSM86 seq heads keyValueHeads)))))
def streamingAttentionKeyGroupedByConstantSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionKeyGroupedWithHeadSourceUncheckedSM86 seq heads keyValueHeads 1 0)))))
def streamingAttentionKeyGroupedByConstantAndTileSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionKeyGroupedWithHeadSourceUncheckedSM86 seq heads keyValueHeads 1 1)))))
-- Independent grid axes select the K/V head and within-group query head;
-- the constant selects only the contiguous key-tile range. This composition
-- keeps the mathematical body identical to the serial per-head schedule.
-- The raw body overwrites a global-load address before its late read has
-- retired (checked scoreboard ordinal 38 at 4096/9/3). Higher concurrency
-- cannot rely on the serial launch's observed latency hiding. Reuse the
-- target-owned read-barrier guard; full-kernel ordering is checked beside
-- the indexing KATs. Physical correctness and its performance cost remain
-- separate gates; this is not an established cause of the prior failures.
-- Wide grids need separate physical qualification: the earlier unsplit grid
-- did not retire on the 3090, and no occupancy/speed claim follows here.
def streamingAttentionKeyGroupedByGridAndTileSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(sm86GuardLateReads
(streamingAttentionKeyGroupedWithHeadSourceUncheckedSM86 seq heads keyValueHeads 0 1))))))
-- Order-1 forms (heads fastest; forward and query tiles counted down). The
-- launch grid is (K/V heads, groups, tiles); see saAxis.
def streamingAttentionForwardHeavyFirstSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionForwardGroupedOrderedUncheckedSM86 1 seq heads keyValueHeads)))))
def streamingAttentionQueryHeavyFirstSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionQueryGroupedOrderedUncheckedSM86 1 seq heads keyValueHeads)))))
-- Key tiles already run longest first (tile j reads the queries from j on);
-- order 1 only puts the heads fastest. The constant still selects the range.
def streamingAttentionKeyHeadsFirstByGridAndTileSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(sm86GuardLateReads
(streamingAttentionKeyGroupedOrderedUncheckedSM86 1 seq heads keyValueHeads 0 1))))))
-- The same with fixed-latency products (SB2 and SB3 freed) before the
-- late-read guard serializes and compacts it.
def streamingAttentionKeyHeadsFirstFixedProductsSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(sm86GuardLateReads
(sm86TensorProductsFixed 12
(streamingAttentionKeyGroupedOrderedUncheckedSM86 1 seq heads keyValueHeads 0 1)))))))
-- The same storing dK into a fused projection's gradient (skProjectionEpilogue).
def streamingAttentionKeyHeadsFirstProjectionSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat . (lambda unrestricted width : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionKeyProjectionAdmitted seq heads keyValueHeads width) 1) .
(sm86GuardLateReads
(sm86TensorProductsFixed 12
(streamingAttentionKeyGroupedOrderedWithSM86 (skProjectionEpilogue width) 1 seq heads keyValueHeads 0 1))))))))
def streamingAttentionKeySM86 = (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(streamingAttentionKeyGroupedUncheckedSM86 seq heads heads)))
def streamingAttentionKeySM86Registers : Nat = skRegisters
def streamingAttentionKeySM86Threads : Nat = skThreads
def streamingAttentionKeySM86ConstantBytes : Nat = (saArgument 12)
def streamingAttentionKeyHeadConstantBytes : Nat = (saArgument 13)
def streamingAttentionKeyTileBaseConstantBytes : Nat = (saArgument 14)
-- the shared pass's parameters with dS^T's plane (argument 14)
def streamingAttentionKeyScoreGradientConstantBytes : Nat = (saArgument 15)
-- Every key tile in one launch. On 2026-09-28 the 3090 full-grid launch did
-- not retire and half the key CTAs a launch did; with the QMD register
-- allocation and late-read fixes since, the full grid retires bit for bit
-- (2026-10-03) and saves a launch and a wave tail a layer.
def streamingAttentionKeyContiguousBlocksPerLaunch : Nat = 128
def streamingAttentionKeyTileBaseArgument : Nat = 13
def streamingAttentionKeyGroupArgument : Nat = 12
def streamingAttentionKeyQueryArgument : Nat = 0
def streamingAttentionKeyKeyArgument : Nat = 1
-- V for dP^T: [KV-head][position][64], unlike forward/query
-- attention's transposed [KV-head][64][position] value operand.
def streamingAttentionKeyRowMajorValueArgument : Nat = 2
def streamingAttentionKeyOutputGradientArgument : Nat = 3
def streamingAttentionKeyTransposedQueryArgument : Nat = 4
def streamingAttentionKeyTransposedOutputGradientArgument : Nat = 5
def streamingAttentionKeyLogSumArgument : Nat = 6
def streamingAttentionKeyRowTermArgument : Nat = 7
def streamingAttentionKeyKeyGradientArgument : Nat = 8
def streamingAttentionKeyValueGradientArgument : Nat = 9
def streamingAttentionKeyBaseTwoScaleArgument : Nat = 10
def streamingAttentionKeyNaturalScaleArgument : Nat = 11
-- Sum the query-head dK or dV planes of one K/V group into its compact
-- gradient plane. Both backward outputs have 64*T binary32 elements per
-- head, though dV is transposed within that plane. A CTA owns 128 elements
-- of one K/V head; no atomics or cross-CTA ordering are needed. The caller
-- supplies disjoint source/destination regions and launches grid
-- (64*T/128, keyValueHeads, 1) for T divisible by 64.
def streamingAttentionReduceGroupedGradientUncheckedSM86 =
(lambda unrestricted seq : Nat .
(lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
(let unrestricted planeWords = (naturalMultiply saHeadWidth seq) in
(let unrestricted planeBytes = (naturalMultiply planeWords 4) in
(saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
(saS2R 2 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdX)
(saS2R 3 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdY)
(saMovImmAfter 7 0 saWait5
(saImad 4 2 saThreads saTid
(saImad 5 3 planeWords 4
(saImad 6 3 (naturalMultiply groups planeWords) 4
(saMovImm 11 4
(saWide 12 6 11 (saArgument 0)
(saWide 14 5 11 (saArgument 1)
(saMovImm 9 0
(saFor groups (lambda unrestricted group : Nat .
(lambda unrestricted rest : (family SM86Program) .
(saLoad 8 12 (naturalMultiply group planeBytes) saSB0 saWaitNone
(saFadd 9 9 8 saWait0 rest))))
(saOp (saStore 14 9 0 saWaitNone)
(saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty))))))))))))))))))))
def streamingAttentionReduceGroupedGradientSM86 =
(lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
(lambda unrestricted keyValueHeads : Nat .
(lambda erased admitted : (equal Nat (streamingAttentionGroupedSM86Admitted seq heads keyValueHeads) 1) .
(streamingAttentionReduceGroupedGradientUncheckedSM86 seq heads keyValueHeads)))))
def streamingAttentionReduceGroupedGradientSM86Registers : Nat = 24
def streamingAttentionReduceGroupedGradientSM86Threads : Nat = saThreads
def streamingAttentionReduceGroupedGradientSM86ConstantBytes : Nat = (saArgument 2)
def streamingAttentionReduceSourceArgument : Nat = 0
def streamingAttentionReduceDestinationArgument : Nat = 1
def streamingAttentionReduceGroupedGradientSM86Blocks =
(lambda unrestricted seq : Nat .
(naturalDivideUnchecked (naturalMultiply saHeadWidth seq) saThreads))
-- These dimensions are part of the realization's contract, not facts for a
-- pairing to restate. The admission above guards the divisor and geometry.
def streamingAttentionGroupedSM86QueryBlocks =
(lambda unrestricted seq : Nat . (naturalDivideUnchecked seq saTile))
def streamingAttentionGroupedSM86KeyBlocks =
(lambda unrestricted seq : Nat . (naturalDivideUnchecked seq skSubTile))
def streamingAttentionGroupedSM86Groups =
(lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
(naturalDivideUnchecked heads keyValueHeads)))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.