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.
1023def streamingAttentionKeyGroupedWithHeadSourceUncheckedSM86 = (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat . (lambda unrestricted fromConstant : Nat . (lambda unrestricted fromTileBase : Nat .
1024 (let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
1025 (let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
1026 (let unrestricted subTiles = (naturalDivideUnchecked seq skSubTile) in
1027 (let unrestricted body = (skTileBody seq sm86ProgramEmpty) in
1028 (saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
1029 (saS2R saTileIndex (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdX)
1030 (skLoadTileBase fromTileBase
1031 (skLoadHeadIndices fromConstant
1032 (saMovImm saOne 1
1033 (saMovImmAfter saScratch 31 saWait5
1034 (saImad skHead skKeyValueHead groups skHead
1035 (saAnd saLane saTid saScratch
1036 (saShr saWarp saTid 5
1037 (saShr saQuad saLane 2
1038 (saMovImm saScratch 3
1039 (saAnd saQuadLane saLane saScratch
1040 -- the thread's key row: 32 tile + 16 warp + lane / 4
1041 (saImad skRow saWarp 16 saQuad
1042 (saImad skRow saTileIndex skSubTile skRow
1043 -- K's and V's A fragments
1044 (saMovImm saScratch 0
1045 (saImad saScratch skRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
1046 (saImad saScratch saQuadLane 4 saScratch
1047 (saImad saScratch skKeyValueHead headBytes saScratch
1048 (saWide saPointer saScratch saOne (saArgument 1)
1049 (saFor 16 (lambda unrestricted i : Nat .
1050 (let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
1051 (saLoad (naturalAdd (skK kt) j) saPointer
1052 (naturalAdd (naturalMultiply kt 32)
1053 (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
1054 (naturalMultiply (naturalDivideUnchecked j 2) 16)))
1055 saSB1 saWaitNone))))
1056 (saWide saPointer saScratch saOne (saArgument 2)
1057 (saFor 16 (lambda unrestricted i : Nat .
1058 (let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
1059 (saLoad (naturalAdd (skV kt) j) saPointer
1060 (naturalAdd (naturalMultiply kt 32)
1061 (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
1062 (naturalMultiply (naturalDivideUnchecked j 2) 16)))
1063 saSB1 saWaitNone))))
1064 -- dK's offset (binary32 rows) in saScratch
1065 (saMovImm saScratch 0
1066 (saImad saScratch skRow (naturalMultiply saHeadWidth 4) saScratch
1067 (saImad saScratch saQuadLane 8 saScratch
1068 (saImad saScratch skHead (naturalMultiply seq (naturalMultiply saHeadWidth 4)) saScratch
1069 -- the queries start at the last sub-tile: Q's and dO's rows (row lane / 4,
1070 -- the quad lane's pair), Q^T's and dO^T's columns (row lane / 4), L's and D's
1071 (saMovImm saKeyOffset 0
1072 (saImad saKeyOffset saQuad (naturalMultiply saHeadWidth saHalfBytes) saKeyOffset
1073 (saImad saKeyOffset saQuadLane 4 saKeyOffset
1074 (saImad saKeyOffset skHead headBytes saKeyOffset
1075 (saAddImm saKeyOffset saKeyOffset (naturalMultiply (naturalSaturatingSubtract subTiles 1) (naturalMultiply skSubTile (naturalMultiply saHeadWidth saHalfBytes)))
1076 (saMovImm saValueOffset 0
1077 (saImad saValueOffset saQuad (naturalMultiply seq saHalfBytes) saValueOffset
1078 (saImad saValueOffset saQuadLane 4 saValueOffset
1079 (saImad saValueOffset skHead headBytes saValueOffset
1080 (saAddImm saValueOffset saValueOffset (naturalMultiply (naturalSaturatingSubtract subTiles 1) (naturalMultiply skSubTile saHalfBytes))
1081 (saMovImm skLogOffset 0
1082 (saImad skLogOffset saQuadLane 8 skLogOffset
1083 (saImad skLogOffset skHead (naturalMultiply seq 4) skLogOffset
1084 (saAddImm skLogOffset skLogOffset (naturalMultiply (naturalSaturatingSubtract subTiles 1) (naturalMultiply skSubTile 4))
1085 -- the sub-tiles from the last down to the diagonal; the mask's base,
1086 -- 2 (lane % 4) - 16 warp - lane / 4 + 128 - 4096
1087 (saMovImm saCount subTiles
1088 (saImad saCount saTileIndex 4294967295 saCount
1089 (saMovImm saMaskBase 0
1090 (saImad saMaskBase saQuadLane 2 saMaskBase
1091 (saImad saMaskBase saWarp 4294967280 saMaskBase
1092 (saImad saMaskBase saQuad 4294967295 saMaskBase
1093 (saAddImm saMaskBase saMaskBase (naturalSaturatingSubtract 4294967296 (naturalSaturatingSubtract 4096 128))
1094 (saMovConst saScale (saArgument 10)
1095 (saFor 64 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 84 i) 0))
1096 (sm86ProgramAppend body
1097 (sm86ProgramAppend (skLoopTail (sm86ProgramCount body))
1098 (skEpilogue seq
1099 (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.