Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

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

1,129 lines205 declarations66.1 KiBSHA-256 aefcbc3b00ca

def · lines 588–607

sqTileBody

Full file
588def sqTileBody = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
589  (saImad saMask saCount 4096 saMaskBase
590  -- S = Q K^T
591  (saWide saKeyPointer saKeyOffset saOne (saArgument 1)
592  (sqLoadRows saKeyPointer saWait3
593  (saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 36 i) 0))
594  (saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 212 i) 0))
595  (sqProducts saS (saQ 0) (naturalAdd saWait0 saWait1)
596  -- dP = dO V^T (V's rows at K's offset)
597  (saWide saKeyPointer saKeyOffset saOne (saArgument 2)
598  (sqLoadRows saKeyPointer saWait2
599  (sqProducts sqDP (sqDOut 0) saWait0
600  -- K^T's fragments for dQ, once dP's products are done
601  (saWide saValuePointer sqKTOffset saOne (saArgument 3)
602  (sqLoadKeysT seq
603  (saMaskTile
604  (sqScoreGradient
605  (saPackP
606  (saValues
607    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.