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 2080–2090

sqdTileBody

Full file
One key tile from stage `buffer`: all warps are done with the stage two on (the barrier: it held the previous tile), the copies of the tile two on go out while one follows (P1), this tile's group retires (at most two newer pending), every copy is visible (the barrier), then once the last tile's products have read their operands (SB3) dS's and K^T's fragments and dQ += dS K.
2080def sqdTileBody = (lambda unrestricted seq : Nat . (lambda unrestricted buffer : Nat .
2081  (lambda unrestricted tail : (family SM86Program) .
2082  (sfBarrier
2083  (saGreater saP1 saCount 2
2084  (sqdCopyTile seq (sqdStage (naturalModuloUnchecked (naturalAdd (naturalDivideUnchecked buffer sqdStageBytes) 2) sqdStages)) 1
2085  (sfWaitGroups 2 saWait5
2086  (sfBarrier
2087  (sqdLoadScores (naturalAdd buffer sqdTileBytes) saWait3
2088  (sqdLoadKeys buffer saWaitNone
2089  (sfValues
2090    tail)))))))))))

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.