Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

packages/realizations/cooperative/nvidia-sm86/src/Realization/Nvidia/SM86/StreamingAttentionSM86.alpha

3,230 lines493 declarations193.3 KiBSHA-256 cc0366b3d857

def · lines 2096–2154

sqdSetupAhead

Full file
the LDSM lanes, the copy role (shared r 128 + 16 (c ^ r % 8)), K's offsets (K/V head, key r + 16 m, chunk c) and dS^T's (head, the block's queries 64 i + 8 c, key r + 16 m), then tiles 0 and 1 (P1 when a second follows)
2096def sqdSetupAhead = (lambda unrestricted ahead : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted headBytes : Nat .
2097  (lambda unrestricted tail : (family SM86Program) .
2098  (saShr sqdT0 saLane 4
2099  (saMovImm sqdT1 7
2100  (saAnd sqdT1 saLane sqdT1
2101  (saShr sqdT2 saLane 3
2102  (saAnd sqdT2 sqdT2 saOne
2103  -- K^T: key 8 ((L / 8) % 2) + L % 8, chunk L / 16 + 2 p, swizzled by L % 8
2104  (saImad (sqdKeyOffset 0) sqdT2 8 sqdT1
2105  (saMovImm (sqdKeyOffset 1) 0
2106  (saImad (sqdKeyOffset 1) (sqdKeyOffset 0) sqdRowBytes (sqdKeyOffset 1)
2107  (saFor 4 (lambda unrestricted p : Nat . (lambda unrestricted rest : (family SM86Program) .
2108    (saAddImm (sqdKeyOffset 2) sqdT0 (naturalMultiply 2 p)
2109    (sqdXor (sqdKeyOffset 2) (sqdKeyOffset 2) sqdT1
2110    (saImad (sqdLaneKeys p) (sqdKeyOffset 2) 16 (sqdKeyOffset 1)
2111      rest)))))
2112  -- dS^T: key 8 (L / 16) + L % 8, chunk 2 warp + (L / 8) % 2, swizzled by L % 8
2113  (saImad (sqdKeyOffset 0) sqdT0 8 sqdT1
2114  (saImad (sqdKeyOffset 2) saWarp 2 sqdT2
2115  (sqdXor (sqdKeyOffset 2) (sqdKeyOffset 2) sqdT1
2116  (saMovImm sqdLaneScores 0
2117  (saImad sqdLaneScores (sqdKeyOffset 0) sqdRowBytes sqdLaneScores
2118  (saImad sqdLaneScores (sqdKeyOffset 2) 16 sqdLaneScores
2119  -- the copy role: row r = tid / 8, chunk c = tid % 8
2120  (saShr sqdT0 saTid 3
2121  (saMovImm sqdT1 7
2122  (saAnd sqdT1 saTid sqdT1
2123  (saMovImm sqdT2 7
2124  (saAnd sqdT2 sqdT0 sqdT2
2125  (sqdXor sqdT2 sqdT2 sqdT1
2126  (saMovImm sqdStore 0
2127  (saImad sqdStore sqdT2 16 sqdStore
2128  (saImad sqdStore sqdT0 sqdRowBytes sqdStore
2129  (saMovImm (sqdKeyOffset 0) 0
2130  (saImad (sqdKeyOffset 0) sqdT1 16 (sqdKeyOffset 0)
2131  (saImad (sqdKeyOffset 0) sqdT0 (naturalMultiply saHeadWidth saHalfBytes) (sqdKeyOffset 0)
2132  (saImad (sqdKeyOffset 0) saKeyValueHead headBytes (sqdKeyOffset 0)
2133  -- dS^T's ([head][sub-tile][key][32 queries]): chunk c is sub-tile
2134  -- 2 i + c / 4, queries 8 (c % 4).., key r
2135  (saShr (sqdScoreOffset 1) sqdT1 2
2136  (saMovImm (sqdScoreOffset 2) 3
2137  (saAnd (sqdScoreOffset 2) sqdT1 (sqdScoreOffset 2)
2138  (saMovImm (sqdScoreOffset 0) 0
2139  (saImad (sqdScoreOffset 0) (sqdScoreOffset 2) 16 (sqdScoreOffset 0)
2140  (saImad (sqdScoreOffset 0) sqdT0 sksScoreGradientRowBytes (sqdScoreOffset 0)
2141  (saImad (sqdScoreOffset 0) (sqdScoreOffset 1) (sksScoreGradientSubTileBytes seq) (sqdScoreOffset 0)
2142  (saImad (sqdScoreOffset 0) saTileIndex (naturalMultiply 2 (sksScoreGradientSubTileBytes seq)) (sqdScoreOffset 0)
2143  (saImad (sqdScoreOffset 0) saHead (naturalMultiply seq (naturalMultiply seq saHalfBytes)) (sqdScoreOffset 0)
2144  (saFor 3 (lambda unrestricted i : Nat .
2145    (saAddImm (sqdKeyOffset (succ i)) (sqdKeyOffset 0) (naturalMultiply (succ i) (naturalMultiply 16 (naturalMultiply saHeadWidth saHalfBytes)))))
2146  (saFor 3 (lambda unrestricted i : Nat .
2147    (saAddImm (sqdScoreOffset (succ i)) (sqdScoreOffset 0) (naturalMultiply (succ i) (naturalMultiply 16 sksScoreGradientRowBytes))))
2148  (sqdCopyTile seq (sqdStage 0) 0
2149  (nat-eliminate (lambda unrestricted n : Nat . (family SM86Program))
2150    tail
2151    (lambda unrestricted p : Nat . (lambda unrestricted ignored : (family SM86Program) .
2152      (saGreater saP1 saCount 1
2153      (sqdCopyTile seq (sqdStage 1) 1 tail))))
2154    (naturalSaturatingSubtract ahead 1))))))))))))))))))))))))))))))))))))))))))))))

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.