Fix Metal fence ordering and sealed collective portability

.github/workflows/fork-tests.yml

@@ -9,7 +9,7 @@ name: Fork Tests
 #                  with the GPU and then with the CPU as the default device,
 #                  and the JACCL tests with simulated verbs (tests/jaccl).
 #                  Then it runs the standalone allocation cap tests, one
-#                  process for each case.
+#                  process for each case, and the bounded Metal fence chains.
 #   python-gpu     runs the Python tests with the GPU as the default device.
 #   python-cpu     runs the Python tests with the CPU as the default device.
 #   metal-build-and-test
@@ -167,6 +167,28 @@ jobs:
           ctest --test-dir build -R '^process_allocation_cap' -V \
             --no-tests=error
 
+      - name: Run the bounded Metal fence chains
+        # metal_fence_chain runs exact GPU -> CPU -> GPU chains with fast
+        # synchronization at three command buffer sizes, one ctest entry for
+        # each, each with a 60 s timeout (tests/CMakeLists.txt).
+        if: ${{ !cancelled() && steps.build.outcome == 'success' }}
+        env:
+          LLVM_PROFILE_FILE: ${{ runner.temp }}/profiles/fence-%p-%m.profraw
+        run: |
+          ctest --test-dir build -R '^metal_fence_chain_' --output-on-failure \
+            --no-tests=error
+
+      - name: Run the quantized matmul rows with MLX_QMV_BATCH_INVARIANT=1
+        # MLX reads the switch once in each process. The C++ test steps above
+        # run the case with the switch off; this ctest entry runs it again
+        # with the switch on (tests/CMakeLists.txt).
+        if: ${{ !cancelled() && steps.build.outcome == 'success' }}
+        env:
+          LLVM_PROFILE_FILE: ${{ runner.temp }}/profiles/qmv-%p-%m.profraw
+        run: |
+          ctest --test-dir build -R '^qmv_batch_invariant$' \
+            --output-on-failure --no-tests=error
+
       - name: Merge the coverage profiles
         if: ${{ !cancelled() && steps.build.outcome == 'success' }}
         run: |

docs/src/usage/environment_variables.rst

@@ -143,6 +143,20 @@ Metal
    it to ``0`` to use the GEMM convolution for all inputs. See
    :doc:`precision`.
 
+.. envvar:: MLX_QMV_BATCH_INVARIANT
+
+   Set it to ``1`` to compute a quantized matrix product ``x @ w.T`` with
+   fewer rows than the vector limit only with the vector kernels: ``qmv`` for
+   one row and ``qmv_wide`` from two rows. Then each row gets the same bits at
+   each row count from 2 to the limit - 1, so the rows of different requests
+   can share one product. The default, ``0``, also uses faster kernels that
+   give a row different bits at some row counts: ``qmv_mma`` for 6 to 16 rows
+   of the floating-point modes, and the affine split-K kernel from 8 rows on
+   GPUs with neural accelerators. The vector limit is 6 to 33 rows; it
+   depends on the GPU and on the shape. A row alone (``qmv``) can get bits
+   that are different from the bits of the same row among two or more rows.
+   MLX reads this variable once in each process.
+
 Advanced tuning
 ---------------
 
@@ -179,6 +193,26 @@ users. Their behavior may change as the implementation evolves.
    dot-product attention kernel. Positive values are rounded up to a multiple
    of ``32``.
 
+.. envvar:: MLX_QMM_STEEL_TILE
+
+   Output tile of the Metal steel affine quantized matmul (``x @ w.T`` with
+   at least 64 rows and a multiple of 64 outputs) where no neural-accelerator
+   kernel serves the product. The default, ``0``, takes 64 rows x 32 columns
+   for 16-bit inputs and 32 x 64 for FP32 inputs from 256 rows, and 32 x 32
+   below. ``32`` (32 x 32, the previous tile), ``64`` (64 x 64), ``6432``
+   (64 rows x 32 columns), ``3264`` (32 x 64) or ``6464`` (64 x 64 at a K
+   depth of 64 for 16-bit inputs in groups of 64 or more) choose one tile for
+   every product from 64 rows. Every tile gives the same bits; only the speed
+   changes. Only builds that compile kernels at run time (``MLX_METAL_JIT=ON``)
+   take a tile other than 32 x 32: the precompiled kernel library of the
+   default build holds the 32 x 32 kernels only, so there this variable has no
+   effect.
+
+.. envvar:: MLX_QMM_STEEL_TILE_MIN_GROUPS
+
+   Fewest 64 x 64 threadgroups a product must fill to take a tile other than
+   32 x 32; smaller products keep 32 x 32. The default is ``0``.
+
 CUDA
 ----
 

