1module Realization.Nvidia.SM86.AttentionHeadLayoutSM86
2
3import Accelerator.SM86.Control
4import Accelerator.SM86.Immediate
5import Accelerator.SM86.Instruction
6import Accelerator.SM86.InstructionEncoding
7import Accelerator.SM86.Program
8import Accelerator.SM86.Types
9import Data.SHA256Digest
10import Std.Natural
11
12-- coppelius S5 rung 4b (D24): the ATTENTION HEAD LAYOUT kernels at coppelius's
13-- geometry — sequence 1024, 8 heads of 64 (half width 32), model width 512,
14-- so the fused QKV projection is a 1024 × 1536 fp32 row-major matrix whose
15-- row is [Q heads 0..7 | K heads 0..7 | V heads 0..7], each head 64 wide.
16--
17-- The tensor-core GEMM (Realization.Nvidia.SM86.GEMM.HMMA.*) computes C = A · Bᵀ
18-- with A (M × K) and B (N × K) both fp16 row-major. So per head h:
19-- scores_h = Q_h · K_hᵀ A = Q_h (1024 × 64), B = K_h (1024 × 64)
20-- out_h = P_h · V_h A = P_h (1024 × 1024), B = V_hᵀ (64 × 1024)
21-- and the layout kernels here produce exactly those operands:
22-- QK ROTARY: for every position p, head h and component c < 32 the HF
23-- half-split rotation (Llama): q'[c] = q[c]·cos − q[c+32]·sin,
24-- q'[c+32] = q[c+32]·cos + q[c]·sin, with cos/sin tables of 1024 × 32 fp32
25-- (angle p · theta^(−2c/64)); the same for K; both written fp16 head-major
26-- [head][position][64]. Grid 1024 (one block per position), block 128:
27-- thread = head (tid >> 4) × component pair (tid & 15 → c = 2g, 2g+1), four
28-- fp32 loads per plane, two packed half2 stores per plane.
29-- VALUE TRANSPOSE: V written fp16 as [head][component][position]. Grid 512
30-- (one block per PAIR of positions, so a thread packs two adjacent
31-- positions into one half2 store), block 512: thread = head (tid >> 6) ×
32-- component (tid & 63).
33-- Constant bank 0 (64-bit pointers): 0x160 first output (Q or V), 0x168 second
34-- output (K; unused by the value kernel), 0x170 QKV input, 0x178 cosine table,
35-- 0x180 sine table.
36-- No host fallback: the image is emitted by the sovereign encoder or the
37-- build refuses with a stable code.
38family AttentionHeadLayoutSM86Kind : Type 0
39constructor AttentionHeadLayoutQKRotary
40constructor AttentionHeadLayoutValueTranspose
41constructor AttentionHeadLayoutHeadToToken
42constructor AttentionHeadLayoutTokenToHead
43constructor AttentionHeadLayoutInverseRoPEQKVMerge
44
45end-family
46
47family AttentionHeadLayoutSM86FailureCode : Type 0
48constructor AttentionHeadLayoutInstructionCountMismatch
49constructor AttentionHeadLayoutEncodingFailed
50constructor AttentionHeadLayoutIdentityInvalid
51
52end-family
53
54family AttentionHeadLayoutSM86Extents : Type 0
55constructor AttentionHeadLayoutSM86ExtentsValue
56field unrestricted attentionHeadLayoutExtentSequence : Nat
57field unrestricted attentionHeadLayoutExtentHeads : Nat
58field unrestricted attentionHeadLayoutExtentHeadWidth : Nat
59field unrestricted attentionHeadLayoutExtentHalfWidth : Nat
60field unrestricted attentionHeadLayoutExtentModelWidth : Nat
61field unrestricted attentionHeadLayoutExtentQKVWidth : Nat
62
63end-family
64
65family AttentionHeadLayoutSM86Telemetry : Type 0
66constructor AttentionHeadLayoutSM86TelemetryValue
67field unrestricted attentionHeadLayoutTelemetryKind : (family AttentionHeadLayoutSM86Kind)
68field unrestricted attentionHeadLayoutTelemetryExpectedInstructions : Nat
69field unrestricted attentionHeadLayoutTelemetryActualInstructions : Nat
70field unrestricted attentionHeadLayoutTelemetryRegisters : Nat
71field unrestricted attentionHeadLayoutTelemetryGridX : Nat
72field unrestricted attentionHeadLayoutTelemetryBlockX : Nat
73field unrestricted attentionHeadLayoutTelemetryExtents : (family AttentionHeadLayoutSM86Extents)
74field unrestricted attentionHeadLayoutTelemetryEncodedBytes : Nat
75field unrestricted attentionHeadLayoutTelemetryEncodedFields : Nat
76field unrestricted attentionHeadLayoutTelemetryEncodedBits : Nat
77
78end-family
79
80family AttentionHeadLayoutSM86ImageResult : Type 0
81constructor AttentionHeadLayoutSM86ImageReady
82field unrestricted attentionHeadLayoutImageKind : (family AttentionHeadLayoutSM86Kind)
83field unrestricted attentionHeadLayoutImageBytes : Bytes
84field unrestricted attentionHeadLayoutImageSHA256 : Bytes
85field unrestricted attentionHeadLayoutImageTelemetry : (family AttentionHeadLayoutSM86Telemetry)
86constructor AttentionHeadLayoutSM86ImageFailed
87field unrestricted attentionHeadLayoutImageFailure : (family AttentionHeadLayoutSM86FailureCode)
88field unrestricted attentionHeadLayoutImageFailureInstruction : Nat
89field unrestricted attentionHeadLayoutImageFailureDetail : Bytes
90field unrestricted attentionHeadLayoutImageFailureTelemetry : (family AttentionHeadLayoutSM86Telemetry)
91
92end-family
93
94-- naturals
95def attentionHeadLayoutSM86N8 =
96 (byte-to-nat (byte 8))
97
98def attentionHeadLayoutSM86N23 =
99 (byte-to-nat (byte 23))
100
101def attentionHeadLayoutSM86N24 =
102 (byte-to-nat (byte 24))
103
104def attentionHeadLayoutSM86N31 =
105 (byte-to-nat (byte 31))
106
107def attentionHeadLayoutSM86N32 =
108 (byte-to-nat (byte 32))
109
110def attentionHeadLayoutSM86N48 =
111 (byte-to-nat (byte 48))
112
113def attentionHeadLayoutSM86N56 =
114 (byte-to-nat (byte 56))
115
116def attentionHeadLayoutSM86N63 =
117 (byte-to-nat (byte 63))
118
119def attentionHeadLayoutSM86N64 =
120 (byte-to-nat (byte 64))
121
122def attentionHeadLayoutSM86N72 =
123 (byte-to-nat (byte 72))
124
125def attentionHeadLayoutSM86N73 =
126 (byte-to-nat (byte 73))
127
128def attentionHeadLayoutSM86N77 =
129 (byte-to-nat (byte 77))
130
131def attentionHeadLayoutSM86N128 =
132 (byte-to-nat (byte 128))
133
134def attentionHeadLayoutSM86N256 =
135 (nat-multiply (byte-to-nat (byte 16)) (byte-to-nat (byte 16)))
136
137def attentionHeadLayoutSM86N512 =
138 (nat-multiply attentionHeadLayoutSM86N256 (byte-to-nat (byte 2)))
139
140def attentionHeadLayoutSM86N1024 =
141 (nat-multiply attentionHeadLayoutSM86N256 (byte-to-nat (byte 4)))
142
143def attentionHeadLayoutSM86N1536 =
144 (nat-multiply attentionHeadLayoutSM86N256 (byte-to-nat (byte 6)))
145
146def attentionHeadLayoutSM86Extents : (family AttentionHeadLayoutSM86Extents) =
147 (constructor
148 AttentionHeadLayoutSM86Extents
149 AttentionHeadLayoutSM86ExtentsValue
150 attentionHeadLayoutSM86N1024
151 attentionHeadLayoutSM86N8
152 attentionHeadLayoutSM86N64
153 attentionHeadLayoutSM86N32
154 attentionHeadLayoutSM86N512
155 attentionHeadLayoutSM86N1536)
156
157-- immediates
158def attentionHeadLayoutSM86U =
159 (lambda unrestricted b0 : Byte .
160 (lambda unrestricted b1 : Byte .
161 (lambda unrestricted b2 : Byte .
162 (lambda unrestricted b3 : Byte . (sm86Unsigned32 b0 b1 b2 b3)))))
163
164def attentionHeadLayoutSM86U0 =
165 (attentionHeadLayoutSM86U (byte 0) (byte 0) (byte 0) (byte 0))
166
167def attentionHeadLayoutSM86U1 =
168 (attentionHeadLayoutSM86U (byte 1) (byte 0) (byte 0) (byte 0))
169
170def attentionHeadLayoutSM86U2 =
171 (attentionHeadLayoutSM86U (byte 2) (byte 0) (byte 0) (byte 0))
172
173def attentionHeadLayoutSM86U4 =
174 (attentionHeadLayoutSM86U (byte 4) (byte 0) (byte 0) (byte 0))
175
176def attentionHeadLayoutSM86U15 =
177 (attentionHeadLayoutSM86U (byte 15) (byte 0) (byte 0) (byte 0))
178
179def attentionHeadLayoutSM86U31 =
180 (attentionHeadLayoutSM86U (byte 31) (byte 0) (byte 0) (byte 0))
181
182def attentionHeadLayoutSM86U32 =
183 (attentionHeadLayoutSM86U (byte 32) (byte 0) (byte 0) (byte 0))
184
185def attentionHeadLayoutSM86U63 =
186 (attentionHeadLayoutSM86U (byte 63) (byte 0) (byte 0) (byte 0))
187
188def attentionHeadLayoutSM86U64 =
189 (attentionHeadLayoutSM86U (byte 64) (byte 0) (byte 0) (byte 0))
190
191def attentionHeadLayoutSM86U128 =
192 (attentionHeadLayoutSM86U (byte 128) (byte 0) (byte 0) (byte 0))
193
194def attentionHeadLayoutSM86U132 =
195 (attentionHeadLayoutSM86U (byte 132) (byte 0) (byte 0) (byte 0))
196
197def attentionHeadLayoutSM86U512 =
198 (attentionHeadLayoutSM86U (byte 0) (byte 2) (byte 0) (byte 0))
199
200def attentionHeadLayoutSM86U1024 =
201 (attentionHeadLayoutSM86U (byte 0) (byte 4) (byte 0) (byte 0))
202
203def attentionHeadLayoutSM86U1088 =
204 (attentionHeadLayoutSM86U (byte 64) (byte 4) (byte 0) (byte 0))
205
206def attentionHeadLayoutSM86U1536 =
207 (attentionHeadLayoutSM86U (byte 0) (byte 6) (byte 0) (byte 0))
208
209def attentionHeadLayoutSM86U2048 =
210 (attentionHeadLayoutSM86U (byte 0) (byte 8) (byte 0) (byte 0))
211
212def attentionHeadLayoutSM86U2052 =
213 (attentionHeadLayoutSM86U (byte 4) (byte 8) (byte 0) (byte 0))
214
215def attentionHeadLayoutSM86U2112 =
216 (attentionHeadLayoutSM86U (byte 64) (byte 8) (byte 0) (byte 0))
217
218def attentionHeadLayoutSM86U2176 =
219 (attentionHeadLayoutSM86U (byte 128) (byte 8) (byte 0) (byte 0))
220
221def attentionHeadLayoutSM86U2180 =
222 (attentionHeadLayoutSM86U (byte 132) (byte 8) (byte 0) (byte 0))
223
224def attentionHeadLayoutSM86U3072 =
225 (attentionHeadLayoutSM86U (byte 0) (byte 12) (byte 0) (byte 0))
226
227def attentionHeadLayoutSM86U4096 =
228 (attentionHeadLayoutSM86U (byte 0) (byte 16) (byte 0) (byte 0))
229
230def attentionHeadLayoutSM86U6144 =
231 (attentionHeadLayoutSM86U (byte 0) (byte 24) (byte 0) (byte 0))
232
233def attentionHeadLayoutSM86U65536 =
234 (attentionHeadLayoutSM86U (byte 0) (byte 0) (byte 1) (byte 0))
235
236def attentionHeadLayoutSM86U131072 =
237 (attentionHeadLayoutSM86U (byte 0) (byte 0) (byte 2) (byte 0))
238
239def attentionHeadLayoutSM86U135168 =
240 (attentionHeadLayoutSM86U (byte 0) (byte 16) (byte 2) (byte 0))
241
242-- constant bank 0 offsets
243def attentionHeadLayoutSM86ConstantBase =
244 (attentionHeadLayoutSM86U (byte 40) (byte 0) (byte 0) (byte 0))
245
246def attentionHeadLayoutSM86FirstOutputPointer =
247 (attentionHeadLayoutSM86U (byte 96) (byte 1) (byte 0) (byte 0))
248
249-- Ordered launch arguments shared by the specialised Q, K and V layouts.
250-- V reads through the input pointer but has no table operands.
251def attentionHeadLayoutSM86FirstOutputArgument : Nat = 0
252def attentionHeadLayoutSM86SecondOutputArgument : Nat = 1
253def attentionHeadLayoutSM86InputArgument : Nat = 2
254def attentionHeadLayoutSM86CosineArgument : Nat = 3
255def attentionHeadLayoutSM86SineArgument : Nat = 4
256
257def attentionHeadLayoutSM86SecondOutputPointer =
258 (attentionHeadLayoutSM86U (byte 104) (byte 1) (byte 0) (byte 0))
259
260def attentionHeadLayoutSM86InputPointer =
261 (attentionHeadLayoutSM86U (byte 112) (byte 1) (byte 0) (byte 0))
262
263def attentionHeadLayoutSM86CosinePointer =
264 (attentionHeadLayoutSM86U (byte 120) (byte 1) (byte 0) (byte 0))
265
266def attentionHeadLayoutSM86SinePointer =
267 (attentionHeadLayoutSM86U (byte 128) (byte 1) (byte 0) (byte 0))
268
269def attentionHeadLayoutSM86ValueGradientPointer =
270 (attentionHeadLayoutSM86U (byte 136) (byte 1) (byte 0) (byte 0))
271
272-- registers and program plumbing
273def attentionHeadLayoutSM86R =
274 (lambda unrestricted value : Byte . (sm86Register value))
275
276def attentionHeadLayoutSM86Next =
277 (lambda unrestricted body : (family SM86InstructionBody) .
278 (lambda unrestricted tail : (family SM86Program) .
279 (constructor
280 SM86Program
281 SM86ProgramNext
282 (constructor
283 SM86Instruction
284 SM86InstructionValue
285 (constructor SM86InstructionGuard SM86InstructionAlways)
286 body)
287 tail)))
288
289def attentionHeadLayoutSM86End =
290 (constructor SM86Program SM86ProgramEnd)
291
292def attentionHeadLayoutSM86SetSpecials =
293 (sm86SetBarrierControl (constructor SM86Barrier SM86Barrier0))
294
295def attentionHeadLayoutSM86SetLoads =
296 (sm86SetBarrierControl (constructor SM86Barrier SM86Barrier1))
297
298-- wait on barrier 0 (the special-register reads): mask bit 0
299def attentionHeadLayoutSM86WaitSpecials : (family SM86Control) =
300 (constructor
301 SM86Control
302 SM86ControlValue
303 (byte 7)
304 (constructor SM86YieldMode SM86Continue)
305 (constructor SM86Barrier SM86BarrierNone)
306 (constructor SM86Barrier SM86BarrierNone)
307 (byte 1)
308 (byte 0))
309
310-- wait on barrier 1 (every global load): mask bit 1
311def attentionHeadLayoutSM86WaitLoads : (family SM86Control) =
312 (constructor
313 SM86Control
314 SM86ControlValue
315 (byte 7)
316 (constructor SM86YieldMode SM86Continue)
317 (constructor SM86Barrier SM86BarrierNone)
318 (constructor SM86Barrier SM86BarrierNone)
319 (byte 2)
320 (byte 0))
321
322def attentionHeadLayoutSM86Move =
323 (lambda unrestricted destination : Byte .
324 (lambda unrestricted value : (family SM86Unsigned32) .
325 (constructor
326 SM86InstructionBody
327 SM86MoveImmediate
328 (attentionHeadLayoutSM86R destination)
329 value
330 sm86SafeControl)))
331
332def attentionHeadLayoutSM86IMAD =
333 (lambda unrestricted destination : Byte .
334 (lambda unrestricted left : Byte .
335 (lambda unrestricted value : (family SM86Unsigned32) .
336 (lambda unrestricted addend : Byte .
337 (constructor
338 SM86InstructionBody
339 SM86IntegerMultiplyAddImmediate
340 (attentionHeadLayoutSM86R destination)
341 (attentionHeadLayoutSM86R left)
342 value
343 (attentionHeadLayoutSM86R addend)
344 sm86SafeControl)))))
345
346def attentionHeadLayoutSM86Wide =
347 (lambda unrestricted destination : Byte .
348 (lambda unrestricted left : Byte .
349 (lambda unrestricted right : Byte .
350 (lambda unrestricted pointer : (family SM86Unsigned32) .
351 (constructor
352 SM86InstructionBody
353 SM86IntegerMultiplyAddWideConstant
354 (attentionHeadLayoutSM86R destination)
355 (attentionHeadLayoutSM86R left)
356 (attentionHeadLayoutSM86R right)
357 (byte 0)
358 pointer
359 sm86SafeControl)))))
360
361def attentionHeadLayoutSM86Load =
362 (lambda unrestricted destination : Byte .
363 (lambda unrestricted address : Byte .
364 (lambda unrestricted offset : (family SM86Unsigned32) .
365 (constructor
366 SM86InstructionBody
367 SM86LoadGlobal
368 (attentionHeadLayoutSM86R destination)
369 (attentionHeadLayoutSM86R address)
370 offset
371 attentionHeadLayoutSM86SetLoads))))
372
373def attentionHeadLayoutSM86Store =
374 (lambda unrestricted address : Byte .
375 (lambda unrestricted value : Byte .
376 (lambda unrestricted offset : (family SM86Unsigned32) .
377 (constructor
378 SM86InstructionBody
379 SM86StoreGlobal
380 (attentionHeadLayoutSM86R address)
381 (attentionHeadLayoutSM86R value)
382 offset
383 sm86SafeControl))))
384
385def attentionHeadLayoutSM86Multiply =
386 (lambda unrestricted destination : Byte .
387 (lambda unrestricted left : Byte .
388 (lambda unrestricted right : Byte .
389 (lambda unrestricted control : (family SM86Control) .
390 (constructor
391 SM86InstructionBody
392 SM86FloatMultiply
393 (attentionHeadLayoutSM86R destination)
394 (attentionHeadLayoutSM86R left)
395 (attentionHeadLayoutSM86R right)
396 control)))))
397
398def attentionHeadLayoutSM86Negate =
399 (lambda unrestricted register : Byte .
400 (constructor
401 SM86InstructionBody
402 SM86FloatNegate
403 (attentionHeadLayoutSM86R register)
404 (attentionHeadLayoutSM86R register)
405 sm86SafeControl))
406
407def attentionHeadLayoutSM86NegateWithControl =
408 (lambda unrestricted register : Byte .
409 (lambda unrestricted control : (family SM86Control) .
410 (constructor
411 SM86InstructionBody
412 SM86FloatNegate
413 (attentionHeadLayoutSM86R register)
414 (attentionHeadLayoutSM86R register)
415 control)))
416
417def attentionHeadLayoutSM86FMA =
418 (lambda unrestricted destination : Byte .
419 (lambda unrestricted left : Byte .
420 (lambda unrestricted right : Byte .
421 (constructor
422 SM86InstructionBody
423 SM86FloatFusedMultiplyAdd
424 (attentionHeadLayoutSM86R destination)
425 (attentionHeadLayoutSM86R left)
426 (attentionHeadLayoutSM86R right)
427 (attentionHeadLayoutSM86R destination)
428 sm86SafeControl))))
429
430def attentionHeadLayoutSM86Pack =
431 (lambda unrestricted destination : Byte .
432 (lambda unrestricted high : Byte .
433 (lambda unrestricted low : Byte .
434 (lambda unrestricted control : (family SM86Control) .
435 (constructor
436 SM86InstructionBody
437 SM86FloatPairToPackedHalfPair
438 (attentionHeadLayoutSM86R destination)
439 (attentionHeadLayoutSM86R high)
440 (attentionHeadLayoutSM86R low)
441 control)))))
442
443-- the rotation of one component: x1 (first half), x2 (second half), cos, sin →
444-- out1 = x1·cos − x2·sin ; out2 = x2·cos + x1·sin. `first` waits on the loads.
445def attentionHeadLayoutSM86Rotate =
446 (lambda unrestricted out1 : Byte .
447 (lambda unrestricted out2 : Byte .
448 (lambda unrestricted x1 : Byte .
449 (lambda unrestricted x2 : Byte .
450 (lambda unrestricted cosine : Byte .
451 (lambda unrestricted sine : Byte .
452 (lambda unrestricted first : (family SM86Control) .
453 (lambda unrestricted tail : (family SM86Program) .
454 (attentionHeadLayoutSM86Next
455 (attentionHeadLayoutSM86Multiply out1 x2 sine first)
456 (attentionHeadLayoutSM86Next
457 (attentionHeadLayoutSM86Negate out1)
458 (attentionHeadLayoutSM86Next
459 (attentionHeadLayoutSM86FMA out1 x1 cosine)
460 (attentionHeadLayoutSM86Next
461 (attentionHeadLayoutSM86Multiply out2 x1 sine sm86SafeControl)
462 (attentionHeadLayoutSM86Next
463 (attentionHeadLayoutSM86FMA out2 x2 cosine)
464 tail)))))))))))))
465
466-- shared prologue: R1 ← c[0x28]; R0 = tid; R1 = block; R4 = 4; R7 = 0; R9 = 2;
467-- R5 = mask; R2 = tid >> shift; R3 = tid & mask
468def attentionHeadLayoutSM86Prologue =
469 (lambda unrestricted shift : Byte .
470 (lambda unrestricted mask : (family SM86Unsigned32) .
471 (lambda unrestricted tail : (family SM86Program) .
472 (attentionHeadLayoutSM86Next
473 (constructor
474 SM86InstructionBody
475 SM86MoveConstant
476 (attentionHeadLayoutSM86R (byte 1))
477 (byte 0)
478 attentionHeadLayoutSM86ConstantBase
479 sm86SafeControl)
480 (attentionHeadLayoutSM86Next
481 (constructor
482 SM86InstructionBody
483 SM86SpecialToRegister
484 (attentionHeadLayoutSM86R (byte 0))
485 (constructor SM86SpecialRegister SM86ThreadIdX)
486 attentionHeadLayoutSM86SetSpecials)
487 (attentionHeadLayoutSM86Next
488 (constructor
489 SM86InstructionBody
490 SM86SpecialToRegister
491 (attentionHeadLayoutSM86R (byte 1))
492 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdX)
493 attentionHeadLayoutSM86SetSpecials)
494 (attentionHeadLayoutSM86Next
495 (attentionHeadLayoutSM86Move (byte 4) attentionHeadLayoutSM86U4)
496 (attentionHeadLayoutSM86Next
497 (attentionHeadLayoutSM86Move (byte 7) attentionHeadLayoutSM86U0)
498 (attentionHeadLayoutSM86Next
499 (attentionHeadLayoutSM86Move (byte 5) mask)
500 (attentionHeadLayoutSM86Next
501 (attentionHeadLayoutSM86Move (byte 9) attentionHeadLayoutSM86U2)
502 (attentionHeadLayoutSM86Next
503 (constructor
504 SM86InstructionBody
505 SM86ShiftRightImmediate
506 (attentionHeadLayoutSM86R (byte 2))
507 (attentionHeadLayoutSM86R (byte 0))
508 shift
509 attentionHeadLayoutSM86WaitSpecials)
510 (attentionHeadLayoutSM86Next
511 (constructor
512 SM86InstructionBody
513 SM86LogicThreeInputTruthTable
514 (attentionHeadLayoutSM86R (byte 3))
515 (attentionHeadLayoutSM86R (byte 0))
516 (attentionHeadLayoutSM86R (byte 5))
517 (byte 192)
518 sm86SafeControl)
519 tail))))))))))))
520
521-- QK rotary: 63 instructions
522def attentionHeadLayoutSM86QKRotaryProgram : (family SM86Program) =
523 (attentionHeadLayoutSM86Prologue
524 (byte 4)
525 attentionHeadLayoutSM86U15
526 (attentionHeadLayoutSM86Next
527 (attentionHeadLayoutSM86IMAD (byte 6) (byte 3) attentionHeadLayoutSM86U2 (byte 7))
528 (attentionHeadLayoutSM86Next
529 (attentionHeadLayoutSM86IMAD (byte 8) (byte 1) attentionHeadLayoutSM86U1536 (byte 7))
530 (attentionHeadLayoutSM86Next
531 (attentionHeadLayoutSM86IMAD (byte 8) (byte 2) attentionHeadLayoutSM86U64 (byte 8))
532 (attentionHeadLayoutSM86Next
533 (attentionHeadLayoutSM86IMAD (byte 8) (byte 6) attentionHeadLayoutSM86U1 (byte 8))
534 (attentionHeadLayoutSM86Next
535 (attentionHeadLayoutSM86Wide
536 (byte 10)
537 (byte 8)
538 (byte 4)
539 attentionHeadLayoutSM86InputPointer)
540 (attentionHeadLayoutSM86Next
541 (attentionHeadLayoutSM86Load (byte 12) (byte 10) attentionHeadLayoutSM86U0)
542 (attentionHeadLayoutSM86Next
543 (attentionHeadLayoutSM86Load (byte 13) (byte 10) attentionHeadLayoutSM86U4)
544 (attentionHeadLayoutSM86Next
545 (attentionHeadLayoutSM86Load (byte 14) (byte 10) attentionHeadLayoutSM86U128)
546 (attentionHeadLayoutSM86Next
547 (attentionHeadLayoutSM86Load (byte 15) (byte 10) attentionHeadLayoutSM86U132)
548 (attentionHeadLayoutSM86Next
549 (attentionHeadLayoutSM86Load
550 (byte 16)
551 (byte 10)
552 attentionHeadLayoutSM86U2048)
553 (attentionHeadLayoutSM86Next
554 (attentionHeadLayoutSM86Load
555 (byte 17)
556 (byte 10)
557 attentionHeadLayoutSM86U2052)
558 (attentionHeadLayoutSM86Next
559 (attentionHeadLayoutSM86Load
560 (byte 18)
561 (byte 10)
562 attentionHeadLayoutSM86U2176)
563 (attentionHeadLayoutSM86Next
564 (attentionHeadLayoutSM86Load
565 (byte 19)
566 (byte 10)
567 attentionHeadLayoutSM86U2180)
568 (attentionHeadLayoutSM86Next
569 (attentionHeadLayoutSM86IMAD
570 (byte 20)
571 (byte 1)
572 attentionHeadLayoutSM86U32
573 (byte 6))
574 (attentionHeadLayoutSM86Next
575 (attentionHeadLayoutSM86Wide
576 (byte 22)
577 (byte 20)
578 (byte 4)
579 attentionHeadLayoutSM86CosinePointer)
580 (attentionHeadLayoutSM86Next
581 (attentionHeadLayoutSM86Load
582 (byte 24)
583 (byte 22)
584 attentionHeadLayoutSM86U0)
585 (attentionHeadLayoutSM86Next
586 (attentionHeadLayoutSM86Load
587 (byte 25)
588 (byte 22)
589 attentionHeadLayoutSM86U4)
590 (attentionHeadLayoutSM86Next
591 (attentionHeadLayoutSM86Wide
592 (byte 26)
593 (byte 20)
594 (byte 4)
595 attentionHeadLayoutSM86SinePointer)
596 (attentionHeadLayoutSM86Next
597 (attentionHeadLayoutSM86Load
598 (byte 28)
599 (byte 26)
600 attentionHeadLayoutSM86U0)
601 (attentionHeadLayoutSM86Next
602 (attentionHeadLayoutSM86Load
603 (byte 29)
604 (byte 26)
605 attentionHeadLayoutSM86U4)
606 (attentionHeadLayoutSM86Rotate
607 (byte 30)
608 (byte 31)
609 (byte 12)
610 (byte 14)
611 (byte 24)
612 (byte 28)
613 attentionHeadLayoutSM86WaitLoads
614 (attentionHeadLayoutSM86Rotate
615 (byte 32)
616 (byte 33)
617 (byte 13)
618 (byte 15)
619 (byte 25)
620 (byte 29)
621 sm86SafeControl
622 (attentionHeadLayoutSM86Next
623 (attentionHeadLayoutSM86Pack
624 (byte 34)
625 (byte 32)
626 (byte 30)
627 sm86SafeControl)
628 (attentionHeadLayoutSM86Next
629 (attentionHeadLayoutSM86Pack
630 (byte 35)
631 (byte 33)
632 (byte 31)
633 sm86SafeControl)
634 (attentionHeadLayoutSM86Next
635 (attentionHeadLayoutSM86IMAD
636 (byte 36)
637 (byte 2)
638 attentionHeadLayoutSM86U65536
639 (byte 7))
640 (attentionHeadLayoutSM86Next
641 (attentionHeadLayoutSM86IMAD
642 (byte 36)
643 (byte 1)
644 attentionHeadLayoutSM86U64
645 (byte 36))
646 (attentionHeadLayoutSM86Next
647 (attentionHeadLayoutSM86IMAD
648 (byte 36)
649 (byte 6)
650 attentionHeadLayoutSM86U1
651 (byte 36))
652 (attentionHeadLayoutSM86Next
653 (attentionHeadLayoutSM86Wide
654 (byte 38)
655 (byte 36)
656 (byte 9)
657 attentionHeadLayoutSM86FirstOutputPointer)
658 (attentionHeadLayoutSM86Next
659 (attentionHeadLayoutSM86Store
660 (byte 38)
661 (byte 34)
662 attentionHeadLayoutSM86U0)
663 (attentionHeadLayoutSM86Next
664 (attentionHeadLayoutSM86Store
665 (byte 38)
666 (byte 35)
667 attentionHeadLayoutSM86U64)
668 (attentionHeadLayoutSM86Rotate
669 (byte 40)
670 (byte 41)
671 (byte 16)
672 (byte 18)
673 (byte 24)
674 (byte 28)
675 sm86SafeControl
676 (attentionHeadLayoutSM86Rotate
677 (byte 42)
678 (byte 43)
679 (byte 17)
680 (byte 19)
681 (byte 25)
682 (byte 29)
683 sm86SafeControl
684 (attentionHeadLayoutSM86Next
685 (attentionHeadLayoutSM86Pack
686 (byte 44)
687 (byte 42)
688 (byte 40)
689 sm86SafeControl)
690 (attentionHeadLayoutSM86Next
691 (attentionHeadLayoutSM86Pack
692 (byte 45)
693 (byte 43)
694 (byte 41)
695 sm86SafeControl)
696 (attentionHeadLayoutSM86Next
697 (attentionHeadLayoutSM86Wide
698 (byte 46)
699 (byte 36)
700 (byte 9)
701 attentionHeadLayoutSM86SecondOutputPointer)
702 (attentionHeadLayoutSM86Next
703 (attentionHeadLayoutSM86Store
704 (byte 46)
705 (byte 44)
706 attentionHeadLayoutSM86U0)
707 (attentionHeadLayoutSM86Next
708 (attentionHeadLayoutSM86Store
709 (byte 46)
710 (byte 45)
711 attentionHeadLayoutSM86U64)
712 (attentionHeadLayoutSM86Next
713 (constructor SM86InstructionBody SM86Exit sm86SafeControl)
714 attentionHeadLayoutSM86End)))))))))))))))))))))))))))))))))))))))
715
716-- value transpose: 23 instructions
717def attentionHeadLayoutSM86ValueTransposeProgram : (family SM86Program) =
718 (attentionHeadLayoutSM86Prologue
719 (byte 6)
720 attentionHeadLayoutSM86U63
721 (attentionHeadLayoutSM86Next
722 (attentionHeadLayoutSM86IMAD (byte 8) (byte 1) attentionHeadLayoutSM86U3072 (byte 7))
723 (attentionHeadLayoutSM86Next
724 (attentionHeadLayoutSM86IMAD (byte 8) (byte 2) attentionHeadLayoutSM86U64 (byte 8))
725 (attentionHeadLayoutSM86Next
726 (attentionHeadLayoutSM86IMAD (byte 8) (byte 3) attentionHeadLayoutSM86U1 (byte 8))
727 (attentionHeadLayoutSM86Next
728 (constructor
729 SM86InstructionBody
730 SM86IntegerAddThreeImmediate
731 (attentionHeadLayoutSM86R (byte 8))
732 (attentionHeadLayoutSM86R (byte 8))
733 attentionHeadLayoutSM86U1024
734 sm86SafeControl)
735 (attentionHeadLayoutSM86Next
736 (attentionHeadLayoutSM86Wide
737 (byte 10)
738 (byte 8)
739 (byte 4)
740 attentionHeadLayoutSM86InputPointer)
741 (attentionHeadLayoutSM86Next
742 (attentionHeadLayoutSM86Load (byte 12) (byte 10) attentionHeadLayoutSM86U0)
743 (attentionHeadLayoutSM86Next
744 (attentionHeadLayoutSM86Load (byte 13) (byte 10) attentionHeadLayoutSM86U6144)
745 (attentionHeadLayoutSM86Next
746 (attentionHeadLayoutSM86IMAD
747 (byte 16)
748 (byte 2)
749 attentionHeadLayoutSM86U64
750 (byte 3))
751 (attentionHeadLayoutSM86Next
752 (attentionHeadLayoutSM86IMAD
753 (byte 16)
754 (byte 16)
755 attentionHeadLayoutSM86U1024
756 (byte 7))
757 (attentionHeadLayoutSM86Next
758 (attentionHeadLayoutSM86IMAD
759 (byte 16)
760 (byte 1)
761 attentionHeadLayoutSM86U2
762 (byte 16))
763 (attentionHeadLayoutSM86Next
764 (attentionHeadLayoutSM86Wide
765 (byte 18)
766 (byte 16)
767 (byte 9)
768 attentionHeadLayoutSM86FirstOutputPointer)
769 (attentionHeadLayoutSM86Next
770 (attentionHeadLayoutSM86Pack
771 (byte 14)
772 (byte 13)
773 (byte 12)
774 attentionHeadLayoutSM86WaitLoads)
775 (attentionHeadLayoutSM86Next
776 (attentionHeadLayoutSM86Store
777 (byte 18)
778 (byte 14)
779 attentionHeadLayoutSM86U0)
780 (attentionHeadLayoutSM86Next
781 (constructor SM86InstructionBody SM86Exit sm86SafeControl)
782 attentionHeadLayoutSM86End)))))))))))))))
783
784-- Join the eight fp32 head-major context matrices [head][position][64] into
785-- one fp16 token-major matrix [position][512]. A 256-thread block owns one
786-- position; each thread converts one adjacent component pair.
787def attentionHeadLayoutSM86HeadToTokenProgram : (family SM86Program) =
788 (attentionHeadLayoutSM86Prologue
789 (byte 5)
790 attentionHeadLayoutSM86U31
791 (attentionHeadLayoutSM86Next
792 (attentionHeadLayoutSM86IMAD (byte 6) (byte 3) attentionHeadLayoutSM86U2 (byte 7))
793 (attentionHeadLayoutSM86Next
794 (attentionHeadLayoutSM86IMAD (byte 8) (byte 2) attentionHeadLayoutSM86U65536 (byte 7))
795 (attentionHeadLayoutSM86Next
796 (attentionHeadLayoutSM86IMAD (byte 8) (byte 1) attentionHeadLayoutSM86U64 (byte 8))
797 (attentionHeadLayoutSM86Next
798 (attentionHeadLayoutSM86IMAD (byte 8) (byte 6) attentionHeadLayoutSM86U1 (byte 8))
799 (attentionHeadLayoutSM86Next
800 (attentionHeadLayoutSM86Wide
801 (byte 10)
802 (byte 8)
803 (byte 4)
804 attentionHeadLayoutSM86InputPointer)
805 (attentionHeadLayoutSM86Next
806 (attentionHeadLayoutSM86Load (byte 12) (byte 10) attentionHeadLayoutSM86U0)
807 (attentionHeadLayoutSM86Next
808 (attentionHeadLayoutSM86Load (byte 13) (byte 10) attentionHeadLayoutSM86U4)
809 (attentionHeadLayoutSM86Next
810 (attentionHeadLayoutSM86IMAD
811 (byte 16)
812 (byte 1)
813 attentionHeadLayoutSM86U512
814 (byte 7))
815 (attentionHeadLayoutSM86Next
816 (attentionHeadLayoutSM86IMAD
817 (byte 16)
818 (byte 2)
819 attentionHeadLayoutSM86U64
820 (byte 16))
821 (attentionHeadLayoutSM86Next
822 (attentionHeadLayoutSM86IMAD
823 (byte 16)
824 (byte 6)
825 attentionHeadLayoutSM86U1
826 (byte 16))
827 (attentionHeadLayoutSM86Next
828 (attentionHeadLayoutSM86Wide
829 (byte 18)
830 (byte 16)
831 (byte 9)
832 attentionHeadLayoutSM86FirstOutputPointer)
833 (attentionHeadLayoutSM86Next
834 (attentionHeadLayoutSM86Pack
835 (byte 14)
836 (byte 13)
837 (byte 12)
838 attentionHeadLayoutSM86WaitLoads)
839 (attentionHeadLayoutSM86Next
840 (attentionHeadLayoutSM86Store
841 (byte 18)
842 (byte 14)
843 attentionHeadLayoutSM86U0)
844 (attentionHeadLayoutSM86Next
845 (constructor SM86InstructionBody SM86Exit sm86SafeControl)
846 attentionHeadLayoutSM86End)))))))))))))))
847
848-- Reverse the preceding layout for the attention-output adjoint: fp32
849-- token-major [position][512] to fp16 head-major [head][position][64].
850def attentionHeadLayoutSM86TokenToHeadProgram : (family SM86Program) =
851 (attentionHeadLayoutSM86Prologue
852 (byte 5)
853 attentionHeadLayoutSM86U31
854 (attentionHeadLayoutSM86Next
855 (attentionHeadLayoutSM86IMAD (byte 6) (byte 3) attentionHeadLayoutSM86U2 (byte 7))
856 (attentionHeadLayoutSM86Next
857 (attentionHeadLayoutSM86IMAD (byte 8) (byte 1) attentionHeadLayoutSM86U512 (byte 7))
858 (attentionHeadLayoutSM86Next
859 (attentionHeadLayoutSM86IMAD (byte 8) (byte 2) attentionHeadLayoutSM86U64 (byte 8))
860 (attentionHeadLayoutSM86Next
861 (attentionHeadLayoutSM86IMAD (byte 8) (byte 6) attentionHeadLayoutSM86U1 (byte 8))
862 (attentionHeadLayoutSM86Next
863 (attentionHeadLayoutSM86Wide
864 (byte 10)
865 (byte 8)
866 (byte 4)
867 attentionHeadLayoutSM86InputPointer)
868 (attentionHeadLayoutSM86Next
869 (attentionHeadLayoutSM86Load (byte 12) (byte 10) attentionHeadLayoutSM86U0)
870 (attentionHeadLayoutSM86Next
871 (attentionHeadLayoutSM86Load (byte 13) (byte 10) attentionHeadLayoutSM86U4)
872 (attentionHeadLayoutSM86Next
873 (attentionHeadLayoutSM86IMAD
874 (byte 16)
875 (byte 2)
876 attentionHeadLayoutSM86U65536
877 (byte 7))
878 (attentionHeadLayoutSM86Next
879 (attentionHeadLayoutSM86IMAD
880 (byte 16)
881 (byte 1)
882 attentionHeadLayoutSM86U64
883 (byte 16))
884 (attentionHeadLayoutSM86Next
885 (attentionHeadLayoutSM86IMAD
886 (byte 16)
887 (byte 6)
888 attentionHeadLayoutSM86U1
889 (byte 16))
890 (attentionHeadLayoutSM86Next
891 (attentionHeadLayoutSM86Wide
892 (byte 18)
893 (byte 16)
894 (byte 9)
895 attentionHeadLayoutSM86FirstOutputPointer)
896 (attentionHeadLayoutSM86Next
897 (attentionHeadLayoutSM86Pack
898 (byte 14)
899 (byte 13)
900 (byte 12)
901 attentionHeadLayoutSM86WaitLoads)
902 (attentionHeadLayoutSM86Next
903 (attentionHeadLayoutSM86Store
904 (byte 18)
905 (byte 14)
906 attentionHeadLayoutSM86U0)
907 (attentionHeadLayoutSM86Next
908 (constructor SM86InstructionBody SM86Exit sm86SafeControl)
909 attentionHeadLayoutSM86End)))))))))))))))
910
911-- The attention adjoint produces dQ and dK as fp32
912-- [head][position][component], and dV transposed as fp32
913-- [head][component][position]. Apply the inverse rotary transform to Q/K,
914-- merge all three planes into the projection's token-major QKV order, and
915-- pack the result to fp16 for the QKV input/weight GEMMs. Pointer ABI:
916-- 0x160 output QKV, 0x168 dQ, 0x170 dK, 0x178 cos, 0x180 sin, 0x188 dV^T.
917def attentionHeadLayoutSM86InverseRoPEQKVMergeProgram : (family SM86Program) =
918 (attentionHeadLayoutSM86Prologue
919 (byte 4)
920 attentionHeadLayoutSM86U15
921 (attentionHeadLayoutSM86Next
922 (attentionHeadLayoutSM86IMAD (byte 6) (byte 3) attentionHeadLayoutSM86U2 (byte 7))
923 (attentionHeadLayoutSM86Next
924 (attentionHeadLayoutSM86IMAD (byte 8) (byte 2) attentionHeadLayoutSM86U65536 (byte 7))
925 (attentionHeadLayoutSM86Next
926 (attentionHeadLayoutSM86IMAD (byte 8) (byte 1) attentionHeadLayoutSM86U64 (byte 8))
927 (attentionHeadLayoutSM86Next
928 (attentionHeadLayoutSM86IMAD (byte 8) (byte 6) attentionHeadLayoutSM86U1 (byte 8))
929 (attentionHeadLayoutSM86Next
930 (attentionHeadLayoutSM86Wide
931 (byte 10)
932 (byte 8)
933 (byte 4)
934 attentionHeadLayoutSM86SecondOutputPointer)
935 (attentionHeadLayoutSM86Next
936 (attentionHeadLayoutSM86Load (byte 12) (byte 10) attentionHeadLayoutSM86U0)
937 (attentionHeadLayoutSM86Next
938 (attentionHeadLayoutSM86Load (byte 13) (byte 10) attentionHeadLayoutSM86U4)
939 (attentionHeadLayoutSM86Next
940 (attentionHeadLayoutSM86Load (byte 14) (byte 10) attentionHeadLayoutSM86U128)
941 (attentionHeadLayoutSM86Next
942 (attentionHeadLayoutSM86Load (byte 15) (byte 10) attentionHeadLayoutSM86U132)
943 (attentionHeadLayoutSM86Next
944 (attentionHeadLayoutSM86Wide
945 (byte 16)
946 (byte 8)
947 (byte 4)
948 attentionHeadLayoutSM86InputPointer)
949 (attentionHeadLayoutSM86Next
950 (attentionHeadLayoutSM86Load
951 (byte 18)
952 (byte 16)
953 attentionHeadLayoutSM86U0)
954 (attentionHeadLayoutSM86Next
955 (attentionHeadLayoutSM86Load
956 (byte 19)
957 (byte 16)
958 attentionHeadLayoutSM86U4)
959 (attentionHeadLayoutSM86Next
960 (attentionHeadLayoutSM86Load
961 (byte 20)
962 (byte 16)
963 attentionHeadLayoutSM86U128)
964 (attentionHeadLayoutSM86Next
965 (attentionHeadLayoutSM86Load
966 (byte 21)
967 (byte 16)
968 attentionHeadLayoutSM86U132)
969 (attentionHeadLayoutSM86Next
970 (attentionHeadLayoutSM86IMAD
971 (byte 22)
972 (byte 1)
973 attentionHeadLayoutSM86U32
974 (byte 6))
975 (attentionHeadLayoutSM86Next
976 (attentionHeadLayoutSM86Wide
977 (byte 24)
978 (byte 22)
979 (byte 4)
980 attentionHeadLayoutSM86CosinePointer)
981 (attentionHeadLayoutSM86Next
982 (attentionHeadLayoutSM86Load
983 (byte 26)
984 (byte 24)
985 attentionHeadLayoutSM86U0)
986 (attentionHeadLayoutSM86Next
987 (attentionHeadLayoutSM86Load
988 (byte 27)
989 (byte 24)
990 attentionHeadLayoutSM86U4)
991 (attentionHeadLayoutSM86Next
992 (attentionHeadLayoutSM86Wide
993 (byte 28)
994 (byte 22)
995 (byte 4)
996 attentionHeadLayoutSM86SinePointer)
997 (attentionHeadLayoutSM86Next
998 (attentionHeadLayoutSM86Load
999 (byte 30)
1000 (byte 28)
1001 attentionHeadLayoutSM86U0)
1002 (attentionHeadLayoutSM86Next
1003 (attentionHeadLayoutSM86Load
1004 (byte 31)
1005 (byte 28)
1006 attentionHeadLayoutSM86U4)
1007 (attentionHeadLayoutSM86Next
1008 (attentionHeadLayoutSM86NegateWithControl
1009 (byte 30)
1010 attentionHeadLayoutSM86WaitLoads)
1011 (attentionHeadLayoutSM86Next
1012 (attentionHeadLayoutSM86Negate (byte 31))
1013 (attentionHeadLayoutSM86Rotate
1014 (byte 32)
1015 (byte 33)
1016 (byte 12)
1017 (byte 14)
1018 (byte 26)
1019 (byte 30)
1020 attentionHeadLayoutSM86WaitLoads
1021 (attentionHeadLayoutSM86Rotate
1022 (byte 34)
1023 (byte 35)
1024 (byte 13)
1025 (byte 15)
1026 (byte 27)
1027 (byte 31)
1028 sm86SafeControl
1029 (attentionHeadLayoutSM86Next
1030 (attentionHeadLayoutSM86Pack
1031 (byte 36)
1032 (byte 34)
1033 (byte 32)
1034 sm86SafeControl)
1035 (attentionHeadLayoutSM86Next
1036 (attentionHeadLayoutSM86Pack
1037 (byte 37)
1038 (byte 35)
1039 (byte 33)
1040 sm86SafeControl)
1041 (attentionHeadLayoutSM86Rotate
1042 (byte 38)
1043 (byte 39)
1044 (byte 18)
1045 (byte 20)
1046 (byte 26)
1047 (byte 30)
1048 sm86SafeControl
1049 (attentionHeadLayoutSM86Rotate
1050 (byte 40)
1051 (byte 41)
1052 (byte 19)
1053 (byte 21)
1054 (byte 27)
1055 (byte 31)
1056 sm86SafeControl
1057 (attentionHeadLayoutSM86Next
1058 (attentionHeadLayoutSM86Pack
1059 (byte 42)
1060 (byte 40)
1061 (byte 38)
1062 sm86SafeControl)
1063 (attentionHeadLayoutSM86Next
1064 (attentionHeadLayoutSM86Pack
1065 (byte 43)
1066 (byte 41)
1067 (byte 39)
1068 sm86SafeControl)
1069 (attentionHeadLayoutSM86Next
1070 (attentionHeadLayoutSM86IMAD
1071 (byte 44)
1072 (byte 1)
1073 attentionHeadLayoutSM86U1536
1074 (byte 7))
1075 (attentionHeadLayoutSM86Next
1076 (attentionHeadLayoutSM86IMAD
1077 (byte 44)
1078 (byte 2)
1079 attentionHeadLayoutSM86U64
1080 (byte 44))
1081 (attentionHeadLayoutSM86Next
1082 (attentionHeadLayoutSM86IMAD
1083 (byte 44)
1084 (byte 6)
1085 attentionHeadLayoutSM86U1
1086 (byte 44))
1087 (attentionHeadLayoutSM86Next
1088 (attentionHeadLayoutSM86Wide
1089 (byte 46)
1090 (byte 44)
1091 (byte 9)
1092 attentionHeadLayoutSM86FirstOutputPointer)
1093 (attentionHeadLayoutSM86Next
1094 (attentionHeadLayoutSM86Store
1095 (byte 46)
1096 (byte 36)
1097 attentionHeadLayoutSM86U0)
1098 (attentionHeadLayoutSM86Next
1099 (attentionHeadLayoutSM86Store
1100 (byte 46)
1101 (byte 37)
1102 attentionHeadLayoutSM86U64)
1103 (attentionHeadLayoutSM86Next
1104 (attentionHeadLayoutSM86Store
1105 (byte 46)
1106 (byte 42)
1107 attentionHeadLayoutSM86U1024)
1108 (attentionHeadLayoutSM86Next
1109 (attentionHeadLayoutSM86Store
1110 (byte 46)
1111 (byte 43)
1112 attentionHeadLayoutSM86U1088)
1113 (attentionHeadLayoutSM86Next
1114 (attentionHeadLayoutSM86IMAD
1115 (byte 48)
1116 (byte 2)
1117 attentionHeadLayoutSM86U65536
1118 (byte 7))
1119 (attentionHeadLayoutSM86Next
1120 (attentionHeadLayoutSM86IMAD
1121 (byte 48)
1122 (byte 3)
1123 attentionHeadLayoutSM86U2048
1124 (byte 48))
1125 (attentionHeadLayoutSM86Next
1126 (attentionHeadLayoutSM86IMAD
1127 (byte 48)
1128 (byte 1)
1129 attentionHeadLayoutSM86U1
1130 (byte 48))
1131 (attentionHeadLayoutSM86Next
1132 (attentionHeadLayoutSM86Wide
1133 (byte 50)
1134 (byte 48)
1135 (byte 4)
1136 attentionHeadLayoutSM86ValueGradientPointer)
1137 (attentionHeadLayoutSM86Next
1138 (attentionHeadLayoutSM86Load
1139 (byte 52)
1140 (byte 50)
1141 attentionHeadLayoutSM86U0)
1142 (attentionHeadLayoutSM86Next
1143 (attentionHeadLayoutSM86Load
1144 (byte 53)
1145 (byte 50)
1146 attentionHeadLayoutSM86U4096)
1147 (attentionHeadLayoutSM86Next
1148 (attentionHeadLayoutSM86Load
1149 (byte 55)
1150 (byte 50)
1151 attentionHeadLayoutSM86U131072)
1152 (attentionHeadLayoutSM86Next
1153 (attentionHeadLayoutSM86Load
1154 (byte 56)
1155 (byte 50)
1156 attentionHeadLayoutSM86U135168)
1157 (attentionHeadLayoutSM86Next
1158 (attentionHeadLayoutSM86Pack
1159 (byte 54)
1160 (byte 53)
1161 (byte 52)
1162 attentionHeadLayoutSM86WaitLoads)
1163 (attentionHeadLayoutSM86Next
1164 (attentionHeadLayoutSM86Pack
1165 (byte 58)
1166 (byte 56)
1167 (byte 55)
1168 sm86SafeControl)
1169 (attentionHeadLayoutSM86Next
1170 (attentionHeadLayoutSM86Store
1171 (byte 46)
1172 (byte 54)
1173 attentionHeadLayoutSM86U2048)
1174 (attentionHeadLayoutSM86Next
1175 (attentionHeadLayoutSM86Store
1176 (byte 46)
1177 (byte 58)
1178 attentionHeadLayoutSM86U2112)
1179 (attentionHeadLayoutSM86Next
1180 (constructor SM86InstructionBody SM86Exit sm86SafeControl)
1181 attentionHeadLayoutSM86End)))))))))))))))))))))))))))))))))))))))))))))))))))))
1182
1183def attentionHeadLayoutSM86Program =
1184 (lambda unrestricted kind : (family AttentionHeadLayoutSM86Kind) .
1185 (eliminate
1186 AttentionHeadLayoutSM86Kind
1187 (lambda unrestricted current : (family AttentionHeadLayoutSM86Kind) . (family SM86Program))
1188 kind
1189 (branch AttentionHeadLayoutQKRotary . attentionHeadLayoutSM86QKRotaryProgram)
1190 (branch AttentionHeadLayoutValueTranspose . attentionHeadLayoutSM86ValueTransposeProgram)
1191 (branch AttentionHeadLayoutHeadToToken . attentionHeadLayoutSM86HeadToTokenProgram)
1192 (branch AttentionHeadLayoutTokenToHead . attentionHeadLayoutSM86TokenToHeadProgram)
1193 (branch
1194 AttentionHeadLayoutInverseRoPEQKVMerge
1195 .
1196 attentionHeadLayoutSM86InverseRoPEQKVMergeProgram)))
1197
1198def attentionHeadLayoutSM86ExpectedInstructions =
1199 (lambda unrestricted kind : (family AttentionHeadLayoutSM86Kind) .
1200 (eliminate
1201 AttentionHeadLayoutSM86Kind
1202 (lambda unrestricted current : (family AttentionHeadLayoutSM86Kind) . Nat)
1203 kind
1204 (branch AttentionHeadLayoutQKRotary . attentionHeadLayoutSM86N63)
1205 (branch AttentionHeadLayoutValueTranspose . attentionHeadLayoutSM86N23)
1206 (branch AttentionHeadLayoutHeadToToken . attentionHeadLayoutSM86N23)
1207 (branch AttentionHeadLayoutTokenToHead . attentionHeadLayoutSM86N23)
1208 (branch AttentionHeadLayoutInverseRoPEQKVMerge . attentionHeadLayoutSM86N77)))
1209
1210def attentionHeadLayoutSM86Registers =
1211 (lambda unrestricted kind : (family AttentionHeadLayoutSM86Kind) .
1212 (eliminate
1213 AttentionHeadLayoutSM86Kind
1214 (lambda unrestricted current : (family AttentionHeadLayoutSM86Kind) . Nat)
1215 kind
1216 (branch AttentionHeadLayoutQKRotary . attentionHeadLayoutSM86N48)
1217 (branch AttentionHeadLayoutValueTranspose . attentionHeadLayoutSM86N24)
1218 (branch AttentionHeadLayoutHeadToToken . attentionHeadLayoutSM86N24)
1219 (branch AttentionHeadLayoutTokenToHead . attentionHeadLayoutSM86N24)
1220 (branch AttentionHeadLayoutInverseRoPEQKVMerge . attentionHeadLayoutSM86N64)))
1221
1222def attentionHeadLayoutSM86GridX =
1223 (lambda unrestricted kind : (family AttentionHeadLayoutSM86Kind) .
1224 (eliminate
1225 AttentionHeadLayoutSM86Kind
1226 (lambda unrestricted current : (family AttentionHeadLayoutSM86Kind) . Nat)
1227 kind
1228 (branch AttentionHeadLayoutQKRotary . attentionHeadLayoutSM86N1024)
1229 (branch AttentionHeadLayoutValueTranspose . attentionHeadLayoutSM86N512)
1230 (branch AttentionHeadLayoutHeadToToken . attentionHeadLayoutSM86N1024)
1231 (branch AttentionHeadLayoutTokenToHead . attentionHeadLayoutSM86N1024)
1232 (branch AttentionHeadLayoutInverseRoPEQKVMerge . attentionHeadLayoutSM86N1024)))
1233
1234def attentionHeadLayoutSM86BlockX =
1235 (lambda unrestricted kind : (family AttentionHeadLayoutSM86Kind) .
1236 (eliminate
1237 AttentionHeadLayoutSM86Kind
1238 (lambda unrestricted current : (family AttentionHeadLayoutSM86Kind) . Nat)
1239 kind
1240 (branch AttentionHeadLayoutQKRotary . attentionHeadLayoutSM86N128)
1241 (branch AttentionHeadLayoutValueTranspose . attentionHeadLayoutSM86N512)
1242 (branch AttentionHeadLayoutHeadToToken . attentionHeadLayoutSM86N256)
1243 (branch AttentionHeadLayoutTokenToHead . attentionHeadLayoutSM86N256)
1244 (branch AttentionHeadLayoutInverseRoPEQKVMerge . attentionHeadLayoutSM86N128)))
1245
1246def attentionHeadLayoutSM86FailureCodeBytes =
1247 (lambda unrestricted code : (family AttentionHeadLayoutSM86FailureCode) .
1248 (eliminate
1249 AttentionHeadLayoutSM86FailureCode
1250 (lambda unrestricted current : (family AttentionHeadLayoutSM86FailureCode) . Bytes)
1251 code
1252 (branch AttentionHeadLayoutInstructionCountMismatch . b"ALPHA-SM86-ATTN-LAYOUT-001")
1253 (branch AttentionHeadLayoutEncodingFailed . b"ALPHA-SM86-ATTN-LAYOUT-002")
1254 (branch AttentionHeadLayoutIdentityInvalid . b"ALPHA-SM86-ATTN-LAYOUT-003")))
1255
1256def attentionHeadLayoutSM86TelemetryFor =
1257 (lambda unrestricted kind : (family AttentionHeadLayoutSM86Kind) .
1258 (lambda unrestricted actual : Nat .
1259 (lambda unrestricted bytes : Nat .
1260 (lambda unrestricted fields : Nat .
1261 (lambda unrestricted bits : Nat .
1262 (constructor
1263 AttentionHeadLayoutSM86Telemetry
1264 AttentionHeadLayoutSM86TelemetryValue
1265 kind
1266 (attentionHeadLayoutSM86ExpectedInstructions kind)
1267 actual
1268 (attentionHeadLayoutSM86Registers kind)
1269 (attentionHeadLayoutSM86GridX kind)
1270 (attentionHeadLayoutSM86BlockX kind)
1271 attentionHeadLayoutSM86Extents
1272 bytes
1273 fields
1274 bits))))))
1275
1276def attentionHeadLayoutSM86Image =
1277 (lambda unrestricted kind : (family AttentionHeadLayoutSM86Kind) .
1278 (app
1279 (lambda unrestricted program : (family SM86Program) .
1280 (app
1281 (lambda unrestricted actual : Nat .
1282 (nat-eliminate
1283 (lambda unrestricted countValid : Nat . (family AttentionHeadLayoutSM86ImageResult))
1284 (constructor
1285 AttentionHeadLayoutSM86ImageResult
1286 AttentionHeadLayoutSM86ImageFailed
1287 (constructor
1288 AttentionHeadLayoutSM86FailureCode
1289 AttentionHeadLayoutInstructionCountMismatch)
1290 actual
1291 (attentionHeadLayoutSM86FailureCodeBytes
1292 (constructor
1293 AttentionHeadLayoutSM86FailureCode
1294 AttentionHeadLayoutInstructionCountMismatch))
1295 (attentionHeadLayoutSM86TelemetryFor kind actual zero zero zero))
1296 (lambda unrestricted countPredecessor : Nat .
1297 (lambda unrestricted countInduction : (family AttentionHeadLayoutSM86ImageResult) .
1298 (eliminate
1299 SM86ProgramEncodingResult
1300 (lambda unrestricted current : (family SM86ProgramEncodingResult) .
1301 (family AttentionHeadLayoutSM86ImageResult))
1302 (sm86EncodeProgram program)
1303 (branch
1304 SM86ProgramEncodingSucceeded
1305 image
1306 encodingTelemetry
1307 .
1308 (eliminate
1309 SM86ProgramEncodingTelemetry
1310 (lambda unrestricted current : (family SM86ProgramEncodingTelemetry) .
1311 (family AttentionHeadLayoutSM86ImageResult))
1312 encodingTelemetry
1313 (branch
1314 SM86ProgramEncodingTelemetryValue
1315 instructions
1316 bytes
1317 fields
1318 bits
1319 highest
1320 .
1321 (app
1322 (lambda unrestricted identity : Bytes .
1323 (nat-eliminate
1324 (lambda unrestricted valid : Nat .
1325 (family AttentionHeadLayoutSM86ImageResult))
1326 (constructor
1327 AttentionHeadLayoutSM86ImageResult
1328 AttentionHeadLayoutSM86ImageFailed
1329 (constructor
1330 AttentionHeadLayoutSM86FailureCode
1331 AttentionHeadLayoutIdentityInvalid)
1332 zero
1333 (attentionHeadLayoutSM86FailureCodeBytes
1334 (constructor
1335 AttentionHeadLayoutSM86FailureCode
1336 AttentionHeadLayoutIdentityInvalid))
1337 (attentionHeadLayoutSM86TelemetryFor
1338 kind
1339 instructions
1340 bytes
1341 fields
1342 bits))
1343 (lambda unrestricted validPredecessor : Nat .
1344 (lambda unrestricted validInduction : (family AttentionHeadLayoutSM86ImageResult) .
1345 (constructor
1346 AttentionHeadLayoutSM86ImageResult
1347 AttentionHeadLayoutSM86ImageReady
1348 kind
1349 image
1350 identity
1351 (attentionHeadLayoutSM86TelemetryFor
1352 kind
1353 instructions
1354 bytes
1355 fields
1356 bits))))
1357 (naturalEqual (bytes-length identity) attentionHeadLayoutSM86N64)))
1358 (sha256HexBytesOrEmpty (sha256Hex image))))))
1359 (branch
1360 SM86ProgramEncodingFailed
1361 index
1362 failure
1363 encodingTelemetry
1364 .
1365 (constructor
1366 AttentionHeadLayoutSM86ImageResult
1367 AttentionHeadLayoutSM86ImageFailed
1368 (constructor
1369 AttentionHeadLayoutSM86FailureCode
1370 AttentionHeadLayoutEncodingFailed)
1371 index
1372 (sm86InstructionEncodingStableCode failure)
1373 (attentionHeadLayoutSM86TelemetryFor kind actual zero zero zero))))))
1374 (naturalEqual actual (attentionHeadLayoutSM86ExpectedInstructions kind))))
1375 (sm86ProgramCount program)))
1376 (attentionHeadLayoutSM86Program kind)))
1377
1378-- One half-split rotation plane from a token-major f32 projection to a
1379-- head-major packed-f16 attention operand. Grid (sequence/16, heads, 1),
1380-- block 256: each group of 16 threads owns one position, with each thread
1381-- owning adjacent components in both 32-wide halves. Query
1382-- and key use the same program with different projection offsets, so K can
1383-- have fewer heads without materializing repeated K.
1384def attentionHeadLayoutSM86RotatePlaneAdmitted =
1385 (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
1386 (lambda unrestricted projectionWidth : Nat . (lambda unrestricted planeOffset : Nat .
1387 (naturalAnd (naturalNonzero seq)
1388 (naturalAnd (naturalEqual (naturalModuloUnchecked seq 16) 0)
1389 (naturalAnd (naturalNonzero heads)
1390 (naturalAnd (naturalLessOrEqual (naturalAdd planeOffset (naturalMultiply heads 64)) projectionWidth)
1391 (naturalAnd (naturalLessOrEqual (naturalMultiply seq (naturalMultiply projectionWidth 4)) 4294967295)
1392 (naturalLessOrEqual (naturalMultiply heads (naturalMultiply seq 128)) 4294967295))))))))))
1393
1394def attentionHeadLayoutSM86RotatePlaneWord =
1395 (lambda unrestricted value : Nat . (sm86Unsigned32FromNaturalTruncated value))
1396
1397def attentionHeadLayoutSM86RotatePlaneSpecial =
1398 (lambda unrestricted destination : Byte .
1399 (lambda unrestricted source : (family SM86SpecialRegister) .
1400 (lambda unrestricted tail : (family SM86Program) .
1401 (attentionHeadLayoutSM86Next
1402 (constructor SM86InstructionBody SM86SpecialToRegister
1403 (attentionHeadLayoutSM86R destination) source attentionHeadLayoutSM86SetSpecials)
1404 tail))))
1405
1406-- R8 is the projection element, R20 the table element, R36 the output
1407-- element. They are derived from the shape and the CTA coordinates.
1408def attentionHeadLayoutSM86RotatePlaneEntry =
1409 (lambda unrestricted seq : Nat . (lambda unrestricted projectionWidth : Nat .
1410 (lambda unrestricted planeOffset : Nat . (lambda unrestricted tail : (family SM86Program) .
1411 (attentionHeadLayoutSM86RotatePlaneSpecial (byte 0)
1412 (constructor SM86SpecialRegister SM86ThreadIdX)
1413 (attentionHeadLayoutSM86RotatePlaneSpecial (byte 1)
1414 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdX)
1415 (attentionHeadLayoutSM86RotatePlaneSpecial (byte 2)
1416 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdY)
1417 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Move (byte 4) attentionHeadLayoutSM86U4)
1418 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Move (byte 7) attentionHeadLayoutSM86U0)
1419 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Move (byte 9) attentionHeadLayoutSM86U2)
1420 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Move (byte 15) attentionHeadLayoutSM86U15)
1421 (attentionHeadLayoutSM86Next
1422 (constructor SM86InstructionBody SM86MoveImmediate
1423 (attentionHeadLayoutSM86R (byte 5)) attentionHeadLayoutSM86U0 attentionHeadLayoutSM86WaitSpecials)
1424 (attentionHeadLayoutSM86Next
1425 (constructor SM86InstructionBody SM86ShiftRightImmediate
1426 (attentionHeadLayoutSM86R (byte 6)) (attentionHeadLayoutSM86R (byte 0)) (byte 4) sm86SafeControl)
1427 (attentionHeadLayoutSM86Next
1428 (constructor SM86InstructionBody SM86LogicThreeInputTruthTable
1429 (attentionHeadLayoutSM86R (byte 3)) (attentionHeadLayoutSM86R (byte 0))
1430 (attentionHeadLayoutSM86R (byte 15)) (byte 192) sm86SafeControl)
1431 (attentionHeadLayoutSM86Next
1432 (attentionHeadLayoutSM86IMAD (byte 1) (byte 1)
1433 (attentionHeadLayoutSM86RotatePlaneWord 16) (byte 6))
1434 (attentionHeadLayoutSM86Next
1435 (attentionHeadLayoutSM86IMAD (byte 3) (byte 3) attentionHeadLayoutSM86U2 (byte 7))
1436 (attentionHeadLayoutSM86Next
1437 (attentionHeadLayoutSM86Move (byte 6) (attentionHeadLayoutSM86RotatePlaneWord planeOffset))
1438 (attentionHeadLayoutSM86Next
1439 (attentionHeadLayoutSM86IMAD (byte 8) (byte 1) (attentionHeadLayoutSM86RotatePlaneWord projectionWidth) (byte 6))
1440 (attentionHeadLayoutSM86Next
1441 (attentionHeadLayoutSM86IMAD (byte 8) (byte 2) attentionHeadLayoutSM86U64 (byte 8))
1442 (attentionHeadLayoutSM86Next
1443 (attentionHeadLayoutSM86IMAD (byte 8) (byte 3) attentionHeadLayoutSM86U1 (byte 8))
1444 (attentionHeadLayoutSM86Next
1445 (attentionHeadLayoutSM86IMAD (byte 20) (byte 1) attentionHeadLayoutSM86U32 (byte 3))
1446 (attentionHeadLayoutSM86Next
1447 (attentionHeadLayoutSM86IMAD (byte 36) (byte 1) attentionHeadLayoutSM86U64 (byte 3))
1448 (attentionHeadLayoutSM86Next
1449 (attentionHeadLayoutSM86IMAD (byte 36) (byte 2)
1450 (attentionHeadLayoutSM86RotatePlaneWord (naturalMultiply seq 64)) (byte 36))
1451 tail)))))))))))))))))))))))
1452
1453def attentionHeadLayoutSM86RotatePlaneRead =
1454 (lambda unrestricted inputIndex : Byte . (lambda unrestricted tail : (family SM86Program) .
1455 (attentionHeadLayoutSM86Next
1456 (attentionHeadLayoutSM86Wide (byte 10) inputIndex (byte 4) attentionHeadLayoutSM86InputPointer)
1457 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Load (byte 12) (byte 10) attentionHeadLayoutSM86U0)
1458 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Load (byte 13) (byte 10) attentionHeadLayoutSM86U4)
1459 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Load (byte 14) (byte 10) attentionHeadLayoutSM86U128)
1460 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Load (byte 15) (byte 10) attentionHeadLayoutSM86U132)
1461 (attentionHeadLayoutSM86Next
1462 (attentionHeadLayoutSM86Wide (byte 22) (byte 20) (byte 4) attentionHeadLayoutSM86CosinePointer)
1463 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Load (byte 24) (byte 22) attentionHeadLayoutSM86U0)
1464 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Load (byte 25) (byte 22) attentionHeadLayoutSM86U4)
1465 (attentionHeadLayoutSM86Next
1466 (attentionHeadLayoutSM86Wide (byte 26) (byte 20) (byte 4) attentionHeadLayoutSM86SinePointer)
1467 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Load (byte 28) (byte 26) attentionHeadLayoutSM86U0)
1468 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Load (byte 29) (byte 26) attentionHeadLayoutSM86U4)
1469 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Move (byte 5) attentionHeadLayoutSM86U0)
1470 tail))))))))))))))
1471
1472def attentionHeadLayoutSM86RotatePlaneWrite =
1473 (lambda unrestricted tail : (family SM86Program) .
1474 (attentionHeadLayoutSM86Rotate (byte 30) (byte 31) (byte 12) (byte 14)
1475 (byte 24) (byte 28) attentionHeadLayoutSM86WaitLoads
1476 (attentionHeadLayoutSM86Rotate (byte 32) (byte 33) (byte 13) (byte 15)
1477 (byte 25) (byte 29) sm86SafeControl
1478 (attentionHeadLayoutSM86Next
1479 (attentionHeadLayoutSM86Pack (byte 34) (byte 32) (byte 30) sm86SafeControl)
1480 (attentionHeadLayoutSM86Next
1481 (attentionHeadLayoutSM86Pack (byte 35) (byte 33) (byte 31) sm86SafeControl)
1482 (attentionHeadLayoutSM86Next
1483 (attentionHeadLayoutSM86Wide (byte 38) (byte 36) (byte 9) attentionHeadLayoutSM86FirstOutputPointer)
1484 (attentionHeadLayoutSM86Next
1485 (attentionHeadLayoutSM86Store (byte 38) (byte 34) attentionHeadLayoutSM86U0)
1486 (attentionHeadLayoutSM86Next
1487 (attentionHeadLayoutSM86Store (byte 38) (byte 35) attentionHeadLayoutSM86U64)
1488 tail))))))))
1489
1490def attentionHeadLayoutSM86RotatePlaneUnchecked =
1491 (lambda unrestricted seq : Nat . (lambda unrestricted projectionWidth : Nat .
1492 (lambda unrestricted planeOffset : Nat .
1493 (attentionHeadLayoutSM86RotatePlaneEntry seq projectionWidth planeOffset
1494 (attentionHeadLayoutSM86RotatePlaneRead (byte 8)
1495 (attentionHeadLayoutSM86RotatePlaneWrite
1496 (attentionHeadLayoutSM86Next
1497 (constructor SM86InstructionBody SM86Exit sm86SafeControl)
1498 attentionHeadLayoutSM86End)))))))
1499
1500def attentionHeadLayoutSM86RotatePlane =
1501 (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
1502 (lambda unrestricted projectionWidth : Nat . (lambda unrestricted planeOffset : Nat .
1503 (lambda erased admitted :
1504 (equal Nat (attentionHeadLayoutSM86RotatePlaneAdmitted seq heads projectionWidth planeOffset) 1) .
1505 (attentionHeadLayoutSM86RotatePlaneUnchecked seq projectionWidth planeOffset))))))
1506
1507def attentionHeadLayoutSM86RotatePlaneRegisters : Nat = 48
1508def attentionHeadLayoutSM86RotatePlaneThreads : Nat = 256
1509def attentionHeadLayoutSM86RotatePlaneConstantBytes : Nat = 0x188
1510def attentionHeadLayoutSM86RotatePlaneBlocks =
1511 (lambda unrestricted seq : Nat . (naturalDivideUnchecked seq 16))
1512
1513-- The pullback uses the transpose of the forward rotation: negating sine
1514-- lets the same checked two-component rotation express it. Its source is a
1515-- head-major f32 gradient plane; the destination is the matching slice of
1516-- the token-major projection gradient. Q and K writes are disjoint.
1517def attentionHeadLayoutSM86InverseRotatePlaneAdmitted =
1518 (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
1519 (lambda unrestricted projectionWidth : Nat . (lambda unrestricted planeOffset : Nat .
1520 (naturalAnd
1521 (attentionHeadLayoutSM86RotatePlaneAdmitted seq heads projectionWidth planeOffset)
1522 (naturalLessOrEqual (naturalMultiply heads (naturalMultiply seq 256)) 4294967295))))))
1523
1524def attentionHeadLayoutSM86InverseRotatePlaneWrite =
1525 (lambda unrestricted tail : (family SM86Program) .
1526 (attentionHeadLayoutSM86Next
1527 (attentionHeadLayoutSM86NegateWithControl (byte 28) attentionHeadLayoutSM86WaitLoads)
1528 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Negate (byte 29))
1529 (attentionHeadLayoutSM86Rotate (byte 30) (byte 31) (byte 12) (byte 14)
1530 (byte 24) (byte 28) sm86SafeControl
1531 (attentionHeadLayoutSM86Rotate (byte 32) (byte 33) (byte 13) (byte 15)
1532 (byte 25) (byte 29) sm86SafeControl
1533 (attentionHeadLayoutSM86Next
1534 (attentionHeadLayoutSM86Wide (byte 38) (byte 8) (byte 4) attentionHeadLayoutSM86FirstOutputPointer)
1535 (attentionHeadLayoutSM86Next
1536 (attentionHeadLayoutSM86Store (byte 38) (byte 30) attentionHeadLayoutSM86U0)
1537 (attentionHeadLayoutSM86Next
1538 (attentionHeadLayoutSM86Store (byte 38) (byte 32) attentionHeadLayoutSM86U4)
1539 (attentionHeadLayoutSM86Next
1540 (attentionHeadLayoutSM86Store (byte 38) (byte 31) attentionHeadLayoutSM86U128)
1541 (attentionHeadLayoutSM86Next
1542 (attentionHeadLayoutSM86Store (byte 38) (byte 33) attentionHeadLayoutSM86U132)
1543 tail))))))))))
1544
1545def attentionHeadLayoutSM86InverseRotatePlaneUnchecked =
1546 (lambda unrestricted seq : Nat . (lambda unrestricted projectionWidth : Nat .
1547 (lambda unrestricted planeOffset : Nat .
1548 (attentionHeadLayoutSM86RotatePlaneEntry seq projectionWidth planeOffset
1549 (attentionHeadLayoutSM86RotatePlaneRead (byte 36)
1550 (attentionHeadLayoutSM86InverseRotatePlaneWrite
1551 (attentionHeadLayoutSM86Next
1552 (constructor SM86InstructionBody SM86Exit sm86SafeControl)
1553 attentionHeadLayoutSM86End)))))))
1554
1555def attentionHeadLayoutSM86InverseRotatePlane =
1556 (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
1557 (lambda unrestricted projectionWidth : Nat . (lambda unrestricted planeOffset : Nat .
1558 (lambda erased admitted :
1559 (equal Nat (attentionHeadLayoutSM86InverseRotatePlaneAdmitted seq heads projectionWidth planeOffset) 1) .
1560 (attentionHeadLayoutSM86InverseRotatePlaneUnchecked seq projectionWidth planeOffset))))))
1561
1562def attentionHeadLayoutSM86InverseRotatePlaneRegisters : Nat = 48
1563def attentionHeadLayoutSM86InverseRotatePlaneThreads : Nat = 256
1564def attentionHeadLayoutSM86InverseRotatePlaneConstantBytes : Nat = 0x188
1565
1566-- Transpose one token-major f32 value plane to head-major [head][component]
1567-- [position] packed BF16/f16. Grid (sequence/2, heads, 1), block 64. Each
1568-- thread packs two adjacent positions of one component into a 32-bit store.
1569def attentionHeadLayoutSM86TransposeValueAdmitted =
1570 (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
1571 (lambda unrestricted projectionWidth : Nat . (lambda unrestricted planeOffset : Nat .
1572 (naturalAnd (naturalNonzero seq)
1573 (naturalAnd (naturalEqual (naturalModuloUnchecked seq 2) 0)
1574 (naturalAnd (naturalNonzero heads)
1575 (naturalAnd (naturalLessOrEqual (naturalAdd planeOffset (naturalMultiply heads 64)) projectionWidth)
1576 (naturalAnd (naturalLessOrEqual (naturalMultiply seq (naturalMultiply projectionWidth 4)) 4294967295)
1577 (naturalLessOrEqual (naturalMultiply heads (naturalMultiply seq 128)) 4294967295))))))))))
1578
1579def attentionHeadLayoutSM86TransposeValueUnchecked =
1580 (lambda unrestricted seq : Nat . (lambda unrestricted projectionWidth : Nat .
1581 (lambda unrestricted planeOffset : Nat .
1582 (attentionHeadLayoutSM86RotatePlaneSpecial (byte 0)
1583 (constructor SM86SpecialRegister SM86ThreadIdX)
1584 (attentionHeadLayoutSM86RotatePlaneSpecial (byte 1)
1585 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdX)
1586 (attentionHeadLayoutSM86RotatePlaneSpecial (byte 2)
1587 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdY)
1588 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Move (byte 4) attentionHeadLayoutSM86U4)
1589 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Move (byte 7) attentionHeadLayoutSM86U0)
1590 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Move (byte 9) attentionHeadLayoutSM86U2)
1591 (attentionHeadLayoutSM86Next
1592 (constructor SM86InstructionBody SM86MoveImmediate
1593 (attentionHeadLayoutSM86R (byte 5)) attentionHeadLayoutSM86U0 attentionHeadLayoutSM86WaitSpecials)
1594 (attentionHeadLayoutSM86Next
1595 (attentionHeadLayoutSM86Move (byte 6) (attentionHeadLayoutSM86RotatePlaneWord planeOffset))
1596 (attentionHeadLayoutSM86Next
1597 (attentionHeadLayoutSM86IMAD (byte 3) (byte 1) attentionHeadLayoutSM86U2 (byte 7))
1598 (attentionHeadLayoutSM86Next
1599 (attentionHeadLayoutSM86IMAD (byte 8) (byte 3)
1600 (attentionHeadLayoutSM86RotatePlaneWord projectionWidth) (byte 6))
1601 (attentionHeadLayoutSM86Next
1602 (attentionHeadLayoutSM86IMAD (byte 8) (byte 2) attentionHeadLayoutSM86U64 (byte 8))
1603 (attentionHeadLayoutSM86Next
1604 (attentionHeadLayoutSM86IMAD (byte 8) (byte 0) attentionHeadLayoutSM86U1 (byte 8))
1605 (attentionHeadLayoutSM86Next
1606 (attentionHeadLayoutSM86Wide (byte 10) (byte 8) (byte 4) attentionHeadLayoutSM86InputPointer)
1607 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Load (byte 12) (byte 10) attentionHeadLayoutSM86U0)
1608 (attentionHeadLayoutSM86Next
1609 (attentionHeadLayoutSM86Load (byte 13) (byte 10)
1610 (attentionHeadLayoutSM86RotatePlaneWord (naturalMultiply projectionWidth 4)))
1611 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Move (byte 5) attentionHeadLayoutSM86U0)
1612 (attentionHeadLayoutSM86Next
1613 (attentionHeadLayoutSM86Pack (byte 14) (byte 13) (byte 12) attentionHeadLayoutSM86WaitLoads)
1614 (attentionHeadLayoutSM86Next
1615 (attentionHeadLayoutSM86IMAD (byte 16) (byte 0)
1616 (attentionHeadLayoutSM86RotatePlaneWord seq) (byte 3))
1617 (attentionHeadLayoutSM86Next
1618 (attentionHeadLayoutSM86IMAD (byte 16) (byte 2)
1619 (attentionHeadLayoutSM86RotatePlaneWord (naturalMultiply seq 64)) (byte 16))
1620 (attentionHeadLayoutSM86Next
1621 (attentionHeadLayoutSM86Wide (byte 18) (byte 16) (byte 9) attentionHeadLayoutSM86FirstOutputPointer)
1622 (attentionHeadLayoutSM86Next
1623 (attentionHeadLayoutSM86Store (byte 18) (byte 14) attentionHeadLayoutSM86U0)
1624 (attentionHeadLayoutSM86Next
1625 (constructor SM86InstructionBody SM86Exit sm86SafeControl)
1626 attentionHeadLayoutSM86End)))))))))))))))))))))))))
1627
1628def attentionHeadLayoutSM86TransposeValue =
1629 (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
1630 (lambda unrestricted projectionWidth : Nat . (lambda unrestricted planeOffset : Nat .
1631 (lambda erased admitted :
1632 (equal Nat (attentionHeadLayoutSM86TransposeValueAdmitted seq heads projectionWidth planeOffset) 1) .
1633 (attentionHeadLayoutSM86TransposeValueUnchecked seq projectionWidth planeOffset))))))
1634
1635def attentionHeadLayoutSM86TransposeValueRegisters : Nat = 24
1636def attentionHeadLayoutSM86TransposeValueThreads : Nat = 64
1637def attentionHeadLayoutSM86TransposeValueConstantBytes : Nat = 0x178
1638def attentionHeadLayoutSM86TransposeValueBlocks =
1639 (lambda unrestricted seq : Nat . (naturalDivideUnchecked seq 2))
1640
1641-- Pull back the transposed V operand into its token-major projection slice.
1642-- The incoming dV^T is f32 [head][component][position]; no rotation applies
1643-- to V. The two positions owned by one thread have disjoint output words.
1644def attentionHeadLayoutSM86InverseTransposeValueAdmitted =
1645 (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
1646 (lambda unrestricted projectionWidth : Nat . (lambda unrestricted planeOffset : Nat .
1647 (naturalAnd
1648 (attentionHeadLayoutSM86TransposeValueAdmitted seq heads projectionWidth planeOffset)
1649 (naturalLessOrEqual (naturalMultiply heads (naturalMultiply seq 256)) 4294967295))))))
1650
1651def attentionHeadLayoutSM86InverseTransposeValueUnchecked =
1652 (lambda unrestricted seq : Nat . (lambda unrestricted projectionWidth : Nat .
1653 (lambda unrestricted planeOffset : Nat .
1654 (attentionHeadLayoutSM86RotatePlaneSpecial (byte 0)
1655 (constructor SM86SpecialRegister SM86ThreadIdX)
1656 (attentionHeadLayoutSM86RotatePlaneSpecial (byte 1)
1657 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdX)
1658 (attentionHeadLayoutSM86RotatePlaneSpecial (byte 2)
1659 (constructor SM86SpecialRegister SM86CooperativeThreadArrayIdY)
1660 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Move (byte 4) attentionHeadLayoutSM86U4)
1661 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Move (byte 7) attentionHeadLayoutSM86U0)
1662 (attentionHeadLayoutSM86Next
1663 (constructor SM86InstructionBody SM86MoveImmediate
1664 (attentionHeadLayoutSM86R (byte 5)) attentionHeadLayoutSM86U0 attentionHeadLayoutSM86WaitSpecials)
1665 (attentionHeadLayoutSM86Next
1666 (attentionHeadLayoutSM86Move (byte 6) (attentionHeadLayoutSM86RotatePlaneWord planeOffset))
1667 (attentionHeadLayoutSM86Next
1668 (attentionHeadLayoutSM86IMAD (byte 3) (byte 1) attentionHeadLayoutSM86U2 (byte 7))
1669 (attentionHeadLayoutSM86Next
1670 (attentionHeadLayoutSM86IMAD (byte 16) (byte 0)
1671 (attentionHeadLayoutSM86RotatePlaneWord seq) (byte 3))
1672 (attentionHeadLayoutSM86Next
1673 (attentionHeadLayoutSM86IMAD (byte 16) (byte 2)
1674 (attentionHeadLayoutSM86RotatePlaneWord (naturalMultiply seq 64)) (byte 16))
1675 (attentionHeadLayoutSM86Next
1676 (attentionHeadLayoutSM86Wide (byte 10) (byte 16) (byte 4) attentionHeadLayoutSM86InputPointer)
1677 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Load (byte 12) (byte 10) attentionHeadLayoutSM86U0)
1678 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Load (byte 13) (byte 10) attentionHeadLayoutSM86U4)
1679 (attentionHeadLayoutSM86Next (attentionHeadLayoutSM86Move (byte 5) attentionHeadLayoutSM86U0)
1680 (attentionHeadLayoutSM86Next
1681 (constructor SM86InstructionBody SM86MoveImmediate
1682 (attentionHeadLayoutSM86R (byte 5)) attentionHeadLayoutSM86U0 attentionHeadLayoutSM86WaitLoads)
1683 (attentionHeadLayoutSM86Next
1684 (attentionHeadLayoutSM86IMAD (byte 8) (byte 3)
1685 (attentionHeadLayoutSM86RotatePlaneWord projectionWidth) (byte 6))
1686 (attentionHeadLayoutSM86Next
1687 (attentionHeadLayoutSM86IMAD (byte 8) (byte 2) attentionHeadLayoutSM86U64 (byte 8))
1688 (attentionHeadLayoutSM86Next
1689 (attentionHeadLayoutSM86IMAD (byte 8) (byte 0) attentionHeadLayoutSM86U1 (byte 8))
1690 (attentionHeadLayoutSM86Next
1691 (attentionHeadLayoutSM86Wide (byte 18) (byte 8) (byte 4) attentionHeadLayoutSM86FirstOutputPointer)
1692 (attentionHeadLayoutSM86Next
1693 (attentionHeadLayoutSM86Store (byte 18) (byte 12) attentionHeadLayoutSM86U0)
1694 (attentionHeadLayoutSM86Next
1695 (attentionHeadLayoutSM86Store (byte 18) (byte 13)
1696 (attentionHeadLayoutSM86RotatePlaneWord (naturalMultiply projectionWidth 4)))
1697 (attentionHeadLayoutSM86Next
1698 (constructor SM86InstructionBody SM86Exit sm86SafeControl)
1699 attentionHeadLayoutSM86End)))))))))))))))))))))))))
1700
1701def attentionHeadLayoutSM86InverseTransposeValue =
1702 (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat .
1703 (lambda unrestricted projectionWidth : Nat . (lambda unrestricted planeOffset : Nat .
1704 (lambda erased admitted :
1705 (equal Nat (attentionHeadLayoutSM86InverseTransposeValueAdmitted seq heads projectionWidth planeOffset) 1) .
1706 (attentionHeadLayoutSM86InverseTransposeValueUnchecked seq projectionWidth planeOffset))))))
1707
1708def attentionHeadLayoutSM86InverseTransposeValueRegisters : Nat = 24
1709def attentionHeadLayoutSM86InverseTransposeValueThreads : Nat = 64
1710def attentionHeadLayoutSM86InverseTransposeValueConstantBytes : Nat = 0x178The compiler supplied declaration spans and resolved links from this source snapshot. This page does not assert that this file belongs to a checked closure.