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 1027–1036

sfpMaskTile

Full file
the softmax steps on `scores` (saMaskTile, saScaleRow, saRowMax, saExponentials, saPackP with the score block a parameter); the first scaling waits for the scores' products (SB2)
1027def sfpMaskTile = (lambda unrestricted scores : (pi unrestricted nt : Nat . Nat) . (lambda unrestricted tail : (family SM86Program) .
1028  (saFor 32 (lambda unrestricted i : Nat .
1029    (let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
1030      (let unrestricted b = (naturalModuloUnchecked i 2) in
1031      (let unrestricted threshold = (naturalSaturatingSubtract (naturalAdd (naturalAdd (naturalMultiply 8 nt) b) 63) (naturalMultiply 8 h)) in
1032      (lambda unrestricted rest : (family SM86Program) .
1033        (saGreater saP0 saMask threshold
1034          (saUnless saP0 (constructor SM86InstructionBody SM86MoveImmediate (saR (saElem (scores nt) h b)) (saU saMinusInfinity) saPlain)
1035            rest))))))))
1036    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.