fix(metal): attnFold slabs did not grow with the wide tail-absorbed prefill chunk — ~52K long-context corruption

Root cause of the long-context needle-sweep failures (argmax -1 / instant EOS /
silent needle loss at ~52K on E2B): attnFold keyed its reallocation on mlpFold's
foldRowCap. mlpFold always runs FIRST and had already raised foldRowCap for a
wider chunk, so attnFold skipped its realloc and every attention slab (attn-norm,
Q, SDPA out, O out, K/V stage) stayed at the previous row count. The batched
prefill's tail absorption produces exactly ONE chunk wider than all before it
(window + tail, when promptLen mod window lands in (0, window/2]), so rows past
the stale capacity read and wrote out of bounds — undefined bytes that varied per
process: NaN cascades through attention (the argmax -1 hard failure), or silently
wrong hiddens (the needle loss at 97K, and non-identical temp-0 outputs).

attnFold now tracks its own attnRowCap/attnDModel. Forensics that found it stay
in-tree, env-gated (LTHN_DEBUG_ARGMAX): the invalid-argmax per-stage dump
(hidden -> normed -> logits -> tiles), the per-chunk boundary-hidden NaN scan with
per-layer KV fingerprinting, and the last-prefill-step seam dump. cmd/lem gains
-prompt-file (long prompts exceed argv).

Receipts: TestDenseBatchScratchAttnFoldGrowsWithRows pins the production call
order (mlpFold first, then attnFold at a grown row count) against the slab sizes;
CLI repro 3/3 clean at 51,924 tokens (was ~2/3 failing) with zero NaN diagnostics;
TestE2BQ4LongContextPrefillCorruption (real-model, E2B_Q4_DIR-gated) now produces
byte-identical temp-0 tokens across reps at the wide-chunk length (previously
diverged per rep — the silent form); full engine/metal suite 1436 green.

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

Author: Snider

go/cmd/lem/generate.go

