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.
2050def cgStreamingAttentionBackward = (lambda unrestricted l : Nat .
2051 (cgThen
2052 (cgLaunch
2053 cgImageAttentionQuery
2054 (naturalDivideUnchecked cgSeq 64)
2055 cgHeads
2056 1
2057 (cgPointer (cgArgument 0) (cgQueriesOf l)
2058 (cgPointer (cgArgument 1) (cgKeysOf l)
2059 (cgPointer (cgArgument 2) cgValuesHeads
2060 (cgPointer (cgArgument 3) cgKeysT
2061 (cgPointer (cgArgument 4) cgContextGradientHeads
2062 (cgPointer (cgArgument 5) (cgMergedHalfOf l)
2063 (cgPointer (cgArgument 6) (cgLogSumExpOf l)
2064 (cgPointer (cgArgument 7) cgQueriesGradient
2065 (cgPointer (cgArgument 8) cgRowTerms
2066 (cgSlot (cgArgument 9) cgStreamingScale
2067 (cgSlot (cgArgument 10) cgScoreScale cgNoSlots))))))))))))
2068 (cgLaunch
2069 cgImageAttentionKey
2070 (naturalDivideUnchecked cgSeq 32)
2071 cgHeads
2072 1
2073 (cgPointer (cgArgument 0) (cgQueriesOf l)
2074 (cgPointer (cgArgument 1) (cgKeysOf l)
2075 (cgPointer (cgArgument 2) cgValuesHeads
2076 (cgPointer (cgArgument 3) cgContextGradientHeads
2077 (cgPointer (cgArgument 4) cgQueriesT
2078 (cgPointer (cgArgument 5) cgContextGradientT
2079 (cgPointer (cgArgument 6) (cgLogSumExpOf l)
2080 (cgPointer (cgArgument 7) cgRowTerms
2081 (cgPointer (cgArgument 8) cgKeysGradient
2082 (cgPointer (cgArgument 9) cgValuesGradient
2083 (cgSlot (cgArgument 10) cgStreamingScale
2084 (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.