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 2051–2059

sqdLoadScores

Full file
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
2051def sqdLoadScores = (lambda unrestricted base : Nat . (lambda unrestricted wait : Nat .
2052  (lambda unrestricted tail : (family SM86Program) .
2053    (saFor 4 (lambda unrestricted kk : Nat .
2054      (saOp (constructor SM86InstructionBody SM86LoadSharedMatrix (saR (saPA kk)) (saR sqdLaneScores)
2055        (saU (naturalAdd base (naturalMultiply kk (naturalMultiply 16 sqdRowBytes))))
2056        (constructor SM86SharedMatrixCount SM86SharedMatrix4)
2057        (constructor SM86SharedMatrixTranspose SM86SharedMatrixTransposed)
2058        (saCtl saSB1 (naturalSelect (naturalIsZero kk) wait saWaitNone)))))
2059      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.