412def streamingAttentionForwardGroupedUncheckedSM86 = (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
413 (let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
414 (let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
415 (let unrestricted rowBytes = (naturalMultiply heads (naturalMultiply saHeadWidth saHalfBytes)) in
416 (let unrestricted body = (saTileBody seq sm86ProgramEmpty) in
417 (saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
418 (saS2R saTileIndex (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdX)
419 (saS2R saKeyValueHead (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdY)
420 (saS2R saHead (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdZ)
421 (saMovImm saOne 1
422 -- the special registers are in once this waits
423 (saMovImmAfter saScratch 31 saWait5
424 (saImad saHead saKeyValueHead groups saHead
425 (saAnd saLane saTid saScratch
426 (saShr saWarp saTid 5
427 (saShr saQuad saLane 2
428 (saMovImm saScratch 3
429 (saAnd saQuadLane saLane saScratch
430 -- the thread's query row: 64 tile + 16 warp + lane / 4
431 (saImad saRow saWarp 16 saQuad
432 (saImad saRow saTileIndex saTile saRow
433 -- Q: head h's plane, the row, the quad lane's pair of columns
434 (saMovImm saScratch 0
435 (saImad saScratch saRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
436 (saImad saScratch saQuadLane 4 saScratch
437 (saImad saScratch saHead headBytes saScratch
438 (saWide saPointer saScratch saOne (saArgument 0)
439 (saFor 16 (lambda unrestricted i : Nat .
440 (let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
441 (saLoad (saQReg kt j) saPointer
442 (naturalAdd (naturalMultiply kt 32)
443 (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
444 (naturalMultiply (naturalDivideUnchecked j 2) 16)))
445 saSB1 saWaitNone))))
446 -- the output's offset in the merged plane (kept in saScratch): the row,
447 -- head h's 64 columns, the quad lane's pair
448 (saMovImm saScratch 0
449 (saImad saScratch saRow rowBytes saScratch
450 (saImad saScratch saHead (naturalMultiply saHeadWidth saHalfBytes) saScratch
451 (saImad saScratch saQuadLane 4 saScratch
452 -- the log-sum-exp's (kept in saRow): 4 (h seq + row)
453 (saImad saRow saHead seq saRow
454 (saMovImm saKeyOffset 0
455 (saImad saRow saRow 4 saKeyOffset
456 -- K from key 0: the quad's row lane / 4, its pair of columns
457 (saImad saKeyOffset saQuad (naturalMultiply saHeadWidth saHalfBytes) saKeyOffset
458 (saImad saKeyOffset saQuadLane 4 saKeyOffset
459 (saImad saKeyOffset saKeyValueHead headBytes saKeyOffset
460 -- V^T from key 0: row lane / 4 of head h's 64 x seq plane
461 (saMovImm saValueOffset 0
462 (saImad saValueOffset saQuad (naturalMultiply seq saHalfBytes) saValueOffset
463 (saImad saValueOffset saQuadLane 4 saValueOffset
464 (saImad saValueOffset saKeyValueHead headBytes saValueOffset
465 -- the tiles up to the diagonal one; the mask's base, 16 warp + lane / 4
466 -- - 2 (lane % 4) + 64 - 4096
467 (saAddImm saCount saTileIndex 1
468 (saImad saMaskBase saWarp 16 saQuad
469 (saImad saMaskBase saQuadLane 4294967294 saMaskBase
470 (saAddImm saMaskBase saMaskBase (naturalSaturatingSubtract 4294967296 (naturalSaturatingSubtract 4096 64))
471 (saMovConst saScale (saArgument 5)
472 (saMovImm (saM 0) saMinusInfinity (saMovImm (saM 1) saMinusInfinity
473 (saMovImm (saL 0) 0 (saMovImm (saL 1) 0
474 (saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
475 (sm86ProgramAppend body
476 (sm86ProgramAppend (saLoopTail seq (sm86ProgramCount body))
477 (saEpilogue seq heads
478 (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.