AlphaGitHub
Source/Packages

Realization.Nvidia.SM86.StreamingAttentionSM86

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

1,203 lines221 declarations70.9 KiBSHA-256 23d4e5e2aa2a
Download exact fileView at revision0dc6acfc09dcaa39d8

def · lines 407–422

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.
407def streamingAttentionGroupedSM86Admitted =
408  (lambda unrestricted seq : Nat .
409    (lambda unrestricted heads : Nat .
410      (lambda unrestricted keyValueHeads : Nat .
411        (naturalAnd (naturalNonzero keyValueHeads)
412          (naturalAnd (naturalNonzero heads)
413            (naturalAnd (naturalNonzero seq)
414              (naturalAnd (naturalEqual (naturalModuloUnchecked seq saTile) 0)
415                (naturalAnd
416                  (naturalEqual
417                    (naturalModuloUnchecked heads
418                      (naturalSelect (naturalNonzero keyValueHeads) keyValueHeads 1)) 0)
419                  (naturalAnd (naturalLessOrEqual heads 65535)
420                    (naturalLessOrEqual
421                      (naturalMultiply heads (naturalMultiply saHeadWidth (naturalMultiply seq 4)))
422                      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.