AlphaGitHub
Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

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

1,129 lines205 declarations66.1 KiBSHA-256 aefcbc3b00ca
Download exact fileView at revisionf5034773ad4bb52133

def · lines 395–410

streamingAttentionGroupedSM86Admitted

Full file
---- the program ---- seq a multiple of 64, heads the number of 64-wide heads Grid Y selects the compact K/V head and grid Z selects one of its query heads. This avoids a device integer divide and keeps K/V planes compact. The equal-head case uses one group and a unit Z dimension.
395def streamingAttentionGroupedSM86Admitted =
396  (lambda unrestricted seq : Nat .
397    (lambda unrestricted heads : Nat .
398      (lambda unrestricted keyValueHeads : Nat .
399        (naturalAnd (naturalNonzero keyValueHeads)
400          (naturalAnd (naturalNonzero heads)
401            (naturalAnd (naturalNonzero seq)
402              (naturalAnd (naturalEqual (naturalModuloUnchecked seq saTile) 0)
403                (naturalAnd
404                  (naturalEqual
405                    (naturalModuloUnchecked heads
406                      (naturalSelect (naturalNonzero keyValueHeads) keyValueHeads 1)) 0)
407                  (naturalAnd (naturalLessOrEqual heads 65535)
408                    (naturalLessOrEqual
409                      (naturalMultiply heads (naturalMultiply saHeadWidth (naturalMultiply seq 4)))
410                      4294967295))))))))))

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.