[BUG] mx.fast.metal_kernel: use-after-free when multiple custom kernels compose in lazy graph

[Ziqiao-git]

Describe the bug

When multiple mx.fast.metal_kernel calls are composed in a single lazy evaluation graph, Metal API Validation reports:

failed assertion `command buffer references deallocated object which previously existed at address 0x...'

Without validation enabled, this manifests as non-deterministic NaN outputs (~2% per forward pass, amplified to ~100% in autoregressive generation).

The root cause is commandBufferWithUnretainedReferences() in Device::get_command_buffer() (device.cpp:392). Metal does not retain GPU buffers referenced by the command buffer, so intermediate arrays whose Python references are dropped before the command buffer finishes have their backing buffers deallocated while still in use by the GPU.

To Reproduce

# Requires: pip install paroquant
# Run with: MTL_DEBUG_LAYER=1 python repro.py
import mlx.core as mx
from paroquant.inference.backends.mlx.load import load
from paroquant.inference.backends.mlx.modules import RotateQuantizedLinear
from mlx_lm.models.cache import make_prompt_cache

model, proc, _ = load("z-lab/Qwen3.5-4B-PARO", force_text=True)
for _, m in model.named_modules():
    if isinstance(m, RotateQuantizedLinear):
        m._force_eval = False

inner = model.language_model.model
h = inner.embed_tokens(mx.array([[0]]))
mx.eval(h)

cache = make_prompt_cache(model.language_model)
out = inner.layers[0](h, mask=None, cache=cache[0])
mx.eval(out)
# MTL_DEBUG_LAYER=1: crashes with "command buffer references deallocated object"
# Without: non-deterministic NaN

Expected behavior

mx.fast.metal_kernel outputs should be correct regardless of lazy graph depth.

Desktop (please complete the following information):

• OS Version: macOS 15.x
• Hardware: Apple M3 Max 48GB
• MLX Version: 0.31.1

Additional context

Proposed fix — one-line change in mlx/backend/metal/device.cpp:392:

-    stream.buffer = stream.queue->commandBufferWithUnretainedReferences();
+    stream.buffer = stream.queue->commandBuffer();

commandBuffer() uses retained references — Metal automatically retains all buffers until the command buffer completes.

Performance impact (M3 Max 48GB):

Model	Unretained (current)	Retained (fix)
Qwen3-8B-4bit (standard model, no custom kernel)	77.0 tok/s	76.9 tok/s
Qwen3.5-35B-A3B MoE (custom kernel + workaround)	21 tok/s	48 tok/s

No regression on standard models. Models using mx.fast.metal_kernel see 2.3x speedup because the mx.eval() workaround after every custom kernel call is no longer needed.

Branch with the fix: https://github.com/Ziqiao-git/mlx/tree/fix/retained-command-buffer

[Ziqiao-git]

MTL_DEBUG_LAYER=1 output

Running the reproduction script with Metal API Validation:

MTL_DEBUG_LAYER=1 python -c "
import mlx.core as mx
from paroquant.inference.backends.mlx.load import load
from paroquant.inference.backends.mlx.modules import RotateQuantizedLinear
from mlx_lm.models.cache import make_prompt_cache

model, proc, _ = load('z-lab/Qwen3.5-4B-PARO', force_text=True)
for _, m in model.named_modules():
    if isinstance(m, RotateQuantizedLinear):
        m._force_eval = False

inner = model.language_model.model
h = inner.embed_tokens(mx.array([[0]]))
mx.eval(h)

cache = make_prompt_cache(model.language_model)
out = inner.layers[0](h, mask=None, cache=cache[0])
mx.eval(out)
"

Output:

Metal API Validation Enabled
-[MTLDebugCommandBuffer preCommit]:1167: failed assertion `command buffer references deallocated object which previously existed at address 0x484350c80.'

This reproduces consistently on every run with MTL_DEBUG_LAYER=1. Without validation, the same code produces non-deterministic NaN (~2% per forward pass).

Environment:

• macOS 15.x, Apple M3 Max 48GB
• MLX 0.31.1 (official pip release)
• pip install paroquant + model z-lab/Qwen3.5-4B-PARO

Note on the reproduction dependency: The bug requires mx.fast.metal_kernel to be called multiple times in a deep lazy graph. ParoQuant's RotateQuantizedLinear (pairwise Givens rotation via mx.fast.metal_kernel + mx.quantized_matmul) inside GatedDeltaNet layers triggers this reliably. If a self-contained repro without ParoQuant would be helpful, I can try to construct one using only mx.fast.metal_kernel directly.

[Ziqiao-git]

Root cause update

