vulkan: scratch pool spills out of the 256 MB BAR heap into host memory — prefill collapses for chunks above 64

[geisten]

Found while tuning the PQ2_0 tensor-core GEMM (#487).

Problem

Prefill throughput collapses with larger chunks: on the RTX 2080 Ti, Ternary-Bonsai-2-27B pp512 is 350 t/s at the default m_max 64 but 139 t/s at GEIST_M_MAX=256 (128 t/s at 128). A GEIST_VK_PROFILE=1 run at 256 shows the elementwise/copy ops exploding while the GEMM stays flat: silu_mul 12 → 522 µs per call, hadamard 11 → 135 µs, copy 4970 calls at 98 µs, attention_f16 0.3 → 3.8 ms. Those buffers are the per-session scratch pool (GEIST_BUFFER_SCRATCH, created host-visible). With a 256 MB BAR heap (no resizable BAR) the pool for m = 256 no longer fits, vk_buffer_create falls back to plain host-visible system memory, and every dispatch that touches the scratch reads/writes it over PCIe.

Proposal

Scratch that only GPU ops touch can live in ordinary device-local VRAM (unmappable): the arch/backend knows which buffers a CPU fallback may map. Options: a GEIST_BUFFER_SCRATCH_DEVICE role for the activation slabs of a fully GPU-resident model (decided from the fused-op probes at plan build), or an explicit "no host access" flag; a CPU fallback then fails loudly (vk_tensor_host already returns nullptr for device-local memory) instead of silently running over PCIe. Also print a one-line note under GEIST_VK_VERBOSE whenever a scratch buffer lands outside device-local memory.

Why it matters

Larger chunks are what a tensor-core GEMM wants (fewer weight passes, wider N tile), and the same spill hits any model whose chunk × width exceeds the BAR — it currently just shows up as a slowdown.

Acceptance

GEIST_M_MAX=256 and 128 are not slower than 64 on the Bonsai (pp512, RTX 2080 Ti); a test or verbose counter that shows where each scratch buffer lives.

Related: the remainder chunk (m % 16 != 0) still uses the register-tiled GEMM — routing m & ~15 rows through the tensor-core kernel and the rest through the tiled one belongs with this (#487, #467).

[geisten]

Update after a closer look (allocation log + GEMM micro-benchmark):

• The spill is confirmed by the allocation log (GEIST_VK_VERBOSE=1): at GEIST_M_MAX=64 both scratch pools (2 × 90 MiB, one per session — the default session's pool and the bench session's) sit in the BAR heap; at 128 the two pools are 180 MiB each, the first takes the BAR and the second lands in host memory (bar=0), and GPU ops on it run over PCIe (silu_mul 30 → 1461 µs, hadamard 18 → 375 µs, add 8 → 376 µs per call).
• Two cheaper fixes exist before the big one: (1) the default session's pool is allocated but unused once an engine session exists — one pool would fit the BAR up to m ≈ 190; (2) keep the pool device-local: not a one-line change — the arch allocates the pool GEIST_MEMORY_MAPPED, calls buffer_map on it and slices it with buffer_create_aliased(host_ptr) (forcing the role device-local makes the load fail with 'scratch pool alloc failed'). A device-only pool needs a slicing API that does not go through a host pointer.
• Why it matters (measured, RTX 2080 Ti, 17408 × 5120 PQ2_0, our coopmat GEMM, weights only): 22.6 TFLOP/s at m = 64, 23.4 at 128, 24.1 at 256, 32.5 at 512; at m ≤ 64 there are only 136 workgroups (one token tile) for 68 SMs. The fork runs the whole 512-token batch per GEMM (41.5 TFLOP/s on the same shape). See the analysis in the Bonsai benchmark record.

[geisten]

Status on main (bce22ae, 2026-10-01). Mitigated: one scratch pool per model with a lazily created default session (62a041f, #493); the default chunk now shrinks toward 64 until the pool fits the BAR via the fast_host_bytes hook (c374358, #496, arch_state.c:1509-1523); verbose logs print bar= per buffer.

Remaining: an explicit GEIST_M_MAX ≥ 128 bypasses the fit loop and still spills to host memory; a device-local scratch role (needs slicing without a host pointer); the "m=256 not slower than 64" acceptance is unmeasured; and the m & 15 remainder is not split onto the coopmat GEMM (vk_linear_cm_route returns early, ops.c:593).

---

Generated by Claude Code

[geisten]

#612 implements the device-local scratch pool as an opt-in (GEIST_VK_SCRATCH_DEVICE=1). It adds slicing by offset without a host pointer, a verbose note when a scratch buffer lands in host memory, placement counters, and a test on lavapipe.

Still open:

• The measurement on the 2080 Ti (GEIST_M_MAX 128/256 vs. 64 on Bonsai), which lavapipe can't provide.
• The m & 15 remainder onto the coopmat GEMM.

Questions about the vtable slot and the default are in the PR.

---

Generated by Claude Code