Repository navigation
Merge master, fix MPS TopK at k=16, a GEMM kernel race and GPU error handling - #1
Merged
Merged
Conversation
Skip the invalidation store when a reduction returns NOT_FOUND. Add regression coverage for exhausted rows, masked values, and multiple k values. Validation: 12 TopK tests passed on Windows CUDA. AI-assisted analysis, implementation, and test execution.
…penNMT#2108) * fix(python): use pybind11 with CPython 3.13+ tp_doc allocator support * fix(wheels): align wheel build prerequisites with allocator-safe pybind11 * ci(python): guard wheel publishing on CPython 3.14 debug allocator exit * ci(python): select regular CPython 3.14 wheel explicitly Distinguish the cp314 ABI from cp314t when both release artifacts exist. Keep the exactly-one wheel check and debug-allocator shutdown gate. Verified filename-only probes for coexistence, missing, and duplicate wheels.
…NMT#2110) * Update replica_pool.h * Update test_load_model_from_memory() * restart tests
* Integrate Clang ThreadSanitizer in tests * Remove unnecessary debug flag from TSAN_FLAGS
* Add UMT5 model loader and support for unique relative attention biases * Declare _shared_position_bias on the Transformer stacks * Update documentation and tests for UMT5 model support * Add runtime tests for UMT5 model to verify per-layer bias usage * Refactor import statements in _save_umt5_encoder for cleaner code --------- Co-authored-by: Jordi Mas <jmas@softcatala.org>
Changes: - Test Linux and Windows ROCm wheels across the Python matrix, installing ROCm runtimes only where needed. - Gate PyPI, Docker, documentation, and ROCm release publishing on relevant tests and style checks. - Upgrade wheel build actions to Node.js 24: checkout@v5, cache@v5, and upload-artifact@v6; local upload tests passed.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Master's TopKWithExhaustedCandidates test (OpenNMT#2104) fails on MPS for k=16: rows with fewer than k finite candidates came back with 0 in the exhausted slots. For 8 < k <= 16 with 16-byte aligned result rows the backend used MPSMatrixFindTopK, which pads exhausted slots with 0 and a bogus index instead of returning the masked input; a beam search with 2 * beam_size = 16 would read that 0 as the best score. small_topk already handles exhaustion (masked candidates beat the -INFINITY/UINT_MAX initial slots on index), so it now covers every 1 < k <= 16 for float32, float16 and bfloat16, and the MPSMatrixFindTopK path and its thread-local kernel cache are removed. k > 16 keeps the CPU fallback. Adds FusedBeamSearchTopKWithExhaustedCandidates, which checks the opt-in fused beam search TopK against the materialized path on rows with fewer than k finite candidates, for (beam, k) = (4, 8), (5, 10), (8, 16). Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A translator on MPS with inter_threads=4 and four callers crashed with SIGSEGV inside MPSSetNDArraysOnComputeEncoder, from mps::gemm. gemm_kernel() kept one MPSMatrixMultiplication per shape for the whole process; the thread_local map only cached pointers to those shared objects. Apple documents MPSKernel objects as unsafe to use from several threads at once, and workers of one translator run the same shapes. The cache is now thread_local and owns its kernels, released when the thread ends. Each encoded kernel is recorded as in use so it outlives the command buffer even if its thread exits first. A nil kernel now throws instead of being cached and silently doing nothing. Adds ConcurrentGemmFloat16/Float32WithTheSameShape and ConcurrentBatchedMatMulFloat16WithTheSameShape: four threads, the same 64x512 by 512x2048 shape, 300 iterations each, checked against the CPU. They crashed 10 times out of 10 before this change. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…from idle workers or destructors A failed asynchronous command buffer recorded its error in one process-wide slot, and whichever thread synchronized next threw it: an innocent replica could fail with another one's error. Worse, the throw could come from ReplicaWorker::idle(), outside any job, ending the worker thread and the process with std::terminate, or from Model::~Model through synchronize_device. CUDA ignores errors on both paths. Each stream now owns its error in a small shared state captured by the completion handler, so a handler that runs after its stream is gone is safe, and only that stream's next command_buffer() or synchronize() throws it, once. A per-command-buffer token stops a late handler from recording a failure a second time. synchronize_all() reports one stream's error and leaves the others to their owners. idle() calls the new synchronize_stream_nothrow(), identical to synchronize_stream() on CPU and CUDA; on MPS it logs once and keeps the error for the stream's next operation. synchronize_device(MPS), which only ~Model and move_variables call, uses the same no-throw barrier; the wait itself is unchanged. tests/mps_error_test.cc injects errors through an internal hook (src/mps/testing.h). On the previous code 5 of its 10 tests fail and 3 abort with std::terminate (ReplicaWorker::idle, a Translator, ~Model). The backend stays experimental and opt-in: a warning when a model loads on MPS, "mps (experimental)" in the device docstrings and the CLI help, a note in README.md and docs/environment_variables.md, and a CHANGELOG entry. WITH_MPS stays OFF by default and out of the wheel scripts. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Owner
|
@nerln I made one small follow-up to the TopK implementation. Using 16 private result slots for every TopK call increased register pressure and made the common beam-4 The full C++ suite passes with 416 tests passed and one existing skip. I preserved your original commits and authorship and added the optimization as a separate follow-up commit. Thanks for identifying and addressing these issues. |
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.
This brings
apple-mps-backendup to today's master and fixes three things that turned up while checking the list in OpenNMT#2077 (comment). Pulling it into your branch should make OpenNMT#2077 mergeable again.Merge with master (51526bf). Two textual conflicts:
src/devices.cc(master's Windows thread-pool cleanup kept, in this branch's indentation) andtests/ops_test.cc(include lists). Master's OpenNMT#2104 broughtOpDeviceFPTest.TopKWithExhaustedCandidates, which failed on MPS.Use the custom MPS TopK kernel up to k=16. For 8 < k <= 16 with 16-byte aligned rows the backend usedMPSMatrixFindTopK, which fills the slots of a row with fewer than k finite candidates with 0 and a bogus index. A beam search with2 * beam_size = 16would read that 0 as the best score.small_topkalready handled exhaustion, so it now covers every 1 < k <= 16 (float32, float16, bfloat16) and theMPSMatrixFindTopKpath and its thread-local cache are gone. AddsFusedBeamSearchTopKWithExhaustedCandidatesfor the opt-in fused path, which was already correct.Give each thread its own MPS matrix multiplication kernels. A translator on MPS withinter_threads=4and four callers died with SIGSEGV inMPSSetNDArraysOnComputeEncoder(frommps::gemm):gemm_kernel()shared oneMPSMatrixMultiplicationper shape across threads, and MPSKernel objects must not be encoded from several threads at once. The cache is now thread-local and owns its kernels; each encoded kernel is kept alive until its command buffer completes. Three newConcurrent*WithTheSameShapetests crashed 10 times out of 10 before the change.Keep MPS command buffer errors on their stream, and never throw them from idle workers or destructors. A failed command buffer stored its error in one process-wide slot, so another thread could receive it, and the throw could come fromReplicaWorker::idle()(endingWorker::runwithstd::terminate) or fromModel::~Model. Errors now live on the stream that submitted the buffer and are raised once by that stream's next operation;idle()uses a newsynchronize_stream_nothrow()(identical tosynchronize_streamon CPU and CUDA) andsynchronize_device(MPS)uses a no-throw barrier.tests/mps_error_test.ccinjects errors through an internal hook (src/mps/testing.h); 8 of its 10 tests fail or abort on the previous code. The same commit adds the opt-in/experimental notes: a warning at model load on MPS,mps (experimental)in the device docstrings and CLI help, notes inREADME.mdanddocs/environment_variables.md, and a CHANGELOG entry crediting you (please check the handle).On an M4 (macOS 27),
-DWITH_MPS=ON -DWITH_ACCELERATE=ON -DWITH_MKL=OFF -DWITH_RUY=ON -DOPENMP_RUNTIME=NONE: C++ suite 417 tests, 416 passed, 1 skipped;CT2_TEST_MPS=1Python MPS tests 6 passed. CI on this head in my fork, all jobs green: https://github.com/nerln/CTranslate2/actions/runs/38498483339🤖 Generated with Claude Code