@@ -39,6 +39,7 @@ func runGenerateCommand(ctx context.Context, args []string, stdout, stderr io.Wr
 	fs := flag.NewFlagSet(cliCommandName("generate"), flag.ContinueOnError)
 	fs.SetOutput(stderr)
 	prompt := fs.String("prompt", "Write a detailed Go function that reverses a singly linked list, with inline comments on every step, then explain the pointer dance.", "user prompt")
+	promptFile := fs.String("prompt-file", "", "read the user prompt from a file (long-context runs exceed argv limits); overrides -prompt")
 	maxTokens := fs.Int("max-tokens", 128, "tokens to generate")
 	draftPath := fs.String("draft", "auto", "MTP drafter: 'auto' detects one beside a Gemma 4 target (assistant/ pair layout, MTP/ gguf), a path forces it, '' disables")
 	draftBlock := fs.Int("draft-block", 0, "MTP draft block (verify forward = carried lead + block-1 proposals); 0 = engine default 4")
@@ -91,10 +92,25 @@ func runGenerateCommand(ctx context.Context, args []string, stdout, stderr io.Wr
 		return 2
 	}
 
+	promptText := *prompt
+	if *promptFile != "" {
+		read := core.ReadFile(*promptFile)
+		if !read.OK {
+			core.Print(stderr, "%s generate: read -prompt-file %s: %s", cliName(), *promptFile, read.Error())
+			return 1
+		}
+		bytes, ok := read.Value.([]byte)
+		if !ok || len(bytes) == 0 {
+			core.Print(stderr, "%s generate: -prompt-file %s is empty", cliName(), *promptFile)
+			return 1
+		}
+		promptText = string(bytes)
+	}
+
 	native.SetPipelinedGPUDecode(*pipeline) // engine-level: -pipeline=false forces the chained serial loop
 	err := generate.RunGenerate(ctx, generate.Config{
 		ModelPath:  fs.Arg(0),
-		Prompt:     *prompt,
+		Prompt:     promptText,
 		MaxTokens:  *maxTokens,
 		Temp:       *temp,
 		Think:      *think,

go/engine/metal/arch_session.go

@@ -1867,6 +1867,7 @@ func (s *ArchSession) prefillRetainedTokensBatchedDense(ids []int32, scope strin
 
 func (s *ArchSession) prefillRetainedTokensBatchedDenseChunks(ids []int32, scope string) ([]byte, bool, error) {
 	var hidden []byte
+	chunk := 0
 	for len(ids) > 0 {
 		n := s.batchedDensePrefillChunkLen(len(ids))
 		if n <= 0 {
@@ -1876,6 +1877,29 @@ func (s *ArchSession) prefillRetainedTokensBatchedDenseChunks(ids []int32, scope
 		if err != nil || !ok {
 			return nil, ok, err
 		}
+		if argmaxDebugEnabled() {
+			if nan, first := bf16NaNScanBytes(nextHidden); nan > 0 {
+				nativeTraceLog(core.Sprintf("argmax-diag: batched prefill chunk %d (rows %d, pos now %d): boundary hidden NaN=%d first=%d\n",
+					chunk, n, s.pos, nan, first))
+				if views, verr := s.stateLayerViews(); verr == nil {
+					for _, v := range views {
+						kc, kf := bf16NaNScanBytes(v.keyBytes)
+						vc, vf := bf16NaNScanBytes(v.valueBytes)
+						if kc > 0 || vc > 0 {
+							sp := s.state.specs[v.layer]
+							at := "sliding"
+							if sp.Attention == model.GlobalAttention {
+								at = "GLOBAL "
+							}
+							rowB := kvHeadsOf(sp, s.arch.KVHeads) * headDimOf(sp, s.arch.HeadDim)
+							nativeTraceLog(core.Sprintf("argmax-diag:   L%2d %s owns=%v shareFrom=%d  K-NaN=%d(first row %d)  V-NaN=%d(first row %d)\n",
+								v.layer, at, sp.OwnsCache(), sp.KVShareFrom, kc, kf/rowB, vc, vf/rowB))
+						}
+					}
+				}
+			}
+		}
+		chunk++
 		hidden = nextHidden
 		ids = ids[n:]
 	}

go/engine/metal/decode_batched_session.go

@@ -71,6 +71,8 @@ type denseBatchScratch struct {
 	attnOutPacked  metal.MTLBuffer // K × dModel O-projection outputs
 	kStagePacked   metal.MTLBuffer // K × kvDimMax staged K rows (ring-wrap landing)
 	vStagePacked   metal.MTLBuffer // K × kvDimMax staged V rows
+	attnRowCap     int // attnFold's OWN row capacity — NOT foldRowCap (mlpFold updates that first, which masked row growth and left these slabs short: the ~52K wide-tail-chunk corruption)
+	attnDModel     int
 	foldQDimCap    int
 	foldKVDimCap   int
 	// per-layer staging for the deferred-landing lane (the big-K staged sliding tail): each
@@ -100,16 +102,22 @@ func (s *denseBatchScratch) mlpFold(k, dModel, dFFMax int) (h, normed, gate, up,
 	return s.hPacked, s.mlpNormPacked, s.gatePacked, s.upPacked, s.gatedPacked, s.downPacked
 }
 
-// attnFold returns the attention-fold slabs, (re)allocating alongside mlpFold's sizing. Call after
-// mlpFold (it owns foldRowCap/foldDModel); qDimMax/kvDimMax are the widest per-layer head geometry.
+// attnFold returns the attention-fold slabs, (re)allocating when the batch width, model width or
+// the widest per-layer head geometry grows. Growth is tracked by attnFold's OWN attnRowCap /
+// attnDModel: the old code keyed on mlpFold's foldRowCap, which mlpFold (always called first)
+// had ALREADY raised for a wider chunk — attnFold then skipped its realloc and every attention
+// slab stayed at the old row count. At ~52K-token prompts the sliding-tail absorption produces
+// one chunk wider than all before it (window + tail), so rows past the stale capacity read and
+// wrote out of bounds: undefined bytes → per-process NaN/garbage → the long-context corruption.
 func (s *denseBatchScratch) attnFold(k, dModel, qDimMax, kvDimMax int) (normed, q, attn, attnOut, kStage, vStage metal.MTLBuffer) {
-	if s.attnNormPacked == nil || s.foldRowCap < k || s.foldDModel != dModel || s.foldQDimCap < qDimMax || s.foldKVDimCap < kvDimMax {
+	if s.attnNormPacked == nil || s.attnRowCap < k || s.attnDModel != dModel || s.foldQDimCap < qDimMax || s.foldKVDimCap < kvDimMax {
 		s.attnNormPacked = scratchBF16(k * dModel)
 		s.attnOutPacked = scratchBF16(k * dModel)
 		s.qPacked = scratchBF16(k * qDimMax)
 		s.attnPacked = scratchBF16(k * qDimMax)
 		s.kStagePacked = scratchBF16(k * kvDimMax)
 		s.vStagePacked = scratchBF16(k * kvDimMax)
+		s.attnRowCap, s.attnDModel = k, dModel
 		s.foldQDimCap, s.foldKVDimCap = qDimMax, kvDimMax
 	}
 	return s.attnNormPacked, s.qPacked, s.attnPacked, s.attnOutPacked, s.kStagePacked, s.vStagePacked

go/engine/metal/decode_batched_session_test.go

@@ -301,3 +301,40 @@ func TestStepTokensBatchedDenseSyncsLinearCacheAfterPagedStep(t *testing.T) {
 		eqBytes(t, core.Sprintf("batched after paged row %d vs stepToken", i), batOut[i], seqOut[i])
 	}
 }
+
+// TestDenseBatchScratchAttnFoldGrowsWithRows pins the ~52K long-context corruption's root
+// cause: attnFold must reallocate its slabs when the batch row count GROWS, independent of
+// mlpFold. The production call order is mlpFold first (which raises the shared-looking
+// foldRowCap), then attnFold — the old attnFold keyed its growth check on foldRowCap and so
+// skipped the realloc for the one wide tail-absorbed chunk (window + tail rows), leaving every
+// attention slab short: rows past the stale capacity read/wrote out of bounds (undefined bytes,
+// NaN/garbage varying per process). The wide-chunk shape only occurs when promptLen mod window
+// lands in (0, window/2], which is why it escaped every fixed-size fixture.
+func TestDenseBatchScratchAttnFoldGrowsWithRows(t *testing.T) {
+	requireNativeRuntime(t)
+	const dModel, qDim, kvDim, dFF = 2048, 2048, 256, 8192
+	s := &denseBatchScratch{}
+	// the steady prompt chunks: window-sized batches
+	s.mlpFold(512, dModel, dFF)
+	normed, q, attn, attnOut, kSt, vSt := s.attnFold(512, dModel, qDim, kvDim)
+	if int(bufferLengthFast(q)) < 512*qDim*bf16Size {
+		t.Fatalf("baseline q slab too small: %d", bufferLengthFast(q))
+	}
+	_, _, _, _, _ = normed, attn, attnOut, kSt, vSt
+	// the wide tail-absorbed chunk: window + tail rows, mlpFold FIRST (production order)
+	const wide = 724
+	s.mlpFold(wide, dModel, dFF)
+	normed, q, attn, attnOut, kSt, vSt = s.attnFold(wide, dModel, qDim, kvDim)
+	check := func(name string, got, want int) {
+		t.Helper()
+		if got < want {
+			t.Fatalf("attnFold %s slab did not grow with the wide chunk: %d bytes, want >= %d — rows past the stale capacity read/write out of bounds", name, got, want)
+		}
+	}
+	check("attnNorm", int(bufferLengthFast(normed)), wide*dModel*bf16Size)
+	check("q", int(bufferLengthFast(q)), wide*qDim*bf16Size)
+	check("attn", int(bufferLengthFast(attn)), wide*qDim*bf16Size)
+	check("attnOut", int(bufferLengthFast(attnOut)), wide*dModel*bf16Size)
+	check("kStage", int(bufferLengthFast(kSt)), wide*kvDim*bf16Size)
+	check("vStage", int(bufferLengthFast(vSt)), wide*kvDim*bf16Size)
+}

go/engine/metal/head_nocopy.go

@@ -5,6 +5,7 @@
 package native
 
 import (
+	"os"
 	"runtime"
 	"slices"
 	"sync"
@@ -862,6 +863,11 @@ func (h *headEncoder) greedyBufferAtInPool(hiddenBuf metal.MTLBuffer, hiddenOff
 	commitCommandBufferFast(cb)
 	waitUntilCompletedFast(cb)
 	token = scratch.token()
+	if (token < 0 || int(token) >= h.vocab) && argmaxDebugEnabled() {
+		diag := h.argmaxInvalidDiag(scratch, hiddenBuf, hiddenOff)
+		h.putGreedyScratch(scratch)
+		return 0, true, core.NewError(core.Sprintf("native.headEncoder.greedy: direct argmax returned invalid token %d for vocab %d%s", token, h.vocab, diag))
+	}
 	h.putGreedyScratch(scratch)
 	if !ok {
 		return 0, false, nil
@@ -872,6 +878,116 @@ func (h *headEncoder) greedyBufferAtInPool(hiddenBuf metal.MTLBuffer, hiddenOff
 	return token, true, nil
 }
 
+// argmaxDebugEnabled gates the invalid-token forensic dump (the ~52K long-context
+// corruption hunt): when the direct argmax returns an out-of-range token, dump per-stage
+// NaN counts and extrema (hidden → normed → logits → tiles) so a failing run localises
+// which stage the garbage entered at. Enable with LTHN_DEBUG_ARGMAX=1.
+func argmaxDebugEnabled() bool { return os.Getenv("LTHN_DEBUG_ARGMAX") != "" }
+
+// bf16NaNScanBytes counts bf16 NaNs in a host byte slice, reporting the first NaN's
+// element index (-1 when none) — the host-side sibling of bf16BufStats.
+func bf16NaNScanBytes(b []byte) (count, firstIdx int) {
+	firstIdx = -1
+	for i := 0; i+1 < len(b); i += 2 {
+		h := uint16(b[i]) | uint16(b[i+1])<<8
+		if h&0x7F80 == 0x7F80 && h&0x007F != 0 {
+			if firstIdx < 0 {
+				firstIdx = i / 2
+			}
+			count++
+		}
+	}
+	return count, firstIdx
+}
+
+// bf16BufStats scans n bf16 elements at a buffer offset: NaN/±Inf counts, finite min/max,
+// and the first NaN's element index (-1 when none).
+func bf16BufStats(buf metal.MTLBuffer, off uint, n int) (nan, inf int, minV, maxV float32, firstNaN int) {
+	firstNaN = -1
+	if buf == nil || n <= 0 {
+		return 0, 0, 0, 0, -1
+	}
+	b := unsafe.Slice((*byte)(unsafe.Add(buf.Contents(), uintptr(off))), n*bf16Size)
+	first := true
+	for i := 0; i < n; i++ {
+		h := uint16(b[i*2]) | uint16(b[i*2+1])<<8
+		if h&0x7F80 == 0x7F80 {
+			if h&0x007F != 0 {
+				nan++
+				if firstNaN < 0 {
+					firstNaN = i
+				}
+			} else {
+				inf++
+			}
+			continue
+		}
+		v := bf16ToF32(b[i*2], b[i*2+1])
+		if first {
+			minV, maxV, first = v, v, false
+			continue
+		}
+		if v < minV {
+			minV = v
+		}
+		if v > maxV {
+			maxV = v
+		}
+	}
+	return nan, inf, minV, maxV, firstNaN
+}
+
+// argmaxInvalidDiag reads the greedy scratch's shared-storage stages after an invalid
+// direct-argmax token and reports where the garbage entered: the input hidden, the
+// RMS-normed hidden, the (quant-path) logits, and the per-tile argmax values/indices.
+func (h *headEncoder) argmaxInvalidDiag(scratch *headGreedyScratch, hiddenBuf metal.MTLBuffer, hiddenOff uint) string {
+	rowsPerTile := bf16LMHeadArgmaxRowsPerTile
+	if h.quant {
+		rowsPerTile = bf16LogitsArgmaxRowsPerTile
+	}
+	tileCount := (h.vocab + rowsPerTile - 1) / rowsPerTile
+	hn, hi, hmin, hmax, hf := bf16BufStats(hiddenBuf, hiddenOff, h.dModel)
+	nn, ni, nmin, nmax, nf := bf16BufStats(scratch.normed, 0, h.dModel)
+	d := core.Sprintf("\n  argmax-diag: hidden NaN=%d Inf=%d min=%.4g max=%.4g firstNaN=%d\n  normed NaN=%d Inf=%d min=%.4g max=%.4g firstNaN=%d",
+		hn, hi, hmin, hmax, hf, nn, ni, nmin, nmax, nf)
+	if h.quant && scratch.logits != nil {
+		ln, li, lmin, lmax, lf := bf16BufStats(scratch.logits, 0, h.vocab)
+		d += core.Sprintf("\n  logits NaN=%d Inf=%d min=%.4g max=%.4g firstNaN=%d", ln, li, lmin, lmax, lf)
+	}
+	if scratch.tileValues != nil && scratch.tileIndices != nil {
+		tv := unsafe.Slice((*float32)(scratch.tileValues.Contents()), tileCount)
+		ti := unsafe.Slice((*int32)(scratch.tileIndices.Contents()), tileCount)
+		tnan, tneg := 0, 0
+		var tmin, tmax float32
+		firstBad := -1
+		for i := 0; i < tileCount; i++ {
+			v := tv[i]
+			if v != v {
+				tnan++
+				if firstBad < 0 {
+					firstBad = i
+				}
+				continue
+			}
+			if i == 0 || v < tmin {
+				tmin = v
+			}
+			if i == 0 || v > tmax {
+				tmax = v
+			}
+			if ti[i] < 0 || int(ti[i]) >= h.vocab {
+				tneg++
+				if firstBad < 0 {
+					firstBad = i
+				}
+			}
+		}
+		d += core.Sprintf("\n  tiles=%d NaN=%d badIdx=%d min=%.4g max=%.4g firstBad=%d  tile0=(%.4g,%d) tileLast=(%.4g,%d)",
+			tileCount, tnan, tneg, tmin, tmax, firstBad, tv[0], ti[0], tv[tileCount-1], ti[tileCount-1])
+	}
+	return d
+}
+
 func (h *headEncoder) logitsSampleUsable() bool {
 	if h.finalNorm.buf == nil || h.weight.buf == nil || !logitsSampleBF16Usable(h.vocab) {
 		return false