Repository navigation
fix: run compiled functions on the task-scoped default device - #17
Merged
Merged
Conversation
The backend keys its compile cache on, and builds the Compiled primitive for, the C++ default stream. Device.withDefaultDevice only sets a Swift TaskLocal, so a function first traced on the GPU was reused inside a CPU-scoped call and ran its fused kernel on the GPU against CPU-stream inputs. Set the C++ default device for the duration of the call. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This was referenced Sep 24, 2026
Closed
solderzzc
pushed a commit
to SharpAI/SwiftLM
that referenced
this pull request
Sep 24, 2026
…eam-experts - TurboKV-on-M5 note → tracking issue #175. - --gpu-layers GPU timeout: reproduced on M5 as well as M6; tracked in #176 with the fix in SharpAI/mlx-swift#17. CPU-resident MoE layers stay very slow after the fix (~0.4 tok/s prefill), so the note points 32 GB users at --stream-experts first. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Compiled functions ignore
Device.withDefaultDevice. A function first traced on the GPU is reused inside a CPU-scoped call and runs its fused kernel on the GPU.Why
Device.withDefaultDevice(_:_:)sets only a SwiftTaskLocal; the C++ default device is not touched. The backend's compile cache, however, is keyed on the C++ default stream (default_stream(default_device())incompile.cpp), and that stream is also the one theCompiledprimitive is built for.In SwiftLM this breaks
--gpu-layers N.partitionedLayerCallruns the CPU layers underwithDefaultDevice(.cpu). Gemma 4's MoE calls compiled helpers there, such asweightedExpertSumand the activation products. Those helpers were traced in a GPU layer, so they run on the GPU, with inputs produced by CPUGatherQMM. The GPU command buffer then waits on the CPU fence for longer than the Metal watchdog allows, and the process aborts withkIOGPUCommandBufferCallbackErrorTimeout.Change
CompiledFunction.innerCallnow sets the C++ default device toDevice.defaultDevice()aroundmlx_closure_apply, and restores the previous device afterwards. This happens whileevalLockis held. With the correct device set, a CPU-scoped call gets its own cache entry and traces for the CPU stream.Tests
testCompileHonorsScopedDefaultDevice: traces on the GPU, then calls underwithDefaultDevice(.cpu), and asserts that the call retraces and that a later GPU call reuses the original entry. It fails without this change (1 trace where 2 are expected).TransformTests: 19/19 pass.gemma-4-26b-a4b-it-4bit --gpu-layers 23on an M5 Pro now completes a 490-token request instead of aborting after ~105 s. Full-GPU Qwen3.8-27B is unchanged.🤖 Generated with Claude Code