Sum the query-head dK or dV planes of one K/V group into its compact
gradient plane. Both backward outputs have 64*T binary32 elements per
head, though dV is transposed within that plane. A CTA owns 128 elements
of one K/V head; no atomics or cross-CTA ordering are needed. The caller
supplies disjoint source/destination regions and launches grid
(64*T/128, keyValueHeads, 1) for T divisible by 64.
3182def streamingAttentionReduceGroupedGradientUncheckedSM86 =
3183 (lambda unrestricted seq : Nat .
3184 (lambda unrestricted heads : Nat .
3185 (lambda unrestricted keyValueHeads : Nat .
3186 (let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
3187 (let unrestricted planeWords = (naturalMultiply saHeadWidth seq) in
3188 (let unrestricted planeBytes = (naturalMultiply planeWords 4) in
3189 (saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
3190 (saS2R 2 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdX)
3191 (saS2R 3 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdY)
3192 (saMovImmAfter 7 0 saWait5
3193 (saImad 4 2 saThreads saTid
3194 (saImad 5 3 planeWords 4
3195 (saImad 6 3 (naturalMultiply groups planeWords) 4
3196 (saMovImm 11 4
3197 (saWide 12 6 11 (saArgument 0)
3198 (saWide 14 5 11 (saArgument 1)
3199 (saMovImm 9 0
3200 (saFor groups (lambda unrestricted group : Nat .
3201 (lambda unrestricted rest : (family SM86Program) .
3202 (saLoad 8 12 (naturalMultiply group planeBytes) saSB0 saWaitNone
3203 (saFadd 9 9 8 saWait0 rest))))
3204 (saOp (saStore 14 9 0 saWaitNone)
3205 (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.