rocm: fix native extensions, kernel JIT builds, and Triton PTX fallbacks for gfx1150

[skywalk1411]

Summary

Gets FreeToken booting and serving on native ROCm (tested on a Ryzen AI 9 HX 470 / Radeon 890M, gfx1150). Verified end to end: server boot, weight load, KV cache allocation, CUDA graph capture at bs=1/2/4, and real chat completions against Qwen/Qwen3-8B (bf16, triton attention backend).

This covers the same ground as #137 (build system, kernel JIT flags, launch_pdl, the decode-attention WMMA tile floor) and independently arrived at the same fixes there, plus two more that block a working boot on this GPU:

• The actual root cause of the silent crash during warmup: kernel/triton/activation.py's _fast_tanh/_fast_ex2 inline raw NVIDIA PTX text (tanh.approx.f32, ex2.approx.f32) via tl.inline_asm_elementwise. HIP's inline-asm path doesn't reject foreign PTX at parse time — it fails much later in register allocation with a generic, misleading diagnostic (error: couldn't allocate output register for constraint 'f') that looks like a matrix-core or register-pressure issue. This is called from every MLP layer's SiLU, so it reproduces on any model, deterministically, and took the longest to isolate (needed py-spy live process inspection to catch the real call site, since the crash bypasses Python's exception handling entirely). Routed through libdevice.tanh / tl.exp2 on HIP instead.
• The split extend/prefill attention kernel (_extend_attention_split_kernel) runs both the cached-KV and newly-computed-KV loops live at once, which is register-heavier than the plain extend kernel and exhausts this GPU's VGPR file at the CUDA-tuned 128x64 tile. Shrinks to 64x32 under HIP.

Every change is gated on HIP detection (torch.version.hip / ROCM_HOME / __HIP_PLATFORM_AMD__) at build or run time; the CUDA path is unchanged.

Details by file

Build system (setup.py, new kernel/csrc/hip_compat.h)

