Repository navigation
fix: reject over-budget Nemotron-H requests before Metal prefill #2291
Description
Problem / Background
Follow-up to #2281 and #2290: bounded SSD removes quadratic prefill growth, but an in-context Nemotron-H request can still exceed Metal memory and terminate serving. Reject predictably oversized requests before GPU forward; do not claim protection against arbitrary external GPU pressure.
Current Behavior
src/server/model_worker.rs:585 clamps unsupported batching to one slot. src/server/batch/scheduler/admission.rs:174 rejects context overflow, not request memory. src/server/batch/scheduler/planned_prefill.rs:62 owns actual prefill geometry. src/execution/memory_estimate.rs:530 estimates startup memory; its global-chunk activation assumption is not necessarily this model's forward span. src/models/nemotron_h.rs:935 bounds SSD internals independently.
Proposed Solution
Add a pure checked-byte estimator and scheduler admission guard for Metal, text-only Nemotron-H at effective width one. Reuse normalized NemotronHConfig, architecture-aware execution::kv_arch, PrefillPlan, runtime budgets, and mlxcel_core::memory counters. Account conservatively for resident allocations, prompt plus effective planned-output KV, recurrent state, actual largest-forward activations/logits, and bounded float32 SSD scratch. Derive dimensions/dtypes from loaded configuration; no checkpoint-specific counts. Preserve existing context rejection precedence.
Use MLXCEL_REQUEST_MEMORY_BUDGET as a positive byte/size override that can only lower the default: minimum positive device recommended working-set, configured wired cap, and allocator cap. Reject invalid explicit values at startup; unknown budget or arithmetic overflow fails closed. Note get_wired_limit() reports device recommendation, not the configured cap (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:4696).
Admission must run before request GPU work and after preceding single-slot work finishes. Count retained caches; never subtract reusable allocations twice or treat allocator free-cache as permanent resident weight. On rejection return HTTP 503 with error.type and error.code both request_memory_budget_exceeded, including estimated/budget bytes; streaming must reject before committing SSE headers. Keep health responsive and queue/cache ownership intact.
Scope
One PR; no other-family estimator expansion, multi-slot reservation scheme, CUDA policy, or speculative/VLM redesign. Document unsupported paths explicitly.
Acceptance Criteria
- Deterministic unit/integration tests cover budget boundaries, output-length growth, overflow, invalid values, cache warm/cold accounting, context precedence, and zero forward calls on rejection.
- Real checkpoint demonstrates within-context rejection, healthy server, then successful short request, for streaming and non-streaming; record estimates versus measured accepted-request peaks and exact commands.
- Smoke-test a second family if shared paths change; unsupported paths retain behavior.
Verification
cargo test --release request_memory
cargo test --release nemotron_h_tests::
cargo fmt --check
cargo clippy --all-targets -- -D warnings
cargo build --release --features metal,accelerate
Activity
fix: close request admission bypasses and lazy residency gaps
docs: record request memory admission validation
fix: avoid copying scheduler in worker setup
docs: verify final admission artifact after scheduler stack fix
fix: publish WebUI harness controls atomically
fix: admit Nemotron-H requests against a checked memory budget (#2294)