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 1672–1694

sqsTileBodyMasked

Full file
1672def sqsTileBodyMasked = (lambda unrestricted mask : (pi unrestricted tail : (family SM86Program) . (family SM86Program)) .
1673  (lambda unrestricted buffer : Nat . (lambda unrestricted other : Nat .
1674  (lambda unrestricted tail : (family SM86Program) .
1675  (sfBarrier
1676  (saGreater saP1 saCount 1
1677  (saFor 4 (lambda unrestricted i : Nat . (saAddImm (sqsOffset i) (sqsOffset i) (naturalMultiply saTile (naturalMultiply saHeadWidth saHalfBytes))))
1678  (sqsCopyTile other 1
1679  (sfWaitGroups 1 saWait5
1680  (sfBarrier
1681  (saImad saMask saCount 4096 saMaskBase
1682  -- S = Q K^T
1683  (sqsLoadFragments sqsLane 0 buffer saWait3
1684  (sqProductsFresh saS (saQ 0) saWait1
1685  -- dP = dO V^T
1686  (sqsLoadFragments sqsLane 0 (naturalAdd buffer sfValueTile) saWait2
1687  (sqProductsFresh sqDP (sqDOut 0) saWait1
1688  -- K^T's fragments for dQ: K's tile transposed, once dP's products are done
1689  (sqsLoadFragments sqsLaneTransposed 1 buffer saWait2
1690  (mask
1691  (sqScoreGradient
1692  (saPackP
1693  (sfValues
1694    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.