5080deploy: rebuild on upstream main d3512b4 (upstream #546 max_tokens fix, #548 row store, #588 eager PLE buffers without its graph fence)

docs/cli.md

@@ -33,6 +33,7 @@ parsers all resolve automatically from the checkpoint and the GPU.
 |---|---|---|
 | `--model-path`, `--model` | required | Local dir, HF repo id, or an FTW dir (auto-detected) |
 | `--served-model-name` | basename of `--model` | Model id reported by `/v1/models` |
+| `--hf-overrides` | — | JSON object applied to the checkpoint's config as vLLM's `--hf-overrides`: a nested config section updates key by key, any other value is replaced whole. A YaRN `rope_parameters` override serves `original_max_position_embeddings * factor` positions |
 
 ### Server & runtime
 

docs/install.md

@@ -1,5 +1,7 @@
 # Install
 
+AMD users: see the [AMD ROCm installation guide (WIP)](install_amd.md).
+
 ## Requirements
 
 - Linux x86_64, NVIDIA GPU, driver r580+ (CUDA 13)

docs/install_amd.md

@@ -0,0 +1,39 @@
+# AMD ROCm installation (WIP)
+
+AMD support is experimental and remains a work in progress. See the
+[AMD Support roadmap](https://github.com/FlashML-org/FreeToken/issues/541) for
+the current integration and qualification status.
+
+## Requirements
+
+- Linux x86_64
+- AMD RDNA3/RDNA4 GPU (`gfx1100`-`gfx1103`, `gfx1200`, or `gfx1201`)
+- ROCm 7.14
+- Python >= 3.10
+
+## Install from source
+
+Use an official ROCm PyTorch image whose PyTorch version satisfies the project's
+`torch>=2.11,<2.12` constraint. For RDNA4, the matching ROCm 7.14 image is:
+
+```bash
+VIDEO_GID="$(getent group video | cut -d: -f3)"
+RENDER_GID="$(getent group render | cut -d: -f3)"
+docker run --rm -it \
+  --device=/dev/kfd --device=/dev/dri \
+  --group-add="$VIDEO_GID" --group-add="$RENDER_GID" --ipc=host \
+  --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
+  -e PYTORCH_ROCM_ARCH=gfx1201 -e FREETOKEN_ROCM_ARCH=gfx1201 \
+  -v "$PWD:/workspace/FreeToken" -w /workspace/FreeToken \
+  rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.11.0 bash
+```
+
+Inside the container, preserve the ROCm-enabled PyTorch already supplied by the
+image and disable build isolation so it is also used to compile the extensions:
+
+```bash
+python -m pip install --no-build-isolation -e .
+```
+
+Set both architecture variables to `gfx1200` for RX 9060 family GPUs, or to the
+actual target reported by `rocminfo`.

install.sh

@@ -20,7 +20,8 @@
 #                                 Defaults to $DEFAULT_KERNEL_CACHE_WHEEL_URL; if
 #                                 unset and FREETOKEN_WHEEL is local, the script
 #                                 auto-detects a sibling freetoken_kernel_cache-*.whl.
-#   FREETOKEN_HOME                install root (default: ~/.freetoken); venv at $FREETOKEN_HOME/venv
+#   FREETOKEN_HOME                install root (default: ~/.freetoken); venv at $FREETOKEN_HOME/venv;
+#                                 a successful install writes it to environment.d
 #   FREETOKEN_PY_VERSION          python for the venv (default: 3.12 — must match the wheel tag)
 #   FREETOKEN_BIN_DIR             where to symlink `ft` (default: ~/.local/bin)
 #   FREETOKEN_ENV_DIR             environment.d dir (default: ~/.config/environment.d)
@@ -248,8 +249,8 @@ ln -sf "$FT_BIN" "$BIN_DIR/ft"
 say "symlinked $BIN_DIR/ft -> $FT_BIN"
 
 mkdir -p "$ENV_DIR"
-printf 'FREETOKEN_FT_BIN=%s\n' "$FT_BIN" > "$ENV_DIR/50-freetoken.conf"
-say "wrote $ENV_DIR/50-freetoken.conf (FREETOKEN_FT_BIN) — GUI picks it up after next login"
+printf 'FREETOKEN_HOME=%s\n' "$(cd "$FT_HOME" && pwd)" >"$ENV_DIR/50-freetoken.conf"
+say "wrote $ENV_DIR/50-freetoken.conf (FREETOKEN_HOME) - FreeToken Desktop reads the file, other programs see the variable after next login"
 
 # --- 5. Self-check ---------------------------------------------------------
 if "$FT_BIN" --help >/dev/null 2>&1; then
@@ -264,10 +265,9 @@ ${C_GREEN}FreeToken engine installed.${C_RESET}
 
   ft binary        $FT_BIN
   on PATH as       $BIN_DIR/ft   (ensure $BIN_DIR is on PATH)
-  Desktop env      FREETOKEN_FT_BIN via environment.d (re-login to apply)
+  Desktop env      FREETOKEN_HOME in $ENV_DIR/50-freetoken.conf
 
-Run in this shell without re-login:
-  export FREETOKEN_FT_BIN="$FT_BIN"
+Run:
   ft serve --model <path> --port 1919
 
 EOF

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",
@@ -59,7 +60,10 @@ dependencies = [
     "torchvision>=0.26,<0.27",
     "tqdm>=4.66,<5",
     "transformers>=5.16,<5.17",
-    "triton==3.6.0; platform_system == 'Linux'",
+    # CUDA torch 2.11 resolves Triton 3.6; AMD's ROCm 7.14 image supplies its
+    # gfx1201-enabled Triton 3.7 build. Keep both supported without replacing the
+    # runtime-specific wheel selected by the PyTorch distribution.
+    "triton>=3.6,<3.8; platform_system == 'Linux'",
     "uvicorn>=0.30,<1",
 ]
 

python/freetoken/attention/qsa_sparse.py

@@ -187,23 +187,22 @@ def _cmp_pages(self, slot: int) -> torch.Tensor:
 
     def _index_rope_cache(self) -> torch.Tensor:
         """cos/sin table of the indexer rope: same rotary_dim and frequencies as the main
-        attention, ``head_size`` 128 instead of 256, so it is a separate get_rope instance.
+        attention with ``head_size`` 128 instead of 256, so the attention's own table.
 
         The table itself (not RotaryEmbedding.forward) because the indexer's norm+rope is one
         fused kernel and the compressed keys rope at their group's position, not the query's."""
         if self._index_cos_sin is None:
-            from freetoken.layers.rotary import get_rope
+            from freetoken.layers.rotary import rope_cos_sin_table
 
             rotary = self.rotary_config
             with torch.device(self.device):
-                rope = get_rope(
-                    head_dim=self.index_head_dim,
-                    rotary_dim=rotary.rotary_dim,
-                    max_position=rotary.max_position,
-                    base=rotary.base,
-                    rope_scaling=tuple(rotary.scaling.items()) if rotary.scaling else None,
+                self._index_cos_sin = rope_cos_sin_table(
+                    self.index_head_dim,
+                    rotary.rotary_dim,
+                    rotary.max_position,
+                    rotary.base,
+                    tuple(rotary.scaling.items()) if rotary.scaling else None,
                 )
-            self._index_cos_sin = rope._cos_sin_cache.to(self.device)
         return self._index_cos_sin
 
     def _index_rope_rows(self, positions: torch.Tensor) -> torch.Tensor:

python/freetoken/engine/config.py

@@ -4,7 +4,7 @@
 import os
 from dataclasses import dataclass, field, replace
 from functools import cached_property
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, Any, List, Mapping
 
 import torch
 from freetoken.distributed import DistributedInfo
@@ -24,6 +24,8 @@ class EngineConfig:
     model_path: str
     tp_info: DistributedInfo
     dtype: torch.dtype
+    # --hf-overrides: applied to the checkpoint config the model is built from (cached_load_hf_config)
+    hf_overrides: Mapping[str, Any] = field(default_factory=dict)
     max_running_req: int = 4
     attention_backend: str = "auto"
     moe_strategy: str = "auto"
@@ -117,7 +119,7 @@ def __post_init__(self):
 
     @cached_property
     def hf_config(self):
-        return cached_load_hf_config(self.model_path)
+        return cached_load_hf_config(self.model_path, self.hf_overrides)
 
     @cached_property
     def model_spec(self) -> ModelSpec:
@@ -156,7 +158,7 @@ def model_config(self) -> ModelConfig:
     def max_seq_len(self) -> int:
         if self.max_seq_len_override is not None:
             return self.max_seq_len_override
-        return self.model_config.rotary_config.max_position
+        return self.model_config.rotary_config.table_positions
 
     @property
     def max_forward_len(self) -> int:

python/freetoken/engine/engine.py

@@ -2129,18 +2129,17 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
             )
         override("num_page_override", config.num_token_override // config.page_size)
 
-    # The rope cos/sin table is baked to rotary_config.max_position, and neither rope kernel
+    # The rope cos/sin table covers rotary_config.table_positions, and neither rope kernel
     # bounds-checks the position it gathers with -- a longer ceiling reads past the table.
     # DSV4 is exempt: it sizes its own table from the resolved max_seq_len (_adjust_dsv4_config).
     rotary = getattr(model_config, "rotary_config", None)
     seq_override = getattr(config, "max_seq_len_override", None)
     if seq_override is not None and rotary is not None and not is_dsv4:
-        if seq_override > rotary.max_position:
+        if seq_override > rotary.table_positions:
             raise ValueError(
                 f"--max-seq-len-override {seq_override} exceeds the model's "
-                f"rope table ({rotary.max_position} positions). Serving past it would read "
-                "out of bounds; extend the checkpoint's rope_scaling / "
-                "max_position_embeddings in config.json instead."
+                f"rope table ({rotary.table_positions} positions). Serving past it would read "
+                "out of bounds; extend the rope with --hf-overrides (YaRN rope_parameters) instead."
             )
 
     # The startup ServerArgs dump is the *requested* config, printed in the frontend process

python/freetoken/kernel/__main__.py

@@ -6,27 +6,32 @@ def generate_clangd():
     import subprocess
 
     from freetoken.kernel.utils import DEFAULT_INCLUDE
-    from freetoken.utils import init_logger
+    from freetoken.utils import get_rocm_gfx_arch, init_logger, is_rocm
     from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path
 
     logger = init_logger(__name__)
     logger.info("Generating .clangd file...")
     include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE
-    status = subprocess.run(
-        args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
-        capture_output=True,
-        check=True,
-    )
-    compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0]
-    major, minor = compute_cap.split(".")
+
+    # TODO(ROCm): hiprtc JIT cache should be separate from nvcc JIT cache to avoid stale binaries.
+    if is_rocm():
+        arch_flags = ["-xhip", f"--offload-arch={get_rocm_gfx_arch() or 'gfx1201'}"]
+    else:
+        try:
+            status = subprocess.run(
+                args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
+                capture_output=True,
+                check=True,
+            )
+            compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0]
+            major, minor = compute_cap.split(".")
+        except (subprocess.CalledProcessError, FileNotFoundError, ValueError):
+            import torch
+
+            major, minor = torch.cuda.get_device_capability()
+        arch_flags = ["-xcuda", f"--cuda-gpu-arch=sm_{major}{minor}"]
     compile_flags = ",\n    ".join(
-        [
-            "-xcuda",
-            f"--cuda-gpu-arch=sm_{major}{minor}",
-            "-std=c++20",
-            "-Wall",
-            "-Wextra",
-        ]
+        arch_flags + ["-std=c++20", "-Wall", "-Wextra"]
         + [f"-isystem{path}" for path in include_paths]
     )
     clangd_content = f"""

python/freetoken/kernel/_toolchain.py

@@ -1,4 +1,4 @@
-"""CUDA toolchain/torch consistency checks.
+"""CUDA/HIP toolchain/torch consistency checks.
 
 Standalone on purpose: setup.py and the kernel-cache build backend load this
 file by path, so it must not import the freetoken package.
@@ -16,6 +16,11 @@
 _TRUE_VALUES = {"1", "true", "yes", "on"}
 
 
+def _is_rocm() -> bool:
+    import torch
+    return getattr(torch.version, "hip", None) is not None
+
+
 def _nvcc_path() -> str | None:
     from torch.utils.cpp_extension import CUDA_HOME
 
@@ -49,6 +54,8 @@ def check_nvcc_matches_torch() -> None:
     nvcc-built binaries link libcudart.so.<nvcc major>; at runtime only the
     torch wheel's own CUDA runtime is guaranteed to be loadable.
     """
+    if _is_rocm():
+        return  # ROCm uses hipcc, not nvcc
     if os.getenv(ALLOW_MISMATCH_ENV, "").strip().lower() in _TRUE_VALUES:
         return
     torch_major = torch_cuda_major()

python/freetoken/kernel/backend.py

@@ -47,6 +47,24 @@ def device_capability() -> tuple[int, int]:
     return (int(major), int(minor))
 
 
+@functools.cache
+def is_rocm() -> bool:
+    """True when torch is built for ROCm (AMD GPU)."""
+    import torch
+    return getattr(torch.version, "hip", None) is not None
+
+
+@functools.cache
+def driver_hip_version() -> int | None:
+    """ROCm driver version, or None if undetermined."""
+    # TODO(ROCm): flashinfer/sgl_kernel have no ROCm builds — Triton fallback is used.
+    try:
+        from freetoken.kernel.pinned import _load_pinned_extension
+        return int(_load_pinned_extension().driver_cuda_version()) or None
+    except Exception:
+        return None
+
+
 @functools.cache
 def driver_cuda_version() -> int | None:
     """Max CUDA version the installed NVIDIA driver supports (``13000`` == CUDA 13.0),

python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp

@@ -29,7 +29,7 @@
 #include <thread>
 #include <vector>
 
-#include <cuda_runtime_api.h>
+#include <freetoken/hip_compat.h>
 #include <torch/extension.h>
 
 #if defined(__linux__)