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 225–233

saLoadKeys

Full file
---- one key tile ---- the tile's K fragments (all 32 slices x tiles x halves), S zeroed
225def saLoadKeys = (lambda unrestricted tail : (family SM86Program) .
226  (saFor 32 (lambda unrestricted i : Nat .
227    (let unrestricted k = (naturalDivideUnchecked i 8) in (let unrestricted nt = (naturalModuloUnchecked i 8) in
228      (lambda unrestricted rest : (family SM86Program) .
229        (saLoad (saB k nt) saKeyPointer (naturalAdd (naturalMultiply nt 1024) (naturalMultiply k 32)) saSB0 (naturalSelect (naturalIsZero i) saWait3 saWaitNone)
230        (saLoad (succ (saB k nt)) saKeyPointer (naturalAdd (naturalMultiply nt 1024) (naturalAdd (naturalMultiply k 32) 16)) saSB0 saWaitNone
231          rest))))))
232  (saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 36 i) 0))
233    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.