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 946–946

sfpScoresB

Full file
V as rows ([K/V head][T][64] half, K's layout, the backward passes' V): the forward needs no transposed copy of V =========================================================================== The shared forward with V as rows, software-pipelined: while a warp works through tile j's softmax it issues tile j+1's score products from K's next tile, so the tensor cores need not idle through the softmax. The scores alternate between two accumulator blocks (S at R36, S' at R100), and K's and V's B fragments are loaded a 16-key slice at a time into two slice buffers (R132, R148) instead of a tile at once (R100..R163). K and V keep separate double buffers in the same shared layout (K at 0 and 18432, V 9216 on): tile j's body copies K's tile j + 2 into the K buffer tile j's scores have finished with and V's tile j + 1 into the V buffer tile j - 1's products have finished with, so every copy has a whole tile to land. Every product, softmax step and rounding is the shared forward's, in the same order; only the issue order interleaves two independent streams.
946def sfpScoresB = (lambda unrestricted nt : Nat . (naturalAdd 100 (naturalMultiply nt 4)))

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.