[ROCm] Add portable HIP build for ZhiLight on AMD GPUs

[jeffdaily]

Summary

Adds a ROCm/HIP build path to ZhiLight, which was CUDA-only before this. The portable inference surface -- elementwise/norm/softmax/rope/embedding, the scalar attention decode path, the cuBLASLt GEMM, the int8 reduce kernels, and the fp16/bf16 MoE expert-routing softmax -- now builds with hipcc and runs correctly on AMD GPUs. Everything is gated behind an opt-in USE_HIP; the NVIDIA build path is unchanged.

Build it with:

ROCM_PATH=${ROCM_PATH:-/opt/rocm}
HIPCXX=$ROCM_PATH/lib/llvm/bin/clang++ ZHILIGHT_USE_HIP=1 \
  CMAKE_ARGS="-DUSE_HIP=ON -DCMAKE_HIP_ARCHITECTURES=gfx90a -DCMAKE_HIP_COMPILER=$ROCM_PATH/lib/llvm/bin/clang++ -DCMAKE_HIP_PLATFORM=amd" \
  CMAKE_BUILD_PARALLEL_LEVEL=32 TESTING=0 python setup.py bdist_wheel

Approach

A central compat header (hip_compat/) plus toolkit-named forwarding shims (cuda_runtime.h, cublasLt.h, nccl.h, ...) on a HIP-only include path, so host .cpp resolve the CUDA spellings to their HIP equivalents without editing include lines. CMake gains an opt-in USE_HIP that enables the HIP language, marks the existing .cu as HIP, swaps cuBLAS/cuBLASLt/cuRAND/NCCL for hipBLAS/hipBLASLt/hipRAND/RCCL, and defaults CMAKE_HIP_ARCHITECTURES to gfx90a (overridable for other GPUs with no source edit).

The substantive correctness work is wave64: 3rd/bmengine's warp/block reductions hardcoded a 32-lane warp (32-lane shuffle offsets, 0xFFFFFFFF masks, %32//32, shared[33]). On a 64-lane wavefront these silently reduced only half the lanes (wrong amax/softmax, no crash). They are rewritten wave-agnostically under USE_HIP; the CUDA path keeps its original spelling and is value-identical. The README documents the ROCm build alongside the existing CUDA instructions. The commit body has a file-by-file review order.

Scope

Validated on real GPUs on Linux: gfx90a (MI250X, CDNA2/wave64) and gfx1100 (RDNA3/wave32). Windows is not part of this PR -- ZhiLight's bundled bmengine host runtime uses Linux-only POSIX APIs (execinfo.h, sched.h, pthread.h, sys/syscall.h) that have no Windows equivalent; the GPU kernels themselves compile for AMD on Windows, but the host runtime would need a separate Windows-compat effort.

The NVIDIA tensor-core / Hopper-only fast paths (Marlin/AWQ/FP8/GPTQ quant kernels, the WMMA attention decode, int8-quantized-KV attention, the deep_gemm FP8-block and flash_mla MLA H20 static libs, and flash-attn prefill) are gated off on HIP and throw a clear error or fall back. The unquantized fp16/bf16 dense + MoE path, the fp16/bf16 MoE routing softmax, and the cuBLASLt Int8Linear GEMM are supported; MLA falls back to the generic attention path.

Test plan

Built for gfx90a and ran the in-tree GPU kernel tests (torch reference, no model weights) on a real MI250X GCD; all pass and are deterministic across runs:

HIPCXX=/opt/rocm/lib/llvm/bin/clang++ ZHILIGHT_USE_HIP=1 ENABLE_NCCL_TP=on TESTING=1 \
cmake -S . -B build -GNinja -DUSE_HIP=ON -DCMAKE_HIP_ARCHITECTURES=gfx90a \
  -DCMAKE_HIP_COMPILER=/opt/rocm/lib/llvm/bin/clang++ -DCMAKE_HIP_PLATFORM=amd \
  -DCMAKE_BUILD_TYPE=Release -DWITH_TESTING=ON -DPYTHON_EXECUTABLE=$(which python3)
cmake --build build --target C internals_ -j16

HIP_VISIBLE_DEVICES=1 PYTHONPATH=. python3 -m pytest \
  tests/test_softmax.py tests/test_attn_softmax.py tests/test_rotary_embedding.py \
  tests/test_feedforward.py tests/test_linear.py tests/test_embedding.py \
  tests/test_arthmetic.py tests/test_concat_tensor.py tests/test_index_along_dim.py \
  tests/test_log_prob.py tests/test_attention.py tests/test_lazy_loader.py
# 77 passed, deterministic across runs

The amax/amin tests fail before the wave64 reduce fix and pass after, confirming the root cause. Also built and validated on gfx1100. Toolchain: ROCm 7.2.1, ROCm PyTorch 2.13.

This work was authored with assistance from Claude (Anthropic).