Skip to content

fix: reject over-budget Nemotron-H requests before Metal prefill #2291

Closed

Description

@inureyes

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

  1. added
    type:bugBug fixes, error corrections, or issue resolutions
    area:inferenceGeneration, sampling, decoding (incl. speculative, DRY)
    platform:macosmacOS (Apple Silicon) specific
    and removed on Oct 13, 2026
  2. added a commit that references this issue on Oct 14, 2026
    5997428
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:inferenceGeneration, sampling, decoding (incl. speculative, DRY)platform:macosmacOS (Apple Silicon) specificpriority:highHigh prioritystatus:doneCompletedtype:bugBug fixes, error corrections, or issue resolutions

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions