Skip to content

ROCm: validate native gfx1151 serving on Strix Halo - #260

Open
dbourdea wants to merge 204 commits into
FlashML-org:mainfrom
dbourdea:amd-rocm-gfx1151
Open

ROCm: validate native gfx1151 serving on Strix Halo#260
dbourdea wants to merge 204 commits into
FlashML-org:mainfrom
dbourdea:amd-rocm-gfx1151

Conversation

@dbourdea

Copy link
Copy Markdown

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:

  • explicitly detect HIP so gfx1151 cannot be interpreted as NVIDIA SM 11.5
  • suppress CUDA-only optional backends when PyTorch reports HIP
  • accept HIP tensors in the TVM JIT and fast indexed-copy paths
  • use the safe native Triton serial NVFP4 prefill route on HIP, avoiding the reproduced grouped-prefill HSA aperture fault
  • make the GGUF JIT locate a system Thrust include directory when the PyTorch wheel omits Thrust
  • make the GGUF JIT locate a ROCm runtime library directory when the wheel lacks the unversioned libamdhip64.so linker name
  • document the code and the reproducible LAN-223 validation in detail

Validation

Validated on LAN-223 only:

  • AMD Radeon 8060S, gfx1151
  • PyTorch 2.13.0+rocm10.0.0, HIP 7.15.26333
  • nvidia/Qwen3.6-35B-A3B-NVFP4: OpenAI-compatible non-streaming and SSE streaming requests returned HTTP 200
  • google/gemma-4-26B-A4B-it-qat-q4_0-gguf revision d1c082be9cf3c8a514acf63b8761f4b41935842e: native GGUF HIP JIT compiled and linked, then non-streaming and SSE streaming requests returned HTTP 200
  • python -m pytest tests/utils/test_rocm_runtime.py -q: 2 passed
  • python -m compileall -q python tests and git diff --check passed locally

The 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.

Copilot AI lite review requested due to automatic review settings August 28, 2026 19:26

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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 gfx11xx capability 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__)

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 54d6ab2. setup.py now passes FREETOKEN_USE_ROCM=1 to both host extensions whenever the active PyTorch build reports torch.version.hip. hip_compat.h selects HIP when that explicit macro is present, in addition to the HIP compiler macros. Verified on LAN-223: setup.py build_ext --inplace compiled and linked both _pinned_tensor and _cpu_moe against libamdhip64.

Comment thread setup.py Outdated
Comment on lines +9 to +11

ROOT = Path(__file__).parent
IS_ROCM = CUDA_HOME is None and ROCM_HOME is not None

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 54d6ab2. IS_ROCM now derives solely from torch.version.hip rather than CUDA_HOME/ROCM_HOME precedence, so a CUDA toolkit installed beside an active HIP PyTorch build cannot select cudart or trigger nvcc checks. Added regression coverage and verified 12 focused ROCm runtime and pinned-extension tests on LAN-223, plus a successful native host-extension build.

skywalk1411 and others added 27 commits August 28, 2026 18:44
…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.
@dbourdea

Copy link
Copy Markdown
Author

Validation update for 42e5be2 (feat(rocm): validate Gemma vision and LAN-223 controls):

  • Repaired the Gemma 4 multimodal path by preserving image tensors through the tokenizer-to-scheduler wire and by emitting channel-planar RGB patches for the projector convolution layout.
  • FreeToken passed 21 of 21 repeated deterministic Gemma visual controls, covering red, green, blue, yellow, horizontal placement, and vertical placement. The matched ROCm llama.cpp control passed the same seven-fixture set.
  • Qwen OpenAI-compatible API controls passed: deterministic text and JSON checks, 10 of 10 multi-turn state-retention sessions, 9 of 9 long-context marker retrieval samples through 7,736 reported prompt tokens, and a clean-memory 30 of 30 multi-turn endurance run.
  • Focused regression tests passed: 21 of 21.
  • Same-host fixed-workload throughput evidence is documented: FreeToken NVFP4 measured 28.15 median decode TPS and llama.cpp ROCm Q4_K_M measured 48.87. This is not quantization-equivalent, so it is reported as an observed configuration result rather than a general runner ranking.

