perf(metal): deep-prompt global-layer SDPA rides a steel GEMM composition

The multiQ vector kernel re-reads the whole K/V stream once per query row.
At depth that traffic outgrows the SLC and the measured per-key cost ramps
0.032µs → 0.065µs (the #345 knee); even its floor runs the prompt n² at a
small fraction of the machine's matmul rate — mlx-lm's gemma4 prompt
attention (the same GEMM chain below, headDim 256 exceeding MLX's fused
steel_attention) prefilled 63K tokens 11× faster.

Past sdpaPromptGEMMMinKV (4096) a global layer's chunk SDPA now runs as
three dispatches per GQA group: S = Q @ Kᵀ (steel nt, batched over the
group's q-heads on grid.z with scalar batch strides, B stride 0 sharing
the group's K), an in-place causal row softmax (new
lthn_softmax_causal_rows_bf16 — per-row valid-prefix cap identical to the
multiQ kernel's, masked tail written zero for the following GEMM), and
O = P @ V (steel nn, batched). K/V are read once per kv head instead of
once per query row. Sliding layers and short prompts keep their existing
kernels; the S slabs double-buffer across kv heads and track their OWN
capacities (the #343 lesson).

Token-identity tier — S stores bf16 between the GEMMs, the same boundary
the fold's qmm and ≥32-row steel projections already trade at; pinned by
the per-row cosine closeness test vs the multiQ kernel.

Receipts (gemma-4-e2b-4bit, ctx 131072, temp 0, quiet GPU):
  63K prefill  95.8s → 29.4s (3.3×) · needle HIT · byte-deterministic
  97K prefill 276.0s → 54.7s (5.0×) · needle HIT · decode after 87.9 tok/s
                                      (= the #344 receipt, decode untouched)
  per-key cost flat 0.0107µs at every depth — the SLC-decay ramp is gone
  26B-A4B 6.6K: 8.7s → 7.2s, needle HIT (quant/MoE arch rides the same lane)
  engine/metal suite 1439 green · module suite 9057 green

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

Author: Snider

go/engine/metal/decode_batched_session.go

@@ -82,6 +82,13 @@ type denseBatchScratch struct {
 	layerVStage      []metal.MTLBuffer
 	layerStageRowCap int
 	layerStageKVCap  int
+	// prompt-attention GEMM score slabs (sdpa_prompt_gemm.go): two K × nCap buffers
+	// double-buffered across heads so head h+1's wide QKᵀ GEMM overlaps head h's skinny
+	// P@V GEMM instead of draining the GPU behind its 32 output tiles.
+	sdpaS0      metal.MTLBuffer
+	sdpaS1      metal.MTLBuffer
+	sdpaSRowCap int // sdpaPromptS's OWN capacities — never another fold's (the attnRowCap lesson)
+	sdpaSNCap   int
 	// moeBatch holds the K-row MoE fold slabs (moe_batch.go) — nil until an MoE chunk runs.
 	moeBatch *moeBatchScratch
 }
@@ -139,6 +146,23 @@ func (s *denseBatchScratch) layerStage(li, layers, k, kvDimMax int) (kSt, vSt me
 	return s.layerKStage[li], s.layerVStage[li]
 }
 
+// sdpaPromptS returns the two prompt-attention score slabs (kRows × nCap bf16 each) for the
+// GEMM SDPA composition, (re)allocating when the chunk row count grows or the attended-length
+// cap rises. nCap is the session's maxLen so the allocation happens ONCE per session instead
+// of every deepening chunk. Growth is tracked by sdpaPromptS's OWN capacity fields — never
+// another fold's (the attnRowCap lesson: capacity consumed outside its owner left slabs short
+// at the wide tail-absorbed chunk).
+func (s *denseBatchScratch) sdpaPromptS(kRows, nCap int) (s0, s1 metal.MTLBuffer) {
+	if s.sdpaS0 == nil || s.sdpaSRowCap < kRows || s.sdpaSNCap < nCap {
+		rows := max(kRows, s.sdpaSRowCap)
+		cols := max(nCap, s.sdpaSNCap)
+		s.sdpaS0 = scratchBF16(rows * cols)
+		s.sdpaS1 = scratchBF16(rows * cols)
+		s.sdpaSRowCap, s.sdpaSNCap = rows, cols
+	}
+	return s.sdpaS0, s.sdpaS1
+}
+
 func (s *denseBatchScratch) Close() {
 	if s == nil {
 		return
@@ -890,6 +914,15 @@ func (s *archDecodeState) stepTokensBatchedDenseResultWithInputViewsPLE(embs [][
 			// (the per-row oracle would have routed 2-pass there), same tier as the fold's qmm.
 			useMultiQ := !sdpaMultiQDisabledForTest && (slideW == 0 || basePos+K <= slideW) &&
 				(basePos+K < sdpa2PassMinKV || K >= steelGEMMMinRows) && gpuHasSDPAMultiQ(lhd)
+			// Deep-prompt global layers route the SDPA to the steel GEMM composition
+			// (sdpa_prompt_gemm.go): K/V read once per HEAD instead of once per query row,
+			// so the traffic no longer scales with the row count and the multiQ kernel's
+			// deep-context SLC-decay ramp never engages. Same emission seam as multiQ —
+			// the per-row loop still runs its staged/rope tail, only the SDPA dispatches
+			// differ. Token-identity tier (S stores bf16 between the GEMMs), the same
+			// boundary the fold's qmm and ≥32-row steel projections already trade at.
+			useGEMMSDPA := useMultiQ && slideW == 0 && basePos+K >= sdpaPromptGEMMMinKV &&
+				K <= sdpaPromptGEMMMaxRows && !sdpaPromptGEMMDisabledForTest && gpuHasPromptSDPAGEMM()
 			if batchedRope {
 				if err = encQKNormRopeRows(enc, qSlab, s.lb[li].qNorm.buf, qSlab, 0, s.lb[li].qNorm.off, 0, qDim, qDim, offBuf[0], layerRopeFreqs, K, s.nHeads, lhd, rotDim, rbase, s.scale, s.eps); err != nil {
 					endEncodingFast(enc)
@@ -980,6 +1013,13 @@ func (s *archDecodeState) stepTokensBatchedDenseResultWithInputViewsPLE(embs [][
 					stagedDeferred[li] = true
 					pendingLandings = append(pendingLandings, ringLanding{li: li, kvDim: kvDim, slideW: slideW})
 				}
+			} else if useGEMMSDPA {
+				sScore0, sScore1 := s.denseBatch.sdpaPromptS((s.nHeads/lkv)*K, s.maxLen)
+				if err = encSDPAPromptGEMM(enc, qSlab, ownerK, ownerV, attnSlab, sScore0, sScore1,
+					s.nHeads, lkv, lhd, K, basePos+K, qDim, kvDim, s.scale); err != nil {
+					endEncodingFast(enc)
+					return nil, false, err
+				}
 			} else if useMultiQ {
 				if err = encSDPAMultiQCausal(enc, qSlab, ownerK, ownerV, attnSlab, s.nHeads, lkv, lhd, K, basePos+K,
 					int64(lhd), int64(kvDim), int64(lhd), int64(kvDim), s.scale); err != nil {

go/engine/metal/gemm_steel.go

@@ -51,7 +51,7 @@ type steelGEMMParams struct {
 	_                      int32 // trailing pad to the struct's 8-byte alignment
 }
 
-type steelGEMMKey struct{ alignM, alignN, alignK bool }
+type steelGEMMKey struct{ transB, alignM, alignN, alignK bool }
 
 var (
 	steelGEMMMu       sync.Mutex
@@ -72,21 +72,32 @@ const (
 // each combo is its own PSO. has_batch/use_out_source/do_axpby are baked false — the batched pass
 // runs plain single-batch D = A @ Bᵀ.
 func steelGEMMPipeline(alignM, alignN, alignK bool) (metal.MTLComputePipelineState, bool) {
+	return steelGEMMPipelineTrans(true, alignM, alignN, alignK)
+}
+
+// steelGEMMPipelineTrans is steelGEMMPipeline with the B-transpose selectable: transB=true loads
+// the nt kernel (B read transposed — weights, K caches), false the nn kernel (B read as-is — the
+// prompt attention's P @ V). Same function-constant layout either way.
+func steelGEMMPipelineTrans(transB, alignM, alignN, alignK bool) (metal.MTLComputePipelineState, bool) {
 	steelGEMMMu.Lock()
 	defer steelGEMMMu.Unlock()
 	if steelGEMMBroken {
 		return nil, false
 	}
-	key := steelGEMMKey{alignM: alignM, alignN: alignN, alignK: alignK}
+	key := steelGEMMKey{transB: transB, alignM: alignM, alignN: alignN, alignK: alignK}
 	if pso, ok := steelGEMMPSOCache[key]; ok {
 		return pso, true
 	}
 	if library == nil || library.GetID() == 0 {
 		steelGEMMBroken = true
 		return nil, false
 	}
-	name := core.Sprintf("steel_gemm_fused_nt_bfloat16_bfloat16_bm%d_bn%d_bk%d_wm%d_wn%d",
-		steelGEMMBM, steelGEMMBN, steelGEMMBK, steelGEMMWM, steelGEMMWN)
+	variant := "nn"
+	if transB {
+		variant = "nt"
+	}
+	name := core.Sprintf("steel_gemm_fused_%s_bfloat16_bfloat16_bm%d_bn%d_bk%d_wm%d_wn%d",
+		variant, steelGEMMBM, steelGEMMBN, steelGEMMBK, steelGEMMWM, steelGEMMWN)
 	fc := metal.NewMTLFunctionConstantValues()
 	off := uint8(0)
 	for _, idx := range []uint{10, 100, 110} { // has_batch, use_out_source, do_axpby
@@ -116,38 +127,62 @@ func steelGEMMPipeline(alignM, alignN, alignK bool) (metal.MTLComputePipelineSta
 // (ldd = outDim). Reports false when the steel pipeline is unavailable (the caller keeps the
 // batched gemv).
 func encGemmBF16NT(enc metal.MTLComputeCommandEncoder, mat, vec, out metal.MTLBuffer, matOff, vecOff, outOff uint, outDim, inDim, rows int) bool {
-	pso, ok := steelGEMMPipeline(rows%steelGEMMBM == 0, outDim%steelGEMMBN == 0, inDim%steelGEMMBK == 0)
+	return encGemmBF16Strided(enc, true, vec, mat, out, vecOff, matOff, outOff,
+		rows, outDim, inDim, inDim, inDim, outDim)
+}
+
+// encGemmBF16Strided encodes D[m × n] = A[m × k] @ B (nt: B[n × k] read transposed; nn: B[k × n]
+// read as-is) as ONE steel GEMM with explicit leading dimensions — the general form behind
+// encGemmBF16NT. lda/ldb/ldd are element strides between consecutive rows of A/B/D, so a caller
+// can address one head's columns inside a wider slab (the prompt attention's Q/K/V/O views).
+// Reports false when the steel pipeline is unavailable (the caller keeps its fallback path).
+func encGemmBF16Strided(enc metal.MTLComputeCommandEncoder, transB bool, a, b, d metal.MTLBuffer,
+	aOff, bOff, dOff uint, m, n, k, lda, ldb, ldd int) bool {
+	return encGemmBF16StridedBatch(enc, transB, a, b, d, aOff, bOff, dOff, m, n, k, lda, ldb, ldd, 1, 0, 0, 0)
+}
+
+// encGemmBF16StridedBatch is encGemmBF16Strided with a uniform-stride batch on grid depth:
+// batch problems at element strides batchA/batchB/batchD (stride 0 broadcasts an operand —
+// the GQA group's shared K/V). The steel fused kernel applies the scalar batch strides off
+// tid.z even with has_batch baked false, so this rides the SAME PSO as the single dispatch.
+func encGemmBF16StridedBatch(enc metal.MTLComputeCommandEncoder, transB bool, a, b, d metal.MTLBuffer,
+	aOff, bOff, dOff uint, m, n, k, lda, ldb, ldd, batch int, batchA, batchB, batchD int64) bool {
+	pso, ok := steelGEMMPipelineTrans(transB, m%steelGEMMBM == 0, n%steelGEMMBN == 0, k%steelGEMMBK == 0)
 	if !ok {
 		return false
 	}
 	if pieceTimingOn {
 		steelGEMMDispatchesForTest++
 	}
-	tilesM := (rows + steelGEMMBM - 1) / steelGEMMBM
-	tilesN := (outDim + steelGEMMBN - 1) / steelGEMMBN
+	tilesM := (m + steelGEMMBM - 1) / steelGEMMBM
+	tilesN := (n + steelGEMMBN - 1) / steelGEMMBN
 	// threadblock swizzle (mlx matmul.cpp): interleave the tile walk so neighbouring threadgroups
 	// share B panels in L2 — 0 for short grids, 2 on this device class for tall ones.
 	swizzle := 0
 	if tilesM > 3 {
 		swizzle = 2
 	}
 	params := steelGEMMParams{
-		M: int32(rows), N: int32(outDim), K: int32(inDim),
-		LDA: int32(inDim), LDB: int32(inDim), LDD: int32(outDim),
+		M: int32(m), N: int32(n), K: int32(k),
+		LDA: int32(lda), LDB: int32(ldb), LDD: int32(ldd),
 		TilesN: int32(tilesN), TilesM: int32(tilesM),
-		SwizzleLog: int32(swizzle), GemmKIterationsAligned: int32(inDim / steelGEMMBK), BatchNDim: 1,
+		BatchStrideA: batchA, BatchStrideB: batchB, BatchStrideD: batchD,
+		SwizzleLog: int32(swizzle), GemmKIterationsAligned: int32(k / steelGEMMBK), BatchNDim: 1,
 	}
 	sink := encSink{enc}
 	sink.setPSO(pso)
-	sink.setBuf(vec, vecOff, 0) // A: the activation rows
-	sink.setBuf(mat, matOff, 1) // B: the weight, read transposed (nt)
-	sink.setBuf(out, outOff, 3) // D
+	sink.setBuf(a, aOff, 0)
+	sink.setBuf(b, bOff, 1)
+	sink.setBuf(d, dOff, 3)
 	setBytes(enc, unsafe.Pointer(&params), uint(unsafe.Sizeof(params)), 4)
 	tile := 1 << swizzle
 	gridX := tilesN * tile
 	gridY := (tilesM + tile - 1) / tile
+	if batch < 1 {
+		batch = 1
+	}
 	sink.dispatchThreadgroups(
-		metal.MTLSize{Width: uint(gridX), Height: uint(gridY), Depth: 1},
+		metal.MTLSize{Width: uint(gridX), Height: uint(gridY), Depth: uint(batch)},
 		metal.MTLSize{Width: 32, Height: steelGEMMWN, Depth: steelGEMMWM},
 	)
 	return true

go/engine/metal/kernels/lthn_softmax_causal.metal

@@ -0,0 +1,86 @@
+// SPDX-Licence-Identifier: EUPL-1.2
+
+#include <metal_stdlib>
+using namespace metal;
+
+// lthn_softmax_causal_rows_bf16 — in-place causal row softmax over the prompt attention's
+// score slab S[rows × N] (bf16, rows packed at ldd = N). One threadgroup per slab row. The
+// grid may stack several heads' K-row score blocks ([head][K][N], the batched-GEMM layout),
+// so the query index within the chunk is row % K.
+//
+// Query s attends keys [0 .. N-K+s] — the same per-row cap the multiQ vector kernel applies
+// (query s uses key i iff i <= N - K + s), expressed as a valid-prefix length so causality
+// needs no mask storage. Entries at or beyond the cap are written 0.0 so the following
+// P @ V steel GEMM reads zero weight from the masked tail.
+//
+// Maths is f32 throughout (scores load bf16 → f32, scale applied at read — equivalent to the
+// vector kernel scaling q before the dot). Three sweeps: max → exp+sum (exp stored bf16
+// in place) → normalise. The stored-P double rounding (unnormalised exp then normalise) is
+// bounded by bf16's 2⁻⁸ relative step — the same tier as the score slab itself.
+kernel void lthn_softmax_causal_rows_bf16(
+    device bfloat* S [[buffer(0)]],
+    const constant int& N [[buffer(1)]],
+    const constant int& K [[buffer(2)]],
+    const constant float& scale [[buffer(3)]],
+    uint row [[threadgroup_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint lsize [[threads_per_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  device bfloat* rowPtr = S + size_t(row) * size_t(N);
+  const int s = int(row) % K;    // query index within the chunk (grid may stack heads)
+  const int valid = N - K + s + 1; // keys [0, valid)
+  const int simds = int(lsize) / 32;
+
+  threadgroup float tgRed[32];
+  threadgroup float tgOut[2];
+
+  // sweep 1: row max over the valid prefix (scaled scores)
+  float m = -MAXFLOAT;
+  for (int i = int(lid); i < valid; i += int(lsize)) {
+    m = max(m, float(rowPtr[i]) * scale);
+  }
+  m = simd_max(m);
+  if (simd_lid == 0) {
+    tgRed[simd_gid] = m;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  m = (int(lid) < simds) ? tgRed[lid] : -MAXFLOAT;
+  m = simd_max(m);
+  if (lid == 0) {
+    tgOut[0] = m;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  m = tgOut[0];
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // sweep 2: exponentiate in place + accumulate the sum; zero the masked tail
+  float sum = 0.0f;
+  for (int i = int(lid); i < N; i += int(lsize)) {
+    if (i < valid) {
+      float p = fast::exp(float(rowPtr[i]) * scale - m);
+      rowPtr[i] = bfloat(p);
+      sum += p;
+    } else {
+      rowPtr[i] = bfloat(0.0f);
+    }
+  }
+  sum = simd_sum(sum);
+  if (simd_lid == 0) {
+    tgRed[simd_gid] = sum;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  sum = (int(lid) < simds) ? tgRed[lid] : 0.0f;
+  sum = simd_sum(sum);
+  if (lid == 0) {
+    tgOut[1] = sum;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  sum = tgOut[1];
+
+  // sweep 3: normalise the valid prefix
+  const float inv = sum > 0.0f ? (1.0f / sum) : 0.0f;
+  for (int i = int(lid); i < valid; i += int(lsize)) {
+    rowPtr[i] = bfloat(float(rowPtr[i]) * inv);
+  }
+}