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 701–707

sfMaskBranch

Full file
One tile from `buffer` while the next is copied into `other`: all warps are done with `other` (the barrier), the copies go out while another tile follows (P1), this tile's group retires (and the copies' register reads), every thread's copies are visible (the barrier), then the register kernel's tile body from shared fragments. `rows` 1: V is staged as rows ([K/V head][T][64], K's layout), the next tile 64 rows on, and O += P V's B fragments are its tile read transposed (LDSM.T through sfT1, as the query pass reads K^T): the same values in the same lanes as V^T's tile read plain. The causal mask can change only the diagonal tile, which every pass walks last: while another tile follows (P1, set at the body's top), a forward branch skips the mask's 32 compares and 32 predicated moves (which would leave every score as it is). The schedule checks follow a straight line, so each body is checked twice: through the mask (the diagonal tile) and with the branch taken (`sfMaskSkipped`: the branch, then what follows the mask).
701def sfMaskBranch = (lambda unrestricted bytes : Nat . (lambda unrestricted tail : (family SM86Program) .
702  (saWhen saP1
703    (constructor SM86InstructionBody SM86Branch
704      (saU bytes)
705      (sm86Unsigned32 (byte 0) (byte 0) (byte 128) (byte 3))
706      saPlain)
707    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.