From 01583588943d5e99b11863d76b32ecedb87e2aa3 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 5 Oct 2026 00:24:24 +0900 Subject: [PATCH 1/7] fix(rocm): pass custom-kernel input shapes and strides by value The overlay's `fast::hip_kernel` declared an input's `_shape` and `_strides` as pointers while `CustomKernel::eval_gpu` passes them by value through `KernelArgs::append_ndim` (padded to `JIT_MAX_NDIM` entries). A kernel that read `_shape` therefore dereferenced the shape values as an address and faulted the queue on its first launch; the HIP Gumbel-max sampler for #2064 hit it at `0x80000000000` on gfx1151. No HIP kernel in mlxcel used either argument before. The generated header now defines `KernelShape` and `KernelStrides` (eight entries, tied to `JIT_MAX_NDIM` by a `static_assert`) and declares both parameters with them by value, as upstream CUDA does with `__grid_constant__ Shape` / `Strides`; `elem_to_loc` gains an overload for them. Recorded as LOCAL_FIXES item 30, an upstreaming candidate for #1813 (not packaged yet). Refs #2064 --- docs/mlxcelverse/upstream/README.md | 1 + src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md | 1 + .../mlx/backend/rocm/custom_kernel.cpp | 39 ++++++++++++++++++- 3 files changed, 39 insertions(+), 2 deletions(-) diff --git a/docs/mlxcelverse/upstream/README.md b/docs/mlxcelverse/upstream/README.md index f90fae2a1..9578fc2cc 100644 --- a/docs/mlxcelverse/upstream/README.md +++ b/docs/mlxcelverse/upstream/README.md @@ -35,6 +35,7 @@ Each directory holds the patch (`git format-patch` output, authored by the maint | 9 (expert-batched gather qmv opt-in) | Held. It disables a path that returns wrong bf16 results on gfx1151 but has no root cause yet; LOCAL_FIXES.md says to propose it only after one. | | 15 (`SearchSorted`) | Not applicable yet: the primitive arrived in ml-explore/mlx#4035 (`5ec30acd`, 2026-08-10), after the fork's MLX base, so the fork has no `SearchSorted` to implement. Package it once the fork merges MLX at or past #4035. | | 29 (bf16 GEMMs leave the WMMA kernel from 128 rows on RDNA 3.5) | Not packaged yet. The row ceiling was measured on gfx1151 only, and the exported `quantized_matmul_runs_dequant_gemm` exists for mlxcel's dense prefill check; a fork PR would carry the ceiling alone, with measurements from at least one more RDNA 3 or RDNA 3.5 device. | +| 30 (custom-kernel `_shape` / `_strides` by value) | Not packaged yet (added 2026-10-04 with lablup/mlxcel#2064). It is a two-line declaration change plus a header struct in `custom_kernel.cpp`; a package needs a repro that launches a `hip_kernel` reading `_shape` on the unpatched fork, which faults the queue. | ## What was verified, and what the submitter still has to do diff --git a/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md b/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md index 017fb0363..825692bdd 100644 --- a/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md +++ b/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md @@ -39,6 +39,7 @@ The fork branched from upstream on 2026-06-26 (`39886de4`). Moving its ROCm code 27. **CPU-stream BLAS runs single-threaded over fine-grained memory.** On an APU the fork's allocator (`mlx/backend/rocm/allocator.cpp`, `unified_malloc`) gives every array fine-grained device memory, CPU-stream arrays included, and `Buffer::raw_ptr()` hands the CPU that same pointer. MLX's CPU backend writes BLAS and LAPACK results straight into it (`cblas_sgemm` in `mlx/backend/cpu/gemms/cblas.cpp`, also `conv.cpp`, `masked_mm.cpp` and the LAPACK primitives), and multithreaded OpenBLAS returned wrong output columns there, differently from call to call. Measured on the gfx1151 host (Ryzen AI MAX+ 395, Debian OpenBLAS 0.3.29 pthread, 32 threads): a standalone C `cblas_sgemm` of `[1, 2880] x [2880, 2880]^T` was wrong in 37 of 50 calls with inputs and output from `hipExtMallocWithFlags(hipDeviceMallocFinegrained)`, in 10 of 300 with only the output there, and exact in every call with malloc or `hipHostMalloc` buffers or with `OPENBLAS_NUM_THREADS` at 1, 2 or 4. Through MLX the wrong columns were each the last column of one OpenBLAS thread's share of the output (shares of 93 columns: 650, 1022, 1859, ...), the column whose 64-byte line the next thread also writes. This was lablup/mlxcel#2072: `tests/rocm_mxfp4_quant.rs` failed `qmm 2880x2880 M=1` in 3 of 8 runs here (7 of 10 in the issue), and the fault was in its CPU reference, not the GPU kernel; with the tensors dumped from a failing run, the GPU result matched an independent dequantize-and-dot of the packed bytes and the CPU `matmul` did not. `unified_malloc` now calls `openblas_set_num_threads(1)` once, on its first fine-grained allocation, which precedes any CPU-stream BLAS call writing such a buffer; the symbol is declared weak, so a build against another BLAS links and changes nothing. The cost is CPU-stream BLAS speed: the CPU reads fine-grained memory slowly, and that gemv went from about 0.2 s to about 1.4 s per call. GPU work and quantized or half-precision CPU matmuls (MLX's own kernels, no BLAS) are unaffected, so model inference under `MLXCEL_DEVICE=cpu` barely touches BLAS. Keeping BLAS multithreaded by giving it a cacheable scratch output would need a change at every BLAS call site in MLX's CPU backend and was not taken. `tests/rocm_cpu_blas_finegrained.rs` repeats that gemv on the CPU stream against an f64 host reference; with this change reverted it failed in 5 of 5 runs, and with it `quantized_matmul_matches_dequantized_reference` passed 50 runs in a row. Applies to the fork; to be proposed there (upstreaming candidate, lablup/mlxcel#1813). 28. **In-flight command batches bounded by what they allocate, and the cache limit enforced.** Measured on the gfx1151 host at pp512/tg128 (lablup/mlxcel#2062), the MLX peak (live buffers, `get_peak_memory`) was 20.60 GB for Meta-Llama-3.1-8B-Instruct-4bit, whose weights take 4.75 GB once loaded, and 23.56 GB for Qwen3-30B-A3B-4bit (17.17 GB), and the device-wide `mem_info_vram_used` rose by 21.19 GB and 24.09 GB. The excess was not buffer cache (0.4 to 1.7 GB at every phase boundary) but live transients inside the measured prefill: every buffer an operation allocates is kept by `CommandEncoder::add_temporary` until its batch's completion handler runs, the eager path (`use_hip_graphs()` is off) committed only every `MLX_MAX_OPS_PER_BUFFER` ops (2000) and never registered those commits with MLX's scheduler, so neither the scheduler's task cap nor `set_memory_limit` ever made the host wait, and the host could encode far ahead of the GPU, holding the transients of many batches at once. For the 8B, an f16 checkpoint, most of it is the f16 weight copy `QuantizedMatmul`'s dequantize-and-GEMM path allocates per matrix (with `MLX_ROCM_QMM_DEQUANT_GEMM=0` the peak was 6.78 GB, at a twentieth of the prefill speed); the MoE model is bf16, whose affine 4-bit matmuls took the fused WMMA kernel at the time and allocated no copy (item 29 moved its prefill projections to the dequantize path) (the variable left its peak at 23.56 GB), so its excess is other operations' outputs held the same way, which the bound below removes too. `MLX_MAX_OPS_PER_BUFFER=50` alone changed nothing, because more commits do not stop the host from running ahead. Now `gpu::eval` (`mlx/backend/rocm/eval.cpp`) adds what each primitive allocated during `eval_gpu` to its encoder's open batch, read from a thread-local byte counter in the allocator (`rocm::thread_allocated_bytes()`, so frees on the worker thread do not disturb it), and `CommandEncoder::maybe_commit` (`mlx/backend/rocm/device.cpp`) commits once the open batch has allocated a quarter of `MLX_ROCM_MAX_INFLIGHT_MB` (default 1024); after that commit and after the commit `gpu::finalize` makes at the end of every eval and `async_eval`, it records a HIP event, and blocks on the oldest such batch while the committed ones still exceed the budget. The wait is `hipEventSynchronize` (the device is in blocking-sync mode, so the thread sleeps), or a polled wait that gives up at `MLX_ROCM_GPU_WATCHDOG_SECS` when that is set; a faulted stream ends either with the fault, which is recorded as the stream's error. Nothing is tracked while a decode-step capture or any stream capture is in progress. The events are raw `hipEvent_t` owned by the encoder, not `HipEvent`: `HipEvent` returns its handle to a function-local static pool when destroyed, an encoder can be destroyed after that pool at process exit, and a first version built on it aborted `tests/rocm_mxfp4_quant.rs` at exit with `malloc_consolidate(): unaligned fastbin chunk detected`. `0` restores the old behavior. `tests/rocm_inflight_bound.rs` runs 40 f16 4-bit 8192x8192 matmuls over 512 rows in one evaluation (5 GiB of dequantized weight copies): the peak rose by 1.75 GiB with the default and by 6.25 GiB with `MLX_ROCM_MAX_INFLIGHT_MB=0`, which fails its 3 GiB bound. With the default the peaks were 6.14 GB and 18.57 to 18.58 GB over three harness runs each, decode throughput was unchanged (median 37.57 against 37.02 tok/s and 61.69 against 61.79 tok/s) and the 8B's prefill lost 1.8% (the MoE model's prefill varies more than that from run to run); budgets from 256 to 4096 MiB traded peak for prefill speed within a few percent (the table is in `docs/benchmark_results/rocm-memory-gfx1151-2026-09-30.md`). Separately, `RocmAllocator::set_cache_limit` stored `max_pool_size_` and nothing read it, so MLX's `set_cache_limit` (mlxcel's `MLXCEL_CACHE_LIMIT`) bounded nothing, while the exact-size cache (`min_utilization` 1.0) keeps every buffer size it has seen until `clear_cache()`. `malloc_async` now trims the cache on a cache miss once it exceeds `max_pool_size_`, the only place the footprint grows, so active plus cache stays within the live set plus the limit plus the one request; it trims to three quarters of the limit, so one round of blocking `hipFree` calls covers several misses of a workload whose shapes keep changing; `free()` and cache hits still never call `hipFree`, which the fork avoids for the training reasons in its comment there. The limit still defaults to the memory limit (76.8 GiB here), so the trim changes nothing until someone sets a limit; mlxcel sets 2 GiB on ROCm builds. `memory::tests::cache_limit_bounds_the_free_buffer_cache` in mlxcel-core frees 64 buffers of distinct sizes under a 1 MiB limit and checks the cache stays under 8 MiB; with the trim removed it held 49.8 MB and failed. The count is of allocations, not only of transients (a batch that grows the KV cache counts too), and a batch's buffers are dropped by the worker thread just after its event completes, so the bound is approximate. Applies to the fork; to be proposed there (upstreaming candidate, lablup/mlxcel#1813). 29. **bf16 GEMMs of 128 rows or more leave the fused WMMA kernel on RDNA 3.5.** `QuantizedMatmul::eval_gpu` (`mlx/backend/rocm/quantized/qmm.hip`) sent every single bf16 affine GEMM with more than one row (4-, 6- or 8-bit, group 64, `N % 16 == 0`, `K % 64 == 0`) to `qmm_wmma_dense_kernel` on a device with native WMMA that is not a low-CU iGPU, however many rows it had. On gfx1151 that kernel lost to the dequantize + hipBLASLt path it bypasses on every bf16 shape measured from 128 rows up (1.0x to 3.1x at 128 rows, 1.6x to 2.4x at 256, 2.4x to 5.1x at 1024 and 2048; it won only on some shapes at 2 to 64 rows), so bf16 prefill ran at a fraction of what the same device did for f16 checkpoints (lablup/mlxcel#2081). It also made `quantized_matmul` and `dequantize` + `matmul` disagree for bf16 (499 of 1,048,576 outputs at `[1024, 2048] x [1024, 2048]`, a reduction-order difference), which is what mlxcel's dense prefill path checks. The dispatch decision is now one function, `select_qmm_route`, used by `eval_gpu`; it keeps the kernel below `MLX_ROCM_WMMA_QMM_MAX_M` rows (default 128 on the RDNA 3.5 tier, no ceiling on the unmeasured tiers) and sends larger GEMMs to the bf16 dequantize + hipBLASLt route, except where that route would be the fp8 path (RDNA 4), which keeps the kernel as before. `MLX_ROCM_WMMA_QMM=1` removes the ceiling, which is the old dispatch on any device that is not a low-CU iGPU (there it forces the kernel on, as it always did). `rocm.h` exports `quantized_matmul_runs_dequant_gemm` (stubbed in `no_rocm.cpp`), which answers from the same function whether a single GEMM of a given shape takes dequantize + hipBLASLt, for mlxcel to decide when its dense prefill path returns the same bytes. Same binary, `MLX_ROCM_WMMA_QMM=1` against the default, prefill tok/s over three ABBA runs: Gemma 3 4B 4-bit 1146 to 2289 at 512 tokens and 961 to 2815 at 2048, Qwen3-0.6B 4-bit 4581 to 7633 and 3059 to 4236, Qwen3-30B-A3B 4-bit (MoE, only the attention projections take this path) 311 to 326 and 283 to 297; decode unchanged, and MLX peak memory within 0.1 GB except Qwen3-0.6B at 512 tokens (1.04 to 1.46 GB), since the route allocates a bf16 weight copy per matrix that item 28 bounds. The route's dequantized-weight LRU (8 matrices or 256 MB, `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE`, `MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES`) gets no hits across a model's projections in prefill but still keeps its last entries alive afterwards, as it already did for f16 checkpoints; `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE=0` turns it off. `docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md` has the tables. Applies to the fork; the ceiling was measured on gfx1151 only. +30. **Custom kernels receive `_shape` and `_strides` by value.** `build_kernel` in the fork's `mlx/backend/rocm/custom_kernel.cpp` declared an input's shape and strides as `const int32_t* _shape` and `const int64_t* _strides`, while `CustomKernel::eval_gpu` passes them with `KernelArgs::append_ndim`, which pads the vector to `JIT_MAX_NDIM` (8) entries and hands `hipModuleLaunchKernel` a pointer to that storage, so the kernel received the values themselves and read the first eight bytes of the shape as an address. Any `fast::hip_kernel` source that mentioned `_shape` faulted the queue on its first launch: the HIP Gumbel-max sampler (lablup/mlxcel#2064), which reads `logits_shape[1]` as the CUDA and Metal ports do, died with `HSA_STATUS_ERROR_MEMORY_FAULT` at address `0x80000000000`, the page of `0x80000000800`, which is the shape `{2048, 2048}` read as a 64-bit pointer, on gfx1151. No HIP kernel in mlxcel used either argument before, which is why it went unseen. The generated header now defines `KernelShape` (8 `int32_t`) and `KernelStrides` (8 `int64_t`) with an `operator[]`, the parameters are declared with them by value, which is what upstream CUDA's `const __grid_constant__ Shape` / `Strides` does, and `elem_to_loc` gains an overload for them; a `static_assert` ties the 8 to `JIT_MAX_NDIM`. `sampling_fixed_key_tests` in mlxcel-core reads `logits_shape` through the Gumbel-max port and faults without this change; the `_strides` declaration has no user in mlxcel and was checked only by compiling it. Applies to the fork; to be proposed there (upstreaming candidate, lablup/mlxcel#1813). ## Fixes to the fork's build diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/custom_kernel.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/custom_kernel.cpp index 2a050e30d..755c05a3a 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/custom_kernel.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/custom_kernel.cpp @@ -17,6 +17,12 @@ namespace mlx::core::fast { namespace { +// `KernelShape` and `KernelStrides` in `default_header` hold this many entries, +// the width `KernelArgs::append_ndim` pads to. +static_assert( + rocm::JIT_MAX_NDIM == 8, + "resize KernelShape/KernelStrides in default_header with JIT_MAX_NDIM"); + // Inline the essential definitions for custom kernels // This avoids the need for include paths in JIT compilation constexpr const char* default_header = R"( @@ -67,6 +73,34 @@ elem_to_loc(IdxT elem, const int* shape, const int64_t* strides, int ndim) { return loc; } +// An input's `_shape` and `_strides` arrive BY VALUE, as the host +// side passes them: `KernelArgs::append_ndim` pads the vector to JIT_MAX_NDIM +// (8) entries and hands the launch a pointer to that storage, which +// `hipModuleLaunchKernel` copies into the parameter. CUDA's custom kernels +// declare the same arguments `const __grid_constant__ Shape` / `Strides`. +struct KernelShape { + int32_t data[8]; + __device__ int32_t operator[](int i) const { + return data[i]; + } +}; + +struct KernelStrides { + int64_t data[8]; + __device__ int64_t operator[](int i) const { + return data[i]; + } +}; + +template +__device__ IdxT elem_to_loc( + IdxT elem, + const KernelShape& shape, + const KernelStrides& strides, + int ndim) { + return elem_to_loc(elem, shape.data, strides.data, ndim); +} + } // namespace mlx::core::rocm )"; @@ -118,10 +152,11 @@ std::string build_kernel( // Add input shape, strides and ndim if present in the source if (arr.ndim() > 0) { if (std::get<0>(shape_infos[i])) { - kernel_source << " const int32_t* " << name << "_shape,\n"; + // By value, not a pointer: see `KernelShape` in the header above. + kernel_source << " const KernelShape " << name << "_shape,\n"; } if (std::get<1>(shape_infos[i])) { - kernel_source << " const int64_t* " << name << "_strides,\n"; + kernel_source << " const KernelStrides " << name << "_strides,\n"; } if (std::get<2>(shape_infos[i])) { kernel_source << " const int " << name << "_ndim,\n"; From 19e6c9a318b68515df04d74a08990f355a9af4e0 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 5 Oct 2026 00:24:37 +0900 Subject: [PATCH 2/7] update(rocm): port the Gumbel-max and rejection samplers to HIP Both fused samplers took the MLX graph fallback on ROCm because their port tables had no `.rocm` entry, and `rejection_sample_supported()` answered with `custom_kernels_available()`, which is Metal-or-CUDA by definition, so a filled slot would have stayed unreachable. - `GUMBEL_MAX_SAMPLE_HIP_SOURCE` and `REJECTION_SAMPLE_HIP_SOURCE` (new headers, data only) are line-for-line ports of the CUDA bodies with the same inputs, outputs, grid, template arguments, Philox-4x32-10 counter and key layout. The only changes are spelling: explicit `(float)` reads, since `hip_bfloat16` converts only explicitly, and `-__builtin_huge_valf()` for `-INFINITY`. The `hip_kernel` holders and launches stay in `sampling.cpp` and `sampling_rejection.cpp`, so the dtype-key pin and its count are unchanged. - No wave32 `#error`: neither kernel has a lane-level operation (every reduction and the scan go through shared memory with a barrier per step), so both are correct on any wavefront size. - `rejection_sample_supported()` reads `has_kernel_port(rejection_ports())`, as `gumbel_max_sample_supported()` already did; true on Metal and CUDA as before. - `random_bits` bridge function, and `sampling_fixed_key_tests`: each kernel's draw recomputed on the host from the key it consumed (Gumbel-max also through the MLX graph, in f32, f16 and bf16; rejection under min-p, where round 0 decides), plus two-sample chi-square tests against `fused_sample_categorical`. The fixed-key tests fail with the Philox counter or word changed in either HIP source. - The dtype-key checker's fixture mutations now move both launches in `sampling.cpp`. Refs #2064 --- docs/environment-variables.md | 4 +- scripts/ci/check_kernel_dtype_keys_test.sh | 3 + src/lib/mlx-cpp/turbo/sampling.cpp | 40 +- src/lib/mlx-cpp/turbo/sampling.h | 4 +- src/lib/mlx-cpp/turbo/sampling_gumbel_hip.h | 141 ++++++ src/lib/mlx-cpp/turbo/sampling_rejection.cpp | 46 +- src/lib/mlx-cpp/turbo/sampling_rejection.h | 4 +- .../mlx-cpp/turbo/sampling_rejection_hip.h | 318 ++++++++++++ src/lib/mlxcel-core/build.rs | 2 + src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp | 10 +- src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h | 1 + src/lib/mlxcel-core/src/lib.rs | 22 +- .../src/sampling_fixed_key_tests.rs | 476 ++++++++++++++++++ .../mlxcel-core/src/sampling_gumbel_tests.rs | 9 +- .../src/sampling_rejection_tests.rs | 9 +- 15 files changed, 1073 insertions(+), 16 deletions(-) create mode 100644 src/lib/mlx-cpp/turbo/sampling_gumbel_hip.h create mode 100644 src/lib/mlx-cpp/turbo/sampling_rejection_hip.h create mode 100644 src/lib/mlxcel-core/src/sampling_fixed_key_tests.rs diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 0bbf87c1b..03cfafc7a 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -559,8 +559,8 @@ recommended as normal deployment settings. | `MLXCEL_DISABLE_FUSED_GLOBAL_SCALE` | `1`/`true`/`on`/`yes` (case-insensitive) disable; unset or any other value keeps the fold | off (fold on) | Rolls back the NVFP4 global-scale fold for Gemma 4 (issues #698/#705). By default the fused MLP and per-layer-input-gate C++ paths fold each per-projection `weight_scale_2` sidecar (from the direct ModelOpt transcode, issue #693/#697) into the fused kernel at the mathematically correct points: the gate scale before the GeGLU activation, the up scale on the up product, and the down scale on the fused output, each reproducing `apply_global_scale` byte-for-byte. Native NVFP4 prefill now uses a shape-specific scaled MLP graph, and standalone `UnifiedLinear` sidecar projections can apply qmm + global scale + dense bias through one C++ helper. When set, sidecar-carrying paths fall back to the op-at-a-time `UnifiedLinear::forward` scalar application (the pre-#698 bypass). Greedy temp-0 decode is token-identical across the two paths on `gemma-4-31b-it-nvfp4`; the fold removes element-wise dispatches on the gemma-4 path where CUDA graphs are disabled (#688). See `docs/benchmark_results/nvfp4-direct-transcode-gb10-2026-07-08.md` and `docs/benchmark_results/nvfp4-native-prefill-m1ultra-2026-07-09.md`. | | `MLXCEL_ONE_BIT_KERNEL` | `0`/`false`/`off`/`no` (case-insensitive) disable; unset or any other value keeps the kernels | on (Metal) | Affine 1-bit checkpoints (`quantization.bits == 1`, issue #1370). MLX has no 1-bit kernel, so the `quantized_matmul`, `quantized_linear_forward`, `dequantize` and `quantized_embedding` bridge primitives route a 1-bit affine triple to mlxcel's own fused Metal qmv (fewer than 16 rows) or tiled qmm kernel. Set to `0` to force the dequantize graph (`x @ dequant(W)^T`) instead, which is the parity oracle for the kernels and the path CUDA and ROCm always take. Read once in C++ at first use and cached for the process lifetime. | | `MLXCEL_FUSED_XIELU` | `0`/`false`/`off`/`no` disable; any other value or unset enables | on | Fused single-launch Metal xIELU kernel for the Apertus MLP activation (#409), on by default since the M5 Max validation. `MLP::forward` routes through one Metal dispatch covering the ~11 elementwise ops in `apertus_xielu` (square, minimum, expm1, where, and neighbors) instead of the per-op graph. Greedy temp-0 decode is byte-identical to the elementwise path on Apple Silicon: every intermediate stays in the input dtype (bf16) and the kernel reproduces MLX's `expm1f` exactly. Measured decode speedup on M1 Ultra (+2.7%, Apertus-8B 83.4 to 85.7 tok/s) and M5 Max (+1.9%, 112.0 to 114.2 tok/s), with no regression. Set to `0` to force the elementwise path. On non-Metal back-ends the FFI falls back to an equivalent elementwise graph, so the flag is safe to set everywhere. Apertus only; no other model family is affected. | -| `MLXCEL_SAMPLING_GUMBEL` | `0`/`false`/`off`/`no` disable (case-insensitive); any other value or unset enables | on | Softmax-free Gumbel-max categorical sampling kernel (issue #900), on by default on Metal and CUDA. On the no-filter stochastic path (`temperature > 0` with no top-k, top-p, or min-p) the fused sampler replaces `mlx::core::random::categorical` with one index-carrying max reduction over the vocabulary: adding i.i.d. Gumbel(0,1) noise to `logits / temperature` and taking the argmax draws exactly from `softmax(logits / temperature)`, so the normalization pass over 32K-152K entries disappears. Noise comes from a stateless in-kernel Philox-4x32-10 counter keyed on `(row, element)`, so the sampled id is a pure function of the launch key and the logits and does not move when the launcher splits a row across threadgroups to keep a small batch wide. One launch covers the whole `[B, vocab]` batch. **This changes the RNG stream**: at an equal seed the sampled tokens differ from the pre-#900 `categorical` streams, because the two consume the shared MLX key sequence differently. The distribution is identical (chi-square goodness of fit against exact softmax is a committed test), and greedy (`temperature == 0`) output is byte-identical either way. Set to a falsy value to restore `categorical` exactly, which is the way to reproduce a token stream recorded before this landed. Filtered configurations (top-k / top-p / min-p) never reach the kernel and are bit-identical with the switch either way; token bias is an upstream `-inf` logit mask that the kernel reproduces exactly, and XTC-active requests never reach this kernel (XTC lives inside the stock filter chain). Read once per process, so set it before starting `mlxcel` or `mlxcel-server`. Inert on CPU-only builds and under `MLXCEL_DEVICE=cpu`. A/B with `examples/gumbel_sampling_microbench.rs`. | -| `MLXCEL_SAMPLING_REJECTION` | `0`/`false`/`off`/`no` disable (case-insensitive); any other value or unset enables | on | Sorting-free top-k / top-p / min-p sampling by dual-pivot rejection (issue #901), on by default on Metal and CUDA. On the filtered stochastic path the fused sampler replaces the `argpartition` top-k, the `argsort` + `cumsum` top-p and the min-p mask, plus the trailing `random::categorical`, with the softmax passes and one custom kernel; since the untempered-chain fix the kernel reads two probability rows (the untempered row resolves every filter, the tempered row weights the draw), which coincide at `temperature 1`. **Routing is restricted to the configurations measured faster.** The kernel replaces a *sort*, so it is routed only where the stock chain sorts, which is when top-p is active; and when top-k is active as well, only at vocabularies up to 32768. Measured on M1 Ultra over three repetitions of vocab {32K, 64K, 152K} x batch {1, 4, 8}: top-p alone 1.28x to 2.35x at every cell, top-k alone 0.31x to 0.97x at every cell, min-p alone 0.47x to 0.88x at 152K, top-k with top-p 1.27x to 1.64x at 32K but 0.71x to 0.83x at 152K. The mechanism is the kernel's round count, which the microbenchmark reports: top-p accepts in one or two rounds because a probability-weighted draw lands in the high-mass head, while top-k needs the draw to land in the top forty of the whole vocabulary and takes two to seven rounds, growing with the vocabulary. Every round is another full-row sweep, and the loop runs on one threadgroup per row because the shrinking interval is carried across sweeps and MLX custom kernels have no grid-wide barrier, so the kernel pays single-core bandwidth per round while `argpartition` and the stock min-p mask run across the whole GPU. top-k alone, min-p alone, top-k with min-p, and top-k with top-p above vocab 32768 therefore stay on the stock chain and are bit-identical to before; the decline is announced once at INFO with the numbers behind it. Vocab 65536 has not been measured for the top-k plus top-p combination and is excluded rather than interpolated. All three filters are threshold filters on the probability value, so the filtered support is `{p > low}` for a single scalar per row; the kernel finds that scalar without sorting, by rejection sampling on a shrinking interval. Each round draws a candidate from the current proposal with a fixed-order block scan, takes the candidate's probability as one pivot and the midpoint of the bracket as the other, reduces `(count, mass)` above both in one vocabulary sweep, and accepts the candidate as soon as it passes every active filter test. The bisection pivot is the midpoint of the IEEE-754 **bit patterns**, not the arithmetic mean: pure integer arithmetic that a fast-math backend with flush-to-zero cannot corrupt for subnormal probabilities, and it isolates a single float in at most 32 rounds at any magnitude. min-p is folded into the initial interval and costs no round at all. One launch covers the whole `[B, vocab]` batch, and the kernel reads `{top_k, top_p, min_p}` per row so rows with different values need no second launch. **This changes the RNG stream**: at an equal seed the sampled tokens differ from the pre-#901 streams, because the kernel consumes the shared MLX key sequence differently from `random::categorical`. **One semantic difference is intentional.** When top-k AND top-p are both active, the stock chain masks to the top-k set and *renormalises* before applying top-p, so its mass target is `top_p * Z_k` with `Z_k` the top-k mass; the kernel applies both tests to the untruncated distribution, so its target is `top_p * total` and its support is a superset. The two agree unless some token's exclusive cumulative mass falls in `(top_p * Z_k, top_p * total]`. Every other configuration (top-k alone, top-p alone, min-p alone, top-k+min-p, top-p+min-p) is exact, and support equality against the `argpartition` mask including ties, plus chi-square goodness of fit against the renormalised truncated distribution, are committed tests. Greedy (`temperature == 0` or `top_k == 1`) is byte-identical either way, and the no-filter path is unaffected (it belongs to `MLXCEL_SAMPLING_GUMBEL`). **The routed path never synchronizes.** Both decode drivers are software pipelines that build step n+1 and `async_eval` it before reading step n, so a sampler that evaluates anything internally drains the queue inside the caller's build phase and collapses the pipeline. The first cut of this kernel read the per-row converged flags back to host inside `fused_sample` and measured 1.7x slower end to end on Qwen3-0.6B, against an op-level matrix that scored the same configuration 1.14x to 1.17x faster. The production path therefore evaluates nothing: the round cap is unreachable by construction (the bit-space bracket bounds the loop at 31 rounds), and the converged flags are stashed and inspected on a later call once `array::is_available()` reports the launch has landed, which in a decode loop is by the next token. A row that somehow exhausted the cap returns its own argmax, which is always inside the filtered support, so the degradation would be one greedy draw rather than an invalid token; the event is counted in `mlxcel_core::rejection_cap_overflow_rows()` and announced once at INFO. Which path a run took is announced at INFO, once per distinct outcome kind, so a benchmark arm can be proven from a log. Set to a falsy value to restore the `argpartition` chain exactly, which is the way to reproduce a token stream recorded before this landed. Read once per process, so set it before starting `mlxcel` or `mlxcel-server`. Inert on CPU-only builds and under `MLXCEL_DEVICE=cpu`. A/B with `examples/rejection_sampling_microbench.rs`. | +| `MLXCEL_SAMPLING_GUMBEL` | `0`/`false`/`off`/`no` disable (case-insensitive); any other value or unset enables | on | Softmax-free Gumbel-max categorical sampling kernel (issue #900), on by default on Metal, CUDA and ROCm (the HIP port is lablup/mlxcel#2064). On the no-filter stochastic path (`temperature > 0` with no top-k, top-p, or min-p) the fused sampler replaces `mlx::core::random::categorical` with one index-carrying max reduction over the vocabulary: adding i.i.d. Gumbel(0,1) noise to `logits / temperature` and taking the argmax draws exactly from `softmax(logits / temperature)`, so the normalization pass over 32K-152K entries disappears. Noise comes from a stateless in-kernel Philox-4x32-10 counter keyed on `(row, element)`, so the sampled id is a pure function of the launch key and the logits and does not move when the launcher splits a row across threadgroups to keep a small batch wide. One launch covers the whole `[B, vocab]` batch. **This changes the RNG stream**: at an equal seed the sampled tokens differ from the pre-#900 `categorical` streams, because the two consume the shared MLX key sequence differently. The distribution is identical (chi-square goodness of fit against exact softmax is a committed test), and greedy (`temperature == 0`) output is byte-identical either way. Set to a falsy value to restore `categorical` exactly, which is the way to reproduce a token stream recorded before this landed. Filtered configurations (top-k / top-p / min-p) never reach the kernel and are bit-identical with the switch either way; token bias is an upstream `-inf` logit mask that the kernel reproduces exactly, and XTC-active requests never reach this kernel (XTC lives inside the stock filter chain). Read once per process, so set it before starting `mlxcel` or `mlxcel-server`. Inert on CPU-only builds and under `MLXCEL_DEVICE=cpu`. A/B with `examples/gumbel_sampling_microbench.rs`. | +| `MLXCEL_SAMPLING_REJECTION` | `0`/`false`/`off`/`no` disable (case-insensitive); any other value or unset enables | on | Sorting-free top-k / top-p / min-p sampling by dual-pivot rejection (issue #901), on by default on Metal, CUDA and ROCm (the HIP port is lablup/mlxcel#2064). On the filtered stochastic path the fused sampler replaces the `argpartition` top-k, the `argsort` + `cumsum` top-p and the min-p mask, plus the trailing `random::categorical`, with the softmax passes and one custom kernel; since the untempered-chain fix the kernel reads two probability rows (the untempered row resolves every filter, the tempered row weights the draw), which coincide at `temperature 1`. **Routing is restricted to the configurations measured faster.** The kernel replaces a *sort*, so it is routed only where the stock chain sorts, which is when top-p is active; and when top-k is active as well, only at vocabularies up to 32768. Measured on M1 Ultra over three repetitions of vocab {32K, 64K, 152K} x batch {1, 4, 8}: top-p alone 1.28x to 2.35x at every cell, top-k alone 0.31x to 0.97x at every cell, min-p alone 0.47x to 0.88x at 152K, top-k with top-p 1.27x to 1.64x at 32K but 0.71x to 0.83x at 152K. The mechanism is the kernel's round count, which the microbenchmark reports: top-p accepts in one or two rounds because a probability-weighted draw lands in the high-mass head, while top-k needs the draw to land in the top forty of the whole vocabulary and takes two to seven rounds, growing with the vocabulary. Every round is another full-row sweep, and the loop runs on one threadgroup per row because the shrinking interval is carried across sweeps and MLX custom kernels have no grid-wide barrier, so the kernel pays single-core bandwidth per round while `argpartition` and the stock min-p mask run across the whole GPU. top-k alone, min-p alone, top-k with min-p, and top-k with top-p above vocab 32768 therefore stay on the stock chain and are bit-identical to before; the decline is announced once at INFO with the numbers behind it. Vocab 65536 has not been measured for the top-k plus top-p combination and is excluded rather than interpolated. All three filters are threshold filters on the probability value, so the filtered support is `{p > low}` for a single scalar per row; the kernel finds that scalar without sorting, by rejection sampling on a shrinking interval. Each round draws a candidate from the current proposal with a fixed-order block scan, takes the candidate's probability as one pivot and the midpoint of the bracket as the other, reduces `(count, mass)` above both in one vocabulary sweep, and accepts the candidate as soon as it passes every active filter test. The bisection pivot is the midpoint of the IEEE-754 **bit patterns**, not the arithmetic mean: pure integer arithmetic that a fast-math backend with flush-to-zero cannot corrupt for subnormal probabilities, and it isolates a single float in at most 32 rounds at any magnitude. min-p is folded into the initial interval and costs no round at all. One launch covers the whole `[B, vocab]` batch, and the kernel reads `{top_k, top_p, min_p}` per row so rows with different values need no second launch. **This changes the RNG stream**: at an equal seed the sampled tokens differ from the pre-#901 streams, because the kernel consumes the shared MLX key sequence differently from `random::categorical`. **One semantic difference is intentional.** When top-k AND top-p are both active, the stock chain masks to the top-k set and *renormalises* before applying top-p, so its mass target is `top_p * Z_k` with `Z_k` the top-k mass; the kernel applies both tests to the untruncated distribution, so its target is `top_p * total` and its support is a superset. The two agree unless some token's exclusive cumulative mass falls in `(top_p * Z_k, top_p * total]`. Every other configuration (top-k alone, top-p alone, min-p alone, top-k+min-p, top-p+min-p) is exact, and support equality against the `argpartition` mask including ties, plus chi-square goodness of fit against the renormalised truncated distribution, are committed tests. Greedy (`temperature == 0` or `top_k == 1`) is byte-identical either way, and the no-filter path is unaffected (it belongs to `MLXCEL_SAMPLING_GUMBEL`). **The routed path never synchronizes.** Both decode drivers are software pipelines that build step n+1 and `async_eval` it before reading step n, so a sampler that evaluates anything internally drains the queue inside the caller's build phase and collapses the pipeline. The first cut of this kernel read the per-row converged flags back to host inside `fused_sample` and measured 1.7x slower end to end on Qwen3-0.6B, against an op-level matrix that scored the same configuration 1.14x to 1.17x faster. The production path therefore evaluates nothing: the round cap is unreachable by construction (the bit-space bracket bounds the loop at 31 rounds), and the converged flags are stashed and inspected on a later call once `array::is_available()` reports the launch has landed, which in a decode loop is by the next token. A row that somehow exhausted the cap returns its own argmax, which is always inside the filtered support, so the degradation would be one greedy draw rather than an invalid token; the event is counted in `mlxcel_core::rejection_cap_overflow_rows()` and announced once at INFO. Which path a run took is announced at INFO, once per distinct outcome kind, so a benchmark arm can be proven from a log. Set to a falsy value to restore the `argpartition` chain exactly, which is the way to reproduce a token stream recorded before this landed. Read once per process, so set it before starting `mlxcel` or `mlxcel-server`. Inert on CPU-only builds and under `MLXCEL_DEVICE=cpu`. A/B with `examples/rejection_sampling_microbench.rs`. | | `MLXCEL_CUDA_F16_NORMALIZE` | `1`/`true`/`on`/`yes` enable; `0`/`false`/`off`/`no` disable | **on** below sm_80, off (opt-in) at sm_80 and later | **CUDA only.** Load-time bf16 -> f16 normalization of weights for the single-dtype decode graph (issue #636). The default is split by compute capability because the hardware is. **Ampere and later (sm_80+)**: off, unchanged from #636. bf16 has native ALUs there, the merged `patches-cuda/dtype.cpp` promotion patch already yields a 0-AsType single-dtype bf16 decode graph, and f16 offers no measured throughput gain (qwen2.5-0.5b-bf16 decode 208 vs 207 tok/s on GB10) while narrowing dynamic range. **Volta and Turing (sm_70, sm_75)**: on, and quantized checkpoints are included. Those parts have no bf16 ALU, no bf16 tensor-core MMA atom, and no cuBLAS bf16 GEMM, so every bf16 operand is converted before it can execute. Measured on a Tesla V100 with qwen3.8-27B-4bit: prefill 85.3 s to 4.8 s and decode 118.1 to 49.2 ms per token, with `qmm_naive` spending 199.1 s against 11.4 s over an identical 994 launches. On a quantized checkpoint the packed planes stay u32 and only the bf16 side-data moves, which is why quantized models are in scope here and out of scope above. Set to `0` to opt out on pre-Ampere and keep bf16. The f16-fragile exception differs by architecture. At sm_80 and later it is the full conservative list from #732 (gemma, cohere/command-r, apertus, gpt-oss, and any config with a nonzero softcap or `logit_scale`). Below sm_80 every entry on that list was measured on a V100 and only **apertus** stays excluded: its xIELU squares its input, and f16 produces NaN at all 508 scored positions from the first token where bf16 scores 24.27. Gemma (+0.09% perplexity over 10208 tokens, decode 6.12 to 9.79 tok/s), gpt-oss (token-weighted -0.02%, with the two window sizes disagreeing in sign) and Cohere (which ships F16 weights already and so runs in f16 regardless, scoring 20.25 and 13.15) all measured clean, as did the generic softcap and `logit_scale` triggers, which is what a cap bounding a value rather than growing it predicts. Metal / Apple Silicon is governed by the separate always-on policy and is untouched. Drafter checkpoints loaded through `mlxcel-core` (DFlash, Muse Glimmer, Inkling MTP, Qwen 3.5 MTP) follow the same rule (issue #1782): before that, every drafter was converted to f16 unconditionally, and on Ampere and later CUDA, where the target keeps bf16, MLX promoted the bf16 residual stream times the f16 drafter weights to float32, so the whole drafter forward ran in f32 and re-upcast all of its weights on every round (about 240 launches and 30 to 38 ms of device time per verify round on GB10 with `qwen3.5-4b-4bit` and its DFlash drafter, the largest fixed term in the round). | | `MLXCEL_CUDA_F16_FRAGILE` | `1`/`true`/`on`/`yes` enable; unset keeps the exclusion | off (opt-in) | **CUDA below sm_80 only.** Converts a family on the f16-fragile exception list anyway. That list is inherited from #732, which adopted it for Ampere and later where f16 gave, in that commit's words, "no AsType reduction and no throughput gain": with no upside, excluding a family on a suspicion costs nothing, and that PR's test plan ran no fragile family in f16. Below sm_80 the list is down to apertus, which is the one entry measurement supported, so this flag now means converting a family that produced NaN at every scored position on this hardware. It remains for a checkpoint whose activations you have reason to believe differ, and for taking the measurement on a family this list does not yet know. | | `MLXCEL_KEEP_BF16` | set to any value to enable; unset disables | off | **Measurement instrument, not a supported configuration.** Skips load-time bf16 -> f16 conversion for every checkpoint, on every backend, so a bf16 model loads uniformly bf16. It exists to A/B a dtype policy against the mixture it replaces, and it is what established the Gemma3n result below. Gemma3n used to keep only its language MLP bf16 and convert the rest to f16, leaving one checkpoint holding both dtypes; MLX promotes f16 with bf16 to f32, so every boundary between the kept MLP and its f16 neighbours promoted and the decode path read promoted weights the whole way down. Loading uniformly bf16 measured, on M5 Max with `gemma3n-e4b-bf16` at a 273-token prompt with an image and 128 generated, decode 38.02 to 47.48 tok/s and prefill 2143 to 2371 tok/s, against mlx-vlm 0.6.17 at 48.72 and 1207: decode moves from 78% of the reference to 97% while prefill stays about 1.95x ahead. Gemma3n now loads uniformly bf16 by default (see `gemma3n_bf16_key`), so this flag no longer changes that family. **It is not evidence that bf16 beats f16.** The control run is the point: on families with no forced-bf16 subset the two dtypes land within noise of each other on the same hardware, `qwen2.5-0.5b-bf16` at 401.8 against 397.3 tok/s and `llama-3.1-8b-bf16` at 33.07 against 33.02. What costs is the mixture, not the dtype, so every other family keeps converting to f16 and setting this flag for them buys nothing while narrowing nothing. Drafter checkpoints loaded through `mlxcel-core` honor it the same way (issue #1782). Read per load. | diff --git a/scripts/ci/check_kernel_dtype_keys_test.sh b/scripts/ci/check_kernel_dtype_keys_test.sh index 8598bec1a..e71f2e26d 100755 --- a/scripts/ci/check_kernel_dtype_keys_test.sh +++ b/scripts/ci/check_kernel_dtype_keys_test.sh @@ -141,7 +141,9 @@ fi # the `template_args` stay behind. The pin names the file that left, and the # helper that arrived. dir="$(make_tree launch-moved-to-helper)" +# The sampler holds a CUDA and a HIP launch (#2064); both move. replace_in "$dir/$sampling" 'mlx::core::fast::cuda_kernel(' 'mlxcel::make_jit_kernel(' +replace_in "$dir/$sampling" 'mlx::core::fast::hip_kernel(' 'mlxcel::make_jit_kernel(' cat >"$dir/src/lib/mlxcel-core/cpp/jit_helper.cpp" <<'CPP' mlx::core::fast::CustomKernelFunction make_jit_kernel(const std::string& name) { return mlx::core::fast::cuda_kernel(name, {"x"}, {"out"}, ""); @@ -156,6 +158,7 @@ fi # A launch that is commented out is not a launch: its file leaves scope. dir="$(make_tree launch-commented-out)" replace_in "$dir/$sampling" 'kernel = mlx::core::fast::cuda_kernel(' '// kernel = mlx::core::fast::cuda_kernel(' +replace_in "$dir/$sampling" 'kernel = mlx::core::fast::hip_kernel(' '// kernel = mlx::core::fast::hip_kernel(' if run_case launch-commented-out "$dir" 1; then assert_contains launch-commented-out "$sampling: pinned in EXPECTED_IN_SCOPE but no longer calls" fi diff --git a/src/lib/mlx-cpp/turbo/sampling.cpp b/src/lib/mlx-cpp/turbo/sampling.cpp index b56025817..dda846136 100644 --- a/src/lib/mlx-cpp/turbo/sampling.cpp +++ b/src/lib/mlx-cpp/turbo/sampling.cpp @@ -15,6 +15,7 @@ #include "sampling.h" #include "gpu_backend.h" #include "kernel_port.h" +#include "sampling_gumbel_hip.h" #include #include @@ -29,6 +30,7 @@ #include #include +#include #include #include #include @@ -351,6 +353,39 @@ inline GumbelKernelHolderCuda& get_gumbel_kernel_cuda() { return *holder; } +// HIP counterpart (issue #2064), reached only on a ROCm build. The body lives +// in `sampling_gumbel_hip.h`; the launch stays in `gumbel_max_sample` below, +// shared with the other two ports. On a build without the ROCm backend +// `fast::hip_kernel` is not declared, and `gumbel_ports()` never resolves this +// entry there (`select_kernel_port` reads the running backend), so the throw is +// unreachable and exists only to keep the table shape uniform. +struct GumbelKernelHolderHip { + std::optional kernel; + std::once_flag init_flag; + + mlx::core::fast::CustomKernelFunction& get() { + std::call_once(init_flag, [this] { +#ifdef MLXCEL_BRIDGE_ROCM_BACKEND + kernel = mlx::core::fast::hip_kernel( + "mlxcel_gumbel_max_sample", + {"logits", "rng_key", "temp"}, + {"vals", "idxs"}, + std::string(GUMBEL_MAX_SAMPLE_HIP_SOURCE)); +#else + throw std::runtime_error( + "[gumbel_max_sample] this build has no ROCm backend"); +#endif + }); + return *kernel; + } +}; + +inline GumbelKernelHolderHip& get_gumbel_kernel_hip() { + // Leaked on purpose, for the reason given at `get_gumbel_kernel`. + static GumbelKernelHolderHip* holder = new GumbelKernelHolderHip(); + return *holder; +} + // This kernel's ports, in one place. `has_kernel_port` and // `select_kernel_port` both read it, so the support predicate below and the @@ -363,8 +398,9 @@ const KernelPorts& gumbel_ports() { .cuda = +[]() -> mlx::core::fast::CustomKernelFunction& { return get_gumbel_kernel_cuda().get(); }, - // No HIP port yet (#1814). - .rocm = nullptr, + .rocm = +[]() -> mlx::core::fast::CustomKernelFunction& { + return get_gumbel_kernel_hip().get(); + }, }; return ports; } diff --git a/src/lib/mlx-cpp/turbo/sampling.h b/src/lib/mlx-cpp/turbo/sampling.h index 4d64b4862..fc88bdc32 100644 --- a/src/lib/mlx-cpp/turbo/sampling.h +++ b/src/lib/mlx-cpp/turbo/sampling.h @@ -43,7 +43,9 @@ namespace mlxcel::turbo { // True when the Gumbel-max kernel can run on the active backend: a GPU default -// device plus an available Metal or CUDA backend. A CPU-only build, or a +// device whose backend has a port in `gumbel_ports()` (Metal, CUDA and, since +// #2064, ROCm). The predicate reads the same table as the dispatch, so the two +// cannot disagree. A CPU-only build, or a // process running with `MLXCEL_DEVICE=cpu`, returns false and the caller must // keep the `categorical` graph path. bool gumbel_max_sample_supported(); diff --git a/src/lib/mlx-cpp/turbo/sampling_gumbel_hip.h b/src/lib/mlx-cpp/turbo/sampling_gumbel_hip.h new file mode 100644 index 000000000..6c26eb937 --- /dev/null +++ b/src/lib/mlx-cpp/turbo/sampling_gumbel_hip.h @@ -0,0 +1,141 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +// HIP body of the Gumbel-max sampling kernel (issue #2064). +// +// Kept apart from `sampling.cpp` only for size: the launch, and the +// `fast::hip_kernel` call that compiles this string, stay in `sampling.cpp` +// next to the CUDA launch, so `make verify-kernel-dtype-keys` keeps its pinned +// scope unchanged. This header holds data, not a launch. +// +// A line-for-line port of `GUMBEL_MAX_SAMPLE_CUDA_SOURCE`: same thread mapping, +// same grid (`(TgSize, NumSplits, B)` over `(TgSize, 1, 1)`, which the ROCm +// `CustomKernel::eval_gpu` ceil-divides into blocks `(1, NumSplits, B)` exactly +// as CUDA does), same Philox-4x32-10 counter `{base/4, 0, row, 0}` and key +// `{rng_key[0], rng_key[1]}`, same 23-bit open-interval uniform, same +// index-carrying halving reduction. A given key and logits row therefore pick +// the same index on Metal, CUDA and ROCm up to `logf` rounding, and +// `mlx::core::random::seed(...)` reproduces a ROCm stream the way it does on +// the other two. `__umulhi` is the HIP intrinsic with CUDA's meaning. +// +// Differences from the CUDA text, all spelling: +// +// - The initial best score is `-__builtin_huge_valf()`. hipRTC compiles this +// body with only the headers `fast::hip_kernel` prepends, and `INFINITY` is a +// `` macro that is not guaranteed among them. +// - Every logits read is an explicit `(float)` conversion. `hip_bfloat16`, the +// type MLX substitutes for bfloat16 on ROCm, converts to float only through +// an `explicit` operator, so the implicit conversion the CUDA text relies on +// for `__nv_bfloat16` would not compile here. +// +// No wave32 guard, deliberately. The guard that the shuffle-based #1814 ports +// carry protects a lane fold that starts at 16 and so silently drops half of a +// 64-lane wave. This kernel has no lane-level operation: every cross-thread +// step goes through `__shared__` memory with a `__syncthreads()` between +// steps, so it is correct for any wavefront size, and an `#error` on wave64 +// would reject a correct kernel. (The guard idiom is also a no-op with ROCm +// 10's AMD clang, which defines neither wavefront macro; see #2067.) + +namespace mlxcel::turbo { + +inline constexpr const char* GUMBEL_MAX_SAMPLE_HIP_SOURCE = R"( + uint32_t t = threadIdx.x; // 0 .. TgSize-1 + uint32_t split = blockIdx.y; // 0 .. NumSplits-1 + uint32_t row = blockIdx.z; // 0 .. B-1 + + uint32_t vocab = (uint32_t)logits_shape[1]; + float temp_v = (float)temp[0]; + uint32_t key0_base = rng_key[0]; + uint32_t key1_base = rng_key[1]; + uint32_t row_off = row * vocab; + + float best = -__builtin_huge_valf(); + uint32_t best_idx = 0u; + + uint32_t chunk = (uint32_t)TgSize * 4u; + uint32_t stride = chunk * (uint32_t)NumSplits; + for (uint32_t base = split * chunk + t * 4u; base < vocab; base += stride) { + // Philox-4x32-10 over counter {base/4, 0, row, 0}. + uint32_t c0 = base >> 2; + uint32_t c1 = 0u; + uint32_t c2 = row; + uint32_t c3 = 0u; + uint32_t k0 = key0_base; + uint32_t k1 = key1_base; + for (uint32_t r = 0u; r < 10u; r++) { + uint32_t hi0 = __umulhi(0xD2511F53u, c0); + uint32_t lo0 = 0xD2511F53u * c0; + uint32_t hi1 = __umulhi(0xCD9E8D57u, c2); + uint32_t lo1 = 0xCD9E8D57u * c2; + uint32_t n0 = hi1 ^ c1 ^ k0; + uint32_t n1 = lo1; + uint32_t n2 = hi0 ^ c3 ^ k1; + uint32_t n3 = lo0; + c0 = n0; + c1 = n1; + c2 = n2; + c3 = n3; + k0 += 0x9E3779B9u; + k1 += 0xBB67AE85u; + } + + for (uint32_t j = 0u; j < 4u; j++) { + uint32_t idx = base + j; + if (idx >= vocab) { + break; + } + uint32_t word = + (j == 0u) ? c0 : ((j == 1u) ? c1 : ((j == 2u) ? c2 : c3)); + // Uniform on the OPEN interval (0, 1); see the Metal source for + // why this takes 23 bits rather than 24. + float u = ((float)(word >> 9) + 0.5f) * (1.0f / 8388608.0f); + float g = -logf(-logf(u)); + float scaled = (float)logits[row_off + idx] / temp_v; + float cand = scaled + g; + if (cand > best || (cand == best && idx < best_idx)) { + best = cand; + best_idx = idx; + } + } + } + + __shared__ float tg_val[TgSize]; + __shared__ uint32_t tg_idx[TgSize]; + tg_val[t] = best; + tg_idx[t] = best_idx; + __syncthreads(); + + for (uint32_t s = (uint32_t)TgSize / 2u; s > 0u; s >>= 1) { + if (t < s) { + float other = tg_val[t + s]; + uint32_t other_idx = tg_idx[t + s]; + if (other > tg_val[t] || + (other == tg_val[t] && other_idx < tg_idx[t])) { + tg_val[t] = other; + tg_idx[t] = other_idx; + } + } + __syncthreads(); + } + + if (t == 0u) { + uint32_t out_off = row * (uint32_t)NumSplits + split; + vals[out_off] = tg_val[0]; + idxs[out_off] = tg_idx[0]; + } +)"; + +} // namespace mlxcel::turbo diff --git a/src/lib/mlx-cpp/turbo/sampling_rejection.cpp b/src/lib/mlx-cpp/turbo/sampling_rejection.cpp index 038a965c3..3b0827d2e 100644 --- a/src/lib/mlx-cpp/turbo/sampling_rejection.cpp +++ b/src/lib/mlx-cpp/turbo/sampling_rejection.cpp @@ -15,6 +15,7 @@ #include "sampling_rejection.h" #include "gpu_backend.h" #include "kernel_port.h" +#include "sampling_rejection_hip.h" #include #include @@ -28,6 +29,7 @@ #include #include +#include #include #include #include @@ -742,6 +744,38 @@ inline RejectionKernelHolderCuda& get_rejection_kernel_cuda() { return *holder; } +// HIP counterpart (issue #2064), reached only on a ROCm build. The body lives +// in `sampling_rejection_hip.h`; the launch stays in `rejection_sample` below, +// shared with the other two ports. Without the ROCm backend `fast::hip_kernel` +// is not declared, and `rejection_ports()` never resolves this entry there, so +// the throw is unreachable and keeps the table shape uniform. +struct RejectionKernelHolderHip { + std::optional kernel; + std::once_flag init_flag; + + mlx::core::fast::CustomKernelFunction& get() { + std::call_once(init_flag, [this] { +#ifdef MLXCEL_BRIDGE_ROCM_BACKEND + kernel = mlx::core::fast::hip_kernel( + "mlxcel_rejection_sample", + {"probs", "probs_draw", "params", "rng_key"}, + {"ids", "ok", "rounds"}, + std::string(REJECTION_SAMPLE_HIP_SOURCE)); +#else + throw std::runtime_error( + "[fused_sample_rejection] this build has no ROCm backend"); +#endif + }); + return *kernel; + } +}; + +inline RejectionKernelHolderHip& get_rejection_kernel_hip() { + // Leaked on purpose, for the reason given at `get_rejection_kernel`. + static RejectionKernelHolderHip* holder = new RejectionKernelHolderHip(); + return *holder; +} + // This kernel's ports, in one place. `has_kernel_port` and // `select_kernel_port` both read it, so a support predicate and the dispatch @@ -754,8 +788,9 @@ const mlxcel::KernelPorts& rejection_ports() { .cuda = +[]() -> mlx::core::fast::CustomKernelFunction& { return get_rejection_kernel_cuda().get(); }, - // No HIP port yet (#1814). - .rocm = nullptr, + .rocm = +[]() -> mlx::core::fast::CustomKernelFunction& { + return get_rejection_kernel_hip().get(); + }, }; return ports; } @@ -766,7 +801,12 @@ bool rejection_sample_supported() { if (mlx::core::default_device() != mlx::core::Device::gpu) { return false; } - return mlxcel::custom_kernels_available(); + // The port table, not the backend kind: a backend-kind test answered + // false on ROCm whatever the table held, so a filled `.rocm` slot stayed + // unreachable (#2064). Reading the table keeps this predicate and the + // dispatch in `rejection_sample` from disagreeing, as in + // `gumbel_max_sample_supported()`. + return mlxcel::has_kernel_port(rejection_ports()); } bool rejection_sample_accepts(const mlx::core::array& probs) { diff --git a/src/lib/mlx-cpp/turbo/sampling_rejection.h b/src/lib/mlx-cpp/turbo/sampling_rejection.h index c99e6f275..0c79abb6a 100644 --- a/src/lib/mlx-cpp/turbo/sampling_rejection.h +++ b/src/lib/mlx-cpp/turbo/sampling_rejection.h @@ -149,7 +149,9 @@ namespace mlxcel::turbo { inline constexpr int REJECTION_MAX_ROUNDS = 32; // True when the rejection kernel can run on the active backend: a GPU default -// device plus an available Metal or CUDA backend. A CPU-only build, or a +// device whose backend has a port in `rejection_ports()` (Metal, CUDA and, since +// #2064, ROCm). The predicate reads the same table as the dispatch, so the two +// cannot disagree. A CPU-only build, or a // process under `MLXCEL_DEVICE=cpu`, returns false and the caller keeps the // `argpartition` chain. bool rejection_sample_supported(); diff --git a/src/lib/mlx-cpp/turbo/sampling_rejection_hip.h b/src/lib/mlx-cpp/turbo/sampling_rejection_hip.h new file mode 100644 index 000000000..297382608 --- /dev/null +++ b/src/lib/mlx-cpp/turbo/sampling_rejection_hip.h @@ -0,0 +1,318 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +// HIP body of the dual-pivot rejection sampler (issue #2064). +// +// Kept apart from `sampling_rejection.cpp` only for size: the launch, and the +// `fast::hip_kernel` call that compiles this string, stay in +// `sampling_rejection.cpp` next to the CUDA launch, so +// `make verify-kernel-dtype-keys` keeps its pinned scope unchanged. This +// header holds data, not a launch. +// +// A line-for-line port of `REJECTION_SAMPLE_CUDA_SOURCE`: same partition (one +// block of `TgSize` threads per row, thread `t` owning entries `t, t + TgSize, +// ...`), same Hillis-Steele scan, same Philox-4x32-10 counter +// `{round, 0, row, 0}` and key `{rng_key[0], rng_key[1]}`, same bit-space +// bisection. The grid `(TgSize, 1, B)` over `(TgSize, 1, 1)` becomes blocks +// `(1, 1, B)` in the ROCm `CustomKernel::eval_gpu` exactly as on CUDA. A given +// key and probability row therefore resolve the same support and draw the same +// index as the CUDA and Metal ports, and `mlx::core::random::seed(...)` +// reproduces a ROCm stream. +// +// Differences from the CUDA text, all spelling: every `probs`, `probs_draw` +// and `params` read is an explicit `(float)` conversion, because `hip_bfloat16` +// (MLX's bfloat16 type on ROCm) converts to float only through an `explicit` +// operator, and `rejection_sample_accepts` places no dtype restriction on the +// inputs. `__umulhi`, `__float_as_uint` and `__uint_as_float` are the HIP +// intrinsics with CUDA's meaning. +// +// No wave32 guard, deliberately. The guard that the shuffle-based #1814 ports +// carry protects a lane fold that starts at 16 and so silently drops half of a +// 64-lane wave. This kernel has no lane-level operation: the max, sum and +// owner reductions and the prefix scan all go through `__shared__` memory with +// a `__syncthreads()` between steps, so it is correct for any wavefront size, +// and an `#error` on wave64 would reject a correct kernel. (The guard idiom is +// also a no-op with ROCm 10's AMD clang, which defines neither wavefront +// macro; see #2067.) Every barrier is reached by every thread for the reason +// the CUDA source gives: the conditionals around them are block-uniform. + +namespace mlxcel::turbo { + +inline constexpr const char* REJECTION_SAMPLE_HIP_SOURCE = R"( + uint32_t t = threadIdx.x; // 0 .. TgSize-1 + uint32_t row = blockIdx.z; // 0 .. B-1 + + const uint32_t tg = (uint32_t)TgSize; + uint32_t vocab = (uint32_t)probs_shape[1]; + uint32_t row_off = row * vocab; + + float top_k_f = (float)params[row * 3u + 0u]; + float top_p = (float)params[row * 3u + 1u]; + float min_p = (float)params[row * 3u + 2u]; + uint32_t top_k = (top_k_f >= 1.0f) ? (uint32_t)top_k_f : 0u; + bool use_k = (top_k > 0u) && (top_k < vocab); + bool use_p = (top_p > 0.0f) && (top_p < 1.0f); + bool use_mp = (min_p > 0.0f) && (min_p < 1.0f); + + __shared__ float sh_max[TgSize]; + __shared__ uint32_t sh_arg[TgSize]; + __shared__ uint32_t sh_own[TgSize]; + __shared__ float sh_scan[TgSize]; + __shared__ float sh_c0[TgSize]; + __shared__ float sh_m0[TgSize]; + __shared__ float sh_c1[TgSize]; + __shared__ float sh_m1[TgSize]; + __shared__ uint32_t bcu[2]; + + // See the Metal source: the draw-row sum is gated on filter membership so + // the proposal mass never counts a token the pick loop cannot reach. + float local_max = -1.0f; + uint32_t local_arg = 0u; + float local_sum = 0.0f; + float local_draw = 0.0f; + for (uint32_t i = t; i < vocab; i += tg) { + float p = (float)probs[row_off + i]; + local_sum += p; + if (p > 0.0f) { + local_draw += (float)probs_draw[row_off + i]; + } + if (p > local_max) { + local_max = p; + local_arg = i; + } + } + + sh_max[t] = local_max; + sh_arg[t] = local_arg; + __syncthreads(); + for (uint32_t s = tg / 2u; s > 0u; s >>= 1) { + if (t < s) { + float o = sh_max[t + s]; + uint32_t oi = sh_arg[t + s]; + if (o > sh_max[t] || (o == sh_max[t] && oi < sh_arg[t])) { + sh_max[t] = o; + sh_arg[t] = oi; + } + } + __syncthreads(); + } + float p_max = sh_max[0]; + uint32_t arg_max = sh_arg[0]; + __syncthreads(); + + sh_scan[t] = local_sum; + __syncthreads(); + for (uint32_t s = tg / 2u; s > 0u; s >>= 1) { + if (t < s) { + sh_scan[t] += sh_scan[t + s]; + } + __syncthreads(); + } + float total = sh_scan[0]; + __syncthreads(); + + float low = 0.0f; + float part = local_draw; + if (use_mp) { + float tau_min = min_p * p_max; + uint32_t tau_bits = __float_as_uint(tau_min); + if (tau_bits > 0u) { + low = __uint_as_float(tau_bits - 1u); + part = 0.0f; + for (uint32_t i = t; i < vocab; i += tg) { + float p = (float)probs[row_off + i]; + if (p > low) { + part += (float)probs_draw[row_off + i]; + } + } + } + } + + float high = p_max; + float pmass_target = use_p ? (top_p * total) : 0.0f; + uint32_t sampled = arg_max; + uint32_t ok_flag = 0u; + uint32_t rounds_used = 0u; + + // See the Metal source: a fully -inf row softmaxes to NaN and is served by + // its argmax without entering the loop. + if (!(p_max > 0.0f)) { + ok_flag = 1u; + } + + for (uint32_t round = 0u; (p_max > 0.0f) && round < (uint32_t)MaxRounds; + round++) { + sh_scan[t] = part; + __syncthreads(); + for (uint32_t off = 1u; off < tg; off <<= 1) { + float add = (t >= off) ? sh_scan[t - off] : 0.0f; + __syncthreads(); + sh_scan[t] += add; + __syncthreads(); + } + float incl = sh_scan[t]; + float mass_prop = sh_scan[tg - 1u]; + float excl = incl - part; + __syncthreads(); + + uint32_t c0 = round; + uint32_t c1 = 0u; + uint32_t c2 = row; + uint32_t c3 = 0u; + uint32_t k0 = rng_key[0]; + uint32_t k1 = rng_key[1]; + for (uint32_t r = 0u; r < 10u; r++) { + uint32_t hi0 = __umulhi(0xD2511F53u, c0); + uint32_t lo0 = 0xD2511F53u * c0; + uint32_t hi1 = __umulhi(0xCD9E8D57u, c2); + uint32_t lo1 = 0xCD9E8D57u * c2; + uint32_t n0 = hi1 ^ c1 ^ k0; + uint32_t n1 = lo1; + uint32_t n2 = hi0 ^ c3 ^ k1; + uint32_t n3 = lo0; + c0 = n0; + c1 = n1; + c2 = n2; + c3 = n3; + k0 += 0x9E3779B9u; + k1 += 0xBB67AE85u; + } + float u = ((float)(c0 >> 9) + 0.5f) * (1.0f / 8388608.0f); + float target = u * mass_prop; + + uint32_t pick = 0xFFFFFFFFu; + if (incl > target) { + float run = excl; + for (uint32_t i = t; i < vocab; i += tg) { + float p = (float)probs[row_off + i]; + if (p > low) { + run += (float)probs_draw[row_off + i]; + if (run > target) { + pick = i; + break; + } + } + } + } + + sh_own[t] = (pick == 0xFFFFFFFFu) ? 0xFFFFFFFFu : t; + __syncthreads(); + for (uint32_t s = tg / 2u; s > 0u; s >>= 1) { + if (t < s) { + uint32_t o = sh_own[t + s]; + if (o < sh_own[t]) { + sh_own[t] = o; + } + } + __syncthreads(); + } + uint32_t owner = sh_own[0]; + __syncthreads(); + + if (t == 0u) { + bcu[0] = arg_max; + } + __syncthreads(); + if (owner != 0xFFFFFFFFu && t == owner) { + bcu[0] = pick; + } + __syncthreads(); + uint32_t cand = bcu[0]; + + float pivot0 = (float)probs[row_off + cand]; + + float hi_eff = (pivot0 > high) ? pivot0 : high; + uint32_t b0 = __float_as_uint(pivot0); + uint32_t bh = __float_as_uint(hi_eff); + uint32_t bm = (b0 >> 1u) + (bh >> 1u) + (b0 & bh & 1u); + float pivot1 = __uint_as_float(bm); + + float cnt0 = 0.0f; + float mass0 = 0.0f; + float cnt1 = 0.0f; + float mass1 = 0.0f; + float d0acc = 0.0f; + float d1acc = 0.0f; + if (use_k || use_p) { + float c0acc = 0.0f; + float c1acc = 0.0f; + float m0acc = 0.0f; + float m1acc = 0.0f; + for (uint32_t i = t; i < vocab; i += tg) { + float p = (float)probs[row_off + i]; + float pd = (float)probs_draw[row_off + i]; + if (p > pivot0) { + c0acc += 1.0f; + m0acc += p; + d0acc += pd; + } + if (p > pivot1) { + c1acc += 1.0f; + m1acc += p; + d1acc += pd; + } + } + sh_c0[t] = c0acc; + sh_m0[t] = m0acc; + sh_c1[t] = c1acc; + sh_m1[t] = m1acc; + __syncthreads(); + for (uint32_t s = tg / 2u; s > 0u; s >>= 1) { + if (t < s) { + sh_c0[t] += sh_c0[t + s]; + sh_m0[t] += sh_m0[t + s]; + sh_c1[t] += sh_c1[t + s]; + sh_m1[t] += sh_m1[t + s]; + } + __syncthreads(); + } + cnt0 = sh_c0[0]; + mass0 = sh_m0[0]; + cnt1 = sh_c1[0]; + mass1 = sh_m1[0]; + __syncthreads(); + } + + bool fail0 = (use_k && cnt0 >= (float)top_k) || + (use_p && mass0 > pmass_target); + bool fail1 = (use_k && cnt1 >= (float)top_k) || + (use_p && mass1 > pmass_target); + + rounds_used = round + 1u; + if (!fail0) { + sampled = cand; + ok_flag = 1u; + break; + } + + if (fail1) { + low = pivot1; + part = d1acc; + } else { + low = pivot0; + high = pivot1; + part = d0acc; + } + } + + if (t == 0u) { + ids[row] = sampled; + ok[row] = ok_flag; + rounds[row] = rounds_used; + } +)"; + +} // namespace mlxcel::turbo diff --git a/src/lib/mlxcel-core/build.rs b/src/lib/mlxcel-core/build.rs index da0bc9e11..bc3dfe8ae 100644 --- a/src/lib/mlxcel-core/build.rs +++ b/src/lib/mlxcel-core/build.rs @@ -301,8 +301,10 @@ fn main() { // Gumbel-max categorical sampling kernel launcher (#900). println!("cargo:rerun-if-changed=../mlx-cpp/turbo/sampling.h"); println!("cargo:rerun-if-changed=../mlx-cpp/turbo/sampling.cpp"); + println!("cargo:rerun-if-changed=../mlx-cpp/turbo/sampling_gumbel_hip.h"); println!("cargo:rerun-if-changed=../mlx-cpp/turbo/sampling_rejection.h"); println!("cargo:rerun-if-changed=../mlx-cpp/turbo/sampling_rejection.cpp"); + println!("cargo:rerun-if-changed=../mlx-cpp/turbo/sampling_rejection_hip.h"); // Fused residual-add RMSNorm and fused RoPE + KV-append kernel launchers (#905). println!("cargo:rerun-if-changed=../mlx-cpp/turbo/fused_norm.h"); println!("cargo:rerun-if-changed=../mlx-cpp/turbo/fused_norm.cpp"); diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index 15300be13..682cc02b9 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -5307,7 +5307,7 @@ namespace { } // True when the no-filter sampling path will use the Gumbel-max kernel: the -// backend supports it (GPU default device with Metal or CUDA available) and +// default device is the GPU, its backend has a port in `gumbel_ports()`, and // `MLXCEL_SAMPLING_GUMBEL` is not falsy. bool sampling_gumbel_available() { static const bool env_enabled = parse_sampling_gumbel_env(); @@ -7393,6 +7393,14 @@ std::unique_ptr random_gumbel(rust::Slice shape, int32_ return std::make_unique(mlx::core::random::gumbel(to_shape(shape), to_dtype(dtype), opt_key)); } +// Raw random words from `random::bits`. With a null key it draws from the +// default key sequence exactly as the fused samplers' launchers do, which is +// how `sampling_fixed_key_tests` recovers the key a launch will use (#2064). +std::unique_ptr random_bits(rust::Slice shape, int32_t width, const MlxArray* key) { + std::optional opt_key = key ? std::make_optional(key->inner) : std::nullopt; + return std::make_unique(mlx::core::random::bits(to_shape(shape), width, opt_key)); +} + std::unique_ptr random_laplace(rust::Slice shape, int32_t dtype, const MlxArray* key) { std::optional opt_key = key ? std::make_optional(key->inner) : std::nullopt; return std::make_unique(mlx::core::random::laplace(to_shape(shape), to_dtype(dtype), opt_key)); diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index e68c271d6..2344c35fa 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -1824,6 +1824,7 @@ std::unique_ptr random_bernoulli_p(float p, rust::Slice std::unique_ptr random_randint(int32_t low, int32_t high, rust::Slice shape, int32_t dtype, const MlxArray* key); std::unique_ptr random_truncated_normal(float lower, float upper, rust::Slice shape, int32_t dtype, const MlxArray* key); std::unique_ptr random_gumbel(rust::Slice shape, int32_t dtype, const MlxArray* key); +std::unique_ptr random_bits(rust::Slice shape, int32_t width, const MlxArray* key); std::unique_ptr random_laplace(rust::Slice shape, int32_t dtype, const MlxArray* key); std::unique_ptr random_permutation(int32_t x, const MlxArray* key); std::unique_ptr random_permutation_array(const MlxArray& a, int32_t axis, const MlxArray* key); diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index bcc7b13c5..355a384eb 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -2496,8 +2496,9 @@ mod ffi { fn gumbel_max_sample(logits: &MlxArray, temperature: f32) -> Result>; /// True when [`fused_sample`]'s no-filter stochastic path takes the - /// Gumbel-max kernel: the backend supports it (GPU default device with - /// Metal or CUDA available) and `MLXCEL_SAMPLING_GUMBEL` is not falsy. + /// Gumbel-max kernel: the default device is the GPU, its backend has a + /// Gumbel-max port (Metal, CUDA or ROCm) and `MLXCEL_SAMPLING_GUMBEL` + /// is not falsy. /// The env value is read once per process. fn sampling_gumbel_available() -> bool; @@ -3033,6 +3034,16 @@ mod ffi { key: *const MlxArray, ) -> UniquePtr; + /// Raw random words (`random::bits`): `width` bytes per element, + /// returned as uint32 / uint16 / uint8. A null `key` draws from the + /// default key sequence, which is how the fused samplers draw their + /// Philox key, so a test can recover the key a launch will use. + unsafe fn random_bits( + shape: &[i32], + width: i32, + key: *const MlxArray, + ) -> UniquePtr; + /// Gumbel random samples unsafe fn random_gumbel( shape: &[i32], @@ -3961,6 +3972,13 @@ mod sampling_gumbel_tests; #[path = "sampling_rejection_tests.rs"] mod sampling_rejection_tests; +// Fixed-key and graph-parity tests shared by every port of the two sampler +// kernels (#2064): the kernel's draw recomputed from the key it consumed, and a +// two-sample test against the graph sampler. GPU-only; they skip without a port. +#[cfg(test)] +#[path = "sampling_fixed_key_tests.rs"] +mod sampling_fixed_key_tests; + // Numeric-parity, Gemma `(1 + w)` convention, kill-switch and greedy-argmax // tests for the fused residual-add RMSNorm kernel (#905). GPU-only; they skip // on CPU-only builds. diff --git a/src/lib/mlxcel-core/src/sampling_fixed_key_tests.rs b/src/lib/mlxcel-core/src/sampling_fixed_key_tests.rs new file mode 100644 index 000000000..6ff245312 --- /dev/null +++ b/src/lib/mlxcel-core/src/sampling_fixed_key_tests.rs @@ -0,0 +1,476 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Port-against-graph tests for the fused sampler kernels (issue #2064). +//! +//! `sampling_gumbel_tests.rs` and `sampling_rejection_tests.rs` hold each +//! kernel to the exact distribution it should draw from. This file adds the +//! two comparisons a new backend port needs on top of that, and runs on every +//! backend whose port table has an entry (Metal, CUDA, ROCm): +//! +//! 1. **Fixed key, deterministic.** Both kernels draw one Philox-4x32-10 key +//! from MLX's default key sequence per call, and given that key their +//! output is a pure function of the input. The key is recovered here by +//! reseeding and calling `random::bits` exactly as the launcher does, the +//! kernel's draw is recomputed on the host from that key, and for the +//! Gumbel-max kernel the same noise is also pushed through the MLX graph +//! (`argmax(logits / T + g)`). A port that changed the counter layout, the +//! key words, the uniform mapping, the partition order or a dtype read +//! disagrees on almost every row. Rows whose decision sits within rounding +//! of a tie are skipped and counted, since `logf` and summation order are +//! allowed to differ from the host in the last bits. +//! 2. **Statistical, against the graph sampler.** Where the kernel and the +//! graph consume the RNG differently (`random::categorical`, the stock +//! filter chain), a fixed key cannot line them up, so a two-sample +//! chi-square homogeneity test compares the kernel's histogram with +//! `fused_sample_categorical`'s, the explicit graph arm, on the same input. +//! +//! GPU-only: each test returns early when the kernel's support predicate says +//! this backend has no port, the convention of the two files above. + +use super::*; +use crate::dtype; + +/// Upper-tail standard normal quantile for p = 1e-6, as in the sibling files. +const CRITICAL_Z: f64 = 4.7534; + +/// Minimum pooled count per bin for the two-sample test. +const MIN_POOLED_PER_BIN: u64 = 20; + +/// Draws per arm for the two-sample tests. +const TWO_SAMPLE_DRAWS: usize = 400_000; + +/// Smallest gap, between the winning and the runner-up Gumbel score, at which +/// a row counts as decided. GPU `logf` and the host's f64 `ln` may differ in +/// the last bits of a score of order ten, far below this. +const GUMBEL_DECIDED_GAP: f64 = 1e-4; + +/// Smallest distance, relative to the proposal mass, between the rejection +/// kernel's draw target and the nearest cumulative boundary at which a row +/// counts as decided. The kernel sums in f32, the host in f64. +const REJECTION_DECIDED_GAP: f64 = 1e-5; + +// -- host reference of the kernels' RNG -- + +/// Philox-4x32-10 exactly as both kernel sources spell it. +fn philox4x32_10(counter: [u32; 4], key: [u32; 2]) -> [u32; 4] { + let [mut c0, mut c1, mut c2, mut c3] = counter; + let [mut k0, mut k1] = key; + for _ in 0..10 { + let p0 = u64::from(0xD251_1F53u32) * u64::from(c0); + let p1 = u64::from(0xCD9E_8D57u32) * u64::from(c2); + let (hi0, lo0) = ((p0 >> 32) as u32, p0 as u32); + let (hi1, lo1) = ((p1 >> 32) as u32, p1 as u32); + let n0 = hi1 ^ c1 ^ k0; + let n2 = hi0 ^ c3 ^ k1; + c0 = n0; + c1 = lo1; + c2 = n2; + c3 = lo0; + k0 = k0.wrapping_add(0x9E37_79B9); + k1 = k1.wrapping_add(0xBB67_AE85); + } + [c0, c1, c2, c3] +} + +/// The kernels' uniform on the open interval (0, 1): the top 23 bits plus a +/// half step, computed in f32 as the kernels do (both steps are exact). +fn uniform_open(word: u32) -> f32 { + ((word >> 9) as f32 + 0.5) * (1.0 / 8_388_608.0) +} + +/// The key the next kernel launch will draw after `random_seed(seed)`. +/// +/// The launchers call `random::bits({2}, 4)` with no key, which takes the next +/// key from the default sequence; this makes the same call, so reseeding with +/// the same value afterwards hands the launch the same key. +fn key_after_seed(seed: u64) -> [u32; 2] { + random_seed(seed); + // SAFETY: a null key pointer asks for MLX's default key sequence; the + // shape slice outlives the call. + let bits = unsafe { random_bits(&[2], 4, std::ptr::null()) }; + let words = u32_values(&bits); + assert_eq!(words.len(), 2, "random_bits returned {} words", words.len()); + [words[0], words[1]] +} + +// -- small helpers -- + +fn u32_values(arr: &MlxArray) -> Vec { + array_to_raw_bytes(arr) + .chunks_exact(4) + .map(|c| u32::from_ne_bytes([c[0], c[1], c[2], c[3]])) + .collect() +} + +fn f32_values(arr: &MlxArray) -> Vec { + let as_f32 = astype(arr, dtype::FLOAT32); + array_to_raw_bytes(&as_f32) + .chunks_exact(4) + .map(|c| f32::from_ne_bytes([c[0], c[1], c[2], c[3]])) + .collect() +} + +/// Deterministic pseudo-random logits in `[-span, span]`, distinct per row, +/// with a handful of raised entries per row so a filter keeps a support of +/// tens of tokens spread over the whole row. +fn row_logits(rows: usize, vocab: usize, span: f32, seed: u64) -> Vec { + let mut state = seed | 1; + let mut next = || { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + ((state >> 40) as f32) / ((1u64 << 24) as f32) + }; + let mut out = Vec::with_capacity(rows * vocab); + for _ in 0..rows { + let start = out.len(); + for _ in 0..vocab { + out.push((next() * 2.0 - 1.0) * span); + } + for _ in 0..24 { + let i = (next() * vocab as f32) as usize % vocab; + out[start + i] = span + 2.0 + next() * 3.0; + } + } + out +} + +/// Two-sample chi-square homogeneity statistic for equal draw counts, with +/// cells pooled (largest first) until each bin holds [`MIN_POOLED_PER_BIN`] +/// draws across both arms. Returns `(statistic, degrees of freedom)`. +fn two_sample_chi_square(a: &[u64], b: &[u64]) -> (f64, f64) { + let mut cells: Vec<(u64, u64)> = a + .iter() + .zip(b) + .map(|(&x, &y)| (x, y)) + .filter(|&(x, y)| x + y > 0) + .collect(); + cells.sort_by_key(|&(x, y)| std::cmp::Reverse(x + y)); + let mut bins: Vec<(u64, u64)> = Vec::new(); + let mut acc = (0u64, 0u64); + for (x, y) in cells { + acc = (acc.0 + x, acc.1 + y); + if acc.0 + acc.1 >= MIN_POOLED_PER_BIN { + bins.push(acc); + acc = (0, 0); + } + } + if acc.0 + acc.1 > 0 { + match bins.last_mut() { + Some(last) => *last = (last.0 + acc.0, last.1 + acc.1), + None => bins.push(acc), + } + } + let stat = bins + .iter() + .map(|&(x, y)| { + let d = x as f64 - y as f64; + d * d / (x + y) as f64 + }) + .sum(); + (stat, (bins.len().max(2) - 1) as f64) +} + +fn chi_square_upper_critical(df: f64, z: f64) -> f64 { + let t = 2.0 / (9.0 * df); + let x = 1.0 - t + z * t.sqrt(); + df * x * x * x +} + +fn histogram_of(vocab: usize, n: usize, mut draw: impl FnMut() -> Vec) -> Vec { + let mut counts = vec![0u64; vocab]; + let mut drawn = 0usize; + while drawn < n { + for id in draw() { + if drawn >= n { + break; + } + counts[id as usize] += 1; + drawn += 1; + } + } + counts +} + +fn tiled(row: &[f32], rows: usize) -> UniquePtr { + let mut data = Vec::with_capacity(rows * row.len()); + for _ in 0..rows { + data.extend_from_slice(row); + } + from_slice_f32(&data, &[rows as i32, row.len() as i32]) +} + +// -- 1. fixed key -- + +#[test] +fn gumbel_kernel_matches_the_keyed_graph_argmax() { + if !sampling_gumbel_available() { + return; + } + // 5003 is not a multiple of four, so the last Philox block is partial, and + // spans several 1024-entry sweeps; 48 rows take two splits per row. + let rows = 48usize; + let vocab = 5003usize; + let cases = [ + (dtype::FLOAT32, 1.0f32), + (dtype::FLOAT32, 0.7), + (dtype::BFLOAT16, 1.3), + (dtype::FLOAT16, 0.9), + ]; + for (case, &(dt, temperature)) in cases.iter().enumerate() { + let host = row_logits(rows, vocab, 4.0, 0x2064 + case as u64); + let mut logits = from_slice_f32(&host, &[rows as i32, vocab as i32]); + if dt != dtype::FLOAT32 { + logits = astype(&logits, dt); + } + // The values the kernel reads, after any rounding to `dt`. + let seen = f32_values(&logits); + + let seed = 0x5EED_2064_0000 + case as u64; + let key = key_after_seed(seed); + random_seed(seed); + let kernel_ids = u32_values( + &gumbel_max_sample(&logits, temperature) + .expect("sampling_gumbel_available() reported a port"), + ); + + // Host scores in f64 from the kernel's own key, and the same noise + // pushed through the MLX graph. + let mut noise = vec![0f32; rows * vocab]; + let mut expected = vec![0usize; rows]; + let mut decided = vec![false; rows]; + for row in 0..rows { + let mut best = (f64::NEG_INFINITY, 0usize); + let mut second = f64::NEG_INFINITY; + for block in 0..vocab.div_ceil(4) { + let words = philox4x32_10([block as u32, 0, row as u32, 0], key); + for (j, &word) in words.iter().enumerate() { + let idx = block * 4 + j; + if idx >= vocab { + break; + } + let u = f64::from(uniform_open(word)); + let g = -(-u.ln()).ln(); + noise[row * vocab + idx] = g as f32; + let scaled = f64::from(seen[row * vocab + idx] / temperature); + let score = scaled + g; + if score > best.0 { + second = best.0; + best = (score, idx); + } else if score > second { + second = score; + } + } + } + expected[row] = best.1; + decided[row] = best.0 - second >= GUMBEL_DECIDED_GAP; + } + let noise = from_slice_f32(&noise, &[rows as i32, vocab as i32]); + let t = from_slice_f32(&[temperature], &[1]); + let scaled = divide(&astype(&logits, dtype::FLOAT32), &t); + let graph_ids = u32_values(&argmax(&add(&scaled, &noise), -1, false)); + + let label = format!("dtype code {dt}, T={temperature}"); + let mut checked = 0usize; + for row in 0..rows { + if !decided[row] { + continue; + } + checked += 1; + assert_eq!( + kernel_ids[row] as usize, expected[row], + "{label}: row {row} kernel drew {} but the keyed host reference drew {}", + kernel_ids[row], expected[row] + ); + assert_eq!( + graph_ids[row] as usize, expected[row], + "{label}: row {row} graph argmax {} disagrees with the host reference {}", + graph_ids[row], expected[row] + ); + } + assert!( + checked * 10 >= rows * 9, + "{label}: only {checked} of {rows} rows were decided; the test lost its power" + ); + } +} + +#[test] +fn rejection_kernel_draws_the_keyed_token_under_min_p() { + // The probe does not park flags in the deferred ring, but the shared guard + // keeps this test from interleaving with one that inspects it. + let _dispatch = crate::sampling_dispatch::dispatch_test_guard(); + if !sampling_rejection_available() { + return; + } + // min-p alone resolves its whole threshold before the first draw, so the + // first round always accepts and the sampled token is fully determined by + // the round-0 Philox word and the kernel's partition order: thread `t` of + // the row's block owns entries `t, t + tg, ...`, and the block scans + // threads in order. That is what the host reproduces below. + let tg = sampling_rejection_threadgroup_size() as usize; + let rows = 40usize; + let vocab = 3001usize; + let min_p = 0.05f32; + for (case, &temperature) in [1.0f32, 0.7].iter().enumerate() { + let host = row_logits(rows, vocab, 3.0, 0x0901 + case as u64); + let logits = from_slice_f32(&host, &[rows as i32, vocab as i32]); + // The kernel's two inputs, built with the bridge's own ops: the filter + // row is the untempered softmax, the draw row the tempered one. + let filter = f32_values(&softmax(&logits, -1)); + let draw = if temperature == 1.0 { + filter.clone() + } else { + let t = from_slice_f32(&[temperature], &[1]); + f32_values(&softmax(÷(&logits, &t), -1)) + }; + + let seed = 0x5EED_2064_0100 + case as u64; + let key = key_after_seed(seed); + random_seed(seed); + let stacked = sampling_rejection_probe(&logits, temperature, 0, 1.0, min_p, 32); + let flat = u32_values(&stacked); + let (ids, ok, used) = (&flat[..rows], &flat[rows..2 * rows], &flat[2 * rows..]); + + let label = format!("min_p={min_p}, T={temperature}"); + let mut checked = 0usize; + for row in 0..rows { + assert_eq!(ok[row], 1, "{label}: row {row} did not converge"); + assert_eq!(used[row], 1, "{label}: row {row} took {} rounds", used[row]); + + let f = &filter[row * vocab..(row + 1) * vocab]; + let d = &draw[row * vocab..(row + 1) * vocab]; + let p_max = f.iter().cloned().fold(-1.0f32, f32::max); + let low = f32::from_bits((min_p * p_max).to_bits() - 1); + + let order: Vec = (0..tg) + .flat_map(|t| (t..vocab).step_by(tg)) + .filter(|&i| f[i] > low) + .collect(); + let mass: f64 = order.iter().map(|&i| f64::from(d[i])).sum(); + let u = uniform_open(philox4x32_10([0, 0, row as u32, 0], key)[0]); + let target = f64::from(u) * mass; + + let mut run = 0.0f64; + let mut pick = None; + let mut gap = f64::INFINITY; + for &i in &order { + gap = gap.min((run - target).abs()); + run += f64::from(d[i]); + if pick.is_none() && run > target { + pick = Some(i); + } + } + gap = gap.min((run - target).abs()); + if gap < REJECTION_DECIDED_GAP * mass { + continue; + } + checked += 1; + let expected = pick.expect("target lies below the proposal mass"); + assert_eq!( + ids[row] as usize, expected, + "{label}: row {row} kernel drew {} but the keyed host reference drew {expected}", + ids[row] + ); + } + assert!( + checked * 10 >= rows * 9, + "{label}: only {checked} of {rows} rows were decided; the test lost its power" + ); + } +} + +// -- 2. statistical, against the graph sampler -- + +#[test] +fn gumbel_kernel_and_graph_categorical_draw_one_distribution() { + if !sampling_gumbel_available() { + return; + } + let vocab = 2048usize; + let row = row_logits(1, vocab, 2.0, 0x900); + let rows = 2048usize; + let batched = tiled(&row, rows); + let temperature = 0.8f32; + + random_seed(0x2064_0900); + let kernel = histogram_of(vocab, TWO_SAMPLE_DRAWS, || { + u32_values( + &gumbel_max_sample(&batched, temperature) + .expect("sampling_gumbel_available() reported a port"), + ) + }); + let graph = histogram_of(vocab, TWO_SAMPLE_DRAWS, || { + u32_values(&fused_sample_categorical( + &batched, + temperature, + 0, + 1.0, + 0.0, + )) + }); + + let (stat, df) = two_sample_chi_square(&kernel, &graph); + let critical = chi_square_upper_critical(df, CRITICAL_Z); + assert!( + stat <= critical, + "Gumbel-max kernel and random::categorical disagree: chi2 {stat:.1} > {critical:.1} (df {df})" + ); +} + +#[test] +fn rejection_kernel_and_graph_chain_draw_one_distribution() { + let _dispatch = crate::sampling_dispatch::dispatch_test_guard(); + if !sampling_rejection_available() { + return; + } + let vocab = 4096usize; + let row = row_logits(1, vocab, 2.0, 0x901); + let rows = 1024usize; + let batched = tiled(&row, rows); + let temperature = 0.8f32; + // (top_k, top_p, min_p): top-p alone (the routed configuration) and + // top-p with min-p. Both are exact in the kernel, unlike top-k with top-p + // (see `sampling_rejection_tests.rs`), so the two arms must agree. + for (top_k, top_p, min_p) in [(0i32, 0.9f32, 0.0f32), (0, 0.95, 0.02)] { + random_seed(0x2064_0901); + let kernel = histogram_of(vocab, TWO_SAMPLE_DRAWS, || { + let stacked = sampling_rejection_probe(&batched, temperature, top_k, top_p, min_p, 32); + let flat = u32_values(&stacked); + assert!( + flat[rows..2 * rows].iter().all(|&ok| ok == 1), + "a row exhausted the round cap" + ); + flat[..rows].to_vec() + }); + let graph = histogram_of(vocab, TWO_SAMPLE_DRAWS, || { + u32_values(&fused_sample_categorical( + &batched, + temperature, + top_k, + top_p, + min_p, + )) + }); + + let (stat, df) = two_sample_chi_square(&kernel, &graph); + let critical = chi_square_upper_critical(df, CRITICAL_Z); + assert!( + stat <= critical, + "top_p={top_p} min_p={min_p}: rejection kernel and the stock chain disagree: \ + chi2 {stat:.1} > {critical:.1} (df {df})" + ); + } +} diff --git a/src/lib/mlxcel-core/src/sampling_gumbel_tests.rs b/src/lib/mlxcel-core/src/sampling_gumbel_tests.rs index 4b46bc024..c52fc1988 100644 --- a/src/lib/mlxcel-core/src/sampling_gumbel_tests.rs +++ b/src/lib/mlxcel-core/src/sampling_gumbel_tests.rs @@ -36,13 +36,18 @@ //! active, greedy stays byte-identical to `argmax`, and every filtered //! configuration stays bit-identical to the pre-#900 categorical path. //! -//! GPU-only: the kernel JITs through `mx.fast.metal_kernel` / -//! `mx.fast.cuda_kernel`, so every test returns early on a CPU-only build, +//! GPU-only: the kernel JITs through `mx.fast.metal_kernel`, +//! `mx.fast.cuda_kernel` or `mx.fast.hip_kernel` (#2064), so every test returns +//! early where the support predicate reports no port (a CPU-only build), //! matching the convention in `fused_moe_parity_tests.rs`. //! //! Run on Apple Silicon: //! cargo test --release -p mlxcel-core --lib --features metal,accelerate \ //! sampling_gumbel_tests:: +//! +//! Run on ROCm: +//! cargo test --release -p mlxcel-core --lib --features rocm \ +//! sampling_gumbel_tests:: -- --test-threads=1 use super::*; diff --git a/src/lib/mlxcel-core/src/sampling_rejection_tests.rs b/src/lib/mlxcel-core/src/sampling_rejection_tests.rs index b19bc8fdf..d2744a881 100644 --- a/src/lib/mlxcel-core/src/sampling_rejection_tests.rs +++ b/src/lib/mlxcel-core/src/sampling_rejection_tests.rs @@ -70,13 +70,18 @@ //! limit that makes top-k slow is ever lifted, widening the policy is a //! one-constant change against tests that already prove the kernel right. //! -//! GPU-only: the kernel JITs through `mx.fast.metal_kernel` / -//! `mx.fast.cuda_kernel`, so every test returns early on a CPU-only build, +//! GPU-only: the kernel JITs through `mx.fast.metal_kernel`, +//! `mx.fast.cuda_kernel` or `mx.fast.hip_kernel` (#2064), so every test returns +//! early where the support predicate reports no port (a CPU-only build), //! matching the convention in `sampling_gumbel_tests.rs`. //! //! Run on Apple Silicon: //! cargo test --release -p mlxcel-core --lib --features metal,accelerate \ //! sampling_rejection_tests:: +//! +//! Run on ROCm: +//! cargo test --release -p mlxcel-core --lib --features rocm \ +//! sampling_rejection_tests:: -- --test-threads=1 use super::*; use std::collections::BTreeSet; From 8b0e858dca0d7a324917d7b022280b5bdb40a8d9 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 5 Oct 2026 01:12:44 +0900 Subject: [PATCH 3/7] docs(rocm): record sampled decode with the HIP samplers on gfx1151 `bench_decode.sh` gains `--temperature` and `--top-p`, forwarded to `mlxcel-bench-decode`, so a sampled decode can be measured with the standard harness; the CSV schema is unchanged and the auto-generated filename is tagged instead. Llama-3.1-8B-4bit on gfx1151, three interleaved guarded runs per arm: `--temperature 0.8 --top-p 0.95` (rejection kernel) went from a median of 33.98 to 37.56 tok/s with no overlap between arms; `--temperature 0.8` alone (Gumbel-max) moved 0.9%, inside the run-to-run spread. The issue's `--top-k 40 --top-p 0.95` command is not routed to the kernel at vocab 128256 and was not measured. The results page also records the fixed-key and two-sample tests and a speculative decode check with a Qwen3 drafter; `docs/installation.md` gains a sampled-decode row. Refs #2064 --- ...gfx1151_2026-10-05_samplers-after_t0.8.csv | 4 ++ ...51_2026-10-05_samplers-after_t0.8p0.95.csv | 4 ++ ...fx1151_2026-10-05_samplers-before_t0.8.csv | 4 ++ ...1_2026-10-05_samplers-before_t0.8p0.95.csv | 4 ++ .../rocm-samplers-gfx1151-2026-10-05.md | 47 +++++++++++++++++++ docs/benchmarks.md | 2 + docs/installation.md | 10 ++-- scripts/bench_decode.sh | 19 ++++++++ 8 files changed, 90 insertions(+), 4 deletions(-) create mode 100644 benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-after_t0.8.csv create mode 100644 benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-after_t0.8p0.95.csv create mode 100644 benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-before_t0.8.csv create mode 100644 benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-before_t0.8p0.95.csv create mode 100644 docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md diff --git a/benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-after_t0.8.csv b/benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-after_t0.8.csv new file mode 100644 index 000000000..c1e5f681f --- /dev/null +++ b/benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-after_t0.8.csv @@ -0,0 +1,4 @@ +model,model_path,prompt_tokens,generated_tokens,prefill_ms,prefill_tok_s,decode_ms,decode_tok_s,date,hardware,mlxcel_version,build_type,max_tokens,prompt,prompt_target_len,mlxcel_commit,mlx_commit,mlx_rocm_overlay_commit,hip_version +Meta-Llama-3.1-8B-Instruct-4bit,models/mlx/Meta-Llama-3.1-8B-Instruct-4bit,512,128,482.58,1060.96,3405.23,37.59,2026-10-05,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,85e39880,81ba1c6a,75915908,7.15.26333 +Meta-Llama-3.1-8B-Instruct-4bit,models/mlx/Meta-Llama-3.1-8B-Instruct-4bit,512,128,479.67,1067.39,3345.17,38.26,2026-10-05,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,85e39880,81ba1c6a,75915908,7.15.26333 +Meta-Llama-3.1-8B-Instruct-4bit,models/mlx/Meta-Llama-3.1-8B-Instruct-4bit,512,128,489.90,1045.10,3261.41,39.25,2026-10-05,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,85e39880,81ba1c6a,75915908,7.15.26333 diff --git a/benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-after_t0.8p0.95.csv b/benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-after_t0.8p0.95.csv new file mode 100644 index 000000000..b3ed72464 --- /dev/null +++ b/benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-after_t0.8p0.95.csv @@ -0,0 +1,4 @@ +model,model_path,prompt_tokens,generated_tokens,prefill_ms,prefill_tok_s,decode_ms,decode_tok_s,date,hardware,mlxcel_version,build_type,max_tokens,prompt,prompt_target_len,mlxcel_commit,mlx_commit,mlx_rocm_overlay_commit,hip_version +Meta-Llama-3.1-8B-Instruct-4bit,models/mlx/Meta-Llama-3.1-8B-Instruct-4bit,512,128,482.36,1061.45,3384.13,37.82,2026-10-05,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,85e39880,81ba1c6a,75915908,7.15.26333 +Meta-Llama-3.1-8B-Instruct-4bit,models/mlx/Meta-Llama-3.1-8B-Instruct-4bit,512,128,482.84,1060.40,3407.90,37.56,2026-10-05,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,85e39880,81ba1c6a,75915908,7.15.26333 +Meta-Llama-3.1-8B-Instruct-4bit,models/mlx/Meta-Llama-3.1-8B-Instruct-4bit,512,128,487.12,1051.07,3453.93,37.06,2026-10-05,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,85e39880,81ba1c6a,75915908,7.15.26333 diff --git a/benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-before_t0.8.csv b/benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-before_t0.8.csv new file mode 100644 index 000000000..056c049f0 --- /dev/null +++ b/benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-before_t0.8.csv @@ -0,0 +1,4 @@ +model,model_path,prompt_tokens,generated_tokens,prefill_ms,prefill_tok_s,decode_ms,decode_tok_s,date,hardware,mlxcel_version,build_type,max_tokens,prompt,prompt_target_len,mlxcel_commit,mlx_commit,mlx_rocm_overlay_commit,hip_version +Meta-Llama-3.1-8B-Instruct-4bit,models/mlx/Meta-Llama-3.1-8B-Instruct-4bit,512,128,506.71,1010.45,3376.66,37.91,2026-10-05,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,57d8ed29,81ba1c6a,75915908,7.15.26333 +Meta-Llama-3.1-8B-Instruct-4bit,models/mlx/Meta-Llama-3.1-8B-Instruct-4bit,512,128,482.35,1061.48,3397.80,37.67,2026-10-05,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,57d8ed29,81ba1c6a,75915908,7.15.26333 +Meta-Llama-3.1-8B-Instruct-4bit,models/mlx/Meta-Llama-3.1-8B-Instruct-4bit,512,128,481.34,1063.71,3340.17,38.32,2026-10-05,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,57d8ed29,81ba1c6a,75915908,7.15.26333 diff --git a/benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-before_t0.8p0.95.csv b/benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-before_t0.8p0.95.csv new file mode 100644 index 000000000..a6da73872 --- /dev/null +++ b/benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-before_t0.8p0.95.csv @@ -0,0 +1,4 @@ +model,model_path,prompt_tokens,generated_tokens,prefill_ms,prefill_tok_s,decode_ms,decode_tok_s,date,hardware,mlxcel_version,build_type,max_tokens,prompt,prompt_target_len,mlxcel_commit,mlx_commit,mlx_rocm_overlay_commit,hip_version +Meta-Llama-3.1-8B-Instruct-4bit,models/mlx/Meta-Llama-3.1-8B-Instruct-4bit,512,128,483.78,1058.33,3529.96,36.26,2026-10-05,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,57d8ed29,81ba1c6a,75915908,7.15.26333 +Meta-Llama-3.1-8B-Instruct-4bit,models/mlx/Meta-Llama-3.1-8B-Instruct-4bit,512,128,483.76,1058.37,3767.19,33.98,2026-10-05,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,57d8ed29,81ba1c6a,75915908,7.15.26333 +Meta-Llama-3.1-8B-Instruct-4bit,models/mlx/Meta-Llama-3.1-8B-Instruct-4bit,512,128,486.86,1051.64,3901.08,32.81,2026-10-05,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,57d8ed29,81ba1c6a,75915908,7.15.26333 diff --git a/docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md b/docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md new file mode 100644 index 000000000..53a4c5ca1 --- /dev/null +++ b/docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md @@ -0,0 +1,47 @@ +# HIP ports of the fused samplers on gfx1151 (2026-10-05) + +lablup/mlxcel#2064, part of #1814. Before this change both fused samplers took the MLX graph on ROCm: `gumbel_ports()` and `rejection_ports()` had no `.rocm` entry, and `rejection_sample_supported()` answered with `custom_kernels_available()`, which is Metal-or-CUDA by definition, so a filled slot would have stayed unreachable. Sampled decode therefore ran `random::categorical` on the no-filter path and the `argpartition` / `argsort` / `cumsum` chain on the filtered path. The decode profile ([rocm-decode-profile-gfx1151-2026-09-30.md](rocm-decode-profile-gfx1151-2026-09-30.md)) put the whole sampler tail at 0 of greedy decode GPU time and 0.4 to 3.8% of sampled decode GPU time. + +The ports (`sampling_gumbel_hip.h`, `sampling_rejection_hip.h` in `src/lib/mlx-cpp/turbo/`) are the CUDA bodies with the same inputs, outputs, grid, template arguments and Philox-4x32-10 counter and key layout, so a seed reproduces a ROCm stream as it does on CUDA and Metal. Neither kernel has a lane-level operation (every reduction and the rejection kernel's scan go through shared memory with a barrier per step), so neither carries a wave32 guard. `rejection_sample_supported()` now reads `has_kernel_port(rejection_ports())`. + +The Gumbel-max port exposed a fault in the vendored `fast::hip_kernel`: it declared an input's `_shape` as a pointer while the launch passes the shape by value, so the first kernel to read one (`logits_shape[1]`) faulted the queue. That is fixed in the overlay (LOCAL_FIXES item 30). + +## Environment + +AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`, RDNA 3.5), 96 GiB VRAM carve-out, Debian 13, ROCm 10.0.0. Before: `main` at `57d8ed29`. After: `85e39880` (the PR's second commit; later commits change tests and docs only). MLX pin `81ba1c6a`, ROCm overlay `75915908`, `cargo build --release --features rocm`. Checkpoint `mlx-community/Meta-Llama-3.1-8B-Instruct-4bit` (vocab 128256). + +## Decode throughput + +`scripts/bench_decode.sh` at its default pp512/tg128 shape with the `--temperature` / `--top-p` options this change adds, three runs per arm, before and after interleaved run by run (before, after, before, after). Every run went through `scripts/rocm_gpu_guard.sh` (the copy fixed in lablup/mlxcel#2098, run from the #2065 worktree) with `--idle-secs 75`: 75 consecutive 1 Hz samples with `/sys/class/kfd/kfd/proc` empty and no compiler, which took 112 to 116 s of wall time per run, then the run under the monitor. A parallel unit benchmarked with the guard's default 90 at the same time; at equal windows both guards started together and rejected each other, so the shorter sample count was used to break the tie. All 12 accepted runs were clean on their first attempt. Raw rows: `benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-{before,after}_{t0.8,t0.8p0.95}.csv`. + +| Configuration | Path after | Before tok/s | After tok/s | Medians | +|---|---|---|---|---| +| `--temperature 0.8` | Gumbel-max kernel | 37.91 / 37.67 / 38.32, median 37.91 | 37.59 / 38.26 / 39.25, median 38.26 | 1.01x, inside the spread | +| `--temperature 0.8 --top-p 0.95` | rejection kernel | 36.26 / 33.98 / 32.81, median 33.98 | 37.82 / 37.56 / 37.06, median 37.56 | 1.11x | + +The Gumbel-max row is not a measurable change: the arms overlap and the medians differ by 0.9%, against a spread of 1.7% before and 4.4% after. The top-p row is: every after run is faster than every before run. The before arm also drifted down across its three runs (36.3 to 32.8), which the after arm did not. The cause was not isolated; the chain runs an `argsort` over all 128256 entries per token for top-p, which is what the rejection kernel replaces. Greedy decode reaches neither kernel and was not rerun. + +The issue's second command (`--top-k 40 --top-p 0.95`) does not reach the rejection kernel at this vocabulary: the routing policy sends top-k with top-p to the kernel only up to vocab 32768 (`REJECTION_JOINT_VOCAB_MAX`, measured on M1 Ultra), so that configuration runs the stock chain before and after. Top-p alone is the routed configuration and is the one measured. Whether the joint cap is right on ROCm was not measured. + +## Correctness + +Sampling is distributional, so the ports are held to the graph in two ways. + +**Fixed key.** Both kernels draw one Philox key per call from MLX's default key sequence, and given the key their output is deterministic. `sampling_fixed_key_tests` (`src/lib/mlxcel-core/src/sampling_fixed_key_tests.rs`) reseeds, reads the key the next launch will take through the new `random_bits` bridge function, reseeds again, launches, and recomputes the draw on the host from that key: + +- Gumbel-max, 48 rows of 5003 logits in f32 (T 1.0 and 0.7), bf16 (T 1.3) and f16 (T 0.9): the kernel's index equals the host's `argmax(logits / T + g)` on every decided row, and so does the MLX graph's `argmax` with the same noise pushed through it. A row counts as decided when its winning score leads by at least 1e-4. +- Rejection under min-p 0.05, 40 rows of 3001 entries, T 1.0 and 0.7: min-p resolves its whole threshold before the first draw, so round 0 accepts and the token follows from the round-0 Philox word and the kernel's thread-major scan order; the kernel matches the host on every decided row (target at least 1e-5 of the mass from a cumulative boundary). + +With the Philox counter (`row + 1`) or the drawn word (`c1` for `c0`) changed in the HIP sources, both tests fail on row 0. + +**Statistical.** Two-sample chi-square tests (pooled to 20 draws per bin, critical value at p = 1e-6) between each kernel and `fused_sample_categorical`, the explicit graph arm, on the same input: 400,000 draws per arm for the Gumbel-max kernel against `random::categorical` at T 0.8, and for the rejection kernel against the stock chain at top-p 0.9, and at top-p 0.95 with min-p 0.02. Both pass. The existing suites, which test each kernel against the exact softmax or truncated distribution, now run on ROCm instead of returning early: `sampling_gumbel_tests` and `sampling_rejection_tests` pass on gfx1151. + +**Speculative decoding.** The classic `SpeculativeGenerator` path (`mlxcel generate --draft-model`) verifies drafted tokens with the target's sampler, and `fused_sample_probs` follows the same routing, so with top-p the target's draws and reported distribution now come from the rejection kernel. Qwen3-30B-A3B-4bit with Qwen3-0.6B-4bit as the drafter, `--temp 0.8 --top-p 0.95 -n 128`, `MLXCEL_SPECULATIVE_ACCEPT_DIAG=1`: the dispatch log names the rejection kernel after and the stock chain before, the output is coherent, and the measured per-position acceptance sits near the closed form for its rule in both arms (after: sampler-match 0.643 against `sum_prod` 0.665, stochastic 0.732 against `sum_min` 0.676; before: 0.709 against 0.678 and 0.673 against 0.699, about 110 tested positions each). These are single runs and show the path works, not a distribution; the tests above carry that. No MTP checkpoint was available locally, so the MTP and DFlash round loops were not run; `docs/environment-variables.md` (`MLXCEL_SPECULATIVE_STOCHASTIC_ACCEPT`) records that they select the target token by argmax. + +## Reproducing + +```bash +cargo build --release --features rocm --bin mlxcel --bin mlxcel-bench-decode +cargo test --release --features rocm -p mlxcel-core --lib sampling_ -- --test-threads=1 +scripts/rocm_gpu_guard.sh -- scripts/bench_decode.sh models/mlx/Meta-Llama-3.1-8B-Instruct-4bit --temperature 0.8 --top-p 0.95 +``` diff --git a/docs/benchmarks.md b/docs/benchmarks.md index a237b8f4a..4bf2b12f9 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -209,6 +209,8 @@ MLXLM_PYTHON=/bin/python LD_LIBRARY_PATH=/opt/rocm/lib \ `bench_mlxlm.py`'s `--big-cooldown` replaces the normal cooldown after a big model where `bench_decode.sh` adds it, so 60 there matches 30 + 30 here. +`bench_decode.sh --temperature T [--top-p P]` measures sampled decode, which is where the fused samplers run (greedy never dispatches them); the CSV has no sampling column, so the auto-generated filename gains `_t` and `_p

