F-mesh-1: KV Cache OOM Root Cause + Fix Status

Executive Summary

• Symptom: vec![0.0; 503 MiB] per worker on gemma4-31b boot, exceeds 128 MiB CF Worker isolate cap, kernel unreachable panic on /split-a.
• Root cause: legacy whole-model KV alloc — every worker (even single-layer trisplit shards) allocated slots for all 60 layers.
• Fix: already in source as wave-12 task #31. KVCache::with_base(base, num_layers, …) constructor + kv_base_layer / kv_num_layers plumbing through WasmGemma4Pipeline::from_backend and called from apps/distributed-inference-worker/src/index.ts:569. Per-worker alloc drops from ~503 MiB → ~8 MiB.
• Status: code-complete and end-to-end verified at the source level (kv_cache.rs → model_gemma4.rs → wasm_bindings.rs → distributed-inference-worker/index.ts:569). Not yet validated on a deployed trisplit worker — the wasm bundle in apps/distributed-inference-worker/wasm/ may still predate the change. Bule cost to ship: ~3 (rebuild wasm + redeploy 188 nodes + smoke).
• Verification (2026-05-03): confirmed (a) KVCache::with_base allocates num_layers * num_kv_heads * max_seq_len * head_dim floats — exactly the legacy formula, but num_layers is now caller-controlled; (b) Gemma4Pipeline::from_knot_with_kv_range defaults kv_layers = cfg.num_layers when None (back-compat preserved); (c) WasmGemma4Pipeline::from_backend exposes both Option<u32> params with the explicit comment "single-layer trisplit worker passing (17, 1) cuts the cache from ~503 MiB to ~8 MiB"; (d) the TS worker computes kvBaseLayer = spec.startLayer; kvNumLayers = spec.endLayer - spec.startLayer and passes both into WasmGemma4Pipeline.from_backend (gemma path) AND WasmLlamaPipeline.from_backend (llama path). No code change required — only build + deploy.

1. Allocation Site

open-source/gnosis/distributed-inference/src/kv_cache.rs:40-49

pub fn with_base(base_layer, num_layers, num_kv_heads, head_dim, max_seq_len) -> Self {
    let size = num_layers * num_kv_heads * max_seq_len * head_dim;
    KVCache {
        ...,
        keys:   vec![0.0; size],   // f32 * size bytes
        values: vec![0.0; size],   // f32 * size bytes
    }
}

Caller at model_gemma4.rs:215-223:

let kv_base   = kv_base_layer.unwrap_or(0);
let kv_layers = kv_num_layers.unwrap_or(cfg.num_layers);    // <-- legacy = full 60
let kv_cache  = KVCache::with_base(
    kv_base, kv_layers,
    cfg.num_kv_heads.max(cfg.global_num_kv_heads.unwrap_or(0)),  // 16
    head_dim_max,                                                 // 512 (global)
    KV_MAX_SEQ_LEN,                                               // 128 on wasm32
);

Formula breakdown (gemma4-31b, wasm32)

Whole-model: 2 × 60 × 16 × 512 × 128 × 4 = 503,316,480 B ≈ 480 MiB (memory log rounds 503 MiB).

2. Per-Worker Math (the win)

Trisplit worker owns 3 layers:

2 × 3 × 16 × 512 × 128 × 4 = 25,165,824 B ≈ 24 MiB

Single-layer fluid-mesh worker:

2 × 1 × 16 × 512 × 128 × 4 ≈ 8 MiB

Either is comfortably under the 128 MiB cap.

KV_MAX_SEQ_LEN=128 on wasm32 (model_gemma4.rs:51) — already trimmed from the native 2048 in a prior pass for prompt-fit headroom; do not bump without re-budgeting.

3. Proposed Fix — already implemented

The exact wave-12 task #31 fix is in source:

1. src/kv_cache.rs — with_base(base_layer, num_layers, …) constructor + local_layer() translates absolute → local slot with debug-assert range check. new(...) keeps back-compat (delegates to with_base(0, num_layers, …)).
2. src/model_gemma4.rs:108-118, 124-128, 215-223 — from_parts and from_knot_with_kv_range accept Option<usize> for kv_base_layer / kv_num_layers; defaults to whole-model alloc when None.
3. src/wasm_bindings.rs:760-814 — WasmGemma4Pipeline::from_backend exposes both as Option<u32> over the wasm-bindgen boundary. Same parameters added to WasmLlamaPipeline::from_backend (line 241).
4. apps/distributed-inference-worker/src/index.ts:569-609 — TS worker computes kvBaseLayer = spec.startLayer, kvNumLayers = spec.endLayer - spec.startLayer and passes both into cls.from_backend(...) for both Gemma and Llama paths.

