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 2335–2346

skLoadTileBase

Full file
A key launch may own a contiguous subset of key tiles. The caller derives the subset's base from its launch schedule; the ordinary unsplit kernel keeps using CTA x directly. This keeps all query-loop and output-address semantics in one implementation while bounding the work of one launch.
2335def skLoadTileBase = (lambda unrestricted fromConstant : Nat .
2336  (lambda unrestricted tail : (family SM86Program) .
2337    (nat-eliminate (lambda unrestricted mode : Nat . (family SM86Program))
2338      tail
2339      (lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86Program) .
2340        -- Tile rebasing consumes CTA x before the later head-selection
2341        -- wait. S2R is variable latency: the immediate-result machine model
2342        -- cannot prove this ordering. Drain SB5 here, before its first use;
2343        -- Proof.CheckedAttentionGridSelection checks the control dependency.
2344        (saMovConstAfter saScratch (saArgument 13) saWait5
2345          (saImad saTileIndex saScratch 1 saTileIndex tail))))
2346      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.