milnor_gpu: drop the global device mutex (thread-local resident, per-thread streams)#15
Draft
JoeyBF wants to merge 6 commits into
Draft
milnor_gpu: drop the global device mutex (thread-local resident, per-thread streams)#15JoeyBF wants to merge 6 commits into
JoeyBF wants to merge 6 commits into
Conversation
…hanism" This reverts commit 71a21b6.
The relaxed wavefront keeps many bidegrees in flight at once, so at any instant it is likely that some job is inside a linear-algebra critical section (ParallelGuard). The scheduler re-spawned a bounced job immediately, which just re-checked is_in_parallel, found it still busy, and bounced again — spawning a whole rayon job per re-check and pegging every core on a retry storm that does no useful work. Instead the receiver checks the flag itself (a cheap atomic load) and parks a bidegree only when the section is genuinely busy. A job acquires and releases its guards many times and spends most of its time outside them, so the section frees far more often than jobs complete; parked bidegrees are therefore re-checked via a short recv_timeout while anything is parked, and re-spawned as soon as the section frees. Incoming messages are still handled the instant they arrive. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MCUtWj6P6suSZqvATCdg6d
is_in_parallel was a global count of active par_iter critical sections, so a step_resolution job bounced whenever *any* thread was in one. Under the relaxed wavefront many bidegrees are in flight, so that flag is almost always set and nearly every job bounced, producing the retry churn the parking mitigation only softened. The priority inversion the guard exists to prevent is narrower: a worker that initiated a par_iter blocks in the join and work-steals, and if it steals another (heavy, nested-parallel) resolution step, that step stalls the section the worker is blocked on. A stolen job runs on the stealer's own OS thread, so a thread-local depth counter reports exactly whether *this* worker is a blocked guard holder. Jobs picked up by a free worker read zero and run, letting independent bidegrees resolve concurrently instead of serializing behind any single critical section. The scheduler thread never holds a guard, so it can no longer read the flag to sense saturation; park bounced bidegrees and retry them on each completion or a short recv_timeout tick. Bounces are now rare (only a genuine steal-onto-a-blocked-holder), so the parking path barely engages. The classical scheduler shares the guard and benefits the same way, so its immediate-respawn no longer storms. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MCUtWj6P6suSZqvATCdg6d
Pins the invariant the previous commit relies on — a ParallelGuard held on one thread reads as absent on another — so a future change that reverts to a shared counter fails loudly instead of silently reintroducing the retry storm. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MCUtWj6P6suSZqvATCdg6d
…-graph-d69k0s' into hpc
…ead-local
The batched multiply serialized every launch behind one RESIDENT mutex held
across the whole marshal+upload+kernel+readback section, and additionally
pinned all work to CUDA stream 0. With the relaxed dependency graph exposing
~max_s-wide bidegree parallelism, that lock collapsed a ~12-core CPU wavefront
to ~2.6 busy cores and left the GPU idle 80% of the time — making NASSAU_GPU=1
a net 1.4x slowdown over CPU-only at stem 130 (193s vs 142s).
cubecl 0.10 does not need the lock: a per-device runner thread already
serializes all server access (concurrent client calls are memory-safe), and
memory pools are per-stream. So:
- RESIDENT becomes a thread_local RefCell: each rayon worker keeps its own
admissible cache and cs/mk device handles, created and consumed only on the
thread (and thus the default per-thread CUDA stream) that owns them, so no
handle ever crosses threads and no cross-stream event sync fires.
- The GPU_STREAM{value:0}.executes pin is removed; each worker launches on its
own default stream, so independent bidegrees marshal and execute
concurrently. memory_cleanup now trims only the calling worker's pool.
Stem 130 (S_2, s<=152, 16-core H200 box): 193s/2.6 cores (old mutex GPU) and
142s/10 cores (CPU-only) -> 44-49s/5.6 cores. Verified bit-identical to the
CPU path with NASSAU_GPU_VERIFY=1 at stem 80 (MIN_WORK=0, every build) and
stem 130 (default gate, all offloaded/chunked launches, concurrent workers).
Note: concurrency raises peak host memory (concurrent marshal buffers across
workers); a 16-worker VERIFY run at stem 130 exceeded a ~48GB cgroup, while
normal runs fit comfortably. Bound RAYON_NUM_THREADS if memory-constrained.
Co-Authored-By: Claude Opus 4.8 <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
2 participants
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.
See commit message: removes the RESIDENT mutex + stream-0 pin; thread-local resident store, per-thread CUDA streams. Stem-130: 193s/2.6 cores -> 44-49s/5.6 cores (CPU-only: 142s/10). Verified bit-identical (VERIFY at stem 80 all-builds and stem 130 default-gate, concurrent). Note: verify at 16 workers exceeded a ~48GB cgroup (normal runs fit); bound RAYON_NUM_THREADS if memory-constrained. Includes the local guard-relaxation commits not yet on origin/hpc.
🤖 Generated with Claude Code