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 431–437

saAxis

Full file
---- block order ---- Order 0: grid x the query tile, y the K/V head, z the query group. Order 1: grid x the K/V head, y the query group, z the tile. Blocks are dispatched in linear order, x fastest, so order 1 issues one tile of every head before the next tile; the forward and query passes also count their tiles down from the diagonal one, so the longest causal blocks start first and the short ones fill the tail. Each block computes exactly what it computed before.
431def saAxis = (lambda unrestricted order : Nat .
432  (lambda unrestricted first : (family SM86SpecialRegister) .
433    (lambda unrestricted second : (family SM86SpecialRegister) .
434      (nat-eliminate (lambda unrestricted mode : Nat . (family SM86SpecialRegister))
435        first
436        (lambda unrestricted predecessor : Nat . (lambda unrestricted induction : (family SM86SpecialRegister) . second))
437        order))))

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.