What is NOT yet done:

• Rebuild + redeploy of the 180 tri-g4 + 8 cobordism-g4 workers with the wasm bundle that contains this code.
• Smoke verify on one trisplit worker that boot completes without unreachable panic.

4. Validation Plan

1. cd open-source/gnosis/distributed-inference && wasm-pack build --target web --out-dir pkg (release; debug-assert off in release so the local_layer panic won't crash prod).
2. Copy pkg/distributed_inference_bg.wasm into apps/distributed-inference-worker/wasm/.
3. Deploy ONE worker first: wrangler deploy --name tri-g4-001 --env production.
4. Boot smoke: curl https://tri-g4-001.taylorbuley.workers.dev/health — expect 200, no unreachable.
5. Inference smoke: curl -X POST .../split-a with the canonical prompt; expect coherent K/V chunk in the response (not garbage-filled, not OOM).
6. Memory delta: peek wrangler tail boot log for [boot] arch=gemma4 pipeline=WasmGemma4Pipeline and watch for the absence of the prior OOM signature.
7. If clean, batch-deploy remaining 179 trisplit + 8 cobordism workers.

5. Hopf-Link Prediction (F-mesh-4)

F-mesh-4 = "KV cache stride mismatch". kv_stride = max(num_kv_heads * head_dim, global_num_kv_heads * global_head_dim) at model_gemma4.rs:210 accommodates both sliding (16×256=4096) and global (4×512=2048) shapes in a single batch_k/v buffer. head_dim_max (512) is what feeds the KV cache row size, and is independent of per-layer kv_dim.

The OOM bug and the stride bug share the same physical buffer (the K/V rows of the per-worker cache) but represent different failure modes:

• F-mesh-1 (this fix): row count is wrong (full 60 instead of owned slice) → boot OOM.
• F-mesh-4: row layout within a slot may be wrong when sliding-vs-global layers interleave on a worker that owns layers spanning a pattern boundary.

Hopf-link prediction: the two are linked but not identical. Cutting num_layers to the owned slice means a worker now hosts 1–3 contiguous absolute layers, so it sees at most one global / sliding boundary. The stride buffer is already sized to the max so writes won't go out-of-bounds. Whether the read side (attention scoring across cached positions) honors per-layer kv_dim is the second test — a pipeline that owns layers [5..8] straddles the global layer 5. If F-mesh-4 was a function of "many layers, mixed types, all aliased into one cache", reducing to 3 layers may collapse the symptom without an explicit stride fix.

Verdict: not a guaranteed joint resolution but a strong test of the wave-13 Hopf-link hypothesis. After the F-mesh-1 deploy, run the F-mesh-4 reproducer (KV scoring divergence vs Ollama oracle) before and after; if F-mesh-4 also clears, the link is confirmed.

6. Risks & Rollback

• Range mistake (worker passes wrong [startLayer, endLayer)): local_layer() debug_assert! catches in debug builds; release builds will silently read wrong cache rows → garbage output. Mitigation: smoke one worker before fanning out.
• Wasm bundle mismatch (deployed worker has stale wasm without the new param positions): wasm-bindgen ABI shift = boot panic. Mitigation: verify wasmgemma4pipeline_from_backend arity in the deployed .wasm.d.ts before deploy (current ABI = 6 params: a-f).
• Rollback: wrangler rollback tri-g4-XXX returns to the previous version that has the legacy whole-model alloc. The previous version OOM'd before but did not corrupt state on disk; rollback is one command.
• R2 knot is unaffected — KV cache is purely worker memory, no on-disk side effect.

Bule Cost

~3 bule total. Source change is done. Remaining: rebuild wasm (~5 min), copy to apps/*/wasm/ (~1 min), deploy first worker + smoke (~10 min), batch-deploy 187 more (~30 min via wrangler script). All low-risk because the source fix is already authored and end-to-end plumbed; we are validating, not designing.