`. + `scripts/rocm_gpu_guard.sh -- ` is that check as a script: it waits for 90 s with `/sys/class/kfd/kfd/proc` empty and no compiler process, runs the command under a 1 Hz monitor, reruns it if any sample shows another GPU process diff --git a/docs/installation.md b/docs/installation.md index 3b9eabc47..554a74826 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -552,7 +552,8 @@ so `mlxcel inspect` and `--estimate-memory` give the same answers as before. | mxfp8 and mxfp4 checkpoints | Run natively, including MoE experts through `gather_qmm` (for example gpt-oss-20b-MXFP4-Q4). The load log says so once per mode (`Quantization mode mxfp4: running on native ROCm kernels, no load-time conversion ...`). mxfp4 `quantize`, `quantized_matmul` and `gather_qmm` are checked against CPU references by `tests/rocm_mxfp4_quant.rs` (lablup/mlxcel#1808). Vendor FP8 block checkpoints (`quant_method: fp8`, 128x128 blocks) are requantized to mxfp8 at load and run natively; verified end to end on a dense 0.8B Qwen3.5 FP8 checkpoint, with mxfp8 `gather_qmm` and `quantized_matmul` checked by `models::switch_layers::mxfp_tests`. The official Qwen3.5 FP8 MoE releases have not been run on ROCm (lablup/mlxcel#1807). | | NVFP4 checkpoints | No native kernel. ModelOpt NVFP4 checkpoints (for example the Gemma 4 NVFP4 exports) are converted to affine 4-bit at load with no environment variables, and the load log names the route and the reason. A layer that cannot be converted fails the load with its name and the reason. MLX-native NVFP4 exports (`"mode": "nvfp4"` in `config.json`, such as `mlx-community/*-nvfp4`) have no load-time conversion and are refused at load; use an affine export instead (lablup/mlxcel#1806). | | Affine MoE models (for example Qwen3-30B-A3B) | Run natively. Single-token decode takes the fused two-kernel MoE path (gate-up, then down) through its HIP port (lablup/mlxcel#2065), as on Metal and CUDA: Qwen3-30B-A3B-4bit decode on `gfx1151` went from 61.3 to 62.5 tok/s, with logits matching Metal's fused kernel at every decided position ([results](benchmark_results/rocm-fused-moe-gfx1151-2026-10-05.md)); `MLXCEL_FUSED_MOE=0` selects the `gather_qmm` graph path instead. Experts wider than `MLXCEL_FUSED_MOE_MAX_DFF` (8192 by default, as on CUDA; Mixtral-8x7B's 14336) stay on `gather_qmm`. | -| mlxcel's fused kernels (sampling, fused norm, RoPE + KV append, paged attention) | Run as MLX graph fallbacks on the paths that have one (lablup/mlxcel#1803); ROCm ports are lablup/mlxcel#1814. | +| mlxcel's other fused kernels (fused norm, RoPE + KV append, paged attention) | Run as MLX graph fallbacks on the paths that have one (lablup/mlxcel#1803); ROCm ports are lablup/mlxcel#1814. | +| Sampled decode (`temperature > 0`) | The fused Gumbel-max sampler (no filter) and the dual-pivot rejection sampler (top-p, and the other filter combinations its routing admits) run as HIP kernels, as on Metal and CUDA (lablup/mlxcel#2064). On `gfx1151` with Meta-Llama-3.1-8B-Instruct-4bit, decode at `--temperature 0.8 --top-p 0.95` went from 33.98 to 37.56 tok/s (medians of three); at `--temperature 0.8` alone the change is inside the noise ([results](benchmark_results/rocm-samplers-gfx1151-2026-10-05.md)). `MLXCEL_SAMPLING_GUMBEL=0` and `MLXCEL_SAMPLING_REJECTION=0` force the graph. | | Hybrid SSM decode (granite-4.0-h, falcon-h1, plamo-2, Nemotron-H) | The single-token Mamba2 SSM update runs the fused HIP kernel instead of the ~55-op SSD graph (lablup/mlxcel#2067): decode on `gfx1151` went from 60.4 to 88.4 tok/s on granite-4.0-h-tiny-4bit and from 51.4 to 74.4 tok/s on Nemotron-3-Nano-30B-A3B-4bit, with model logits matching the graph and Metal's kernel at every decided position ([results](benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md)). Measured on granite and Nemotron-H; falcon-h1 and plamo-2 take the same gate but have not been run on ROCm. `MLXCEL_SSM_KERNEL=0` forces the graph. | | GPU faults | Reported as errors (lablup/mlxcel#1804). A launch HIP rejects (an oversized block, no code object for the device) fails the evaluation that issued it and the device stays usable. An asynchronous fault (an out-of-bounds access) fails the evaluation waiting on it within about a second instead of hanging or returning NaN; the server fails that request and answers later ones with the same error, because after a queue fault the HIP runtime rejects every call for the rest of the process, so the process has to restart, and its shutdown may need a SIGKILL (HIP's teardown waits on callbacks the faulted queue never runs). A kernel that never finishes is not detected; `MLX_ROCM_GPU_WATCHDOG_SECS` (default off) fails any single host wait that outlives it. | | Memory estimation on UMA hosts | Correct. Measured on the tested configuration: the ROCm allocator reports a nonzero cap (76.80 GiB of the 96 GiB carve-out), which the estimator reads before it would ever reach host RAM, so nothing that fits the carve-out is refused for that reason (lablup/mlxcel#1805). | @@ -881,6 +882,7 @@ check that the user is in the `video` and `render` groups so `rocminfo` can open **ROCm: `[cuda_kernel] No CUDA back-end.` followed by an abort**: a fused path without a ROCm fallback was reached. lablup/mlxcel#1803 routed the fused paths that had a graph fallback, MoE included, so this should no longer happen on a -routed path; report the model and the traceback. The known remaining case is a -direct call to the Gumbel-max or rejection sampler entry point, which aborts -with `[metal_kernel] No Metal back-end.` instead (lablup/mlxcel#1885). +routed path; report the model and the traceback. A direct call to a fused +kernel entry point that has no ROCm port returns an error naming the kernel and +its fallback rather than aborting (lablup/mlxcel#1885); the Gumbel-max and +rejection samplers have HIP ports (lablup/mlxcel#2064). diff --git a/scripts/bench_decode.sh b/scripts/bench_decode.sh index 6c77b8bce..150f52b32 100755 --- a/scripts/bench_decode.sh +++ b/scripts/bench_decode.sh @@ -155,6 +155,12 @@ NO_CHAT_TEMPLATE=0 NO_DEDUP=0 OUTPUT="" SUFFIX="" +# Sampling for both passes (issue #2064). Empty means the runner's default, +# greedy argmax, which never dispatches a sampler kernel; a sampled run is how +# the Gumbel-max and rejection kernels are measured. The CSV schema does not +# record either, so a non-greedy run also tags the auto-generated filename. +TEMPERATURE="" +TOP_P="" DATE=$(date '+%Y-%m-%d') # Version recorded in the CSV `mlxcel_version` column. This is the mlxcel # version from Cargo.toml (the /update-benchmarks staleness check compares @@ -850,6 +856,13 @@ Options: --jit-preheat-timeout N Timeout for CUDA JIT warmup in seconds (default: 600) --output PATH Write CSV to specific file (overrides auto-naming) --suffix TAG Append suffix to auto-generated filename (e.g. --suffix baseline) + --temperature T Sampling temperature for both passes (default: the + runner's 0, greedy). A positive value reaches the + sampler kernels; the auto-generated filename gains + `_t` because the CSV has no sampling column. + --top-p P Nucleus threshold, effective with a positive + --temperature (default: the runner's 1.0, off). Tags + the filename with `_p

`. --cooldown N Sleep N seconds after every model to let the GPU cool down (default: 0). Use on thermally constrained hardware such as the MacBook Pro M5 Max where back-to- @@ -941,6 +954,8 @@ default_output_path() { name="${name}_vlm" fi name="${name}_${DATE}" + [[ -n "$TEMPERATURE" ]] && name="${name}_t${TEMPERATURE}" + [[ -n "$TOP_P" ]] && name="${name}_p${TOP_P}" if [[ -n "$SUFFIX" ]]; then name="${name}_${SUFFIX}" fi @@ -1055,6 +1070,8 @@ bench_one() { if [[ "$NO_CHAT_TEMPLATE" -eq 1 ]]; then extra_args+=(--no-chat-template) fi + [[ -n "$TEMPERATURE" ]] && extra_args+=(--temperature "$TEMPERATURE") + [[ -n "$TOP_P" ]] && extra_args+=(--top-p "$TOP_P") >&2 printf '>>> [bench] %s (same-process warmup=%s) ...\n' "$model_name" "$WARMUP_TOKENS" local raw rc=0 @@ -1142,6 +1159,8 @@ while [[ $# -gt 0 ]]; do --jit-preheat-timeout) JIT_PREHEAT_TIMEOUT="$2"; shift 2 ;; --output) OUTPUT="$2"; shift 2 ;; --suffix) SUFFIX="$2"; shift 2 ;; + --temperature) TEMPERATURE="$2"; shift 2 ;; + --top-p) TOP_P="$2"; shift 2 ;; --cooldown) COOLDOWN_SECS="$2"; shift 2 ;; --big-cooldown) BIG_MODEL_COOLDOWN_SECS="$2"; shift 2 ;; --big-threshold-gb) From fc2556dcd2e46b96de1479279d488e4173aa5f66 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 5 Oct 2026 01:48:57 +0900 Subject: [PATCH 4/7] fix(rocm): ask the sampler port tables in the kill-switch tests With the HIP ports, `rejection_sample_supported()` is true on ROCm while `custom_kernels_available()` stays Metal-or-CUDA, so `tests/sampling_rejection_kill_switch.rs` expected the "no rejection-sampling kernel port" decline where `fused_sample` now reports the kill switch. Both kill-switch tests now ask new env-independent bridge predicates, `sampling_gumbel_backend_supported()` and `sampling_rejection_backend_supported()`, which read the kernels' own support predicates. Also from review: - The overlay's custom-kernel launch appends an input's shape, strides and ndim only for ndim > 0, the condition under which `build_kernel` declares them (LOCAL_FIXES item 30). - `sampling_fixed_key_tests` checks that rejection draws under top-p and top-k (multi-round rows included) converge and stay inside the support computed from the kernel's own probabilities, and asserts ids are in range when building histograms. - `gpu_backend.h`, the bridge comment on `custom_kernels_available()` and the `MLXCEL_DEBUG_KERNEL_BACKEND` line no longer say ROCm has no ports; the probe-parser fixture follows the new line. - Corrected the Gumbel HIP header note (the CUDA text already casts the logits read) and the results page wording. Refs #2064 --- .../rocm-samplers-gfx1151-2026-10-05.md | 5 +- src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md | 2 +- .../mlx/backend/rocm/custom_kernel.cpp | 10 ++- src/lib/mlx-cpp/turbo/gpu_backend.cpp | 2 +- src/lib/mlx-cpp/turbo/gpu_backend.h | 17 +++--- src/lib/mlx-cpp/turbo/sampling_gumbel_hip.h | 8 +-- src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp | 22 +++++-- src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h | 2 + src/lib/mlxcel-core/src/lib.rs | 10 +++ .../src/sampling_fixed_key_tests.rs | 61 +++++++++++++++++++ tests/sampling_gumbel_kill_switch.rs | 2 +- tests/sampling_rejection_kill_switch.rs | 4 +- tests/test_bench_rocm_detection.py | 2 +- 13 files changed, 121 insertions(+), 26 deletions(-) diff --git a/docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md b/docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md index 53a4c5ca1..2fb9861ab 100644 --- a/docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md +++ b/docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md @@ -1,6 +1,6 @@ # HIP ports of the fused samplers on gfx1151 (2026-10-05) -lablup/mlxcel#2064, part of #1814. Before this change both fused samplers took the MLX graph on ROCm: `gumbel_ports()` and `rejection_ports()` had no `.rocm` entry, and `rejection_sample_supported()` answered with `custom_kernels_available()`, which is Metal-or-CUDA by definition, so a filled slot would have stayed unreachable. Sampled decode therefore ran `random::categorical` on the no-filter path and the `argpartition` / `argsort` / `cumsum` chain on the filtered path. The decode profile ([rocm-decode-profile-gfx1151-2026-09-30.md](rocm-decode-profile-gfx1151-2026-09-30.md)) put the whole sampler tail at 0 of greedy decode GPU time and 0.4 to 3.8% of sampled decode GPU time. +lablup/mlxcel#2064, part of #1814. Before this change both fused samplers took the MLX graph on ROCm: `gumbel_ports()` and `rejection_ports()` had no `.rocm` entry, and `rejection_sample_supported()` answered with `custom_kernels_available()`, which is Metal-or-CUDA by definition, so a filled slot would have stayed unreachable. Sampled decode therefore ran `random::categorical` on the no-filter path and the `argpartition` / `argsort` / `cumsum` chain on the filtered path. The decode profile ([rocm-decode-profile-gfx1151-2026-09-30.md](rocm-decode-profile-gfx1151-2026-09-30.md)) put the whole sampler tail at 0% of greedy decode GPU time and 0.4 to 3.8% of sampled decode GPU time. The ports (`sampling_gumbel_hip.h`, `sampling_rejection_hip.h` in `src/lib/mlx-cpp/turbo/`) are the CUDA bodies with the same inputs, outputs, grid, template arguments and Philox-4x32-10 counter and key layout, so a seed reproduces a ROCm stream as it does on CUDA and Metal. Neither kernel has a lane-level operation (every reduction and the rejection kernel's scan go through shared memory with a barrier per step), so neither carries a wave32 guard. `rejection_sample_supported()` now reads `has_kernel_port(rejection_ports())`. @@ -8,7 +8,7 @@ The Gumbel-max port exposed a fault in the vendored `fast::hip_kernel`: it decla ## Environment -AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`, RDNA 3.5), 96 GiB VRAM carve-out, Debian 13, ROCm 10.0.0. Before: `main` at `57d8ed29`. After: `85e39880` (the PR's second commit; later commits change tests and docs only). MLX pin `81ba1c6a`, ROCm overlay `75915908`, `cargo build --release --features rocm`. Checkpoint `mlx-community/Meta-Llama-3.1-8B-Instruct-4bit` (vocab 128256). +AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`, RDNA 3.5), 96 GiB VRAM carve-out, Debian 13, ROCm 10.0.0. Before: `main` at `57d8ed29`. After: `85e39880` (the PR's second commit; later commits change docs, the benchmark script and tests only). MLX pin `81ba1c6a`, ROCm overlay `75915908`, `cargo build --release --features rocm`. Checkpoint `mlx-community/Meta-Llama-3.1-8B-Instruct-4bit` (vocab 128256). ## Decode throughput @@ -43,5 +43,6 @@ With the Philox counter (`row + 1`) or the drawn word (`c1` for `c0`) changed in ```bash cargo build --release --features rocm --bin mlxcel --bin mlxcel-bench-decode cargo test --release --features rocm -p mlxcel-core --lib sampling_ -- --test-threads=1 +cargo test --release --features rocm --test sampling_gumbel_kill_switch --test sampling_rejection_kill_switch scripts/rocm_gpu_guard.sh -- scripts/bench_decode.sh models/mlx/Meta-Llama-3.1-8B-Instruct-4bit --temperature 0.8 --top-p 0.95 ``` diff --git a/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md b/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md index 825692bdd..c7da29c02 100644 --- a/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md +++ b/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md @@ -39,7 +39,7 @@ The fork branched from upstream on 2026-06-26 (`39886de4`). Moving its ROCm code 27. **CPU-stream BLAS runs single-threaded over fine-grained memory.** On an APU the fork's allocator (`mlx/backend/rocm/allocator.cpp`, `unified_malloc`) gives every array fine-grained device memory, CPU-stream arrays included, and `Buffer::raw_ptr()` hands the CPU that same pointer. MLX's CPU backend writes BLAS and LAPACK results straight into it (`cblas_sgemm` in `mlx/backend/cpu/gemms/cblas.cpp`, also `conv.cpp`, `masked_mm.cpp` and the LAPACK primitives), and multithreaded OpenBLAS returned wrong output columns there, differently from call to call. Measured on the gfx1151 host (Ryzen AI MAX+ 395, Debian OpenBLAS 0.3.29 pthread, 32 threads): a standalone C `cblas_sgemm` of `[1, 2880] x [2880, 2880]^T` was wrong in 37 of 50 calls with inputs and output from `hipExtMallocWithFlags(hipDeviceMallocFinegrained)`, in 10 of 300 with only the output there, and exact in every call with malloc or `hipHostMalloc` buffers or with `OPENBLAS_NUM_THREADS` at 1, 2 or 4. Through MLX the wrong columns were each the last column of one OpenBLAS thread's share of the output (shares of 93 columns: 650, 1022, 1859, ...), the column whose 64-byte line the next thread also writes. This was lablup/mlxcel#2072: `tests/rocm_mxfp4_quant.rs` failed `qmm 2880x2880 M=1` in 3 of 8 runs here (7 of 10 in the issue), and the fault was in its CPU reference, not the GPU kernel; with the tensors dumped from a failing run, the GPU result matched an independent dequantize-and-dot of the packed bytes and the CPU `matmul` did not. `unified_malloc` now calls `openblas_set_num_threads(1)` once, on its first fine-grained allocation, which precedes any CPU-stream BLAS call writing such a buffer; the symbol is declared weak, so a build against another BLAS links and changes nothing. The cost is CPU-stream BLAS speed: the CPU reads fine-grained memory slowly, and that gemv went from about 0.2 s to about 1.4 s per call. GPU work and quantized or half-precision CPU matmuls (MLX's own kernels, no BLAS) are unaffected, so model inference under `MLXCEL_DEVICE=cpu` barely touches BLAS. Keeping BLAS multithreaded by giving it a cacheable scratch output would need a change at every BLAS call site in MLX's CPU backend and was not taken. `tests/rocm_cpu_blas_finegrained.rs` repeats that gemv on the CPU stream against an f64 host reference; with this change reverted it failed in 5 of 5 runs, and with it `quantized_matmul_matches_dequantized_reference` passed 50 runs in a row. Applies to the fork; to be proposed there (upstreaming candidate, lablup/mlxcel#1813). 28. **In-flight command batches bounded by what they allocate, and the cache limit enforced.** Measured on the gfx1151 host at pp512/tg128 (lablup/mlxcel#2062), the MLX peak (live buffers, `get_peak_memory`) was 20.60 GB for Meta-Llama-3.1-8B-Instruct-4bit, whose weights take 4.75 GB once loaded, and 23.56 GB for Qwen3-30B-A3B-4bit (17.17 GB), and the device-wide `mem_info_vram_used` rose by 21.19 GB and 24.09 GB. The excess was not buffer cache (0.4 to 1.7 GB at every phase boundary) but live transients inside the measured prefill: every buffer an operation allocates is kept by `CommandEncoder::add_temporary` until its batch's completion handler runs, the eager path (`use_hip_graphs()` is off) committed only every `MLX_MAX_OPS_PER_BUFFER` ops (2000) and never registered those commits with MLX's scheduler, so neither the scheduler's task cap nor `set_memory_limit` ever made the host wait, and the host could encode far ahead of the GPU, holding the transients of many batches at once. For the 8B, an f16 checkpoint, most of it is the f16 weight copy `QuantizedMatmul`'s dequantize-and-GEMM path allocates per matrix (with `MLX_ROCM_QMM_DEQUANT_GEMM=0` the peak was 6.78 GB, at a twentieth of the prefill speed); the MoE model is bf16, whose affine 4-bit matmuls took the fused WMMA kernel at the time and allocated no copy (item 29 moved its prefill projections to the dequantize path) (the variable left its peak at 23.56 GB), so its excess is other operations' outputs held the same way, which the bound below removes too. `MLX_MAX_OPS_PER_BUFFER=50` alone changed nothing, because more commits do not stop the host from running ahead. Now `gpu::eval` (`mlx/backend/rocm/eval.cpp`) adds what each primitive allocated during `eval_gpu` to its encoder's open batch, read from a thread-local byte counter in the allocator (`rocm::thread_allocated_bytes()`, so frees on the worker thread do not disturb it), and `CommandEncoder::maybe_commit` (`mlx/backend/rocm/device.cpp`) commits once the open batch has allocated a quarter of `MLX_ROCM_MAX_INFLIGHT_MB` (default 1024); after that commit and after the commit `gpu::finalize` makes at the end of every eval and `async_eval`, it records a HIP event, and blocks on the oldest such batch while the committed ones still exceed the budget. The wait is `hipEventSynchronize` (the device is in blocking-sync mode, so the thread sleeps), or a polled wait that gives up at `MLX_ROCM_GPU_WATCHDOG_SECS` when that is set; a faulted stream ends either with the fault, which is recorded as the stream's error. Nothing is tracked while a decode-step capture or any stream capture is in progress. The events are raw `hipEvent_t` owned by the encoder, not `HipEvent`: `HipEvent` returns its handle to a function-local static pool when destroyed, an encoder can be destroyed after that pool at process exit, and a first version built on it aborted `tests/rocm_mxfp4_quant.rs` at exit with `malloc_consolidate(): unaligned fastbin chunk detected`. `0` restores the old behavior. `tests/rocm_inflight_bound.rs` runs 40 f16 4-bit 8192x8192 matmuls over 512 rows in one evaluation (5 GiB of dequantized weight copies): the peak rose by 1.75 GiB with the default and by 6.25 GiB with `MLX_ROCM_MAX_INFLIGHT_MB=0`, which fails its 3 GiB bound. With the default the peaks were 6.14 GB and 18.57 to 18.58 GB over three harness runs each, decode throughput was unchanged (median 37.57 against 37.02 tok/s and 61.69 against 61.79 tok/s) and the 8B's prefill lost 1.8% (the MoE model's prefill varies more than that from run to run); budgets from 256 to 4096 MiB traded peak for prefill speed within a few percent (the table is in `docs/benchmark_results/rocm-memory-gfx1151-2026-09-30.md`). Separately, `RocmAllocator::set_cache_limit` stored `max_pool_size_` and nothing read it, so MLX's `set_cache_limit` (mlxcel's `MLXCEL_CACHE_LIMIT`) bounded nothing, while the exact-size cache (`min_utilization` 1.0) keeps every buffer size it has seen until `clear_cache()`. `malloc_async` now trims the cache on a cache miss once it exceeds `max_pool_size_`, the only place the footprint grows, so active plus cache stays within the live set plus the limit plus the one request; it trims to three quarters of the limit, so one round of blocking `hipFree` calls covers several misses of a workload whose shapes keep changing; `free()` and cache hits still never call `hipFree`, which the fork avoids for the training reasons in its comment there. The limit still defaults to the memory limit (76.8 GiB here), so the trim changes nothing until someone sets a limit; mlxcel sets 2 GiB on ROCm builds. `memory::tests::cache_limit_bounds_the_free_buffer_cache` in mlxcel-core frees 64 buffers of distinct sizes under a 1 MiB limit and checks the cache stays under 8 MiB; with the trim removed it held 49.8 MB and failed. The count is of allocations, not only of transients (a batch that grows the KV cache counts too), and a batch's buffers are dropped by the worker thread just after its event completes, so the bound is approximate. Applies to the fork; to be proposed there (upstreaming candidate, lablup/mlxcel#1813). 29. **bf16 GEMMs of 128 rows or more leave the fused WMMA kernel on RDNA 3.5.** `QuantizedMatmul::eval_gpu` (`mlx/backend/rocm/quantized/qmm.hip`) sent every single bf16 affine GEMM with more than one row (4-, 6- or 8-bit, group 64, `N % 16 == 0`, `K % 64 == 0`) to `qmm_wmma_dense_kernel` on a device with native WMMA that is not a low-CU iGPU, however many rows it had. On gfx1151 that kernel lost to the dequantize + hipBLASLt path it bypasses on every bf16 shape measured from 128 rows up (1.0x to 3.1x at 128 rows, 1.6x to 2.4x at 256, 2.4x to 5.1x at 1024 and 2048; it won only on some shapes at 2 to 64 rows), so bf16 prefill ran at a fraction of what the same device did for f16 checkpoints (lablup/mlxcel#2081). It also made `quantized_matmul` and `dequantize` + `matmul` disagree for bf16 (499 of 1,048,576 outputs at `[1024, 2048] x [1024, 2048]`, a reduction-order difference), which is what mlxcel's dense prefill path checks. The dispatch decision is now one function, `select_qmm_route`, used by `eval_gpu`; it keeps the kernel below `MLX_ROCM_WMMA_QMM_MAX_M` rows (default 128 on the RDNA 3.5 tier, no ceiling on the unmeasured tiers) and sends larger GEMMs to the bf16 dequantize + hipBLASLt route, except where that route would be the fp8 path (RDNA 4), which keeps the kernel as before. `MLX_ROCM_WMMA_QMM=1` removes the ceiling, which is the old dispatch on any device that is not a low-CU iGPU (there it forces the kernel on, as it always did). `rocm.h` exports `quantized_matmul_runs_dequant_gemm` (stubbed in `no_rocm.cpp`), which answers from the same function whether a single GEMM of a given shape takes dequantize + hipBLASLt, for mlxcel to decide when its dense prefill path returns the same bytes. Same binary, `MLX_ROCM_WMMA_QMM=1` against the default, prefill tok/s over three ABBA runs: Gemma 3 4B 4-bit 1146 to 2289 at 512 tokens and 961 to 2815 at 2048, Qwen3-0.6B 4-bit 4581 to 7633 and 3059 to 4236, Qwen3-30B-A3B 4-bit (MoE, only the attention projections take this path) 311 to 326 and 283 to 297; decode unchanged, and MLX peak memory within 0.1 GB except Qwen3-0.6B at 512 tokens (1.04 to 1.46 GB), since the route allocates a bf16 weight copy per matrix that item 28 bounds. The route's dequantized-weight LRU (8 matrices or 256 MB, `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE`, `MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES`) gets no hits across a model's projections in prefill but still keeps its last entries alive afterwards, as it already did for f16 checkpoints; `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE=0` turns it off. `docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md` has the tables. Applies to the fork; the ceiling was measured on gfx1151 only. -30. **Custom kernels receive `_shape` and `_strides` by value.** `build_kernel` in the fork's `mlx/backend/rocm/custom_kernel.cpp` declared an input's shape and strides as `const int32_t* _shape` and `const int64_t* _strides`, while `CustomKernel::eval_gpu` passes them with `KernelArgs::append_ndim`, which pads the vector to `JIT_MAX_NDIM` (8) entries and hands `hipModuleLaunchKernel` a pointer to that storage, so the kernel received the values themselves and read the first eight bytes of the shape as an address. Any `fast::hip_kernel` source that mentioned `_shape` faulted the queue on its first launch: the HIP Gumbel-max sampler (lablup/mlxcel#2064), which reads `logits_shape[1]` as the CUDA and Metal ports do, died with `HSA_STATUS_ERROR_MEMORY_FAULT` at address `0x80000000000`, the page of `0x80000000800`, which is the shape `{2048, 2048}` read as a 64-bit pointer, on gfx1151. No HIP kernel in mlxcel used either argument before, which is why it went unseen. The generated header now defines `KernelShape` (8 `int32_t`) and `KernelStrides` (8 `int64_t`) with an `operator[]`, the parameters are declared with them by value, which is what upstream CUDA's `const __grid_constant__ Shape` / `Strides` does, and `elem_to_loc` gains an overload for them; a `static_assert` ties the 8 to `JIT_MAX_NDIM`. `sampling_fixed_key_tests` in mlxcel-core reads `logits_shape` through the Gumbel-max port and faults without this change; the `_strides` declaration has no user in mlxcel and was checked only by compiling it. Applies to the fork; to be proposed there (upstreaming candidate, lablup/mlxcel#1813). +30. **Custom kernels receive `_shape` and `_strides` by value.** `build_kernel` in the fork's `mlx/backend/rocm/custom_kernel.cpp` declared an input's shape and strides as `const int32_t* _shape` and `const int64_t* _strides`, while `CustomKernel::eval_gpu` passes them with `KernelArgs::append_ndim`, which pads the vector to `JIT_MAX_NDIM` (8) entries and hands `hipModuleLaunchKernel` a pointer to that storage, so the kernel received the values themselves and read the first eight bytes of the shape as an address. Any `fast::hip_kernel` source that mentioned `_shape` faulted the queue on its first launch: the HIP Gumbel-max sampler (lablup/mlxcel#2064), which reads `logits_shape[1]` as the CUDA and Metal ports do, died with `HSA_STATUS_ERROR_MEMORY_FAULT` at address `0x80000000000`, the page of `0x80000000800`, which is the shape `{2048, 2048}` read as a 64-bit pointer, on gfx1151. No HIP kernel in mlxcel used either argument before, which is why it went unseen. The generated header now defines `KernelShape` (8 `int32_t`) and `KernelStrides` (8 `int64_t`) with an `operator[]`, the parameters are declared with them by value, which is what upstream CUDA's `const __grid_constant__ Shape` / `Strides` does, and `elem_to_loc` gains an overload for them; a `static_assert` ties the 8 to `JIT_MAX_NDIM`. The launch also appends an input's shape, strides and ndim only when its ndim is above zero, the condition under which `build_kernel` declares them; before, a 0-d input whose name appeared as `_shape` in the source shifted every later argument (no mlxcel kernel has such an input). `sampling_fixed_key_tests` in mlxcel-core reads `logits_shape` through the Gumbel-max port and faults without this change; the `_strides` declaration has no user in mlxcel and was checked only by compiling it. Applies to the fork; to be proposed there (upstreaming candidate, lablup/mlxcel#1813). ## Fixes to the fork's build diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/custom_kernel.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/custom_kernel.cpp index 755c05a3a..7a9d5df45 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/custom_kernel.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/custom_kernel.cpp @@ -360,13 +360,17 @@ void CustomKernel::eval_gpu( const array& in = checked_inputs[i]; auto& shape_info = shape_infos_[i]; args.append(in); - if (std::get<0>(shape_info)) { + // `build_kernel` declares these only for an input with ndim > 0; append + // them under the same condition so a scalar input cannot shift every + // later argument. + const bool has_dims = in.ndim() > 0; + if (has_dims && std::get<0>(shape_info)) { args.append_ndim(in.shape()); } - if (std::get<1>(shape_info)) { + if (has_dims && std::get<1>(shape_info)) { args.append_ndim(in.strides()); } - if (std::get<2>(shape_info)) { + if (has_dims && std::get<2>(shape_info)) { args.append(in.ndim()); } } diff --git a/src/lib/mlx-cpp/turbo/gpu_backend.cpp b/src/lib/mlx-cpp/turbo/gpu_backend.cpp index 7fdf11d7d..dc74d073c 100644 --- a/src/lib/mlx-cpp/turbo/gpu_backend.cpp +++ b/src/lib/mlx-cpp/turbo/gpu_backend.cpp @@ -77,7 +77,7 @@ GpuKernelBackend gpu_kernel_backend() { backend_name(resolved), custom_kernels_available_for(resolved) ? "" - : " (no custom kernel ports; using MLX graph fallbacks)"); + : " (not every kernel family is ported; the rest use MLX graph fallbacks)"); } return resolved; }(); diff --git a/src/lib/mlx-cpp/turbo/gpu_backend.h b/src/lib/mlx-cpp/turbo/gpu_backend.h index 5995aa322..91b03027e 100644 --- a/src/lib/mlx-cpp/turbo/gpu_backend.h +++ b/src/lib/mlx-cpp/turbo/gpu_backend.h @@ -38,16 +38,19 @@ enum class GpuKernelBackend { None = 0, Metal = 1, Cuda = 2, - // ROCm is a real GPU backend with `fast::hip_kernel` ports for some kernels - // only (issue #1814). Backend-wide checks treat it like `None`; a ported - // kernel's own `has_kernel_port` predicate opens it there. Callers must not - // call `fast::cuda_kernel`. + // ROCm is a real GPU backend whose `fast::hip_kernel` ports arrive one + // kernel at a time (issue #1814). A family-wide gate + // (`custom_kernels_available()`) treats it like `None` and takes the graph + // fallback; a kernel with a HIP port is gated on its own predicate, which + // reads its port table (`has_kernel_port`). Nothing may call + // `fast::cuda_kernel` here. Rocm = 3, }; -// True when this backend has custom kernel ports, that is Metal or CUDA. This -// is the condition the launchers gate on; it is deliberately not "a GPU is -// present", which is what the old idiom conflated it with. +// True when this backend has every custom kernel family ported, that is Metal +// or CUDA. Family-wide gates use it; a kernel with per-backend ports answers +// through `has_kernel_port` on its own table instead. It is deliberately not +// "a GPU is present", which is what the old idiom conflated it with. constexpr bool custom_kernels_available_for(GpuKernelBackend backend) { return backend == GpuKernelBackend::Metal || backend == GpuKernelBackend::Cuda; diff --git a/src/lib/mlx-cpp/turbo/sampling_gumbel_hip.h b/src/lib/mlx-cpp/turbo/sampling_gumbel_hip.h index 6c26eb937..746651e0e 100644 --- a/src/lib/mlx-cpp/turbo/sampling_gumbel_hip.h +++ b/src/lib/mlx-cpp/turbo/sampling_gumbel_hip.h @@ -36,10 +36,10 @@ // - The initial best score is `-__builtin_huge_valf()`. hipRTC compiles this // body with only the headers `fast::hip_kernel` prepends, and `INFINITY` is a // `` macro that is not guaranteed among them. -// - Every logits read is an explicit `(float)` conversion. `hip_bfloat16`, the -// type MLX substitutes for bfloat16 on ROCm, converts to float only through -// an `explicit` operator, so the implicit conversion the CUDA text relies on -// for `__nv_bfloat16` would not compile here. +// - `temp[0]` gains a redundant `(float)` (it is float32). The logits read +// keeps the CUDA text's explicit `(float)`, which is required here: +// `hip_bfloat16`, the type MLX substitutes for bfloat16 on ROCm, converts to +// float only through an `explicit` operator. // // No wave32 guard, deliberately. The guard that the shuffle-based #1814 ports // carry protects a lane fold that starts at 16 and so silently drops half of a diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index 682cc02b9..db3011180 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -5132,10 +5132,11 @@ bool gpu_backend_available() { return mlx::core::device_count(mlx::core::Device::gpu) > 0; } -// True when the resolved GPU backend has mlxcel's fused kernel ports, that is -// Metal or CUDA (issue #1803). ROCm answers false until #1814 ports them, so -// the Rust callers take the graph fallback each family already has instead of -// reaching a `fast::cuda_kernel` that throws. +// True when the resolved GPU backend has every fused kernel family ported, +// that is Metal or CUDA (issue #1803). ROCm answers false: its ports arrive one +// kernel at a time under #1814, so a gate on this predicate keeps the graph +// fallback on ROCm, and a kernel that has a HIP port is gated on its own +// predicate, which reads its port table (`has_kernel_port`). bool custom_kernels_available() { return mlxcel::custom_kernels_available(); } @@ -5314,6 +5315,13 @@ bool sampling_gumbel_available() { return env_enabled && mlxcel::turbo::gumbel_max_sample_supported(); } +// The backend half of `sampling_gumbel_available`, without the env switch: +// the default device is the GPU and its backend has a Gumbel-max port. Lets a +// test that sets `MLXCEL_SAMPLING_GUMBEL=0` still ask whether a port exists. +bool sampling_gumbel_backend_supported() { + return mlxcel::turbo::gumbel_max_sample_supported(); +} + // Env kill switch for the dual-pivot rejection sampling kernels (#901). Falsy // disables and restores the `argpartition` + `argsort` + `cumsum` filter chain // exactly. Read once per process, for the same reason as the #900 switch. @@ -5558,6 +5566,12 @@ bool sampling_rejection_available() { return env_enabled && mlxcel::turbo::rejection_sample_supported(); } +// The backend half of `sampling_rejection_available`, without the env switch. +// It decides which decline `fused_sample` reports (no port, or the switch). +bool sampling_rejection_backend_supported() { + return mlxcel::turbo::rejection_sample_supported(); +} + // Threads per threadgroup the rejection kernel launches with. int32_t sampling_rejection_threadgroup_size() { return mlxcel::turbo::rejection_threadgroup_size(); diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 2344c35fa..1b9b88ab0 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -1524,6 +1524,7 @@ std::unique_ptr gumbel_max_sample( // True when `fused_sample`'s no-filter path takes the Gumbel-max kernel: // backend support plus a non-falsy `MLXCEL_SAMPLING_GUMBEL`. bool sampling_gumbel_available(); +bool sampling_gumbel_backend_supported(); // Threadgroups the Gumbel-max kernel puts on one row for this launch shape. // Exposed so tests can pin that the sampled id does not depend on it. @@ -1572,6 +1573,7 @@ std::unique_ptr sampling_rejection_probe( // True when `fused_sample`'s filtered path takes the rejection kernel: backend // support plus a non-falsy `MLXCEL_SAMPLING_REJECTION`. bool sampling_rejection_available(); +bool sampling_rejection_backend_supported(); // Pure routing policy: would this configuration go to the rejection kernel, // ignoring backend support and the env switch? The kernel replaces a sort, so diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index 355a384eb..da46780b1 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -2502,6 +2502,11 @@ mod ffi { /// The env value is read once per process. fn sampling_gumbel_available() -> bool; + /// The backend half of [`sampling_gumbel_available`], ignoring + /// `MLXCEL_SAMPLING_GUMBEL`: true when the default device is the GPU + /// and its backend has a Gumbel-max port. + fn sampling_gumbel_backend_supported() -> bool; + /// Threadgroups the Gumbel-max kernel cooperates on one row with, for a /// `[batch, vocab]` launch. Always a power of two in `[1, 64]`. Exposed /// so tests can pin that the sampled id does not depend on it. @@ -2559,6 +2564,11 @@ mod ffi { /// and `MLXCEL_SAMPLING_REJECTION` is not falsy. Read once per process. fn sampling_rejection_available() -> bool; + /// The backend half of [`sampling_rejection_available`], ignoring + /// `MLXCEL_SAMPLING_REJECTION`: true when the default device is the GPU + /// and its backend has a rejection-sampling port. + fn sampling_rejection_backend_supported() -> bool; + /// Pure routing policy (#901): would `fused_sample` send this /// configuration to the rejection kernel, ignoring backend support and /// the env switch? The kernel replaces a sort, so it is routed only diff --git a/src/lib/mlxcel-core/src/sampling_fixed_key_tests.rs b/src/lib/mlxcel-core/src/sampling_fixed_key_tests.rs index 6ff245312..2739f6921 100644 --- a/src/lib/mlxcel-core/src/sampling_fixed_key_tests.rs +++ b/src/lib/mlxcel-core/src/sampling_fixed_key_tests.rs @@ -197,6 +197,10 @@ fn histogram_of(vocab: usize, n: usize, mut draw: impl FnMut() -> Vec) -> V if drawn >= n { break; } + assert!( + (id as usize) < vocab, + "sampler returned id {id} for vocab {vocab}" + ); counts[id as usize] += 1; drawn += 1; } @@ -392,6 +396,63 @@ fn rejection_kernel_draws_the_keyed_token_under_min_p() { } } +#[test] +fn rejection_kernel_draws_stay_inside_the_filtered_support() { + let _dispatch = crate::sampling_dispatch::dispatch_test_guard(); + if !sampling_rejection_available() { + return; + } + // Rounds past the first (the pivot and bisection logic) are not + // reproducible on the host bit for bit, because the kernel's f32 sums may + // differ in their last bits. What is deterministic whatever the round + // count: every row converges and every drawn token is inside the filter's + // support, computed here from the same probabilities the kernel reads. The + // host side is a slight superset (a relative 1e-5 slack on the top-p + // cutoff), so rounding at the boundary cannot fail a correct kernel. + let rows = 40usize; + let vocab = 3001usize; + let host = row_logits(rows, vocab, 3.0, 0x0902); + let logits = from_slice_f32(&host, &[rows as i32, vocab as i32]); + let filter = f32_values(&softmax(&logits, -1)); + let mut multi_round = 0usize; + for (top_k, top_p) in [(0i32, 0.9f32), (40, 1.0), (0, 0.5)] { + for launch in 0..4u64 { + random_seed(0x5EED_2064_0200 + launch); + let stacked = sampling_rejection_probe(&logits, 1.0, top_k, top_p, 0.0, 32); + let flat = u32_values(&stacked); + let (ids, ok, used) = (&flat[..rows], &flat[rows..2 * rows], &flat[2 * rows..]); + for row in 0..rows { + let label = format!("top_k={top_k} top_p={top_p} launch {launch} row {row}"); + assert_eq!(ok[row], 1, "{label}: did not converge"); + multi_round += usize::from(used[row] > 1); + let f = &filter[row * vocab..(row + 1) * vocab]; + let id = ids[row] as usize; + assert!(id < vocab, "{label}: id {id} out of range"); + let p = f[id]; + if top_k > 0 { + let above = f.iter().filter(|&&q| q > p).count(); + assert!( + above < top_k as usize, + "{label}: id {id} has {above} tokens above it" + ); + } + if top_p < 1.0 { + let total: f64 = f.iter().map(|&q| f64::from(q)).sum(); + let exclusive: f64 = f.iter().filter(|&&q| q > p).map(|&q| f64::from(q)).sum(); + assert!( + exclusive <= f64::from(top_p) * total * (1.0 + 1e-5), + "{label}: id {id} starts at mass {exclusive}, past top_p {top_p}" + ); + } + } + } + } + assert!( + multi_round > 0, + "no row needed a second round; the bisection went untested" + ); +} + // -- 2. statistical, against the graph sampler -- #[test] diff --git a/tests/sampling_gumbel_kill_switch.rs b/tests/sampling_gumbel_kill_switch.rs index d8ea0b914..b08be3490 100644 --- a/tests/sampling_gumbel_kill_switch.rs +++ b/tests/sampling_gumbel_kill_switch.rs @@ -105,7 +105,7 @@ fn falsy_env_restores_the_categorical_sampling_path() { } Err(error) => { assert!( - !mlxcel_core::custom_kernels_available(), + !mlxcel_core::sampling_gumbel_backend_supported(), "the launcher refused on a backend that reports a kernel port: {error}" ); let message = error.to_string(); diff --git a/tests/sampling_rejection_kill_switch.rs b/tests/sampling_rejection_kill_switch.rs index a9a5cc3c7..9e9af1bd7 100644 --- a/tests/sampling_rejection_kill_switch.rs +++ b/tests/sampling_rejection_kill_switch.rs @@ -111,7 +111,7 @@ fn falsy_env_restores_the_argpartition_filter_chain() { // that reason before the kill switch is ever consulted, so asserting the // kill-switch text there would be asserting something that cannot happen // (issue #1885). Both arms are checked rather than skipping one. - let expected = if mlxcel_core::custom_kernels_available() { + let expected = if mlxcel_core::sampling_rejection_backend_supported() { "MLXCEL_SAMPLING_REJECTION" } else { "no rejection-sampling kernel port" @@ -136,7 +136,7 @@ fn falsy_env_restores_the_argpartition_filter_chain() { } Err(error) => { assert!( - !mlxcel_core::custom_kernels_available(), + !mlxcel_core::sampling_rejection_backend_supported(), "the launcher refused on a backend that reports a kernel port: {error}" ); let message = error.to_string(); diff --git a/tests/test_bench_rocm_detection.py b/tests/test_bench_rocm_detection.py index 6e7683fef..7d46113d1 100644 --- a/tests/test_bench_rocm_detection.py +++ b/tests/test_bench_rocm_detection.py @@ -103,7 +103,7 @@ def test_kernel_backend(self) -> None: self.assertEqual(_run(["probe_kernel_backend"], "probe_kernel_backend", ROCM_PROBE_OUTPUT), "rocm") def test_kernel_backend_with_fallback_suffix(self) -> None: - line = "[mlxcel] custom kernel backend: rocm (no custom kernel ports; using MLX graph fallbacks)\n" + line = "[mlxcel] custom kernel backend: rocm (not every kernel family is ported; the rest use MLX graph fallbacks)\n" self.assertEqual(_run(["probe_kernel_backend"], "probe_kernel_backend", line), "rocm") def test_gfx_target(self) -> None: From e034e8300b6336034725afee18cfaafe18e1a951 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 5 Oct 2026 02:15:26 +0900 Subject: [PATCH 5/7] test(sampling): allow 2 ulp on the routed T=1 snapshot rows `temperature_one_support_unchanged` skipped its two routed cases wherever the rejection kernel had no port, so they first ran on ROCm with this change. Their rows are the graph's softmax over the kernel's support (`fused_sample_probs` does not launch the kernel), and gfx1151 lands 2 ulp from the Metal capture at token 26 of the (40, 0.9) row with the support unchanged, in `make verify-rocm` and in three test-fast and one release rerun. The bound goes from 1 to 2 ulp; the support check is unchanged and the stock case keeps its 4. Refs #2064 --- src/lib/mlxcel-core/src/sampling.rs | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/lib/mlxcel-core/src/sampling.rs b/src/lib/mlxcel-core/src/sampling.rs index 011afa66c..1026705f7 100644 --- a/src/lib/mlxcel-core/src/sampling.rs +++ b/src/lib/mlxcel-core/src/sampling.rs @@ -4316,7 +4316,8 @@ mod tests { /// decision boundary the sampler observes, and that has shown up /// intermittently in nightly `make verify-test` runs. The routed kernel /// cases therefore hold the exact support and keep every positive - /// probability within one f32 ULP of the saved row. The non-routed stock + /// probability within two f32 ULPs of the saved row (one until gfx1151 + /// ran them, see the assertion). The non-routed stock /// chain case uses the same support-preserving snapshot helper with a /// two-ULP ceiling because local CPU validation observed a two-ULP drift /// on one low-probability entry while preserving the same filtered @@ -4398,7 +4399,14 @@ mod tests { .collect(); let ctx = format!("T=1.0 fused_sample_probs drifted for top_k={top_k} top_p={top_p}"); if routed { - assert_probs_match_snapshot_within_ulp(&got, &expected, 1, &ctx); + // 2, not 1. These rows are the graph's softmax over the + // kernel's support (`fused_sample_probs` never launches the + // kernel), and they first ran on ROCm when the rejection + // kernel got its HIP port (#2064): gfx1151 lands 2 ulp from + // the Metal capture at token 26 of the (40, 0.9) row, with the + // support unchanged. The same reassociation drift the stock + // case below allows for. + assert_probs_match_snapshot_within_ulp(&got, &expected, 2, &ctx); } else { // 4, not 2. The bound is a drift tripwire on the fused softmax, // not the correctness gate: the support set is asserted just From c0b836f5adc098fb25ec2ca5f03bf1ac6e57f6e5 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 5 Oct 2026 02:32:17 +0900 Subject: [PATCH 6/7] docs(rocm): tighten the sampler results page and bench option input From review of #2100: the results page now says what the commits after the measured binary change and quotes `compare_bench_csv.py` (1.024x and 1.130x on the third runs) next to the medians and the per-pair ratios (1.04x, 1.11x, 1.13x); `docs/installation.md` notes the before arm's drift; the upstream README row for LOCAL_FIXES item 30 describes the whole change. `bench_decode.sh` rejects a `--temperature` or `--top-p` value that is not a plain decimal, since both go into the output filename. Two header comments are reflowed. Refs #2064 --- .../rocm-samplers-gfx1151-2026-10-05.md | 7 +++++-- docs/installation.md | 2 +- docs/mlxcelverse/upstream/README.md | 2 +- scripts/bench_decode.sh | 9 +++++++++ src/lib/mlx-cpp/turbo/sampling.h | 6 +++--- src/lib/mlx-cpp/turbo/sampling_rejection.h | 5 ++--- 6 files changed, 21 insertions(+), 10 deletions(-) diff --git a/docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md b/docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md index 2fb9861ab..a7d8c5ec3 100644 --- a/docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md +++ b/docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md @@ -8,7 +8,7 @@ The Gumbel-max port exposed a fault in the vendored `fast::hip_kernel`: it decla ## Environment -AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`, RDNA 3.5), 96 GiB VRAM carve-out, Debian 13, ROCm 10.0.0. Before: `main` at `57d8ed29`. After: `85e39880` (the PR's second commit; later commits change docs, the benchmark script and tests only). MLX pin `81ba1c6a`, ROCm overlay `75915908`, `cargo build --release --features rocm`. Checkpoint `mlx-community/Meta-Llama-3.1-8B-Instruct-4bit` (vocab 128256). +AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`, RDNA 3.5), 96 GiB VRAM carve-out, Debian 13, ROCm 10.0.0. Before: `main` at `57d8ed29`. After: `85e39880` (the PR's second commit). Later commits add the benchmark script's `--temperature` / `--top-p` options (the runs used them, with each binary copied into place), docs, tests, a debug log string, two test-only bridge predicates, and the overlay's 0-d argument gating, none of which a sampler launch with 1-d and 2-d inputs reaches. MLX pin `81ba1c6a`, ROCm overlay `75915908`, `cargo build --release --features rocm`. Checkpoint `mlx-community/Meta-Llama-3.1-8B-Instruct-4bit` (vocab 128256). ## Decode throughput @@ -19,7 +19,9 @@ AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`, RDNA 3.5), 96 GiB VRAM carve | `--temperature 0.8` | Gumbel-max kernel | 37.91 / 37.67 / 38.32, median 37.91 | 37.59 / 38.26 / 39.25, median 38.26 | 1.01x, inside the spread | | `--temperature 0.8 --top-p 0.95` | rejection kernel | 36.26 / 33.98 / 32.81, median 33.98 | 37.82 / 37.56 / 37.06, median 37.56 | 1.11x | -The Gumbel-max row is not a measurable change: the arms overlap and the medians differ by 0.9%, against a spread of 1.7% before and 4.4% after. The top-p row is: every after run is faster than every before run. The before arm also drifted down across its three runs (36.3 to 32.8), which the after arm did not. The cause was not isolated; the chain runs an `argsort` over all 128256 entries per token for top-p, which is what the rejection kernel replaces. Greedy decode reaches neither kernel and was not rerun. +`compare_bench_csv.py --before --after `, which keeps the last row per model and so compares the third runs, reports 1.024x and 1.130x. + +The Gumbel-max row is not a measurable change: the arms overlap and the medians differ by 0.9%, against a spread of 1.7% before and 4.4% after. The top-p row is: every after run is faster than every before run, by 1.04x, 1.11x and 1.13x in the three interleaved pairs. The before arm also drifted down across its three runs (36.3 to 32.8), which the after arm did not. The cause was not isolated; the chain runs an `argsort` over all 128256 entries per token for top-p, which is what the rejection kernel replaces. Greedy decode reaches neither kernel and was not rerun. The issue's second command (`--top-k 40 --top-p 0.95`) does not reach the rejection kernel at this vocabulary: the routing policy sends top-k with top-p to the kernel only up to vocab 32768 (`REJECTION_JOINT_VOCAB_MAX`, measured on M1 Ultra), so that configuration runs the stock chain before and after. Top-p alone is the routed configuration and is the one measured. Whether the joint cap is right on ROCm was not measured. @@ -44,5 +46,6 @@ With the Philox counter (`row + 1`) or the drawn word (`c1` for `c0`) changed in cargo build --release --features rocm --bin mlxcel --bin mlxcel-bench-decode cargo test --release --features rocm -p mlxcel-core --lib sampling_ -- --test-threads=1 cargo test --release --features rocm --test sampling_gumbel_kill_switch --test sampling_rejection_kill_switch +python3 scripts/compare_bench_csv.py --before benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-before_t0.8p0.95.csv --after benchmarks/rocm_strixhalo-gfx1151_2026-10-05_samplers-after_t0.8p0.95.csv scripts/rocm_gpu_guard.sh -- scripts/bench_decode.sh models/mlx/Meta-Llama-3.1-8B-Instruct-4bit --temperature 0.8 --top-p 0.95 ``` diff --git a/docs/installation.md b/docs/installation.md index 554a74826..ee682e27c 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -553,7 +553,7 @@ so `mlxcel inspect` and `--estimate-memory` give the same answers as before. | NVFP4 checkpoints | No native kernel. ModelOpt NVFP4 checkpoints (for example the Gemma 4 NVFP4 exports) are converted to affine 4-bit at load with no environment variables, and the load log names the route and the reason. A layer that cannot be converted fails the load with its name and the reason. MLX-native NVFP4 exports (`"mode": "nvfp4"` in `config.json`, such as `mlx-community/*-nvfp4`) have no load-time conversion and are refused at load; use an affine export instead (lablup/mlxcel#1806). | | Affine MoE models (for example Qwen3-30B-A3B) | Run natively. Single-token decode takes the fused two-kernel MoE path (gate-up, then down) through its HIP port (lablup/mlxcel#2065), as on Metal and CUDA: Qwen3-30B-A3B-4bit decode on `gfx1151` went from 61.3 to 62.5 tok/s, with logits matching Metal's fused kernel at every decided position ([results](benchmark_results/rocm-fused-moe-gfx1151-2026-10-05.md)); `MLXCEL_FUSED_MOE=0` selects the `gather_qmm` graph path instead. Experts wider than `MLXCEL_FUSED_MOE_MAX_DFF` (8192 by default, as on CUDA; Mixtral-8x7B's 14336) stay on `gather_qmm`. | | mlxcel's other fused kernels (fused norm, RoPE + KV append, paged attention) | Run as MLX graph fallbacks on the paths that have one (lablup/mlxcel#1803); ROCm ports are lablup/mlxcel#1814. | -| Sampled decode (`temperature > 0`) | The fused Gumbel-max sampler (no filter) and the dual-pivot rejection sampler (top-p, and the other filter combinations its routing admits) run as HIP kernels, as on Metal and CUDA (lablup/mlxcel#2064). On `gfx1151` with Meta-Llama-3.1-8B-Instruct-4bit, decode at `--temperature 0.8 --top-p 0.95` went from 33.98 to 37.56 tok/s (medians of three); at `--temperature 0.8` alone the change is inside the noise ([results](benchmark_results/rocm-samplers-gfx1151-2026-10-05.md)). `MLXCEL_SAMPLING_GUMBEL=0` and `MLXCEL_SAMPLING_REJECTION=0` force the graph. | +| Sampled decode (`temperature > 0`) | The fused Gumbel-max sampler (no filter) and the dual-pivot rejection sampler (top-p, and the other filter combinations its routing admits) run as HIP kernels, as on Metal and CUDA (lablup/mlxcel#2064). On `gfx1151` with Meta-Llama-3.1-8B-Instruct-4bit, decode at `--temperature 0.8 --top-p 0.95` went from 33.98 to 37.56 tok/s (medians of three; the before runs drifted from 36.26 down to 32.81, the after runs held 37.06 to 37.82); at `--temperature 0.8` alone the change is inside the noise ([results](benchmark_results/rocm-samplers-gfx1151-2026-10-05.md)). `MLXCEL_SAMPLING_GUMBEL=0` and `MLXCEL_SAMPLING_REJECTION=0` force the graph. | | Hybrid SSM decode (granite-4.0-h, falcon-h1, plamo-2, Nemotron-H) | The single-token Mamba2 SSM update runs the fused HIP kernel instead of the ~55-op SSD graph (lablup/mlxcel#2067): decode on `gfx1151` went from 60.4 to 88.4 tok/s on granite-4.0-h-tiny-4bit and from 51.4 to 74.4 tok/s on Nemotron-3-Nano-30B-A3B-4bit, with model logits matching the graph and Metal's kernel at every decided position ([results](benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md)). Measured on granite and Nemotron-H; falcon-h1 and plamo-2 take the same gate but have not been run on ROCm. `MLXCEL_SSM_KERNEL=0` forces the graph. | | GPU faults | Reported as errors (lablup/mlxcel#1804). A launch HIP rejects (an oversized block, no code object for the device) fails the evaluation that issued it and the device stays usable. An asynchronous fault (an out-of-bounds access) fails the evaluation waiting on it within about a second instead of hanging or returning NaN; the server fails that request and answers later ones with the same error, because after a queue fault the HIP runtime rejects every call for the rest of the process, so the process has to restart, and its shutdown may need a SIGKILL (HIP's teardown waits on callbacks the faulted queue never runs). A kernel that never finishes is not detected; `MLX_ROCM_GPU_WATCHDOG_SECS` (default off) fails any single host wait that outlives it. | | Memory estimation on UMA hosts | Correct. Measured on the tested configuration: the ROCm allocator reports a nonzero cap (76.80 GiB of the 96 GiB carve-out), which the estimator reads before it would ever reach host RAM, so nothing that fits the carve-out is refused for that reason (lablup/mlxcel#1805). | diff --git a/docs/mlxcelverse/upstream/README.md b/docs/mlxcelverse/upstream/README.md index 9578fc2cc..40e564fe2 100644 --- a/docs/mlxcelverse/upstream/README.md +++ b/docs/mlxcelverse/upstream/README.md @@ -35,7 +35,7 @@ Each directory holds the patch (`git format-patch` output, authored by the maint | 9 (expert-batched gather qmv opt-in) | Held. It disables a path that returns wrong bf16 results on gfx1151 but has no root cause yet; LOCAL_FIXES.md says to propose it only after one. | | 15 (`SearchSorted`) | Not applicable yet: the primitive arrived in ml-explore/mlx#4035 (`5ec30acd`, 2026-08-10), after the fork's MLX base, so the fork has no `SearchSorted` to implement. Package it once the fork merges MLX at or past #4035. | | 29 (bf16 GEMMs leave the WMMA kernel from 128 rows on RDNA 3.5) | Not packaged yet. The row ceiling was measured on gfx1151 only, and the exported `quantized_matmul_runs_dequant_gemm` exists for mlxcel's dense prefill check; a fork PR would carry the ceiling alone, with measurements from at least one more RDNA 3 or RDNA 3.5 device. | -| 30 (custom-kernel `_shape` / `_strides` by value) | Not packaged yet (added 2026-10-04 with lablup/mlxcel#2064). It is a two-line declaration change plus a header struct in `custom_kernel.cpp`; a package needs a repro that launches a `hip_kernel` reading `_shape` on the unpatched fork, which faults the queue. | +| 30 (custom-kernel `_shape` / `_strides` by value) | Not packaged yet (added 2026-10-04 with lablup/mlxcel#2064). It is confined to `custom_kernel.cpp`: the two declarations, header structs with an `elem_to_loc` overload, and the matching launch condition; a package needs a repro that launches a `hip_kernel` reading `_shape` on the unpatched fork, which faults the queue. | ## What was verified, and what the submitter still has to do diff --git a/scripts/bench_decode.sh b/scripts/bench_decode.sh index 150f52b32..1696674f9 100755 --- a/scripts/bench_decode.sh +++ b/scripts/bench_decode.sh @@ -1177,6 +1177,15 @@ while [[ $# -gt 0 ]]; do esac done +# Both values end up in the auto-generated filename, so only plain decimals. +for _sampling_opt in "--temperature:$TEMPERATURE" "--top-p:$TOP_P"; do + _value="${_sampling_opt#*:}" + if [[ -n "$_value" && ! "$_value" =~ ^[0-9]+(\.[0-9]+)?$ ]]; then + echo "Error: ${_sampling_opt%%:*} takes a non-negative decimal, got '$_value'" >&2 + exit 1 + fi +done + if [[ -z "$MODEL_ARG" ]]; then echo "Error: model path or 'all' required" >&2 usage >&2 diff --git a/src/lib/mlx-cpp/turbo/sampling.h b/src/lib/mlx-cpp/turbo/sampling.h index fc88bdc32..c333f5947 100644 --- a/src/lib/mlx-cpp/turbo/sampling.h +++ b/src/lib/mlx-cpp/turbo/sampling.h @@ -45,9 +45,9 @@ namespace mlxcel::turbo { // True when the Gumbel-max kernel can run on the active backend: a GPU default // device whose backend has a port in `gumbel_ports()` (Metal, CUDA and, since // #2064, ROCm). The predicate reads the same table as the dispatch, so the two -// cannot disagree. A CPU-only build, or a -// process running with `MLXCEL_DEVICE=cpu`, returns false and the caller must -// keep the `categorical` graph path. +// cannot disagree. A CPU-only build, or a process running with +// `MLXCEL_DEVICE=cpu`, returns false and the caller must keep the +// `categorical` graph path. bool gumbel_max_sample_supported(); // True when `logits` has a shape and dtype the kernel accepts: 2-D `[B, V]` diff --git a/src/lib/mlx-cpp/turbo/sampling_rejection.h b/src/lib/mlx-cpp/turbo/sampling_rejection.h index 0c79abb6a..478fa9c39 100644 --- a/src/lib/mlx-cpp/turbo/sampling_rejection.h +++ b/src/lib/mlx-cpp/turbo/sampling_rejection.h @@ -151,9 +151,8 @@ inline constexpr int REJECTION_MAX_ROUNDS = 32; // True when the rejection kernel can run on the active backend: a GPU default // device whose backend has a port in `rejection_ports()` (Metal, CUDA and, since // #2064, ROCm). The predicate reads the same table as the dispatch, so the two -// cannot disagree. A CPU-only build, or a -// process under `MLXCEL_DEVICE=cpu`, returns false and the caller keeps the -// `argpartition` chain. +// cannot disagree. A CPU-only build, or a process under `MLXCEL_DEVICE=cpu`, +// returns false and the caller keeps the `argpartition` chain. bool rejection_sample_supported(); // True when `probs` has a shape the kernel accepts: 2-D `[B, V]` with `B > 0` From e2be53a0330b0bae6656b5b5f78932fa38d9559f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 5 Oct 2026 03:02:48 +0900 Subject: [PATCH 7/7] docs: add technical report for PR #2100 Bilingual (en/ko) report on the HIP ports of the Gumbel-max and rejection samplers: the line-for-line ports, the rejection predicate reading the port table, the fast::hip_kernel shape/strides overlay fix (LOCAL_FIXES item 30), the fixed-key, support and chi-square evidence, the gfx1151 measurements, and the 2 ulp snapshot bound. --- ...2100-rocm-samplers-hip-port-20261005.en.md | 191 ++++++++++++++++++ ...2100-rocm-samplers-hip-port-20261005.ko.md | 191 ++++++++++++++++++ 2 files changed, 382 insertions(+) create mode 100644 TECHNICAL_REPORTS/2100-rocm-samplers-hip-port-20261005.en.md create mode 100644 TECHNICAL_REPORTS/2100-rocm-samplers-hip-port-20261005.ko.md diff --git a/TECHNICAL_REPORTS/2100-rocm-samplers-hip-port-20261005.en.md b/TECHNICAL_REPORTS/2100-rocm-samplers-hip-port-20261005.en.md new file mode 100644 index 000000000..247444190 --- /dev/null +++ b/TECHNICAL_REPORTS/2100-rocm-samplers-hip-port-20261005.en.md @@ -0,0 +1,191 @@ +# Technical Report: PR #2100 - Port the Gumbel-max and rejection samplers to HIP + +**Date**: 2026-10-05 + +**Status**: Implemented and measured on the gfx1151 host; head `74624ea7` on origin/main `57d8ed29`, pending merge. + +**Languages**: C++ (HIP kernel sources, kernel holders, support predicate, bridge functions, vendored ROCm overlay), Rust (fixed-key and two-sample tests, kill-switch tests, snapshot bound), Bash (`bench_decode.sh` options, dtype-key checker test), Markdown (environment variables, installation, benchmarks, LOCAL_FIXES, benchmark results), CSV (bench rows) + +**Risk Level**: Medium (it changes the default sampled-decode path on ROCm for every model, patches the vendored `fast::hip_kernel` argument layout that every HIP custom kernel goes through, and touches the rejection predicate, bridge functions and a snapshot bound that Metal and CUDA share, neither of which could be run on the development host) + +## Executive Summary + +Issue #2064 (part of #1814, epic #1801) asked for HIP ports of the two fused samplers that run once per generated token whenever temperature is above 0: the Gumbel-max sampler on the no-filter path and the dual-pivot rejection sampler on the filtered path. Before this PR both took the MLX graph on ROCm. `gumbel_ports()` and `rejection_ports()` had no `.rocm` entry, and `rejection_sample_supported()` returned `custom_kernels_available()`, which is Metal-or-CUDA by definition, so a filled rejection slot would have stayed unreachable. + +The PR adds `sampling_gumbel_hip.h` and `sampling_rejection_hip.h`, line-for-line ports of the CUDA bodies with the same inputs, outputs, grid, template arguments and Philox-4x32-10 counter and key layout, fills both `.rocm` slots, and makes `rejection_sample_supported()` return `has_kernel_port(rejection_ports())`. Neither port carries a wave32 `#error` guard, because neither kernel has a lane-level operation: every reduction and the rejection kernel's scan go through shared memory with a barrier per step. + +The Gumbel port exposed a bug in the vendored ROCm overlay: `fast::hip_kernel` declared `_shape` and `_strides` as pointers while the launch passed them by value, so the first kernel that read one (`logits_shape[1]`) faulted the queue. That is fixed as LOCAL_FIXES item 30, a fork-side change and an upstreaming candidate for #1813. + +On gfx1151 with Llama-3.1-8B-Instruct-4bit at pp512/tg128, top-p 0.95 decode went from a median of 33.98 to 37.56 tok/s (1.11x), with a downward drift in the before arm that the report states. The Gumbel path moved from 37.91 to 38.26 tok/s (+0.9%), inside the run-to-run spread, and is not claimed as a gain. The issue's top-k plus top-p command was not measured because at vocab 128256 it never reaches the kernel. + +## 1. Problem Statement + +### 1.1 Two empty slots, one predicate that could not see them + +`gumbel_max_sample_supported()` already read `has_kernel_port(gumbel_ports())`, so filling `.rocm` there was enough to make it reachable. `rejection_sample_supported()` did not: + +```cpp +return mlxcel::custom_kernels_available(); +``` + +`custom_kernels_available_for` answers true only for Metal and CUDA, so on ROCm the predicate stayed false whatever `rejection_ports()` held. It feeds `sampling_rejection_available()` and the routing in the bridge. Checker rule 4 of `verify-kernel-port-dispatch` did not catch it because it is not spelled `metal_is_available() || cuda_is_available()`. + +### 1.2 What the profile said it was worth + +The decode profile (`rocm-decode-profile-gfx1151-2026-09-30.md`) put the whole sampler tail at 0% of greedy decode GPU time and 0.4 to 3.8% of sampled decode GPU time. Without the ports, sampled decode on ROCm ran `random::categorical` on the no-filter path and the `argpartition` / `argsort` / `cumsum` chain on the filtered path. The expected gain was small on the no-filter path and larger only where the graph sorts the full vocabulary. + +### 1.3 No test held a port to the graph on the same randomness + +The existing suites, `sampling_gumbel_tests` and `sampling_rejection_tests`, test each kernel against its exact target distribution. They returned early on ROCm. Nothing checked that a port consumes MLX's key sequence the same way as the other ports, which is what makes `mlx::core::random::seed(...)` reproduce a stream across backends. + +## 2. Change Summary + +| Area | Change | +|---|---| +| `sampling_gumbel_hip.h` (new) | `GUMBEL_MAX_SAMPLE_HIP_SOURCE`, the CUDA body with spelling changes only | +| `sampling_rejection_hip.h` (new) | `REJECTION_SAMPLE_HIP_SOURCE`, the CUDA body with spelling changes only | +| `sampling.cpp` | `GumbelKernelHolderHip` calling `fast::hip_kernel` under `MLXCEL_BRIDGE_ROCM_BACKEND`; `.rocm` getter in `gumbel_ports()` | +| `sampling_rejection.cpp` | `RejectionKernelHolderHip`; `.rocm` getter in `rejection_ports()`; `rejection_sample_supported()` returns `has_kernel_port(rejection_ports())` | +| ROCm overlay `custom_kernel.cpp` | `KernelShape` / `KernelStrides` by-value structs, an `elem_to_loc` overload, a `static_assert` on `JIT_MAX_NDIM`, and shape/strides/ndim appended only for ndim > 0 (LOCAL_FIXES item 30) | +| Bridge (`mlx_cxx_bridge.cpp/.h`, `lib.rs`) | `sampling_gumbel_backend_supported()`, `sampling_rejection_backend_supported()` (env-independent), `random_bits()`; reworded comments | +| `gpu_backend.h/.cpp` | Comments on what `custom_kernels_available()` now means; the `MLXCEL_DEBUG_KERNEL_BACKEND` suffix reworded | +| Tests | New `sampling_fixed_key_tests.rs` (five tests); kill-switch tests ask the new predicates; `temperature_one_support_unchanged` routed bound from 1 to 2 ulp; dtype-key checker test moves both launches; Python probe test follows the new suffix | +| `bench_decode.sh` | `--temperature` and `--top-p`, plain decimals only, tagging the auto-generated filename | +| Docs and data | `environment-variables.md`, `installation.md`, `benchmarks.md`, upstream README row for item 30, `LOCAL_FIXES.md` item 30, new `rocm-samplers-gfx1151-2026-10-05.md`, four bench CSVs | + +The branch has six commits: the overlay fix (`cc93789f`), the ports and predicate (`85e39880`), the results doc (`ce4105b7`), the kill-switch predicate fix (`f7bec56d`), the snapshot bound (`6e336ec8`), and doc and `bench_decode.sh` input tightening (`74624ea7`). The diff is 32 files, 1,341 insertions and 46 deletions. + +## 3. The Ports + +### 3.1 Line for line, with three kinds of spelling change + +Both HIP sources keep the CUDA kernels' thread mapping, grid, inputs (`logits`, `rng_key`, `temp`; `probs`, `probs_draw`, `params`, `rng_key`), outputs (`vals`, `idxs`; `ids`, `ok`, `rounds`) and template arguments, dtype keys included. The Philox-4x32-10 counter (`{base/4, 0, row, 0}` for Gumbel, `{round, 0, row, 0}` for rejection), the key `{rng_key[0], rng_key[1]}` and the 23-bit open-interval uniform are unchanged, so a seed reproduces a ROCm stream the way it does on CUDA and Metal. `__umulhi` has the same meaning in HIP. The differences are: + +- explicit `(float)` reads of `hip_bfloat16` values, because MLX's bfloat16 substitute on ROCm converts to float only through an `explicit` operator; +- `-__builtin_huge_valf()` in place of `-INFINITY`, because hipRTC compiles the body with only the headers `fast::hip_kernel` prepends and `` is not guaranteed among them; +- a redundant `(float)` on `temp[0]`. + +The holders and launches stay in `sampling.cpp` and `sampling_rejection.cpp`, next to the CUDA ones. The headers hold data, not a launch, so `make verify-kernel-dtype-keys` keeps `EXPECTED_IN_SCOPE` and its count unchanged. The checker's own test, which moves the sampler's launch to a helper and comments it out, now does the same to the HIP launch. + +### 3.2 Why no wave32 guard + +The #1814 port requirements tell every port that reduces or shuffles across lanes to carry two `#error` checks on the wavefront-size macros, because a lane fold that starts at 16 drops half of a 64-lane CDNA wave and still returns a finite, wrong result. These two kernels have no lane-level operation. The Gumbel kernel's index-carrying halving reduction and the rejection kernel's reductions and Hillis-Steele scan all go through `__shared__` memory with a `__syncthreads()` between steps, so the result does not depend on how many threads share a wave. An `#error` on wave64 would reject a correct kernel. PR #2099 also found that the guard idiom is inert with ROCm 10's AMD clang, which defines neither macro, so it would catch nothing here in any case. The issue's acceptance criterion was ticked with that reason written next to it. + +## 4. The Predicate and the Tests That Read It + +`rejection_sample_supported()` keeps its `default_device() != Device::gpu` check and then returns `mlxcel::has_kernel_port(rejection_ports())`, the same table `select_kernel_port` reads in `rejection_sample`, which mirrors `gumbel_max_sample_supported()`. On Metal and CUDA the answer is true as before. `fused_sample`, `fused_sample_probs` and the speculative paths reach the HIP kernels with no caller change. + +Two integration tests depended on the old meaning. `tests/sampling_gumbel_kill_switch.rs` and `tests/sampling_rejection_kill_switch.rs` set the env kill switch and choose which decline message to expect with `custom_kernels_available()`. That function still answers Metal-or-CUDA, so on ROCm, where a port now exists, the tests would have expected "no port". Under the env switch, `sampling_*_available()` is false whatever the backend, so the tests needed a predicate without the switch. The PR adds `sampling_gumbel_backend_supported()` and `sampling_rejection_backend_supported()` to the bridge for that purpose (`f7bec56d`). The `gpu_backend.h` comments now say that `custom_kernels_available()` is a family-wide gate and that a kernel with a HIP port is gated on its own predicate. The `MLXCEL_DEBUG_KERNEL_BACKEND` suffix changes from "no custom kernel ports" to "not every kernel family is ported". + +## 5. The Overlay Bug: Shapes and Strides by Value + +`build_kernel` in the fork's `mlx/backend/rocm/custom_kernel.cpp` declared an input's shape as `const int32_t* _shape` and its strides as `const int64_t* _strides`. `CustomKernel::eval_gpu` passes them with `KernelArgs::append_ndim`, which pads the vector to `JIT_MAX_NDIM` (8) entries and hands `hipModuleLaunchKernel` a pointer to that storage. The kernel therefore received the values and read the first eight bytes of the shape as an address. The HIP Gumbel port reads `logits_shape[1]`, as the CUDA and Metal ports do, and died with `HSA_STATUS_ERROR_MEMORY_FAULT` at `0x80000000000`, the page of `0x80000000800`, which is the shape `{2048, 2048}` read as a 64-bit pointer. No HIP kernel in mlxcel had read either argument before, which is why it went unseen. + +The fix (`cc93789f`) adds `KernelShape` (8 `int32_t`) and `KernelStrides` (8 `int64_t`) with an `operator[]` to the generated header, declares the parameters with them by value, which is what upstream CUDA's `const __grid_constant__ Shape` / `Strides` does, and adds an `elem_to_loc` overload. A `static_assert` ties the 8 to `JIT_MAX_NDIM`. The launch now appends an input's shape, strides and ndim only when its ndim is above zero, the same condition under which `build_kernel` declares them. Before, a 0-d input whose name appeared as `_shape` in the source would have shifted every later argument. No mlxcel kernel has such an input. + +The fix is recorded as LOCAL_FIXES item 30. It is a fork-side change confined to `custom_kernel.cpp`, and an upstreaming candidate under #1813; the upstream README marks it "not packaged yet" because a package needs a repro on the unpatched fork. `sampling_fixed_key_tests` faults without it. The `_strides` half has no user in mlxcel and was checked only by compiling it. + +## 6. Correctness Evidence + +A sampler has no elementwise fallback to diff against, so the PR holds the ports to the graph in four ways. The two existing suites also run on ROCm now instead of returning early, and pass on gfx1151. + +### 6.1 Fixed key, recomputed on the host + +Both kernels draw one Philox key per call from MLX's default key sequence, and given the key their output is deterministic. `sampling_fixed_key_tests` reseeds, reads the key the next launch will take through the new `random_bits` bridge function, reseeds again, launches, and recomputes the draw on the host. + +- `gumbel_kernel_matches_the_keyed_graph_argmax`: 48 rows of 5003 logits in f32 (T 1.0 and 0.7), bf16 (T 1.3) and f16 (T 0.9). The kernel's index equals the host's `argmax(logits / T + g)` on every decided row, and so does the MLX graph's `argmax` with the same noise. A row counts as decided when its winning score leads by at least 1e-4. +- `rejection_kernel_draws_the_keyed_token_under_min_p`: min-p 0.05, 40 rows of 3001 entries, T 1.0 and 0.7. Min-p resolves its threshold before the first draw, so round 0 accepts and the token follows from the round-0 Philox word and the kernel's thread-major scan order. The kernel matches the host on every decided row (target at least 1e-5 of the mass from a cumulative boundary). + +With the Philox counter (`row + 1`) or the drawn word (`c1` for `c0`) changed in the HIP sources, both tests fail on row 0. So the tests detect a port that consumes the key differently, not only one that draws from the wrong distribution. + +### 6.2 Support membership + +`rejection_kernel_draws_stay_inside_the_filtered_support` covers rounds past the first, which the host cannot reproduce bit for bit because the kernel's f32 sums may differ in their last bits. For 40 rows of 3001 entries under (top-k, top-p) = (0, 0.9), (40, 1.0) and (0, 0.5), four launches each, every row must converge and every drawn token must lie inside the support computed on the host from the same probabilities, with a relative 1e-5 slack on the top-p cutoff. The test also asserts that at least one row needed a second round, so the bisection is exercised. + +### 6.3 Chi-square against the graph sampler + +Where the kernel and the graph consume the RNG differently, a fixed key cannot line them up. Two-sample chi-square tests compare each kernel's histogram with `fused_sample_categorical`, the explicit graph arm, on the same input, at 400,000 draws per arm, pooled to 20 draws per bin, with the critical value at p = 1e-6: + +- the Gumbel kernel against `random::categorical` at T 0.8; +- the rejection kernel against the stock chain at top-p 0.9, and at top-p 0.95 with min-p 0.02. + +Both pass. Top-k with top-p is left out because the kernel is not exact there (see `sampling_rejection_tests.rs`). + +### 6.4 End-to-end speculative decoding + +The classic `SpeculativeGenerator` path verifies drafted tokens with the target's sampler, and `fused_sample_probs` follows the same routing, so with top-p the target's draws come from the rejection kernel. Qwen3-30B-A3B-4bit with Qwen3-0.6B-4bit as drafter, `--temp 0.8 --top-p 0.95 -n 128`, `MLXCEL_SPECULATIVE_ACCEPT_DIAG=1`: the dispatch log names the rejection kernel after and the stock chain before, the output is coherent, and per-position acceptance sits near the closed form for each rule. After: sampler-match 0.643 against `sum_prod` 0.665, stochastic 0.732 against `sum_min` 0.676. Before: 0.709 against 0.678 and 0.673 against 0.699, about 110 tested positions each. These are single runs that show the path works; the tests above carry the distributional claim. No MTP checkpoint was available, so the MTP and DFlash round loops were not run. + +## 7. Measured Results + +### 7.1 Decode throughput + +`scripts/bench_decode.sh` at pp512/tg128 with the new `--temperature` / `--top-p` options, Meta-Llama-3.1-8B-Instruct-4bit (vocab 128256), three runs per arm interleaved run by run (before, after, before, after). Before is `main` at `57d8ed29`; after is `85e39880`. Later commits do not touch any code that a sampler launch with 1-d and 2-d inputs reaches. + +| Configuration | Path after | Before tok/s | After tok/s | Medians | +|---|---|---|---|---| +| `--temperature 0.8` | Gumbel-max kernel | 37.91 / 37.67 / 38.32, median 37.91 | 37.59 / 38.26 / 39.25, median 38.26 | 1.01x, inside the spread | +| `--temperature 0.8 --top-p 0.95` | rejection kernel | 36.26 / 33.98 / 32.81, median 33.98 | 37.82 / 37.56 / 37.06, median 37.56 | 1.11x | + +`compare_bench_csv.py`, which keeps the last row per model and so compares the third runs, reports 1.024x and 1.130x. + +**Gumbel, not claimed.** The medians differ by 0.9%, against a spread of 1.7% in the before arm and 4.4% in the after arm, and the arms overlap. This is consistent with the profile's 0.4 to 3.8% share for the whole sampler tail. The PR reports no gain on this path. + +**Rejection, 1.11x with a drifting before arm.** Every after run is faster than every before run, by 1.04x, 1.11x and 1.13x in the three interleaved pairs. The before arm drifted down across its runs (36.26, 33.98, 32.81), while the after arm stayed within 37.06 to 37.82. The cause of the drift was not isolated. The pair ratios grow with the drift, so the 1.11x median ratio depends partly on it; the smallest pair, 1.04x, is the conservative reading. The graph chain runs an `argsort` over all 128256 entries per token for top-p, which is the work the rejection kernel replaces. + +### 7.2 Why top-k plus top-p was not measured + +The issue's second command was `--temp 0.8 --top-k 40 --top-p 0.95`. The routing policy sends top-k combined with top-p to the rejection kernel only up to vocab 32768 (`REJECTION_JOINT_VOCAB_MAX`, measured on M1 Ultra). Llama-3.1's vocabulary is 128256, so that configuration runs the stock chain both before and after, and any difference would be noise. Top-p alone is the routed configuration and is the one measured. Whether the 32768 cap is right on ROCm was not measured. + +### 7.3 Idle-GPU guard and the `--idle-secs 75` deviation + +Every run went through `scripts/rocm_gpu_guard.sh`, the copy fixed in #2098 run from the #2065 worktree and not committed in this PR. The issue's method asks for 90 s with `/sys/class/kfd/kfd/proc` empty and no compiler. A parallel unit (#2065) benchmarked at the same time with the guard's default 90. At equal windows both guards started their idle windows together and each rejected the other's run, in lockstep. This PR's runs used `--idle-secs 75`, 75 consecutive 1 Hz samples, which took 112 to 116 s of wall time per run and broke the tie. All 12 accepted runs were clean on their first attempt. The deviation lowers the required sample count; each run's guard phase still took 112 to 116 s of wall time, more than the method's 90 s. + +## 8. The Snapshot Bound: 1 ulp to 2 ulp + +`sampling::tests::temperature_one_support_unchanged` compares `fused_sample_probs` at T 1.0 with saved Metal rows. The routed cases had skipped wherever the rejection kernel had no port, so they first ran on ROCm with this PR. The gate at `f7bec56d` failed only this test. These rows are the graph's softmax over the kernel's support; `fused_sample_probs` never launches the kernel. gfx1151 lands 2 ulp from the Metal capture at token 26 of the (top-k 40, top-p 0.9) row, with the support unchanged. Commit `6e336ec8` raises the routed bound from 1 to 2 ulp, allowing the same kind of reassociation drift as the non-routed stock case, which allows 4. + +The bound is not per-backend. The change widens it on Metal and CUDA too, where the routed rows had run at 1 ulp before. A 2 ulp drift there now passes silently. The support-set assertion in the same test, which is the correctness gate, is unchanged. + +## 9. Technical Decisions + +- **Port the CUDA bodies line for line.** CUDA and HIP share the grid semantics, `template_args` and hipRTC-style runtime compilation. Keeping the bodies identical except for spelling leaves one place to look when they disagree, and keeps the Philox layout bit-identical across three backends. +- **No wave32 guard.** Shared-memory reductions with barriers are wave-size independent. A guard would reject a correct kernel on wave64 and, per #2067, does not fire on ROCm 10 clang anyway. +- **Keep the launches in the existing files.** Moving the HIP body to a header for size, but leaving the `fast::hip_kernel` call next to the CUDA one, keeps the dtype-key checker's pinned scope and count unchanged. +- **Make the rejection predicate read the table.** `has_kernel_port(rejection_ports())` cannot disagree with `select_kernel_port`, and the next backend port becomes reachable with no predicate edit. +- **Fix the overlay rather than work around it.** Reading the vocab from a scalar input would have avoided `logits_shape`, but the overlay's declaration was wrong for every future HIP kernel that reads a shape. The by-value structs match upstream CUDA and the fix is upstreamable. +- **Add env-independent bridge predicates for the kill-switch tests.** The tests need to know whether a port exists while the kill switch is set; `custom_kernels_available()` answers a different question. +- **Hold the ports to the graph on the same randomness.** The fixed-key tests catch a port that draws from the right distribution but consumes the key differently, which the existing chi-square suites would pass. + +## 10. Validation + +From the PR body, on gfx1151 (Radeon 8060S, ROCm 10.0.0): + +- `make verify-rocm` at `6e336ec8`: OK; 11861 passed, 0 failed, 378 ignored across 146 test binaries; ROCm smoke OK. The later commit `74624ea7` changes docs, two header comments and `bench_decode.sh` input validation; the fast gates, fmt and `dead_doc_pointers` pass on it. +- `cargo test --release --features rocm -p mlxcel-core --lib sampling_ -- --test-threads=1`: 53 passed, plus both kill-switch integration tests. +- `cargo clippy -p mlxcel-core --features rocm --lib --tests -- -D warnings`, `make verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-rocm-overlay verify-fmt`: pass. + +Orchestrator verification on head `74624ea7`, up to date with origin/main `57d8ed29`: + +- `make verify-rocm` passed every step: 11,861 tests passed, 0 failed, 378 ignored; smoke OK. + +## 11. Residual Risks and What Was Not Verified + +- **Metal and CUDA were not run.** They are touched by the rejection predicate (`has_kernel_port`, true on both as before), the shared bridge functions (`random_bits`, `sampling_*_backend_supported`), the kill-switch tests' predicate, the routed snapshot bound and `sampling_fixed_key_tests`, which is written to run there too. +- **The snapshot bound is looser everywhere.** A 2 ulp drift on Metal or CUDA in the routed rows no longer trips the test. +- **The `_strides` half of the overlay fix has no user.** No kernel reads `_strides`, so it is checked only by the struct compiling in every kernel header. +- **The before-arm drift in the top-p measurement is unexplained.** The 1.11x median ratio leans on it; the per-pair range is 1.04x to 1.13x. +- **The joint top-k plus top-p cap was not re-measured on ROCm.** At vocab above 32768 that configuration stays on the stock chain. +- **MTP and DFlash speculative loops were not run.** No MTP checkpoint was available locally. +- **The guard copy is not in this PR.** The runs used the #2098 fix from the #2065 worktree, with `--idle-secs 75` instead of the method's 90 samples. +- **One device, one session, shared GPU.** All numbers are from gfx1151. No wave64 hardware was available, so wave-size independence on CDNA is reasoned from the kernels' structure, not run. + +## 12. Learning Points + +- **A predicate that enumerates backends hides a filled slot.** `custom_kernels_available()` answered a family-wide question; per-kernel predicates must read their own port table. +- **The first kernel to use an argument finds the bug in how it is passed.** The overlay's shape declaration was wrong since the fork, and nothing failed until a kernel read `logits_shape`. A memory fault at an address that decodes to a shape is a sign that a by-value argument is being read as a pointer. +- **Distribution tests do not prove a port is the same sampler.** A port with a different Philox counter layout draws from the right distribution and passes chi-square. Recomputing the draw from the consumed key is what pins the RNG contract. +- **A wave32 guard belongs only where lanes talk to each other.** Shared-memory reductions with barriers need no guard, and adding one would reject a correct kernel. +- **Check that the measured configuration reaches the kernel.** The issue's top-k plus top-p command routes to the stock chain at this vocabulary; measuring it would have reported noise as a result. +- **Concurrent guards with equal windows can lock each other out.** Two idle guards with the same sample count start together and reject each other; a different count breaks the tie. + +Refs: #2064, #1814, #1801, #1813, #2061, #2067, #2099, #2065, #2098, #1862, #1870, #1885, #900, #901, #1803. diff --git a/TECHNICAL_REPORTS/2100-rocm-samplers-hip-port-20261005.ko.md b/TECHNICAL_REPORTS/2100-rocm-samplers-hip-port-20261005.ko.md new file mode 100644 index 000000000..f8b4c5317 --- /dev/null +++ b/TECHNICAL_REPORTS/2100-rocm-samplers-hip-port-20261005.ko.md @@ -0,0 +1,191 @@ +# 기술 보고서: PR #2100 - Gumbel-max sampler와 rejection sampler를 HIP으로 port + +**날짜**: 2026-10-05 + +**상태**: gfx1151 호스트에서 구현 및 측정 완료. head `74624ea7`, origin/main `57d8ed29` 위, 머지 대기 중. + +**언어**: C++ (HIP kernel source, kernel holder, support predicate, bridge 함수, vendored ROCm overlay), Rust (fixed-key 및 two-sample test, kill-switch test, snapshot bound), Bash (`bench_decode.sh` 옵션, dtype-key checker test), Markdown (환경 변수, 설치, benchmark, LOCAL_FIXES, benchmark 결과), CSV (bench 행) + +**위험도**: 중간 (ROCm에서 모든 모델의 기본 sampled decode 경로가 바뀝니다. 모든 HIP custom kernel이 거치는 vendored `fast::hip_kernel`의 argument 배치를 고치고, Metal과 CUDA가 공유하는 rejection predicate, bridge 함수, snapshot bound도 건드리며, 두 backend는 개발 호스트에서 실행할 수 없었습니다) + +## 요약 + +Issue #2064 (#1814의 일부, epic #1801)는 temperature가 0보다 클 때 생성 token마다 한 번씩 실행되는 두 fused sampler의 HIP port를 요청했습니다. Filter 없는 경로의 Gumbel-max sampler와 filter 경로의 dual-pivot rejection sampler입니다. 이 PR 이전에는 둘 다 ROCm에서 MLX graph로 실행되었습니다. `gumbel_ports()`와 `rejection_ports()`에는 `.rocm` 항목이 없었고, `rejection_sample_supported()`는 정의상 Metal 또는 CUDA만 뜻하는 `custom_kernels_available()`을 반환했으므로, rejection slot을 채워도 도달할 수 없었습니다. + +이 PR은 CUDA 본문을 줄 단위로 옮긴 `sampling_gumbel_hip.h`와 `sampling_rejection_hip.h`를 추가합니다. Input, output, grid, template argument, Philox-4x32-10 counter와 key 배치는 CUDA와 같습니다. 두 `.rocm` slot을 채우고, `rejection_sample_supported()`가 `has_kernel_port(rejection_ports())`를 반환하게 합니다. 두 port 모두 wave32 `#error` guard를 두지 않습니다. 두 kernel에는 lane 수준 연산이 없고, 모든 reduction과 rejection kernel의 scan이 step마다 barrier를 두는 shared memory로 이루어지기 때문입니다. + +Gumbel port는 vendored ROCm overlay의 버그를 드러냈습니다. `fast::hip_kernel`은 `_shape`와 `_strides`를 pointer로 선언했지만 launch는 값으로 넘겼고, 그래서 이를 처음 읽은 kernel (`logits_shape[1]`)이 queue를 fault시켰습니다. 이는 LOCAL_FIXES item 30으로 고쳤으며, fork 쪽 변경이고 #1813의 upstream 후보입니다. + +gfx1151에서 Llama-3.1-8B-Instruct-4bit, pp512/tg128로 측정한 결과 top-p 0.95 decode의 중앙값은 33.98에서 37.56 tok/s (1.11x)가 되었고, 이전 arm이 아래로 drift한 사실도 함께 적습니다. Gumbel 경로는 37.91에서 38.26 tok/s (+0.9%)로, 실행 간 편차 안이므로 향상으로 주장하지 않습니다. Issue의 top-k + top-p 명령은 vocab 128256에서 kernel에 도달하지 않으므로 측정하지 않았습니다. + +## 1. 문제 정의 + +### 1.1 빈 slot 두 개, 그것을 보지 못하는 predicate 하나 + +`gumbel_max_sample_supported()`는 이미 `has_kernel_port(gumbel_ports())`를 읽었으므로, `.rocm`만 채우면 도달 가능해졌습니다. `rejection_sample_supported()`는 그렇지 않았습니다. + +```cpp +return mlxcel::custom_kernels_available(); +``` + +`custom_kernels_available_for`는 Metal과 CUDA에서만 true이므로, ROCm에서는 `rejection_ports()`에 무엇이 있든 predicate가 false였습니다. 이 predicate는 `sampling_rejection_available()`과 bridge의 routing에 쓰입니다. `verify-kernel-port-dispatch`의 checker rule 4는 이것이 `metal_is_available() || cuda_is_available()` 형태로 쓰이지 않았기 때문에 잡지 못했습니다. + +### 1.2 Profile이 말한 기대 효과 + +Decode profile (`rocm-decode-profile-gfx1151-2026-09-30.md`)은 sampler 꼬리 전체를 greedy decode GPU 시간의 0%, sampled decode GPU 시간의 0.4 ~ 3.8%로 보고했습니다. Port가 없을 때 ROCm의 sampled decode는 filter 없는 경로에서 `random::categorical`, filter 경로에서 `argpartition` / `argsort` / `cumsum` chain을 실행했습니다. 따라서 filter 없는 경로의 기대 이득은 작고, graph가 vocabulary 전체를 정렬하는 경로에서만 더 클 것으로 보였습니다. + +### 1.3 같은 난수 위에서 port를 graph에 맞춰 보는 test가 없었음 + +기존 suite인 `sampling_gumbel_tests`와 `sampling_rejection_tests`는 각 kernel을 정확한 목표 분포와 비교합니다. 이들은 ROCm에서 일찍 반환했습니다. Port가 MLX의 key sequence를 다른 port와 같은 방식으로 소비하는지 확인하는 test는 없었습니다. 그것이 `mlx::core::random::seed(...)`가 backend 사이에서 stream을 재현하게 하는 조건입니다. + +## 2. 변경 요약 + +| 영역 | 변경 | +|---|---| +| `sampling_gumbel_hip.h` (신규) | `GUMBEL_MAX_SAMPLE_HIP_SOURCE`, 표기만 바꾼 CUDA 본문 | +| `sampling_rejection_hip.h` (신규) | `REJECTION_SAMPLE_HIP_SOURCE`, 표기만 바꾼 CUDA 본문 | +| `sampling.cpp` | `MLXCEL_BRIDGE_ROCM_BACKEND` 아래에서 `fast::hip_kernel`을 호출하는 `GumbelKernelHolderHip`; `gumbel_ports()`의 `.rocm` getter | +| `sampling_rejection.cpp` | `RejectionKernelHolderHip`; `rejection_ports()`의 `.rocm` getter; `rejection_sample_supported()`가 `has_kernel_port(rejection_ports())` 반환 | +| ROCm overlay `custom_kernel.cpp` | 값으로 전달되는 `KernelShape` / `KernelStrides` struct, `elem_to_loc` overload, `JIT_MAX_NDIM`에 대한 `static_assert`, ndim > 0일 때만 shape/strides/ndim 추가 (LOCAL_FIXES item 30) | +| Bridge (`mlx_cxx_bridge.cpp/.h`, `lib.rs`) | env와 무관한 `sampling_gumbel_backend_supported()`, `sampling_rejection_backend_supported()`, 그리고 `random_bits()`; 주석 수정 | +| `gpu_backend.h/.cpp` | `custom_kernels_available()`의 현재 의미에 대한 주석; `MLXCEL_DEBUG_KERNEL_BACKEND` 접미사 문구 수정 | +| Test | 신규 `sampling_fixed_key_tests.rs` (test 다섯 개); kill-switch test가 새 predicate를 물음; `temperature_one_support_unchanged`의 routed bound 1 ulp에서 2 ulp로; dtype-key checker test가 launch 둘을 모두 옮김; Python probe test가 새 접미사를 따름 | +| `bench_decode.sh` | `--temperature`, `--top-p` (단순 십진수만 허용, 자동 생성 파일명에 tag) | +| 문서와 데이터 | `environment-variables.md`, `installation.md`, `benchmarks.md`, item 30에 대한 upstream README 행, `LOCAL_FIXES.md` item 30, 신규 `rocm-samplers-gfx1151-2026-10-05.md`, bench CSV 네 개 | + +Branch에는 commit이 여섯 개 있습니다. Overlay 수정 (`cc93789f`), port와 predicate (`85e39880`), 결과 문서 (`ce4105b7`), kill-switch predicate 수정 (`f7bec56d`), snapshot bound (`6e336ec8`), 문서와 `bench_decode.sh` 입력 검증 정리 (`74624ea7`)입니다. Diff는 파일 32개, 추가 1,341줄, 삭제 46줄입니다. + +## 3. Port + +### 3.1 줄 단위 이식, 세 종류의 표기 변경 + +두 HIP source는 CUDA kernel의 thread mapping, grid, input (`logits`, `rng_key`, `temp`; `probs`, `probs_draw`, `params`, `rng_key`), output (`vals`, `idxs`; `ids`, `ok`, `rounds`), template argument를 dtype key까지 그대로 유지합니다. Philox-4x32-10 counter (Gumbel은 `{base/4, 0, row, 0}`, rejection은 `{round, 0, row, 0}`), key `{rng_key[0], rng_key[1]}`, 23-bit open-interval uniform도 같으므로, seed가 CUDA와 Metal에서처럼 ROCm stream을 재현합니다. `__umulhi`는 HIP에서도 같은 뜻입니다. 차이는 다음과 같습니다. + +- `hip_bfloat16` 값을 읽을 때 명시적 `(float)`. ROCm에서 MLX가 bfloat16 대신 쓰는 이 type은 `explicit` 연산자로만 float로 변환됩니다. +- `-INFINITY` 대신 `-__builtin_huge_valf()`. hipRTC는 `fast::hip_kernel`이 앞에 붙이는 header만으로 본문을 compile하며, 그 안에 ``가 있다는 보장이 없습니다. +- `temp[0]`에 붙은 불필요한 `(float)`. + +Holder와 launch는 `sampling.cpp`와 `sampling_rejection.cpp`의 CUDA 것 옆에 남습니다. Header는 launch가 아니라 데이터만 담으므로 `make verify-kernel-dtype-keys`의 `EXPECTED_IN_SCOPE`와 그 개수가 바뀌지 않습니다. Sampler의 launch를 helper로 옮기거나 주석 처리하는 checker 자체의 test는 이제 HIP launch에도 같은 일을 합니다. + +### 3.2 Wave32 guard를 두지 않은 이유 + +#1814 port 요구사항은 lane 사이에서 reduction이나 shuffle을 하는 port마다 wavefront 크기 macro에 대한 `#error` 검사 두 개를 두게 합니다. Lane 16에서 시작하는 fold는 64-lane CDNA wave의 절반을 버리고도 유한하고 틀린 결과를 내기 때문입니다. 이 두 kernel에는 lane 수준 연산이 없습니다. Gumbel kernel의 index를 함께 나르는 halving reduction, rejection kernel의 reduction과 Hillis-Steele scan은 모두 step 사이에 `__syncthreads()`를 두는 `__shared__` memory를 거치므로, 결과가 한 wave에 몇 thread가 있는지에 의존하지 않습니다. Wave64에서 `#error`를 내면 올바른 kernel을 거부하게 됩니다. 또한 PR #2099는 ROCm 10의 AMD clang이 두 macro를 모두 정의하지 않아 guard 관용구가 아무 효과가 없다는 것을 발견했으므로, 여기서 guard를 두어도 잡아낼 것이 없습니다. Issue의 acceptance criterion은 이 이유를 옆에 적고 체크되었습니다. + +## 4. Predicate와 그것을 읽는 test + +`rejection_sample_supported()`는 `default_device() != Device::gpu` 검사를 유지한 뒤 `mlxcel::has_kernel_port(rejection_ports())`를 반환합니다. 이는 `rejection_sample`의 `select_kernel_port`가 읽는 바로 그 table이며, `gumbel_max_sample_supported()`와 같은 형태입니다. Metal과 CUDA에서는 이전처럼 true입니다. `fused_sample`, `fused_sample_probs`, speculative 경로는 호출자 변경 없이 HIP kernel에 도달합니다. + +Integration test 두 개가 예전 의미에 의존했습니다. `tests/sampling_gumbel_kill_switch.rs`와 `tests/sampling_rejection_kill_switch.rs`는 env kill switch를 켜고, 어떤 거절 메시지를 기대할지를 `custom_kernels_available()`로 골랐습니다. 이 함수는 여전히 Metal 또는 CUDA만 true이므로, 이제 port가 있는 ROCm에서 test는 "port 없음"을 기대했을 것입니다. Env switch가 켜져 있으면 `sampling_*_available()`은 backend와 무관하게 false이므로, test에는 switch를 빼고 묻는 predicate가 필요했습니다. 이를 위해 PR은 bridge에 `sampling_gumbel_backend_supported()`와 `sampling_rejection_backend_supported()`를 추가했습니다 (`f7bec56d`). `gpu_backend.h` 주석은 이제 `custom_kernels_available()`이 family 단위 gate이고 HIP port가 있는 kernel은 자기 predicate로 gate된다고 설명합니다. `MLXCEL_DEBUG_KERNEL_BACKEND` 접미사는 "no custom kernel ports"에서 "not every kernel family is ported"로 바뀌었습니다. + +## 5. Overlay 버그: shape와 stride를 값으로 + +Fork의 `mlx/backend/rocm/custom_kernel.cpp`에 있는 `build_kernel`은 input의 shape를 `const int32_t* _shape`, stride를 `const int64_t* _strides`로 선언했습니다. `CustomKernel::eval_gpu`는 이들을 `KernelArgs::append_ndim`으로 넘기는데, 이 함수는 vector를 `JIT_MAX_NDIM` (8) 항목으로 채운 뒤 그 저장소를 가리키는 pointer를 `hipModuleLaunchKernel`에 줍니다. 그래서 kernel은 값 자체를 받았고, shape의 처음 8 byte를 주소로 읽었습니다. HIP Gumbel port는 CUDA, Metal port처럼 `logits_shape[1]`을 읽으며, `HSA_STATUS_ERROR_MEMORY_FAULT`로 `0x80000000000`에서 죽었습니다. 이는 `0x80000000800`이 속한 page이고, 이 값은 shape `{2048, 2048}`을 64-bit pointer로 읽은 것입니다. mlxcel의 어떤 HIP kernel도 이전에 두 argument를 읽은 적이 없었기 때문에 드러나지 않았습니다. + +수정 (`cc93789f`)은 생성되는 header에 `operator[]`가 있는 `KernelShape` (`int32_t` 8개)와 `KernelStrides` (`int64_t` 8개)를 추가하고, parameter를 이들로 값 전달하도록 선언합니다. Upstream CUDA의 `const __grid_constant__ Shape` / `Strides`와 같은 방식입니다. `elem_to_loc` overload도 추가하고, `static_assert`로 8을 `JIT_MAX_NDIM`에 묶습니다. Launch는 이제 input의 ndim이 0보다 클 때만 shape, stride, ndim을 추가하며, 이는 `build_kernel`이 그것들을 선언하는 조건과 같습니다. 이전에는 이름이 source에 `_shape`로 나오는 0-d input이 이후 모든 argument를 밀어냈을 것입니다. mlxcel kernel 중 그런 input을 가진 것은 없습니다. + +이 수정은 LOCAL_FIXES item 30으로 기록됩니다. `custom_kernel.cpp`에 한정된 fork 쪽 변경이며 #1813 아래의 upstream 후보입니다. Upstream README는 package를 만들려면 패치되지 않은 fork에서의 재현이 필요하므로 "not packaged yet"으로 표시합니다. `sampling_fixed_key_tests`는 이 수정 없이 fault합니다. `_strides` 쪽은 mlxcel에 사용처가 없어 compile로만 확인되었습니다. + +## 6. 정확성 근거 + +Sampler에는 원소 단위로 비교할 fallback이 없으므로, PR은 네 가지 방식으로 port를 graph에 맞춰 봅니다. 기존 suite 두 개도 이제 ROCm에서 일찍 반환하지 않고 실행되며 gfx1151에서 통과합니다. + +### 6.1 고정 key, host에서 다시 계산 + +두 kernel은 호출마다 MLX의 기본 key sequence에서 Philox key 하나를 뽑으며, key가 주어지면 출력이 결정적입니다. `sampling_fixed_key_tests`는 reseed하고, 새 `random_bits` bridge 함수로 다음 launch가 가져갈 key를 읽고, 다시 reseed한 뒤 launch하고, host에서 draw를 다시 계산합니다. + +- `gumbel_kernel_matches_the_keyed_graph_argmax`: 5003개 logit 48행, f32 (T 1.0, 0.7), bf16 (T 1.3), f16 (T 0.9). 모든 decided 행에서 kernel의 index가 host의 `argmax(logits / T + g)`와 같고, 같은 noise를 넣은 MLX graph의 `argmax`와도 같습니다. 이기는 score가 1e-4 이상 앞설 때 decided로 봅니다. +- `rejection_kernel_draws_the_keyed_token_under_min_p`: min-p 0.05, 3001개 항목 40행, T 1.0, 0.7. Min-p는 첫 draw 전에 threshold를 확정하므로 round 0이 accept하고, token은 round 0 Philox word와 kernel의 thread-major scan 순서로 정해집니다. 모든 decided 행 (target이 누적 경계에서 mass의 1e-5 이상 떨어진 행)에서 kernel이 host와 일치합니다. + +HIP source에서 Philox counter (`row + 1`)나 뽑는 word (`c0` 대신 `c1`)를 바꾸면 두 test 모두 row 0에서 실패합니다. 즉 이 test는 잘못된 분포에서 뽑는 port뿐 아니라 key를 다르게 소비하는 port도 잡아냅니다. + +### 6.2 Support 소속 + +`rejection_kernel_draws_stay_inside_the_filtered_support`는 첫 round 이후를 다룹니다. Kernel의 f32 합이 마지막 bit에서 다를 수 있어 host가 bit 단위로 재현할 수 없는 부분입니다. 3001개 항목 40행에 대해 (top-k, top-p) = (0, 0.9), (40, 1.0), (0, 0.5)로 각각 네 번 launch하며, 모든 행이 수렴하고 모든 draw가 같은 확률로 host에서 계산한 support 안에 있어야 합니다. Top-p cutoff에는 상대 1e-5의 여유를 둡니다. 또한 최소 한 행이 두 번째 round를 필요로 했는지 확인하므로 bisection이 실제로 실행됩니다. + +### 6.3 Graph sampler와의 chi-square + +Kernel과 graph가 RNG를 다르게 소비하는 곳에서는 고정 key로 둘을 맞출 수 없습니다. Two-sample chi-square test가 같은 input에서 각 kernel의 histogram을 명시적 graph arm인 `fused_sample_categorical`과 비교합니다. Arm당 400,000번 draw, bin당 20 draw로 pooling, critical value는 p = 1e-6입니다. + +- Gumbel kernel 대 `random::categorical`, T 0.8; +- rejection kernel 대 stock chain, top-p 0.9, 그리고 top-p 0.95와 min-p 0.02. + +둘 다 통과합니다. Top-k와 top-p의 조합은 kernel이 정확하지 않은 경우이므로 제외했습니다 (`sampling_rejection_tests.rs` 참고). + +### 6.4 End-to-end speculative decoding + +Classic `SpeculativeGenerator` 경로는 target의 sampler로 draft token을 검증하고, `fused_sample_probs`도 같은 routing을 따르므로, top-p에서는 target의 draw가 rejection kernel에서 나옵니다. Qwen3-30B-A3B-4bit에 Qwen3-0.6B-4bit drafter, `--temp 0.8 --top-p 0.95 -n 128`, `MLXCEL_SPECULATIVE_ACCEPT_DIAG=1`로 실행했습니다. Dispatch log는 이후에는 rejection kernel, 이전에는 stock chain을 가리키고, 출력은 자연스러우며, 위치별 acceptance는 각 규칙의 closed form 근처에 있습니다. 이후: sampler-match 0.643 대 `sum_prod` 0.665, stochastic 0.732 대 `sum_min` 0.676. 이전: 0.709 대 0.678, 0.673 대 0.699, 각각 약 110개 위치. 단일 실행이므로 경로가 동작함을 보일 뿐이고, 분포에 대한 주장은 위 test가 맡습니다. MTP checkpoint가 없어 MTP와 DFlash round loop는 실행하지 않았습니다. + +## 7. 측정 결과 + +### 7.1 Decode throughput + +`scripts/bench_decode.sh`를 pp512/tg128과 새 `--temperature` / `--top-p` 옵션으로 실행했습니다. Meta-Llama-3.1-8B-Instruct-4bit (vocab 128256), arm당 세 번, 실행 단위로 교차 (before, after, before, after)했습니다. Before는 `main`의 `57d8ed29`, after는 `85e39880`입니다. 이후 commit은 1-d, 2-d input으로 sampler를 launch할 때 거치는 코드를 건드리지 않습니다. + +| 구성 | 이후 경로 | Before tok/s | After tok/s | 중앙값 | +|---|---|---|---|---| +| `--temperature 0.8` | Gumbel-max kernel | 37.91 / 37.67 / 38.32, 중앙값 37.91 | 37.59 / 38.26 / 39.25, 중앙값 38.26 | 1.01x, 편차 안 | +| `--temperature 0.8 --top-p 0.95` | rejection kernel | 36.26 / 33.98 / 32.81, 중앙값 33.98 | 37.82 / 37.56 / 37.06, 중앙값 37.56 | 1.11x | + +모델마다 마지막 행만 남겨 세 번째 실행끼리 비교하는 `compare_bench_csv.py`는 1.024x와 1.130x를 보고합니다. + +**Gumbel, 주장하지 않음.** 중앙값 차이는 0.9%이고, before arm의 편차는 1.7%, after arm은 4.4%이며 두 arm이 겹칩니다. 이는 sampler 꼬리 전체가 0.4 ~ 3.8%라는 profile과도 맞습니다. PR은 이 경로에서 향상을 보고하지 않습니다. + +**Rejection, drift하는 before arm과 함께 1.11x.** 모든 after 실행이 모든 before 실행보다 빠르며, 교차된 세 쌍에서 1.04x, 1.11x, 1.13x입니다. Before arm은 실행을 거듭하며 아래로 drift했고 (36.26, 33.98, 32.81), after arm은 37.06 ~ 37.82 안에 머물렀습니다. Drift의 원인은 밝히지 못했습니다. 쌍별 비율이 drift와 함께 커지므로 중앙값 비율 1.11x는 일부 drift에 기대고 있으며, 가장 작은 쌍인 1.04x가 보수적인 해석입니다. Graph chain은 top-p를 위해 token마다 128256개 항목 전체에 `argsort`를 실행하며, rejection kernel이 대체하는 것이 바로 그 작업입니다. + +### 7.2 Top-k + top-p를 측정하지 않은 이유 + +Issue의 두 번째 명령은 `--temp 0.8 --top-k 40 --top-p 0.95`였습니다. Routing 정책은 top-k와 top-p의 조합을 vocab 32768까지만 rejection kernel로 보냅니다 (`REJECTION_JOINT_VOCAB_MAX`, M1 Ultra에서 측정). Llama-3.1의 vocabulary는 128256이므로 이 구성은 before와 after 모두 stock chain을 실행하며, 차이가 나와도 noise일 뿐입니다. Top-p 단독이 routing되는 구성이므로 그것을 측정했습니다. 32768 상한이 ROCm에서 맞는지는 측정하지 않았습니다. + +### 7.3 Idle-GPU guard와 `--idle-secs 75` 편차 + +모든 실행은 `scripts/rocm_gpu_guard.sh`를 거쳤습니다. #2098에서 고친 copy를 #2065 worktree에서 실행했으며 이 PR에 commit되지 않았습니다. Issue의 방법은 `/sys/class/kfd/kfd/proc`가 비어 있고 compiler가 없는 상태 90초를 요구합니다. 병렬 unit (#2065)이 같은 시간에 guard 기본값 90으로 benchmark했습니다. 같은 길이에서는 두 guard가 idle window를 함께 시작했고, 서로의 실행을 lockstep으로 거부했습니다. 이 PR의 실행은 `--idle-secs 75` (1 Hz sample 75개 연속)를 썼고, 이는 실행당 112 ~ 116초의 wall time이 걸렸으며 동률을 깼습니다. 받아들여진 12번의 실행은 모두 첫 시도에서 깨끗했습니다. 이 편차는 요구 sample 수를 줄인 것이며, 각 실행의 guard 단계는 여전히 112 ~ 116초로 방법의 90초보다 길었습니다. + +## 8. Snapshot bound: 1 ulp에서 2 ulp로 + +`sampling::tests::temperature_one_support_unchanged`는 T 1.0에서 `fused_sample_probs`를 저장된 Metal 행과 비교합니다. Routed case는 rejection kernel에 port가 없는 곳에서 건너뛰었으므로, 이 PR로 처음 ROCm에서 실행되었습니다. `f7bec56d`의 gate는 이 test 하나만 실패했습니다. 이 행들은 kernel의 support 위에서 graph가 계산한 softmax이며, `fused_sample_probs`는 kernel을 launch하지 않습니다. gfx1151은 (top-k 40, top-p 0.9) 행의 token 26에서 Metal capture와 2 ulp 차이가 나고, support는 바뀌지 않습니다. Commit `6e336ec8`은 routed bound를 1 ulp에서 2 ulp로 올리며, 4를 허용하는 non-routed stock case와 같은 종류의 reassociation drift를 허용합니다. + +이 bound는 backend별이 아닙니다. 이 변경은 routed 행이 이전에 1 ulp로 실행되던 Metal과 CUDA에서도 bound를 넓힙니다. 그곳에서 2 ulp drift는 이제 조용히 통과합니다. 같은 test의 정확성 gate인 support 집합 assertion은 바뀌지 않았습니다. + +## 9. 기술적 선택과 그 이유 + +- **CUDA 본문을 줄 단위로 이식.** CUDA와 HIP은 grid 의미, `template_args`, hipRTC 방식의 runtime compile을 공유합니다. 표기 외에는 본문을 같게 두면 둘이 다를 때 볼 곳이 한 군데이고, Philox 배치가 세 backend에서 bit 단위로 같게 유지됩니다. +- **Wave32 guard 없음.** Barrier가 있는 shared memory reduction은 wave 크기와 무관합니다. Guard는 wave64에서 올바른 kernel을 거부하고, #2067에 따르면 ROCm 10 clang에서는 어차피 발동하지 않습니다. +- **Launch를 기존 파일에 유지.** 크기 때문에 HIP 본문만 header로 옮기고 `fast::hip_kernel` 호출은 CUDA 호출 옆에 두어, dtype-key checker의 고정된 범위와 개수를 바꾸지 않았습니다. +- **Rejection predicate가 table을 읽게 함.** `has_kernel_port(rejection_ports())`는 `select_kernel_port`와 어긋날 수 없고, 다음 backend port는 predicate 수정 없이 도달 가능해집니다. +- **우회하지 않고 overlay를 고침.** Vocab을 scalar input으로 읽으면 `logits_shape`를 피할 수 있었겠지만, overlay의 선언은 shape를 읽는 모든 향후 HIP kernel에 대해 틀렸습니다. 값 전달 struct는 upstream CUDA와 맞고 upstream으로 보낼 수 있습니다. +- **Kill-switch test용 env 무관 bridge predicate 추가.** Test는 kill switch가 켜진 상태에서 port가 있는지 알아야 하며, `custom_kernels_available()`은 다른 질문에 답합니다. +- **같은 난수 위에서 port를 graph에 맞춰 봄.** Fixed-key test는 올바른 분포에서 뽑지만 key를 다르게 소비하는 port를 잡아내며, 기존 chi-square suite는 그런 port를 통과시킵니다. + +## 10. 검증 + +PR 본문 기준, gfx1151 (Radeon 8060S, ROCm 10.0.0): + +- `6e336ec8`에서 `make verify-rocm`: OK. 146개 test binary에서 11861 통과, 0 실패, 378 ignored. ROCm smoke OK. 이후 commit `74624ea7`은 문서, header 주석 두 개, `bench_decode.sh` 입력 검증만 바꾸며, fast gate, fmt, `dead_doc_pointers`가 통과합니다. +- `cargo test --release --features rocm -p mlxcel-core --lib sampling_ -- --test-threads=1`: 53 통과, kill-switch integration test 둘도 통과. +- `cargo clippy -p mlxcel-core --features rocm --lib --tests -- -D warnings`, `make verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-rocm-overlay verify-fmt`: 통과. + +Orchestrator 검증, head `74624ea7`, origin/main `57d8ed29`과 최신 상태: + +- `make verify-rocm` 모든 단계 통과: 11,861 test 통과, 0 실패, 378 ignored. Smoke OK. + +## 11. 남은 위험과 검증하지 못한 부분 + +- **Metal과 CUDA는 실행하지 않았습니다.** Rejection predicate (`has_kernel_port`, 둘 다 이전처럼 true), 공유 bridge 함수 (`random_bits`, `sampling_*_backend_supported`), kill-switch test의 predicate, routed snapshot bound, 그곳에서도 실행되도록 쓰인 `sampling_fixed_key_tests`가 이 backend들에 닿습니다. +- **Snapshot bound가 모든 곳에서 느슨해졌습니다.** Metal이나 CUDA의 routed 행에서 2 ulp drift는 더 이상 test를 실패시키지 않습니다. +- **Overlay 수정의 `_strides` 쪽에는 사용자가 없습니다.** `_strides`를 읽는 kernel이 없으므로, 모든 kernel header에서 struct가 compile되는 것으로만 확인됩니다. +- **Top-p 측정에서 before arm의 drift는 설명되지 않았습니다.** 중앙값 비율 1.11x가 이에 기대며, 쌍별 범위는 1.04x ~ 1.13x입니다. +- **Top-k + top-p 결합 상한은 ROCm에서 다시 측정하지 않았습니다.** Vocab이 32768을 넘으면 그 구성은 stock chain에 남습니다. +- **MTP와 DFlash speculative loop는 실행하지 않았습니다.** 로컬에 MTP checkpoint가 없었습니다. +- **Guard copy는 이 PR에 없습니다.** 실행은 #2065 worktree의 #2098 수정을 썼고, 방법의 90 sample 대신 `--idle-secs 75`를 썼습니다. +- **장치 하나, 세션 하나, 공유 GPU.** 모든 수치는 gfx1151에서 나왔습니다. Wave64 하드웨어가 없었으므로 CDNA에서 wave 크기와 무관하다는 것은 실행이 아니라 kernel 구조로부터 추론한 것입니다. + +## 12. 학습 포인트 + +- **Backend를 나열하는 predicate는 채워진 slot을 숨깁니다.** `custom_kernels_available()`은 family 단위 질문에 답했습니다. Kernel별 predicate는 자기 port table을 읽어야 합니다. +- **Argument 전달 방식의 버그는 그 argument를 처음 쓰는 kernel이 찾습니다.** Overlay의 shape 선언은 fork 때부터 틀렸고, kernel이 `logits_shape`를 읽기 전까지 아무것도 실패하지 않았습니다. Shape로 해석되는 주소에서의 memory fault는 값으로 전달된 argument를 pointer로 읽고 있다는 신호입니다. +- **분포 test만으로는 port가 같은 sampler임을 증명하지 못합니다.** Philox counter 배치가 다른 port도 올바른 분포에서 뽑고 chi-square를 통과합니다. 소비한 key로 draw를 다시 계산해야 RNG 계약이 고정됩니다. +- **Wave32 guard는 lane끼리 통신하는 곳에만 필요합니다.** Barrier가 있는 shared memory reduction에는 guard가 필요 없고, 넣으면 올바른 kernel을 거부합니다. +- **측정하는 구성이 kernel에 도달하는지 확인합니다.** Issue의 top-k + top-p 명령은 이 vocabulary에서 stock chain으로 갑니다. 그것을 측정했다면 noise를 결과로 보고했을 것입니다. +- **같은 window를 쓰는 동시 guard는 서로를 막을 수 있습니다.** Sample 수가 같은 두 idle guard는 함께 시작해 서로를 거부하며, 다른 개수를 쓰면 동률이 깨집니다. + +Refs: #2064, #1814, #1801, #1813, #2061, #2067, #2099, #2065, #2098, #1862, #1870, #1885, #900, #901, #1803.