Each query head writes a separate dK/dV plane. A later reduction sums
those planes into the compact K/V gradient; concurrent CTAs never race on
one compact gradient address.
2618def streamingAttentionKeyGroupedOrderedWithSM86 = (lambda unrestricted epilogue : (pi unrestricted seq : Nat . (pi unrestricted tail : (family SM86Program) . (family SM86Program))) . (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat . (lambda unrestricted fromConstant : Nat . (lambda unrestricted fromTileBase : Nat .
2619 (let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
2620 (let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
2621 (let unrestricted subTiles = (naturalDivideUnchecked seq skSubTile) in
2622 (let unrestricted body = (skTileBody seq sm86ProgramEmpty) in
2623 (saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
2624 (saS2R saTileIndex (saTileAxis order)
2625 (skLoadTileBase fromTileBase
2626 (skLoadHeadIndicesOrdered order fromConstant
2627 (saMovImm saOne 1
2628 (saMovImmAfter saScratch 31 saWait5
2629 (saImad skHead skKeyValueHead groups skHead
2630 (saAnd saLane saTid saScratch
2631 (saShr saWarp saTid 5
2632 (saShr saQuad saLane 2
2633 (saMovImm saScratch 3
2634 (saAnd saQuadLane saLane saScratch
2635 -- the thread's key row: 32 tile + 16 warp + lane / 4
2636 (saImad skRow saWarp 16 saQuad
2637 (saImad skRow saTileIndex skSubTile skRow
2638 -- K's and V's A fragments
2639 (saMovImm saScratch 0
2640 (saImad saScratch skRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
2641 (saImad saScratch saQuadLane 4 saScratch
2642 (saImad saScratch skKeyValueHead headBytes saScratch
2643 (saWide saPointer saScratch saOne (saArgument 1)
2644 (saFor 16 (lambda unrestricted i : Nat .
2645 (let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
2646 (saLoad (naturalAdd (skK kt) j) saPointer
2647 (naturalAdd (naturalMultiply kt 32)
2648 (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
2649 (naturalMultiply (naturalDivideUnchecked j 2) 16)))
2650 saSB1 saWaitNone))))
2651 (saWide saPointer saScratch saOne (saArgument 2)
2652 (saFor 16 (lambda unrestricted i : Nat .
2653 (let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
2654 (saLoad (naturalAdd (skV kt) j) saPointer
2655 (naturalAdd (naturalMultiply kt 32)
2656 (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
2657 (naturalMultiply (naturalDivideUnchecked j 2) 16)))
2658 saSB1 saWaitNone))))
2659 -- dK's offset (binary32 rows) in saScratch
2660 (saMovImm saScratch 0
2661 (saImad saScratch skRow (naturalMultiply saHeadWidth 4) saScratch
2662 (saImad saScratch saQuadLane 8 saScratch
2663 (saImad saScratch skHead (naturalMultiply seq (naturalMultiply saHeadWidth 4)) saScratch
2664 -- the queries start at the last sub-tile: Q's and dO's rows (row lane / 4,
2665 -- the quad lane's pair), Q^T's and dO^T's columns (row lane / 4), L's and D's
2666 (saMovImm saKeyOffset 0
2667 (saImad saKeyOffset saQuad (naturalMultiply saHeadWidth saHalfBytes) saKeyOffset
2668 (saImad saKeyOffset saQuadLane 4 saKeyOffset
2669 (saImad saKeyOffset skHead headBytes saKeyOffset
2670 (saAddImm saKeyOffset saKeyOffset (naturalMultiply (naturalSaturatingSubtract subTiles 1) (naturalMultiply skSubTile (naturalMultiply saHeadWidth saHalfBytes)))
2671 (saMovImm saValueOffset 0
2672 (saImad saValueOffset saQuad (naturalMultiply seq saHalfBytes) saValueOffset
2673 (saImad saValueOffset saQuadLane 4 saValueOffset
2674 (saImad saValueOffset skHead headBytes saValueOffset
2675 (saAddImm saValueOffset saValueOffset (naturalMultiply (naturalSaturatingSubtract subTiles 1) (naturalMultiply skSubTile saHalfBytes))
2676 (saMovImm skLogOffset 0
2677 (saImad skLogOffset saQuadLane 8 skLogOffset
2678 (saImad skLogOffset skHead (naturalMultiply seq 4) skLogOffset
2679 (saAddImm skLogOffset skLogOffset (naturalMultiply (naturalSaturatingSubtract subTiles 1) (naturalMultiply skSubTile 4))
2680 -- the sub-tiles from the last down to the diagonal; the mask's base,
2681 -- 2 (lane % 4) - 16 warp - lane / 4 + 128 - 4096
2682 (saMovImm saCount subTiles
2683 (saImad saCount saTileIndex 4294967295 saCount
2684 (saMovImm saMaskBase 0
2685 (saImad saMaskBase saQuadLane 2 saMaskBase
2686 (saImad saMaskBase saWarp 4294967280 saMaskBase
2687 (saImad saMaskBase saQuad 4294967295 saMaskBase
2688 (saAddImm saMaskBase saMaskBase (naturalSaturatingSubtract 4294967296 (naturalSaturatingSubtract 4096 128))
2689 (saMovConst saScale (saArgument 10)
2690 (saFor 64 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 84 i) 0))
2691 (sm86ProgramAppend body
2692 (sm86ProgramAppend (skLoopTail (sm86ProgramCount body))
2693 (epilogue seq
2694 (saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))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.