update(rocm): route top-k with top-p up to vocab 152064 on ROCm

[inureyes]

Summary

Re-measures the joint top-k + top-p rejection vocab cap (REJECTION_JOINT_VOCAB_MAX) on gfx1151. A ROCm build now routes top-k with top-p to the rejection kernel up to vocab 152064; Metal and CUDA keep 32768. With --temperature 0.8 --top-k 40 --top-p 0.95, Llama 3.1 8B decode on this branch goes from 33.15 to 37.73 tok/s (1.14x).

Where the cap lives (for epic #2166 Phase 2): REJECTION_JOINT_VOCAB_MAX in src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp, read only by sampling_rejection_routes and reported by the new bridge accessor sampling_rejection_joint_vocab_max(). A unified per-row sampling step should call sampling_rejection_routes (or the accessor) rather than copy the number. The per-build value is pinned by the_joint_vocab_cap_is_the_measured_crossover in src/lib/mlxcel-core/src/sampling_rejection_tests.rs.

Decision and evidence

Full tables: docs/benchmark_results/rocm-rejection-joint-cap-gfx1151-2026-10-07.md.

• Microbenchmark (rejection_sampling_microbench, three guarded runs, a measurement build with the cap lifted so every joint cell routes): batches 4 and 8 win at every vocabulary (pipelined 1.17x to 2.02x), but the synthetic batch-1 cell does not clear 1.0 above 32768 (0.90x to 0.98x pipelined, 0.92x to 0.96x isolated at 128256 and up). A bf16 rerun reads the same, so the logits dtype is not the cause.
• End-to-end decode, the confirmation the issue's rule calls for at 128256 and above, contradicts that cell. Measurement build, three interleaved guarded pairs each: Meta-Llama-3.1-8B-Instruct-4bit (128256) 32.67 to 37.57 tok/s (1.15x), Qwen2.5-7B-Instruct-4bit (152064) 41.91 to 46.92 tok/s (1.12x), with every kernel run faster than every chain run. Rebased final build: Llama 33.15 to 37.73 tok/s (1.14x, every pair 1.08x to 1.14x).
• Read literally, the issue's rule stops at 32768 because of the synthetic batch-1 cells. I chose 152064, the largest vocabulary measured both ways, because the end-to-end win on the configuration the issue was opened for is well outside the run-to-run spread. Why the synthetic batch-1 cell does not predict real decode was not isolated. Vocabularies above 152064 (gpt-oss 201088, Gemma 262144) are unmeasured and stay on the stock chain.

Changes

• mlx_cxx_bridge.cpp: REJECTION_JOINT_VOCAB_MAX is 152064 under #ifdef MLXCEL_BRIDGE_ROCM_BACKEND and 32768 otherwise. This uses the build-flag pattern of the fused-MoE SGY default (update(rocm): port the fused MoE decode kernels to HIP #2098) and the cfg!(feature = "rocm") defaults of update(rocm): enable fused add-RMSNorm and RoPE-append by default on ROCm #2189, with no runtime backend comparison, so verify-kernel-port-dispatch passes. The gfx1151 tables are in the comment above the constant. The not-routed message names the build's cap.
• New bridge accessor sampling_rejection_joint_vocab_max().
• Tests: the cap is pinned per build; the routing matrix covers 65536, 128256, 151936 and 152064 per build, with 201088 and 262144 declining everywhere; the large-vocabulary decline test picks a vocabulary above the build's cap.
• examples/rejection_sampling_microbench.rs: adds vocab 128256 and 151936 and the --dtype f32|bf16|f16 and --config <label> options, and prints the build's cap. It also fixes the port gate: it checked custom_kernels_available(), which is Metal-or-CUDA only, so it refused to run on ROCm. It now checks sampling_rejection_backend_supported().
• mlxcel-bench-decode --top-k and scripts/bench_decode.sh --top-k (filename tag _k<K>, validated), so the joint filter can be benchmarked end to end.
• Docs: new results page, linked from rocm-samplers-gfx1151-2026-10-05.md. The ROCm sampled-decode row in installation.md now lists which filter combinations route at the Llama 3 and Qwen vocabularies, and the MLXCEL_SAMPLING_REJECTION row in environment-variables.md is updated. Raw CSVs are under benchmarks/.

Verification

All on gfx1151, branch rebased on main at ad844354:

• make -k verify-rocm: OK. 12187 passed, 0 failed, 398 ignored across 161 test binaries; fmt, clippy (-D warnings), the overlay and kernel-dispatch checks, and the smoke run pass.
• cargo test --release --features rocm -p mlxcel-core --lib -- --test-threads=1 sampling_rejection_tests sampling_fixed_key_tests: 34 passed, covering the support, distribution and fixed-key tests.
• Kill-switch tests: cargo test --release --features rocm --test sampling_gumbel_kill_switch --test sampling_rejection_kill_switch pass.
• cargo test --test dead_doc_pointers passes.
• Dispatch logs: on the final build, mlxcel generate logs rejection kernel (batch 1, vocab 128256, ... top_k 40, top_p 0.95 ...) by default and argpartition chain: pinned by MLXCEL_SAMPLING_REJECTION with the switch set.
• Every benchmark ran through scripts/rocm_gpu_guard.sh. Every accepted attempt was clean, and one microbenchmark run was retried after a rejection.

Not verified: Metal and CUDA are not available on this host. Their cap value (32768) is unchanged and the routing code is shared, but the non-ROCm arm of the per-build tests was not run here. Qwen2.5-7B was measured on the measurement build only, not re-run on the rebased build.

Closes #2157