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 2320–2328

skLoadHeadIndicesOrdered

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.
2320def skLoadHeadIndicesOrdered = (lambda unrestricted order : Nat . (lambda unrestricted fromConstant : Nat .
2321  (lambda unrestricted tail : (family SM86Program) .
2322    (nat-eliminate (lambda unrestricted mode : Nat . (family SM86Program))
2323      (saS2R skKeyValueHead (saKeyValueAxis order)
2324        (saS2R skHead (saGroupAxis order) tail))
2325      (lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86Program) .
2326        (saMovConst skKeyValueHead (naturalAdd (saArgument 12) 4)
2327          (saMovConst skHead (saArgument 12) tail))))
2328      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.