Fix Metal resource leak under high concurrency

[janhilgard]

Summary

Fixes #91 — Metal buffer leak under high concurrency.

Root cause: batch.tokens grows via mx.concatenate() each generation step but is never evaluated, so computation graph nodes hold AGXAllocation handles indefinitely. Under high concurrency this exhausts Metal resource handles.

Changes:

• Add mx.async_eval(*batch.tokens) after each generation step to eagerly evaluate accumulated token concatenations and release Metal buffers
• Make cache clear interval adaptive: scales inversely with active sequence count (min interval 8, base interval 32) to prevent Metal resource handle exhaustion under high-concurrency workloads
• Add explicit mx.eval(*tokens) during periodic cache clear to collapse any remaining lazy concatenation chains

How it works

Fix A — Eager evaluation (line ~202):

mx.async_eval(batch.y, batch.logprobs)
# NEW: evaluate accumulated tokens to prevent Metal buffer buildup
if batch.tokens:
    mx.async_eval(*batch.tokens)

Fix B — Adaptive cache clearing (line ~2239):

• Base interval: 32 steps (unchanged for single-user)
• At 8+ concurrent sequences: interval halves to 16
• At 16+ concurrent sequences: interval drops to 8 (minimum)
• Each clear also evaluates batch.tokens to collapse lazy chains

Test plan

• Verify ioclasscount | grep AGXAllocation stays stable under sustained high-concurrency load
• Benchmark single-user throughput (should be unchanged)
• Benchmark multi-user throughput with 8+ concurrent requests
• Verify no regression in generation quality

🤖 Generated with Claude Code

[dinnyosz]

I've tried this and seems to be working. Tanked my memory (48GB) with 2 4b qwen modell loaded without this, with the fix it is around 6gb memory usage. can you please check and merge if looking good?

[waybarrios]

2 issues found. Do you mind if you take a look @janhilgard ? Other than that, the pr is good!

[waybarrios]

mx.async_eval(*batch.tokens) may be redundant here. batch.tokens[i] is built via mx.concatenate((toks, y[i:i+1])) which depends on batch.y. Two lines below, y = y.tolist() forces synchronous evaluation of y, which cascades to evaluate these token concatenation chains. This async_eval either (a) is a no-op since the sync point already forces evaluation, or (b) inserts extra async work between the batch.y async eval above and its consumption via y.tolist(), potentially causing premature materialization of shared computation graph nodes.

[waybarrios]

The hardcoded floor of 8 causes effective_interval to collapse to 8 for any active_seqs >= 32 (since 32 // (32//8) = 8) — a 4x increase in cache-clear frequency. The comment at lines 1051-1052 documents _clear_cache_interval as the operator-configurable knob, but this formula bypasses it. Consider deriving the floor from the configured value (e.g. self._clear_cache_interval // 4) instead of a magic constant. Also, recalculating effective_interval every step() with a monotonic _step_count means the modulo trigger can be skipped for extended periods when concurrency drops suddenly.

[waybarrios]

The pr is good, just minor edits

[waybarrios]

pushed a small fix (30c5ad4):

• removed the mx.async_eval(*batch.tokens) — turns out y.tolist() right below already forces sync eval of the token chains, so it was redundant
• changed the interval floor from a hardcoded 8 to self._clear_cache_interval // 4 so it respects whatever the operator configures