fix(rocm): complete RDNA3 runtime path

[nekomario28]

Follow-up to FlashML-org#23, based directly on feat/rocm-support at 27c0977b6f2ffd476b85de116e2db839b614d76a.

The initial draft used a correctness-first PyTorch fallback for fast_index_copy_jit. Real RX 7800 XT / gfx1101 validation exposed additional CUDA-only assumptions, so the published branch now uses the native HIP path throughout the supported single-GPU RDNA3 scope. History is intentionally retained rather than force-pushed.

Current published head: 48146760e6d2dfc92dd35233447c9602ad57054b — 9 commits after PR FlashML-org#23 head.

What this fixes

• package the FreeToken C++ include directory for extension builds;
• route the CPU MoE extension and common JIT utilities through hip_compat.h;
• provide HIP runtime aliases needed by pinned-host and JIT code;
• omit CUDA-only PDL/grid-constant launch features on HIP;
• use HIP libdevice activation math and the software E4M3 path on ROCm;
• replace NVIDIA PTX non-caching vector loads/stores with portable AMD vector accesses while preserving the existing CUDA PTX branch;
• avoid passing CUDA-only launch_pdl options to Triton on HIP;
• emit one --offload-arch= argument per requested ROCm architecture instead of embedding a semicolon-separated architecture list in one shell argument;
• explicitly accept kDLROCM / kDLROCMHost in the supported single-GPU JIT tensor matchers (fast_index_copy, index, store).

The earlier broad #define kDLCUDA kDLROCM / kDLCUDAHost kDLROCMHost compatibility alias was removed. Explicit matchers are easier to audit and avoid silently broadening the still-NCCL-only pynccl path. pynccl remains CUDA-only here; RCCL/multi-GPU migration is intentionally outside this PR.

The original Python row-copy fallback and its tests are no longer present in the final diff.

Real hardware validation

Hardware/runtime: RX 7800 XT (gfx1101), PyTorch 2.11.0+rocm7.2, HIP 7.2.26015, Triton ROCm 3.6.0.

Earlier native-runtime validation on this branch lineage:

• HIP device/matmul and native pinned extensions: PASS;
• production-shape fused pinned-host-to-VRAM JIT copy (1024 threads x 8 blocks): byte-exact PASS;
• NVFP4 MoE overlap/SwiGLU/decode/cache-stomp: 4 passed;
• NVFP4 dense row-major/transposed paths at rows 1/8/200: six numerical comparisons PASS;
• Qwen3.5 paged/decode/extend attention: 3 passed;
• Qwen3.5 GDN 65-token prefill output and recurrent-state contract: matched the pure-PyTorch reference;
• per-tensor FP8 W8A16: 22 passed, 11 skipped; skips are native-FP8-only cases intentionally excluded on this RDNA3 target;
• targeted regression from the previously validated native tree: 29 passed, 11 skipped.

Direct validation of the current published head 48146760... was then run through the Dormant Giant A0 composition harness with no composition-local DLPack compatibility patch applied. Only the separate staged-source seam was layered on top of FreeToken.

Exact validation evidence:

• Dormant Giant exact head: 856bc758b55355bac81acbe2bda4e53462088429;
• workflow run: 32726063241;
• job: 97427500852;
• job conclusion: SUCCESS;
• source contract: PASS rocm-dlpack-source-contract;
• ROCM_MULTI_ARCH_FLAGS=SHELL_SAFE_PER_ARCH;
• ROCM_DLPACK_DEVICE_MATCHING=EXPLICIT_SINGLE_GPU_NO_GLOBAL_ALIAS;
• ROCM_PYNCCL_SCOPE=CUDA_ONLY_RCCL_NOT_CLAIMED;
• ROCM_A0_COPY_CLASSIFICATION=PASS_NATIVE_HIP_SINGLE_BANK_COPY;
• full staged movement: ROCM_A0_PHYSICAL_CLASSIFICATION=PASS_NATIVE_FREETOKEN_A0_ON_ROCM;
• resident-vs-staged FreeToken BF16 expert decode output: ROCM_A0_BF16_OUTPUT_CLASSIFICATION=PASS_RESIDENT_STAGED_BIT_EXACT with logical IDs [[5,2],[7,1]], remapped cache slots [[5,3],[6,1]], non-contiguous victims (1,3,5,6), and lease_cleanup=0.

The last two checks are composition evidence rather than a claim that Dormant Giant is part of this PR; their relevance here is that the published FreeToken head itself successfully compiled and executed the native HIP copy/JIT path and FreeToken expert decode path under real gfx1101 load without the former external DLPack patch.

Scope

This is single-GPU RDNA3 correctness work. It does not claim:

• RCCL/multi-GPU tensor parallelism;
• RDNA3 performance tuning or optimal Triton configs;
• physical NVIDIA/CUDA regression coverage for these follow-up commits;
• full production-model end-to-end serving acceptance.

CUDA-specific behavior remains on the existing non-HIP paths, and the explicit DLPack matchers avoid changing unsupported multi-GPU semantics by preprocessor side effect.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 6f8240a0-91ce-4265-b6d3-0aaaf9c6430e

📥 Commits

Reviewing files that changed from the base of the PR and between 27c0977 and 4e11ed0.

📒 Files selected for processing (8)

• python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp
• python/freetoken/kernel/csrc/include/freetoken/hip_compat.h
• python/freetoken/kernel/csrc/include/freetoken/utils.cuh
• python/freetoken/kernel/csrc/jit/fast_index_copy.cuh
• python/freetoken/kernel/triton/activation.py
• python/freetoken/kernel/triton/e4m3_compat.py
• python/freetoken/kernel/triton/norm.py
• setup.py

Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review.

---

📝 Walkthrough

Walkthrough

The change adds HIP compatibility mappings for C++ extensions and introduces ROCm-specific paths for memory operations, activation math, fp8 handling, and kernel launch configuration.

Changes

ROCm and HIP support

Layer / File(s)	Summary
HIP C++ compatibility and build wiring python/freetoken/kernel/csrc/..., setup.py	C++ sources use hip_compat.h. The header maps CUDA-named APIs and attributes to HIP equivalents. Extension builds include the kernel header directory.
ROCm kernel primitives and numeric paths python/freetoken/kernel/csrc/jit/fast_index_copy.cuh, python/freetoken/kernel/triton/activation.py, python/freetoken/kernel/triton/e4m3_compat.py	AMD HIP builds use direct vector loads and stores, libdevice math functions, and emulated e4m3 handling.
ROCm launch configuration python/freetoken/kernel/csrc/include/freetoken/utils.cuh, python/freetoken/kernel/triton/activation.py, python/freetoken/kernel/triton/norm.py	ROCm launches omit unsupported launch_pdl arguments and CUDA PDL attributes. Non-ROCm launches retain the existing behavior.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 4e11e

The PR completes localized ROCm/RDNA3 runtime compatibility changes, and no actionable merge-blocking risk remains beyond normal checks and review.

Suggested reviewers: jason-fxz

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring check was indeterminate for this PR — some files could not be analyzed in time. Not blocking.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely summarizes the PR's main change: completing the RDNA3 ROCm runtime path.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[nekomario28]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.