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 1130–1146

sfpTileBodyMasked

Full file
One tile, j, with its scores in `scores` and tile j + 1's to come in `next`; `b` is j % 2. All warps are done with K's buffer b (tile j's products ran in the previous body) and V's buffer 1 - b (tile j - 1's): K's tile j + 2 (P0: two tiles follow) and V's tile j + 1 (P1: one does) are copied there; the previous body's copies retire and are visible to every thread (the barrier); the diagonal mask, then tile j's softmax with tile j + 1's products interleaved (on P1), then O += P V from V's buffer b.
1130def sfpTileBodyMasked = (lambda unrestricted mask : (pi unrestricted scores : (pi unrestricted nt : Nat . Nat) . (pi unrestricted tail : (family SM86Program) . (family SM86Program))) .
1131  (lambda unrestricted scores : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted next : (pi unrestricted nt : Nat . Nat) .
1132  (lambda unrestricted b : Nat . (lambda unrestricted tail : (family SM86Program) .
1133  (let unrestricted other = (naturalSaturatingSubtract 1 b) in
1134  (sfBarrier
1135  (saGreater saP1 saCount 1
1136  (saGreater saP0 saCount 2
1137  (saFor 4 (lambda unrestricted i : Nat . (saAddImm (sfKeyOffset i) (sfKeyOffset i) (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes))))
1138  (saFor 4 (lambda unrestricted i : Nat . (saAddImm (sfValueOffset i) (sfValueOffset i) (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes))))
1139  (sfpCopy (sfpKeyBuffer b) 2 (sfpValueBuffer other) 1
1140  (sfWaitGroups 1 saWait5
1141  (sfBarrier
1142  (saImad saMask saCount 4096 saMaskBase
1143  (mask scores
1144  (sm86ProgramAppend (sfpInterleave sfpEvery (sfpSoftmax scores) (sfpScores 1 next (sfpKeyBuffer other) sm86ProgramEmpty))
1145  (sfpValues (sfpValueBuffer b)
1146    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.