Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

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

1,203 lines221 declarations70.9 KiBSHA-256 23d4e5e2aa2a

def · lines 836–844

skLoadHeadIndices

Full file
The wide grouped launch can leave the 3090's channel waiting with no retired fence. A per-head launch uses the same body and output planes but reads its K/V head and within-group query head from one CB0 slot. Keeping this selection in the realization avoids copying the backward kernel into a system; the launch schedule determines the head traversal.
836def skLoadHeadIndices = (lambda unrestricted fromConstant : Nat .
837  (lambda unrestricted tail : (family SM86Program) .
838    (nat-eliminate (lambda unrestricted mode : Nat . (family SM86Program))
839      (saS2R skKeyValueHead (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdY)
840        (saS2R skHead (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdZ) tail))
841      (lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86Program) .
842        (saMovConst skKeyValueHead (naturalAdd (saArgument 12) 4)
843          (saMovConst skHead (saArgument 12) tail))))
844      fromConstant)))

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.