474def streamingAttentionForwardGroupedOrderedUncheckedSM86 = (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
475 (let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
476 (let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
477 (let unrestricted rowBytes = (naturalMultiply heads (naturalMultiply saHeadWidth saHalfBytes)) in
478 (let unrestricted body = (saTileBody seq sm86ProgramEmpty) in
479 (saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
480 (saS2R saTileIndex (saTileAxis order)
481 (saS2R saKeyValueHead (saKeyValueAxis order)
482 (saS2R saHead (saGroupAxis order)
483 (saMovImm saOne 1
484 -- the special registers are in once this waits
485 (saMovImmAfter saScratch 31 saWait5
486 (saImad saHead saKeyValueHead groups saHead
487 (saCountDownTile order (naturalDivideUnchecked seq saTile)
488 (saAnd saLane saTid saScratch
489 (saShr saWarp saTid 5
490 (saShr saQuad saLane 2
491 (saMovImm saScratch 3
492 (saAnd saQuadLane saLane saScratch
493 -- the thread's query row: 64 tile + 16 warp + lane / 4
494 (saImad saRow saWarp 16 saQuad
495 (saImad saRow saTileIndex saTile saRow
496 -- Q: head h's plane, the row, the quad lane's pair of columns
497 (saMovImm saScratch 0
498 (saImad saScratch saRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
499 (saImad saScratch saQuadLane 4 saScratch
500 (saImad saScratch saHead headBytes saScratch
501 (saWide saPointer saScratch saOne (saArgument 0)
502 (saFor 16 (lambda unrestricted i : Nat .
503 (let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
504 (saLoad (saQReg kt j) saPointer
505 (naturalAdd (naturalMultiply kt 32)
506 (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
507 (naturalMultiply (naturalDivideUnchecked j 2) 16)))
508 saSB1 saWaitNone))))
509 -- the output's offset in the merged plane (kept in saScratch): the row,
510 -- head h's 64 columns, the quad lane's pair
511 (saMovImm saScratch 0
512 (saImad saScratch saRow rowBytes saScratch
513 (saImad saScratch saHead (naturalMultiply saHeadWidth saHalfBytes) saScratch
514 (saImad saScratch saQuadLane 4 saScratch
515 -- the log-sum-exp's (kept in saRow): 4 (h seq + row)
516 (saImad saRow saHead seq saRow
517 (saMovImm saKeyOffset 0
518 (saImad saRow saRow 4 saKeyOffset
519 -- K from key 0: the quad's row lane / 4, its pair of columns
520 (saImad saKeyOffset saQuad (naturalMultiply saHeadWidth saHalfBytes) saKeyOffset
521 (saImad saKeyOffset saQuadLane 4 saKeyOffset
522 (saImad saKeyOffset saKeyValueHead headBytes saKeyOffset
523 -- V^T from key 0: row lane / 4 of head h's 64 x seq plane
524 (saMovImm saValueOffset 0
525 (saImad saValueOffset saQuad (naturalMultiply seq saHalfBytes) saValueOffset
526 (saImad saValueOffset saQuadLane 4 saValueOffset
527 (saImad saValueOffset saKeyValueHead headBytes saValueOffset
528 -- the tiles up to the diagonal one; the mask's base, 16 warp + lane / 4
529 -- - 2 (lane % 4) + 64 - 4096
530 (saAddImm saCount saTileIndex 1
531 (saImad saMaskBase saWarp 16 saQuad
532 (saImad saMaskBase saQuadLane 4294967294 saMaskBase
533 (saAddImm saMaskBase saMaskBase (naturalSaturatingSubtract 4294967296 (naturalSaturatingSubtract 4096 64))
534 (saMovConst saScale (saArgument 5)
535 (saMovImm (saM 0) saMinusInfinity (saMovImm (saM 1) saMinusInfinity
536 (saMovImm (saL 0) 0 (saMovImm (saL 1) 0
537 (saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
538 (sm86ProgramAppend body
539 (sm86ProgramAppend (saLoopTail seq (sm86ProgramCount body))
540 (saEpilogue seq heads
541 (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.