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 2251–2261

sqd2TileBody

Full file
---- the same in two shared stages ---- One key tile from `buffer` while the next is copied into `other` (the forward pass's double buffer): 32 KB of shared stages, so with the program's own registers (Realization.Nvidia.SM86.RegisterCompactionSM86: 160) three blocks fit an SM. Every product, fragment and order is the three-stage pass's.
2251def sqd2TileBody = (lambda unrestricted seq : Nat . (lambda unrestricted buffer : Nat . (lambda unrestricted other : Nat .
2252  (lambda unrestricted tail : (family SM86Program) .
2253  (sfBarrier
2254  (saGreater saP1 saCount 1
2255  (sqdCopyTile seq other 1
2256  (sfWaitGroups 1 saWait5
2257  (sfBarrier
2258  (sqdLoadScores (naturalAdd buffer sqdTileBytes) saWait3
2259  (sqdLoadKeys buffer saWaitNone
2260  (sfValues
2261    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.