Realization.Nvidia.SM86.StreamingAttentionSM86

1527def streamingAttentionQueryGroupedOrderedUncheckedSM86 = (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
1528  (let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
1529  (let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
1530  (let unrestricted rowBytes = (naturalMultiply heads (naturalMultiply saHeadWidth saHalfBytes)) in
1531  (let unrestricted body = (sqTileBody seq sm86ProgramEmpty) in
1532  (saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
1533  (saS2R saTileIndex (saTileAxis order)
1534  (saS2R saKeyValueHead (saKeyValueAxis order)
1535  (saS2R saHead (saGroupAxis order)
1536  (saMovImm saOne 1
1537  (saMovImmAfter saScratch 31 saWait5
1538  (saImad saHead saKeyValueHead groups saHead
1539  (saCountDownTile order (naturalDivideUnchecked seq saTile)
1540  (saAnd saLane saTid saScratch
1541  (saShr saWarp saTid 5
1542  (saShr saQuad saLane 2
1543  (saMovImm saScratch 3
1544  (saAnd saQuadLane saLane saScratch
1545  (saImad saRow saWarp 16 saQuad
1546  (saImad saRow saTileIndex saTile saRow
1547  -- Q's and dO's A fragments: head h's plane, the row, the quad lane's pair
1548  (saMovImm saScratch 0
1549  (saImad saScratch saRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
1550  (saImad saScratch saQuadLane 4 saScratch
1551  (saImad saScratch saHead headBytes saScratch
1552  (saWide saPointer saScratch saOne (saArgument 0)
1553  (saFor 16 (lambda unrestricted i : Nat .
1554    (let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
1555      (saLoad (saQReg kt j) saPointer
1556        (naturalAdd (naturalMultiply kt 32)
1557          (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
1558            (naturalMultiply (naturalDivideUnchecked j 2) 16)))
1559        saSB1 saWaitNone))))
1560  (saWide sqOutputGradientLoadPointer saScratch saOne (saArgument 4)
1561  (saFor 16 (lambda unrestricted i : Nat .
1562    (let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
1563      (saLoad (naturalAdd (sqDOut kt) j) sqOutputGradientLoadPointer
1564        (naturalAdd (naturalMultiply kt 32)
1565          (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
1566            (naturalMultiply (naturalDivideUnchecked j 2) 16)))
1567        saSB1 saWaitNone))))
1568  -- dQ's offset (binary32 rows), kept in saKeyPointer's partner until the
1569  -- epilogue: saScratch becomes O's (merged plane) for the row terms
1570  (saMovImm saKeyOffset 0
1571  (saImad saKeyOffset saRow (naturalMultiply saHeadWidth 4) saKeyOffset
1572  (saImad saKeyOffset saQuadLane 8 saKeyOffset
1573  (saImad saKeyOffset saHead (naturalMultiply seq (naturalMultiply saHeadWidth 4)) saKeyOffset
1574  (saMovImm saScratch 0
1575  (saImad saScratch saRow rowBytes saScratch
1576  (saImad saScratch saHead (naturalMultiply saHeadWidth saHalfBytes) saScratch
1577  (saImad saScratch saQuadLane 4 saScratch
1578  -- the row terms' offset: 4 (h seq + row)
1579  (saImad saRow saHead seq saRow
1580  (saMovImm saValueOffset 0
1581  (saImad saRow saRow 4 saValueOffset
1582  (sqRowTerms seq heads
1583  -- dQ's offset into saScratch; K's and V's from key 0, K^T's
1584  (saAddImm saScratch saKeyOffset 0
1585  (saMovImm saKeyOffset 0
1586  (saImad saKeyOffset saQuad (naturalMultiply saHeadWidth saHalfBytes) saKeyOffset
1587  (saImad saKeyOffset saQuadLane 4 saKeyOffset
1588  (saImad saKeyOffset saKeyValueHead headBytes saKeyOffset
1589  (saMovImm sqKTOffset 0
1590  (saImad sqKTOffset saQuad (naturalMultiply seq saHalfBytes) sqKTOffset
1591  (saImad sqKTOffset saQuadLane 4 sqKTOffset
1592  (saImad sqKTOffset saKeyValueHead headBytes sqKTOffset
1593  (saAddImm saCount saTileIndex 1
1594  (saImad saMaskBase saWarp 16 saQuad
1595  (saImad saMaskBase saQuadLane 4294967294 saMaskBase
1596  (saAddImm saMaskBase saMaskBase (naturalSaturatingSubtract 4294967296 (naturalSaturatingSubtract 4096 64))
1597  (saMovConst saScale (saArgument 9)
1598  (saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
1599  (sm86ProgramAppend body
1600    (sm86ProgramAppend (sqLoopTail (sm86ProgramCount body))
1601      (sqEpilogue
1602        (saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))