docs/src/usage/using_streams.rst

@@ -16,3 +16,26 @@ the operation is run on the default stream of the default device:
 be a :obj:`Device` (e.g. ``stream=my_device``) in which case the operation is
 run on the default stream of the provided device
 ``mx.default_stream(my_device)``.
+
+Cross-device fence ordering
+---------------------------
+
+With fast Metal synchronization, a GPU consumer can wait for work on a CPU
+stream. Its wait must follow both earlier command passes and earlier kernels
+in the current concurrent pass. The CPU work can depend on those GPU kernels.
+``Fence::wait`` therefore inserts both dependencies before the GPU wait.
+
+A GPU stream never spins on another GPU stream. The fast wait is a kernel that
+spins until its producer's timestamp arrives, and Metal does not guarantee that
+a second command queue runs while a pass of the first one spins: the producing
+queue can starve and both stop. When an evaluation's graph spans two GPU
+streams, a consumer on one waits for the other on a shared event, and so does a
+GPU stream waiting for CPU work that can itself wait for another GPU stream.
+The other CPU-to-GPU and GPU-to-CPU waits keep the fast path.
+
+The standalone ``metal_fence_chain`` test checks exact GPU-to-CPU-to-GPU chains
+with shared side roots, chains where two GPU streams feed the CPU, and
+request-array retirement. CTest runs separate
+processes with fast synchronization and several command-buffer sizes, with a
+60-second timeout per process. This component test does not qualify distributed
+collective order, a transport, or an installed model service.

fork.yaml

@@ -85,6 +85,69 @@ def:
         - "mlx/backend/common/command_buffer_limits.h"
         - "mlx/backend/metal/device.h"
         - "mlx/backend/metal/device.cpp"