After further debugging I found this is not a use-after-free caused by commandBufferWithUnretainedReferences() (as I initially suspected in #3378). The real root cause is in custom_kernel.cpp's CustomKernelCache.

What actually happens

1. mx.fast.metal_kernel returns a callable that regenerates the full kernel source (write_signature() + user body) on every call, using the current inputs' dtype to form the function signature.
2. The same kernel name can therefore produce different source strings across calls — e.g. if one call has a float16 input and another has a float32 input, write_signature() emits const device float16_t* x vs const device float* x.
3. CustomKernel::eval_gpu() looks up the kernel in CustomKernelCache by name only. When it detects it->second != source_ it calls d.clear_library(name_), which drops the MTL::Library and all cached NS::SharedPtr<MTL::ComputePipelineState> for that library.
4. If a previous dispatch with the old pipeline state is still sitting in an uncommitted command buffer, clear_library releases it out from under the command buffer → Metal validation reports command buffer references deallocated object. Without validation you get non-deterministic NaN.

Concrete trigger in ParoQuant

GET_KERNEL paro_rotate_r1 lib=A kernel=K1          # in_proj_qkv: x is float16
GET_KERNEL paro_rotate_r1 lib=A kernel=K1          # in_proj_z:   x is float16
GET_KERNEL gated_delta_step lib=B kernel=K2
SOURCE_MISMATCH paro_rotate_r1 cached=3131 new=3123 # out_proj now receives float32 x
CLEAR_LIBRARY paro_rotate_r1                       # releases K1 while a command buffer still references it
GET_KERNEL paro_rotate_r1 lib=A' kernel=K3         # new float-variant pipeline state

The dtype change comes from ParoQuant storing scales as float32 (biases as float16): mx.quantized_matmul(float16 x, uint32 w, float32 scales, float16 biases) promotes output to float32. That propagates through conv1d → gated_delta_update → RMSNormGated → out_proj's input. Same rotation.metal body, different signature.

Proposed fix

Keep per-source variants coexisting in the cache instead of clearing the old one:

-  {
-    // Clear kernels from the device library cache if needed
-    auto& kernel_cache = cache();
-    if (auto it = kernel_cache.libraries.find(name_);
-        it != kernel_cache.libraries.end()) {
-      if (it->second != source_) {
-        auto& d = metal::device(s.device);
-        d.clear_library(name_);
-        it->second = source_;
-      }
-    } else {
-      kernel_cache.libraries.emplace(name_, source_);
-    }
-  }
-
-  auto lib = d.get_library(name_, [this] { return metal::utils() + source_; });
-  auto kernel = d.get_kernel(name_, lib);
+  // Use a source-dependent library key so different source variants
+  // (e.g. from write_signature picking different dtype qualifiers across
+  // calls) coexist without evicting each other. Clearing a library while
+  // its pipeline states are still referenced by an in-flight command
+  // buffer causes use-after-free.
+  auto source_hash = std::hash<std::string>{}(source_);
+  auto lib_key = name_ + "_" + std::to_string(source_hash);
+
+  auto lib = d.get_library(lib_key, [this] { return metal::utils() + source_; });
+  auto kernel = d.get_kernel(name_, lib);

This:

• Removes the clear_library path entirely, so no pipeline state is released while still in flight.
• Keeps name_ as the kernel-function symbol (no ABI / dispatch change).
• Doesn't touch commandBufferWithUnretainedReferences, so no performance trade-off on older hardware.

Trade-offs / open questions

• Cache growth: each unique source variant now keeps a live MTL::Library and pipeline state. In practice, variants are bounded by the small number of (dtype, shape-info) combinations used in a model (ParoQuant sees 2–3 variants per model). But for pathological callers this is unbounded — happy to add an LRU bound if preferred.
• Hash collisions: std::hash<std::string> is 64-bit. Collisions are astronomically unlikely for real source strings but not impossible. Using the source string itself as the map key (instead of a hash) is safer at the cost of memory; I defaulted to hash for compactness.
• Why name_ for get_kernel: the Metal function symbol inside each library is custom_kernel_<name>, which is the same across variants, so get_kernel(name_, lib) correctly resolves per-library.

Verification

• MTL_DEBUG_LAYER=1 reproduction from [BUG] mx.fast.metal_kernel: use-after-free when multiple custom kernels compose in lazy graph #3347 no longer triggers the deallocated object assertion.
• Qwen3.5-4B-PARO: 5/5 correct, 91 tok/s (previously 30 tok/s with the _force_eval workaround).
• Qwen3.5-35B-A3B-PARO: 5/5 correct, 48 tok/s (previously 21 tok/s).
• Qwen3-8B-4bit (no custom kernels): 75.1 tok/s patched vs 75.8 tok/s on 0.31.1 — within measurement noise.

Happy to open a fresh PR with this custom_kernel.cpp-only change (distinct from the commandBuffer() direction in #3378) if the approach looks acceptable.

[Ziqiao-git]

Why only Qwen3.5 (GatedDeltaNet / mamba hybrid) triggers this

The bug requires the same custom kernel name to be invoked with different input dtypes inside a single lazy graph — that is what causes write_signature to emit different source strings and trigger the cache eviction path.

Standard transformer attention avoids this because all four projections see the same input dtype:

normed (float16)
  → q_proj(normed)   # float16 input
  → k_proj(normed)   # float16 input
  → v_proj(normed)   # float16 input
  → SDPA (dtype-preserving)
  → o_proj(...)      # float16 input

All four RotateQuantizedLinear calls receive float16 → same generated source → no cache miss.

GatedDeltaNet breaks dtype homogeneity within a single layer:

normed (float16)
  → in_proj_qkv(normed)     # float16 input
  → in_proj_z(normed)       # float16 input
  → mx.quantized_matmul promotes to float32 (scales are float32)
  → conv1d → gated_delta_update → RMSNormGated (all preserve float32)
  → out_proj(norm_out)      # float32 input   ← source mismatch!

The first two rotations get float16, the last one gets float32. RMSNormGated has no implicit downcast back to float16, so the dtype propagates all the way to out_proj. That's the single point where the cache sees it->second != source_ and calls clear_library.

So the bug is not strictly "mamba-only" — it can manifest anywhere the same metal_kernel is called with mixed dtypes in one graph — but GDN's internal pipeline is the natural case that triggers it in practice. PR #3434 fixes the cache so any such dtype variation works without use-after-free.