Repository navigation
fix(attention): keep SDPA256 prefill on the bounded route - #3685
Conversation
|
Making routing reproducible is useful, but I’d like to preserve the existing OOM protection as well. The current live-memory check was added in #2204 to keep the fast unfused path when it fits, and use the bounded path otherwise. This PR’s static budget leaves out other loaded models and active requests. Since the prefill guard assumes the smaller bounded allocation, it could approve a chunk that then takes the larger unfused path, potentially bringing back the OOM issue from #2025. Could you check whether the remaining guards safely cover this case? If this concern is valid, I’d appreciate an update that preserves that protection while improving reproducibility. |
c5f008d to
7e72951
Compare
tc3oliver
changed the title
|
Thanks — the concern is valid, and checking it changed the fix. I traced the remaining admission and memory guards. The prefill estimator still prices a qualifying head_dim-256 chunk as O(L) for every shape the route covers, while since #2204 the route can select the unfused path from live headroom. So the estimator assumption and the route that actually runs can diverge, and a chunk admitted on the bounded transient can materialize the full score matrix. The static budget I had used does not close that — it prices only the requesting model and the current request, so other loaded models and other in-flight work are invisible to it. I've dropped that design. I tried a route-aware estimator first. It restores the accounting, but admission's baseline is live residency and its throttle also reads live usage, so charging the real unfused transient brought process-history-dependent partitioning back — five histories, five partitions. With no request-level reservation available, I went with the smaller option: keep qualifying SDPA256 prefill on the bounded route, which matches both the original #2025 safety semantics and the estimator's assumption. On the same 68k reproduction (M4 Max 64GB, 24:4 GQA head_dim 256, temperature 0, three fresh processes per group): prefill throughput 201.8 → 196.6 tok/s (−2.6%), peak MLX memory 35.57 → 24.70 GiB (−30.6%), decode unchanged, and TTFT within run-to-run noise (358.7 s vs 356.1 s) — no material regression, and I'm not claiming it is faster. A separate four-history qualification on the final tree gave one route, one chunk partition, one logits hash and one output hash. I also checked the second-model-resident and concurrent-request cases. One thing I did not fix: admission has no reservation for approved-but-not-yet-materialized allocations, so two chunks admitted against the same snapshot both pass. Keeping the route bounded reduces the exposure by more than an order of magnitude but does not remove that race; it looks like it belongs to the admission design rather than here, so I've left it characterized in the tests rather than changed. The PR and tests are updated accordingly and rebased onto current main. |
7e72951 to
387a09e
Compare
head_dim 256 has no fused MLX kernel for multi-token prefill, so the stock fallback materializes the whole [n_q, query_len, kv_len] fp32 score matrix. and registered the head dim with the memory monitor so prefill admission prices the bounded transient instead of the matrix. unfused path wherever its matrix fits. Two things follow from deciding per call from live process memory. The route stops being reproducible: the same request saw headroom differ by gigabytes between otherwise identical processes (resident pages of the mmap'd weights, a previous request's pooled buffers), switched routes at a different kv_len, and — because the bounded and unfused paths are different floating-point reductions — produced different temperature-0 output. And the estimator was not made conditional alongside it: memory_monitor still prices a registered head dim as O(L) for every shape the route covers, so the route and the guard's accounting can diverge, and a chunk admitted on the bounded transient can run the unfused one. Decide from the call shape alone. Qualifying calls take the bounded route in every process, which restores jundot#2025's semantics and makes the estimator's assumption true by construction. OMLX_SDPA256_TILED=1/0 still forces or disables the route, and a disabled route still withdraws the O(L) registration. The scheduler's headroom provider, its static-residency helpers and the propagation marker that existed only to feed them are removed; the route now takes no memory input at all, so there is no failure mode left in which it can fall back to the unfused path. Measured on a 68,034-token prefill of a head_dim-256 GQA model (24:4, 16 full-attention layers), same machine and settings, three fresh processes per group: TTFT 358.7s unfused vs 356.1s bounded (within run-to-run noise), prefill 201.8 -> 196.6 tok/s (-2.6%), peak MLX memory 35.57 -> 24.70 GiB (-30.6%), decode unchanged. Five different process histories produced one chunk partition, one first-token logits hash and one output hash. This does not change prefill admission itself: nothing reserves the bytes a chunk has been admitted to allocate, so two chunks admitted before either transient lands still both pass. Keeping the route bounded shrinks the per-request attention transient by more than an order of magnitude, which reduces the exposure, but that race belongs to the admission design and is left alone here.
387a09e to
a17a28b
Compare
jundot
left a comment
jundot
left a comment
There was a problem hiding this comment.
Keeping qualifying prefill on the bounded route makes execution match the memory estimate. I did not reproduce the earlier large causal-prefill regression on M5.
The tested image workload showed about a 2% TTFT increase with array tiling. I consider that an acceptable tradeoff for predictable memory use.
I'll preserve the FP32 fallback when resolving the merge conflict and clean up the stale routing references afterward.
…oute jundot's review on the conditional q-split route (jundot#2991) reproduced retained-pool growth from differently-sized causal sub-tile score buffers and a laziness pile-up from missing per-sub-tile evaluation. The conditional route is gone (jundot#3685 keeps SDPA256 prefill on the bounded route, decided from call shape alone); inside that bounded route these two hazards are structural invariants, not patches: - every per-kv-tile score buffer is uniformly clamped to (_Q_TILE, _KV_TILE), so Metal's pool reuses buffers tile-to-tile (uniform shapes => no retained growth); - mx.eval fires at least once per kv tile, bounding the live graph. This regression test instruments both invariants so re-adding a growth-shaped or lazily-evaluated attention path fails loudly, preserving the audit's intent on the architecture that superseded the PR.
…eaving jundot#3099 fixed a log flap where 16 full-attention layers (plus concurrent requests) deduped route-transition logging against one module-global last-decision, so per-call variation chattered the log. The conditional route that produced per-call variation is gone — jundot#3685 keeps SDPA256 prefill on the always-bounded route and logs engagement per static reason. This regression test ports the audit's intent onto that design: interleaved caches, layers, and call-site detail variations must never produce more than one engagement record per reason.
Problem
The SDPA256 route chooses between two different floating-point reductions. Deciding per call from live guard headroom makes that choice a function of live process memory, which differs between otherwise identical processes — resident pages of the mmap'd weights, a previous request's pooled buffers. The same request then switches routes at a different
kv_lenand produces different temperature-0 output.Safety audit
The concern raised on this PR is valid, and tracing it changed the fix.
memory_monitorprices a registered head dim as O(L) for every shape the route covers, but since #2204 the route can still select the unfused path based on live headroom. The existing estimator assumption and the live-memory route can therefore diverge, and a chunk admitted on the bounded transient can go on to materialize the full[n_q, query_len, kv_len]fp32 matrix.The static-budget approach this PR previously used does not address that: its budget prices only the requesting model and the current request, so other loaded models and other in-flight work are invisible to it. That approach has been dropped.
Fix
Decide from the call shape alone. Qualifying head_dim-256 prefill takes the bounded route in every process, which restores #2025's semantics and makes the estimator's assumption true by construction.
OMLX_SDPA256_TILED=1/0still forces or disables the route, and a disabled route still withdraws the O(L) registration. The scheduler's headroom provider, its helpers and the propagation marker that existed only to feed them are removed, so the route now takes no memory input at all.Why this design
A route-aware live estimator was tried first. It restores the accounting, but prefill admission's baseline is live residency and its throttle also reads live usage, so charging the real unfused transient pulls the request into the throttle and chunk partitioning becomes process-history dependent again — five histories produced five partitions. There is no request-level reservation to lean on instead. Unconditional bounded routing is the smallest change found that keeps both deterministic routing and #2025's safety semantics without redesigning scheduler admission.
Validation
Apple M4 Max, 64 GB, macOS 15.6. head_dim-256 GQA model (24 query heads : 4 KV heads, 16 full-attention layers, 4-bit weights), 68,034-token prompt, 32 output tokens, temperature 0 / top_k 1,
max_concurrent_requests: 1, prefix cache disabled, memory guardbalanced,prefill_min_chunk_tokens: 256.Reproducibility, four process histories (two cold starts, one small preceding request, one 68K preceding request):
Memory behaviour:
Tests: full suite green; the new regression module fails on the parent commit and passes here.
Performance
Controlled A/B on the same machine, model, prompt and settings, three fresh processes per group:
TTFT shows no material regression — the difference is within run-to-run noise on this machine. The measured cost is prefill throughput, −2.6%.
Known limitation
Prefill admission has no reservation for approved-but-not-materialized allocations: two chunks admitted against the same memory snapshot both pass. Keeping the route bounded shrinks the per-request attention transient by more than an order of magnitude, which reduces the exposure, but this PR does not claim to fix that scheduler-wide race. It is characterized in the tests rather than changed.