Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

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

1,203 lines221 declarations70.9 KiBSHA-256 23d4e5e2aa2a

def · lines 995–1018

skEpilogue

Full file
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
995def skEpilogue = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
996  (saMovConst (skTmp 0) (saArgument 11)
997  (saFor 32 (lambda unrestricted i : Nat .
998    (saFmul (naturalAdd 116 i) (naturalAdd 116 i) (skTmp 0) (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
999  (saWide skKeyGradientStorePointer saScratch saOne (saArgument 8)
1000  (saFor 16 (lambda unrestricted i : Nat .
1001    (let unrestricted nt = (naturalDivideUnchecked i 2) in (let unrestricted h = (naturalModuloUnchecked i 2) in
1002      (saOp (sbStore64 skKeyGradientStorePointer (saElem (skDK nt) h 0)
1003        (naturalAdd (naturalMultiply nt 32) (naturalMultiply h (naturalMultiply 8 (naturalMultiply saHeadWidth 4))))
1004        saWaitNone)))))
1005  -- dV^T's offset: head h's 64 x seq plane, the quad lane's first column,
1006  -- the key
1007  (saMovImm (skTmp 1) 0
1008  (saImad (skTmp 1) saQuadLane (naturalMultiply 8 seq) (skTmp 1)
1009  (saImad (skTmp 1) skHead (naturalMultiply saHeadWidth (naturalMultiply seq 4)) (skTmp 1)
1010  (saImad (skTmp 1) skRow 4 (skTmp 1)
1011  (saWide saPointer (skTmp 1) saOne (saArgument 9)
1012  (saFor 32 (lambda unrestricted i : Nat .
1013    (let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
1014      (let unrestricted b = (naturalModuloUnchecked i 2) in
1015      (saOp (saStore saPointer (saElem (skDV nt) h b)
1016        (naturalAdd (naturalMultiply (naturalAdd (naturalMultiply 8 nt) b) (naturalMultiply seq 4)) (naturalMultiply h 32))
1017        saWaitNone))))))
1018    tail))))))))))))

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.