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 671–677

sfScores

Full file
S = Q K^T and O += P V as the register kernel's, the first product waiting for the fragments (SB1, which also ordered Q's loads) The first slice reads RZ for its accumulator, so no tile zeroes the scores first: +0 + the products is the same word either way.
671def sfScores = (lambda unrestricted tail : (family SM86Program) .
672  (saFor 4 (lambda unrestricted k : Nat .
673    (saFor 8 (lambda unrestricted nt : Nat .
674      (saHmmaFrom (saS nt) (saQ k) (saB k nt)
675        (naturalSelect (naturalIsZero k) sm86RegisterZero (saS nt)) saSB2
676        (naturalSelect (naturalIsZero nt) (naturalSelect (naturalIsZero k) saWait1 saWait2) saWaitNone)))))
677    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.