Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

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

3,230 lines493 declarations193.3 KiBSHA-256 cc0366b3d857

def · lines 240–248

saLoadKeys

Full file
---- one key tile ---- the tile's K fragments (all 32 slices x tiles x halves), S zeroed
240def saLoadKeys = (lambda unrestricted tail : (family SM86Program) .
241  (saFor 32 (lambda unrestricted i : Nat .
242    (let unrestricted k = (naturalDivideUnchecked i 8) in (let unrestricted nt = (naturalModuloUnchecked i 8) in
243      (lambda unrestricted rest : (family SM86Program) .
244        (saLoad (saB k nt) saKeyPointer (naturalAdd (naturalMultiply nt 1024) (naturalMultiply k 32)) saSB0 (naturalSelect (naturalIsZero i) saWait3 saWaitNone)
245        (saLoad (succ (saB k nt)) saKeyPointer (naturalAdd (naturalMultiply nt 1024) (naturalAdd (naturalMultiply k 32) 16)) saSB0 saWaitNone
246          rest))))))
247  (saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 36 i) 0))
248    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.