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 2458–2496

skTileBody

Full file
2458def skTileBody = (lambda unrestricted seq : Nat . (lambda unrestricted tail : (family SM86Program) .
2459  (saImad saMask saCount 4096 saMaskBase
2460  -- the sub-tile's L and D (their negatives)
2461  (saWide saPointer skLogOffset saOne (saArgument 6)
2462  (saFor 8 (lambda unrestricted i : Nat .
2463    (saLoad (skTmp i) saPointer (naturalAdd (naturalMultiply (naturalDivideUnchecked i 2) 32) (naturalMultiply (naturalModuloUnchecked i 2) 4))
2464      saSB1 (naturalSelect (naturalIsZero i) saWait3 saWaitNone)))
2465  (saFor 8 (lambda unrestricted i : Nat .
2466    (sbNegate (skNegL (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2)) (skTmp i) (naturalSelect (naturalIsZero i) saWait1 saWaitNone)))
2467  (saWide saPointer skLogOffset saOne (saArgument 7)
2468  (saFor 8 (lambda unrestricted i : Nat .
2469    (saLoad (skTmp i) saPointer (naturalAdd (naturalMultiply (naturalDivideUnchecked i 2) 32) (naturalMultiply (naturalModuloUnchecked i 2) 4))
2470      saSB1 saWaitNone))
2471  (saFor 8 (lambda unrestricted i : Nat .
2472    (sbNegate (skNegD (naturalDivideUnchecked i 2) (naturalModuloUnchecked i 2)) (skTmp i) (naturalSelect (naturalIsZero i) saWait1 saWaitNone)))
2473  -- S^T = K Q^T
2474  (saWide saKeyPointer saKeyOffset saOne (saArgument 0)
2475  (skLoadRows saKeyPointer saWait3
2476  -- Keep the qualified ordinary schedule. A first-product RZ candidate
2477  -- removed these clears but gave overlapping 100-step timings and higher
2478  -- measured energy; see the 2026-10-01 attention-product-origin fixture.
2479  (saFor 16 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 52 i) 0))
2480  (saFor 16 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
2481  (skProducts 0 skS skK 4 4 saSB2 saWait2 (naturalAdd saWait0 saWait1)
2482  -- dP^T = V dO^T
2483  (saWide saKeyPointer saKeyOffset saOne (saArgument 3)
2484  (skLoadRows saKeyPointer saWait2
2485  (skProducts 0 skDP skV 4 4 saSB2 saWait2 saWait0
2486  -- dO^T's fragments for dV, once dP^T's products are done
2487  (saWide saValuePointer saValueOffset saOne (saArgument 5)
2488  (skLoadColumns seq saValuePointer saWait2
2489  (skMaskTile
2490  (skScoreGradient
2491  (skProducts 0 skDV skPA 2 8 saSB3 saWait3 saWait0
2492  -- Q^T's fragments for dK, once dV's products have read dO^T's
2493  (saWide saValuePointer saValueOffset saOne (saArgument 4)
2494  (skLoadColumns seq saValuePointer saWait3
2495  (skProducts 0 skDK skSA 2 8 saSB3 saWait3 saWait0
2496    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.