+    - title: "Steel affine quantized matmul: prompt-sized output tiles"
+      description: |
+        On a GPU without neural accelerators the steel `affine_qmm_t` kernel computed every
+        product in 32 x 32 output tiles (2 x 2 fragments a SIMD group). A product of at least
+        256 rows with a 64-aligned output width now takes a 64-row x 32-column tile for 16-bit
+        inputs and a 32 x 64 tile for FP32 inputs by default (measured on the M3 Ultra at 512
+        rows: BF16 products 4-6% faster; FP32's 64 x 64 tile was slower), JIT-instantiated from
+        the same `qmm_t_impl`. Every output element still accumulates the same 8-wide K slabs
+        in the same order with the same simdgroup multiply-accumulates over the same
+        dequantized values, so no bit changes; the split-K decision (made from 32 x 32 tiles)
+        and the NAX route are unchanged.
+        `MLX_QMM_STEEL_TILE` picks 32 (as before), 64, 6432, 3264 or 6464 (64 x 64 with K depth
+        64 for 16-bit inputs in groups of 64 or more; FP32 tiles that deep would exceed 32 KiB of
+        threadgroup memory); `MLX_QMM_STEEL_TILE_MIN_GROUPS` keeps 32 x 32 below a 64 x 64
+        threadgroup count; `mlx_metal_set_thread_qmm_steel_tile` overrides the tile for the
+        calling thread (tests and benches compare the tiles' bytes in one process). Both
+        variables are documented in `docs/src/usage/environment_variables.rst`.
+        Only builds that compile kernels at run time take these tiles: the precompiled
+        library instantiates `affine_qmm_t` at 32 x 32 only, so CMake builds with the default
+        `MLX_METAL_JIT=OFF` keep 32 x 32 (`kernels_jit_compiled()`, defined by
+        `jit_kernels.cpp` and `nojit_kernels.cpp`). `tests/quantized_tile_tests.cpp` runs
+        512 x 1024 products (one qmm, no split-K) at K = 96 (no NAX route) and K = 192 in
+        groups of 64 at every tile, checks the 32 x 32 result against the dequantized
+        product, and every other tile against its bytes. A wide tile passes upstream's
+        template arguments in order (global scale, then BM, BK, BN); upstream's own 64 x 64
+        FP4 tile (half-loader `fp_qmm_t`, precompiled) keeps its name and its selection.
+      globs:
+        - "mlx/backend/common/command_buffer_limits.h"
+        - "mlx/backend/metal/device.h"
+        - "mlx/backend/metal/device.cpp"
+        - "mlx/backend/metal/quantized.cpp"
+        - "mlx/backend/metal/kernels.h"
+        - "mlx/backend/metal/jit_kernels.cpp"
+        - "mlx/backend/metal/nojit_kernels.cpp"
+        - "mlx/utils.h"
+        - "docs/src/usage/environment_variables.rst"
+        - "tests/quantized_tile_tests.cpp"
+    - title: "Batch-invariant quantized matmul rows (opt-in)"
+      description: |
+        Opt-in (`MLX_QMV_BATCH_INVARIANT=1`, read once in each process). A quantized `x @ w.T`
+        with fewer rows than the vector limit (`get_qmv_batch_limit`) then runs only the vector
+        kernels, as upstream did before two of its routes: `qmv` for one row, and `qmv_wide`
+        from two rows (`qmv` where `qmv_wide` does not serve the mode on the GPU). These
+        routes do not run below the limit: `qmv_mma`/`qmv_nax` for 6 to 16 rows of the
+        floating-point modes ([ml-explore/mlx#4712](https://github.com/ml-explore/mlx/pull/4712))
+        and the staged NAX split-K kernel for affine products from 8 rows
+        ([ml-explore/mlx#4682](https://github.com/ml-explore/mlx/pull/4682)). Each row then has
+        the same bits at each row count from 2 to the limit - 1, so the rows of several
+        requests in one launch (for example their speculative verify windows) get the bits
+        that each request gets alone. A row alone (`qmv`) can have other bits than the same
+        row among two or more rows, as before. From the limit on, and without the switch,
+        the routes do not change. `tests/qmv_batch_invariant_tests.cpp` compares each window
+        of rows with the same rows in a larger launch at GLM and Qwen tensor-parallel shapes,
+        for affine and NVFP4 weights with bfloat16 and float32 inputs. The tests binary runs
+        it without the switch (up to 5 rows), and the `qmv_batch_invariant` ctest entry, which
+        Fork Tests runs, runs it with the switch (every count below the limit).
+      globs:
+        - "mlx/utils.h"
+        - "mlx/backend/metal/quantized.cpp"
+        - "docs/src/usage/environment_variables.rst"
+        - "tests/qmv_batch_invariant_tests.cpp"
+        - "tests/CMakeLists.txt"
+        - ".github/workflows/fork-tests.yml"
     - title: "GPT-OSS MXFP4 decode: masked-tail gathered qmv"
       description: |
         GPT-OSS 20B gathered MXFP4 decode has K=2880, which misses the fast vector kernel's
@@ -353,6 +416,57 @@ def:
         - "mlx/distributed/sealed.h"
         - "mlx/distributed/sealed_stream.h"
         - "tests/distributed_sealed_stream_tests.cpp"
+    - title: "Fast fence waits order earlier kernels in the same Metal pass"
+      description: |
+        An ordinary fast-synchronization GPU fence wait follows both pending
+        command passes and earlier kernels in the current concurrent encoder.
+        The in-pass buffer barrier matches the existing fenced-input consumer
+        path. Without it, a spinning wait can run before a GPU producer needed
+        by the CPU work it awaits. Arithmetic and command-buffer limits do not
+        change. A standalone bounded test checks exact GPU-to-CPU-to-GPU chains,
+        shared side roots, chains where two GPU streams feed the CPU, and request-array
+        retirement at several buffer sizes.
+        Fork Tests executes the three CTest processes on its Metal runner and
+        rejects a missing test registration. This component coverage does not
+        establish an installed RDMA model fix.
+      globs:
+        - "mlx/backend/metal/fence.cpp"
+        - "tests/metal_fence_chain.cpp"
+        - "tests/CMakeLists.txt"
+        - "docs/src/usage/using_streams.rst"
+        - ".github/workflows/fork-tests.yml"
+    - title: "GPU streams wait for each other on events"
+      description: |
+        With fast synch (`MLX_METAL_FAST_SYNCH=1`) a fence wait on the GPU is a one-thread
+        kernel that spins until the producer's timestamp arrives. That is safe when the
+        producer is a CPU stream (a thread always progresses), but not when it is another
+        GPU stream: Metal does not promise that a second command queue runs while a pass of
+        the first one spins, so the producing queue can starve and both stop. An eval whose
+        graph spans two GPU streams therefore waits for the other GPU stream on a shared
+        event (the non-fast path), while CPU <-> GPU waits keep the fast path. An array read
+        both by the CPU and by another GPU stream updates both fences, the fast fence first:
+        its cross-device update can move the array (CUDA moves it to unified memory), and the
+        event must be recorded after the move. `MLX_EAGER_FENCE_COMMIT` commits after both
+        updates. Synchronization only: results do not change.
+      globs:
+        - "mlx/fence.h"
+        - "mlx/transforms.cpp"
+        - "mlx/backend/metal/fence.cpp"
+        - "mlx/backend/no_gpu/fence.cpp"
+        - "mlx/backend/cuda/fence.cpp"
+        - "docs/src/usage/using_streams.rst"
+    - title: "Helper ops keep the caller's stream"
+      description: |
+        `quantized_matmul` cast `x`, `scales` and `biases`, `gather` and `scatter` broadcast
+        their indices, the gather matmuls cast their indices, and `affine_quantize` selected
+        its scales on the thread's default stream instead of the stream the caller passed.
+        mlx-swift passes its own process-wide stream to every op, so each of these helpers
+        put a second GPU stream into the graph: a 4-bit projection with float32 input and
+        bfloat16 scales crossed streams twice, and every crossing is a fence wait. They now
+        use the caller's stream. Where an op runs does not change what it computes.
+      globs:
+        - "mlx/ops.cpp"
+        - "tests/ops_tests.cpp"
     - title: "Fused cross-device fence update"
       description: |
         Opt-in (`MLX_METAL_FUSED_FENCE_UPDATE=1`, fast synch only). A GPU array that the CPU
@@ -470,18 +584,18 @@ def:
       globs:
         - "mlx/backend/metal/kernels/quantized_nax.h"
         - "mlx/backend/metal/kernels/fp_quantized_nax.h"
-    - title: "JIT instantiation of the NAX head-dim split attention kernel"
-      description: |
-        `get_steel_attention_nax_kernel` builds the JIT source of both `attention_nax` and
-        `attention_nax_dsplit`, and upstream passes `bv` to both. `attention_nax_dsplit` has
-        no `BV` template parameter (its value head dim is `BD`), so every later argument
-        shifts by one, an integer lands in the `MaskType` slot, and the library does not
-        compile. JIT builds such as mlx-swift's compile this kernel at runtime for full
-        attention with head dim 256 or 512 and at least 1,024 query tokens on NAX GPUs, and
-        the caller's process exited with "Unable to build metal library from source". The
-        fork passes the split kernel the same arguments as its instantiation in
-        `steel_attention_nax.metal`; the non-split kernel is unchanged. Builds that use the
-        precompiled metallib were not affected. This is a latent upstream bug (PR #48).
+    - title: "Run-time template of the head-dim split NAX attention"
+      description: |
+        Upstream's head-dim split attention (`attention_nax_dsplit`, the NAX full-attention
+        kernel for head dims 256 and 512) has no BV template parameter: its value width is
+        BD. The precompiled instances (`instantiate_attn_dsplit`) leave BV out, but the
+        run-time (JIT) build passed the template arguments of `attention_nax`, BV included,
+        so every value after BD moved one parameter to the left and the source did not
+        compile ("Unable to build metal library from source"). A JIT build (mlx-swift
+        builds this way) therefore failed every head-dim 256 or 512 prefill that takes the
+        NAX full-attention kernel. `get_steel_attention_nax_kernel` now leaves BV out for
+        the split kernel. Builds with the precompiled library did not have the bug, and the
+        kernel names do not change.
       globs:
         - "mlx/backend/metal/jit_kernels.cpp"
     - title: "64-bit products in the JACCL ring and the vector SDPA layout check"

mlx/backend/common/command_buffer_limits.h

@@ -23,6 +23,16 @@ MLX_API void mlx_metal_set_thread_command_buffer_limits(
     int* previous_ops,
     int* previous_mb);
 
+/// Output tile (32, 64, 6432, 3264 or 6464: rows x columns[, K depth]) of the
+/// steel affine quantized matmul for the
+/// products the calling thread encodes on a GPU without neural accelerators;
+/// any other value restores MLX_QMM_STEEL_TILE. Stores the thread's previous
+/// override (0 for none) in `previous_tile` when it is not null. The tile
+/// changes which threadgroup computes an element, never its bits. Builds whose
+/// kernels come only from the precompiled library (MLX_METAL_JIT=OFF) keep
+/// 32 x 32 whatever the tile.
+MLX_API void mlx_metal_set_thread_qmm_steel_tile(int tile, int* previous_tile);
+
 #ifdef __cplusplus
 }
 #endif

mlx/backend/cuda/fence.cpp

@@ -25,6 +25,8 @@ Fence::Fence(Stream s) {
   fence_ = std::make_shared<FenceImpl>(0, s);
 }
 
+Fence::Fence(Stream s, bool) : Fence(s) {}
+
 void Fence::wait(Stream s, const array&, uint32_t value) {
   auto& f = cast<FenceImpl>();
   if (value == 0) {

mlx/backend/metal/device.cpp

@@ -542,6 +542,9 @@ namespace {
 // Per-thread commit limits; zero selects the device default.
 thread_local int thread_max_ops_per_buffer = 0;
 thread_local int thread_max_mb_per_buffer = 0;
+// Per-thread steel quantized-matmul output tile (32, 64, 6432, 3264 or 6464);
+// zero selects MLX_QMM_STEEL_TILE (see `qmm` in quantized.cpp).
+thread_local int thread_qmm_steel_tile_override = 0;
 
 void set_thread_command_buffer_limits(
     int max_ops,
@@ -559,6 +562,10 @@ void set_thread_command_buffer_limits(
 }
 } // namespace
 
+int thread_qmm_steel_tile() {
+  return thread_qmm_steel_tile_override;
+}
+
 std::tuple<int, int> Device::get_max_ops_mb_per_buffer() const {
   return std::make_tuple(
       thread_max_ops_per_buffer > 0 ? thread_max_ops_per_buffer
@@ -1130,6 +1137,18 @@ extern "C" void mlx_metal_set_thread_command_buffer_limits(
       max_ops, max_mb, previous_ops, previous_mb);
 }
 
+extern "C" void mlx_metal_set_thread_qmm_steel_tile(
+    int tile,
+    int* previous_tile) {
+  if (previous_tile != nullptr) {
+    *previous_tile = mlx::core::metal::thread_qmm_steel_tile_override;
+  }
+  mlx::core::metal::thread_qmm_steel_tile_override =
+      (tile == 32 || tile == 64 || tile == 6432 || tile == 3264 || tile == 6464)
+      ? tile
+      : 0;
+}
+
 extern "C" void mlx_metal_gemma4_expert_qmm_diagnostics_snapshot(
     mlx_metal_gemma4_expert_qmm_diagnostics* diagnostics) {
   gemma4_expert_qmm_diagnostics_snapshot(diagnostics, false);

mlx/backend/metal/device.h

@@ -460,6 +460,11 @@ class MLX_API Device {
 };
 
 MLX_API Device& device(mlx::core::Device);
+
+// The calling thread's steel quantized-matmul output tile (32, 64, 6432,
+// 3264 or 6464), or 0
+// for MLX_QMM_STEEL_TILE (mlx_metal_set_thread_qmm_steel_tile).
+int thread_qmm_steel_tile();
 MLX_API CommandEncoder& get_command_encoder(Stream s);
 
 std::unordered_map<int, CommandEncoder>& get_command_encoders();

mlx/backend/metal/fence.cpp

@@ -12,12 +12,13 @@
 namespace mlx::core {
 
 struct FenceImpl {
-  FenceImpl(Stream stream) : cpu_stream(stream.device == Device::cpu) {
+  FenceImpl(Stream stream, bool allow_fast = true)
+      : cpu_stream(stream.device == Device::cpu) {
     auto d = metal::device(stream.device).mtl_device();
     if (!d->supportsFamily(MTL::GPUFamilyMetal3)) {
       use_fast = false;
     } else if (__builtin_available(macOS 15, iOS 18, *)) {
-      use_fast = env::metal_fast_synch();
+      use_fast = allow_fast && env::metal_fast_synch();
     }
 
     if (!use_fast) {
@@ -94,6 +95,10 @@ Fence::Fence(Stream stream) {
   fence_ = std::make_shared<FenceImpl>(stream);
 }
 
+Fence::Fence(Stream stream, bool fast) {
+  fence_ = std::make_shared<FenceImpl>(stream, fast);
+}
+
 void Fence::wait(Stream stream, const array& x, uint32_t value) {
   auto& f = *static_cast<FenceImpl*>(fence_.get());