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 2388–2397

skMaskTile

Full file
the diagonal sub-tile's mask: element (key r + 8h, query 8 nt + 2 (lane % 4) + b) is kept when the query is not before the key: saMask holds 2 (lane % 4) - 16 warp - lane / 4 + 128, plus 4096 for every sub-tile before the diagonal one
2388def skMaskTile = (lambda unrestricted tail : (family SM86Program) .
2389  (saFor 16 (lambda unrestricted i : Nat .
2390    (let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
2391      (let unrestricted b = (naturalModuloUnchecked i 2) in
2392      (let unrestricted threshold = (naturalSaturatingSubtract (naturalAdd 127 (naturalMultiply 8 h)) (naturalAdd (naturalMultiply 8 nt) b)) in
2393      (lambda unrestricted rest : (family SM86Program) .
2394        (saGreater saP0 saMask threshold
2395          (saUnless saP0 (constructor SM86InstructionBody SM86MoveImmediate (saR (saElem (skS nt) h b)) (saU saMinusInfinity) saPlain)
2396            rest))))))))
2397    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.