The commit adds the LAN-223 validation report, reproducible local runner controls, and regression coverage. It does not include private model files or host artifacts.

@dbourdea

Copy link
Copy Markdown
Author

Validation follow-up: full-context MoE cache telemetry is now recorded in commit 5c7f0fd.

A temporary loopback-only Qwen instance enabled --moe-collect-stats with the validated native ROCm/HIP configuration, 0.35 memory ratio, and 8,192-token KV reservation. It resolved 8,903 MoE cache slots and 8,224 KV pages. The fixed warmup plus three-sample scheduler workload completed 3/3 at 28.035 mean visible decode TPS with 0.0066 TPS standard deviation. Across 40,800 decode-layer calls, the device counters reported eight active experts per layer, 0.586 missing experts per layer, and a 7.33% cache-miss rate.

This supports the prior measured rejection of a larger static cache: reducing misses by raising cache capacity did not improve sustained end-to-end TPS. The normal no-counter Qwen service was immediately restored and passed the deterministic AIME quality gate with the required SHA-1 0acef4eab6f4 at 28.60 visible decode TPS, 399.08 ms TTFT, and 38.49 ms p99 stream-event gap.

The report documents the exact methodology, limitation, and LAN-223 artifact root. This is evidence for the native AMD port, not a claim of quantization-equivalent parity with the separate Q4_K_M llama.cpp control.

@dbourdea

Copy link
Copy Markdown
Author

Exact-format ROCm comparison update in commit d8a2dd6.

The branch now serves the exact same Qwen3.6-35B-A3B-UD-Q4_K_M.gguf used by the llama.cpp control through native FreeToken ROCm/HIP. The Qwen hybrid GGUF path keeps Q8_0, Q6_K, Q4_K, and Q5_K weights packed, uses the native Gated-DeltaNet path, and was launched with a source-revision-specific gfx1151 cache. A strict no-JIT verifier successfully resolved all 82 helper modules before model startup.

Matched protocol: same model file, same host and GPU, auto DPM policy, 8,192-token context, one request, same fixed prompt, greedy decoding, 256 requested output tokens, warmup, and three scored samples.

Runtime Mean visible decode TPS Median TPS Quality suite
FreeToken ROCm/HIP 48.444 48.450 3/3 pass
llama.cpp ROCm 10 49.125 49.131 3/3 pass

The fresh same-format difference is 0.680 TPS, or 1.39% in favor of this llama.cpp control. This removes the previous NVFP4-versus-Q4_K_M representation mismatch. FreeToken is near parity but does not yet meet or exceed llama.cpp under this strict workload. A temporary high DPM policy was also screened and rejected because it reduced FreeToken Q4 throughput to 47.287 TPS while quality remained correct.

The normal NVFP4 FreeToken endpoint was restored after the time-share test and passed its deterministic AIME hash gate. The report includes artifact locations, startup geometry, the 512-token raw-prompt check, and the full method.

@dbourdea

Copy link
Copy Markdown
Author

Final LAN-223 native ROCm/HIP Q4 recovery evidence is now on commit a937862.

  • The one-hour endurance battery completed 60 of 60 deterministic three-turn sessions with every visible answer correct.
  • The FreeToken runner process group stayed at 0 KiB swap for the entire run. Whole-host swap was retained separately as diagnostic telemetry, ranging from 33.07 MiB to 38.17 MiB.
  • Maximum-turn TTFT: mean 0.424 s, p95 0.414 s, p99 and maximum 1.184 s.
  • Maximum visible-token gap: mean 24.95 ms, p95 25.98 ms, p99 and maximum 27.17 ms.
  • The normal NVFP4 service was restored after the battery. It completed an approximately eight-minute cold expert initialization, then passed a live OpenAI-compatible API request with the correct visible answer.

The matched exact-Q4 comparison remains transparent: FreeToken recovery configuration was 47.960 TPS versus fresh llama.cpp ROCm 10 at 48.831 TPS, 1.78 percent lower. This PR does not claim throughput parity where the measured result is below it.

Raw artifacts and the machine-checkable endurance summary are documented in docs/lan223-rocm-validation-2026-08-30.md.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants