fix(nemotron_h): prefill memory grows roughly quadratically with prompt length and aborts with Metal out-of-memory at ~6K tokens #2281

Prefilling a Nemotron-H model on Metal needs transient memory that grows roughly with the square of the prompt length. A prompt of about 6,000 tokens aborts the server with a Metal out-of-memory error on a 48 GB Mac, and about 3,000 tokens is enough to exhaust a 16-24 GB Mac.

After each request the footprint returns to about 11 GB, so this is a transient peak, not a leak. Backend.AI GO users hit it as "load the model, send 안녕, and memory overflows": the app's chat request carries a system prompt and tool schemas that already run to several thousand tokens.

Environment

• mlxcel-server 0.7.0 (Backend.AI GO bundle), Apple M4 Pro 48 GB, macOS. The startup log reports Wired memory limit: 37.4 GB.

• Model: hancheolp/mlx_test.

  • Nemotron 3.5 Lightning 30B-A3B, MLX affine 2-bit (group 64), 10.2 GB on disk.
  • model_type: nemotron_h.
  • 52 layers: Mamba2, MoE (128 experts, top 6), and 6 attention layers.
  • mamba_ssm_cache_dtype: float32, chunk_size: 128, max_position_embeddings: 262144.

• Launch (as Backend.AI GO launches it):

  mlxcel-server --model <dir> --ctx-size 8192 --jinja --host 127.0.0.1 --port 39913 --alias repro --parallel 1
  

Measurements

Each request was one non-streaming /v1/chat/completions call:

• a system message of N tokens plus the user message "안녕"
• max_tokens: 16
• chat_template_kwargs.enable_thinking: false

The peak is Physical footprint (peak) from vmmap --summary <pid>, read after each request. The footprint is 10.1 GB right after loading.

Doubling the prompt from 1,494 to 2,980 tokens multiplies the extra memory by about 3.3. The abort:

libc++abi: terminating due to uncaught exception of type std::runtime_error: [METAL] Command buffer execution failed: Insufficient Memory (00000008:kIOGPUCommandBufferCallbackErrorOutOfMemory).

What does not change it

• Prefill chunk size. --prefill-chunk-size 128 gives the same peaks: 15.0 GB at 1,494 tokens and 26.5 GB at 2,980, with an abort at ~6,000. With --parallel 1 (hybrid models are clamped to one slot) the chunk setting does not seem to bound this path.
• Prompt cache. The prompt cache is on by default. The peak is reached during prefill of a request with no cache hit, so the cache is not the cause.

Comparison

choppedgarlic/Qwen3.8-27B-AEON-ULTIMATE-UNCENSORED-4bit-MLX (qwen3_5, gated delta hybrid) under the same launch peaks at only about 2.2 GB over loaded for a 4,484-token prompt. The growth looks specific to the Nemotron-H path. The likely suspect is the Mamba2 SSD prefill materializing an L×L (per head) intermediate over the whole prompt instead of per chunk_size block, or keeping all layers' intermediates alive in one lazy graph.

Expected

• Prefill memory for nemotron_h grows about linearly with prompt length: chunked SSD over chunk_size blocks, and per-layer evaluation so intermediates are freed.
• A prompt that does not fit returns an error instead of aborting the process.

Repro script

A self-contained Python script (start the server, send N-token prompts, read vmmap peaks) is available on request; the table above is its output.