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 3046–3057

streamingAttentionKeySharedProjectionSM86

Full file
The same storing dK and dV into a fused projection's gradient (skProjectionEpilogue) after the +0 block below the diagonal.
3046def streamingAttentionKeySharedProjectionSM86 =
3047  (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
3048    (lambda unrestricted keyValueHeads : Nat . (lambda unrestricted width : Nat .
3049      (lambda erased admitted : (equal Nat (streamingAttentionKeyProjectionAdmitted seq heads keyValueHeads width) 1) .
3050        (sksPrologue 1 seq heads keyValueHeads 0 1
3051          (sm86ProgramAppend (sksBodyA seq)
3052          (sm86ProgramAppend (sfFirstTail (sm86ProgramCount (sksBodyB seq)))
3053          (sm86ProgramAppend (sksBodyB seq)
3054          (sm86ProgramAppend (sfSecondTail (sm86ProgramCount (sksBodyA seq)) (sm86ProgramCount (sksBodyB seq)))
3055            (sksZeroBelowDiagonal seq
3056              (skProjectionEpilogue width seq
3057                (saOp (constructor SM86InstructionBody SM86Exit (saAfter saWaitAll)) sm86ProgramEmpty)))))))))))))

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.