feat(rocm): AMD ROCm support (gfx1100) + Qwen3.5-MoE GGUF loader

.gitignore

@@ -227,3 +227,13 @@ benchmarks/cross_framework
 
 # local e2e/bench artifacts (harnesses may run with repo cwd)
 /results/
+
+# torch.utils.cpp_extension ROCm hipify build artifacts (preprocessed .cu -> .hip,
+# and includes rewritten to *_hip.cuh) written next to the GGUF sources.
+python/freetoken/kernel/csrc/gguf/*.hip
+python/freetoken/kernel/csrc/gguf/*_hip.cuh
+python/freetoken/kernel/csrc/gguf/ggml-common_hip.h
+
+# local session junk
+nohup.out
+.plans/

.vscode/settings.json

@@ -0,0 +1,19 @@
+{
+    // Prevent VS Code from re-introducing CRLF endings into the serve script —
+    // CRLF was what broke the backslash-newline continuations before.
+    "files.eol": "\n",
+    "files.associations": {
+        "*.sh": "shellscript"
+    },
+    "[shellscript]": {
+        "files.eol": "\n",
+        "editor.tabSize": 4,
+        "editor.insertSpaces": false
+    },
+    "terminal.integrated.env.linux": {
+        "PYTHONPATH": "${workspaceFolder}/python"
+    },
+    "python.defaultInterpreterPath": "${workspaceFolder}/.venv-rocm/bin/python",
+    "python.testing.pytestEnabled": true,
+    "python.testing.pytestArgs": ["tests"]
+}

.vscode/tasks.json

@@ -0,0 +1,63 @@
+{
+    // FreeToken server lifecycle tasks. All of them call scripts/serve-qwen-moe.sh,
+    // which builds its argv as a bash array and launches on one line — safe for the
+    // VS Code integrated shell (no backslash-newline continuations to mangle).
+    "version": "2.0.0",
+    "tasks": [
+        {
+            "label": "FreeToken: serve (start)",
+            "type": "shell",
+            "command": "${workspaceFolder}/scripts/serve-qwen-moe.sh",
+            "args": ["start"],
+            "options": {
+                "cwd": "${workspaceFolder}",
+                "env": {}
+            },
+            "isBackground": true,
+            "problemMatcher": [],
+            "presentation": {
+                "reveal": "always",
+                "panel": "dedicated",
+                "clear": true
+            },
+            "detail": "Launch the Qwen MoE server (~3-4 min model load), then wait for readiness."
+        },
+        {
+            "label": "FreeToken: serve (stop)",
+            "type": "shell",
+            "command": "${workspaceFolder}/scripts/serve-qwen-moe.sh",
+            "args": ["stop"],
+            "options": { "cwd": "${workspaceFolder}" },
+            "problemMatcher": [],
+            "presentation": { "reveal": "always", "panel": "dedicated" }
+        },
+        {
+            "label": "FreeToken: serve (status)",
+            "type": "shell",
+            "command": "${workspaceFolder}/scripts/serve-qwen-moe.sh",
+            "args": ["status"],
+            "options": { "cwd": "${workspaceFolder}" },
+            "problemMatcher": [],
+            "presentation": { "reveal": "always", "panel": "shared" }
+        },
+        {
+            "label": "FreeToken: serve (follow log)",
+            "type": "shell",
+            "command": "${workspaceFolder}/scripts/serve-qwen-moe.sh",
+            "args": ["log"],
+            "options": { "cwd": "${workspaceFolder}" },
+            "isBackground": true,
+            "problemMatcher": [],
+            "presentation": { "reveal": "always", "panel": "dedicated" }
+        },
+        {
+            "label": "FreeToken: tests (fast)",
+            "type": "shell",
+            "command": "${workspaceFolder}/.venv-rocm/bin/python",
+            "args": ["-m", "pytest", "-q", "-m", "not slow"],
+            "options": { "cwd": "${workspaceFolder}", "env": { "PYTHONPATH": "${workspaceFolder}/python" } },
+            "group": "test",
+            "problemMatcher": []
+        }
+    ]
+}

docs/install-amd.md

@@ -0,0 +1,83 @@
+# AMD GPU (ROCm) support
+
+FreeToken targets Linux + NVIDIA CUDA by default. AMD (ROCm) is a supported, tested
+configuration with a **single-GPU** milestone: correct functional path first, performance
+recovered via HIP ports where safe. This page covers installing and running on RX 7000.
+
+> Status: **experimental.** The default and best-tested path remains CUDA. AMD brings up a
+> correct functional path (Triton attention + offload/CPU MoE + portable quant) and is
+> recovering performance via the HIP kernel ports.
+
+## Requirements
+
+| Component | Requirement |
+| --- | --- |
+| OS | Linux x86_64 (Windows WDDM pinned-memory is a known edge, not supported yet) |
+| GPU | AMD RX 7000 (RDNA 3, `gfx1100`); RX 9000 (`gfx1201`) is future work |
+| ROCm | ROCm toolkit with `hipcc` (`/opt/rocm/bin/hipcc` or on `PATH`) |
+| torch | ROCm build, e.g. `torch==2.5.1+rocm6.2` |
+
+The build refuses to mix toolchains: it will **not** silently fall back to `nvcc`/`libcudart`
+when only the ROCm toolkit is present, and vice versa.
+
+## Install
+
+```bash
+# ROCm torch (PyTorch official ROCm wheels) -- must satisfy the repo's torch>=2.11,<2.12
+# build pin, so use the rocm7.2 index (rocm6.2 only carries torch up to 2.5.1).
+pip install --index-url https://download.pytorch.org/whl/rocm7.2 \
+    "torch==2.11.0+rocm7.2" torchvision triton-rocm==3.6.0
+
+# FreeToken with the ROCm extra (builds the native extensions with hipcc)
+uv pip install -e ".[rocm]" --no-build-isolation
+```
+
+`pip install ".[rocm]"` pulls ROCm-compatible `torch`/`triton`; the NVIDIA-only `[accel]`
+packages (`flashinfer`, `sgl-kernel`, `triton_kernels`, Marlin) are **not** installed on AMD
+and their backends are rejected with a clean error if requested.
+
+## Verified feature matrix
+
+| Feature | On AMD | Notes |
+| --- | --- | --- |
+| Attention | `--attention-backend triton` | flashinfer/fa/trtllm are NVIDIA-only and rejected |
+| MoE | `--moe-backend offload / cpu / hybrid` | offload needs pinned host memory |
+| Quant | BF16, MXFP4, GGUF (Q4_K/Q8_0), Triton inline-dequant NVFP4 | Marlin INT4 / native NVFP4 SASS unavailable |
+| NVFP4 checkpoints with no MXFP4 variant | converted to MXFP4 on load (auto) | `--nvfp4-backend auto` → triton/MXFP4 |
+| CUDA graphs (decode) | HIP graph capture **if** the capture probe passes | otherwise kernel-launch decode |
+| Multi-GPU (RCCL) | out of scope (single-GPU milestone) | |
+
+## CLI behavior on AMD
+
+* `--nvfp4-backend marlin` / `flashinfer` → error (NVIDIA-only). Use `triton` / `auto`.
+* `--attention-backend fi` / `fa` / `trtllm` → error (NVIDIA-only). Use `triton` / `auto`.
+* `--moe-backend fused` → warning (fused MoE is CUDA-only; falls back to offload/cpu).
+* `--nvfp4-backend auto` → resolves to the portable Triton inline-dequant path (or MXFP4
+  for a converted checkpoint).
+
+## Verify
+
+```bash
+ft version            # prints an AMD / ROCm banner
+ft serve --model Qwen3.6-35B-A3B \
+  --moe-backend offload --attention-backend triton --nvfp4-backend auto
+```
+
+`ldd` of the built `.so` should show `hiprt`/`amdhip64`, not `libcudart`.
+
+## AOT kernel cache
+
+Build the prebuilt `+rocm` kernel-cache wheel (no nvcc needed on the target):
+
+```bash
+scripts/build-release-wheels.sh   # on a ROCm torch + hipcc box; tags the cache +rocm
+```
+
+The runtime refuses to pair a `+rocm` cache with a `+cu130` runtime (and vice versa).
+
+## Notes / limitations
+
+* `nvtx_annotate` is a no-op on ROCm; roctx profiling is future work.
+* FP8 / NVFP4-class formats: BF16 / MXFP4 / GGUF are the supported AMD matrix; performance
+  parity vs CUDA is not guaranteed for NVFP4-class formats.
+* Windows AMD is not yet supported (WDDM zero-copy semantics differ).

freetoken-kernel-cache/build_backend.py

@@ -40,6 +40,11 @@ def _cuda_version_suffix() -> str:
         return ""
 
     cuda_version = getattr(torch.version, "cuda", None)
+    hip_version = getattr(torch.version, "hip", None)
+    if hip_version:
+        # ROCm torch: torch.version.cuda is None; tag the cache with +rocm so it pairs
+        # only with a ROCm runtime (kernel/utils.py._arch_tags enforces the match).
+        return "+rocm"
     if not cuda_version:
         return ""
     # The tag advertises torch's CUDA; the cache .so link nvcc's libcudart.
@@ -110,7 +115,18 @@ def _build_jit_cache() -> None:
     #   12.0 -> RTX 50 series, RTX PRO 6000 Blackwell (Blackwell, consumer / workstation)
     # Override with FREETOKEN_KERNEL_CACHE_ARCHES (space-separated maj.min) or
     # TVM_FFI_CUDA_ARCH_LIST directly. Needs an nvcc that supports every listed arch.
-    if "TVM_FFI_CUDA_ARCH_LIST" not in os.environ:
+    try:
+        import torch  # noqa: PLC0415
+
+        is_rocm_build = bool(getattr(torch.version, "hip", None))
+    except Exception:
+        is_rocm_build = False
+    if is_rocm_build:
+        # ROCm: the CUDA arch-list is meaningless; the gfx arch is passed through
+        # kernel/utils.py._arch_flags (--offload-arch), defaulting to the RX 7000
+        # (gfx1100) target. Env override for other RX 7000 SKUs / future archs.
+        os.environ.setdefault("FREETOKEN_KERNEL_CACHE_GFX", "gfx1100")
+    elif "TVM_FFI_CUDA_ARCH_LIST" not in os.environ:
         os.environ["TVM_FFI_CUDA_ARCH_LIST"] = os.getenv(
             "FREETOKEN_KERNEL_CACHE_ARCHES", "8.0 8.6 8.9 9.0 10.0 12.0"
         )

pyproject.toml

@@ -23,6 +23,7 @@ classifiers = [
     "Intended Audience :: Science/Research",
     "Operating System :: POSIX :: Linux",
     "Environment :: GPU :: NVIDIA CUDA",
+    "Environment :: GPU :: AMD ROCm",
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
@@ -78,6 +79,13 @@ fi = ["flashinfer-python[cu13]>=0.6,<0.7"]
 # renamed from sgl-kernel at 0.4; still imports as `sgl_kernel`, so never co-install both
 sgl = ["sglang-kernel==0.4.5"]
 accel = ["freetoken[fi,sgl]"]
+# ROCm (AMD) install: the NVIDIA-only fi/sgl/Marlin packages are NOT pulled in. torch must
+# come from the ROCm wheel index (e.g. `pip install torch==2.11.0+rocm7.2` from
+# https://download.pytorch.org/whl/rocm7.2) so the native extensions build against the HIP
+# runtime; this extra pins the rest. See docs/install-amd.md.
+rocm = [
+    "triton==3.6.0; platform_system == 'Linux'",
+]
 # NOTE: the Marlin W4A16 NVFP4 expert-GEMM path (sm_80-99) borrows vLLM's AOT wheel
 # (vllm>=0.14,<0.15), which pins transformers>=4.56,<5 and so is INCOMPATIBLE with the
 # core transformers>=5.5 requirement. It is therefore not a lockable extra and is left

python/freetoken/attention/__init__.py

@@ -96,6 +96,19 @@ def create_triton_backend(config: ModelConfig):
     return TritonAttentionBackend(config)
 
 
+@SUPPORTED_ATTENTION_BACKENDS.register(
+    "torch",
+    BackendInfo(
+        supported_types=frozenset({AttnType.FULL}),
+        # Debugging/eager ground-truth backend; no package/arch requirements.
+    ),
+)
+def create_torch_backend(config: ModelConfig):
+    from .torch import TorchAttentionBackend
+
+    return TorchAttentionBackend(config)
+
+
 @SUPPORTED_ATTENTION_BACKENDS.register(
     "dsv4_sparse",
     BackendInfo(supported_types=frozenset({AttnType.DSV4})),