extro.rook.host
Fork for thermals request add-json-schema-dpeq
Serialize GPU-stream drivers to fix concurrent-encoder crashes
MLX is not thread-safe: mlx_eval/mlx_async_eval run gpu::eval inline on the
calling thread, encoding into and committing the shared Metal command buffer,
and the C++ stream mutex only guards the stream-map lookup (not the encoder).
Since evalLock was dropped from the eval hot path, a request's prefill
async_eval could run concurrently with a teardown/unload Stream.synchronize()
on the same command buffer.
Live-reproduced SIGABRT (disconnect-during-cold-load + retry): 'Completed
handler provided after commit' in addCompletedHandler, and AGX 'command
encoder is already encoding'. Two threads: unload->Stream.gpu.synchronize()
committing the buffer while startSoloFastPath->TokenIterator.prepare->async_eval
added a completion handler to it.
Fix: hold evalLock across the brief CPU-side encode+commit in eval/asyncEval/
MLXArray.eval (Stream.synchronize/clearCache already take it). The GPU executes
asynchronously after the lock releases, so CPU->GPU pipeline parallelism is
preserved; the lock is uncontended in steady-state single-producer decode and
only serializes during teardown/cancel/cross-model.
Also guard compiled activations (safeGeluApproximate/compiledSwiGLU/compiledGeGLU)
with CompiledDecodeTrace.isActive to avoid an illegal nested compile inside the
compiled-decode trace, and drain the stream on generateLoopTask cancel/error
early-exits.
Validated live: disconnect repro 35/35 clean (was crash at iter 1-2), concurrent
GPU 15/15 clean, decode 76.6 tok/s (no regression), multiturn coherent.
author
Eric
date
(Jul 5, 2026, 6:11 PM -0700)
commit
e0e2eb6e
e0e2eb6e6caffa198a3377651c9c03c4b2ec3aef
parent
de5f185c
de5f185c61baba84836861f96da214c4cad1968f
+75
-17
+15
Libraries/MLXLMCommon/Evaluate.swift
+15
Libraries/MLXLMCommon/Evaluate.swift
···
3328
3328
// the chunked prepare loops bail between chunks. Not an error;
3329
3329
// finish the stream with a `.cancelled` info like any other
3330
3330
// cancellation.
3331
+
//
3332
+
// Drain the shared GPU stream before finishing: the chunked
3333
+
// `prepare` may have already enqueued prompt-prefill command
3334
+
// buffers on the default stream before it bailed. Closing the
3335
+
// stream lets the consumer (or an unload/teardown) start the
3336
+
// next producer immediately; if we return without draining, that
3337
+
// producer opens an encoder while our half-built prefill buffer
3338
+
// is still live on the same stream — the cold-load-disconnect
3339
+
// "command encoder is already encoding" / end_encoding races.
3340
+
// The normal completion path below drains twice for the same
3341
+
// reason; the early-exit paths must match it.
3342
+
Stream().synchronize()
3331
3343
handler.onGenerationEnd(emit: continuation.yield)
3332
3344
_ = continuation.yield(handler.infoEvent(GenerateCompletionInfo(
3333
3345
promptTokenCount: promptTokenCount,
···
3341
3353
} catch {
3342
3354
Logger(subsystem: "vmlx", category: "generateLoopTask").error(
3343
3355
"Iterator construction failed: \(error.localizedDescription, privacy: .public)")
3356
+
// Drain any prefill work enqueued before the failure before
3357
+
// closing the stream (see the CancellationError branch above).
3358
+
Stream().synchronize()
3344
3359
handler.onGenerationEnd(emit: continuation.yield)
3345
3360
_ = continuation.yield(handler.infoEvent(GenerateCompletionInfo(
3346
3361
promptTokenCount: promptTokenCount,
+21
-12
Libraries/MLXLMCommon/SwitchLayers.swift
+21
-12
Libraries/MLXLMCommon/SwitchLayers.swift
···
12
12
let body: @Sendable (MLXArray) -> MLXArray = { (x: MLXArray) -> MLXArray in
13
13
0.5 * x * (1 + tanh(sqrt(2 / Float.pi) * (x + 0.044715 * x * x * x)))
14
14
}
15
-
if HardwareInfo.isCompiledDecodeSupported {
16
-
return compile(shapeless: true, body)
17
-
}
18
-
return body
15
+
guard HardwareInfo.isCompiledDecodeSupported else { return body }
16
+
let compiled = compile(shapeless: true, body)
17
+
// When this activation is invoked *inside* the outer compiled-decode trace
18
+
// (`setupCompiledDecode` → `CompiledDecodeTrace.withActive`), calling a
19
+
// separately-compiled function is a nested compile — illegal, exactly like
20
+
// `eval` during a trace (see the `!CompiledDecodeTrace.isActive` guards in
21
+
// Gemma4Text). The inner `compileState.call` returns an empty result and
22
+
// `[0]` traps (Transforms+Compile.swift). Run the plain body while tracing:
23
+
// its ops are captured into the outer graph and fused there, so there is no
24
+
// throughput loss — the inner compile was both illegal and redundant.
25
+
return { x in CompiledDecodeTrace.isActive ? body(x) : compiled(x) }
19
26
}()
20
27
21
28
/// Drop-in replacement for MLXNN.GELU that avoids the Power primitive crash.
···
35
42
(gate: MLXArray, x: MLXArray) -> MLXArray in
36
43
silu(gate) * x
37
44
}
38
-
if HardwareInfo.isCompiledDecodeSupported {
39
-
return compile(shapeless: true, body)
40
-
}
41
-
return body
45
+
guard HardwareInfo.isCompiledDecodeSupported else { return body }
46
+
let compiled = compile(shapeless: true, body)
47
+
// Fall back to the plain body inside the outer compiled-decode trace to
48
+
// avoid an illegal nested compile (see `safeGeluApproximate`).
49
+
return { g, x in CompiledDecodeTrace.isActive ? body(g, x) : compiled(g, x) }
42
50
}()
43
51
44
52
private let compiledGeGLU: @Sendable (MLXArray, MLXArray) -> MLXArray = {
···
46
54
(gate: MLXArray, x: MLXArray) -> MLXArray in
47
55
(0.5 * gate * (1 + tanh(sqrt(2 / Float.pi) * (gate + 0.044715 * gate * gate * gate)))) * x
48
56
}
49
-
if HardwareInfo.isCompiledDecodeSupported {
50
-
return compile(shapeless: true, body)
51
-
}
52
-
return body
57
+
guard HardwareInfo.isCompiledDecodeSupported else { return body }
58
+
let compiled = compile(shapeless: true, body)
59
+
// Fall back to the plain body inside the outer compiled-decode trace to
60
+
// avoid an illegal nested compile (see `safeGeluApproximate`).
61
+
return { g, x in CompiledDecodeTrace.isActive ? body(g, x) : compiled(g, x) }
53
62
}()
54
63
55
64
public func gatherSort(x: MLXArray, indices: MLXArray) -> (MLXArray, MLXArray, MLXArray) {
+8
-1
Source/MLX/MLXArray.swift
+8
-1
Source/MLX/MLXArray.swift
···
554
554
/// MLX is lazy and arrays are not fully realized until they are evaluated. This method is typically
555
555
/// not needed as all reads ensure the contents are evaluated.
556
556
public func eval() {
557
-
// No evalLock — C++ scheduler is thread-safe internally
557
+
// Serialize the CPU-side encode+commit against Stream.synchronize() and
558
+
// other GPU-stream drivers: MLX is NOT thread-safe (mlx_array_eval runs
559
+
// gpu::eval inline on this thread, mutating the shared command buffer),
560
+
// and the C++ stream mutex only guards the stream-map lookup, not the
561
+
// encoder. See the evalLock note in Transforms+Eval.swift. Held only
562
+
// across the brief commit, so GPU pipeline parallelism is preserved.
563
+
evalLock.lock()
558
564
mlx_array_eval(ctx)
565
+
evalLock.unlock()
559
566
}
560
567
561
568
/// Replace the contents with a reference to a new array (INTERNAL).
+31
-4
Source/MLX/Transforms+Eval.swift
+31
-4
Source/MLX/Transforms+Eval.swift
···
3
3
import Cmlx
4
4
import Foundation
5
5
6
-
/// Lock for operations that modify MLX global state (compile, stream creation).
7
-
/// NOT used for eval/asyncEval — the C++ scheduler has its own std::mutex.
8
-
/// Removing evalLock from the eval hot path allows asyncEval and item() to
9
-
/// overlap, enabling Metal command pipeline parallelism.
6
+
/// Serializes every CPU-side driver of the shared Metal command stream:
7
+
/// `eval`/`asyncEval`/`item` here, plus `Stream.synchronize()`,
8
+
/// `Memory.clearCache()`, `compile`, and stream lifecycle elsewhere.
9
+
///
10
+
/// This is REQUIRED for correctness, not just for the global-state ops.
11
+
/// MLX is not thread-safe: `mlx_eval`/`mlx_async_eval` run `gpu::eval`
12
+
/// *inline on the calling thread* (the per-stream `StreamThread` stays idle
13
+
/// for the default GPU stream), encoding into and committing the stream's
14
+
/// single `MTLCommandBuffer`. `Stream.synchronize()` ends+commits that same
15
+
/// buffer. The fork's C++ `stream_map_mtx_` only guards the stream-map lookup
16
+
/// (fix 9dabb6c4), NOT the encoder, so two Swift threads — e.g. a request's
17
+
/// prefill `async_eval` and a concurrent unload/`strictEvict`
18
+
/// `Stream.gpu.synchronize()` — otherwise mutate the same command buffer
19
+
/// concurrently. Live-reproduced SIGABRT: "Completed handler provided after
20
+
/// commit" (`addCompletedHandler` on a buffer synchronize just committed) and
21
+
/// AGX "command encoder is already encoding".
22
+
///
23
+
/// The lock is held ONLY across the brief CPU-side encode+commit
24
+
/// (`mlx_eval`/`mlx_async_eval` return once the buffer is committed; the GPU
25
+
/// executes asynchronously afterward with the lock already released), so the
26
+
/// CPU→GPU pipeline parallelism that motivated dropping the lock from this
27
+
/// path is preserved: in steady-state single-producer decode the lock is
28
+
/// uncontended, and it only serializes during teardown/cancel/cross-model —
29
+
/// exactly when a second thread would otherwise corrupt the encoder.
30
+
/// Recursive so a lock-holding path that re-enters eval does not self-deadlock.
10
31
let evalLock = NSRecursiveLock()
11
32
12
33
/// Evaluate one or more `MLXArray`
···
15
36
/// - <doc:lazy-evaluation>
16
37
public func eval(_ arrays: MLXArray...) {
17
38
let vector_array = newEvalVectorArray(arrays)
39
+
evalLock.lock()
18
40
mlx_eval(vector_array)
41
+
evalLock.unlock()
19
42
mlx_vector_array_free(vector_array)
20
43
}
21
44
···
25
48
/// - <doc:lazy-evaluation>
26
49
public func eval(_ arrays: some Collection<MLXArray>) {
27
50
let vector_array = newEvalVectorArray(arrays)
51
+
evalLock.lock()
28
52
mlx_eval(vector_array)
53
+
evalLock.unlock()
29
54
mlx_vector_array_free(vector_array)
30
55
}
31
56
···
36
61
/// - ``asyncEval(_:)-(Collection<MLXArray>)``
37
62
public func asyncEval(_ arrays: some Collection<MLXArray>) {
38
63
let vector_array = newEvalVectorArray(arrays)
64
+
evalLock.lock()
39
65
mlx_async_eval(vector_array)
66
+
evalLock.unlock()
40
67
mlx_vector_array_free(vector_array)
41
68
}
42
69