ROCm: validate native gfx1151 serving on Strix Halo - #260
Open
dbourdea wants to merge 16 commits into
Open
Conversation
…cks for gfx1150
FreeToken built and ran only against CUDA. On a native ROCm install (tested
on a Ryzen AI 9 HX 470 / Radeon 890M, gfx1150) it failed at every stage:
build, JIT compile, and finally a silent native crash mid-warmup with no
Python traceback.
Build system (setup.py, hip_compat.h):
- _pinned_tensor and _cpu_moe link against HIP instead of cudart when
ROCM_HOME is present (CUDA_HOME stays required on the CUDA path).
- kernel/csrc/hip_compat.h aliases the CUDA Runtime API calls those two
files use onto their HIP equivalents, 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
cuda_runtime.h in 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 instead, with
with_attr() as a no-op since there is no attribute to carry.
- 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 (never on ROCm).
- attention.py: the decode kernel's GQA head-tile floors at 16 under HIP
(RDNA WMMA has no instruction below M=16); the kernel already masks
padded head lanes for non-power-of-two groups, so this is a safe
widening. Falls back to broadcast-multiply-reduce instead of tl.dot for
that tile as a second-layer guard. The split extend/prefill kernel's
tile shrinks from 128x64 to 64x32 under HIP -- running both the cached-
and newly-computed-KV loops live at once is register-heavier than the
plain extend kernel, and exhausts this GPU's VGPR file at the CUDA-tuned
tile size.
- activation.py (the actual root cause of the crash above): _fast_tanh and
_fast_ex2 inline raw 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 ("couldn't allocate output register for
constraint 'f'") that looks like a matrix-core or register-pressure
issue and sent debugging down that path for a while. Routed through
libdevice.tanh / tl.exp2 on HIP instead.
pyproject.toml: loosen the torch/triton ceilings so ROCm builds (which
carry a local version segment such as +rocm7.14.0...) can satisfy them.
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.
Verified end to end on gfx1150: server boot, weight load, KV cache alloc,
CUDA graph capture at bs=1/2/4, and real chat completions against
Qwen/Qwen3-8B (bf16, triton attention backend).
…t-copy kernel Two more real bugs found while running an actual MoE model (Qwen3.6-35B-A3B-FP8, --moe-backend offload) end to end on gfx1150, past what the first commit covered. e4m3_compat.py: e4m3_native() decides whether kernels get raw fp8 tensors or a uint8 view by checking torch.cuda.get_device_capability() >= (8, 9). On a HIP build that call returns the GPU's RDNA generation number, not a CUDA compute capability -- gfx1150 reports (11, 5), and (11, 5) >= (8, 9) is True by plain tuple comparison (11 > 8), so this incorrectly claimed native fp8 support on AMD. Triton's own compile-time twin, e4m3_native_cx() (target_info. cuda_capability_geq, which checks target.backend != "cuda" first), correctly said False, so the kernel compiled for the emulated uint8 path while the host side hands it an untouched fp8 tensor -- IncompatibleTypeErrorImpl inside e4m3_u8_to_f32's bitwise ops. Fixed by checking torch.version.hip first. fast_index_copy.cuh (the offload cache's fast host->device expert-copy kernel, only exercised once a real MoE model with --moe-backend offload actually streams experts): same two problems as the first commit's fixes elsewhere in this file family, just not caught until this path actually ran. - Missing HIP aliases for cudaGetDevice/cudaDeviceGetAttribute/ cudaHostGetDevicePointer/the two cudaDevAttr* constants it uses -- added to utils.cuh's existing HIP block alongside the ones from the first commit. - load_nc/store_nc inline raw PTX (ld.global.L1::no_allocate, st.global.wt -- cache-policy hints, no HIP equivalent). Falls back to plain loads/stores under HIP; correctness unchanged, only the cache hint is lost. Verified: Qwen3.6-35B-A3B-FP8 (256 experts/layer x 40 layers, 3B active) boots and serves real chat completions with --moe-backend offload --moe-cache-size 2560 (25% of the model's 10240 total experts resident, LRU-evicting the rest from host RAM on every miss) -- ft ctl cache confirms the pool is live at the requested size, not silently falling back to full residency.
Third loading path verified: google/gemma-4-26B-A4B-it-qat-q4_0-gguf (native GGUF, MoE offload) now boots and serves on gfx1150, alongside the dense bf16 and FP8 MoE paths from the earlier commits. kernel/gguf.py: same nvcc-only-flag problem as elsewhere in this port, in 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). --expt-relaxed-constexpr is rejected outright, and the -ccbin/CXX-forcing block exists only to work around an nvcc+libtorch-headers compiler mismatch that doesn't apply under hipcc (its own bundled clang already is the host compiler). Both dropped on HIP. kernel/csrc/gguf/dispatch.h: the donor's SGLANG_SHFL_XOR_SYNC(_WIDTH) macros forward a CUDA-style 32-bit mask straight into __shfl_xor_sync. HIP's amd_warp_sync_functions.h static_asserts the mask must be 64 bits unconditionally (regardless of actual wavefront width) -- widened the cast on HIP only. .gitignore: torch's ROCm auto-hipify (a real, working translation pass built into torch.utils.cpp_extension -- unlike the other two JIT paths, this one needed no manual porting for the .cu/.cuh sources themselves) writes translated copies next to the CUDA sources it processes (gguf_kernel.cu -> .hip, *.cuh -> *_hip.cuh). Regenerated every build, never hand-edited; ignore rather than track. Not in this commit, environment-only: the pip ROCm nightly distribution used here (rocm.nightlies.amd.com) ships no thrust/rocprim headers, which torch's own extension headers pull in transitively. Ubuntu's librocthrust-dev is one fix, but it depends on libamdhip64-dev, which drops a second, conflicting HIP header set into /usr/include/hip that silently wins over the correct pip-bundled ones for any plain -I (though not -isystem) -- diagnosed by hand with `clang++ -v` and a minimal reproducer. Worked around locally by extracting just the thrust/rocprim headers (dpkg -x, no install) into the pip package's own include dir and removing the conflicting system packages; ROCM_PATH/HIP_PATH/HIP_DEVICE_LIB_PATH also had to point at the pip package for this JIT path's device-bitcode-library lookup. Left out of the diff since there's no source change to make -- noting it here for the next person on this distribution.
Both were made mid-investigation, before the real cause of a since-fixed crash (the raw-PTX bug in activation.py, and separately the e4m3_native() tuple- comparison bug) was actually found. Re-tested each in isolation -- eager, batched, and inside real CUDA graph capture+replay -- now that those are fixed, and both work fine at the original, CUDA-tuned settings: - decode_paged_attention: the block_h>=16 floor (kept -- RDNA WMMA genuinely has no instruction below M=16, confirmed independently and matches upstream FlashML-org#137) was sufficient on its own. The USE_TL_DOT broadcast-sum fallback this PR had added on top was solving a problem that was actually in a different kernel; removed, restoring real matrix-core-accelerated decode attention. - _select_extend_tile: the 128x64 -> 64x32 shrink on HIP was diagnosed as a VGPR-exhaustion issue via a py-spy trace mid-investigation, before the session had isolated the actual crash to activation.py. Re-verified end-to-end against Qwen3.6-35B-A3B-FP8's GDN/split-extend path (the kernel this shrink targeted) at the original tile size: no crash, correct output. Reverted to the CUDA-tuned tile. Both re-verified against real chat completions (Qwen3-8B for the decode path, Qwen3.6-35B-A3B-FP8 for the extend/split path) after reverting, not just the isolated kernel tests.
There was a problem hiding this comment.
🟡 Changes recommended
There are build-detection/compilation-guard issues in the new ROCm linking path that can cause ROCm builds to incorrectly include CUDA headers or link against the wrong runtime.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds a native ROCm/HIP execution path targeting AMD gfx1151 (Strix Halo / Radeon 8060S) while keeping NVIDIA CUDA behavior as a separately gated path, including build/link fixes, runtime feature gating, and validation docs.
Changes:
- Add explicit HIP/ROCm runtime detection to prevent
gfx11xxcapability tuples from being treated as NVIDIA SM versions and to disable CUDA-only optional backends on HIP. - Extend native/JIT paths (TVM JIT matchers, fast index copy, Triton kernels, GGUF JIT) to accept ROCm devices and avoid CUDA-only flags/launch kwargs/PTX.
- Add regression tests plus detailed LAN-223 validation and ROCm deployment documentation.
File summaries
| File | Description |
|---|---|
| tests/utils/test_rocm_runtime.py | Adds regression tests for HIP gating and CUDA-only optional backend suppression. |
| tests/kernels/test_pinned_tensor.py | Adjusts UVA identity test behavior for HIP runtime semantics. |
| setup.py | Links native extensions against HIP runtime when ROCm is detected; skips nvcc toolchain checks on ROCm. |
| python/freetoken/utils/arch.py | Adds is_rocm_runtime() and gates CUDA capability queries on non-HIP builds. |
| python/freetoken/utils/init.py | Exports is_rocm_runtime from utils. |
| python/freetoken/moe/fused_nvfp4.py | Routes HIP prefill to safer serial Triton NVFP4 path. |
| python/freetoken/kernel/utils.py | Drops nvcc-only flags when compiling under HIP. |
| python/freetoken/kernel/triton/norm.py | Avoids passing CUDA-only launch_pdl kwarg to AMD Triton backend. |
| python/freetoken/kernel/triton/e4m3_compat.py | Forces e4m3 “native” detection off on HIP to avoid tuple false-positives. |
| python/freetoken/kernel/triton/attention.py | Adjusts HIP decode tiling to satisfy RDNA WMMA constraints. |
| python/freetoken/kernel/triton/activation.py | Avoids inlining PTX on HIP and avoids CUDA-only launch kwargs. |
| python/freetoken/kernel/gguf.py | Adds ROCm fallback discovery for Thrust headers and libamdhip64.so linker path. |
| python/freetoken/kernel/csrc/pinned_tensor.cpp | Switches pinned tensor extension to HIP/CUDA compat shim header. |
| python/freetoken/kernel/csrc/jit/store.cu | Allows ROCm tensors in TVM JIT store path device checks. |
| python/freetoken/kernel/csrc/jit/index.cu | Allows ROCm tensors in TVM JIT index path device checks. |
| python/freetoken/kernel/csrc/jit/fast_index_copy.cuh | Adds HIP-safe load/store fallbacks and broadens accepted device types. |
| python/freetoken/kernel/csrc/include/freetoken/utils.cuh | Adds HIP runtime includes/aliases and HIP LaunchKernel implementation. |
| python/freetoken/kernel/csrc/hip_compat.h | Introduces CUDA-runtime API name shims for HIP linking in host C++ extensions. |
| python/freetoken/kernel/csrc/gguf/dispatch.h | Fixes HIP shuffle mask width expectations by widening masks for HIP. |
| python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp | Switches CPU MoE extension to HIP/CUDA compat shim header. |
| python/freetoken/kernel/backend.py | Adds HIP detection and disables CUDA-only optional backend probes on ROCm. |
| python/freetoken/engine/engine.py | Skips optional Triton prefill warmup by default on ROCm unless explicitly enabled. |
| pyproject.toml | Loosens torch/triton constraints and adds ROCm classifier. |
| docs/lan223-rocm-validation-2026-08-28.md | Records LAN-223 reproducible ROCm validation evidence and known limitations. |
| docs/amd-rocm-gfx1151.md | Adds ROCm/GFX1151 port documentation and validation procedure. |
| .gitignore | Ignores ROCm hipify-generated artifacts (*.hip, *_hip.*). |
Review details
- Files reviewed: 25/26 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| // Lets pinned_tensor.cpp and cpu_moe_ext.cpp call the CUDA Runtime API names they | ||
| // were written against while actually linking HIP on ROCm builds. Only the calls | ||
| // those two files use are covered -- this is not a general CUDA/HIP compat layer. | ||
| #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__) |
Comment on lines
9
to
+11
|
|
||
| ROOT = Path(__file__).parent | ||
| IS_ROCM = CUDA_HOME is None and ROCM_HOME is not None |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds and validates a native ROCm and HIP execution path for AMD Radeon 8060S
gfx1151, while retaining the NVIDIA CUDA path as a separate path.This branch includes the current upstream ROCm work from #241 plus focused follow-up fixes found while completing full-model serving validation on Strix Halo:
gfx1151cannot be interpreted as NVIDIA SM 11.5libamdhip64.solinker nameValidation
Validated on LAN-223 only:
gfx11512.13.0+rocm10.0.0, HIP7.15.26333nvidia/Qwen3.6-35B-A3B-NVFP4: OpenAI-compatible non-streaming and SSE streaming requests returned HTTP 200google/gemma-4-26B-A4B-it-qat-q4_0-ggufrevisiond1c082be9cf3c8a514acf63b8761f4b41935842e: native GGUF HIP JIT compiled and linked, then non-streaming and SSE streaming requests returned HTTP 200python -m pytest tests/utils/test_rocm_runtime.py -q: 2 passedpython -m compileall -q python testsandgit diff --checkpassed locallyThe full command shapes, raw artifact locations, API behavior, limitations, and telemetry are recorded in
docs/lan223-rocm-validation-2026-08-28.md.Scope
This PR does not add llama-swap integration. CUDA graph capture remains disabled for the HIP MVP, and the safe serial Triton NVFP4 prefill path is intentionally used until the grouped path is stable on this hardware.