AlphaGitHub
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
Download exact fileView at revisionbcc11a4e7647d31628

def · lines 457–472

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.
457def streamingAttentionGroupedSM86Admitted =
458  (lambda unrestricted seq : Nat .
459    (lambda unrestricted heads : Nat .
460      (lambda unrestricted keyValueHeads : Nat .
461        (naturalAnd (naturalNonzero keyValueHeads)
462          (naturalAnd (naturalNonzero heads)
463            (naturalAnd (naturalNonzero seq)
464              (naturalAnd (naturalEqual (naturalModuloUnchecked seq saTile) 0)
465                (naturalAnd
466                  (naturalEqual
467                    (naturalModuloUnchecked heads
468                      (naturalSelect (naturalNonzero keyValueHeads) keyValueHeads 1)) 0)
469                  (naturalAnd (naturalLessOrEqual heads 65535)
470                    (naturalLessOrEqual
471                      (naturalMultiply heads (naturalMultiply saHeadWidth (naturalMultiply seq 4)))
472                      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.