===========================================================================
The query pass split in two around the key pass, which stores the score
gradients it forms anyway (sksStoreScoreGradients):
the row terms (streamingAttentionRowTermsSM86), before the key pass: the
query pass's prologue without Q, forming D for its 64 rows by the
same instructions (sqRowTermsWith) from the same fragments;
dQ (streamingAttentionQueryFromScoresSM86), after it: for each key tile
up to the diagonal, in the query pass's order, dQ += dS K from dS^T's
tile read transposed (LDSM.T) into the A registers saPackP fills and K's
read as the query pass reads it -- the same HMMAs on the same words --
then the query pass's epilogue.
S, P and dP are formed once (in the key pass) instead of twice.
the row terms: grid and parameters as the query pass's (dO 4, O 5, L 6,
D 8), one block of four warps per 64 rows of a query head
1926def streamingAttentionRowTermsOrderedUncheckedSM86 = (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat .
1927 (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
1928 (let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
1929 (let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
1930 (let unrestricted rowBytes = (naturalMultiply heads (naturalMultiply saHeadWidth saHalfBytes)) in
1931 (saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
1932 (saS2R saTileIndex (saTileAxis order)
1933 (saS2R saKeyValueHead (saKeyValueAxis order)
1934 (saS2R saHead (saGroupAxis order)
1935 (saMovImm saOne 1
1936 (saMovImmAfter saScratch 31 saWait5
1937 (saImad saHead saKeyValueHead groups saHead
1938 (saCountDownTile order (naturalDivideUnchecked seq saTile)
1939 (saAnd saLane saTid saScratch
1940 (saShr saWarp saTid 5
1941 (saShr saQuad saLane 2
1942 (saMovImm saScratch 3
1943 (saAnd saQuadLane saLane saScratch
1944 (saImad saRow saWarp 16 saQuad
1945 (saImad saRow saTileIndex saTile saRow
1946 -- dO's A fragments: head h's plane, the row, the quad lane's pair
1947 (saMovImm saScratch 0
1948 (saImad saScratch saRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
1949 (saImad saScratch saQuadLane 4 saScratch
1950 (saImad saScratch saHead headBytes saScratch
1951 (saWide sqOutputGradientLoadPointer saScratch saOne (saArgument 4)
1952 (saFor 16 (lambda unrestricted i : Nat .
1953 (let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
1954 (saLoad (naturalAdd (sqDOut kt) j) sqOutputGradientLoadPointer
1955 (naturalAdd (naturalMultiply kt 32)
1956 (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
1957 (naturalMultiply (naturalDivideUnchecked j 2) 16)))
1958 saSB1 saWaitNone))))
1959 -- O's offset (merged plane) and the row terms', as the query pass's
1960 (saMovImm saScratch 0
1961 (saImad saScratch saRow rowBytes saScratch
1962 (saImad saScratch saHead (naturalMultiply saHeadWidth saHalfBytes) saScratch
1963 (saImad saScratch saQuadLane 4 saScratch
1964 (saImad saRow saHead seq saRow
1965 (saMovImm saValueOffset 0
1966 (saImad saRow saRow 4 saValueOffset
1967 -- D through its own pair: the dO loads' address pair is never re-formed
1968 (sqRowTermsWith seq heads saKeyPointer
1969 (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.