perf(metal): recorded ICB fans global-layer SDPA over the 2-pass pair — deep decode un-collapsed

The arch ICB recorded MLX's single-pass sdpa_vector once at session open: one
threadgroup per head serially streaming the WHOLE cache. The re-encode lane's
sdpa2PassMinKV switch cannot happen inside a recorded command buffer, so decode
at depth collapsed far below bandwidth physics (E2B @52k: 44 tok/s measured vs
~115 expected from the ~6% attention-read overhead; E4B worse — more globals).

GLOBAL layers now record the 2-pass pair: sdpa_vector_2pass_1 (blocks baked as
function constant 26, supportIndirectCommandBuffers; N bound via the rebindable
nGlobalBuf; K/V at binds 1/2 exactly as the single-pass op so the replay's rebind
indices are untouched) + the 2pass_2 merge. Sliding layers keep single-pass —
their window bounds n below the knee. blocks = sdpa2PassBlocks(maxLen), fixed at
record time and safe at ANY smaller n: a block whose strided key walk starts past
N writes finite_min/0 partials the merge zeroes. Gated on maxLen >= 1024 so short
sessions — and every synthetic byte-parity fixture — keep the pure single-pass
layout. Shared per-replay partials/sums/maxs + scalarI32(blocks) are registered
resident (an ICB op reading a non-resident buffer is undefined — the #343 class).

New plumbing: emitSDPA2Pass1NAt (N-from-buffer variant of the shared emitter),
sdpaVector2Pass1PipelineICB / sdpaVector2Pass2PipelineICB.

Receipts (E2B 4bit, tg128 temp0 think, quiet GPU): 6.6K 126.4 -> 149.7, 26K 71.0
-> 136.7, 52K 44.0 -> 114.4 (the physics prediction), 97K 25.7 -> 87.7 (+241%).
Zero-depth unregressed: E2B 175.2, 26B MoE 146.0. Deep-session (131072) tokens
byte-identical to the single-pass receipt and deterministic across reps; ICB
parity suites green; full engine/metal suite 1436 green.

Co-Authored-By: Virgil <virgil@lethean.io>

Author: Snider

go/engine/metal/decode_forward_arch_icb.go

@@ -1210,6 +1210,43 @@ func recordArchICB(
 			sdpaPSOByHd[hd] = pso
 		}
 	}
+	// Deep-decode: GLOBAL layers record the 2-pass SDPA pair instead of the single-pass
+	// kernel. The single-pass sdpa_vector runs ONE threadgroup per head over the whole
+	// cache — recorded once, it can never re-parallelise as the KV grows, which collapsed
+	// deep decode far below bandwidth physics (E2B @52K: 44 tok/s measured vs ~115
+	// expected). blocks is fixed at record time from maxLen and is safe at ANY smaller n:
+	// a block whose strided key walk starts past N writes finite_min/0 partials that the
+	// pass-2 merge zeroes. Sliding layers stay single-pass — their window bounds n below
+	// the 2-pass knee. Sessions too short to ever cross the knee keep the pure
+	// single-pass layout (and the existing byte-parity fixtures with it).
+	sdpa2PassICBBlocks := 0
+	if maxLen >= sdpa2PassMinKV {
+		sdpa2PassICBBlocks = int(sdpa2PassBlocks(maxLen))
+	}
+	sdpa2Pass1PSOByHd := make(map[int]metal.MTLComputePipelineState)
+	sdpa2Pass2PSOByHd := make(map[int]metal.MTLComputePipelineState)
+	nGlobal2Pass := 0
+	if sdpa2PassICBBlocks > 0 {
+		for li := range nLayers {
+			if specs[li].Attention != model.GlobalAttention {
+				continue
+			}
+			nGlobal2Pass++
+			hd := hdOf(li)
+			if _, ok := sdpa2Pass1PSOByHd[hd]; !ok {
+				p1, e := sdpaVector2Pass1PipelineICB(hd, int32(sdpa2PassICBBlocks))
+				if e != nil {
+					return nil, e
+				}
+				p2, e2 := sdpaVector2Pass2PipelineICB(hd)
+				if e2 != nil {
+					return nil, e2
+				}
+				sdpa2Pass1PSOByHd[hd] = p1
+				sdpa2Pass2PSOByHd[hd] = p2
+			}
+		}
+	}
 	addPSO, err := pipelineForICB("vv_Addbfloat16")
 	if err != nil {
 		return nil, err
@@ -1411,6 +1448,20 @@ func recordArchICB(
 			offBuf, nGlobalBuf, nSlidingBuf, epsBuf, axisBuf, wsBuf,
 			ropeScaleB, ropeBaseB, ropeLocalBaseB, freqStride1B, sdpaScaleB, addModelB,
 		}
+		// 2-pass SDPA intermediates for the GLOBAL layers (shared across layers — the replay's
+		// dependency barriers already serialise each layer's attention on the shared scratch,
+		// exactly as the single-row attn scratch). f32 per the kernel ABI; sized at the widest
+		// head dim. Owned by the replay via residentRes for the session's lifetime.
+		var p2Partials, p2Sums, p2Maxs metal.MTLBuffer
+		if nGlobal2Pass > 0 {
+			p2Partials = device.NewBufferWithLengthOptions(uint(nHeads*sdpa2PassICBBlocks*maxHd*4), metal.MTLResourceStorageModeShared)
+			p2Sums = device.NewBufferWithLengthOptions(uint(nHeads*sdpa2PassICBBlocks*4), metal.MTLResourceStorageModeShared)
+			p2Maxs = device.NewBufferWithLengthOptions(uint(nHeads*sdpa2PassICBBlocks*4), metal.MTLResourceStorageModeShared)
+			// pass-2 binds blocks via the memoised scalar — a value no other op declares, so
+			// register it resident explicitly (an ICB op reading a non-resident buffer is
+			// undefined; the strides/scale/N binds all reuse scalars already listed above).
+			resident = append(resident, p2Partials, p2Sums, p2Maxs, scalarI32(int32(sdpa2PassICBBlocks)))
+		}
 		if !hasFusedGELU {
 			resident = append(resident, x2, x3, x3s, inner, scaled, tnh, onePlus, halfG, gelu, c044, c079, c1c, c05)
 		}
@@ -1538,7 +1589,8 @@ func recordArchICB(
 				opsPerLayer--
 			}
 		}
