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 2439–2456

skScoreGradient

Full file
P^T = 2^(t - L), dS^T = P^T (dP^T - D), in S^T's registers
2439def skScoreGradient = (lambda unrestricted tail : (family SM86Program) .
2440  (saFor 16 (lambda unrestricted i : Nat .
2441    (let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
2442      (let unrestricted b = (naturalModuloUnchecked i 2) in
2443      (let unrestricted e = (saElem (skS nt) h b) in
2444      (lambda unrestricted rest : (family SM86Program) .
2445        (saFfma e e saScale (skNegL nt b)
2446          (saEx2 e e saWaitNone rest))))))))
2447  -- P^T's fragments, before dS^T replaces it
2448  (skPackRowFragments skS skPA saWait4
2449  (saFor 16 (lambda unrestricted i : Nat .
2450    (let unrestricted nt = (naturalDivideUnchecked i 4) in (let unrestricted h = (naturalDivideUnchecked (naturalModuloUnchecked i 4) 2) in
2451      (let unrestricted b = (naturalModuloUnchecked i 2) in
2452      (let unrestricted e = (saElem (skS nt) h b) in (let unrestricted g = (saElem (skDP nt) h b) in
2453      (lambda unrestricted rest : (family SM86Program) .
2454        (skWeightedDifference e e g (skNegD nt b) saWaitNone rest))))))))
2455  (skPackRowFragments skS skSA saWaitNone
2456    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.