From 3a1261f86160e94d470ae07b5c77860681bc739d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 7 Oct 2026 19:38:17 +0900 Subject: [PATCH 1/6] update(sampling): measure the joint rejection cap on gfx1151 and pin it Re-measure the top-k + top-p rejection vocab cap (#2157) on gfx1151 and keep 32768, which is the same crossover M1 Ultra measured. Batch 1 does not clear 1.0 above 32768 (0.98x pipelined at 65536, 0.96x isolated at 128256), so by the issue's rule the ROCm value equals the default and the cap stays one constant with no build-flag branch. - `REJECTION_JOINT_VOCAB_MAX` in `mlx_cxx_bridge.cpp` carries the gfx1151 table in its comment and the rule for moving it per build (`#ifdef MLXCEL_BRIDGE_ROCM_BACKEND`, never a runtime backend check). - New bridge accessor `sampling_rejection_joint_vocab_max()`. `the_joint_vocab_cap_is_the_measured_crossover` pins the value and checks the policy turns exactly there; the routing matrix test adds the 128256 cells. - The not-routed reason string now names the cap and both measurements instead of the M1-only ratio. - `rejection_sampling_microbench` gates on `sampling_rejection_backend_supported()` instead of `custom_kernels_available()`, which is Metal-or-CUDA only and kept the harness from running on ROCm. It adds vocab 128256 and 151936 and prints the build's cap. - `mlxcel-bench-decode --top-k` and `bench_decode.sh --top-k` (filename tag `_k`) so the joint filter can be benchmarked end to end. --- examples/rejection_sampling_microbench.rs | 28 +++++++----- scripts/bench_decode.sh | 11 +++++ src/bin/bench_decode.rs | 10 ++++- src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp | 41 ++++++++++++++--- src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h | 4 ++ src/lib/mlxcel-core/src/lib.rs | 7 +++ .../src/sampling_rejection_tests.rs | 44 ++++++++++++++++++- 7 files changed, 126 insertions(+), 19 deletions(-) diff --git a/examples/rejection_sampling_microbench.rs b/examples/rejection_sampling_microbench.rs index d55554b3d..4841df88b 100644 --- a/examples/rejection_sampling_microbench.rs +++ b/examples/rejection_sampling_microbench.rs @@ -92,6 +92,8 @@ //! --example rejection_sampling_microbench //! Run (CUDA): //! cargo run --release --features cuda --example rejection_sampling_microbench +//! Run (ROCm): +//! cargo run --release --features rocm --example rejection_sampling_microbench //! //! Options: //! --iters N timed repetitions per point (default 200) @@ -105,12 +107,13 @@ use std::fmt::Write as _; use std::time::{Duration, Instant}; use mlxcel_core::{ - MlxArray, UniquePtr, array_to_raw_bytes, async_eval_pair, custom_kernels_available, eval, - from_slice_f32, fused_sample, fused_sample_categorical, fused_sample_rejection, - gpu_backend_available, matmul, random_seed, rejection_cap_overflow_launches, - rejection_cap_overflow_rows, reset_sampling_dispatch, sampling_dispatch_drain_pending, - sampling_dispatch_recorded_report, sampling_rejection_available, sampling_rejection_max_rounds, - sampling_rejection_probe, sampling_rejection_routes, synchronize_default, + MlxArray, UniquePtr, array_to_raw_bytes, async_eval_pair, eval, from_slice_f32, fused_sample, + fused_sample_categorical, fused_sample_rejection, gpu_backend_available, matmul, random_seed, + rejection_cap_overflow_launches, rejection_cap_overflow_rows, reset_sampling_dispatch, + sampling_dispatch_drain_pending, sampling_dispatch_recorded_report, + sampling_rejection_available, sampling_rejection_backend_supported, + sampling_rejection_joint_vocab_max, sampling_rejection_max_rounds, sampling_rejection_probe, + sampling_rejection_routes, synchronize_default, }; /// Target duration for the synthetic forward in the pipelined mode. Roughly a @@ -121,7 +124,9 @@ const PIPELINE_FORWARD_TARGET_US: f64 = 2000.0; /// Side length of the synthetic forward's square matmul. const FORWARD_DIM: i32 = 1024; -const VOCABS: [i32; 3] = [32_768, 65_536, 152_064]; +/// 128256 (Llama 3) and 151936 (Qwen3) were added for the ROCm joint-cap +/// measurement (#2157); 152064 is Qwen 2.5's. +const VOCABS: [i32; 5] = [32_768, 65_536, 128_256, 151_936, 152_064]; const BATCHES: [i32; 3] = [1, 4, 8]; const TEMPERATURE: f32 = 1.0; @@ -355,7 +360,9 @@ fn main() { // port predicate alone, deliberately not `sampling_rejection_available()`, // because the whole point of the isolated arm is to measure the kernel with // routing switched off, and that predicate folds in the routing kill switch. - if !custom_kernels_available() { + // `custom_kernels_available()` is Metal-or-CUDA by definition, so it must + // not be the gate: the ROCm port (#2064) would never be measured (#2157). + if !sampling_rejection_backend_supported() { eprintln!( "This GPU backend has no rejection-sampling kernel port; there is nothing to measure." ); @@ -479,8 +486,9 @@ fn main() { ); println!( "routing policy: the kernel replaces a sort, so it is routed only where the stock chain \ - sorts (top-p active), and top-k + top-p only at vocab <= 32768. A routed=no row is a \ - path production does not take." + sorts (top-p active), and top-k + top-p only at vocab <= {} on this build. A routed=no \ + row is a path production does not take.", + sampling_rejection_joint_vocab_max() ); println!( "read pipe_x, not iso_x: iso synchronizes around every iteration and is structurally \ diff --git a/scripts/bench_decode.sh b/scripts/bench_decode.sh index 44c8de839..8ae951cc9 100755 --- a/scripts/bench_decode.sh +++ b/scripts/bench_decode.sh @@ -164,6 +164,7 @@ SUFFIX="" # record either, so a non-greedy run also tags the auto-generated filename. TEMPERATURE="" TOP_P="" +TOP_K="" 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 @@ -866,6 +867,9 @@ Options: --top-p P Nucleus threshold, effective with a positive --temperature (default: the runner's 1.0, off). Tags the filename with `_p

`. + --top-k K Top-k cutoff, effective with a positive --temperature + (default: the runner's 0, off). Tags the filename with + `_k`. --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- @@ -959,6 +963,7 @@ default_output_path() { name="${name}_${DATE}" [[ -n "$TEMPERATURE" ]] && name="${name}_t${TEMPERATURE}" [[ -n "$TOP_P" ]] && name="${name}_p${TOP_P}" + [[ -n "$TOP_K" ]] && name="${name}_k${TOP_K}" if [[ -n "$SUFFIX" ]]; then name="${name}_${SUFFIX}" fi @@ -1084,6 +1089,7 @@ bench_one() { fi [[ -n "$TEMPERATURE" ]] && extra_args+=(--temperature "$TEMPERATURE") [[ -n "$TOP_P" ]] && extra_args+=(--top-p "$TOP_P") + [[ -n "$TOP_K" ]] && extra_args+=(--top-k "$TOP_K") >&2 printf '>>> [bench] %s (same-process warmup=%s) ...\n' "$model_name" "$WARMUP_TOKENS" local raw rc=0 @@ -1175,6 +1181,7 @@ while [[ $# -gt 0 ]]; do --suffix) SUFFIX="$2"; shift 2 ;; --temperature) TEMPERATURE="$2"; shift 2 ;; --top-p) TOP_P="$2"; shift 2 ;; + --top-k) TOP_K="$2"; shift 2 ;; --cooldown) COOLDOWN_SECS="$2"; shift 2 ;; --big-cooldown) BIG_MODEL_COOLDOWN_SECS="$2"; shift 2 ;; --big-threshold-gb) @@ -1199,6 +1206,10 @@ for _sampling_opt in "--temperature:$TEMPERATURE" "--top-p:$TOP_P"; do exit 1 fi done +if [[ -n "$TOP_K" && ! "$TOP_K" =~ ^[0-9]+$ ]]; then + echo "Error: --top-k takes a non-negative integer, got '$TOP_K'" >&2 + exit 1 +fi if [[ -z "$MODEL_ARG" ]]; then echo "Error: model path or 'all' required" >&2 diff --git a/src/bin/bench_decode.rs b/src/bin/bench_decode.rs index 35dce8e02..7f70441ba 100644 --- a/src/bin/bench_decode.rs +++ b/src/bin/bench_decode.rs @@ -93,6 +93,12 @@ struct Args { #[arg(long, default_value_t = 1.0)] top_p: f32, + /// Top-k cutoff for both passes. 0 disables it. Only has an effect with a + /// positive `--temperature`; with `--top-p` it measures the joint filter, + /// whose rejection-kernel routing is capped by vocabulary (#2157). + #[arg(long, default_value_t = 0)] + top_k: i32, + /// Generated tokens in the warmup pass. #[arg(long, default_value_t = 20)] warmup_tokens: usize, @@ -284,11 +290,12 @@ fn sampling_config( model: &LoadedModel, ignore_eos: bool, temperature: f32, + top_k: i32, top_p: f32, ) -> SamplingConfig { let mut config = build_sampling_config(ResolvedSamplingParams { temperature, - top_k: 0, + top_k, top_p, min_p: 0.0, seed: None, @@ -517,6 +524,7 @@ fn main() -> Result<()> { &model, args.ignore_eos, args.temperature, + args.top_k, args.top_p, ); diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index 905e26298..fe1a3b61d 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -5790,12 +5790,37 @@ static bool gumbel_sample_applies( // Vocabulary ceiling for routing top-k and top-p together. // -// 32768 is measured (1.27x - 1.64x across three repetitions); 152064 is measured -// as a loss (0.71x - 0.83x); 65536 has not been measured for this combination and -// is therefore excluded rather than interpolated. Raise it only against a -// measurement, never to widen coverage. +// M1 Ultra (#901): 32768 is measured (1.27x - 1.64x across three repetitions); +// 152064 is measured as a loss (0.71x - 0.83x); 65536 has not been measured for +// this combination and is therefore excluded rather than interpolated. +// +// gfx1151 (#2157, docs/benchmark_results/rocm-rejection-joint-cap-gfx1151-2026-10-07.md): +// re-measured with the cap lifted on a measurement build, so `fused_sample` +// routed every cell. Medians of three runs, top-k 40 + top-p 0.9, pipe_x by +// batch {1, 4, 8} (iso_x for batch 1 in brackets): +// +// vocab 32768 1.00 1.17 1.44 [2.23] +// vocab 65536 0.98 1.18 1.64 [1.59] +// vocab 128256 0.98 1.28 1.80 [0.96] +// vocab 151936 0.90 1.36 1.90 [0.92] +// vocab 152064 0.93 1.29 2.02 [0.92] +// +// Batches 4 and 8 win at every vocabulary, but batch 1, the single-stream +// decode shape, does not clear 1.0 above 32768, and at 128256 and up the kernel +// is slower in isolation too, because the round count (4 to 6) times a +// single-threadgroup sweep outgrows the chain's whole-GPU `argpartition`. By the +// rule the issue set (the largest vocabulary where every batch cell at and below +// it has pipe_x >= 1.0), the ROCm value is 32768, the same as Metal's, so the +// cap stays one constant for every backend. A backend that measures a different +// crossover gets a build-flag branch here (`#ifdef MLXCEL_BRIDGE_ROCM_BACKEND`, +// as the SGY default in `mlx_cxx_kernels.cpp` does), never a runtime backend +// comparison. Raise it only against a measurement, never to widen coverage. constexpr int32_t REJECTION_JOINT_VOCAB_MAX = 32768; +int32_t sampling_rejection_joint_vocab_max() { + return REJECTION_JOINT_VOCAB_MAX; +} + // Pure routing policy: would `fused_sample` send this configuration to the // rejection kernel, ignoring backend support and the env switch? Exposed so a // test can pin every measured cell of the matrix above without a GPU. @@ -5829,8 +5854,12 @@ static std::string rejection_not_routed_reason( "measured slower: top-k alone 0.31x-0.97x, min-p alone " "0.47x-0.88x"; } else { - os << "top-k and top-p together measured 0.71x-0.83x above vocab " - << REJECTION_JOINT_VOCAB_MAX << ", and this row has vocab " << vocab; + os << "top-k and top-p together route only up to vocab " + << REJECTION_JOINT_VOCAB_MAX + << ", the largest vocabulary measured to win at every batch size " + "(M1 Ultra 0.71x-0.83x at 152064; gfx1151 0.90x-0.98x at batch 1 " + "from 65536 up), and this row has vocab " + << vocab; } os << "; top_k " << top_k << ", top_p " << top_p << ", min_p " << min_p << ")"; diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 68385fb5f..8f1122b3c 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -1630,6 +1630,10 @@ bool sampling_rejection_routes( float min_p ); +// The vocabulary ceiling `sampling_rejection_routes` applies to top-k and +// top-p together on this build (`REJECTION_JOINT_VOCAB_MAX`, #2157). +int32_t sampling_rejection_joint_vocab_max(); + // Threads per threadgroup the rejection kernel launches with. int32_t sampling_rejection_threadgroup_size(); diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index dfc5f06be..8a6c45ef9 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -2666,6 +2666,13 @@ mod ffi { /// Pure host arithmetic, so it answers on a CPU-only build too. fn sampling_rejection_routes(vocab: i32, top_k: i32, top_p: f32, min_p: f32) -> bool; + /// The vocabulary ceiling [`sampling_rejection_routes`] applies to + /// top-k and top-p together on this build. A compile-time value + /// (`REJECTION_JOINT_VOCAB_MAX` in `mlx_cxx_bridge.cpp`): ROCm builds + /// carry the gfx1151 measurement, every other build the M1 Ultra one + /// (#2157). + fn sampling_rejection_joint_vocab_max() -> i32; + /// Threads per threadgroup the rejection kernel launches with. The /// determinism argument rests on this being fixed, so it is exposed. fn sampling_rejection_threadgroup_size() -> i32; diff --git a/src/lib/mlxcel-core/src/sampling_rejection_tests.rs b/src/lib/mlxcel-core/src/sampling_rejection_tests.rs index d2744a881..96aa2d169 100644 --- a/src/lib/mlxcel-core/src/sampling_rejection_tests.rs +++ b/src/lib/mlxcel-core/src/sampling_rejection_tests.rs @@ -983,6 +983,12 @@ fn a_starved_round_cap_falls_back_and_the_event_is_counted() { /// | min-p alone | 0.47x - 0.88x at 152K | no | /// | top-k + top-p | 1.27x - 1.64x at 32K, 0.71x - 0.83x at 152K | only at vocab <= 32768 | /// +/// gfx1151 re-measured the joint row (#2157, vocab {32K, 64K, 128256, 151936, +/// 152064}): batches 4 and 8 win everywhere, but batch 1 does not clear 1.0 +/// above 32768 (pipelined 0.90x - 0.98x), so ROCm keeps the same cap. The cap is +/// one constant on every build; [`the_joint_vocab_cap_is_the_measured_crossover`] +/// pins it. +/// /// The kernel replaces a sort, so it wins exactly where the stock chain sorts, /// which is when top-p is active. The `rounds` column of the microbenchmark /// shows why the joint case degrades with vocabulary: top-p accepts in one or @@ -1021,13 +1027,16 @@ fn the_routing_policy_matches_the_measured_matrix() { } // The joint case is capped at the vocabulary where it was measured to win. - // 65536 is excluded because it has not been measured for this combination, - // not because it was measured to lose. + // On M1 Ultra 65536 is excluded because it has not been measured for this + // combination; on gfx1151 it was measured and batch 1 did not win (#2157). + // 128256 is the Llama 3 vocabulary that prompted the ROCm re-measurement. assert!(sampling_rejection_routes(4096, 40, 0.9, 0.0)); assert!(sampling_rejection_routes(32_768, 40, 0.9, 0.0)); assert!(!sampling_rejection_routes(65_536, 40, 0.9, 0.0)); + assert!(!sampling_rejection_routes(128_256, 40, 0.9, 0.0)); assert!(!sampling_rejection_routes(152_064, 40, 0.9, 0.0)); assert!(sampling_rejection_routes(32_768, 40, 0.9, 0.05)); + assert!(!sampling_rejection_routes(128_256, 40, 0.95, 0.0)); assert!(!sampling_rejection_routes(152_064, 40, 0.9, 0.05)); // A top-k that cannot bind is not a top-k: it leaves the chain's cost @@ -1038,6 +1047,37 @@ fn the_routing_policy_matches_the_measured_matrix() { assert!(sampling_rejection_routes(152_064, 1, 0.9, 0.0)); } +/// The joint top-k + top-p cap is the measured crossover on this build, and the +/// routing policy turns exactly at it (#2157). +/// +/// Both measurements put the crossover at 32768: M1 Ultra (#901) and gfx1151, +/// where batch 1 read 0.98x pipelined at 65536 and 0.96x isolated at 128256. +/// The value is a compile-time constant, so a backend that measures a different +/// crossover changes it by build flag; this test is where that build's value +/// gets pinned. Pure host arithmetic, so it runs on a CPU-only build too. +#[test] +fn the_joint_vocab_cap_is_the_measured_crossover() { + const MEASURED_CAP: i32 = 32_768; + let cap = sampling_rejection_joint_vocab_max(); + assert_eq!( + cap, + MEASURED_CAP, + "REJECTION_JOINT_VOCAB_MAX moved without a measurement (rocm build: {})", + cfg!(feature = "rocm") + ); + + // The policy reads the same constant the accessor reports: the cells on + // either side of it decide differently. + assert!(sampling_rejection_routes(cap, 40, 0.9, 0.0)); + assert!(!sampling_rejection_routes(cap + 1, 40, 0.9, 0.0)); + assert!(sampling_rejection_routes(cap, 40, 0.95, 0.02)); + assert!(!sampling_rejection_routes(cap + 1, 40, 0.95, 0.02)); + // Above the cap top-p alone still routes, and a top-k that cannot bind + // does not bring the cap in. + assert!(sampling_rejection_routes(cap + 1, 0, 0.9, 0.0)); + assert!(sampling_rejection_routes(cap + 1, cap + 1, 0.9, 0.0)); +} + /// Which path `fused_sample` took for one configuration, decided by comparing /// its stream against both arms at one seed. /// From defe7065e3bb419e61e713a889204d0c0f46e6dc Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 8 Oct 2026 12:05:30 +0900 Subject: [PATCH 2/6] update(sampling): add bf16 and per-config runs to the rejection microbench - `rejection_sampling_microbench --dtype f32|bf16|f16` casts the synthetic logits, since a bf16 checkpoint hands the sampler bf16 logits and the stock chain's sort cost depends on the dtype; `--config