-		total := opsPerLayer * nLayers
+		// GLOBAL layers' 2-pass SDPA is pass-1 + pass-2 where the single-pass was one op.
+		total := opsPerLayer*nLayers + nGlobal2Pass
 		icbDesc := metal.NewMTLIndirectCommandBufferDescriptor()
 		icbDesc.SetCommandTypes(metal.MTLIndirectCommandTypeConcurrentDispatch)
 		icbDesc.SetInheritBuffers(false)
@@ -1743,9 +1795,23 @@ func recordArchICB(
 			// gqaOf/sdpaStrideOf/sdpaScaleB hold. attendK read offset rebound/token if sliding.
 			hd, kv := hdOf(li), kvOf(li)
 			kvd := int64(kv * hd)
-			emitSDPA(fastICBSink{emit()}, sdpaPSOByHd[hd], qr, attendK, attendV, attn, 0, nBufForLayer,
-				nHeads, kv, 0, int64(hd), kvd, int64(hd), kvd, scale)
-			sdpaIdx[li] = opIdx - 1
+			if sdpa2PassICBBlocks > 0 && specs[li].Attention == model.GlobalAttention {
+				// GLOBAL layer deep-decode: the 2-pass pair fans the growing-cache reduction over
+				// blocks threadgroups (pass 1) and merges the partials (pass 2) — the recorded
+				// replacement for the single-pass kernel that serialised the whole cache on one
+				// threadgroup per head. N binds the same rebindable nGlobalBuf; K/V bind at slots
+				// 1/2 exactly as the single-pass op, so the replay's rebind indices are unchanged.
+				emitSDPA2Pass1NAt(fastICBSink{emit()}, sdpa2Pass1PSOByHd[hd], qr, 0, attendK, attendV,
+					p2Partials, p2Sums, p2Maxs, 0, nBufForLayer, 1, nHeads, kv, 0, sdpa2PassICBBlocks,
+					int64(hd), kvd, int64(hd), kvd, scale)
+				sdpaIdx[li] = opIdx - 1
+				emitSDPA2Pass2(fastICBSink{emit()}, sdpa2Pass2PSOByHd[hd], p2Partials, p2Sums, p2Maxs,
+					attn, 1, nHeads, sdpa2PassICBBlocks)
+			} else {
+				emitSDPA(fastICBSink{emit()}, sdpaPSOByHd[hd], qr, attendK, attendV, attn, 0, nBufForLayer,
+					nHeads, kv, 0, int64(hd), kvd, int64(hd), kvd, scale)
+				sdpaIdx[li] = opIdx - 1
+			}
 			recordProj(li, emit(), attn, attnOut, 0, projO)
 			if hasPA && useFusedResRMS { // fused: hBuf = inBuf + rms(Wo·attn) — one op, one fewer barrier
 				setRMSResidual(emit(), attnOut, postAttnBufs[li], inBuf, hBuf)
@@ -1846,7 +1912,7 @@ func recordArchICB(
 		// the recorded layout diverged from opsPerLayer·nLayers — a recorder bug, not a numeric
 		// drift; fail loud rather than replay a misaligned ICB.
 		if opIdx != total {
-			coreErr = core.NewError(core.Sprintf("native.decodeForwardArchICBCore: recorded %d ops, expected %d (opsPerLayer=%d × %d layers) — heterogeneous layout misaligned", opIdx, total, opsPerLayer, nLayers))
+			coreErr = core.NewError(core.Sprintf("native.decodeForwardArchICBCore: recorded %d ops, expected %d (opsPerLayer=%d × %d layers + %d global 2-pass) — heterogeneous layout misaligned", opIdx, total, opsPerLayer, nLayers, nGlobal2Pass))
 			return
 		}
 

go/engine/metal/dispatch_sink.go

@@ -385,20 +385,34 @@ func emitSDPAAt[S dispatchSink](sink S, pso metal.MTLComputePipelineState, q met
 // emitSDPA2Pass1 records the first long-context SDPA pass. It writes one partial
 // weighted-V sum plus online-softmax sum/max per (batch, kv-head, block).
 func emitSDPA2Pass1[S dispatchSink](sink S, pso metal.MTLComputePipelineState, q, k, v, partials, sums, maxs metal.MTLBuffer, kvByteOff uint, batch, nHeads, nKVHeads, n, blocks int, kHeadStride, kSeqStride, vHeadStride, vSeqStride int64, scale float32) {
-	emitSDPA2Pass1At(sink, pso, q, 0, k, v, partials, sums, maxs, kvByteOff, batch, nHeads, nKVHeads, n, blocks, kHeadStride, kSeqStride, vHeadStride, vSeqStride, scale)
+	emitSDPA2Pass1NAt(sink, pso, q, 0, k, v, partials, sums, maxs, kvByteOff, nil, batch, nHeads, nKVHeads, n, blocks, kHeadStride, kSeqStride, vHeadStride, vSeqStride, scale)
 }
 
 // emitSDPA2Pass1At is emitSDPA2Pass1 with the query bound at a byte offset (the attention fold's
 // slab rows). The partials/sums/maxs stay whole-buffer — they are per-dispatch scratch.
 func emitSDPA2Pass1At[S dispatchSink](sink S, pso metal.MTLComputePipelineState, q metal.MTLBuffer, qOff uint, k, v, partials, sums, maxs metal.MTLBuffer, kvByteOff uint, batch, nHeads, nKVHeads, n, blocks int, kHeadStride, kSeqStride, vHeadStride, vSeqStride int64, scale float32) {
+	emitSDPA2Pass1NAt(sink, pso, q, qOff, k, v, partials, sums, maxs, kvByteOff, nil, batch, nHeads, nKVHeads, n, blocks, kHeadStride, kSeqStride, vHeadStride, vSeqStride, scale)
+}
+
+// emitSDPA2Pass1NAt is emitSDPA2Pass1At with the attended length optionally taken from a
+// BUFFER: nBuf != nil binds it at index 7 (the recorded arch ICB rebinds that buffer per
+// token — N is the one per-token-varying scalar, exactly as emitSDPAAt's nBuf); nBuf == nil
+// inlines n (the live paths know the length). blocks stays baked in the PSO + grid, which is
+// safe at ANY n: a block whose strided key walk starts past N contributes finite_min/0
+// partials that the pass-2 merge zeroes out.
+func emitSDPA2Pass1NAt[S dispatchSink](sink S, pso metal.MTLComputePipelineState, q metal.MTLBuffer, qOff uint, k, v, partials, sums, maxs metal.MTLBuffer, kvByteOff uint, nBuf metal.MTLBuffer, batch, nHeads, nKVHeads, n, blocks int, kHeadStride, kSeqStride, vHeadStride, vSeqStride int64, scale float32) {
 	sink.setPSO(pso)
 	sink.setBuf(q, qOff, 0)
 	sink.setBuf(k, kvByteOff, 1)
 	sink.setBuf(v, kvByteOff, 2)
 	sink.setBuf(partials, 0, 3)
 	sink.setBuf(sums, 0, 4)
 	sink.setBuf(maxs, 0, 5)
-	sink.setI32(int32(n), 7)
+	if nBuf != nil {
+		sink.setBuf(nBuf, 0, 7) // ICB: the N buffer, rebound per token at replay
+	} else {
+		sink.setI32(int32(n), 7)
+	}
 	sink.setI64(kHeadStride, 8)
 	sink.setI64(kSeqStride, 9)
 	sink.setI64(vHeadStride, 10)

go/engine/metal/icb.go

@@ -463,6 +463,78 @@ func sdpaVectorPipelineICB(name string) (metal.MTLComputePipelineState, error) {
 	return pso, nil
 }
 
+// sdpaVector2Pass1PipelineICB builds the ICB-recordable sdpa_vector_2pass_1 — the same
+// function constants as the live builder (20-25 false: decode-time no mask/transpose/
+// causal/sinks; 26 = blocks, baked because the kernel strides its key walk and indexes
+// the partials by it) with supportIndirectCommandBuffers set. The recorded arch ICB
+// replays this pair for GLOBAL layers so deep decode keeps the fanned-out reduction the
+// single-pass kernel cannot give (one threadgroup per head serialising the whole cache
+// — the deep-decode collapse).
+func sdpaVector2Pass1PipelineICB(headDim int, blocks int32) (metal.MTLComputePipelineState, error) {
+	key := core.Sprintf("sdpa_vector_2pass_1_bfloat16_t_%d_%d|b%d|icb", headDim, headDim, blocks)
+	icbPSOMu.Lock()
+	defer icbPSOMu.Unlock()
+	if pso, ok := icbPSOCache[key]; ok {
+		return pso, nil
+	}
+	if library == nil || library.GetID() == 0 {
+		return nil, core.NewError("native.sdpaVector2Pass1PipelineICB: library unavailable")
+	}
+	fc := metal.NewMTLFunctionConstantValues()
+	off := uint8(0)
+	for _, idx := range []uint{20, 21, 22, 23, 24, 25} {
+		fc.SetConstantValueTypeAtIndex(unsafe.Pointer(&off), metal.MTLDataTypeBool, idx)
+	}
+	blk := blocks
+	fc.SetConstantValueTypeAtIndex(unsafe.Pointer(&blk), metal.MTLDataTypeInt, 26)
+	name := core.Sprintf("sdpa_vector_2pass_1_bfloat16_t_%d_%d", headDim, headDim)
+	fn, err := library.NewFunctionWithNameConstantValuesError(name, fc)
+	if err != nil {
+		return nil, core.E("native.sdpaVector2Pass1PipelineICB", name, err)
+	}
+	if fn == nil || fn.GetID() == 0 {
+		return nil, core.NewError("native.sdpaVector2Pass1PipelineICB: kernel " + name + " not found")
+	}
+	desc := metal.NewMTLComputePipelineDescriptor()
+	desc.SetComputeFunction(fn)
+	desc.SetSupportIndirectCommandBuffers(true)
+	pso, err := device.NewComputePipelineStateWithDescriptorOptionsReflectionError(desc, 0, nil)
+	if err != nil {
+		return nil, core.E("native.sdpaVector2Pass1PipelineICB", "pipeline "+name, err)
+	}
+	icbPSOCache[key] = pso
+	return pso, nil
+}
+
+// sdpaVector2Pass2PipelineICB builds the ICB-recordable sdpa_vector_2pass_2 combine
+// kernel (plain function, no constants — blocks arrives inlined at bind 4) with
+// supportIndirectCommandBuffers set.
+func sdpaVector2Pass2PipelineICB(headDim int) (metal.MTLComputePipelineState, error) {
+	name := core.Sprintf("sdpa_vector_2pass_2_bfloat16_t_%d", headDim)
+	key := name + "|icb"
+	icbPSOMu.Lock()
+	defer icbPSOMu.Unlock()
+	if pso, ok := icbPSOCache[key]; ok {
+		return pso, nil
+	}
+	if library == nil || library.GetID() == 0 {
+		return nil, core.NewError("native.sdpaVector2Pass2PipelineICB: library unavailable")
+	}
+	fn := library.NewFunctionWithName(name)
+	if fn == nil || fn.GetID() == 0 {
+		return nil, core.NewError("native.sdpaVector2Pass2PipelineICB: kernel " + name + " not found")
+	}
+	desc := metal.NewMTLComputePipelineDescriptor()
+	desc.SetComputeFunction(fn)
+	desc.SetSupportIndirectCommandBuffers(true)
+	pso, err := device.NewComputePipelineStateWithDescriptorOptionsReflectionError(desc, 0, nil)
+	if err != nil {
+		return nil, core.E("native.sdpaVector2Pass2PipelineICB", "pipeline "+name, err)
+	}
+	icbPSOCache[key] = pso
+	return pso, nil
+}
+
 func sdpaVectorPipelineICBForHeadDim(headDim int) (metal.MTLComputePipelineState, error) {
 	icbPSOMu.Lock()
 	if pso, ok := sdpaVectorICBHeadDimPSOCache[headDim]; ok {