Skip to content

fix(rocm): gate CUDA-only optional backends on ROCm - #134

Draft
zihaomu wants to merge 11 commits into
FlashML-org:mainfrom
zihaomu:feat/rocm-rdna4-backends
Draft

zihaomu wants to merge 11 commits into
FlashML-org:mainfrom
zihaomu:feat/rocm-rdna4-backends

Conversation

@zihaomu

@zihaomu zihaomu commented Aug 24, 2026 •

Copy link
Copy Markdown
Contributor

Draft follow-up to #132. This branch is based on the #132 head. Until #132 merges, GitHub Files changed also includes the foundation diff; the incremental review scope is three files.

Summary

  • never select flashinfer, sgl_kernel, OpenAI triton_kernels, Marlin, or other NVIDIA donor binaries on ROCm;
  • explicitly route NVFP4 auto/triton requests to Triton on ROCm;
  • fail clearly for forced marlin/flashinfer requests without probing or importing donor packages;
  • avoid querying CUDA driver/capability paths on an AMD runtime;
  • remove the unused driver_hip_version() helper;
  • retain existing Triton/PyTorch fallback paths and CUDA selection behavior.

Incremental review scope

Relative to zihaomu:feat/rocm-rdna3-rdna4-foundation:

  • python/freetoken/kernel/backend.py
  • python/freetoken/moe/nvfp4_backends.py
  • tests/kernels/test_backend.py

Diff size: 107 additions, 11 deletions.

Dependency

Validation

  • backend-selection tests: 8 passed
  • ROCm NVFP4 tests cover auto, triton, marlin, and flashinfer;
  • tests assert ROCm never executes package probes or donor imports;
  • gfx1201 Triton attention, norm, and activation fallback tests: PASS.

Compatibility

The CUDA path continues to probe and select the same optional packages. This remains Draft until #132 establishes the shared ROCm foundation.

nekomario28 added a commit to nekomario28/FreeToken that referenced this pull request Aug 25, 2026
@zihaomu
zihaomu force-pushed the feat/rocm-rdna4-backends branch from 9a523dc to 32dfb7b Compare August 31, 2026 07:55
@zihaomu
zihaomu force-pushed the feat/rocm-rdna4-backends branch from 32dfb7b to e35a0ec Compare September 16, 2026 07:49
@zihaomu

zihaomu commented Sep 16, 2026

Copy link
Copy Markdown
Contributor Author

Dependency and architecture refresh (2026-09-16): rebased this PR onto the conflict-resolved #132 head c2af4c6 and ported the CUDA-only backend gates to the current quantization architecture. The old python/freetoken/moe/nvfp4_backends.py remains deleted; ROCm now gates flashinfer, sgl_kernel, and the current vLLM probe, while tests exercise the new layers.quantization kernel selector for Triton fallback and forced Marlin/b12x rejection. Validation: tests/kernels/test_backend.py (5 passed), adjacent MiniMax NVFP4 selector test (1 passed), AST and git diff --check. New head: e35a0ec; GitHub reports clean/mergeable. No GPU hardware tests were rerun for this conflict refresh.

LZ-QWQ added a commit to LZ-QWQ/FreeToken that referenced this pull request Sep 19, 2026
bouclem and others added 11 commits September 21, 2026 10:54
- Add hip_compat.h shim mapping CUDA runtime API to HIP equivalents
- Update pinned_tensor.cpp to compile under both nvcc and hipcc
- Add ROCm detection in arch.py (is_rocm, get_rocm_gfx_arch, is_gfx11xx_family)
- Guard NVIDIA arch checks to return None on ROCm
- Skip nvcc version check in _toolchain.py when on ROCm
- Add ROCm build path in setup.py (ROCM_HOME, amdhip64, --offload-arch)
- Add _hip_cflags() in kernel/utils.py for JIT compilation on ROCm
- Add is_rocm() and driver_hip_version() in backend.py
- Add rocm-smi fallback in __main__.py for clangd generation
- Add TODO(ROCm) for NCCL->RCCL, flashinfer/sgl_kernel ROCm builds,
  Triton autotune RDNA3 tuning, PDL equivalent, hiprtc JIT cache
- Add AMD ROCm classifier in pyproject.toml
@zihaomu
zihaomu force-pushed the feat/rocm-rdna4-backends branch from e35a0ec to c4f4a36 Compare September 21, 2026 03:01
LZ-QWQ pushed a commit to LZ-QWQ/FreeToken that referenced this pull request Sep 22, 2026
LZ-QWQ pushed a commit to LZ-QWQ/FreeToken that referenced this pull request Sep 22, 2026
LZ-QWQ pushed a commit to LZ-QWQ/FreeToken that referenced this pull request Sep 22, 2026
@zihaomu zihaomu mentioned this pull request Sep 23, 2026
29 tasks

This branch has not been deployed

No deployments
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