perf(rocm): recalibrate the joint top-k + top-p rejection vocab cap #2157

Part of #1801

Problem / Background

The rejection sampler routes top-k together with top-p only up to vocabulary 32768 (REJECTION_JOINT_VOCAB_MAX). That cap was measured on M1 Ultra: joint top-k + top-p at 32K measured 1.27x to 1.64x, and at 152K 0.71x to 0.83x, because top-k needs more rejection rounds as the vocabulary grows and each round is a full-row sweep on one threadgroup. On ROCm, where #2064 / PR #2100 ported the kernel, the cap was never measured: docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md:26 notes that --top-k 40 --top-p 0.95 on a vocab-128256 model (Llama 3.1) never reaches the kernel, and "Whether the joint cap is right on ROCm was not measured." gfx1151's per-threadgroup bandwidth and launch cost differ from M1 Ultra's, so the crossover can sit elsewhere.

Current Behavior

• src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:5665 constexpr int32_t REJECTION_JOINT_VOCAB_MAX = 32768;, one value for every backend; measurement notes above it (:5622-5664).
• sampling_rejection_routes (:5670-5683) returns !top_k_active || vocab <= REJECTION_JOINT_VOCAB_MAX once top-p is active; rejection_not_routed_reason (:5686-5704) prints the cap.
• src/lib/mlxcel-core/src/sampling_rejection_tests.rs:1023-1030 pins literal cells (32768 routes, 65536 and 152064 do not).
• examples/rejection_sampling_microbench.rs measures the matrix; its vocabularies are VOCABS = [32_768, 65_536, 152_064] (:124), batches [1, 4, 8], and it reports iso_x, pipe_x, rounds and routed per cell.

Proposed Solution

1. Measure on gfx1151. Add 128_256 to VOCABS (Llama 3 vocabulary, the case that prompted this) and run the microbench on a ROCm build for the joint config. Read pipe_x, not iso_x (the harness explains why). Also measure end-to-end decode on Meta-Llama-3.1-8B-Instruct-4bit with --temperature 0.8 --top-k 40 --top-p 0.95, a build with the candidate ROCm cap against the same build with MLXCEL_SAMPLING_REJECTION=0, medians of 3 through scripts/rocm_gpu_guard.sh.
2. Decision rule for the ROCm value: the largest measured vocabulary V such that every batch cell at V and at every measured vocabulary below V has pipe_x >= 1.0, confirmed by the end-to-end run when V >= 128256. If 32768 itself loses on ROCm, the ROCm value is 0 (joint never routes). Unmeasured vocabularies are never included, as the existing comment requires.
3. Make the cap per build the way PR update(rocm): port the fused MoE decode kernels to HIP #2098 made the MLXCEL_FUSED_MOE_SGY default per build (mlx_cxx_kernels.cpp:3154-3158): #ifdef MLXCEL_BRIDGE_ROCM_BACKEND selects the ROCm value, #else keeps 32768. A ROCm build has no Metal or CUDA backend, so no runtime backend comparison is needed, and scripts/ci/check_kernel_port_dispatch.py stays satisfied. Add the gfx1151 table to the comment block above the constant.
4. Expose the active value through the bridge (sampling_rejection_joint_vocab_max() -> i32) and rewrite the tests at sampling_rejection_tests.rs:1023-1030 to pin the per-build value and the cells either side of it, so the ROCm and non-ROCm builds each assert their own measured cells.

Implementation Notes

• Reuse: the existing microbench, routing function and reason string; only the constant, the vocab list and the tests change.
• Edge cases: a top-k that cannot bind (top_k >= vocab) still routes regardless of the cap, unchanged. min-p rides along unchanged.
• Constraints: Metal and CUDA routing must not change; make verify on macOS shows the same pinned cells as before.

Acceptance Criteria

• docs/benchmark_results/rocm-samplers-gfx1151-2026-10-05.md (or a new dated page linked from it) holds the gfx1151 joint top-k + top-p table at vocab {32768, 65536, 128256, 152064} x batch {1, 4, 8} with pipe_x, iso_x, rounds, plus the end-to-end decode result, and states the chosen value by the rule above.
• REJECTION_JOINT_VOCAB_MAX is that value on ROCm builds and 32768 elsewhere; rejection_not_routed_reason prints the active value.
• Tests pin the per-build cap and pass on gfx1151 and on a non-ROCm build.
• docs/installation.md's ROCm sampled-decode row states which filter combinations route at Llama 3 vocabulary.

Verification

cargo run --release --features rocm --example rejection_sampling_microbench -- --csv /tmp/rejection-gfx1151.csv
cargo test --release --features rocm -p mlxcel-core --lib -- --test-threads=1 sampling_rejection_tests
make verify-rocm