• _pinned_tensor/_cpu_moe link against HIP instead of cudart when ROCM_HOME is present.
• hip_compat.h aliases the CUDA Runtime API calls those two files use onto HIP, including CUDART_CB (undefined under HIP, which otherwise corrupts the surrounding declaration's parse).
• The pip-vendored ROCm SDK ships versioned sonames (libamdhip64.so.7) with no bare .so dev symlink, so link the exact file via -l:; the dynamic linker dedupes by SONAME at runtime against whatever libamdhip64 torch itself already loaded.

Shared kernel header (kernel/csrc/include/freetoken/utils.cuh)

• Explicit HIP runtime include + CUDA Runtime API aliases (nvcc pulls in cuda_runtime.h implicitly for .cu files; hipcc does not).
• __grid_constant__ has no HIP equivalent; falls back to an ordinary by-value kernel parameter.
• LaunchKernel has no HIP path for cudaLaunchKernelEx/cudaLaunchConfig_t (that API only exists to carry Hopper PDL attributes) — added a HIP variant that launches via plain triple-chevron syntax, with with_attr() as a no-op.
• The griddepcontrol PDL asm is now unconditionally a no-op under HIP, not just when kUsePDL is false, so a stray HIP-side call can't try to assemble Hopper-only PTX.

Kernel JIT (kernel/utils.py): drop --expt-relaxed-constexpr on HIP; hipcc/clang rejects it outright.

Triton kernels

• norm.py, activation.py: launch_pdl is a CUDA-Hopper-only kwarg; the AMD arg-packer raises KeyError on it even when passed as False, so it's only included when pdl is actually true.
• attention.py: decode kernel's GQA head-tile floors at 16 under HIP (RDNA WMMA has no instruction below M=16); falls back to broadcast-multiply-reduce instead of tl.dot for that tile as a second-layer guard. Extend/prefill split kernel tile shrinks 128x64 → 64x32 under HIP (see above).
• activation.py: _fast_tanh/_fast_ex2 PTX → libdevice.tanh/tl.exp2 on HIP (see above).

pyproject.toml: loosen the torch/triton version ceilings so ROCm builds (which carry a local version segment such as +rocm7.14.0...) can satisfy them.

Test plan

• pip install -e . builds _pinned_tensor/_cpu_moe cleanly against ROCm
• ft serve --model Qwen/Qwen3-8B boots, loads weights, allocates KV cache
• CUDA graph capture completes at bs=1/2/4
• /v1/chat/completions returns correct, coherent output
• MoE model end to end: Qwen/Qwen3.6-35B-A3B-FP8 (256 experts/layer x 40 layers, 3B active) with --moe-backend offload --moe-cache-size 2560 (25% of experts resident, LRU-evicting the rest from host RAM) -- boots, captures CUDA graphs, serves correct chat completions, ft ctl cache confirms the pool is live at the requested size
• GGUF end to end: google/gemma-4-26B-A4B-it-qat-q4_0-gguf (native GGUF dequant kernels, MoE offload) -- boots and serves correct chat completions. See the third commit for the source fixes and a documented environment gotcha (pip ROCm nightly ships no thrust/rocprim headers; the obvious apt fix pulls in a second, conflicting HIP header set)
• Not tested: multi-GPU (pynccl)

[skywalk1411]

Pushed a second commit — two more real bugs, found by actually running a MoE model end to end (the first commit only covered the dense Qwen3-8B boot path):

1. e4m3_native() (host-side fp8 capability check) disagrees with Triton's own compile-time check on HIP. It compares torch.cuda.get_device_capability() against (8, 9) — on ROCm that call returns the GPU's RDNA generation number ((11, 5) for gfx1150), and (11, 5) >= (8, 9) is True by plain tuple comparison. So the host side thinks native fp8 is available while the kernel (compiled via Triton's target_info.cuda_capability_geq, which correctly checks backend != "cuda" first) took the emulated path — a fp8/uint8 type mismatch. Fixed by checking torch.version.hip first.

2. fast_index_copy.cuh (the offload cache's fast host→device expert-copy kernel — only reachable once a MoE model actually streams experts with --moe-backend offload): missing HIP aliases for a few CUDA Runtime calls, and the same raw-PTX streaming-load/store pattern already fixed elsewhere in this file family, just not caught until this path ran for the first time.

Verified against Qwen/Qwen3.6-35B-A3B-FP8 (256 experts/layer × 40 layers) with --moe-backend offload --moe-cache-size 2560 — 25% of experts resident, the rest genuinely LRU-evicted from host RAM on miss. Boots, captures CUDA graphs, serves correct completions. Test plan checklist updated in the PR description.

[skywalk1411]

Pushed a third commit — GGUF now works too (google/gemma-4-26B-A4B-it-qat-q4_0-gguf, native GGUF dequant kernels + MoE offload, verified with real chat completions).

Two small source fixes, same pattern as before:

• kernel/gguf.py builds via a third JIT mechanism (torch.utils.cpp_extension.load, distinct from both setup.py's CppExtension and the tvm-ffi JIT the rest of kernel/ uses) — same nvcc-only-flag problem (--expt-relaxed-constexpr), plus a -ccbin/CXX-forcing block that exists only to work around an nvcc+libtorch-headers mismatch that doesn't apply under hipcc.
• kernel/csrc/gguf/dispatch.h: the donor's SGLANG_SHFL_XOR_SYNC macros pass a CUDA-style 32-bit warp mask into __shfl_xor_sync; HIP's header statically asserts the mask must be 64 bits unconditionally. Widened on HIP only.

Worth calling out: this JIT path uses torch's own ROCm auto-hipify (torch.utils.cpp_extension.load's built-in CUDA→HIP translation), and it worked correctly with zero manual porting of the actual .cu/.cuh kernel sources — the gguf_kernel.cu/mmq.cuh/vecdotq.cuh/dequantize.cuh dequant/GEMM kernels translated automatically. Added a .gitignore entry for the translated copies it writes next to the sources (regenerated every build, shouldn't be tracked).

One thing that isn't in the diff, flagged for the next person on this distribution: the pip ROCm nightly build (rocm.nightlies.amd.com) ships no thrust/rocprim headers, which torch's own extension headers pull in transitively. The obvious fix — apt install librocthrust-dev — pulls in libamdhip64-dev as a dependency, which drops a second, older HIP header set into /usr/include/hip that silently wins over the correct pip-bundled ones for a plain -I (though not -isystem) — took a clang++ -v + minimal reproducer to pin down. No source fix needed; the workaround is extracting just the needed headers (dpkg -x, no install) into the pip package's own include dir and removing the conflicting system packages, plus pointing ROCM_PATH/HIP_PATH/HIP_DEVICE_LIB_PATH at the pip package for this JIT path's device-bitcode lookup.

Test plan checklist updated. Only multi-GPU (pynccl) is left untested.

[skywalk1411]

Pushed a fourth commit — performance follow-up, no new hardware coverage this time.

Ran ft bench bw for real hardware/kernel bandwidth numbers (CPU STREAM ~71 GB/s, PCIe-equivalent linear copy ~40 GB/s on this unified-memory APU), which both exercises and validates the CPU MoE executor and the offload cache's gather kernel under real load, and gets --moe-backend auto picking hybrid correctly for bf16/nvfp4 experts on this box going forward (was previously falling back to offload with no calibration data).

That prompted a second look at two defensive fixes from the earlier commits, both made mid-investigation before the real crash cause (the activation.py raw-PTX bug, e4m3_native()'s tuple-comparison bug) was actually isolated:

• decode_paged_attention's USE_TL_DOT broadcast-sum fallback — added on top of the block_h>=16 floor as a second-layer guard against what turned out to be a different kernel's bug. Re-tested the real tl.dot/WMMA path directly (eager, batched, and inside actual CUDA graph capture+replay) now that the real cause is fixed: works fine. Removed — the block_h>=16 floor alone is sufficient (that part's still needed and correct; RDNA WMMA genuinely has no instruction below M=16).
• _select_extend_tile's 128x64→64x32 shrink on HIP, diagnosed via a py-spy trace as VGPR exhaustion before the session had found the actual culprit. Re-verified at the original CUDA-tuned 128x64 tile against the exact kernel it targeted (Qwen3.6-35B-A3B-FP8's GDN/split-extend path) — no crash, correct output. Reverted.

Both changes reduce the diff (net -26 lines) and restore full matrix-core acceleration + the original tile size rather than leaving unnecessary slow-path code in place. Re-verified against real chat completions on both models after reverting, not just the isolated kernel tests.