layer l's heads' attention backward as two streaming launches
(Realization.Nvidia.SM86.StreamingAttentionSM86): the query launch forms
each row's terms D from dO and the forward's output O (the layer's merged
half) and dQ; the key launch dK and dV^T. P is recomputed from the
layer's log-sum-exp.
2051def cgStreamingAttentionBackward = (lambda unrestricted l : Nat .
2052 (cgThen
2053 (cgLaunch
2054 cgImageAttentionQuery
2055 (naturalDivideUnchecked cgSeq 64)
2056 cgHeads
2057 1
2058 (cgPointer (cgArgument 0) (cgQueriesOf l)
2059 (cgPointer (cgArgument 1) (cgKeysOf l)
2060 (cgPointer (cgArgument 2) cgValuesHeads
2061 (cgPointer (cgArgument 3) cgKeysT
2062 (cgPointer (cgArgument 4) cgContextGradientHeads
2063 (cgPointer (cgArgument 5) (cgMergedHalfOf l)
2064 (cgPointer (cgArgument 6) (cgLogSumExpOf l)
2065 (cgPointer (cgArgument 7) cgQueriesGradient
2066 (cgPointer (cgArgument 8) cgRowTerms
2067 (cgSlot (cgArgument 9) cgStreamingScale
2068 (cgSlot (cgArgument 10) cgScoreScale cgNoSlots))))))))))))
2069 (cgLaunch
2070 cgImageAttentionKey
2071 (naturalDivideUnchecked cgSeq 32)
2072 cgHeads
2073 1
2074 (cgPointer (cgArgument 0) (cgQueriesOf l)
2075 (cgPointer (cgArgument 1) (cgKeysOf l)
2076 (cgPointer (cgArgument 2) cgValuesHeads
2077 (cgPointer (cgArgument 3) cgContextGradientHeads
2078 (cgPointer (cgArgument 4) cgQueriesT
2079 (cgPointer (cgArgument 5) cgContextGradientT
2080 (cgPointer (cgArgument 6) (cgLogSumExpOf l)
2081 (cgPointer (cgArgument 7) cgRowTerms
2082 (cgPointer (cgArgument 8) cgKeysGradient
2083 (cgPointer (cgArgument 9) cgValuesGradient
2084 (cgSlot (cgArgument 10) cgStreamingScale
2085 (cgSlot (cgArgument 11) cgScoreScale cgNoSlots)))))))))))))))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.