fix(rocm): refuse wave32-only kernel ports on wave64 devices at runtime #2147

Part of #1801

Problem / Background

Every shuffle-based HIP port of an mlxcel fused kernel protects itself against a 64-lane wavefront with a preprocessor guard, as #1814 requires ("a wave64 target fails to compile rather than returning a plausible wrong answer"):

#if defined(__AMDGCN_WAVEFRONT_SIZE__) && __AMDGCN_WAVEFRONT_SIZE__ != 32
#error "... assumes a 32-lane wavefront"
#endif
#if defined(__AMDGCN_WAVEFRONT_SIZE) && __AMDGCN_WAVEFRONT_SIZE != 32
#error ...
#endif

With ROCm 10's AMD clang 23 (HIP 7.15) neither macro is defined for gfx1151, gfx942 or gfx90a (checked with hipcc -E -dM in #2065 and #2067), so the guard never fires on any current toolchain. The ports' own comments already say so (src/lib/mlx-cpp/turbo/fused_norm_hip.h:65-73, paged_attention_hip.h:47-56, src/lib/mlxcel-core/src/fused_moe_parity_tests.rs:559-562). The explicit shuffle width of 32 keeps each butterfly inside 32 lanes, which should be correct on wave64 if nothing else in the kernel assumes a 32-wide wavefront (lane and warp indices from threadIdx.x % 32 / / 32, per-warp shared-memory slots, warp-synchronous steps without a barrier). None of that has been run on a wave64 device (#2107: "Wave64 (CDNA) untested"). On CDNA (MI200 gfx90a, MI300 gfx942), which docs/installation.md:384 says compiles, every one of these ports would run unvalidated.

Current Behavior

Guarded HIP sources (all compiled through fast::hip_kernel): BITLINEAR_HIP_SOURCE (src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:159, #1870), XIELU_HIP_SOURCE (:432), SSM_HIP_SOURCE (:735, #2099), MAMBA1_SCAN_HIP_SOURCE (:2107), MOE_GATEUP_HIP_SOURCE (:2686), MOE_DOWN_HIP_SOURCE (:2823, both #2098), MOE_FC1_RELU2_HIP_SOURCE (:3003), ADD3_LN_HIP_SOURCE (:3768), FUSED_ADD_RMS_NORM_HIP_SOURCE (fused_norm_hip.h:79, #2107), and the paged attention v1 decode and v2 partial bodies (paged_attention_hip.h:66,205, #2103). The RoPE append, paged merge and both sampler bodies (#2100) have no lane-level operation and deliberately carry no guard (sampling_gumbel_hip.h:44-50, sampling_rejection_hip.h:42-48).

Port choice goes through one helper: port_for() in src/lib/mlx-cpp/turbo/kernel_port.cpp:30 maps the resolved GpuKernelBackend to a KernelPorts entry, and both has_kernel_port() (every *_available() predicate, hence every Rust gate) and select_kernel_port() read it. The device's real wavefront width is known (hipDeviceProp_t::warpSize, logged at bind in src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/device.cpp:120-139) but nothing on mlxcel's side reads it. Device::warp_size() is not usable for this decision because MLX_ROCM_FORCE_WARP_SIZE overrides it (device.cpp:146-158).

Proposed Solution

Refuse wave32-only ports on a device whose wavefront is not 32, at port selection, so callers take the graph fallback they already have.

1. Overlay: add MLX_API int device_warp_size(); to patches-rocm/mlx/backend/rocm/rocm.h, implemented in rocm.cpp with hipGetDevice plus hipDeviceGetAttribute(&w, hipDeviceAttributeWarpSize, dev), returning the hardware value (never the MLX_ROCM_FORCE_WARP_SIZE override) and 0 on any HIP error. Record it as a new LOCAL_FIXES.md entry.
2. KernelPorts (src/lib/mlx-cpp/turbo/kernel_port.h:62) gains bool rocm_any_wave_size = false; after rocm. A null default means wave32-only, so a new port is refused on wave64 unless its author says otherwise.
3. In port_for, the GpuKernelBackend::Rocm case returns ports.rocm only when ports.rocm_any_wave_size || rocm_device_is_wave32(). rocm_device_is_wave32() lives in gpu_backend.cpp under MLXCEL_BRIDGE_ROCM_BACKEND, caches mlx::core::rocm::device_warp_size() == 32 once per process, and treats 0 (query failed) as not wave32. Factor the decision into a pure bool rocm_port_allowed(bool any_wave_size, int warp_size) so it can be tested without a wave64 device. select_kernel_port's refusal message names the wave size when that is the reason.
4. Set .rocm_any_wave_size = true on exactly fused_rope_ports() (fused_rope_append.cpp:411), paged_merge_ports() (paged_attention_v2_merge.cpp:222), gumbel_ports() (sampling.cpp:393) and rejection_ports() (sampling_rejection.cpp:783), each with a one-line comment citing the existing no-lane-operation note. The other eleven tables (bitlinear, xielu, ssm, mamba1_scan, moe_gateup, moe_down, moe_fc1_relu2, add3_layer_norm, fused_norm, paged_attention, paged_v2_partial) stay wave32-only.
5. Keep the #error guards (harmless, and live on a compiler that defines the macro) but rewrite the shared comments so they name the host-side check as the protection.

Rejected: a compile-time check through warpSize (static_assert(warpSize == 32) does not compile in HIP, as the existing comments record), and a per-launcher check (the reason select_kernel_port exists).

Scope

In scope: kernel_port.h/.cpp, gpu_backend.h/.cpp, the four table edits, overlay rocm.h/.cpp plus LOCAL_FIXES.md, guard comments, a test, docs. Out of scope: validating any port on wave64 and flipping it to rocm_any_wave_size; that needs a CDNA host and is per kernel. Multi-GPU hosts mixing wave sizes (the cached answer is for the current device at first use); note it in the comment.

Acceptance Criteria

• On a wave64 device every wave32-only *_available() predicate returns false and inference runs on the graph paths; on gfx1151 nothing changes (same ports taken, same logits).
• A test pins rocm_port_allowed for (false, 32) = true, (false, 64) = false, (true, 64) = true, (false, 0) = false, exposed through the cxx bridge and run in the default (CPU-capable) test set.
• A ROCm test on gfx1151 asserts device_warp_size() == 32 and that MLX_ROCM_FORCE_WARP_SIZE=64 in a child process does not change it.
• A source test fails if any table with rocm_any_wave_size = true points at a HIP source containing __shfl.
• docs/installation.md ROCm section states that on CDNA (wave64) the shuffle-based fused kernels fall back to the graph until validated.

Verification

cargo test --release --features rocm -p mlxcel-core --lib -- --test-threads=1 kernel_port
make verify-rocm verify-rocm-overlay verify-kernel-port-dispatch

Manual on gfx1151: Qwen3-30B-A3B-4bit and granite-4.0-h-tiny-4bit decode unchanged (fused MoE and SSM ports still taken, visible in rocprofv3 kernel names).

Implementation note

Implemented in #2179. Checking the tables against the current code found a fifth body with no lane-level operation: XIELU_HIP_SOURCE is purely elementwise, so xielu_ports() is also marked rocm_any_wave_size and its inert guard is dropped, as the samplers and RoPE have none. The other ten tables stay wave32-only. The wave64 refusal is proven through a test-only seam (set_rocm_port_warp_size_for_tests), since no wave64 device is available.