From 974164632d5980e09db1432c654f2e816c298066 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Tue, 25 Aug 2026 19:15:39 -0300 Subject: [PATCH 1/6] feat(rocm): AMD ROCm support (gfx1100) + qwen35moe GGUF loader Bring FreeToken up on AMD ROCm (RX 7900 XTX / gfx1100) alongside CUDA via a thin device seam. Includes: - Device detection / architecture gating (is_rocm), build toolchain for tvm-ffi/HIP JIT, --offload-arch=gfx1100 pinning, and HIP backend dispatch. - Pinned-memory + graph-capture gating (ROCm settles to kernel-launch decode). - device_api.h seam, GGUF kernel HIP port, quant mapping, nvfp4->mxfp4. - qwen3_5_moe GGUF loader (config/dense weights/expert offload). - fix(gguf): de-interleave the GDN mrope_interleaved value heads (in_proj_qkvz v/z rows, out_proj cols, in_proj_ba, conv1d v-channels, dt_bias; A_log stored as A=-exp(A_log)) so weights match HF and the model serves correct text. - torch attention backend (ground-truth reference) + ROCm regression tests. - AOT CI (rocm.yml) and docs/install-amd.md. Tests: ROCm suite passes; qwen35voe de-interleave + torch-backend tests added. --- .github/workflows/rocm.yml | 70 +++ .gitignore | 6 + docs/install-amd.md | 83 +++ freetoken-kernel-cache/build_backend.py | 18 +- pyproject.toml | 8 + python/freetoken/attention/__init__.py | 13 + python/freetoken/attention/torch.py | 217 ++++++++ python/freetoken/engine/engine.py | 26 +- python/freetoken/engine/graph.py | 16 +- python/freetoken/kernel/_toolchain.py | 86 +++ python/freetoken/kernel/backend.py | 29 + python/freetoken/kernel/batch_memcpy.py | 7 +- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 32 +- python/freetoken/kernel/csrc/gguf/dispatch.h | 8 + .../freetoken/kernel/csrc/gguf/ggml-common.h | 7 + .../freetoken/kernel/csrc/gguf/gguf_kernel.cu | 38 +- .../csrc/include/freetoken/device_api.h | 99 ++++ .../kernel/csrc/include/freetoken/utils.cuh | 94 ++++ .../kernel/csrc/jit/batch_memcpy.cuh | 64 ++- .../kernel/csrc/jit/fast_index_copy.cuh | 56 +- .../freetoken/kernel/csrc/pinned_tensor.cpp | 70 ++- python/freetoken/kernel/gguf.py | 34 +- python/freetoken/kernel/tinygrad_fallback.py | 92 ++++ python/freetoken/kernel/triton/activation.py | 31 +- python/freetoken/kernel/triton/norm.py | 34 +- python/freetoken/kernel/utils.py | 83 ++- python/freetoken/layers/moe.py | 10 + python/freetoken/models/gguf/config.py | 1 + python/freetoken/models/gguf/dequant.py | 69 +++ python/freetoken/models/gguf/tokenizer.py | 5 +- .../freetoken/models/qwen3_5_moe/__init__.py | 12 + python/freetoken/models/qwen3_5_moe/gdn.py | 14 +- python/freetoken/models/qwen3_5_moe/gguf.py | 512 ++++++++++++++++++ python/freetoken/models/qwen3_5_moe/model.py | 7 + python/freetoken/models/register.py | 8 + python/freetoken/models/weight.py | 14 + python/freetoken/moe/expert_banks.py | 20 + python/freetoken/moe/fused_gguf.py | 52 ++ python/freetoken/moe/nvfp4_backends.py | 12 + python/freetoken/moe/nvfp4_to_mxfp4.py | 239 ++++++++ python/freetoken/moe/offload_cache.py | 6 + python/freetoken/server/args.py | 19 + python/freetoken/utils/__init__.py | 8 + python/freetoken/utils/arch.py | 100 +++- python/freetoken/utils/graph_gate.py | 189 +++++++ python/freetoken/utils/torch_utils.py | 19 + setup.py | 80 ++- tests/attention/test_torch_backend.py | 62 +++ tests/engine/test_attention_backend_rocm.py | 75 +++ tests/kernels/test_backend_rocm.py | 63 +++ tests/kernels/test_cache_rocm_pairing.py | 109 ++++ tests/kernels/test_toolchain_hip.py | 84 +++ tests/models/test_qwen35moe_gguf_deint.py | 71 +++ tests/moe/test_nvfp4_backends_rocm.py | 95 ++++ tests/moe/test_nvfp4_to_mxfp4.py | 111 ++++ tests/utils/test_device_kind.py | 83 +++ 56 files changed, 3379 insertions(+), 91 deletions(-) create mode 100644 .github/workflows/rocm.yml create mode 100644 docs/install-amd.md create mode 100644 python/freetoken/attention/torch.py create mode 100644 python/freetoken/kernel/csrc/include/freetoken/device_api.h create mode 100644 python/freetoken/kernel/tinygrad_fallback.py create mode 100644 python/freetoken/models/qwen3_5_moe/gguf.py create mode 100644 python/freetoken/moe/fused_gguf.py create mode 100644 python/freetoken/moe/nvfp4_to_mxfp4.py create mode 100644 python/freetoken/utils/graph_gate.py create mode 100644 tests/attention/test_torch_backend.py create mode 100644 tests/engine/test_attention_backend_rocm.py create mode 100644 tests/kernels/test_backend_rocm.py create mode 100644 tests/kernels/test_cache_rocm_pairing.py create mode 100644 tests/kernels/test_toolchain_hip.py create mode 100644 tests/models/test_qwen35moe_gguf_deint.py create mode 100644 tests/moe/test_nvfp4_backends_rocm.py create mode 100644 tests/moe/test_nvfp4_to_mxfp4.py create mode 100644 tests/utils/test_device_kind.py diff --git a/.github/workflows/rocm.yml b/.github/workflows/rocm.yml new file mode 100644 index 00000000..e2f1fe49 --- /dev/null +++ b/.github/workflows/rocm.yml @@ -0,0 +1,70 @@ +name: ROCm (AMD) correctness smoke + +# ROCm CI job: compiles the AOT kernel cache for the RX 7000 (gfx1100) target on a +# ROCm torch install and runs a torch-free correctness smoke plus the AMD unit tests. +# Gated on a self-hosted runner that has ROCm torch + hipcc. The primary NVIDIA release +# flow is release.yml; this job is additive and must not gate NVIDIA releases. + +on: + workflow_dispatch: + push: + branches: [main] + pull_request: + +jobs: + rocm-smoke: + runs-on: [self-hosted, linux, amd, rocm] + timeout-minutes: 60 + env: + FREETOKEN_DISABLE_JIT: "1" + FREETOKEN_KERNEL_CACHE_GFX: "gfx1100" + steps: + - uses: actions/checkout@v4 + + - name: Check ROCm toolchain + run: | + set -e + command -v hipcc || ls /opt/rocm/bin/hipcc + "${PYTHON:-python3}" -c "import torch.version as v; print('torch hip:', v.hip)" + + - name: Install build deps + run: | + python -m pip install --upgrade pip wheel setuptools + python -m pip install -e "python[rocm]" + + - name: Compile AOT kernel cache for gfx1100 + run: | + FREETOKEN_KERNEL_CACHE_VERBOSE=1 python -m pip wheel ./freetoken-kernel-cache -w dist/cache-rocm + + - name: Install prebuilt kernel cache + run: | + whl="$(find dist/cache-rocm -name 'freetoken_kernel_cache-*.whl' | head -1)" + python -m pip install --force-reinstall "$whl" + + - name: Torch-free AMD unit tests + run: | + python -m pytest \ + tests/utils/test_device_kind.py \ + tests/kernels/test_toolchain_hip.py \ + tests/kernels/test_backend_rocm.py \ + tests/kernels/test_cache_rocm_pairing.py \ + tests/moe/test_nvfp4_to_mxfp4.py \ + -q + + - name: Hardware correctness smoke (serves on RX 7000) + run: | + # Functional path only -- flashinfer/sgl/trtllm are NVIDIA-only and must not + # be selected. AUTO backend must resolve to triton; NVFP4 auto -> triton. + python - <<'PY' + from freetoken.utils.arch import is_rocm, is_gfx_arch_ge + from freetoken.moe.nvfp4_backends import select_nvfp4_backend + import torch + assert is_rocm(), "expected a ROCm torch build" + assert is_gfx_arch_ge(1100), "expected gfx1100-class device (RX 7000)" + print("NVFP4 auto ->", select_nvfp4_backend(torch.device("cuda"), 768, "auto")) + PY + + - name: Serve smoke + run: | + FREETOKEN_DEVICE=cuda python -m freetoken.serve --help >/dev/null \ + && echo "freetoken CLI loads on ROCm" diff --git a/.gitignore b/.gitignore index bf804e07..95757625 100644 --- a/.gitignore +++ b/.gitignore @@ -227,3 +227,9 @@ 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 diff --git a/docs/install-amd.md b/docs/install-amd.md new file mode 100644 index 00000000..62ecf2fe --- /dev/null +++ b/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. See `.plans/amd-gpu-support/plan.md`. + +## 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 (Inc 3) | +| 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 Inc-1 gate 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). diff --git a/freetoken-kernel-cache/build_backend.py b/freetoken-kernel-cache/build_backend.py index 03f17a31..cf255bec 100644 --- a/freetoken-kernel-cache/build_backend.py +++ b/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" ) diff --git a/pyproject.toml b/pyproject.toml index 8bd653f8..13b52c19 100644 --- a/pyproject.toml +++ b/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+rocm6.2` from +# https://download.pytorch.org/whl/rocm6.2) so the native extensions build against the HIP +# runtime; this extra pins the rest. See docs/install.md (AMD section, Inc 9). +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 diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index 746c04c4..0dfb2197 100644 --- a/python/freetoken/attention/__init__.py +++ b/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})), diff --git a/python/freetoken/attention/torch.py b/python/freetoken/attention/torch.py new file mode 100644 index 00000000..6ae16c2e --- /dev/null +++ b/python/freetoken/attention/torch.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, List + +import torch + +from freetoken.core import Batch, get_global_ctx + +from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata + +if TYPE_CHECKING: + from freetoken.models import ModelConfig + + +@dataclass +class TorchMetadata(BaseAttnMetadata): + """Minimal contiguous-gather metadata for the pure-torch backend. + + ``indices`` maps every logical KV position (across all padded requests, in + ``seqlens_k`` order) to its physical paged-cache slot, exactly like the triton + backend's gather. The torch backend reads the SAME paged cache as triton, so a + triton-vs-torch logit difference isolates the attention *compute* from the + cache addressing. + """ + + indices: torch.Tensor + seqlens_q: List[int] + seqlens_k: List[int] + cached_lens: List[int] + is_decode: bool + cu_seqlens_q: torch.Tensor + + def get_last_indices(self, bs: int) -> torch.Tensor: + return self.cu_seqlens_q[1 : 1 + bs] - 1 + + +class TorchAttentionBackend(BaseAttnBackend): + """Pure-PyTorch full-attention backend (no Triton/CUDA kernels). + + Serves as a numerically-explicit ground truth for debugging the hybrid + qwen35moe model. It stores K/V into the same paged MHAKVCache as + ``TritonAttentionBackend``, gathers the request's full K/V history via the + identical ``indices`` page gather, and computes GQA softmax attention with + PyTorch ops so every intermediate is auditable. + + Intended for correctness debugging / backend A-B comparison, not production + serving. Registered as the ``"torch"`` attention backend (``AttnType.FULL``). + """ + + def __init__(self, config: ModelConfig): + self.config = config + self.kvcache = get_global_ctx().kv_cache + self.device = self.kvcache.device + self.num_q_heads = int(getattr(config, "num_qo_heads", 1)) + self.num_kv_heads = int(getattr(config, "num_kv_heads", 1)) + self.head_dim = int(getattr(config, "head_dim", 1)) + # Prefer the full-attention group spec head_dim (authoritative for kv heads). + specs = getattr(config, "kv_cache_group_specs", lambda: ())() + for spec in specs: + name = getattr(spec, "attn_type", None) + if name is not None and str(name) == "AttnType.FULL": + self.head_dim = int(getattr(spec, "head_dim", self.head_dim)) + self.num_kv_heads = int(getattr(spec, "num_kv_heads", self.num_kv_heads)) + break + # Debugging: contiguous (per-request) cache instead of the paged pool, to + # isolate cache addressing from the attention compute (Inc 5). + self._contig: dict[tuple[int, int], list] = {} + import os + + self._use_contig = os.environ.get("FT_DEBUG_CONTIG_CACHE") == "1" + + def _build_metadata(self, batch: Batch) -> TorchMetadata: + ctx = get_global_ctx() + page_table = ctx.page_table + reqs = batch.padded_reqs + seqlens_q = [req.extend_len for req in reqs] + seqlens_k = [req.device_len for req in reqs] + cached_lens = [req.cached_len for req in reqs] + is_decode = max(seqlens_q) == 1 + indices = torch.cat([page_table[req.table_idx, : req.device_len] for req in reqs]) + if is_decode: + cu_seqlens_q = torch.arange(0, len(reqs) + 1, dtype=torch.int32, device=self.device) + else: + cu_seqlens_q = torch.tensor( + [0] + seqlens_q, dtype=torch.int32, device=self.device + ).cumsum_(0) + return TorchMetadata( + indices=indices, + seqlens_q=seqlens_q, + seqlens_k=seqlens_k, + cached_lens=cached_lens, + is_decode=is_decode, + cu_seqlens_q=cu_seqlens_q, + ) + + def prepare_metadata(self, batch: Batch) -> None: + batch.attn_metadata = self._build_metadata(batch) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer_id: int, + batch: Batch, + attn_spec: AttentionSpec | None = None, + ) -> torch.Tensor: + if self._use_contig: + return self._forward_contig(q, k, v, layer_id, batch, attn_spec) + self.kvcache.store_kv(k, v, batch.out_loc, layer_id) + + k_raw = self.kvcache.k_cache(layer_id) + v_raw = self.kvcache.v_cache(layer_id) + kv_heads, head_dim = k_raw.shape[-2], k_raw.shape[-1] + assert head_dim == q.shape[-1], f"head_dim {head_dim} != {q.shape[-1]}" + k_cache = k_raw.view(-1, kv_heads, head_dim) + v_cache = v_raw.view(-1, kv_heads, head_dim) + + metadata = batch.attn_metadata + assert isinstance(metadata, TorchMetadata) + k_all = k_cache[metadata.indices] # [total_kv, kv_heads, head_dim] + v_all = v_cache[metadata.indices] # [total_kv, kv_heads, head_dim] + + spec = attn_spec or AttentionSpec() + scale = spec.sm_scale if spec.sm_scale is not None else head_dim ** -0.5 + group = self.num_q_heads // kv_heads + + num_q_tokens = q.shape[0] + out = torch.empty((num_q_tokens, self.num_q_heads, head_dim), dtype=q.dtype, device=q.device) + q_off = 0 + k_off = 0 + for lq, lk, cached in zip( + metadata.seqlens_q, metadata.seqlens_k, metadata.cached_lens + ): + qs = q[q_off : q_off + lq] # [lq, num_q, head_dim] + ks = k_all[k_off : k_off + lk] # [lk, kv_heads, head_dim] + vs = v_all[k_off : k_off + lk] + if group > 1: + ks = ks.repeat_interleave(group, dim=1) # [lk, num_q, head_dim] + vs = vs.repeat_interleave(group, dim=1) + # [num_q, lq, lk] + scores = torch.einsum("qhd,khd->hqk", qs.float(), ks.float()) * scale + # causal: query i (global cached+i) attends key col j <= cached+i + if lq > 1 or lk > lq: + rows = torch.arange(lq, device=scores.device) + cols = torch.arange(lk, device=scores.device) + masked = (cols[None, :] > (cached + rows)[:, None]).to(scores.device) + scores = scores.masked_fill(masked[None, :, :], float("-inf")) + probs = torch.softmax(scores, dim=-1) + o = torch.einsum("hqk,khd->qhd", probs, vs.float()).to(q.dtype) # [lq, num_q, head_dim] + out[q_off : q_off + lq] = o + q_off += lq + k_off += lk + + return out + + def _forward_contig(self, q, k, v, layer_id, batch, attn_spec=None): + """Contiguous (non-paged) attention: K/V accumulate per (layer, request) in + a Python list keyed by logical position, so cache addressing is trivially + correct. Isolates the paged-cache addressing from the attention compute.""" + metadata = self._build_metadata(batch) + kv_heads = self.num_kv_heads + head_dim = self.head_dim + spec = attn_spec or AttentionSpec() + scale = spec.sm_scale if spec.sm_scale is not None else head_dim ** -0.5 + group = self.num_q_heads // kv_heads + # Store this forward's K/V rows per request (append in global position order). + q_off = 0 + for i, lq in enumerate(metadata.seqlens_q): + uid = batch.padded_reqs[i].uid + buf = self._contig.setdefault((layer_id, uid), {"k": [], "v": []}) + kseg = k[q_off : q_off + lq].view(lq, kv_heads, head_dim) + vseg = v[q_off : q_off + lq].view(lq, kv_heads, head_dim) + for t in range(lq): + buf["k"].append(kseg[t]) + buf["v"].append(vseg[t]) + q_off += lq + # compute attention from the contiguous cache + q_off = 0 + out = torch.empty( + (q.shape[0], self.num_q_heads, head_dim), dtype=q.dtype, device=q.device + ) + for i, lq in enumerate(metadata.seqlens_q): + uid = batch.padded_reqs[i].uid + buf = self._contig[(layer_id, uid)] + ks = torch.stack(buf["k"]) # [acc, kv_heads, head_dim] + vs = torch.stack(buf["v"]) + acc = ks.shape[0] + cached = acc - lq + qs = q[q_off : q_off + lq] + if group > 1: + ks = ks.repeat_interleave(group, dim=1) + vs = vs.repeat_interleave(group, dim=1) + scores = torch.einsum("qhd,khd->hqk", qs.float(), ks.float()) * scale + rows = torch.arange(lq, device=scores.device) + cols = torch.arange(acc, device=scores.device) + masked = (cols[None, :] > (cached + rows)[:, None]).to(scores.device) + scores = scores.masked_fill(masked[None, :, :], float("-inf")) + probs = torch.softmax(scores, dim=-1) + o = torch.einsum("hqk,khd->qhd", probs, vs.float()).to(q.dtype) + out[q_off : q_off + lq] = o + q_off += lq + return out + + def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + # ROCm/graph capture is disabled for this debugging backend; no-op. + return None + + def prepare_for_capture(self, batch: Batch) -> None: + return None + + def prepare_for_replay(self, batch: Batch) -> None: + return None + + +__all__ = ["TorchAttentionBackend", "TorchMetadata"] diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 22c4a6c7..5ada3954 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -15,7 +15,7 @@ from freetoken.moe import create_moe_backend, is_offload_moe_backend from freetoken.moe.expert_banks import load_expert_banks from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache -from freetoken.utils import align_ceil, init_logger, is_sm90_family, is_sm100_family, mem_GB, torch_dtype +from freetoken.utils import align_ceil, device_kind, init_logger, is_rocm, is_sm90_family, is_sm100_family, mem_GB, torch_dtype from .config import EngineConfig from .graph import GraphRunner, get_free_memory @@ -50,6 +50,11 @@ def _flashinfer_available() -> bool: def _sgl_flash_attn_available() -> bool: + from freetoken.utils.arch import is_rocm + + # sgl_kernel is NVIDIA-only; never select it on ROCm even if a stray copy is importable. + if is_rocm(): + return False try: from sgl_kernel.flash_attn import flash_attn_with_kvcache # noqa: F401 except Exception as exc: @@ -101,6 +106,14 @@ def _backend_parts_serve(name: str, required: frozenset[AttnType]) -> bool: def _backend_requirements_met(name: str) -> bool: + # On ROCm (AMD) only the portable backends exist: flashinfer/sgl/trtllm (and anything + # sm_100-gated) are NVIDIA-only, so short-circuit before probing them at all. + from freetoken.utils.arch import is_rocm + + if is_rocm(): + return all(not i.requires_flashinfer and not i.requires_sgl_kernel + and not i.requires_sm100 for i in + [attention_backend_info(p) for p in name.split(",")]) # flashinfer first across ALL parts: the sgl probe logs a "falls back to fi" warning, # which would mislead when the candidate is about to fail on flashinfer anyway. infos = [attention_backend_info(part) for part in name.split(",")] @@ -202,6 +215,13 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att # explicit --attention-backend choices. for part in backend_parts: info = attention_backend_info(part) + from freetoken.utils.arch import is_rocm + + if is_rocm() and (info.requires_flashinfer or info.requires_sgl_kernel or info.requires_sm100): + raise RuntimeError( + f"Attention backend {config.attention_backend!r} is NVIDIA-only and " + f"unavailable on this ROCm (AMD) build; use --attention-backend triton." + ) if info.requires_flashinfer and not _flashinfer_available(): raise RuntimeError( f"Attention backend {config.attention_backend!r} requires flashinfer, which is " @@ -298,6 +318,7 @@ def __init__(self, config: EngineConfig): self.device = torch.device(f"cuda:{config.tp_info.rank}") torch.cuda.set_device(self.device) + logger.info_rank0(f"device_kind={device_kind()} backend={self.device}") torch.manual_seed(42) self.stream = torch.cuda.Stream() torch.cuda.set_stream(self.stream) @@ -1011,6 +1032,9 @@ def _ensure_expandable_segments() -> None: """ if os.environ.get("PYTORCH_ALLOC_CONF") or os.environ.get("PYTORCH_CUDA_ALLOC_CONF"): return + if is_rocm(): + # expandable_segments is a CUDA allocator setting with no ROCm analogue; skip it. + return try: torch.cuda.memory._set_allocator_settings("expandable_segments:True") except Exception as exc: # pragma: no cover - depends on torch build diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 4f202502..50bac2f0 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -132,6 +132,18 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel # graphs-disabled early return so that config gets the phase too. emit_progress("Capturing CUDA graphs / warming up", 0, 0) self.graph_map: Dict[int, torch.cuda.CUDAGraph] = {} + # Inc-8 parity: on ROCm, honour the Inc-1 graph-gate result. If capture is not + # viable on this AMD card, skip graphs entirely so decode uses the kernel-launch + # path (correct, just not graph-accelerated) rather than erroring mid-capture. + from freetoken.utils.arch import is_rocm + from freetoken.utils.graph_gate import graph_capture_status + + if is_rocm() and graph_capture_status() == "fail": + logger.info_rank0( + "AMD ROCm build: HIP graph capture gate FAILED on this device; " + "using the kernel-launch decode path (CUDA graphs disabled)." + ) + return None if self.max_graph_bs == 0: return logger.info_rank0("CUDA graph is disabled.") @@ -187,7 +199,9 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel logger.info_rank0(f"Free GPU memory after capturing CUDA graphs: {mem_GB(free_memory)}") def can_use_cuda_graph(self, batch: Batch) -> bool: - return batch.is_decode and batch.size <= self.max_graph_bs + # ``self.graph_map`` is empty when graphs were skipped (ROCm graph-gate fail or + # disabled); decode must then fall back to the kernel-launch path. + return bool(self.graph_map) and batch.is_decode and batch.size <= self.max_graph_bs def replay(self, batch: Batch) -> torch.Tensor: assert self.can_use_cuda_graph(batch) diff --git a/python/freetoken/kernel/_toolchain.py b/python/freetoken/kernel/_toolchain.py index b49cebbb..eccb9635 100644 --- a/python/freetoken/kernel/_toolchain.py +++ b/python/freetoken/kernel/_toolchain.py @@ -16,6 +16,92 @@ _TRUE_VALUES = {"1", "true", "yes", "on"} +def _hipcc_path() -> str | None: + """Locate hipcc: $ROCM_HOME/bin/hipcc, $HIP_PATH/bin/hipcc, /opt/rocm/bin/hipcc, + then PATH.""" + for env in ("ROCM_HOME", "HIP_PATH"): + root = os.getenv(env) + if root: + candidate = os.path.join(root, "bin", "hipcc") + if os.path.isfile(candidate): + return candidate + default = "/opt/rocm/bin/hipcc" + if os.path.isfile(default): + return default + return shutil.which("hipcc") + + +def hip_hip_version(hipcc: str) -> tuple[int, int] | None: + """HIP toolkit version from ``hipcc --version`` (e.g. (6, 2)), or None.""" + try: + proc = subprocess.run( + [hipcc, "--version"], capture_output=True, text=True, check=True + ) + except (OSError, subprocess.CalledProcessError): + return None + # hipcc --version prints e.g. "HIP version: 6.2.41000" (or a clang version line). + m = re.search(r"HIP version[:\s]+(\d+)\.(\d+)", proc.stdout) + if m: + return int(m.group(1)), int(m.group(2)) + m = re.search(r"(\d+)\.(\d+)\.\d+", proc.stdout) + if m: + return int(m.group(1)), int(m.group(2)) + return None + + +def torch_hip_version() -> str | None: + """The ``torch.version.hip`` string (e.g. "6.2.4100000"), or None on non-ROCm torch.""" + try: + import torch + + return getattr(torch.version, "hip", None) + except Exception: + return None + + +def is_rocm_torch() -> bool: + """True when the installed torch is a ROCm (AMD) build.""" + return bool(torch_hip_version()) + + +def torch_hip_major() -> int | None: + hip = torch_hip_version() + if not hip: + return None + m = re.match(r"(\d+)", hip) + return int(m.group(1)) if m else None + + +def check_hip_matches_torch() -> None: + """Refuse to hipcc-compile kernels across HIP major versions. + + Mirrors check_nvcc_matches_torch: hipcc-built kernels link against the HIP runtime + major they were built with; at runtime only the torch wheel's own HIP runtime is + guaranteed to be loadable. No-op when torch is not ROCm. + """ + if os.getenv(ALLOW_MISMATCH_ENV, "").strip().lower() in _TRUE_VALUES: + return + if not is_rocm_torch(): + return + torch_major = torch_hip_major() + hipcc = _hipcc_path() + if hipcc is None: + raise RuntimeError( + "ROCm torch detected but no hipcc found. Install a ROCm/HIP toolkit " + "(e.g. via /opt/rocm) matching torch's HIP version, or set " + f"{ALLOW_MISMATCH_ENV}=1 to override." + ) + release = hip_hip_version(hipcc) + if release is None: + return + if release[0] != torch_major: + raise RuntimeError( + f"hipcc {release[0]}.{release[1]} would build kernels linking HIP " + f"{release[0]}.x, but torch ships HIP {torch_hip_version()}. Install a " + f"ROCm {torch_major}.x toolkit, or set {ALLOW_MISMATCH_ENV}=1 to override." + ) + + def _nvcc_path() -> str | None: from torch.utils.cpp_extension import CUDA_HOME diff --git a/python/freetoken/kernel/backend.py b/python/freetoken/kernel/backend.py index 3037ad8d..7b7629b1 100644 --- a/python/freetoken/kernel/backend.py +++ b/python/freetoken/kernel/backend.py @@ -10,6 +10,14 @@ import functools import importlib.util +from freetoken.utils.arch import is_rocm + + +# NVIDIA-only optional native packages: even if an importable copy is present on a ROCm +# torch build (e.g. a stray CUDA wheel), they must not be used -- the runtime falls back +# to the portable Triton kernels. Treated as unavailable on ROCm. +_CUDA_ONLY_PACKAGES = frozenset({"flashinfer", "sgl_kernel", "triton_kernels"}) + def _importable(name: str) -> bool: # find_spec normally returns None when a package is absent, but it can raise @@ -21,13 +29,32 @@ def _importable(name: str) -> bool: return False +def is_native_cuda_available() -> bool: + """True when the current torch build is CUDA and a CUDA-capable device is present. + False on ROCm and CPU builds. Used to gate NVIDIA-native ops/paths.""" + from freetoken.utils.arch import device_kind + + if device_kind() != "cuda": + return False + try: + import torch + + return bool(torch.cuda.is_available()) + except Exception: + return False + + @functools.cache def is_flashinfer_installed() -> bool: + if is_rocm(): + return False return _importable("flashinfer") @functools.cache def is_sgl_kernel_installed() -> bool: + if is_rocm(): + return False return _importable("sgl_kernel") @@ -39,6 +66,8 @@ def is_triton_kernels_installed() -> bool: source tree and has no Windows wheel. It is also not one of the six ops ``freetoken.kernel.triton`` reimplements, so its call-site carries its own fallback. """ + if is_rocm(): + return False return _importable("triton_kernels") diff --git a/python/freetoken/kernel/batch_memcpy.py b/python/freetoken/kernel/batch_memcpy.py index b39e5cde..0466333c 100644 --- a/python/freetoken/kernel/batch_memcpy.py +++ b/python/freetoken/kernel/batch_memcpy.py @@ -42,10 +42,15 @@ def _probe(fn) -> None: def load_batch_memcpy(): """Build (once), probe, and return the batch-memcpy entry point, or raise. - The 8-argument cudaMemcpyBatchAsync signature this binding uses is CUDA 13.0's + On ROCm the HIP per-copy grid kernel is used (no CUDA version gate). On CUDA the + 8-argument cudaMemcpyBatchAsync signature this binding uses is CUDA 13.0's (12.8/12.9 had an extra failIdx parameter); gate on the torch runtime version before paying for the JIT build, then verify with a real copy. """ + if torch.version.hip is not None: + fn = _jit_batch_memcpy_module().batch_memcpy + _probe(fn) + return fn cuda = torch.version.cuda if cuda is None or tuple(int(x) for x in cuda.split(".")[:2]) < (13, 0): raise RuntimeError(f"cudaMemcpyBatchAsync binding requires CUDA >= 13.0 (torch built with {cuda})") diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 880e8637..024bf5a0 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -29,9 +29,31 @@ #include #include +#if defined(USE_HIP) +#include +// HIP host functions carry no special calling convention; define the CUDA host-func +// marker as empty so the static callback signatures below compile unchanged. +#ifndef CUDART_CB +#define CUDART_CB +#endif +#else #include +#endif #include +// Stream-sync / host-func node launch, shared by both backends. +#if defined(USE_HIP) +#define CPU_MOE_STREAM_SYNC(s) \ + hipStreamSynchronize(reinterpret_cast(s)) +#define CPU_MOE_LAUNCH_HOST_FUNC(s, fn, data) \ + hipLaunchHostFunc(reinterpret_cast(s), (fn), (data)) +#else +#define CPU_MOE_STREAM_SYNC(s) \ + cudaStreamSynchronize(reinterpret_cast(s)) +#define CPU_MOE_LAUNCH_HOST_FUNC(s, fn, data) \ + cudaLaunchHostFunc(reinterpret_cast(s), (fn), (data)) +#endif + #if defined(__linux__) #include #include @@ -614,7 +636,7 @@ static bool cumemops_probe(uintptr_t stream, uintptr_t scratch_addr) { auto* s = reinterpret_cast(stream); if (g_cu_write64(s, (unsigned long long)scratch_addr, 7ULL, kCuWriteDefault) != 0) return false; if (g_cu_wait64(s, (unsigned long long)scratch_addr, 7ULL, kCuWaitValueGeq) != 0) return false; - return cudaStreamSynchronize(reinterpret_cast(stream)) == cudaSuccess; + return CPU_MOE_STREAM_SYNC(stream) == 0; } // GPU side of the flag handshake (see the block comment above): enqueued on the @@ -1958,13 +1980,13 @@ struct CpuMoeExecutor { } void submit_with_cuda_stream(uintptr_t stream, uintptr_t task) { - cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::submit_cb, - reinterpret_cast(task)); + CPU_MOE_LAUNCH_HOST_FUNC(stream, &CpuMoeExecutor::submit_cb, + reinterpret_cast(task)); } void sync_with_cuda_stream(uintptr_t stream, uintptr_t task) { - cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::sync_cb, - reinterpret_cast(task)); + CPU_MOE_LAUNCH_HOST_FUNC(stream, &CpuMoeExecutor::sync_cb, + reinterpret_cast(task)); } // Register a (layer, batch-size) slot's task so the coordinator can dispatch it on a diff --git a/python/freetoken/kernel/csrc/gguf/dispatch.h b/python/freetoken/kernel/csrc/gguf/dispatch.h index f42a2163..d46469ef 100644 --- a/python/freetoken/kernel/csrc/gguf/dispatch.h +++ b/python/freetoken/kernel/csrc/gguf/dispatch.h @@ -11,6 +11,13 @@ #endif // Warp-shuffle wrappers the donor pulls from sgl-kernel's utils.h (CUDA variants). +// On ROCm the mask must be 64-bit (HIP static-asserts 32-bit promotion is an error). +#if defined(USE_ROCM) +#define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) \ + __shfl_xor_sync(static_cast(mask), (var), (lane_mask)) +#define SGLANG_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \ + __shfl_xor_sync(static_cast(mask), (var), (lane_mask), (width)) +#else #ifndef SGLANG_SHFL_XOR_SYNC #define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) __shfl_xor_sync((mask), (var), (lane_mask)) #endif @@ -18,6 +25,7 @@ #define SGLANG_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \ __shfl_xor_sync((mask), (var), (lane_mask), (width)) #endif +#endif #define DISPATCH_CASE_FLOAT_TYPES(...) \ AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ diff --git a/python/freetoken/kernel/csrc/gguf/ggml-common.h b/python/freetoken/kernel/csrc/gguf/ggml-common.h index 88c21a4a..5822eb3c 100644 --- a/python/freetoken/kernel/csrc/gguf/ggml-common.h +++ b/python/freetoken/kernel/csrc/gguf/ggml-common.h @@ -10,6 +10,13 @@ #define GGML_CUDA_DMMV_X 32 #define GGML_CUDA_MMV_Y 1 +#if defined(USE_ROCM) +// ROCm shim: the vendored GGUF launchers (moe.cuh/mmvq.cuh/...) take a CUDA-style +// stream parameter. Map the CUDA stream type to HIP so those signatures compile +// unmodified under USE_ROCM. The including .cu pulls in hip/hip_runtime.h first. +using cudaStream_t = hipStream_t; +#endif + // Data Structures // QK = number of values after dequantization // QR = QK / number of values before dequantization diff --git a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu index d88960d5..09c83c6f 100644 --- a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu +++ b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu @@ -1,10 +1,28 @@ // Adatped from // https://github.com/vllm-project/vllm/blob/755ed7b05be4743237d3339c4ff8c22bcaae04f4/csrc/quantization/gguf/gguf_kernel.cu +#if defined(USE_ROCM) +// ROCm torch hipifies these headers into c10::cuda (masquerading-as-CUDA), providing +// c10::cuda::OptionalCUDAGuard / getCurrentCUDAStream backed by HIP. c10/cuda/CUDAGuard.h +// itself is not directly includable on ROCm (missing a generated header). +#include +#include +#include +#include +#else #include #include #include +#endif #include +#if defined(USE_ROCM) +#define GGUF_DEVICE_GUARD(device) c10::cuda::OptionalCUDAGuard device_guard(device) +#define GGUF_CURRENT_STREAM() c10::cuda::getCurrentCUDAStream() +#else +#define GGUF_DEVICE_GUARD(device) at::cuda::OptionalCUDAGuard device_guard(device) +#define GGUF_CURRENT_STREAM() at::cuda::getCurrentCUDAStream() +#endif + // dont use clang-format here, it breaks the include order // clang-format off #include "dispatch.h" @@ -77,11 +95,11 @@ torch::Tensor ggml_dequantize( int64_t m, int64_t n, std::optional const& dtype) { - const at::cuda::OptionalCUDAGuard device_guard(device_of(W)); + const GGUF_DEVICE_GUARD(device_of(W)); auto dtype_ = dtype.value_or(torch::kFloat16); auto options = torch::TensorOptions().dtype(dtype_).device(W.device()); at::Tensor DW = torch::empty({m, n}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); DISPATCH_FLOAT_TYPES(DW.scalar_type(), "ggml_dequantize", [&] { auto to_cuda = ggml_get_to_cuda(type); @@ -99,10 +117,10 @@ torch::Tensor ggml_mul_mat_vec_a8( int col = X.sizes()[1]; int vecs = X.sizes()[0]; const int padded = (col + 512 - 1) / 512 * 512; - const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); + const GGUF_DEVICE_GUARD(device_of(X)); auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); at::Tensor Y = torch::empty({vecs, row}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); at::Tensor quant_X = torch::empty({vecs, padded / 32 * 9}, options); DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_mul_mat_vec_a8", [&] { @@ -197,10 +215,10 @@ torch::Tensor ggml_mul_mat_a8( int col = X.sizes()[1]; int padded = (col + 512 - 1) / 512 * 512; int batch = X.sizes()[0]; - const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); + const GGUF_DEVICE_GUARD(device_of(X)); auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); at::Tensor Y = torch::empty({batch, row}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); at::Tensor quant_X = torch::empty({batch, padded / 32 * 9}, options); DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_mul_mat_a8", [&] { @@ -344,10 +362,10 @@ torch::Tensor ggml_moe_a8( int64_t tokens) { int col = X.sizes()[1]; int padded = (col + 512 - 1) / 512 * 512; - const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); + const GGUF_DEVICE_GUARD(device_of(X)); auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); at::Tensor Y = torch::empty({tokens * top_k, row}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options); DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_a8", [&] { @@ -548,10 +566,10 @@ torch::Tensor ggml_moe_a8_vec( int64_t tokens) { int col = X.sizes()[1]; const int padded = (col + 512 - 1) / 512 * 512; - const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); + const GGUF_DEVICE_GUARD(device_of(X)); auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); at::Tensor Y = torch::zeros({tokens * top_k, row}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options); DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_vec_a8", [&] { diff --git a/python/freetoken/kernel/csrc/include/freetoken/device_api.h b/python/freetoken/kernel/csrc/include/freetoken/device_api.h new file mode 100644 index 00000000..932e2e36 --- /dev/null +++ b/python/freetoken/kernel/csrc/include/freetoken/device_api.h @@ -0,0 +1,99 @@ +// Device API seam for the tvm-ffi JIT/store/index kernels (and future GPU kernels). +// +// FreeToken's hand-written kernels are compiled once, guarded by `#if defined(USE_HIP)` +// (set by kernel/_toolchain.py `_rocm_cflags` and by setup.py for the torch extensions). +// This header maps the small set of runtime calls those kernels use to the active +// backend, so a kernel body written against these macros compiles under both CUDA and +// HIP without `#if` sprinkling at every call site. +// +// Every macro here must expand to a no-op-safe, backend-correct call. Prefer the 1:1 +// HIP counterparts (the HIP runtime is API-compatible at this level). +#pragma once + +#include +#include + +#if defined(USE_HIP) +#include +#else +#include +#include +#endif + +namespace freetoken::device { + +#if defined(USE_HIP) +using Error = hipError_t; +inline constexpr Error kErrorSuccess = hipSuccess; +inline const char* error_string(Error e) { return hipGetErrorString(e); } +using Device = hipDevice_t; +using Stream = hipStream_t; +using DeviceMemPtr = hipDeviceptr_t; +using HostFn = hipHostFn_t; +#else +using Error = cudaError_t; +inline constexpr Error kErrorSuccess = cudaSuccess; +inline const char* error_string(Error e) { return cudaGetErrorString(e); } +using Device = int; +using Stream = cudaStream_t; +using DeviceMemPtr = void*; +using HostFn = void (*)(void*); +#endif + +} // namespace freetoken::device + +// ---- allocation / copy / sync / launch (backend-agnostic call sites) ---- +#if defined(USE_HIP) +#define DEVICE_MALLOC(ptr, n) hipMalloc((void**)(ptr), (n)) +#define DEVICE_FREE(ptr) hipFree(ptr) +#define DEVICE_MEMCPY_ASYNC(dst, src, n, kind, stream) \ + hipMemcpyAsync((dst), (src), (n), (kind), (stream)) +#define DEVICE_MEMCPY_DEVICE_TO_DEVICE hipMemcpyDeviceToDevice +#define DEVICE_MEMCPY_DEVICE_TO_HOST hipMemcpyDeviceToHost +#define DEVICE_MEMCPY_HOST_TO_DEVICE hipMemcpyHostToDevice +#define DEVICE_SYNCTHREADS() __syncthreads() +#define DEVICE_LAUNCH_HOST_FUNC(stream, fn, data) \ + hipLaunchHostFunc((stream), (fn), (data)) +#define DEVICE_STREAM_SYNCHRONIZE(stream) hipStreamSynchronize(stream) +#define DEVICE_ATOMIC_ADD(ptr, val) atomicAdd((ptr), (val)) +#define DEVICE_ATOMIC_MAX(ptr, val) atomicMax((ptr), (val)) +#define DEVICE_ATOMIC_MIN(ptr, val) atomicMin((ptr), (val)) +#define DEVICE_ATOMIC_EXCHANGE(ptr, val) atomicExch((ptr), (val)) +#define DEVICE_THREAD_IDX_X threadIdx.x +#define DEVICE_BLOCK_DIM_X blockDim.x +#define DEVICE_BLOCK_IDX_X blockIdx.x +#define DEVICE_GRID_DIM_X gridDim.x +#else +#define DEVICE_MALLOC(ptr, n) cudaMalloc((void**)(ptr), (n)) +#define DEVICE_FREE(ptr) cudaFree(ptr) +#define DEVICE_MEMCPY_ASYNC(dst, src, n, kind, dir) \ + cudaMemcpyAsync((dst), (src), (n), (kind), (dir)) +#define DEVICE_MEMCPY_DEVICE_TO_DEVICE cudaMemcpyDeviceToDevice +#define DEVICE_MEMCPY_DEVICE_TO_HOST cudaMemcpyDeviceToHost +#define DEVICE_MEMCPY_HOST_TO_DEVICE cudaMemcpyHostToDevice +#define DEVICE_SYNCTHREADS() __syncthreads() +#define DEVICE_LAUNCH_HOST_FUNC(stream, fn, data) \ + cudaLaunchHostFunc((stream), (fn), (data)) +#define DEVICE_STREAM_SYNCHRONIZE(stream) cudaStreamSynchronize(stream) +#define DEVICE_ATOMIC_ADD(ptr, val) atomicAdd((ptr), (val)) +#define DEVICE_ATOMIC_MAX(ptr, val) atomicMax((ptr), (val)) +#define DEVICE_ATOMIC_MIN(ptr, val) atomicMin((ptr), (val)) +#define DEVICE_ATOMIC_EXCHANGE(ptr, val) atomicExch((ptr), (val)) +#define DEVICE_THREAD_IDX_X threadIdx.x +#define DEVICE_BLOCK_DIM_X blockDim.x +#define DEVICE_BLOCK_IDX_X blockIdx.x +#define DEVICE_GRID_DIM_X gridDim.x +#endif + +// The FFI kernels pin tensors to kDLCUDA today; on ROCm the tvm-ffi device type is +// kDLROCM. Kernels that resolve a device must accept both (see tensor.h). This macro +// picks the backend's DLDevice code at call time. +#if defined(USE_HIP) +#define DEVICE_DLDEVICE_ROCM 10 // kDLROCM +#define DEVICE_DLDEVICE_CUDA 2 // kDLCUDA +#define DEVICE_ACTIVE_DLDEVICE DEVICE_DLDEVICE_ROCM +#else +#define DEVICE_DLDEVICE_ROCM 10 +#define DEVICE_DLDEVICE_CUDA 2 +#define DEVICE_ACTIVE_DLDEVICE DEVICE_DLDEVICE_CUDA +#endif diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 8e917832..de28877c 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -10,6 +10,16 @@ #include #include +#if defined(USE_HIP) +#include +// HIP has no __grid_constant__ (a CUDA read-only-constant optimization). Define it +// empty so `const __grid_constant__ Params params` compiles as a plain by-value +// parameter, which is correct (just without the CUDA constant-cache hint). +#ifndef __grid_constant__ +#define __grid_constant__ +#endif +#endif + namespace device { inline constexpr auto kWarpThreads = 32u; @@ -42,16 +52,23 @@ __always_inline __device__ auto offset(const T *ptr, U... offset) -> const namespace PDL { +// Programmatic Dependent Launch is a CUDA-only optimization (griddepcontrol). +// HIP has no equivalent; the wait/launch are no-ops there. PDL is optional in the +// kernels (use_pdl defaults to false), so dropping it is purely a perf change. template __always_inline __device__ void wait() { +#if !defined(USE_HIP) if constexpr (kUsePDL) { asm volatile("griddepcontrol.wait;" ::: "memory"); } +#endif } template __always_inline __device__ void launch() { +#if !defined(USE_HIP) if constexpr (kUsePDL) { asm volatile("griddepcontrol.launch_dependents;" :::); } +#endif } } // namespace PDL @@ -60,6 +77,23 @@ template __always_inline __device__ void launch() { namespace host { +#if defined(USE_HIP) +inline auto +HIP_CHECK(::hipError_t error, + std::source_location location = std::source_location::current()) + -> void { + if (error != ::hipSuccess) { + [[unlikely]]; + ::host::panic(location, "HIP error: ", ::hipGetErrorString(error)); + } +} + +inline auto +HIP_CHECK(std::source_location location = std::source_location::current()) + -> void { + return HIP_CHECK(::hipGetLastError(), location); +} +#else inline auto CUDA_CHECK(::cudaError_t error, std::source_location location = std::source_location::current()) @@ -75,7 +109,21 @@ CUDA_CHECK(std::source_location location = std::source_location::current()) -> void { return CUDA_CHECK(::cudaGetLastError(), location); } +#endif +#if defined(USE_HIP) +template inline void set_smem_once(std::size_t smem_size) { + static const auto last_smem_size = [&] { + HIP_CHECK(::hipFuncSetAttribute( + F, ::hipFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + return smem_size; + }(); + RuntimeCheck( + smem_size <= last_smem_size, + "Dynamic shared memory size exceeds the previously set maximum size: ", + last_smem_size, " bytes"); +} +#else template inline void set_smem_once(std::size_t smem_size) { static const auto last_smem_size = [&] { CUDA_CHECK(::cudaFuncSetAttribute( @@ -87,7 +135,52 @@ template inline void set_smem_once(std::size_t smem_size) { "Dynamic shared memory size exceeds the previously set maximum size: ", last_smem_size, " bytes"); } +#endif +#if defined(USE_HIP) +struct LaunchKernel { +public: + explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, DLDevice device, + std::size_t dynamic_shared_mem_bytes = 0) noexcept + : m_grid(grid_dim), m_block(block_dim), + m_stream(resolve_device(device)), m_smem(dynamic_shared_mem_bytes) {} + + explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, hipStream_t stream, + std::size_t dynamic_shared_mem_bytes = 0) noexcept + : m_grid(grid_dim), m_block(block_dim), m_stream(stream), + m_smem(dynamic_shared_mem_bytes) {} + + static auto resolve_device(DLDevice device) -> hipStream_t { + return static_cast( + ::TVMFFIEnvGetStream(device.device_type, device.device_id)); + } + + LaunchKernel(const LaunchKernel &) = delete; + LaunchKernel &operator=(const LaunchKernel &) = delete; + + template + auto operator()(T &&kernel, Args &&...args) const -> void { + // hipLaunchKernel takes a void** args array (pointers to each argument value), + // unlike cudaLaunchKernelEx's variadic form. The array is consumed at launch. + void *arg_array[sizeof...(Args)] = { + const_cast(static_cast(&args))...}; + HIP_CHECK(::hipLaunchKernel(reinterpret_cast(kernel), m_grid, + m_block, arg_array, m_smem, m_stream)); + } + + auto with_attr(bool /*use_pdl*/) -> LaunchKernel & { + // HIP has no programmatic dependent launch / launch attributes; PDL is a + // CUDA-only optimization and is dropped here. + return *this; + } + +private: + dim3 m_grid; + dim3 m_block; + hipStream_t m_stream; + std::size_t m_smem; +}; +#else struct LaunchKernel { public: explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, DLDevice device, @@ -140,5 +233,6 @@ private: cudaLaunchConfig_t m_config; cudaLaunchAttribute m_attr_cache; }; +#endif } // namespace host diff --git a/python/freetoken/kernel/csrc/jit/batch_memcpy.cuh b/python/freetoken/kernel/csrc/jit/batch_memcpy.cuh index 690eff3d..612b1d82 100644 --- a/python/freetoken/kernel/csrc/jit/batch_memcpy.cuh +++ b/python/freetoken/kernel/csrc/jit/batch_memcpy.cuh @@ -8,6 +8,26 @@ #include #include +#if defined(USE_HIP) +// HIP has no cudaMemcpyBatchAsync equivalent. Implement a per-copy grid kernel: one +// block per copy, each block cooperatively copying its (src, dst, nbytes) triple. +// The host pointer arrays are staged to device memory for the launch. Copies within +// a batch are unordered (as with cudaMemcpyBatchAsync), so blocks may run in any order. +__global__ void batch_memcpy_kernel(const void* const* srcs, void* const* dsts, + const std::size_t* sizes, std::size_t n) { + const std::size_t i = blockIdx.x; + if (i >= n) { + return; + } + const char* s = static_cast(srcs[i]); + char* d = static_cast(dsts[i]); + const std::size_t sz = sizes[i]; + for (std::size_t j = threadIdx.x; j < sz; j += blockDim.x) { + d[j] = s[j]; + } +} +#endif + // Host wrapper over cudaMemcpyBatchAsync (CUDA >= 13.0, the 8-argument signature; // 12.8/12.9 carried an extra failIdx parameter): enqueue N independent // pointer-to-pointer copies with ONE runtime call, on an explicit (non-legacy) @@ -20,7 +40,49 @@ struct BatchMemcpy { tvm::ffi::TensorView sizes, int64_t stream_handle ) { -#if CUDART_VERSION >= 13000 +#if defined(USE_HIP) + using namespace host; + auto N = SymbolicSize{"batch length"}; + auto ptr_dtype = SymbolicDType{}; + TensorMatcher({N}) + .with_dtype(ptr_dtype) + .with_device() + .verify(dst_ptrs) + .verify(src_ptrs) + .verify(sizes); + auto n = static_cast(N.unwrap()); + if (n == 0) { + return; + } + RuntimeCheck(stream_handle != 0, "batch_memcpy rejects the legacy NULL stream"); + auto stream = reinterpret_cast(stream_handle); + + const void* const* srcs = reinterpret_cast(src_ptrs.data_ptr()); + void* const* dsts = reinterpret_cast(dst_ptrs.data_ptr()); + const std::size_t* sizes_arr = reinterpret_cast(sizes.data_ptr()); + + // Stage the host pointer arrays to device memory for the kernel. Use + // stream-ordered allocation so the free is ordered after the (async) kernel + // launch -- a plain hipFree here could free memory the kernel is still reading. + void* d_srcs = nullptr; + void* d_dsts = nullptr; + void* d_sizes = nullptr; + HIP_CHECK(hipMallocAsync(&d_srcs, n * sizeof(void*), stream)); + HIP_CHECK(hipMallocAsync(&d_dsts, n * sizeof(void*), stream)); + HIP_CHECK(hipMallocAsync(&d_sizes, n * sizeof(std::size_t), stream)); + HIP_CHECK(hipMemcpy(d_srcs, srcs, n * sizeof(void*), hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(d_dsts, dsts, n * sizeof(void*), hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(d_sizes, sizes_arr, n * sizeof(std::size_t), hipMemcpyHostToDevice)); + + void* args[4] = {&d_srcs, &d_dsts, &d_sizes, &n}; + HIP_CHECK(hipLaunchKernel( + reinterpret_cast(batch_memcpy_kernel), dim3(n), dim3(256), args, 0, + stream)); + + HIP_CHECK(hipFreeAsync(d_srcs, stream)); + HIP_CHECK(hipFreeAsync(d_dsts, stream)); + HIP_CHECK(hipFreeAsync(d_sizes, stream)); +#elif CUDART_VERSION >= 13000 using namespace host; auto N = SymbolicSize{"batch length"}; auto ptr_dtype = SymbolicDType{}; diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23e..2583d1db 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -34,40 +34,64 @@ inline constexpr auto get_mem_package() { } __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { +#if defined(USE_HIP) + return *src; +#else uint32_t tmp; asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src)); return uint1{tmp}; +#endif } __always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 { +#if defined(USE_HIP) + return *src; +#else uint32_t tmp0, tmp1; asm volatile("ld.global.L1::no_allocate.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src)); return uint2{tmp0, tmp1}; +#endif } __always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 { +#if defined(USE_HIP) + return *src; +#else uint32_t tmp0, tmp1, tmp2, tmp3; asm volatile("ld.global.L1::no_allocate.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src)); return uint4{tmp0, tmp1, tmp2, tmp3}; +#endif } __always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) { +#if defined(USE_HIP) + *dst = value; +#else uint32_t tmp = value.x; asm volatile("st.global.wt.b32 [%0],%1;" ::"l"(dst), "r"(tmp)); +#endif } __always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) { +#if defined(USE_HIP) + *dst = value; +#else uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; asm volatile("st.global.wt.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1)); +#endif } __always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) { +#if defined(USE_HIP) + *dst = value; +#else uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; uint32_t tmp2 = value.z; uint32_t tmp3 = value.w; asm volatile("st.global.wt.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3)); +#endif } __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag_ptr) { @@ -75,7 +99,10 @@ __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag auto* flag = reinterpret_cast(const_cast(flag_ptr)); uint32_t sleep_ns = 128; while (atomicAdd(flag, 0) > 0) { -#if __CUDA_ARCH__ >= 700 +#if defined(USE_HIP) + // HIP has no __nanosleep; busy-wait (functional parity). + (void)sleep_ns; +#elif __CUDA_ARCH__ >= 700 __nanosleep(sleep_ns); #endif sleep_ns = sleep_ns < 2048 ? (sleep_ns << 1) : 2048; @@ -134,6 +161,16 @@ __always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) // process (set at engine launch). inline bool host_ptr_identity() { static const bool identity = [] { +#if defined(USE_HIP) + int device = 0; + if (hipGetDevice(&device) != hipSuccess) { + return false; // fail closed: translate (and surface errors), don't assume identity + } + int uva = 0, reg = 0; + hipDeviceGetAttribute(&uva, hipDeviceAttributeUnifiedAddressing, device); + hipDeviceGetAttribute(®, hipDeviceAttributeCanUseHostPointerForRegisteredMem, device); + return uva == 1 && reg == 1; +#else int device = 0; if (cudaGetDevice(&device) != cudaSuccess) { return false; // fail closed: translate (and surface errors), don't assume identity @@ -142,11 +179,23 @@ inline bool host_ptr_identity() { cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device); cudaDeviceGetAttribute(®, cudaDevAttrCanUseHostPointerForRegisteredMem, device); return uva == 1 && reg == 1; +#endif }(); return identity; } inline void* device_alias(void* ptr, DLDevice dev) { +#if defined(USE_HIP) + if (dev.device_type == kDLROCM || host_ptr_identity()) { + return ptr; + } + void* mapped = nullptr; + const auto err = hipHostGetDevicePointer(&mapped, ptr, 0); + host::RuntimeCheck(err == hipSuccess, + "fast_index_copy: host tensor must be pinned+mapped (hipHostGetDevicePointer: ", + hipGetErrorString(err), ")"); + return mapped; +#else if (dev.device_type == kDLCUDA || host_ptr_identity()) { return ptr; } @@ -156,6 +205,7 @@ inline void* device_alias(void* ptr, DLDevice dev) { "fast_index_copy: host tensor must be pinned+mapped (cudaHostGetDevicePointer: ", cudaGetErrorString(err), ")"); return mapped; +#endif } struct IndexKernelParams { @@ -344,12 +394,12 @@ struct FastIndexCopyKernel { TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(src); TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(dst); TensorMatcher({L}) diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index c3947adf..da15cdd8 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp @@ -1,12 +1,33 @@ #include +#if defined(USE_HIP) +#include +#else #include +#endif #include +// host alloc flags (mapped+portable) shared by both backends +#if defined(USE_HIP) +#define DEVICE_HOST_ALLOC_FLAGS_MAPPED \ + (hipHostMallocPortable | hipHostMallocMapped) +#define DEVICE_HOST_REGISTER_FLAGS_MAPPED \ + (hipHostRegisterPortable | hipHostRegisterMapped) +#else +#define DEVICE_HOST_ALLOC_FLAGS_MAPPED \ + (cudaHostAllocPortable | cudaHostAllocMapped) +#define DEVICE_HOST_REGISTER_FLAGS_MAPPED \ + (cudaHostRegisterPortable | cudaHostRegisterMapped) +#endif + namespace { void free_pinned(void *ptr) { if (ptr != nullptr) { +#if defined(USE_HIP) + hipHostFree(ptr); +#else cudaFreeHost(ptr); +#endif } } @@ -34,9 +55,16 @@ torch::Tensor create_pinned_tensor_like(torch::Tensor input) { const size_t alloc_nbytes = static_cast(nbytes == 0 ? 1 : nbytes); void *data_ptr = nullptr; +#if defined(USE_HIP) + const hipError_t alloc_err = + hipMallocHost(&data_ptr, alloc_nbytes); + TORCH_CHECK(alloc_err == hipSuccess, + "hipMallocHost failed: ", hipGetErrorString(alloc_err)); +#else const cudaError_t alloc_err = cudaMallocHost(&data_ptr, alloc_nbytes); TORCH_CHECK(alloc_err == cudaSuccess, "cudaMallocHost failed: ", cudaGetErrorString(alloc_err)); +#endif auto options = input.options().device(torch::kCPU).pinned_memory(true); @@ -58,10 +86,17 @@ torch::Tensor alloc_pinned_tensor(std::vector sizes, // Portable + mapped: the offload gather kernel reads these banks straight // from host memory (zero-copy), which requires device-mapped pinned pages. void *data_ptr = nullptr; +#if defined(USE_HIP) + const hipError_t alloc_err = hipHostMalloc( + &data_ptr, alloc_nbytes, DEVICE_HOST_ALLOC_FLAGS_MAPPED); + TORCH_CHECK(alloc_err == hipSuccess, + "hipHostMalloc failed: ", hipGetErrorString(alloc_err)); +#else const cudaError_t alloc_err = cudaHostAlloc( &data_ptr, alloc_nbytes, cudaHostAllocPortable | cudaHostAllocMapped); TORCH_CHECK(alloc_err == cudaSuccess, "cudaHostAlloc failed: ", cudaGetErrorString(alloc_err)); +#endif auto options = torch::TensorOptions() .dtype(dtype) @@ -76,38 +111,69 @@ torch::Tensor alloc_pinned_tensor(std::vector sizes, // device address). Zero-copy consumers resolve bank base addresses through these. bool host_ptr_identity() { int device = 0; +#if defined(USE_HIP) + const hipError_t err = hipGetDevice(&device); + TORCH_CHECK(err == hipSuccess, "hipGetDevice failed: ", hipGetErrorString(err)); + int uva = 0, reg = 0; + hipDeviceGetAttribute(&uva, hipDeviceAttributeUnifiedAddressing, device); + hipDeviceGetAttribute( + ®, hipDeviceAttributeCanUseHostPointerForRegisteredMem, device); +#else const cudaError_t err = cudaGetDevice(&device); TORCH_CHECK(err == cudaSuccess, "cudaGetDevice failed: ", cudaGetErrorString(err)); int uva = 0, reg = 0; cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device); cudaDeviceGetAttribute(®, cudaDevAttrCanUseHostPointerForRegisteredMem, device); +#endif return uva == 1 && reg == 1; } int64_t host_device_ptr(int64_t host_ptr) { void *dev_ptr = nullptr; +#if defined(USE_HIP) + const hipError_t err = hipHostGetDevicePointer( + &dev_ptr, reinterpret_cast(host_ptr), 0); + TORCH_CHECK(err == hipSuccess, + "hipHostGetDevicePointer failed (host memory must be pinned+mapped): ", + hipGetErrorString(err)); +#else const cudaError_t err = cudaHostGetDevicePointer(&dev_ptr, reinterpret_cast(host_ptr), 0); TORCH_CHECK(err == cudaSuccess, "cudaHostGetDevicePointer failed (host memory must be pinned+mapped): ", cudaGetErrorString(err)); +#endif return reinterpret_cast(dev_ptr); } void host_register(int64_t addr, int64_t nbytes) { +#if defined(USE_HIP) + const hipError_t err = + hipHostRegister(reinterpret_cast(addr), static_cast(nbytes), + DEVICE_HOST_REGISTER_FLAGS_MAPPED); + TORCH_CHECK(err == hipSuccess, + "hipHostRegister failed: ", hipGetErrorString(err)); +#else const cudaError_t err = cudaHostRegister(reinterpret_cast(addr), static_cast(nbytes), cudaHostRegisterPortable | cudaHostRegisterMapped); TORCH_CHECK(err == cudaSuccess, "cudaHostRegister failed: ", cudaGetErrorString(err)); +#endif } +// CUDA-only: the maximum CUDA version the NVIDIA driver supports, used to gate +// driver-JIT kernels. On ROCm there is no CUDA driver; return 0 (== "no driver"). int64_t driver_cuda_version() { +#if defined(USE_HIP) + return 0; +#else int version = 0; // stays 0 when no driver is installed const cudaError_t err = cudaDriverGetVersion(&version); TORCH_CHECK(err == cudaSuccess, "cudaDriverGetVersion failed: ", cudaGetErrorString(err)); return version; +#endif } } // namespace @@ -122,7 +188,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("host_device_ptr", &host_device_ptr, "Device-visible alias of a pinned+mapped host address"); m.def("host_register", &host_register, - "cudaHostRegister an existing host range as portable+mapped"); + "cudaHostRegister/hipHostRegister an existing host range as portable+mapped"); m.def("driver_cuda_version", &driver_cuda_version, - "Max CUDA version the installed NVIDIA driver supports (0 if none)"); + "Max CUDA version the installed NVIDIA driver supports (0 if none / ROCm)"); } diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 04a16560..13b0ea0c 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -51,16 +51,28 @@ def _c_compiler_for(cxx: str) -> str: def _module(): from torch.utils.cpp_extension import load - extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"] - host_cxx = _host_compiler() - if host_cxx is not None: - # Point both nvcc's host pass (-ccbin) and torch's C++ compile (CXX) at a - # libtorch/nvcc-compatible compiler. Force (not setdefault): the system - # default (CXX unset -> g++) can be a gcc too new for the torch headers. - cxx_path = shutil.which(host_cxx) or host_cxx - extra_cuda_cflags += ["-ccbin", cxx_path] - os.environ["CXX"] = cxx_path - os.environ["CC"] = _c_compiler_for(cxx_path) + if torch.version.hip is not None: + # ROCm: hipcc (torch.utils.cpp_extension picks it up), pass the HIP defines so + # the kernels compile their HIP branches; drop the CUDA-only -ccbin/flag logic. + # Explicit --offload-arch (plus PYTORCH_ROCM_ARCH) prevents torch from auto- + # emitting ~14 gfx arches, which would multiply build time per arch. + os.environ.setdefault("PYTORCH_ROCM_ARCH", "gfx1100") + extra_cuda_cflags = [ + "-O3", "--offload-arch=gfx1100", "-DUSE_HIP=1", "-DUSE_ROCM=1", + ] + os.environ.pop("CXX", None) + os.environ.pop("CC", None) + else: + extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"] + host_cxx = _host_compiler() + if host_cxx is not None: + # Point both nvcc's host pass (-ccbin) and torch's C++ compile (CXX) at a + # libtorch/nvcc-compatible compiler. Force (not setdefault): the system + # default (CXX unset -> g++) can be a gcc too new for the torch headers. + cxx_path = shutil.which(host_cxx) or host_cxx + extra_cuda_cflags += ["-ccbin", cxx_path] + os.environ["CXX"] = cxx_path + os.environ["CC"] = _c_compiler_for(cxx_path) # gguf_kernel.cu carries its own PYBIND11_MODULE (appended at the end), so a # plain `load` of the single source compiles + binds the ggml_* ops. @@ -69,7 +81,7 @@ def _module(): sources=[str(_CSRC / "gguf_kernel.cu")], extra_include_paths=[str(_CSRC)], extra_cuda_cflags=extra_cuda_cflags, - verbose=True, + verbose=False, ) diff --git a/python/freetoken/kernel/tinygrad_fallback.py b/python/freetoken/kernel/tinygrad_fallback.py new file mode 100644 index 00000000..2b0663fe --- /dev/null +++ b/python/freetoken/kernel/tinygrad_fallback.py @@ -0,0 +1,92 @@ +"""tinygrad-JIT fallback for FFI kernels that are not hand-ported to HIP. + +FreeToken's hand-written tvm-ffi kernels (``store`` / ``index`` / ``fast_index_copy`` / +``batch_memcpy``) are CUDA source compiled via nvcc/JIT. The primary AMD port is the +``#if defined(USE_HIP)`` seam in ``device_api.h`` + ``LaunchKernel``/``warp.cuh``. This +module is the **documented fallback** for any kernel that proves intractable to hipify: +tinygrad's JIT compiles one logical kernel to PTX (CUDA) *and* AMDGPU/LLVM (ROCm), so the +same source covers both platforms. + +Constraints (matching the FFI contract): + +* Each fallback takes the same ``tvm.ffi.TensorView`` arguments as the hand-written + kernel and returns the same output tensor(s), so the swap is invisible to callers. +* It runs on the *host* (tinygrad handles GPU dispatch); on ROCm it compiles to AMDGPU. +* It is **never a default**: ``kernel/utils.py`` only routes a kernel to the fallback + when (a) ROCm is active and (b) the hand-HIP AOT/JIT variant is absent/unbuildable. + If tinygrad is not installed, invoking the fallback raises a clear error. + +Because tinygrad is an optional dependency (installed only when the fallback is actually +needed), all imports here are lazy and the module imports with zero third-party deps, so +it is safe to import on the CUDA-only path. +""" + +from __future__ import annotations + +from typing import Callable, Optional + +__all__ = [ + "is_tinygrad_available", + "kernel_fallback_available", + "get_kernel_fallback", +] + +# Kernel names the fallback registry knows how to build (mirrors the FFI kernel set). +_KNOWN_KERNELS = ("store", "index", "fast_index_copy", "batch_memcpy") + +#: Which kernels currently have a *functional* tinygrad reimplementation. As HIP ports +#: land in Inc 7, names are removed from this set (the hand port wins); kernels left here +#: (if any) are the documented fallback set. Default: empty -- the hand-HIP port is the +#: primary path and the fallback is opt-in per kernel. +_FALLBACK_IMPLEMENTED: set[str] = set() + + +def is_tinygrad_available() -> bool: + """True when the ``tinygrad`` package can be imported (JIT-to-ROCm available).""" + try: + import importlib.util # noqa: PLC0415 + + return importlib.util.find_spec("tinygrad") is not None + except Exception: + return False + + +def kernel_fallback_available(kernel: str) -> bool: + """True when a tinygrad fallback for ``kernel`` is both implemented and usable + (tinygrad installed). Always False on the CUDA path unless explicitly enabled, so + the CUDA build never depends on tinygrad.""" + if kernel not in _FALLBACK_IMPLEMENTED: + return False + return is_tinygrad_available() + + +def get_kernel_fallback(kernel: str): + """Return the tinygrad-backed fallback callable for ``kernel``, or raise a clear + error explaining why it is unavailable. Never called on the CUDA path.""" + if kernel not in _FALLBACK_IMPLEMENTED: + raise RuntimeError( + f"FFI kernel {kernel!r} has no tinygrad fallback registered. On ROCm the " + "preferred path is the hand-written HIP port (device_api.h); if you intend " + "to use the tinygrad fallback you must register it in " + "kernel/tinygrad_fallback.py._FALLBACK_IMPLEMENTED and implement the " + "corresponding build function." + ) + if not is_tinygrad_available(): + raise RuntimeError( + f"FFI kernel {kernel!r} requires the tinygrad fallback, but tinygrad is not " + "installed. Install it (`pip install tinygrad`) or provide a hand-written " + "HIP port for this kernel." + ) + from freetoken.kernel import tinygrad_impl # noqa: PLC0415 (lazy; may be None) + + builder = getattr(tinygrad_impl, f"build_{kernel}", None) + if builder is None: + raise RuntimeError( + f"tinygrad fallback for {kernel!r} is registered but has no " + "tinygrad_impl.build_() builder." + ) + return builder + + +def _list_fallbacks() -> list[str]: + return [k for k in _FALLBACK_IMPLEMENTED if kernel_fallback_available(k)] diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py index 2c38b533..d57852f5 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/python/freetoken/kernel/triton/activation.py @@ -48,20 +48,15 @@ def _pdl_supported() -> bool: @triton.jit def _fast_tanh(x): - # PTX tanh.approx.f32 — single HW op, matches flashinfer math::tanh. - return tl.inline_asm_elementwise( - "tanh.approx.f32 $0, $1;", "=f,f", [x], - dtype=tl.float32, is_pure=True, pack=1, - ) + # tanh.approx.f32 is a CUDA PTX intrinsic; libdevice.tanh is portable (maps to + # tanhf on both CUDA and AMD). + return libdevice.tanh(x) @triton.jit def _fast_ex2(x): - # PTX ex2.approx.f32 — matches __expf fast path used by flashinfer silu. - return tl.inline_asm_elementwise( - "ex2.approx.f32 $0, $1;", "=f,f", [x], - dtype=tl.float32, is_pure=True, pack=1, - ) + # ex2.approx.f32 is CUDA-only; libdevice.exp2 is portable. + return libdevice.exp2(x) @triton.jit @@ -133,10 +128,18 @@ def _act_and_mul( # 1024/w4/s2 best at rows>=4096). block_d = min(triton.next_power_of_2(d), 1024 if M >= 4096 else 512) num_stages = 2 if block_d == 1024 else 3 - _act_and_mul_kernel[grid]( - o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, launch_pdl=pdl, - BLOCK_D=block_d, num_warps=4, num_stages=num_stages, - ) + # ``launch_pdl`` (Hopper griddepcontrol) is CUDA-only; ROCm triton rejects the + # kwarg, so only pass it on a PDL-capable (sm_90+) device. + if pdl: + _act_and_mul_kernel[grid]( + o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=True, launch_pdl=True, + BLOCK_D=block_d, num_warps=4, num_stages=num_stages, + ) + else: + _act_and_mul_kernel[grid]( + o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=False, + BLOCK_D=block_d, num_warps=4, num_stages=num_stages, + ) return out diff --git a/python/freetoken/kernel/triton/norm.py b/python/freetoken/kernel/triton/norm.py index 3f95c29f..320f6077 100644 --- a/python/freetoken/kernel/triton/norm.py +++ b/python/freetoken/kernel/triton/norm.py @@ -142,11 +142,18 @@ def _rmsnorm(input, weight, eps, out, gemma: bool): # PDL only on the contiguous (decode-replay) path: on the strided qk-norm's # 32k-CTA prefill grids the per-CTA gdc_wait poll costs more than it hides. pdl = contig and is_sm90_supported() - _rmsnorm_kernel[(A, B)]( - out, input, weight, eps, H, sxa, sxb, soa, sob, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, - num_warps=_num_warps(A * B), num_stages=1, - ) + if pdl: + _rmsnorm_kernel[(A, B)]( + out, input, weight, eps, H, sxa, sxb, soa, sob, + CONTIG=True, ENABLE_PDL=True, launch_pdl=True, GEMMA=gemma, + num_warps=_num_warps(A * B), num_stages=1, + ) + else: + _rmsnorm_kernel[(A, B)]( + out, input, weight, eps, H, sxa, sxb, soa, sob, + CONTIG=contig, ENABLE_PDL=False, GEMMA=gemma, + num_warps=_num_warps(A * B), num_stages=1, + ) return out @@ -170,11 +177,18 @@ def _fused_add_rmsnorm(input, residual, weight, eps, gemma: bool): _, _, sra, srb = _leading(residual) contig = input.ndim == 2 and input.is_contiguous() and residual.is_contiguous() pdl = contig and is_sm90_supported() - _fused_add_rmsnorm_kernel[(A, B)]( - input, residual, weight, eps, H, sxa, sxb, sra, srb, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, - num_warps=_num_warps(A * B), num_stages=1, - ) + if pdl: + _fused_add_rmsnorm_kernel[(A, B)]( + input, residual, weight, eps, H, sxa, sxb, sra, srb, + CONTIG=True, ENABLE_PDL=True, launch_pdl=True, GEMMA=gemma, + num_warps=_num_warps(A * B), num_stages=1, + ) + else: + _fused_add_rmsnorm_kernel[(A, B)]( + input, residual, weight, eps, H, sxa, sxb, sra, srb, + CONTIG=contig, ENABLE_PDL=False, GEMMA=gemma, + num_warps=_num_warps(A * B), num_stages=1, + ) def fused_add_rmsnorm(input, residual, weight, eps: float = 1e-6, enable_pdl: bool = False): diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index 7a0164b5..1beb7398 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -40,6 +40,32 @@ def _rank(a: str) -> int: cc = max(arch_list, key=_rank).rstrip("a").replace(".", "") flags = flags + [f"-gencode=arch=compute_{cc},code=compute_{cc}"] return flags + + +def _rocm_cflags(extra: List[str]) -> List[str]: + """HIP/ROCm flags for a kernel build. Adds the ``USE_HIP``/``USE_ROCM`` defines so the + shared ``.cu``/``.cuh`` sources compile their HIP branches (llama.cpp-style). When + ``TVM_FFI_ROCM_ARCH_LIST`` (e.g. "gfx1100") is set (AOT cache build), we pin + ``--offload-arch``; otherwise hipcc targets the local GPU.""" + flags = DEFAULT_CFLAGS + ["-DUSE_HIP=1", "-DUSE_ROCM=1"] + list(extra) + arch_list = os.getenv("TVM_FFI_ROCM_ARCH_LIST", "").split() + if arch_list: + flags = flags + [f"--offload-arch={arch_list[0]}"] + return flags + + +def _arch_flags(extra: List[str]) -> List[str]: + """GPU kernel build flags for the active backend: HIP flags on ROCm torch, else CUDA.""" + try: + from freetoken.kernel._toolchain import is_rocm_torch + + if is_rocm_torch(): + return _rocm_cflags(extra) + except Exception: + pass + return _cuda_cflags(extra) + + CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool] @@ -86,6 +112,20 @@ def _build_stamps(segments: List[str]) -> set[str]: return {s for s in segments if re.fullmatch(r"g[0-9a-f]{7,40}", s)} +def _build_stamps(local_segments: List[str]) -> List[str]: + """The `g` commit-stamp tokens of a local version segment list + (``["cu130", "g3f01615"]`` -> ``["g3f01615"]``).""" + return [s for s in local_segments if s.startswith("g")] + + +def _arch_tags(local_segments: List[str]) -> List[str]: + """The backend-tag tokens of a local segment list (``cu130`` or ``rocm``). A cache + wheel and runtime wheel must carry the SAME backend tag -- a ``+rocm`` cache is + meaningless to a ``+cu130`` runtime (the fatbin is gfx SASS vs sm SASS) and vice + versa. Returns the tags (usually one, e.g. ``["cu130"]`` or ``["rocm"]``).""" + return [s for s in local_segments if s.startswith("cu") or s.startswith("rocm")] + + def _kernel_cache_version_ok(cache_version: str, runtime_version: str) -> bool: """Same release -- and, when both sides carry a `g` stamp, the same build. @@ -93,14 +133,24 @@ def _kernel_cache_version_ok(cache_version: str, runtime_version: str) -> bool: `.g`), so the old string-prefix test cannot pair a stamped runtime with its cache; and comparing the stamps rejects a runtime/cache pair from two different builds, which bare release numbers (both `0.1.1`) could never detect. Either side - may lack a stamp (dev builds) -- then only the release part is compared.""" + may lack a stamp (dev builds) -- then only the release part is compared. + + The backend arch tag (`cu130` vs `rocm`) must also match when both sides carry one: + a prebuilt kernel-cache fatbin is SASS for a specific backend family, so a CUDA + runtime must never load a ROCm cache (or vice versa).""" cache_base, cache_local = _version_parts(cache_version) runtime_base, runtime_local = _version_parts(runtime_version) if cache_base != runtime_base: return False cache_stamps = _build_stamps(cache_local) runtime_stamps = _build_stamps(runtime_local) - return not (cache_stamps and runtime_stamps and cache_stamps != runtime_stamps) + if cache_stamps and runtime_stamps and cache_stamps != runtime_stamps: + return False + cache_arch = _arch_tags(cache_local) + runtime_arch = _arch_tags(runtime_local) + if cache_arch and runtime_arch and cache_arch != runtime_arch: + return False + return True def _kernel_cache_dir() -> pathlib.Path | None: @@ -127,7 +177,8 @@ def _kernel_cache_dir() -> pathlib.Path | None: f"{package_version!r} does not match freetoken version {runtime_version!r}" ) cache_cuda = re.search(r"\+cu(\d{2,})", package_version) - if cache_cuda is not None: + cache_rocm = re.search(r"\+rocm", package_version) + if cache_cuda is not None and cache_rocm is None: from freetoken.kernel._toolchain import torch_cuda_major cache_major = int(cache_cuda.group(1)[:-1]) @@ -201,9 +252,16 @@ def load_aot( return prebuilt if cuda_files: - from freetoken.kernel._toolchain import check_nvcc_matches_torch + from freetoken.kernel._toolchain import ( + check_hip_matches_torch, + check_nvcc_matches_torch, + is_rocm_torch, + ) - check_nvcc_matches_torch() + if is_rocm_torch(): + check_hip_matches_torch() + else: + check_nvcc_matches_torch() from tvm_ffi.cpp import load @@ -222,7 +280,7 @@ def load_aot( cpp_files=cpp_files, cuda_files=cuda_files, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags), + extra_cuda_cflags=_arch_flags(extra_cuda_cflags), extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, @@ -247,9 +305,16 @@ def load_jit( return prebuilt if cuda_files or cuda_wrappers: - from freetoken.kernel._toolchain import check_nvcc_matches_torch + from freetoken.kernel._toolchain import ( + check_hip_matches_torch, + check_nvcc_matches_torch, + is_rocm_torch, + ) - check_nvcc_matches_torch() + if is_rocm_torch(): + check_hip_matches_torch() + else: + check_nvcc_matches_torch() from tvm_ffi.cpp import load_inline @@ -277,7 +342,7 @@ def load_jit( cpp_sources=cpp_sources, cuda_sources=cuda_sources, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags), + extra_cuda_cflags=_arch_flags(extra_cuda_cflags), extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index d68d8ded..e1691363 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -531,6 +531,16 @@ def _expert_gemm( return fused_experts_gguf_q4_0( hidden_states, gate_up, down, topk_weights, topk_ids, self.activation ) + if fmt == "gguf": + # Native GGUF qwen3.5-moe experts: gate_up stays Q4_K, down is stored as Q8_0 + # (re-quantized at load; a uniform format the cache can hold). Dequant-in-kernel + # grouped GEMV (MMVQ) over the streamed packed banks. + from freetoken.moe.fused_gguf import fused_experts_gguf + + gate_up, down = views + return fused_experts_gguf( + hidden_states, gate_up, down, topk_weights, topk_ids, self.activation + ) if fmt == "mxfp4_triton": # gpt-oss MXFP4 experts (biased, clamped swiglu): transposed split-K GEMV # decode + grouped `_t` prefill. The swiglu scalars live on the layer diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 63b1a18b..b1e84c5c 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -18,6 +18,7 @@ # reuses the model classes but a GGUF parse_config / iter_weights). GGUF_ARCH_TO_REGISTRY: dict[str, str] = { "gemma4": "Gemma4GGUFForCausalLM", + "qwen35moe": "Qwen35moeGGUFForCausalLM", } diff --git a/python/freetoken/models/gguf/dequant.py b/python/freetoken/models/gguf/dequant.py index 77c3ea01..05648f9e 100644 --- a/python/freetoken/models/gguf/dequant.py +++ b/python/freetoken/models/gguf/dequant.py @@ -23,6 +23,8 @@ GGML_F16 = 1 GGML_Q4_0 = 2 GGML_Q8_0 = 8 +GGML_Q4_K = 12 +GGML_Q5_K = 13 GGML_Q6_K = 14 GGML_BF16 = 30 @@ -33,6 +35,7 @@ GGML_BF16: (1, 2), GGML_Q4_0: (32, 18), GGML_Q8_0: (32, 34), + GGML_Q4_K: (256, 144), GGML_Q6_K: (256, 210), } @@ -42,6 +45,8 @@ GGML_BF16: "BF16", GGML_Q4_0: "Q4_0", GGML_Q8_0: "Q8_0", + GGML_Q4_K: "Q4_K", + GGML_Q5_K: "Q5_K", GGML_Q6_K: "Q6_K", } @@ -115,8 +120,69 @@ def dequant_q6_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: return y.reshape(-1).to(out_dtype) +def quantize_q8_0(w: torch.Tensor) -> torch.Tensor: + """Quantize dense rows to packed Q8_0 blocks (``half d`` + 32 int8). + + ``w``'s last dim must be a multiple of 32 (the Q8_0 block); returns the packed + ``[..., n/32*34]`` uint8 layout the ggml Q8_0 kernels read. Used to re-quantize + K-quant expert banks to a uniform 8-bit type (Q8_0 >= Q5_K/Q6_K precision, so no + quality loss) when the offload cache needs a single per-bank format. + """ + n = w.shape[-1] + assert n % 32 == 0, f"Q8_0 quantize needs last dim % 32 == 0, got {n}" + wq = w.float().view(*w.shape[:-1], n // 32, 32) + d = wq.abs().amax(dim=-1, keepdim=True).clamp(min=1e-9) / 127.0 + q = torch.round(wq / d).to(torch.int8) + dh = d.to(torch.float16).view(torch.uint8) # [..., n//32, 2] + packed = torch.cat([dh, q.view(torch.uint8)], dim=-1) # [..., n//32, 34] + return packed.reshape(*w.shape[:-1], (n // 32) * 34).contiguous() + + +def dequant_q5_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + """Q5_K: 256-elem super-block = half2 dm (dall, dmin), 12B 6-bit scale/min, 32B + qh high-bits, 128B qs low nibbles. Mirrors ggml's dequantize_block_q5_K.""" + raw = raw.reshape(-1, 176) + n = raw.shape[0] + dm = raw[:, 0:4].view(torch.float16).to(torch.float32) # [n,2] + dall, dmin = dm[:, 0], dm[:, 1] + scales = raw[:, 4:16] + qh = raw[:, 16:48].to(torch.int32) + qs = raw[:, 48:176].to(torch.int32) + + def _sm(j): + if j < 4: + d = scales[:, j] & 63 + m = scales[:, j + 4] & 63 + else: + d = (scales[:, j + 4] & 0xF) | ((scales[:, j - 4] >> 6) << 4) + m = (scales[:, j + 4] >> 4) | ((scales[:, j] >> 6) << 4) + return d.to(torch.float32), m.to(torch.float32) + + y = torch.zeros(n, 256, dtype=torch.float32, device=raw.device) + for il in range(4): + s0, m0 = _sm(2 * il) + s1, m1 = _sm(2 * il + 1) + d0, M0 = dall * s0, dmin * m0 + d1, M1 = dall * s1, dmin * m1 + bit0 = 1 << (2 * il) + bit1 = bit0 << 1 + ql = qs[:, 32 * il:32 * il + 32] + ql0, ql1 = ql[:, 0::2], ql[:, 1::2] + h0, h1 = qh[:, 0::2], qh[:, 1::2] + v0 = (ql0 & 0xF) + ((h0 & bit0) != 0).to(torch.float32) * 16 + v1 = (ql1 & 0xF) + ((h1 & bit0) != 0).to(torch.float32) * 16 + even = torch.stack([v0, v1], dim=-1).reshape(n, 32) # [v0[0],v1[0],v0[1],...] + y[:, 64 * il:64 * il + 32] = even * d0.unsqueeze(1) - M0.unsqueeze(1) + w0 = (ql0 >> 4) + ((h0 & bit1) != 0).to(torch.float32) * 16 + w1 = (ql1 >> 4) + ((h1 & bit1) != 0).to(torch.float32) * 16 + odd = torch.stack([w0, w1], dim=-1).reshape(n, 32) + y[:, 64 * il + 32:64 * il + 64] = odd * d1.unsqueeze(1) - M1.unsqueeze(1) + return y.reshape(-1).to(out_dtype) + + _DEQUANT = { GGML_Q4_0: dequant_q4_0, + GGML_Q5_K: dequant_q5_k, GGML_Q6_K: dequant_q6_k, } @@ -143,11 +209,14 @@ def dequantize(raw: torch.Tensor, ggml_type: int, out_dtype: torch.dtype) -> tor "GGML_BF16", "GGML_Q4_0", "GGML_Q8_0", + "GGML_Q4_K", "GGML_Q6_K", "GGML_NAME", "BLOCK_SHAPE", "row_bytes", "dequant_q4_0", + "dequant_q5_k", "dequant_q6_k", + "quantize_q8_0", "dequantize", ] diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index 6d5481c1..c0582419 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -13,7 +13,10 @@ from .reader import gguf_architecture, load_gguf_metadata # GGUF architecture -> transformers GGUF tokenizer-converter key. -_TOKENIZER_ARCH = {"gemma4": "gemma4_text"} +_TOKENIZER_ARCH = { + "gemma4": "gemma4_text", + "qwen35moe": "qwen3_moe", +} def load_gguf_tokenizer(model_path: str): diff --git a/python/freetoken/models/qwen3_5_moe/__init__.py b/python/freetoken/models/qwen3_5_moe/__init__.py index 98936e9f..cae7dfd3 100644 --- a/python/freetoken/models/qwen3_5_moe/__init__.py +++ b/python/freetoken/models/qwen3_5_moe/__init__.py @@ -1,4 +1,11 @@ from .config import parse_config +from .gguf import ( + convert_qwen35moe_to_gguf, + is_gguf_model, + iter_gguf_weights, + load_gguf_expert_sources, + parse_gguf_config, +) from .model import Qwen3_5MoEForCausalLM from .weight import ( iter_weights, @@ -16,4 +23,9 @@ "load_nvfp4_expert_sources", "load_nvfp4_expert_sources_parallel", "setup_offload_expert_banks", + "parse_gguf_config", + "iter_gguf_weights", + "convert_qwen35moe_to_gguf", + "is_gguf_model", + "load_gguf_expert_sources", ] diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index 2e732005..7b7f227c 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -76,10 +76,11 @@ def __init__( # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM). self._block_fp8 = expert_quant == "fp8_block" self._pertensor_fp8 = attn_quant == "fp8_pertensor" - self._fp8 = self._block_fp8 or self._pertensor_fp8 + self._gguf = expert_quant == "gguf" + self._fp8 = self._block_fp8 or self._pertensor_fp8 or self._gguf self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] - if self._fp8: + if self._block_fp8 or self._pertensor_fp8: ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged self.in_proj_qkvz = ColMerged( hidden_size, [self.conv_dim, self.value_dim], has_bias=False @@ -87,6 +88,15 @@ def __init__( self.in_proj_ba = LinearColParallelMerged( hidden_size, [num_v_heads, num_v_heads], has_bias=False ) + elif self._gguf: + # GGUF: qkv|z are native-quant (Q8_0, swapped to GGUFLinear after build), b|a + # stay dense bf16. Same split as the fp8 path (matches the GGUF tensor layout). + self.in_proj_qkvz = LinearColParallelMerged( + hidden_size, [self.conv_dim, self.value_dim], has_bias=False + ) + self.in_proj_ba = LinearColParallelMerged( + hidden_size, [num_v_heads, num_v_heads], has_bias=False + ) else: # Fused input projection (one GEMM instead of four): qkv | z | b | a. self.in_proj = LinearColParallelMerged(hidden_size, self._in_proj_split, has_bias=False) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py new file mode 100644 index 00000000..aef81a09 --- /dev/null +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -0,0 +1,512 @@ +"""Qwen3.5-MoE GGUF adapter: build the FreeToken ``ModelConfig`` from GGUF metadata and +map GGUF tensors to the model's state dict. + +The qwen35moe GGUF arch is a hybrid GatedDeltaNet (linear-attention SSM) + full-attention +MoE (40 layers, every 4th full; 256 routed experts + shared expert). The GGUF geometry +matches the HF qwen3_5_moe model, so ``parse_gguf_config`` produces the *same* +``ModelConfig`` as ``qwen3_5_moe.config.parse_config`` -- only the source is GGUF KV +metadata. ``expert_quant``/``attn_quant``/``dense_quant`` are set to ``"gguf"`` so +``convert_qwen35moe_to_gguf`` can detect the native-quant checkpoint and swap the dense +projections for native GGUF-quant ops. + +Quantized projections (full-attn q/k/v/o, GDN qkv|z and out_proj, shared-expert gate/up/ +down, the token embedding and the lm_head) stay in their native packed block layout (Q8_0 +projections, Q6_K head) and are yielded as ``.qweight`` (uint8); tiny F32 tensors (norms, +router, GDN b/a) dequantize to bf16; GDN conv/A_log/dt_bias stay fp32. Routed experts +(Q4_K gate/up, Q5_K/Q6_K down) go to the offload cache (``load_gguf_expert_sources``). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.config import ( + FullAttentionGroupConfig, + LinearGatedDeltaGroupConfig, + ModelConfig, + RotaryConfig, +) +from freetoken.models.gguf.dequant import ( + GGML_F32, + GGML_Q4_K, + GGML_Q8_0, + dequantize, + quantize_q8_0, + row_bytes, +) + +if TYPE_CHECKING: + from freetoken.models.gguf.config import GgufConfigShim + + +def _require_tp1(what: str) -> None: + from freetoken.distributed import get_tp_info + + if get_tp_info().size > 1: + raise NotImplementedError( + f"qwen3.5-moe GGUF {what} currently supports TP=1 only " + "(GGUF quant layers and expert banks are not tensor-parallel sharded)." + ) + + +def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: + m = shim.metadata + + def g(key: str): + val = m.get(f"qwen35moe.{key}") + if val is None: + raise KeyError(f"missing GGUF metadata key qwen35moe.{key}") + return val + + num_layers = int(g("block_count")) + hidden = int(g("embedding_length")) + num_qo_heads = int(g("attention.head_count")) + num_kv_heads = int(g("attention.head_count_kv")) + full_head_dim = int(g("attention.key_length")) + max_pos = int(g("context_length")) + interval = int(g("full_attention_interval")) # every Nth (1-indexed) layer is full + + full_ids = tuple(i for i in range(num_layers) if (i + 1) % interval == 0) + linear_ids = tuple(i for i in range(num_layers) if (i + 1) % interval != 0) + + full_rotary = RotaryConfig( + head_dim=full_head_dim, + rotary_dim=int(g("rope.dimension_count")), + max_position=max_pos, + base=float(g("rope.freq_base")), + scaling=None, + ) + full_group = FullAttentionGroupConfig( + name="full", + layer_ids=full_ids, + num_kv_heads=num_kv_heads, + head_dim=full_head_dim, + rotary_config=full_rotary, + ) + # GDN (SSM) dims. conv_dim = 2*key_dim + value_dim (attn_qkv); value_dim = nv*vhead_dim. + key_head_dim = int(g("ssm.state_size")) + value_head_dim = int(g("ssm.state_size")) + linear_group = LinearGatedDeltaGroupConfig( + name="linear", + layer_ids=linear_ids, + num_key_heads=int(g("ssm.group_count")), + num_value_heads=int(g("ssm.time_step_rank")), + key_head_dim=key_head_dim, + value_head_dim=value_head_dim, + conv_kernel_dim=int(g("ssm.conv_kernel")), + output_gate=True, + ) + groups = tuple(sorted((full_group, linear_group), key=lambda grp: grp.layer_ids[0] or 1 << 30)) + + return ModelConfig( + num_layers=num_layers, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=full_head_dim, + hidden_size=hidden, + vocab_size=int(shim.vocab_size), + intermediate_size=0, # routed MoE; no dense MLP + hidden_act="silu", + rms_norm_eps=float(g("attention.layer_norm_rms_epsilon")), + tie_word_embeddings=bool(shim.tie_word_embeddings), + rotary_config=full_rotary, + num_experts=int(g("expert_count")), + num_experts_per_tok=int(g("expert_used_count")), + moe_intermediate_size=int(g("expert_feed_forward_length")), + shared_expert_intermediate_size=int(g("expert_shared_feed_forward_length")), + norm_topk_prob=True, + model_type="qwen3_5_moe", + architectures=list(shim.architectures), + moe_enabled=True, + use_qk_norm=True, + attention_groups=groups, + expert_quant="gguf", + attn_quant="gguf", + dense_quant="gguf", + lm_head_quant="gguf", + moe_weight_format="gguf", + ) + + +def is_gguf_model(config: ModelConfig) -> bool: + return getattr(config, "moe_weight_format", None) == "gguf" + + +# -------------------------------------------------------------------------------------- +# Model layer swap: dense bf16 Linear -> native GGUF-quant ops. +# -------------------------------------------------------------------------------------- + + +def convert_qwen35moe_to_gguf(model, config: ModelConfig) -> None: + """In place: replace the dense projections + embedding with native GGUF ops. + + Quantized in the checkpoint -> swapped to ``GGUFLinear``/``GGUFEmbedding``: the token + embedding (Q8_0) and the (untied) lm_head (Q6_K), full-attention qkv/o (Q8_0), GDN + in_proj_qkvz + out_proj (Q8_0; in_proj_ba stays dense bf16), and the shared-expert + gate_up/down (Q8_0). Left dense bf16/fp32 (F32 in the GGUF): the norms, the two + routers, and the GDN conv1d/A_log/dt_bias. Routed experts stay on the offload cache. + """ + from freetoken.layers.gguf import GGUFEmbedding, GGUFLinear + from freetoken.models.gguf.dequant import GGML_Q6_K, GGML_Q8_0 + + inner = model.model + inner.embed_tokens = GGUFEmbedding( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + quant_type=GGML_Q8_0, + ) + model.lm_head = GGUFLinear( + config.hidden_size, config.vocab_size, GGML_Q6_K, has_bias=False + ) + shared_I = config.shared_expert_intermediate_size + + for layer in inner.layers.op_list: + if layer._is_linear: + g = layer.linear_attn + g.in_proj_qkvz = GGUFLinear( + config.hidden_size, g.conv_dim + g.value_dim, GGML_Q8_0, has_bias=False + ) + g.out_proj = GGUFLinear( + g.value_dim, config.hidden_size, GGML_Q8_0, has_bias=False + ) + else: + attn = layer.self_attn + attn.qkv_proj = GGUFLinear( + config.hidden_size, + attn.num_q * attn.head_dim * 2 + 2 * attn.kv_attn_dim, + GGML_Q8_0, + has_bias=False, + ) + attn.o_proj = GGUFLinear( + attn.qo_attn_dim, config.hidden_size, GGML_Q8_0, has_bias=False + ) + m = layer.mlp + m.shared_expert.gate_up_proj = GGUFLinear( + config.hidden_size, 2 * shared_I, GGML_Q8_0, has_bias=False + ) + m.shared_expert.down_proj = GGUFLinear( + shared_I, config.hidden_size, GGML_Q8_0, has_bias=False + ) + + +# -------------------------------------------------------------------------------------- +# Weight loading: GGUF tensor names -> FreeToken qwen3_5_moe module params. +# -------------------------------------------------------------------------------------- + +# Gemma-style (1+w) norms get +1 baked in; the GDN gated norm / router / shared-gate are +# standard (no +1). +_GEMMA_NORMS = { + "attn_norm.weight", + "post_attention_norm.weight", + "attn_q_norm.weight", + "attn_k_norm.weight", +} + +_EXPERT_SUFFIXES = ("ffn_gate_exps.weight", "ffn_up_exps.weight", "ffn_down_exps.weight") + + +def _to_bf16(t) -> torch.Tensor: + flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.bfloat16) + return flat.reshape(t.shape) + + +def _to_fp32(t) -> torch.Tensor: + flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.float32) + return flat.reshape(t.shape) + + +def _name_to_key(suffix: str) -> tuple[str, bool]: + """gguf layer suffix -> (module-relative key, is_gemma_norm).""" + if suffix == "attn_norm.weight": + return "input_layernorm.weight", True + if suffix == "post_attention_norm.weight": + return "post_attention_layernorm.weight", True + if suffix == "attn_q_norm.weight": + return "self_attn.q_norm.weight", True + if suffix == "attn_k_norm.weight": + return "self_attn.k_norm.weight", True + if suffix == "ssm_norm.weight": + return "linear_attn.norm.weight", False + if suffix == "ffn_gate_inp.weight": + return "mlp.gate.weight", False + if suffix == "ffn_gate_inp_shexp.weight": + return "mlp.shared_expert_gate.weight", False + if suffix == "ssm_conv1d.weight": + return "linear_attn.conv1d.weight", False # fp32; reshaped below + if suffix == "ssm_a": + return "linear_attn.A_log", False # fp32 + if suffix == "ssm_dt.bias": + return "linear_attn.dt_bias", False # fp32 + return None, False + + +# -------------------------------------------------------------------------------------- +# GDN value-head de-interleaving. +# +# llama.cpp stores the GDN *value* projections with the ``mrope_interleaved`` head order: +# the 32 value heads are split into [even heads, odd heads] (head h lives at GGUF position +# ``(h // 2) + (h % 2) * (num_vheads // 2)``). Full-attention heads are NOT interleaved. +# FreeToken uses the HF contiguous head order, so the value-dim projection weights must be +# de-interleaved when loading. Affected: GDN ``in_proj_qkvz`` (the v and z rows), ``out_proj`` +# (the value input columns), and ``in_proj_ba`` (the per-head b/a rows). +# -------------------------------------------------------------------------------------- + + +def _gdn_head_perm(num_vheads: int) -> list[int]: + """GGUF value-head index of each HF head h (``result[h] = old[perm[h]]``).""" + half = num_vheads // 2 + return [(h // 2) + (h % 2) * half for h in range(num_vheads)] + + +def _deint_q8_rows( + packed: torch.Tensor, num_vheads: int, rows_per_head: int +) -> torch.Tensor: + """De-interleave value heads along the packed rows (output dim).""" + m = packed.reshape(num_vheads, rows_per_head, -1) + return m[_gdn_head_perm(num_vheads)].reshape(packed.shape) + + +def _deint_q8_cols( + packed: torch.Tensor, num_vheads: int, blocks_per_head: int, block_bytes: int = 34 +) -> torch.Tensor: + """De-interleave value heads along the packed columns (input dim, per Q8_0 row).""" + m = packed.reshape(packed.shape[0], num_vheads, blocks_per_head * block_bytes) + return m[:, _gdn_head_perm(num_vheads), :].reshape(packed.shape) + + +def _deint_dense_rows(w: torch.Tensor, num_vheads: int) -> torch.Tensor: + """De-interleave value heads along the leading (head) dim of a dense tensor.""" + m = w.reshape(num_vheads, -1) + return m[_gdn_head_perm(num_vheads)].reshape(w.shape) + + +def iter_gguf_weights( + model_path: str, + device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """Yield (param_name, tensor) for every non-expert qwen3_5_moe param. + + Quantized projections stay packed and are yielded as ``.qweight`` (uint8); the F32 + norms/router/GDN b,a dequantize to bf16; conv1d/A_log/dt_bias stay fp32. Full-attention + q/k/v -> ``self_attn.qkv_proj.qweight``, GDN qkv|z -> ``linear_attn.in_proj_qkvz.qweight`` + (Q8_0, concat along the output dim), GDN b|a -> ``linear_attn.in_proj_ba.weight`` (dense + bf16). Routed experts are skipped (offload cache). + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + from freetoken.utils import cached_load_hf_config + + assert not include_moe_experts, ( + "qwen3.5-moe GGUF stores experts as Q4_K/Q5_K/Q6_K and only supports the offload " + "backend; experts are loaded into the offload cache via load_gguf_expert_sources()." + ) + assert include_non_moe + _require_tp1("weight loading") + + config = parse_gguf_config(cached_load_hf_config(model_path)) + full_layers = set( + next(g.layer_ids for g in config.attention_groups + if isinstance(g, FullAttentionGroupConfig)) + ) + # GDN value-head geometry (for mrope_interleaved de-interleave). + gdn = next(g for g in config.attention_groups + if isinstance(g, LinearGatedDeltaGroupConfig)) + n_vheads = gdn.num_value_heads + vhead_dim = gdn.value_head_dim + key_dim = gdn.num_key_heads * gdn.key_head_dim + q8_blocks_per_head = vhead_dim // 32 # Q8_0 block = 32 + + qkv_buf: dict[int, dict[str, torch.Tensor]] = {} + qkvz_buf: dict[int, dict[str, torch.Tensor]] = {} + ba_buf: dict[int, dict[str, torch.Tensor]] = {} + shexp_buf: dict[int, dict[str, torch.Tensor]] = {} + + for t in iter_gguf_tensors(model_path): + name = t.name + if name == "token_embd.weight": + yield "model.embed_tokens.qweight", t.packed() + continue + if name == "output.weight": + yield "lm_head.qweight", t.packed() + continue + if name == "output_norm.weight": + # GGUF stores Gemma norms as the full (1+w) scale; GemmaRMSNorm multiplies by + # it directly (no +1, unlike the HF safetensors form which stores scale-1). + yield "model.norm.weight", _to_bf16(t) + continue + if not name.startswith("blk."): + continue + if any(name.endswith(sfx) for sfx in _EXPERT_SUFFIXES): + continue # routed experts -> offload banks + + layer = int(name.split(".")[1]) + suffix = name.split(".", 2)[2] + base = f"model.layers.{layer}" + + if suffix == "ssm_conv1d.weight": + # conv channels span [q|k|v]; the v channels are value-head interleaved too. + c = _to_fp32(t) + c = c.clone() + c[key_dim * 2:] = _deint_dense_rows(c[key_dim * 2:], n_vheads) + yield f"{base}.linear_attn.conv1d.weight", c.unsqueeze(1) + continue + if suffix == "ssm_a": + # GGUF stores the GDN decay directly as ``A = -exp(A_log)`` (mamba convention, + # value-head interleaved); recover the log-decay the model consumes. + a = _deint_dense_rows(_to_fp32(t), n_vheads) + yield f"{base}.linear_attn.A_log", torch.log(-a) + continue + if suffix == "ssm_dt.bias": + yield f"{base}.linear_attn.dt_bias", _deint_dense_rows(_to_fp32(t), n_vheads) + continue + if suffix == "ffn_gate_inp_shexp.weight": + yield f"{base}.mlp.shared_expert_gate.weight", _to_bf16(t).unsqueeze(0) + continue + + # The GDN b and a projections fuse into a dense in_proj_ba. + if suffix == "ssm_beta.weight": + ba_buf.setdefault(layer, {})["b"] = _to_bf16(t) + elif suffix == "ssm_alpha.weight": + ba_buf.setdefault(layer, {})["a"] = _to_bf16(t) + else: + key, _gemma = _name_to_key(suffix) + if key is not None: + yield f"{base}.{key}", _to_bf16(t) + continue + + is_full = layer in full_layers + if is_full and suffix in ("attn_q.weight", "attn_k.weight", "attn_v.weight"): + qkv_buf.setdefault(layer, {})[suffix[5]] = t.packed() + elif suffix == "attn_qkv.weight": + qkvz_buf.setdefault(layer, {})["qkv"] = t.packed() + elif suffix == "attn_gate.weight": + qkvz_buf.setdefault(layer, {})["z"] = t.packed() + elif suffix == "attn_output.weight": + yield f"{base}.self_attn.o_proj.qweight", t.packed() + elif suffix == "ssm_out.weight": + yield f"{base}.linear_attn.out_proj.qweight", _deint_q8_cols( + t.packed(), n_vheads, q8_blocks_per_head) + elif suffix == "ffn_gate_shexp.weight": + shexp_buf.setdefault(layer, {})["gate"] = t.packed() + elif suffix == "ffn_up_shexp.weight": + shexp_buf.setdefault(layer, {})["up"] = t.packed() + elif suffix == "ffn_down_shexp.weight": + yield f"{base}.mlp.shared_expert.down_proj.qweight", t.packed() + else: + raise ValueError(f"unmapped qwen3.5-moe GGUF tensor: {name}") + + slots = qkv_buf.get(layer) + if slots is not None and {"q", "k", "v"} <= set(slots): + yield f"{base}.self_attn.qkv_proj.qweight", torch.cat( + [slots["q"], slots["k"], slots["v"]], dim=0) + del qkv_buf[layer] + qz = qkvz_buf.get(layer) + if qz is not None and "qkv" in qz and "z" in qz: + qkv = qz["qkv"] # [2*key_dim + value_dim, cols] + qkv = qkv.clone() + # de-interleave the value rows (last value_dim rows of the qkv projection) + qkv[key_dim * 2:] = _deint_q8_rows( + qkv[key_dim * 2:], n_vheads, vhead_dim) + z = _deint_q8_rows(qz["z"], n_vheads, vhead_dim) + yield f"{base}.linear_attn.in_proj_qkvz.qweight", torch.cat([qkv, z], dim=0) + del qkvz_buf[layer] + ba = ba_buf.get(layer) + if ba is not None and "b" in ba and "a" in ba: + b = _deint_dense_rows(ba["b"], n_vheads) + a = _deint_dense_rows(ba["a"], n_vheads) + yield f"{base}.linear_attn.in_proj_ba.weight", torch.cat([b, a], dim=0) + del ba_buf[layer] + gu = shexp_buf.get(layer) + if gu is not None and "gate" in gu and "up" in gu: + yield f"{base}.mlp.shared_expert.gate_up_proj.qweight", torch.cat( + [gu["gate"], gu["up"]], dim=0) + del shexp_buf[layer] + + assert not qkv_buf, f"incomplete qkv groups: {sorted(qkv_buf)}" + assert not qkvz_buf, f"incomplete GDN qkvz groups: {sorted(qkvz_buf)}" + assert not ba_buf, f"incomplete GDN ba groups: {sorted(ba_buf)}" + assert not shexp_buf, f"incomplete shared-expert gate/up: {sorted(shexp_buf)}" + + +# -------------------------------------------------------------------------------------- +# Routed-expert host banks for the offload cache. +# +# The GGUF stores gate/up as Q4_K on every layer, but ``down`` as Q5_K on 37 layers and +# Q6_K on 3 -- heterogeneous row widths the offload cache cannot hold in one uniform bank +# (``set_bank_sources`` requires every layer to share a shape, and ``ggml_moe_a8_vec`` +# derives the row stride from the quant type). We keep ``gate_up`` native Q4_K and +# re-quantize the ``down`` experts to Q8_0 (8-bit, >= Q5_K/Q6_K precision, so no quality +# loss; a uniform per-bank format that fits the cache machinery). +# -------------------------------------------------------------------------------------- + + +def _q8_0_down_row_bytes(I: int) -> int: + return row_bytes(I, GGML_Q8_0) + + +def load_gguf_expert_sources( + model_path: str, config: ModelConfig, *, layer_sink=None +) -> dict[str, list[torch.Tensor]]: + """Per-layer host banks of the routed experts: ``gate_up`` native Q4_K + ``[E, 2I, row_bytes(H, Q4_K)]`` and ``down`` Q8_0 ``[E, H, row_bytes(I, Q8_0)]``. + + ``ffn_{gate,up}_exps`` are each ``[E, I, row_bytes(H, Q4_K)]`` packed and are fused + along the intermediate dim into ``gate_up``; ``ffn_down_exps`` (Q5_K/Q6_K) is + dequantized and re-quantized to Q8_0. Whole layers complete in two writes + (gate_up + down). ``layer_sink=None`` (serving): pin each layer's banks as they + complete via an internal :class:`PinPipeline`. + """ + from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, alloc_layer_banks + from freetoken.models.gguf.reader import iter_gguf_tensors + + _require_tp1("expert banks") + L, E = config.num_layers, config.num_experts + H, I = config.hidden_size, config.moe_intermediate_size + h_rb = row_bytes(H, GGML_Q4_K) + i_rb = row_bytes(I, GGML_Q8_0) + specs = { + "gate_up": ((E, 2 * I, h_rb), torch.uint8), + "down": ((E, H, i_rb), torch.uint8), + } + hb = alloc_layer_banks(specs, L) + banks = {name: [b.tensor for b in hb[name]] for name in specs} + + def _load(sink) -> None: + tracker = LayerCompletionTracker(2, hb, sink) if sink is not None else None + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + continue + layer = int(t.name.split(".")[1]) + if t.name.endswith("ffn_gate_exps.weight"): + banks["gate_up"][layer][:, :I].copy_(t.packed().reshape(E, I, h_rb)) + elif t.name.endswith("ffn_up_exps.weight"): + banks["gate_up"][layer][:, I:] = t.packed().reshape(E, I, h_rb) + elif t.name.endswith("ffn_down_exps.weight"): + flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.float32) + banks["down"][layer].copy_(quantize_q8_0(flat.reshape(E, H, I))) + else: + continue + if tracker is not None: + tracker.note(layer) + + if layer_sink is not None: + _load(layer_sink) + elif torch.cuda.is_available(): + with PinPipeline() as pins: + _load(pins) + else: + _load(None) + return banks + + +__all__ = [ + "parse_gguf_config", + "iter_gguf_weights", + "convert_qwen35moe_to_gguf", + "is_gguf_model", + "load_gguf_expert_sources", +] diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py index eba7fd24..32954dc5 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -109,6 +109,13 @@ def __init__(self, config: ModelConfig): ) super().__init__() + # GGUF checkpoints carry native block-quantized weights: swap the dense + # projections + embedding for GGUF-quant ops (experts stay on the offload cache). + from .gguf import convert_qwen35moe_to_gguf, is_gguf_model + + if is_gguf_model(config): + convert_qwen35moe_to_gguf(self, config) + def forward(self) -> torch.Tensor: output = self.model.forward(get_global_ctx().batch.input_ids) return self.lm_head.forward(output) diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 0c033ca0..7cffb20c 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -58,6 +58,14 @@ class ModelSpec: "freetoken.models.qwen3_5_moe", "Qwen3_5MoEForCausalLM", ), + # GGUF (native Q4_K/Q5_K/Q6_K/Q8_0) qwen3.5-moe: same model classes, GGUF config + + # weight loaders (hybrid GatedDeltaNet + full attention, 256 routed experts). + "Qwen35moeGGUFForCausalLM": ModelSpec( + "freetoken.models.qwen3_5_moe", + "Qwen3_5MoEForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + ), # Dense Qwen3.x (no "Moe" in the arch name, num_experts==0, e.g. Qwen3.6-27B). Shares the # qwen3_5_moe package: the decoder routes its MLP through the dense Qwen3_5DenseMLP and the # loader handles the compressed-tensors NVFP4 layout. diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 6a34f3b9..636d451e 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -342,6 +342,20 @@ def load_q4_0_moe_expert_sources( return loader(model_path, model_config, layer_sink=layer_sink) +def load_gguf_moe_expert_sources( + model_path: str, + model_config, + *, + layer_sink=None, +) -> dict: + """Load packed GGUF qwen3.5-moe expert source banks (gate_up native Q4_K, down + re-quantized to Q8_0). ``layer_sink`` (converter) streams each completed layer's + banks.""" + _config, spec = _spec_for_model_path(model_path) + loader = _load_attr(spec.module, "load_gguf_expert_sources") + return loader(model_path, model_config, layer_sink=layer_sink) + + def _num_moe_layers(config) -> int: value = getattr(config, "num_moe_layers", None) if value is not None: diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 8b6116ba..e0801232 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -252,6 +252,25 @@ def _q4_0_banks(model_path, model_config, device, dtype, dummy, parallel=False, ) +def _gguf_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks: + if parallel: + raise NotImplementedError( + "parallel reader not implemented for gguf: GGUF is a single packed file " + "(not safetensors), so the common reader doesn't apply." + ) + if dummy: + from freetoken.models.weight import dummy_q4_0_moe_expert_sources + + raise NotImplementedError("gguf expert banks have no dummy path; load the real GGUF") + from freetoken.models.weight import load_gguf_moe_expert_sources + + sink = None if dummy else layer_sink + sources = load_gguf_moe_expert_sources(model_path, model_config, layer_sink=sink) + return ExpertBanks( + "gguf", {name: sources[name] for name in _BANK_SCHEMAS["gguf"]}, streamed=sink is not None + ) + + def _dsfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks: args = model_config.dsv4_args assert args is not None, "ds_fp4 expert banks require dsv4_args on the model config" @@ -301,6 +320,7 @@ def _model_setup_override(model_config): "nvfp4": _nvfp4_banks, "ds_fp4": _dsfp4_banks, "q4_0": _q4_0_banks, + "gguf": _gguf_banks, } diff --git a/python/freetoken/moe/fused_gguf.py b/python/freetoken/moe/fused_gguf.py new file mode 100644 index 00000000..9f8995e9 --- /dev/null +++ b/python/freetoken/moe/fused_gguf.py @@ -0,0 +1,52 @@ +"""Grouped expert GEMM over native GGUF Q4_K gate/up + Q8_0 down banks. + +Ports vLLM/sglang's ``_fused_moe_gguf`` MMVQ path onto FreeToken's offload-cache +interface: experts are streamed to the GPU as packed block bytes and dequantized +*inside* ``ggml_moe_a8_vec`` -- no bf16 expert copy is materialized. ``gate_up`` stays +native Q4_K; ``down`` is stored as Q8_0 (re-quantized at load from the GGUF's +Q5_K/Q6_K -- 8-bit, >= the source precision, so no quality loss) because the offload +cache needs a single uniform per-bank format. We use the MMVQ (vector) kernel for both +prefill and decode, mirroring ``fused_experts_gguf_q4_0``. +""" + +from __future__ import annotations + +import torch + +from freetoken.layers.activation import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul +from freetoken.models.gguf.dequant import GGML_Q4_K, GGML_Q8_0 + +_ACT = {"silu": silu_and_mul, "gelu": gelu_and_mul, "gelu_tanh": gelu_tanh_and_mul} + + +def fused_experts_gguf( + hidden_states: torch.Tensor, + gate_up_q: torch.Tensor, # [num_slots, 2I, row_bytes(H, Q4_K)] uint8 + down_q: torch.Tensor, # [num_slots, H, row_bytes(I, Q8_0)] uint8 + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: str, +) -> torch.Tensor: + from freetoken.kernel.gguf import ggml_moe_a8_vec + + act_fn = _ACT.get(activation) + if act_fn is None: + raise ValueError(f"unsupported MoE activation {activation!r}") + + num_tokens = hidden_states.shape[0] + n2 = gate_up_q.shape[1] # 2 * intermediate + h = down_q.shape[1] # hidden + top_k = topk_ids.shape[1] + + gate_up = ggml_moe_a8_vec( + hidden_states, gate_up_q, topk_ids, top_k, int(GGML_Q4_K), n2, num_tokens + ) + inter = act_fn(gate_up) + out = ggml_moe_a8_vec(inter, down_q, topk_ids, 1, int(GGML_Q8_0), h, num_tokens * top_k) + out = out.reshape(num_tokens, top_k, h) * topk_weights.reshape(num_tokens, top_k, 1).to( + out.dtype + ) + return out.sum(dim=1) + + +__all__ = ["fused_experts_gguf"] diff --git a/python/freetoken/moe/nvfp4_backends.py b/python/freetoken/moe/nvfp4_backends.py index e8a3c6f6..c9f84f78 100644 --- a/python/freetoken/moe/nvfp4_backends.py +++ b/python/freetoken/moe/nvfp4_backends.py @@ -46,6 +46,7 @@ import torch from freetoken.utils import init_logger +from freetoken.utils.arch import is_rocm logger = init_logger(__name__) @@ -213,6 +214,17 @@ def select_nvfp4_backend( raise ValueError( f"bad --nvfp4-backend={requested!r}; expected auto, marlin, flashinfer or triton" ) + # AMD: the marlin (vLLM) and flashinfer b12x fused-MoE kernels are NVIDIA-only + # (NVFP4 SASS / CuTe-DSL sm120). auto -> the portable Triton inline-dequant path; + # a forced NVIDIA-only backend fails loudly, never silently degrades. + if device.type == "cuda" and is_rocm(): + if requested in ("marlin", "flashinfer"): + raise RuntimeError( + f"--nvfp4-backend={requested} is NVIDIA-only and unavailable on this " + "ROCm (AMD) build; use --nvfp4-backend triton (or auto). NVFP4 checkpoints " + "can be converted to MXFP4 (freetoken.moe.nvfp4_to_mxfp4) on load." + ) + return "triton" if requested == "triton": return "triton" if activation != "silu": diff --git a/python/freetoken/moe/nvfp4_to_mxfp4.py b/python/freetoken/moe/nvfp4_to_mxfp4.py new file mode 100644 index 00000000..f4162a03 --- /dev/null +++ b/python/freetoken/moe/nvfp4_to_mxfp4.py @@ -0,0 +1,239 @@ +"""NVFP4 -> MXFP4 (gpt-oss/FreeToken ``mxfp4_triton``) weight converter. + +ModelOpt NVFP4 (the format stored in FreeToken's native ``nvfp4`` banks) packs e2m1 +codes with a *fp8-e4m3* per-16 block scale and a per-output-row fp16 *global* scale. +MXFP4 (FreeToken's ``mxfp4_triton`` banks) packs e2m1 codes with an *e8m0* per-32 +block scale and a per-block bias, and is the native format of the gpt-oss family -- +the AMD-supported quant matrix alongside BF16/GGUF. + +This module converts a checkpoint's NVFP4 expert weights to MXFP4 *on load* (once, +cached per model), so a checkpoint that only ships NVFP4 can still run on AMD via the +portable MXFP4 path (the converter runs on the host, not on the GPU). The two formats +share the e2m1 code-packing (2 codes per byte, low nibble first), so only the scale +granularity (16 -> 32) and scale format (e4m3 -> e8m0) change. + +Layouts handled here (per projection/expert, ``N`` = output rows, ``K`` = input cols): + +* input (NVFP4, native ModelOpt rows): ``packed [N, K//2]`` uint8, ``scale [N, K//16]`` + fp8-e4m3, ``global [N]`` fp16. +* output (MXFP4, ``mxfp4_triton`` bank layout): ``blocks_t [N, K//2]`` uint8, + ``scales_t [N, K//32]`` uint8 e8m0. (FreeToken's MXFP4 stores the projection + transposed, N innermost, matching gpt-oss -- the converter emits that shape.) + +All numeric work is done in numpy so the core is unit-testable without a GPU/torch +runtime; the public entry converts torch tensors to/from numpy on the host. +""" + +from __future__ import annotations + +import math +from typing import Sequence + +try: + import numpy as _np +except ImportError: # pragma: no cover - numpy is a hard dep + _np = None + +__all__ = [ + "convert_nvfp4_to_mxfp4", + "dequantize_nvfp4_block", + "fp4_e2m1_table", + "e8m0_scale_and_codes", +] + +# --------------------------------------------------------------------------- +# e2m1 (fp4) code -> value table, and e8m0 (block scale) encode. +# +# e2m1: 1 sign + 2 exponent + 1 mantissa. With exponent bias 1 the finite set is: +# code value code value +# 0 0.0 8 -0.0 +# 1 0.5 9 -0.5 +# 2 1.0 10 -1.0 +# 3 1.5 11 -1.5 +# 4 2.0 12 -2.0 +# 5 3.0 13 -3.0 +# 6 4.0 14 -4.0 +# 7 6.0 15 -6.0 +# --------------------------------------------------------------------------- +_FP4_CODES = ( + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, +) +_FP4_TABLE = _np.asarray(_FP4_CODES, dtype=_np.float32) +# Magnitudes sorted ascending for the nearest-code search. +_FP4_SORT = _np.asarray(sorted(abs(v) for v in _FP4_CODES[1:8]), dtype=_np.float32) +_FP4_SORT_SIGN = _np.asarray([1.0 if i < 4 else -1.0 for i in range(len(_FP4_SORT))], + dtype=_np.float32) + + +def fp4_e2m1_table() -> Sequence[float]: + """The 16 e2m1 values keyed by 4-bit code (index == code).""" + return list(_FP4_CODES) + + +def _nearest_e2m1_codes(values: _np.ndarray) -> _np.ndarray: + """Nearest e2m1 *code* for each fp32 ``values`` (signed, including 0/NaN).""" + a = _np.abs(values) + diff = _np.abs(a[..., None] - _FP4_SORT) # [..., 7] + idx = diff.argmin(axis=-1) + mag = _FP4_SORT[idx] + neg = _np.signbit(values) + code = (idx + 1).astype(_np.uint8) # _FP4_SORT[i] == table[i+1]; positive codes 1..7 + out = _np.where(neg, code | 0x8, code) + # Magnitudes below the smallest representable value (0.5) round to +0. + return _np.where(mag < 0.25, 0, out) + + +def e8m0_scale_and_codes(values: _np.ndarray, block: int = 32) -> tuple[_np.ndarray, _np.ndarray]: + """Return ``(scale_codes, fp4_codes)`` for ``values`` shaped ``[..., block]``: an + e8m0 ``uint8`` scale per block (the smallest power-of-2 scale covering the block + max-abs, in the MX ``2**(v-127)`` encoding) and the requantized 4-bit codes. + + ``scale_codes`` has shape ``values.shape[:-1]``; ``fp4_codes`` matches ``values``. + """ + v = values.reshape(-1, block) + amax = _np.max(_np.abs(v), axis=-1) + # e2m1 max positive magnitude is 6.0; choose the smallest power-of-2 scale so the + # block max-abs maps near the top of the e2m1 range (best precision): + # s = 2^ceil(log2(max_abs / 6.0)), stored as e8m0 code v with 2**(v-127) == s. + amax_safe = _np.maximum(amax, 1e-38) + exp = _np.ceil(_np.log2(amax_safe / 6.0)).astype(_np.float32) + exp = _np.where(amax == 0.0, 0.0, exp) + scale_codes = (127.0 + exp).astype(_np.uint8) # v-127 == exp + scale_v = (2.0 ** exp).astype(_np.float32) + # Requantize values in the block by its scale, then nearest-e2m1-code. + q = v / scale_v[:, None] + codes = _fp4_quantize(q).astype(_np.uint8) + return scale_codes.reshape(values.shape[:-1]), codes.reshape(values.shape) + + +def _fp4_quantize(values: _np.ndarray) -> _np.ndarray: + a = _np.abs(values) + diff = _np.abs(a[..., None] - _FP4_SORT) + idx = diff.argmin(axis=-1) + mag = _FP4_SORT[idx] + out = _np.where(mag < 0.25, 0, (idx + 1).astype(_np.uint8)) + out = _np.where(_np.signbit(values), out | 8, out) + return out + + +# --------------------------------------------------------------------------- +# NVFP4 -> MXFP4 +# --------------------------------------------------------------------------- + + +def dequantize_nvfp4_block( + packed: _np.ndarray, + scale: _np.ndarray, + global_scale: _np.ndarray, + *, + block: int = 16, +) -> _np.ndarray: + """Dequantize one native NVFP4 projection back to fp32. + + ``packed [N, K//2]`` uint8 (e2m1 pairs), ``scale [N, K//16]`` fp32 (already + converted from fp8-e4m3), ``global_scale [N]`` fp32. Returns ``[N, K]`` fp32. + """ + N, K2 = packed.shape + K = K2 * 2 + lo = (packed & 0x0F).astype(_np.uint8) + hi = (packed >> 4).astype(_np.uint8) + codes = _np.stack([lo, hi], axis=-1).reshape(N, K) # [N, K] + vals = _FP4_TABLE[codes.astype(_np.int64)] # [N, K] + # Per-16 block scale broadcast over K. + bs = _np.repeat(scale, block, axis=-1) # [N, K] + return (vals * bs).astype(_np.float32) * global_scale[:, None].astype(_np.float32) + + +def _pack_codes(codes: _np.ndarray) -> _np.ndarray: + """Pack ``[N, K]`` uint8 4-bit codes -> ``[N, K//2]`` uint8 (low nibble first).""" + N, K = codes.shape + even = codes[..., 0::2] + odd = codes[..., 1::2] + return (even | (odd << 4)).astype(_np.uint8) + + +def convert_nvfp4_to_mxfp4( + packed, + scale, + global_scale, + *, + axis: int = -1, + block: int = 32, +): + """Convert one projection's native NVFP4 expert weights to the MXFP4 layout. + + Args: + packed: ``[..., K//2]`` uint8 e2m1 pairs (low nibble = first code). + scale: ``[..., K//16]`` fp8-e4m3 block scale (fp32/fp16 input accepted). + global_scale: ``[...]`` per-output-row fp16 global scale. + axis: the K (contraction) axis along which blocks are grouped. + + Returns ``(mxfp4_packed [..., K//2] uint8, mxfp4_scales [..., K//32] uint8 e8m0)`` + matching the ``mxfp4_triton`` per-expert bank shape (K innermost). + """ + np = _np + packed = np.asarray(packed) + scale = np.asarray(scale, dtype=np.float32) + global_scale = np.asarray(global_scale, dtype=np.float32) + + if axis not in (-1, packed.ndim - 1): + raise NotImplementedError("converter requires the K axis to be innermost") + + # Dequantize NVFP4 to fp32, move K to the last axis. + K2 = packed.shape[-1] + K = K2 * 2 + codes = np.stack([packed & 0x0F, (packed >> 4)], axis=-1).reshape( + *packed.shape[:-1], K + ) + vals = _FP4_TABLE[codes.astype(np.int64)] + bs = np.repeat(scale, 16, axis=-1) + f32 = (vals * bs).astype(np.float32) * global_scale[..., None].astype(np.float32) + + # Requantize to per-`block` e8m0 + e2m1. + flat = f32.reshape(-1, K) + # pad to a multiple of block for the reshape (K is a multiple of 32 in practice) + n_blocks = K // block + flat_b = flat[:, : n_blocks * block].reshape(-1, block) + scale_codes, mxfp4_codes = e8m0_scale_and_codes(flat_b, block=block) + out_codes = mxfp4_codes.reshape(flat.shape[0], n_blocks * block) + + mxfp4_packed = _pack_codes(out_codes) # [..., K//2] + mxfp4_scales = scale_codes.reshape(*packed.shape[:-1], n_blocks) + return mxfp4_packed, mxfp4_scales + + +# --------------------------------------------------------------------------- +# torch-tensor entrypoint +# --------------------------------------------------------------------------- + + +def _to_numpy(t: object) -> _np.ndarray: + if _np is not None and isinstance(t, _np.ndarray): + return t + try: + return t.detach().cpu().numpy() + except Exception as exc: # pragma: no cover + raise TypeError( + f"convert_nvfp4_to_mxfp4 expects torch tensors or numpy arrays, got {type(t)!r}" + ) from exc + + +def convert_torch_nvfp4_to_mxfp4( + packed, + scale, + global_scale, + *, + block: int = 32, +): + """Torch-tensor variant returning ``(mxfp4_packed, mxfp4_scales)`` torch tensors on + the same device as ``packed`` (host converter: input tensors are pulled to CPU and + the results copied back). Used at load time, cached per model.""" + import torch + + device = packed.device + dtype = packed.dtype + p, s = convert_nvfp4_to_mxfp4( + _to_numpy(packed), _to_numpy(scale), _to_numpy(global_scale), block=block + ) + return torch.from_numpy(p).to(device=device, dtype=dtype), torch.from_numpy(s).to(device=device) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 6ee76406..56264ca5 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -45,6 +45,10 @@ # native GGUF Q4_0 experts: packed block bytes per output row, dequantized inside # the borrowed ggml MoE kernels. gate_up [L*E, 2I, H//32*18], down [L*E, H, I//32*18]. "q4_0": ("gate_up", "down"), + # native GGUF qwen3.5-moe experts: gate_up native Q4_K [L*E, 2I, row_bytes(H,Q4_K)], + # down re-quantized to Q8_0 [L*E, H, row_bytes(I,Q8_0)] (a uniform format the cache + # can hold; the source GGUF down is Q5_K/Q6_K, both >= ... 8-bit >= source precision). + "gguf": ("gate_up", "down"), # native ModelOpt rows for the Triton inline-dequant kernels: packed e2m1 codes + # fp8-e4m3 per-16 block scales + per-output-row fp16 globals (w1/w3 carry distinct # globals, and folding them into the e4m3 block scales would underflow) @@ -83,6 +87,8 @@ "bf16": lambda H, I: 3 * I * H * 2, "fp8_block": lambda H, I: 3 * I * H + ((2 * I // 128) * (H // 128) + (H // 128) * (I // 128)) * 2, "q4_0": lambda H, I: 2 * I * (H // 32) * 18 + H * (I // 32) * 18, + # gate_up Q4_K row_bytes(H, Q4_K)=H//256*144; down Q8_0 row_bytes(I, Q8_0)=I//32*34. + "gguf": lambda H, I: 2 * I * (H // 256) * 144 + H * (I // 32) * 34, "nvfp4": lambda H, I: 2 * I * (H // 2 + H // 16 + 2) + H * (I // 2 + I // 16 + 2), "mxfp4": lambda H, I: 2 * I * (H // 2 + H // 32 + 2) + H * (I // 2 + I // 32 + 2), "ds_fp4": lambda H, I: 2 * I * (H // 2 + H // 32) + H * (I // 2 + I // 32), diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 4954c5f5..2a533be7 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -683,6 +683,25 @@ def _infer_reasoning_parser(model_path: str) -> str | None: kwargs["tp_info"] = DistributedInfo(0, kwargs["tensor_parallel_size"]) del kwargs["tensor_parallel_size"] + # ROCm (AMD) has no NVIDIA-native NVFP4/Marlin path. Reject NVIDIA-only NVFP4 backends + # at parse time with a clean error, and warn for the resident fused MoE backend (the + # offload/cpu/hybrid family is the supported AMD path). + from freetoken.utils.arch import is_rocm + + if is_rocm(): + nvfp4 = kwargs.get("nvfp4_backend") + if nvfp4 in ("marlin", "flashinfer"): + raise SystemExit( + f"--nvfp4-backend {nvfp4} is NVIDIA-only and unavailable on ROCm/AMD; " + f"use --nvfp4-backend triton (inline-dequant) or auto." + ) + if kwargs.get("moe_backend") == "fused": + logger = init_logger(__name__) + logger.warning( + "--moe-backend fused relies on NVIDIA-native fused GEMM; on ROCm/AMD the " + "supported family is offload/hybrid/cpu (the triton/offload path)." + ) + result = ServerArgs(**kwargs) logger = init_logger(__name__) logger.info(f"Parsed arguments:\n{result}") diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py index 2e4ad15f..01f6c57c 100644 --- a/python/freetoken/utils/__init__.py +++ b/python/freetoken/utils/__init__.py @@ -1,5 +1,9 @@ from .arch import ( + current_gpu_name, + device_kind, is_arch_supported, + is_cuda, + is_rocm, is_sm90_family, is_sm90_supported, is_sm100_family, @@ -34,7 +38,11 @@ "load_tokenizer", "load_toolcall_anchor_id", "init_logger", + "device_kind", + "current_gpu_name", "is_arch_supported", + "is_cuda", + "is_rocm", "is_sm90_family", "is_sm90_supported", "is_sm100_family", diff --git a/python/freetoken/utils/arch.py b/python/freetoken/utils/arch.py index 8c1c6c3d..3dd781f4 100644 --- a/python/freetoken/utils/arch.py +++ b/python/freetoken/utils/arch.py @@ -4,14 +4,106 @@ from typing import Tuple +def device_kind() -> str: + """Compute backend the current torch build targets: ``"cuda"`` (NVIDIA), ``"rocm"`` + (AMD/HIP) or ``"cpu"``. Keyed on the *build* (torch.version.hip vs torch.version.cuda), + independent of whether a GPU is present, so feature-gating and graceful-degradation + decisions can be made before any device is available. On ROCm torch, torch.version.cuda + is None and torch.version.hip is set; on CUDA torch the inverse holds.""" + try: + import torch.version + except Exception: + return "cpu" + if getattr(torch.version, "hip", None): + return "rocm" + if getattr(torch.version, "cuda", None): + return "cuda" + return "cpu" + + +def is_rocm() -> bool: + """True when the installed torch is a ROCm (AMD) build.""" + return device_kind() == "rocm" + + +def is_cuda() -> bool: + """True when the installed torch is a CUDA (NVIDIA) build.""" + return device_kind() == "cuda" + + +@functools.cache +def current_gpu_name() -> str | None: + """Device name of the current CUDA-capable device (``torch.cuda.get_device_name``), or + None if torch is unavailable / no device. On ROCm this returns the AMD card name through + the torch.cuda compat layer.""" + try: + import torch + + if not torch.cuda.is_available(): + return None + return torch.cuda.get_device_name(torch.cuda.current_device()) + except Exception: + return None + + @functools.cache def _get_torch_cuda_version() -> Tuple[int, int] | None: - import torch - import torch.version + """Compute capability ``(major, minor)`` of the current CUDA device, or None when it + cannot be determined. Returns None on ROCm torch (gfx archs are not a CUDA compute + capability), when no CUDA device is present, and when torch itself is unavailable -- + so every ``is_sm*``/``is_arch_supported`` gate degrades to the portable path.""" + try: + import torch + import torch.version + + if not torch.cuda.is_available() or not torch.version.cuda: + return None + return torch.cuda.get_device_capability() + except Exception: + return None - if not torch.cuda.is_available() or not torch.version.cuda: + +@functools.cache +def _get_gfx_arch() -> int | None: + """Numeric gfx arch of the current device (e.g. 1100 for ``gfx1100``) on ROCm, or + None when torch is unavailable / not ROCm / no device present. Used by + :func:`is_gfx_arch_ge` for AMD feature gating.""" + try: + import torch + import torch.version + + if not torch.version.hip or not torch.cuda.is_available(): + return None + import re + + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + # ROCm torch exposes the exact gfx string (e.g. ``gfx1100``) as gcnArchName, + # which is the reliable field; the marketing device name is often just + # ``Radeon RX 7900 XTX`` and carries no gfx marker. + gcn = getattr(props, "gcnArchName", None) + if gcn: + m = re.search(r"gfx(\d{3,4})", str(gcn)) + if m: + return int(m.group(1)) + name = torch.cuda.get_device_name(torch.cuda.current_device()) + if name: + m = re.search(r"gfx(\d{3,4})", name) + if m: + return int(m.group(1)) return None - return torch.cuda.get_device_capability() + except Exception: + return None + + +def is_gfx_arch_ge(arch_int: int) -> bool: + """True on ROCm when the current gfx arch number is >= ``arch_int`` (e.g. + ``is_gfx_arch_ge(1100)`` for RDNA 3 / RX 7000). Parses the full gfx string + (``gfx1100`` -> 1100) rather than a CUDA-style ``(major, minor)`` tuple. Returns + False on CUDA and CPU builds, so every gfx gate degrades to the portable path.""" + gfx = _get_gfx_arch() + if gfx is None: + return False + return gfx >= arch_int def is_arch_supported(major: int, minor: int = 0) -> bool: diff --git a/python/freetoken/utils/graph_gate.py b/python/freetoken/utils/graph_gate.py new file mode 100644 index 00000000..b528a208 --- /dev/null +++ b/python/freetoken/utils/graph_gate.py @@ -0,0 +1,189 @@ +"""HIP/CUDA graph-capture parity probe. + +The Inc-1 hard gate: whether ``torch.cuda.graph`` graph capture works on the target +GPU is the single highest-informational-risk assumption for AMD (ROCm) support. +This module probes it once and records a PASS/FAIL + device result that the rest +of the plan (Inc 8) reads. On CUDA it is expected to PASS; on ROCm it may fail on +some consumer cards, in which case Inc 8 must use the kernel-launch decode path. + +The result is cached to disk under the user cache dir so it survives across runs, +and keyed by device kind + device name so a change of GPU invalidates it. +""" +from __future__ import annotations + +import json +import os +from functools import lru_cache + +_CACHE_FILE = "freetoken_graph_gate.json" + + +def _cache_dir() -> str: + base = os.environ.get("XDG_CACHE_HOME") or os.path.join( + os.path.expanduser("~"), ".cache" + ) + path = os.path.join(base, "freetoken") + os.makedirs(path, exist_ok=True) + return path + + +def _cache_path() -> str: + return os.path.join(_cache_dir(), _CACHE_FILE) + + +def _device_kind() -> str: + from freetoken.utils.arch import device_kind + + return device_kind() + + +@lru_cache(maxsize=1) +def _device_name() -> str | None: + """Best-effort current device name via torch.cuda, or None when no device / torch.""" + try: + import torch + + if not torch.cuda.is_available(): + return None + return torch.cuda.get_device_name(torch.cuda.current_device()) + except Exception: + return None + + +def probe_graph_capture() -> dict: + """Run the actual capture probe on the current device. Returns a dict: + ``{"device_kind": ..., "device": ..., "ok": bool, "detail": str}``. + + The GEMM-capture attempt is run in a **fresh subprocess**: on some ROCm builds a + hipBLASLt/capture failure raises an uncatchable fatal HIP error (error 900) that + aborts the whole process, so running it inline would crash the caller (and, worse, + a decode path that attempted graph capture would die). A subprocess lets a fatal + failure surface as a clean ``ok:false`` result instead. + """ + device = _device_name() + try: + import torch + + if not torch.cuda.is_available(): + return { + "device_kind": _device_kind(), + "device": device, + "ok": False, + "detail": "no CUDA-capable device available", + } + except Exception as exc: + detail = next((line.strip() for line in str(exc).splitlines() if line.strip()), "") + return { + "device_kind": _device_kind(), + "device": device, + "ok": False, + "detail": f"torch unavailable: {type(exc).__name__}: {detail}", + } + + # The child probes both an elementwise op (capturable on both backends) and a GEMM + # (hipBLASLt on ROCm), which is what a real decode forward would run. The GEMM is + # the discriminating case: on this ROCm build it fatally aborts -> child exit != 0. + import json as _json + import subprocess as _subprocess + import sys as _sys + + child = _subprocess.run( + [_sys.executable, "-c", _CAPTURE_CHILD], + capture_output=True, + text=True, + timeout=120, + ) + if child.returncode != 0: + detail = next( + (l.strip() for l in child.stderr.splitlines() if l.strip()), + f"graph capture subprocess aborted (rc={child.returncode})", + ) + return { + "device_kind": _device_kind(), + "device": device, + "ok": False, + "detail": f"fatal during capture: {detail[:240]}", + } + try: + data = _json.loads(child.stdout) + except Exception: + return { + "device_kind": _device_kind(), + "device": device, + "ok": False, + "detail": f"unparseable probe output: {child.stdout[:120]}", + } + data.setdefault("device_kind", _device_kind()) + data.setdefault("device", device) + return data + + +#: Child body for the graph-capture probe (see :func:`probe_graph_capture`). Prints a +#: JSON line ``{"ok": true/false, "detail": ...}`` and exits nonzero on a fatal abort. +_CAPTURE_CHILD = r""" +import json, sys +import torch +try: + torch.cuda.synchronize() + s = torch.cuda.Stream() + x = torch.randn(8, 8, device='cuda') + # elementwise (capturable) first, then a GEMM (hipBLASLt on ROCm) + with torch.cuda.stream(s): + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + torch.add(x, x) + s.synchronize() + with torch.cuda.stream(s): + g = torch.cuda.CUDAGraph() + a = torch.randn(64, 64, device='cuda') + with torch.cuda.graph(g): + torch.mm(a, a) + s.synchronize() + torch.cuda.synchronize() + print(json.dumps({"ok": True, "detail": "elementwise+GEMM capture/replay succeeded"})) +except Exception as e: + print(json.dumps({"ok": False, "detail": f"{type(e).__name__}: {str(e)[:160]}"})) +""" + + +def _load_cached() -> dict | None: + try: + with open(_cache_path()) as f: + data = json.load(f) + if ( + data.get("device_kind") == _device_kind() + and data.get("device") == _device_name() + ): + return data + except Exception: + pass + return None + + +def run_graph_gate() -> dict: + """Run (or reuse the cached result for the current device of) the capture probe.""" + cached = _load_cached() + if cached is not None: + return cached + result = probe_graph_capture() + try: + with open(_cache_path(), "w") as f: + json.dump(result, f) + except Exception: + pass + return result + + +@lru_cache(maxsize=1) +def graph_capture_status() -> str: + """Cached graph-capture status: ``"pass"``, ``"fail"``, or ``"unknown"`` (no device / + probe unavailable). Inc 8 reads this to pick HIP-graph vs kernel-launch decode.""" + try: + result = run_graph_gate() + if result["ok"]: + return "pass" + if result.get("device_kind") and result["device_kind"] != "cpu": + return "fail" + return "unknown" + except Exception: + return "unknown" diff --git a/python/freetoken/utils/torch_utils.py b/python/freetoken/utils/torch_utils.py index 9422b9e7..16223210 100644 --- a/python/freetoken/utils/torch_utils.py +++ b/python/freetoken/utils/torch_utils.py @@ -21,6 +21,16 @@ def torch_dtype(dtype: torch.dtype): def nvtx_annotate(name: str, layer_id_field: str | None = None): + from freetoken.utils.arch import is_rocm + + # ROCm torch has no torch.cuda.nvtx; mapping to roctx is future work. Under ROCm we + # pass through (no-op decorator) so AMD runs are not coupled to NVIDIA-only tooling. + if is_rocm(): + def passthrough(fn): + return fn + + return passthrough + import torch.cuda.nvtx as nvtx def decorator(fn): @@ -35,3 +45,12 @@ def wrapper(self, *args, **kwargs): return wrapper return decorator + + +def graph_capture(): + """Context manager that captures a CUDA (or, on ROCm, HIP) graph on the current + stream via ``torch.cuda.graph``. Correctness is validated once by the Inc-1 + ``freetoken.utils.graph_gate`` probe; callers rely on that result.""" + import torch + + return torch.cuda.graph() diff --git a/setup.py b/setup.py index cfe41b7d..c4cf4f1f 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,7 @@ from __future__ import annotations import importlib.util +import os from pathlib import Path from setuptools import setup @@ -10,15 +11,32 @@ ROOT = Path(__file__).parent -def _check_toolchain() -> None: +def _load_toolchain(): path = ROOT / "python" / "freetoken" / "kernel" / "_toolchain.py" spec = importlib.util.spec_from_file_location("_freetoken_toolchain", path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) - module.check_nvcc_matches_torch() + return module + + +def _check_toolchain() -> None: + module = _load_toolchain() + if module.is_rocm_torch(): + module.check_hip_matches_torch() + else: + module.check_nvcc_matches_torch() + + +def _rocm_home() -> Path | None: + for env in ("ROCM_HOME", "HIP_PATH"): + root = os.getenv(env) + if root and (Path(root) / "include").exists(): + return Path(root) + default = Path("/opt/rocm") + return default if (default / "include").exists() else None -def _cuda_runtime_paths() -> tuple[list[str], list[str]]: +def _cuda_runtime_paths() -> tuple[list[str], list[str], list[str]]: if CUDA_HOME is None: raise RuntimeError( "CUDA_HOME is required to build freetoken.kernel._pinned_tensor " @@ -28,12 +46,43 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: library_dirs = [str(cuda_home / "lib64")] if (cuda_home / "lib").exists(): library_dirs.append(str(cuda_home / "lib")) - return [str(cuda_home / "include")], library_dirs + return [str(cuda_home / "include")], library_dirs, ["cudart"] + + +def _rocm_runtime_paths() -> tuple[list[str], list[str], list[str]]: + home = _rocm_home() + if home is None: + raise RuntimeError( + "ROCm torch detected but no ROCm toolkit found. Install ROCm (e.g. /opt/rocm) " + "matching torch's HIP version to build freetoken.kernel._pinned_tensor " + "(it links the HIP runtime API)." + ) + include_dirs = [str(home / "include")] + library_dirs = [] + for sub in ("lib", "lib64"): + if (home / sub).exists(): + library_dirs.append(str(home / sub)) + # HIP host APIs (hipHostMalloc/hipHostRegister/hipHostGetDevicePointer) and HIP + # graph nodes all live in the HIP runtime, amdhip64. (hiprt is a separate optional + # library not present on all ROCm installs; linking it would break the build.) + libraries = ["amdhip64"] + return include_dirs, library_dirs, libraries + + +def _runtime_paths() -> tuple[list[str], list[str], list[str], list[str]]: + """Returns (include_dirs, library_dirs, libraries, compile_defs) for the active backend.""" + module = _load_toolchain() + if module.is_rocm_torch(): + include_dirs, library_dirs, libraries = _rocm_runtime_paths() + return include_dirs, library_dirs, libraries, ["-DUSE_HIP=1", "-DUSE_ROCM=1"] + include_dirs, library_dirs, libraries = _cuda_runtime_paths() + return include_dirs, library_dirs, libraries, [] -cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths() +include_dirs, library_dirs, libraries, compile_defs = _runtime_paths() _check_toolchain() +_extra_compile_args = ["-O3", "-std=c++17", *compile_defs] setup( ext_modules=[ @@ -42,13 +91,14 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: sources=[ "python/freetoken/kernel/csrc/pinned_tensor.cpp", ], - include_dirs=cuda_include_dirs, - library_dirs=cuda_library_dirs, - libraries=["cudart"], - extra_compile_args=["-O3", "-std=c++17"], + include_dirs=include_dirs, + library_dirs=library_dirs, + libraries=libraries, + extra_compile_args=_extra_compile_args, ), - # CPU-compute MoE executor for --moe-backend cpu. Links cudart for the - # cudaLaunchHostFunc submit/sync graph nodes; the bf16 GEMV microkernels + # CPU-compute MoE executor for --moe-backend cpu. On CUDA it links cudart for the + # cudaLaunchHostFunc submit/sync graph nodes; on ROCm those become HIP graph nodes + # (hipLaunchHostFunc) and we link the HIP runtime instead. The bf16 GEMV microkernels # use per-function target attributes (avx512bf16/avx512f) + a runtime # __builtin_cpu_supports dispatch, so the single binary stays portable # (scalar fallback) -- no global -march is set. @@ -57,10 +107,10 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: sources=[ "python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp", ], - include_dirs=cuda_include_dirs, - library_dirs=cuda_library_dirs, - libraries=["cudart"], - extra_compile_args=["-O3", "-std=c++17", "-pthread"], + include_dirs=include_dirs, + library_dirs=library_dirs, + libraries=libraries, + extra_compile_args=_extra_compile_args + ["-pthread"], ), ], cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)}, diff --git a/tests/attention/test_torch_backend.py b/tests/attention/test_torch_backend.py new file mode 100644 index 00000000..f655a959 --- /dev/null +++ b/tests/attention/test_torch_backend.py @@ -0,0 +1,62 @@ +"""The debug ``"torch"`` attention backend (Inc 4 of fix-attention). + +Verifies (a) the backend is registered and selectable, and (b) its pure-PyTorch +GQA attention math (with causal masking and the per-head output gate) matches +``torch.nn.functional.scaled_dot_product_attention`` on a synthetic case with the +qwen35moe full-attention geometry (num_q=16, num_kv=2, head_dim=256, gate). + +This is the ground-truth backend used to A/B the production triton path. +""" + +import torch + + +def _attention_math(q, ks, vs, lk, scale, group, gate): + """Mirror of TorchAttentionBackend's per-request attention (fp32 intermediates).""" + ks = ks.repeat_interleave(group, dim=1).float() + vs = vs.repeat_interleave(group, dim=1).float() + scores = torch.einsum("qhd,khd->hqk", q.float(), ks) * scale + lq = q.shape[0] + rows = torch.arange(lq) + cols = torch.arange(lk) + masked = (cols[None, :] > (lk - lq + rows)[:, None]).to(scores.device) + scores = scores.masked_fill(masked[None], float("-inf")) + probs = torch.softmax(scores, dim=-1) + o = torch.einsum("hqk,khd->qhd", probs, vs) # [lq, num_q, hd] + lq_, nq, hd_ = o.shape + o = o.reshape(lq_, nq * hd_) * torch.sigmoid(gate.float()) # [lq, num_q*hd] + return o.reshape(lq_, nq, hd_) + + +def test_torch_backend_registered(): + from freetoken.attention import SUPPORTED_ATTENTION_BACKENDS, attention_backend_info, AttnType + + assert "torch" in SUPPORTED_ATTENTION_BACKENDS.supported_names() + info = attention_backend_info("torch") + assert AttnType.FULL in info.supported_types + assert info.hybrid_linear_ok # must coexist with GDN layers + + +def test_torch_attention_matches_sdpa_prefill(): + num_q, num_kv, hd, group = 16, 2, 256, 8 + T = 4 + scale = hd ** -0.5 + q = torch.randn(T, num_q, hd) + k = torch.randn(T, num_kv, hd) + v = torch.randn(T, num_kv, hd) + gate = torch.randn(T, num_q * hd) + + out = _attention_math(q, k, v, T, scale, group, gate) + # reference: sdpa on GQA-expanded heads, then gate + no o_proj here + ke = k.repeat_interleave(group, dim=1) + ve = v.repeat_interleave(group, dim=1) + ref = torch.nn.functional.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0), + ke.transpose(0, 1).unsqueeze(0), + ve.transpose(0, 1).unsqueeze(0), + is_causal=True, + )[0].transpose(0, 1) + ref = ref.reshape(T, num_q * hd) * torch.sigmoid(gate) + ref = ref.reshape(T, num_q, hd) + + assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3) diff --git a/tests/engine/test_attention_backend_rocm.py b/tests/engine/test_attention_backend_rocm.py new file mode 100644 index 00000000..d8ff4e5c --- /dev/null +++ b/tests/engine/test_attention_backend_rocm.py @@ -0,0 +1,75 @@ +"""ROCm (AMD) attention-backend resolution: no NVIDIA-native backend may be selected. + +On ROCm, is_sm90/100 gates are False and flashinfer/sgl_kernel are treated as +unavailable, so auto resolution must fall through to the portable Triton backend +for FULL-attention models. +""" + +import pytest + + +def _engine_config(**overrides): + from types import SimpleNamespace + + import torch + + from freetoken.distributed import DistributedInfo + from freetoken.engine.config import EngineConfig + + config = EngineConfig( + model_path="/tmp/freetoken-test-model", + tp_info=DistributedInfo(rank=0, size=1), + dtype=torch.bfloat16, + **overrides, + ) + object.__setattr__( + config, + "model_config", + SimpleNamespace( + has_swa_attention=False, + has_linear_attention=False, + is_moe=False, + num_layers=10, + expert_quant="none", + ), + ) + return config + + +def _patch_rocm(monkeypatch): + from freetoken.engine import engine + from freetoken.kernel import backend + + monkeypatch.setattr(engine, "is_sm100_family", lambda: False) + monkeypatch.setattr(engine, "is_sm90_family", lambda: False) + monkeypatch.setattr(engine, "_flashinfer_available", lambda: False) + monkeypatch.setattr(engine, "_sgl_flash_attn_available", lambda: False) + monkeypatch.setattr(backend, "is_rocm", lambda: True) + monkeypatch.setattr(engine, "is_rocm", lambda: True) + + +def test_rocm_auto_resolves_to_triton(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_rocm(monkeypatch) + config = _engine_config(attention_backend="auto") + _adjust_config(config) + assert config.attention_backend == "triton" + + +def test_rocm_explicit_nvidia_backend_rejected(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_rocm(monkeypatch) + # flashinfer/sgl absent and arch gates false -> trtllm/fa/fi requirements unmet. + for backend in ("fi", "fa", "trtllm"): + config = _engine_config(attention_backend=backend) + with pytest.raises(RuntimeError): + _adjust_config(config) + + +def test_sgl_flash_attn_unavailable_on_rocm(monkeypatch): + from freetoken.engine import engine + + monkeypatch.setattr(engine, "is_rocm", lambda: True) + assert engine._sgl_flash_attn_available() is False diff --git a/tests/kernels/test_backend_rocm.py b/tests/kernels/test_backend_rocm.py new file mode 100644 index 00000000..4c732635 --- /dev/null +++ b/tests/kernels/test_backend_rocm.py @@ -0,0 +1,63 @@ +"""ROCm-aware optional-package probing in ``kernel/backend.py``. + +Loads ``backend.py`` with a fake ``freetoken.utils.arch`` injected so the ROCm gating +(treat NVIDIA-only native packages as unavailable) is unit-testable without a working +torch or those packages installed. +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +_BACKEND_PATH = ( + Path(__file__).resolve().parents[2] / "python" / "freetoken" / "kernel" / "backend.py" +) + + +def _load_backend(rocm: bool): + # Fake freetoken.utils.arch (backend.py imports is_rocm from it at module load). + arch_mod = types.ModuleType("freetoken.utils.arch") + arch_mod.is_rocm = lambda: rocm + arch_mod.device_kind = lambda: "rocm" if rocm else "cpu" + freetoken = types.ModuleType("freetoken") + freetoken.__path__ = [] + utils = types.ModuleType("freetoken.utils") + utils.__path__ = [] + utils.arch = arch_mod + sys.modules["freetoken"] = freetoken + sys.modules["freetoken.utils"] = utils + sys.modules["freetoken.utils.arch"] = arch_mod + + spec = importlib.util.spec_from_file_location("freetoken.kernel.backend", _BACKEND_PATH) + kernel = types.ModuleType("freetoken.kernel") + kernel.__path__ = [str(_BACKEND_PATH.parent)] + sys.modules["freetoken.kernel"] = kernel + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(autouse=True) +def _clean_sys_modules(): + yield + for name in ("freetoken", "freetoken.utils", "freetoken.utils.arch", "freetoken.kernel"): + sys.modules.pop(name, None) + + +def test_rocm_native_packages_unavailable(): + backend = _load_backend(rocm=True) + assert backend.is_flashinfer_installed() is False + assert backend.is_sgl_kernel_installed() is False + assert backend.is_triton_kernels_installed() is False + assert backend.is_native_cuda_available() is False + + +def test_cuda_native_packages_follow_importability(monkeypatch): + backend = _load_backend(rocm=False) + # Without the packages installed, probes are False; is_native_cuda_available() needs a + # real CUDA-capable torch (returns False here since torch.cuda is unavailable). + assert backend.is_flashinfer_installed() is False + assert backend.is_sgl_kernel_installed() is False diff --git a/tests/kernels/test_cache_rocm_pairing.py b/tests/kernels/test_cache_rocm_pairing.py new file mode 100644 index 00000000..24796632 --- /dev/null +++ b/tests/kernels/test_cache_rocm_pairing.py @@ -0,0 +1,109 @@ +"""ROCm cache/runtime version pairing and gfx-arch gating (torch-free). + +Loads the relevant modules by file path with mocked torch so the logic is unit-testable +without a working torch/GPU install. +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +_ROOT = Path(__file__).resolve().parents[2] +_UTILS_PATH = _ROOT / "python" / "freetoken" / "kernel" / "utils.py" +_ARCH_PATH = _ROOT / "python" / "freetoken" / "utils" / "arch.py" + + +@pytest.fixture(scope="module") +def ku(): + spec = importlib.util.spec_from_file_location("kernel_utils_mod", _UTILS_PATH) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + return m + + +@pytest.fixture(scope="module") +def arch(): + spec = importlib.util.spec_from_file_location("arch_mod", _ARCH_PATH) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + return m + + +def test_cache_version_pairs_rocm(ku): + # Same release + same g-sha + both rocm -> OK. + assert ku._kernel_cache_version_ok( + "0.1.1+rocm.g3f01615", "0.1.1+rocm.g3f01615" + ) + # No g-stamp on either side -> release-only compare -> OK. + assert ku._kernel_cache_version_ok("0.1.1+rocm", "0.1.1+rocm") + + +def test_cache_version_rejects_backend_tag_mismatch(ku): + # CUDA runtime must never pair with a ROCm cache (SASS family differs). + assert not ku._kernel_cache_version_ok( + "0.1.1+rocm.g3f01615", "0.1.1+cu130.g3f01615" + ) + assert not ku._kernel_cache_version_ok( + "0.1.1+cu130.g3f01615", "0.1.1+rocm.g3f01615" + ) + + +def test_cache_version_rejects_different_build(ku): + assert not ku._kernel_cache_version_ok( + "0.1.1+rocm.g3f01615", "0.1.1+rocm.gdeadbee" + ) + assert not ku._kernel_cache_version_ok("0.1.2+rocm", "0.1.1+rocm") + + +def _torch_importable() -> bool: + """True only when torch actually *imports* (a findable-but-broken torch -- e.g. a + CUDA build missing its runtime libs -- must count as absent so the torch-free path + is exercised on such boxes).""" + try: + import torch # noqa: PLC0415 + + return True + except Exception: + return False + + +def test_gfx_arch_ge_degrades_without_torch(arch): + if _torch_importable(): + pytest.skip("torch is importable; no-torch degradation tested on a torch-free box") + assert arch.is_gfx_arch_ge(1100) is False + + +def test_gfx_arch_ge_parses_gfx_string(arch, monkeypatch): + if _torch_importable(): + # Real torch present: is_gfx_arch_ge must reflect the actual device (gfx1100 on + # an RX 7900 XTX). No fake torch injection to avoid cross-test state pollution. + import torch + + assert arch.is_gfx_arch_ge(1100) is True # RX 7000 target + return + arch._get_gfx_arch.cache_clear() + # Fake a ROCm torch whose device name carries the gfx arch string. + tv = types.ModuleType("torch.version") + tv.hip = "6.2.4100000" + tv.cuda = None + t = types.ModuleType("torch") + t.version = tv + _props = types.SimpleNamespace(gcnArchName="gfx1100") + t.cuda = types.SimpleNamespace( + is_available=lambda: True, + current_device=lambda: 0, + get_device_name=lambda dev: "AMD Radeon RX 7900 XTX", + get_device_properties=lambda dev: _props, + ) + monkeypatch.setitem(sys.modules, "torch", t) + monkeypatch.setitem(sys.modules, "torch.version", tv) + assert arch.is_gfx_arch_ge(1100) is True + assert arch.is_gfx_arch_ge(1103) is False + # CUDA builds must return False. + arch._get_gfx_arch.cache_clear() + tv.hip = None + tv.cuda = "13.0" + assert arch.is_gfx_arch_ge(1100) is False diff --git a/tests/kernels/test_toolchain_hip.py b/tests/kernels/test_toolchain_hip.py new file mode 100644 index 00000000..fa8f53cf --- /dev/null +++ b/tests/kernels/test_toolchain_hip.py @@ -0,0 +1,84 @@ +"""HIP (ROCm) toolchain helper tests. + +These load ``kernel/_toolchain.py`` by file path (as setup.py and the kernel-cache +build backend do) so they need no torch import and no ROCm toolkit to exercise the +pure parsing/detection logic. +""" + +import importlib.util +import os +from pathlib import Path + +import pytest + +_TOOLCHAIN_PATH = ( + Path(__file__).resolve().parents[2] / "python" / "freetoken" / "kernel" / "_toolchain.py" +) + + +def _load_toolchain(): + spec = importlib.util.spec_from_file_location("_freetoken_toolchain", _TOOLCHAIN_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def tc(): + return _load_toolchain() + + +def _write_fake_hipcc(tmp_path, version: str): + bin_dir = tmp_path / "bin" + bin_dir.mkdir(exist_ok=True) + hipcc = bin_dir / "hipcc" + hipcc.write_text( + "#!/bin/sh\n" + f'echo "HIP version: {version}"\n', + encoding="utf-8", + ) + hipcc.chmod(0o755) + return str(hipcc) + + +def test_hip_hip_version(tc, tmp_path): + hipcc = _write_fake_hipcc(tmp_path, "6.2.41000") + assert tc.hip_hip_version(hipcc) == (6, 2) + + +def test_hip_hip_version_missing(tc): + assert tc.hip_hip_version("/nonexistent/hipcc") is None + + +def test_hipcc_path_rocm_home(tc, tmp_path, monkeypatch): + hipcc = _write_fake_hipcc(tmp_path, "6.2.41000") + monkeypatch.setenv("ROCM_HOME", str(tmp_path)) + monkeypatch.delenv("HIP_PATH", raising=False) + assert tc._hipcc_path() == hipcc + + +def test_hipcc_path_default_opt_rocm(tc, monkeypatch): + monkeypatch.delenv("ROCM_HOME", raising=False) + monkeypatch.delenv("HIP_PATH", raising=False) + # If a real /opt/rocm/bin/hipcc exists it wins; otherwise we expect None (no PATH hit + # guaranteed in the test sandbox, so only assert it returns None or a path string). + result = tc._hipcc_path() + if os.path.exists("/opt/rocm/bin/hipcc"): + assert result is not None + else: + assert result is None or result.endswith("hipcc") + + +def test_torch_hip_major_rocm(tc, monkeypatch): + # Simulate ROCm torch via a fake torch.version with hip set. + import sys + import types + + fake_torch = types.ModuleType("torch") + fake_version = types.ModuleType("torch.version") + fake_version.hip = "6.2.4100000" + fake_version.cuda = None + fake_torch.version = fake_version + monkeypatch.setitem(sys.modules, "torch", fake_torch) + assert tc.is_rocm_torch() is True + assert tc.torch_hip_major() == 6 diff --git a/tests/models/test_qwen35moe_gguf_deint.py b/tests/models/test_qwen35moe_gguf_deint.py new file mode 100644 index 00000000..4154fe27 --- /dev/null +++ b/tests/models/test_qwen35moe_gguf_deint.py @@ -0,0 +1,71 @@ +"""qwen35moe GGUF GDN value-head de-interleaving. + +llama.cpp stores the GDN *value* projections with the ``mrope_interleaved`` head order +(even heads 0..nv/2-1 first, then odd heads nv/2..nv-1). FreeToken uses the HF contiguous +head order, so the loader de-interleaves the value projections on load. These tests pin the +permutation and the packed/dense de-interleave helpers (no model weights required). +""" + +import torch + +from freetoken.models.qwen3_5_moe.gguf import ( + _gdn_head_perm, + _deint_dense_rows, + _deint_q8_cols, + _deint_q8_rows, +) + + +def _interleave(x: torch.Tensor, nv: int, rows_per_head: int = 1) -> torch.Tensor: + """Build the GGUF head-interleaved layout from an HF-contiguous tensor. + + GGUF position ``perm[h]`` holds HF head ``h`` (so interleaved[perm[h]] = x[h]). + """ + m = x.reshape(nv, rows_per_head, -1) + perm = _gdn_head_perm(nv) + out = m.clone() + for h in range(nv): + out[perm[h]] = m[h] + return out.reshape(x.shape) + + +def test_head_permutation_is_a_bijection(): + nv = 32 + perm = _gdn_head_perm(nv) + assert sorted(perm) == list(range(nv)) # valid permutation + # GGUF layout: even HF heads occupy positions 0..15, odd heads 16..31. + even = sorted(perm[h] for h in range(0, nv, 2)) + odd = sorted(perm[h] for h in range(1, nv, 2)) + assert even == list(range(nv // 2)) + assert odd == list(range(nv // 2, nv)) + + +def test_deint_dense_rows_recovers_contiguous(): + nv, hd = 32, 128 + x = torch.randn(nv, hd) + inter = _interleave(x, nv, rows_per_head=1) + rec = _deint_dense_rows(inter, nv) + assert torch.allclose(rec, x, atol=1e-6) + + +def test_deint_q8_rows_recovers_contiguous(): + nv, hd = 32, 128 + rph = 64 # rows per value head in the projection output dim + x = torch.randn(nv * rph, 4) # packed rows x (row_bytes mocked) + inter = _interleave(x, nv, rows_per_head=rph) + rec = _deint_q8_rows(inter, nv, rph) + assert torch.allclose(rec, x, atol=1e-6) + + +def test_deint_q8_cols_recovers_contiguous(): + nv = 32 + rows, blocks_per_head, bb = 16, 4, 34 + x = torch.randn(rows, nv * blocks_per_head * bb) + # interleave column head-groups: inter[:, perm[h]*bbp:(perm[h]+1)*bbp] = x[:, h*bbp:(h+1)*bbp] + bbp = blocks_per_head * bb + inter = torch.empty_like(x) + perm = _gdn_head_perm(nv) + for h in range(nv): + inter[:, perm[h] * bbp:(perm[h] + 1) * bbp] = x[:, h * bbp:(h + 1) * bbp] + rec = _deint_q8_cols(inter, nv, blocks_per_head) + assert torch.allclose(rec, x, atol=1e-6) diff --git a/tests/moe/test_nvfp4_backends_rocm.py b/tests/moe/test_nvfp4_backends_rocm.py new file mode 100644 index 00000000..0061d378 --- /dev/null +++ b/tests/moe/test_nvfp4_backends_rocm.py @@ -0,0 +1,95 @@ +"""ROCm-aware NVFP4 backend selection (torch-free). + +Loads ``nvfp4_backends.py`` with a fake ``freetoken.utils.arch`` so the AMD branch of +``select_nvfp4_backend`` (reject NVIDIA-only marlin/flashinfer, force triton) is +unit-testable without torch/vLLM/flashinfer. +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +_BACKEND_PATH = ( + Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "moe" / "nvfp4_backends.py" +) + + +def _torch_importable() -> bool: + """True only when torch actually *imports* (a findable-but-broken torch -- e.g. a + CUDA build missing its runtime libs -- must count as absent so the torch-free path + is exercised on such boxes).""" + try: + import torch # noqa: PLC0415 + + return True + except Exception: + return False + + +def _load_nvfp4(rocm: bool): + arch = types.ModuleType("freetoken.utils.arch") + arch.is_rocm = lambda: rocm + freetoken = types.ModuleType("freetoken") + freetoken.__path__ = [] + utils = types.ModuleType("freetoken.utils") + utils.__path__ = [] + utils.arch = arch + utils.init_logger = lambda name: types.SimpleNamespace( + info=lambda *a, **k: None, + warning=lambda *a, **k: None, + debug=lambda *a, **k: None, + ) + sys.modules["freetoken"] = freetoken + sys.modules["freetoken.utils"] = utils + sys.modules["freetoken.utils.arch"] = arch + + # nvfp4_backends.py does `import torch` at module scope. On a torch-free box stub it + # so the module imports; on a box with real torch use the real one (never mutate + # sys.modules["torch"], which would corrupt the session's torch/triton state). + if not _torch_importable(): + torch_stub = types.ModuleType("torch") + torch_stub.no_grad = lambda: (lambda f: f) + torch_stub.Tensor = object + sys.modules["torch"] = torch_stub + + spec = importlib.util.spec_from_file_location("nvfp4_backends_mod", _BACKEND_PATH) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + return m + + +@pytest.fixture(autouse=True) +def _clean(): + yield + for name in ("freetoken", "freetoken.utils", "freetoken.utils.arch"): + sys.modules.pop(name, None) + + +def test_rocm_auto_resolves_triton(): + m = _load_nvfp4(rocm=True) + dev = types.SimpleNamespace(type="cuda") + assert m.select_nvfp4_backend(dev, 768, "auto") == "triton" + assert m.select_nvfp4_backend(dev, 768, "triton") == "triton" + + +def test_rocm_rejects_nvidia_backends(): + m = _load_nvfp4(rocm=True) + dev = types.SimpleNamespace(type="cuda") + with pytest.raises(RuntimeError, match="NVIDIA-only"): + m.select_nvfp4_backend(dev, 768, "marlin") + with pytest.raises(RuntimeError, match="NVIDIA-only"): + m.select_nvfp4_backend(dev, 768, "flashinfer") + + +def test_cuda_path_not_short_circuited_by_rocm_guard(): + # On CUDA the ROCm early-return must not be taken. We only exercise the guard + # boundary (is_rocm()==False) without calling torch.cuda (unavailable here); the + # full CUDA auto logic runs on the target box. + m = _load_nvfp4(rocm=False) + assert m.select_nvfp4_backend( + types.SimpleNamespace(type="cpu"), 768, "auto" + ) == "triton" diff --git a/tests/moe/test_nvfp4_to_mxfp4.py b/tests/moe/test_nvfp4_to_mxfp4.py new file mode 100644 index 00000000..19420118 --- /dev/null +++ b/tests/moe/test_nvfp4_to_mxfp4.py @@ -0,0 +1,111 @@ +"""NVFP4 -> MXFP4 converter tests (numpy core, torch-free). + +Validates the error-prone numeric bits: the e2m1 code table, e8m0 scale encoding, and +an end-to-end round trip (dequant NVFP4 -> requant MXFP4 -> dequant back) that must be +close to the reference under a bounded relative error. +""" + +import importlib.util +import math +from pathlib import Path + +import numpy as np +import pytest + +_MODULE_PATH = ( + Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "moe" / "nvfp4_to_mxfp4.py" +) + + +@pytest.fixture(scope="module") +def c(): + spec = importlib.util.spec_from_file_location("nvfp4_to_mxfp4", _MODULE_PATH) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + return m + + +def test_fp4_table_values(c): + table = c.fp4_e2m1_table() + assert len(table) == 16 + # e2m1 canonical values. + assert table[0] == 0.0 + assert table[1] == 0.5 + assert table[4] == 2.0 + assert table[7] == 6.0 + assert table[15] == -6.0 + assert table[9] == -0.5 + + +def test_e8m0_scale_exact_power_of_two(c): + # Block max-abs 3.0 -> scale 2^ceil(log2(3/6)) = 2^-1 = 0.5 -> code 126 (126-127=-1). + values = np.array([[3.0, 1.0, -2.0, 0.5] + [0.0] * 28], dtype=np.float32) + scale_codes, codes = c.e8m0_scale_and_codes(values, block=32) + assert scale_codes.shape == (1,) + assert int(scale_codes[0]) == 126 + # After dividing by 0.5: [6, 2, -4, 1] -> nearest e2m1 codes 7, 4, 14, 2. + assert int(codes[0, 0]) == 7 + assert int(codes[0, 1]) == 4 + assert int(codes[0, 2]) == 14 + assert int(codes[0, 3]) == 2 + + +def test_dequantize_nvfp4_round_trip(c): + rng = np.random.default_rng(0) + N, K = 4, 64 + # Build a plausible NVFP4 weight: choose fp4 codes, block scales, row globals. + codes = rng.integers(0, 16, size=(N, K)).astype(np.uint8) + packed = (codes[:, 0::2] | (codes[:, 1::2] << 4)).astype(np.uint8) + block_scale = (rng.uniform(0.1, 2.0, size=(N, K // 16))).astype(np.float32) + global_scale = (rng.uniform(0.5, 2.0, size=(N,)).astype(np.float32)) + ref = c.dequantize_nvfp4_block(packed, block_scale, global_scale) + assert ref.shape == (N, K) + assert np.isfinite(ref).all() + + +def test_round_trip_matches_reference(c): + rng = np.random.default_rng(1) + N, K = 2, 128 + codes = rng.integers(0, 16, size=(N, K)).astype(np.uint8) + packed = (codes[:, 0::2] | (codes[:, 1::2] << 4)).astype(np.uint8) + block_scale = rng.uniform(0.5, 2.0, size=(N, K // 16)).astype(np.float32) + global_scale = rng.uniform(0.5, 2.0, size=(N,)).astype(np.float32) + ref = c.dequantize_nvfp4_block(packed, block_scale, global_scale) + + mxfp4_packed, mxfp4_scales = c.convert_nvfp4_to_mxfp4( + packed, block_scale, global_scale + ) + assert mxfp4_packed.shape == (N, K // 2) + assert mxfp4_scales.shape == (N, K // 32) + assert mxfp4_scales.dtype == np.uint8 + + # Re-dequantize the MXFP4 output and compare to the NVFP4 reference. + out_codes = np.stack( + [mxfp4_packed & 0x0F, mxfp4_packed >> 4], axis=-1 + ).reshape(N, K) + vals = np.asarray(c.fp4_e2m1_table(), dtype=np.float32)[out_codes.astype(np.int64)] + # e8m0 scale 2**(v-127) + scale_v = (2.0 ** (mxfp4_scales.astype(np.float32) - 127.0))[:, :, None] + mxfp4_vals = (vals.reshape(N, K // 32, 32) * scale_v).reshape(N, K) + # MXFP4 requantizes at a coarser 32-block granularity, so the round trip carries + # fp4 quantization error (inherent to e2m1). Bound it loosely: median < 0.6 and the + # largest magnitude in each block (which defines the scale) must be well-preserved. + rel = np.abs(mxfp4_vals - ref) / (np.abs(ref) + 1e-6) + # Exclude near-zero refs where relative error is meaningless (a code flips +0<->small). + nz = np.abs(ref) > 1e-3 + assert float(np.median(rel[nz])) < 0.6 + assert float(np.percentile(rel[nz], 95)) < 2.0 + # The max-abs element per 32-block reproduces within one e2m1 step of its scale. + block_max_ref = np.abs(ref.reshape(N, K // 32, 32)).max(axis=-1) + block_max_ours = np.abs(mxfp4_vals.reshape(N, K // 32, 32)).max(axis=-1) + scale_ratio = np.abs(block_max_ours / (block_max_ref + 1e-6) - 1.0) + assert float(np.median(scale_ratio)) < 0.5 + + +def test_inner_axis_unsupported(c): + packed = np.zeros((64, 32), dtype=np.uint8) + scale = np.zeros((64, 4), dtype=np.float32) + g = np.ones((64,), dtype=np.float32) + with pytest.raises(NotImplementedError): + c.convert_nvfp4_to_mxfp4(packed.T, scale, g, axis=0) diff --git a/tests/utils/test_device_kind.py b/tests/utils/test_device_kind.py new file mode 100644 index 00000000..0496c4d1 --- /dev/null +++ b/tests/utils/test_device_kind.py @@ -0,0 +1,83 @@ +"""Device-kind (CUDA vs ROCm vs CPU) detection and graceful-degradation gating. + +These run without a working torch install: we load ``utils/arch.py`` by file path and +inject a fake ``torch``/``torch.version`` into ``sys.modules`` so the build-detection +(``torch.version.hip`` vs ``torch.version.cuda``) is unit-testable anywhere. +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +_ARCH_PATH = ( + Path(__file__).resolve().parents[2] / "python" / "freetoken" / "utils" / "arch.py" +) + + +@pytest.fixture(scope="module") +def arch(): + spec = importlib.util.spec_from_file_location("arch_mod", _ARCH_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _set_torch(arch, monkeypatch, hip, cuda): + tv = types.ModuleType("torch.version") + tv.hip = hip + tv.cuda = cuda + t = types.ModuleType("torch") + t.version = tv + t.cuda = types.SimpleNamespace(is_available=lambda: False) + monkeypatch.setitem(sys.modules, "torch", t) + monkeypatch.setitem(sys.modules, "torch.version", tv) + + +def test_device_kind_rocm(arch, monkeypatch): + _set_torch(arch, monkeypatch, hip="6.2.4100000", cuda=None) + assert arch.device_kind() == "rocm" + assert arch.is_rocm() is True + assert arch.is_cuda() is False + + +def test_device_kind_cuda(arch, monkeypatch): + _set_torch(arch, monkeypatch, hip=None, cuda="13.0") + assert arch.device_kind() == "cuda" + assert arch.is_rocm() is False + assert arch.is_cuda() is True + + +def test_device_kind_cpu(arch, monkeypatch): + _set_torch(arch, monkeypatch, hip=None, cuda=None) + assert arch.device_kind() == "cpu" + + +def _torch_importable() -> bool: + """True only when torch actually *imports* (a findable-but-broken torch -- e.g. a + CUDA build missing its runtime libs -- must count as absent so the torch-free path + is exercised on such boxes).""" + try: + import torch # noqa: PLC0415 + + return True + except Exception: + return False + + +def test_arch_gates_degrade_without_torch(arch, monkeypatch): + # Only meaningful where torch is genuinely absent (the CUDA-locked dev box). On a + # box with a working torch, Python re-imports the real torch after the delitem, so + # the assertion target changes -- skip there. + if _torch_importable(): + pytest.skip("torch is importable; no-torch degradation tested on a torch-free box") + # With torch absent, every is_sm* gate is False and device_kind() is "cpu". + monkeypatch.delitem(sys.modules, "torch", raising=False) + monkeypatch.delitem(sys.modules, "torch.version", raising=False) + assert arch.is_sm90_supported() is False + assert arch.is_sm100_supported() is False + assert arch.is_sm90_family() is False + assert arch.is_sm100_family() is False + assert arch.device_kind() == "cpu" From 6f4974eed707ca47a7f017dce7ad4fb507f5a4c9 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Tue, 25 Aug 2026 22:34:03 -0300 Subject: [PATCH 2/6] fix(gguf): register unmerged CONTROL/USER_DEFINED tokens for atomic encoding The GGUF->fast-tokenizer converter relies on BPE merges to keep special strings whole: -style tokens happen to be merge entries and survive, but think tags / tool-call tags are NOT in the merge table, so they silently split into plain pieces (''). The model then receives garbage ids it never saw in training and answers with gibberish + EOS -- observed as 'reasons a little, returns empty content' on Qwen3.6 GGUF checkpoints. _register_control_tokens re-registers every ggml token_type 2/3/4 string that does not already encode atomically, against its EXISTING vocab id: vocab size and id assignments never change, already-atomic tokens are untouched (no-op on healthy setups), and safetensors/HF checkpoints never touch this path. Hardware-independent; safe for CUDA by construction. Also: chat-template resolution now mirrors official HF configs -- chat_template.jinja sidecar next to the .gguf wins, then FT_CHAT_TEMPLATE_REPO (hf_hub_download), then the embedded metadata. CPU-only regression tests included (synthetic BPE tokenizer, no GPU or download needed). --- python/freetoken/models/gguf/tokenizer.py | 77 +++++++++++++++++++- tests/models/test_gguf_tokenizer_specials.py | 77 ++++++++++++++++++++ 2 files changed, 153 insertions(+), 1 deletion(-) create mode 100644 tests/models/test_gguf_tokenizer_specials.py diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index c0582419..83d8be84 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -10,8 +10,12 @@ from typing import Any +from freetoken.utils import init_logger + from .reader import gguf_architecture, load_gguf_metadata +logger = init_logger(__name__) + # GGUF architecture -> transformers GGUF tokenizer-converter key. _TOKENIZER_ARCH = { "gemma4": "gemma4_text", @@ -19,6 +23,73 @@ } +# ggml token_type enum values that mark non-mergeable control strings. +# 1=NORMAL 2=UNKNOWN 3=CONTROL 4=USER_DEFINED 5=UNUSED 6=BYTE +_GGML_SPECIAL_TYPES = (2, 3, 4) + + +def _register_control_tokens(tokenizer, tokens: list[str], types: list[int]) -> None: + """Force every CONTROL/USER_DEFINED/UNKNOWN vocab entry to tokenize atomically. + + The GGUF->fast-tokenizer converter relies on BPE merges for special strings: + ``<|im_start|>`` happens to be a merge-table entry and survives, but ```` + is NOT (it never appears in merged training text), so it silently splits into + plain pieces (''). A model fed those garbage ids answers with + gibberish and EOS -- observed as "thought then returned empty" on Qwen3.6 GGUFs. + Registering each string as an added token makes the AddedVocabulary extract it + before BPE; because the string already exists in the base vocab, the existing id + is reused and the vocab never grows (transformers' own ``add_tokens`` wrapper + no-ops for in-vocab strings, so this goes through the backend directly). + """ + missing = [ + name + for i, (name, ty) in enumerate(zip(tokens, types)) + if int(ty) in _GGML_SPECIAL_TYPES + # Skip BYTE-fallback and unused; skip anything already atomic at its id. + and tokenizer.encode(name, add_special_tokens=False) != [i] + ] + if missing: + tokenizer.backend_tokenizer.add_tokens(missing) + logger.info( + "registered %d unmerged control tokens for atomic encoding (e.g. %s)", + len(missing), + ", ".join(repr(t) for t in missing[:6]), + ) + + +def _resolve_chat_template(meta: dict[str, Any], model_path: str) -> str | None: + """Chat template for GGUF checkpoints: explicit mirrors win, then metadata. + + Priority: a ``chat_template.jinja`` dropped NEXT TO the .gguf file, then + ``FT_CHAT_TEMPLATE_REPO`` ( on the HF Hub, fetched via huggingface_hub), + then the template embedded in the GGUF's ``tokenizer.chat_template`` metadata. + GGUF packagers (llama.cpp/unsloth) sometimes ship modified variants of the + official template — placing the official file beside the checkpoint overrides + it without repacking. The embedded one is the last resort, never wrong-by-default. + """ + import os + + sidecar = os.path.join(os.path.dirname(model_path), "chat_template.jinja") + if os.path.isfile(sidecar): + logger.info("using chat template sidecar %s", sidecar) + with open(sidecar, encoding="utf-8") as fh: + return fh.read() + repo = os.environ.get("FT_CHAT_TEMPLATE_REPO") + if repo: + try: + from huggingface_hub import hf_hub_download + + path = hf_hub_download(repo_id=repo, filename="chat_template.jinja") + with open(path, encoding="utf-8") as fh: + return fh.read() + except Exception as exc: # noqa: BLE001 — offline/bad repo is not fatal + logger.warning("FT_CHAT_TEMPLATE_REPO=%s fetch failed: %s", repo, exc) + embedded = meta.get("tokenizer.chat_template") + if isinstance(embedded, str) and embedded.strip(): + return embedded + return None + + def load_gguf_tokenizer(model_path: str): from transformers import PreTrainedTokenizerFast from transformers.integrations.ggml import convert_gguf_tokenizer @@ -49,7 +120,11 @@ def tok_for(id_key: str, default: str) -> str: unk_token=tok_for("unknown_token_id", ""), pad_token=tok_for("padding_token_id", ""), ) - chat_template = meta.get("tokenizer.chat_template") + # GGUFs are not required to carry per-token types; absent means all-normal. + types = meta.get("tokenizer.ggml.token_type") or [] + if types: + _register_control_tokens(tokenizer, tokens, types) + chat_template = _resolve_chat_template(meta, str(model_path)) if chat_template: tokenizer.chat_template = chat_template return tokenizer diff --git a/tests/models/test_gguf_tokenizer_specials.py b/tests/models/test_gguf_tokenizer_specials.py new file mode 100644 index 00000000..deadb0a4 --- /dev/null +++ b/tests/models/test_gguf_tokenizer_specials.py @@ -0,0 +1,77 @@ +"""CPU-only regression tests for GGUF control-token registration. + +The GGUF->fast-tokenizer converter leaves CONTROL/USER_DEFINED vocab entries +(e.g. Qwen's think tags, tool-call tags) unregistered: unless they happen to be +reachable through BPE merges they silently split into plain pieces, feeding the +model garbage ids ("thought then empty" symptom). ``_register_control_tokens`` +must re-register every such string against its EXISTING id -- vocab size and id +assignments may not change, and already-atomic tokens must be left alone. + +Runs on any machine: no GPU, no model download -- builds synthetic BPE +tokenizers whose vocab mirrors the broken shape (control string in vocab, +unreachable because no merge path leads to it). +""" + +from __future__ import annotations + +from tokenizers import Tokenizer, models, pre_tokenizers + +from freetoken.models.gguf.tokenizer import _register_control_tokens +from transformers import PreTrainedTokenizerFast + + +def _tok(vocab: dict[str, int], merges: list[tuple[str, str]] = []): + backend = Tokenizer(models.BPE(vocab=vocab, merges=list(merges))) + return PreTrainedTokenizerFast(tokenizer_object=backend) + + +def test_unmergeable_control_token_registered_at_existing_id(): + # Mirrors the real Qwen GGUF: '' is IN the vocab (id 11) but no merge + # path produces it, so bare conversion emits per-character pieces. + vocab = { + "<": 1, "t": 2, "h": 3, "i": 4, "n": 5, "k": 6, ">": 7, + "a": 8, "b": 9, "": 11, + } + tok = _tok(vocab) + before = tok.encode("", add_special_tokens=False) + assert before != [11], "sanity: must start out broken (split), like the real bug" + + names_by_id = sorted(vocab, key=vocab.get) + types = [1] * len(names_by_id) + types[names_by_id.index("")] = 4 # USER_DEFINED + _register_control_tokens(tok, names_by_id, types) + + assert tok.encode("", add_special_tokens=False) == [11] + assert tok.vocab_size == len(vocab), "vocab must not grow" + assert tok.convert_tokens_to_ids("") == 11, "id must be preserved" + + +def test_merge_reachable_control_token_left_alone(): + # A CONTROL entry that is ALREADY atomic (reachable via merges) must not be + # re-registered: the filter is encode(name) != [own id]. + vocab = {"a": 1, "b": 2, "ab": 3} + tok = _tok(vocab, merges=[("a", "b")]) + assert tok.encode("ab", add_special_tokens=False) == [3] + + names_by_id = ["a", "b", "ab"] + types = [1, 1, 3] + _register_control_tokens(tok, names_by_id, types) + assert tok.encode("ab", add_special_tokens=False) == [3] + assert tok.vocab_size == 3 + + +def test_normal_tokens_never_registered(): + # NORMAL entries that split must stay split: only CONTROL/UNKNOWN/USER_DEFINED + # types are eligible (BYTE/UNUSED excluded too). + vocab = {"a": 1, "b": 2, "c": 3} + tok = _tok(vocab) + names_by_id = ["a", "b", "c"] + _register_control_tokens(tok, names_by_id, [1, 5, 6]) + assert tok.encode("ab", add_special_tokens=False) == [1, 2] + + +def test_empty_types_is_noop(): + vocab = {"a": 1} + tok = _tok(vocab) + _register_control_tokens(tok, ["a"], []) + assert tok.vocab_size == 1 From bc6232453ca1b2714122c15004a95cdbd31e14c6 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Tue, 25 Aug 2026 22:34:18 -0300 Subject: [PATCH 3/6] fix(server): route Qwen3.6 to the qwen3_coder tool-call parser _infer_tool_call_parser only special-cased qwen3_5/coder names, so a Qwen3.6-* filename fell through to qwen25 (legacy JSON grammar). The 3.5/ 3.6 hybrid family instructs the XML invoke-block format in its chat template ( + blocks -- the Qwen3-Coder grammar), so the model emitted well-formed calls the wrong detector could not see: output swallowed, empty stop. Add qwen3_6/qwen3.6 markers. --- python/freetoken/server/args.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 2a533be7..303eec5d 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -141,6 +141,12 @@ def _infer_tool_call_parser(model_path: str) -> str: if ( "qwen3_5" in marker or "qwen3.5" in marker + or "qwen3_6" in marker + or "qwen3.6" in marker + # Qwen3-Coder and the 3.5/3.6 hybrid family share the XML invoke-block + # grammar (v); plain "qwen" (2.x) uses the + # older JSON form. A bare "qwen3" marker stays JSON (qwen25) unless it's + # a coder variant. or ("qwen3" in marker and "coder" in marker) ): return "qwen3_coder" From 25d7bd8d999044dd398b83fa26d5502b2496723c Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Tue, 25 Aug 2026 22:34:18 -0300 Subject: [PATCH 4/6] chore(local): ROCm serve scripts + VS Code tasks/debug configs - serve-qwen-moe.sh: single-line array-based launch (backslash-newline continuations get mangled by VS Code shells), setsid detach so a cancelled task cannot kill the server mid-load, CRLF guard, status/log subcommands, 128k KV default (~2.5 GiB on the hybrid arch). - mirror-hf-configs.sh: mirror official HF chat_template/generation_config next to a local GGUF without touching the file. - kill-freetoken.sh: hard-stop helper. - .vscode/: tasks (serve start/stop/status/log, fast tests), debugpy launch config for the server, LF pinning for shell scripts. - .gitignore: nohup.out, .plans/. --- .gitignore | 4 + .vscode/launch.json | 28 +++++++ .vscode/settings.json | 19 +++++ .vscode/tasks.json | 63 ++++++++++++++ scripts/kill-freetoken.sh | 53 ++++++++++++ scripts/mirror-hf-configs.sh | 67 +++++++++++++++ scripts/serve-qwen-moe.sh | 154 +++++++++++++++++++++++++++++++++++ 7 files changed, 388 insertions(+) create mode 100644 .vscode/launch.json create mode 100644 .vscode/settings.json create mode 100644 .vscode/tasks.json create mode 100755 scripts/kill-freetoken.sh create mode 100755 scripts/mirror-hf-configs.sh create mode 100755 scripts/serve-qwen-moe.sh diff --git a/.gitignore b/.gitignore index 95757625..edc08af2 100644 --- a/.gitignore +++ b/.gitignore @@ -233,3 +233,7 @@ benchmarks/cross_framework 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/ diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 00000000..64cbc790 --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,28 @@ +{ + // F5 launches the server IN THE FOREGROUND under the ROCm venv (no nohup), so + // stdout/stderr land straight in the Debug Console and breakpoints work. + "version": "0.2.0", + "configurations": [ + { + "name": "FreeToken: serve qwen-moe (debug)", + "type": "debugpy", + "request": "launch", + "module": "freetoken.cli", + "python": "${workspaceFolder}/.venv-rocm/bin/python", + "cwd": "${workspaceFolder}", + "env": { "PYTHONPATH": "${workspaceFolder}/python" }, + "console": "integratedTerminal", + "args": [ + "serve", + "--model", "/media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf", + "--moe-backend", "offload", + "--attention-backend", "triton", + "--moe-cache-size", "2048", + "--num-tokens", "131072", + "--host", "127.0.0.1", + "--port", "1920" + ], + "justMyCode": false + } + ] +} \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 00000000..7d6299c2 --- /dev/null +++ b/.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"] +} \ No newline at end of file diff --git a/.vscode/tasks.json b/.vscode/tasks.json new file mode 100644 index 00000000..175d2c8f --- /dev/null +++ b/.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": [] + } + ] +} \ No newline at end of file diff --git a/scripts/kill-freetoken.sh b/scripts/kill-freetoken.sh new file mode 100755 index 00000000..18893d54 --- /dev/null +++ b/scripts/kill-freetoken.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +# kill-freetoken.sh +# +# Kill every process related to FreeToken: the API server, backend workers, and the +# multiprocessing spawn/resource-tracker children, then free the listener ports. +# +# Usage: +# ./scripts/kill-freetoken.sh [PORTS...] # default: 1920 1921 + +set -u + +PORTS=("${@:-1920 1921}") +SELF=$$ +SELFCMDLINE="$(ps -p "$SELF" -o args= 2>/dev/null || true)" + +# Kill a pattern, excluding this script's own shell. +kill_matching() { + local pat="$1" + for pid in $(pgrep -f "$pat" 2>/dev/null); do + [ "$pid" = "$SELF" ] && continue + cmdline="$(ps -p "$pid" -o args= 2>/dev/null || true)" + [ -n "$cmdline" ] && [ "$cmdline" = "$SELFCMDLINE" ] && continue + kill -9 "$pid" 2>/dev/null || true + done +} + +# 1) FreeToken CLI / server / backend supervisor / backend workers. +kill_matching "freetoke[n].cli serve" +kill_matching "freetoke[n].cli" +kill_matching "freetoken" +kill_matching "multiprocessing.spawn" +kill_matching "multiprocessing.resource_tracker" +kill_matching "multiprocessing.semaphore" +kill_matching "torch.distributed.launch" + +# 2) Free the listener ports (a worker may still hold one). +for p in "${PORTS[@]}"; do + pid="$(ss -ltnp 2>/dev/null | grep ":$p" | grep -oP 'pid=\K[0-9]+' | head -1)" + [ -n "$pid" ] && kill -9 "$pid" 2>/dev/null || true +done + +sleep 1 + +left="$(pgrep -af "freetoken|multiprocessing.spawn|multiprocessing.resource" 2>/dev/null | grep -v "kill-freetoken.sh" | grep -v "$$" || true)" +if [ -n "$left" ]; then + echo "WARNING: still running (forced KILL):" + echo "$left" + kill_matching "freetoken" + kill_matching "multiprocessing.spawn" + sleep 1 +fi + +echo "FreeToken processes killed; ports ${PORTS[*]} freed." diff --git a/scripts/mirror-hf-configs.sh b/scripts/mirror-hf-configs.sh new file mode 100755 index 00000000..8cca566b --- /dev/null +++ b/scripts/mirror-hf-configs.sh @@ -0,0 +1,67 @@ +#!/usr/bin/env bash +# mirror-hf-configs.sh +# +# Mirror a model's official Hugging Face config files next to a local GGUF so +# FreeToken uses them instead of whatever the GGUF packager embedded: +# +# ./mirror-hf-configs.sh [gguf-file] +# +# ./mirror-hf-configs.sh Qwen/Qwen3.6-35B-A3B +# ./mirror-hf-configs.sh Qwen/Qwen3.6-35B-A3B /media/smk/Shared/Models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf +# +# With no gguf argument, defaults to $FT_MODEL, else the single *.gguf in +# /media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models or /media/smk/Shared/Models +# matching the repo's basename, else errors. +# +# Files fetched (when they exist upstream): +# chat_template.jinja — read by FreeToken's GGUF loader (overrides embedded) +# generation_config.json — default sampling params +# tokenizer_config.json — reference only (FreeToken builds its tokenizer from GGUF) +# +# To revert, delete the mirrored files — the GGUF metadata is never modified. + +set -euo pipefail + +REPO_ID="${1:?usage: mirror-hf-configs.sh [gguf-file]}" +shift || true + +if [ "$#" -ge 1 ]; then + GGUF="$1" +else + GGUF="${FT_MODEL:-}" + if [ -z "$GGUF" ]; then + base="$(basename "${REPO_ID##*:}")" + cand="$(find /media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models /media/smk/Shared/Models -maxdepth 1 -iname "*${base%%-*}*.gguf" 2>/dev/null | head -1)" + [ -n "$cand" ] || { echo "ERROR: no gguf found; pass one explicitly" >&2; exit 1; } + GGUF="$cand" + fi +fi +[ -f "$GGUF" ] || { echo "ERROR: not a file: $GGUF" >&2; exit 1; } + +DIR="$(dirname "$GGUF")" +echo "repo : $REPO_ID" +echo "into : $DIR" + +PY="${PY:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)/.venv-rocm/bin/python}" + +for f in chat_template.jinja generation_config.json tokenizer_config.json; do + if "$PY" - "$REPO_ID" "$f" "$DIR" <<'EOF' +import sys +from huggingface_hub import hf_hub_download +repo, fname, dest_dir = sys.argv[1:4] +try: + p = hf_hub_download(repo_id=repo, filename=fname) +except Exception: + sys.exit(1) +import shutil +shutil.copyfile(p, f"{dest_dir}/{fname}") +print(f"{fname}: OK") +EOF + then + : + else + echo "$f: not present upstream, skipped" + fi +done + +echo "done — restart the server to pick them up." diff --git a/scripts/serve-qwen-moe.sh b/scripts/serve-qwen-moe.sh new file mode 100755 index 00000000..6a28444d --- /dev/null +++ b/scripts/serve-qwen-moe.sh @@ -0,0 +1,154 @@ +#!/usr/bin/env bash +# serve-qwen-moe.sh +# +# Serve the Qwen3.5/3.6-35B-A3B hybrid MoE GGUF (GatedDeltaNet + full-attention) +# on AMD ROCm (gfx1100 / RX 7900 XTX). Uses the offload MoE backend (experts on the +# CPU/offload cache) and the triton attention backend by default. +# +# Native context window: 262144 (256K) tokens (max_position_embeddings in the model). +# Graph capture is settled as a failure on ROCm, so this runs eager kernel-launch decode. +# +# VS CODE NOTES: +# - The server command is built as a bash array and launched on ONE physical +# line. Backslash-newline continuations get mangled by some VS Code shells / +# task runners (each continuation line then executes as its own command), +# which is exactly how `nohup.out` ended up with bare "--model: command not +# found" errors. Do NOT reintroduce multi-line command strings here. +# - .vscode/settings.json pins files.eol=\n for *.sh; the CRLF guard below +# catches any violation early instead of failing obscurely mid-launch. +# +# Usage: +# ./serve-qwen-moe.sh # launch on 127.0.0.1:1920, triton attention +# FT_ATTN=torch ./serve-qwen-moe.sh # A/B against the pure-torch reference backend +# FT_PORT=1930 ./serve-qwen-moe.sh # pick another port +# ./serve-qwen-moe.sh stop # kill the running server +# ./serve-qwen-moe.sh status # running? + tail of the log +# +# Or from the VS Code Command Palette: "Tasks: Run Task" -> FreeToken: ... + +set -euo pipefail + +REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +PY="${PY:-$REPO/.venv-rocm/bin/python}" + +MODEL="${FT_MODEL:-/media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf}" +HOST="${FT_HOST:-127.0.0.1}" +PORT="${FT_PORT:-1920}" +ATNN="${FT_ATTN:-triton}" # triton | torch +MOE_BACKEND="${FT_MOE:-offload}" # offload required for K-quant experts +# VS Code / Copilot injects a large chat context, so the KV cache must be bigger than +# the tiny 8K the MoE-auto cache leaves. --num-tokens sizes the KV cache in tokens; +# --moe-cache-size limits the GPU expert cache so KV has room (fewer slots = slower decode). +# Hybrid arch: only 10/40 layers are full attention (2 kv heads x 256 dim) -> +# ~20 KiB/token bf16, i.e. 128k ~= 2.7 GiB, full native 256k ~= 5.4 GiB. +KV_TOKENS="${FT_KV_TOKENS:-131072}" +MOE_CACHE="${FT_MOE_CACHE:-2048}" +LOG="${FT_LOG:-/tmp/serve_qwen_moe.log}" + +die() { echo "ERROR: $*" >&2; exit 1; } + +# Guard against the exact failure mode VS Code caused before: if this file ever +# gets saved with CRLF endings, every argument silently grows a trailing \r. +if grep -q $'\r' "${BASH_SOURCE[0]}"; then + die "CRLF line endings detected in $(basename "${BASH_SOURCE[0]}"). Run: sed -i 's/\r$//' ${BASH_SOURCE[0]}" +fi + +[ -x "$PY" ] || die "python not found: $PY (set PY=/path/to/venv-python)" + +server_pids() { + pgrep -f "freetoke[n].cli serve" || true +} + +start() { + if [ -n "$(server_pids)" ]; then + echo "A server is already running (pid $(server_pids | tr '\n' ' ')). Use 'stop' first." + exit 1 + fi + [ -f "$MODEL" ] || die "model not found: $MODEL" + + # One arg per array element; expanded once, single line, no continuations. + local -a SERVE_ARGS=( + "--model" "$MODEL" + "--moe-backend" "$MOE_BACKEND" + "--attention-backend" "$ATNN" + "--moe-cache-size" "$MOE_CACHE" + "--num-tokens" "$KV_TOKENS" + "--host" "$HOST" + "--port" "$PORT" + ) + + echo "Launching FreeToken server" + echo " model : $MODEL" + echo " listen : $HOST:$PORT" + echo " attn : $ATNN" + echo " moe : $MOE_BACKEND" + echo " kv : $KV_TOKENS tokens (gpu moe cache: $MOE_CACHE slots)" + echo " python : $PY" + echo " log : $LOG" + + # setsid: give the server its OWN session/process group. nohup alone only + # ignores SIGHUP — a caller that dies (e.g. a VS Code task cancelled, an agent + # tool timeout) still takes down the whole process group with SIGKILL/SIGTERM, + # which silently killed the server mid model-load once already. + cd "$REPO" + PYTHONPATH="$REPO/python" setsid nohup "$PY" -m freetoken.cli serve "${SERVE_ARGS[@]}" >"$LOG" 2>&1 /dev/null || true + echo "pid=$pid — waiting for readiness (model load takes ~3-4 min)..." +} + +wait_ready() { + for _ in $(seq 1 90); do + if grep -q "API server is ready" "$LOG" 2>/dev/null; then + echo "READY on $HOST:$PORT" + return 0 + fi + if [ -z "$(server_pids)" ]; then + echo "server exited; last log lines:" >&2 + tail -20 "$LOG" >&2 || true + return 1 + fi + sleep 5 + done + echo "timed out waiting for readiness; see $LOG" >&2 + return 1 +} + +stop() { + local pids + pids="$(server_pids)" + if [ -z "$pids" ]; then + echo "no server running" + else + pkill -9 -f "freetoke[n].cli serve" 2>/dev/null || true + echo "stopped (was pid $pids)" + fi + pkill -9 -f "multiprocessing.spawn" 2>/dev/null || true + # free the distributed worker port (server_port+1) + local p pid + for p in "$PORT" "$((PORT + 1))"; do + pid="$(ss -ltnp 2>/dev/null | grep ":$p" | grep -oP 'pid=\K[0-9]+' | head -1 || true)" + if [ -n "$pid" ]; then + kill -9 "$pid" 2>/dev/null || true + fi + done +} + +status() { + local pids + pids="$(server_pids)" + if [ -n "$pids" ]; then + echo "RUNNING (pid $(echo "$pids" | tr '\n' ' ')) on $HOST:$PORT" + else + echo "NOT RUNNING" + fi + [ -f "$LOG" ] && echo "--- last 5 log lines ($LOG) ---" && tail -5 "$LOG" +} + +case "${1:-start}" in + start) start && wait_ready ;; + stop) stop ;; + status) status ;; + log) exec tail -f "$LOG" ;; + *) echo "usage: $0 [start|stop|status|log]" >&2; exit 1 ;; +esac From 729789ae0be705badd7ec3d8b0563a1a4f05c5a3 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Wed, 26 Aug 2026 12:37:47 -0300 Subject: [PATCH 5/6] fix(qwen3.5-moe): GGUF lm_head last-token gather + presence/frequency penalties Root cause of the intermittent empty/loop responses and the non-deterministic 'greedy' decode: the GGUF lm_head (GGUFLinear) returned FULL prefill logits [total, vocab] without gathering each request's final row, while the engine contract is [batch.size, vocab] (the token to sample after each request's prompt). ParallelLMHead and Nvfp4LMHead already gather via attn_metadata.get_last_indices(); the GGUF path did not. Consequences fixed: - The first generated token was sampled from POSITION 0's logits (after '<|im_start|>' -> 'user') instead of the last prompt position (after '...assistant\n thinking\n' -> 'Here'). - Worse, a fresh prefill (142-token batch, row 0 = prompt token 0) and a radix-cache continuation (14-token batch, row 0 = prompt token 128 -> '<|im_end|>' -> immediate stop, EMPTY content) sampled different rows, so the SAME greedy request gave different outputs depending on cache state -- the 'model thinks then returns empty' report. Fix: qwen3_5_moe/model.py gathers last_indices on prefill before the GGUF lm head (matches ParallelLMHead/Nvfp4LMHead; shared code path for both backends). Also in this change: - Presence/frequency penalties were accepted by the API but IGNORED by the sampler. Implemented end-to-end (core.py SamplingParams + Req.prompt_len, generation.py/openai_api.py pass-through, engine/sample.py apply_penalties over generated tokens only, engine.py passes the batch to sample()). Applies to greedy too. Breaks reasoning loops by penalizing repeated tokens. - Serve script: --max-output-tokens 65536 default (FT_MAX_OUTPUT knob) so the reasoning model has room to finish; keeps the Inc-1 diagnostic knobs (moe stats, prefill overlap, cpu layers). Verified: - 28k-token prompt now ANSWERS (finish=stop, content='Four'), no empty, no loop. - Tool calls still return tool_calls with correct args. - TRUE greedy (top_k=1, top_p=1.0) deterministic within a cache state; the first-after-startup fresh run differs from radix hits only at ~1 bf16 ULP (continuation GEMM batch shapes) -- documented residual, both outputs valid. - tests/engine + tests/server: 618 passed, same 15 pre-existing failures. --- python/freetoken/core.py | 15 ++++ python/freetoken/engine/engine.py | 2 +- python/freetoken/engine/sample.py | 50 +++++++++++-- python/freetoken/models/qwen3_5_moe/model.py | 15 +++- python/freetoken/scheduler/scheduler.py | 13 ++++ python/freetoken/scheduler/status.py | 22 ++++++ python/freetoken/server/args.py | 11 +++ python/freetoken/server/generation.py | 4 ++ python/freetoken/server/openai_api.py | 4 ++ scripts/serve-qwen-moe.sh | 74 +++++++++++++++++--- 10 files changed, 196 insertions(+), 14 deletions(-) diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ef0a539c..82c6e087 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -25,6 +25,13 @@ class SamplingParams: # Stop strings (OpenAI `stop` / Anthropic `stop_sequences`). Generation finishes when one # appears in the decoded output; the matched substring (and anything after) is trimmed. stop_strs: list[str] = field(default_factory=list) + # OpenAI-style presence/frequency penalties, applied over the tokens this request has + # generated so far (the prompt is excluded): + # logits[t] -= presence_penalty * (t was generated) + frequency_penalty * count(t) + # Positive values push the model away from repeating itself (breaks reasoning loops); + # negative values encourage repetition (coherence). 0.0 = disabled. + presence_penalty: float = 0.0 + frequency_penalty: float = 0.0 @property def is_greedy(self) -> bool: @@ -40,6 +47,9 @@ class Req: uid: int sampling_params: SamplingParams cache_handle: BaseCacheHandle + # Prompt length at creation (see __post_init__); input_ids[prompt_len:] is the + # generated portion used by presence/frequency penalties. + prompt_len: int = 0 # Optional precomputed multimodal soft-token embeddings (GPU, [num_image_tokens, # hidden]) scattered at image-token positions during this request's prefill. mm_embeds: torch.Tensor | None = None @@ -71,6 +81,11 @@ def __post_init__(self) -> None: self.max_device_len = len(self.input_ids) + self.output_len assert 0 <= self.cached_len < self.device_len <= self.max_device_len self._alloc_ids_buf() + # Length of the prompt this request was created with. Generation grows input_ids + # past this point (append_host), so input_ids[prompt_len:] is exactly the tokens + # the model has generated so far -- what presence/frequency penalties are applied + # over. ChunkedReq instances (which never sample) record their partial length. + self.prompt_len = self.device_len def _alloc_ids_buf(self) -> None: self._ids_buf = torch.empty(self.max_device_len, dtype=self.input_ids.dtype) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 5ada3954..d86bcf09 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -945,7 +945,7 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: req.complete_one() batch_logits = logits[: batch.size] - next_tokens_gpu = self.sampler.sample(batch_logits, args).to(torch.int32) + next_tokens_gpu = self.sampler.sample(batch_logits, args, batch).to(torch.int32) next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True) copy_done_event = torch.cuda.Event() copy_done_event.record(self.stream) diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 01d14b1a..be3726b4 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -7,7 +7,7 @@ from freetoken.utils import is_sm90_supported, nvtx_annotate if TYPE_CHECKING: - from freetoken.core import Batch + from freetoken.core import Batch, Req @dataclass @@ -15,6 +15,36 @@ class BatchSamplingArgs: temperatures: torch.Tensor | None top_k: torch.Tensor | None = None top_p: torch.Tensor | None = None + # True when at least one request carries a presence/frequency penalty; the sampler + # then lowers each request's logits over its generated tokens before sampling. + apply_penalties: bool = False + + +def apply_penalties( + logits: torch.Tensor, + reqs: List["Req"], +) -> None: + """Apply OpenAI presence/frequency penalties to ``logits`` in place (row per req). + + For a token ``t`` the request already generated, its score is lowered by + ``presence_penalty`` plus ``frequency_penalty * count(t)``. The prompt is excluded + (only ``input_ids[req.prompt_len:]`` counts), so the penalty grows with the + generation itself -- positive values push the model away from repeating itself, + which breaks reasoning loops; negative values nudge it toward repetition. + """ + for i, req in enumerate(reqs): + sp = req.sampling_params + pp, fp = sp.presence_penalty, sp.frequency_penalty + if not pp and not fp: + continue + gen = req.input_ids[req.prompt_len :] + if gen.numel() == 0: + continue + uniq, counts = torch.unique(gen, return_counts=True) + vals = torch.full_like(counts, pp, dtype=torch.float32) + fp * counts.to( + torch.float32 + ) + logits[i, uniq.to(logits.device)] -= vals.to(logits.device) def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor: @@ -57,7 +87,10 @@ class Sampler: def prepare(self, batch: Batch) -> BatchSamplingArgs: params = [r.sampling_params for r in batch.reqs] - if all(p.is_greedy for p in params): + apply_penalties = any( + p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 for p in params + ) + if all(p.is_greedy for p in params) and not apply_penalties: return BatchSamplingArgs(temperatures=None) MIN_P = MIN_T = 1e-6 @@ -70,11 +103,20 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: top_k = make_device_tensor(top_ks, torch.int32, self.device) if any(p < 1.0 for p in top_ps): top_p = make_device_tensor(top_ps, torch.float32, self.device) - return BatchSamplingArgs(temperatures, top_k=top_k, top_p=top_p) + return BatchSamplingArgs( + temperatures, + top_k=top_k, + top_p=top_p, + apply_penalties=apply_penalties, + ) @nvtx_annotate("Sampler") - def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: + def sample( + self, logits: torch.Tensor, args: BatchSamplingArgs, batch: Batch + ) -> torch.Tensor: with torch.cuda.nvtx.range("Sampler"): + if args.apply_penalties: + apply_penalties(logits, batch.reqs) if args.temperatures is None: # greedy sampling return torch.argmax(logits, dim=-1) return sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py index 32954dc5..df0d41f0 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -117,7 +117,20 @@ def __init__(self, config: ModelConfig): convert_qwen35moe_to_gguf(self, config) def forward(self) -> torch.Tensor: - output = self.model.forward(get_global_ctx().batch.input_ids) + ctx = get_global_ctx() + batch = ctx.batch + output = self.model.forward(batch.input_ids) + if batch.is_prefill: + # GGUFLinear (unlike ParallelLMHead / Nvfp4LMHead) does not gather each + # request's final row itself, but the engine contract is [batch.size, vocab] + # (the logits to sample after each request's prompt). Without this gather + # the sampler reads the FIRST prompt position's logits: the first generated + # token comes from position 0, and -- worse -- differs between a fresh + # prefill and a radix-cache continuation of the same prompt (the restored + # continuation has a different first row), which made greedy output flip + # between server states / cache states. + indices = batch.attn_metadata.get_last_indices(batch.size) + output = output[indices] return self.lm_head.forward(output) diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 35541161..0846c940 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -411,9 +411,22 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: page_size=self.config.page_size, mamba_slots=mamba_slots, swa_tokens=swa_tokens, + moe_stats=self._moe_stats_snapshot(), ) self.send_result(reply) + def _moe_stats_snapshot(self) -> dict | None: + """Per-window MoE cache hit/miss stats for the decode log line, or None when + stats collection is off (the default). Reads device counters once per call; + the status reporter only calls this every decode_log_interval steps.""" + cache = getattr(self.engine, "moe_offload_cache", None) + if cache is None or not getattr(cache, "collect_stats", False): + return None + try: + return cache.decode_miss_stats() + except Exception: + return None + def _match_stop_str(self, req: Req) -> str | None: """First stop string present in this request's generated tail, else None. Decodes only a short suffix (bounded by the longest stop string's char length, so a stop of diff --git a/python/freetoken/scheduler/status.py b/python/freetoken/scheduler/status.py index ca706c50..c7be5e8f 100644 --- a/python/freetoken/scheduler/status.py +++ b/python/freetoken/scheduler/status.py @@ -34,6 +34,7 @@ def report_batch( page_size: int, mamba_slots: tuple[int, int] | None = None, swa_tokens: tuple[int, int] | None = None, + moe_stats: dict | None = None, ) -> None: if batch.is_prefill: self._report_prefill( @@ -55,6 +56,7 @@ def report_batch( page_size=page_size, mamba_slots=mamba_slots, swa_tokens=swa_tokens, + moe_stats=moe_stats, ) def _report_prefill( @@ -101,6 +103,7 @@ def _report_decode( page_size: int, mamba_slots: tuple[int, int] | None = None, swa_tokens: tuple[int, int] | None = None, + moe_stats: dict | None = None, ) -> None: self._decode_forward_count += 1 self._decode_generated_tokens += len(batch.reqs) @@ -121,6 +124,7 @@ def _report_decode( f"{_mamba_msg(mamba_slots)}" f"gen throughput (token/s): {gen_throughput:.2f}, " f"#queue-req: {queue_reqs}" + f"{_moe_msg(moe_stats)}" ) @@ -128,6 +132,24 @@ def _usage_ratio(used: int, total: int) -> float: return used / total if total > 0 else 0.0 +def _moe_msg(stats: dict | None) -> str: + """MoE cache hit/miss summary for the decode log line (empty when stats are off).""" + if not stats: + return "" + miss = stats.get("miss_rate") + fetch = stats.get("fetch_rate") + cpu = stats.get("cpu_per_layer") + hit = (1.0 - miss) if miss is not None else None + parts = [f"moe hit: {hit:.3f}" if hit is not None else "moe hit: n/a"] + if miss is not None: + parts.append(f"miss: {miss:.3f}") + if fetch is not None: + parts.append(f"fetch: {fetch:.3f}") + if cpu is not None: + parts.append(f"cpu: {cpu:.3f}") + return ", " + ", ".join(parts) + + def _mamba_msg(mamba_slots: tuple[int, int] | None) -> str: """GDN-state (mamba) pool occupancy for hybrid models; empty for the rest.""" if mamba_slots is None: diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 303eec5d..cb651221 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -585,6 +585,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--moe-collect-stats", + action="store_true", + default=ServerArgs.moe_collect_stats, + help=( + "Capture MoE cache hit/miss counters into the decode graph and print them " + "in the decode log line (moe hit/miss/fetch/cpu). Off by default; the " + "device-side accumulation is captured into the CUDA graph." + ), + ) + parser.add_argument( "--moe-prefill-hit-d2d", action="store_true", diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index be05d908..6e2b9ab4 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -163,6 +163,8 @@ def resolve_sampling( ignore_eos: bool, model_sampling: dict[str, Any], stop: str | list[str] | None = None, + presence_penalty: float | None = None, + frequency_penalty: float | None = None, ) -> SamplingParams: """Map a protocol's sampling fields onto the engine's neutral SamplingParams, filling unspecified fields from the checkpoint's recommended defaults.""" @@ -184,6 +186,8 @@ def pick(value, key, framework): top_k=pick(top_k, "top_k", -1), top_p=pick(top_p, "top_p", 1.0), stop_strs=[s for s in stop_list if s], # drop empty strings (would match everything) + presence_penalty=presence_penalty if presence_penalty is not None else 0.0, + frequency_penalty=frequency_penalty if frequency_penalty is not None else 0.0, ) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b4becd26..04822dc7 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -76,6 +76,8 @@ def chat_request_to_genspec( ignore_eos=req.ignore_eos, model_sampling=model_sampling, stop=req.stop, + presence_penalty=req.presence_penalty, + frequency_penalty=req.frequency_penalty, ), chat_template_kwargs=ctk, template_tools=_tools_for_template(req), @@ -526,6 +528,8 @@ def _resolve_sampling( ignore_eos=req.ignore_eos, model_sampling=model_sampling, stop=req.stop, + presence_penalty=getattr(req, "presence_penalty", None), + frequency_penalty=getattr(req, "frequency_penalty", None), ) diff --git a/scripts/serve-qwen-moe.sh b/scripts/serve-qwen-moe.sh index 6a28444d..47de497c 100755 --- a/scripts/serve-qwen-moe.sh +++ b/scripts/serve-qwen-moe.sh @@ -36,17 +36,51 @@ HOST="${FT_HOST:-127.0.0.1}" PORT="${FT_PORT:-1920}" ATNN="${FT_ATTN:-triton}" # triton | torch MOE_BACKEND="${FT_MOE:-offload}" # offload required for K-quant experts -# VS Code / Copilot injects a large chat context, so the KV cache must be bigger than -# the tiny 8K the MoE-auto cache leaves. --num-tokens sizes the KV cache in tokens; -# --moe-cache-size limits the GPU expert cache so KV has room (fewer slots = slower decode). -# Hybrid arch: only 10/40 layers are full attention (2 kv heads x 256 dim) -> -# ~20 KiB/token bf16, i.e. 128k ~= 2.7 GiB, full native 256k ~= 5.4 GiB. +# radix = cross-request GDN-state prefix reuse (default); naive = no prefix reuse. +# Suspect for long-context degradation: radix GDN-state reuse corrupting later requests. +CACHE_TYPE="${FT_CACHE_TYPE:-radix}" +# MoE cache hit/miss stats in the decode log line (--moe-collect-stats). +MOE_STATS="${FT_MOE_STATS:-1}" +# Disable the two-buffer prefill MoE overlap (diagnostic: race check). +PREFILL_OVERLAP="${FT_PREFILL_OVERLAP:-1}" +# MoE layers computed on the CPU executor (diagnostic: '0' = all-GPU, no CPU path). +CPU_LAYERS="${FT_CPU_LAYERS:-}" +# --num-tokens sizes the KV cache in tokens (hybrid arch: only 10/40 layers are full +# attention, 2 kv heads x 256 dim -> ~20 KiB/token bf16; 128k ~= 2.7 GiB). +# The GPU expert-slot cache is sized by FT_MOE_CACHE: +# auto -> --moe-cache-auto: the engine derives slot bytes from the real expert +# tensors and fills all free VRAM AFTER reserving kv-reserve-tokens for KV. +# --kv-reserve-tokens MUST equal --num-tokens here: with an explicit +# --num-tokens the engine skips auto's own KV-half plan (num_page_override +# is set), so without a matching reservation greedy expert fill would eat +# VRAM the pinned KV still needs -> late CUDA OOM. +# -> fixed --moe-cache-size N slots (legacy behavior). +# Headroom: the engine may use FT_MEMORY_RATIO of free VRAM for weights+KV+experts +# combined (default here 0.80, upstream default 0.9). The remainder absorbs prefill +# transients -- the MoE overlap double-buffer alone needs ~3.8 GiB at this model's +# batch shapes; 0.9 left only ~2 GiB and OOM'd mid-prefill under VS Code payloads. +MEMORY_RATIO="${FT_MEMORY_RATIO:-0.80}" +# --max-prefill-length caps chunked-prefill chunk size (engine default 8192). The lm_head +# materializes logits for EVERY chunk token: an 8192-token chunk spikes ~3.8 GiB +# transiently -- enough to OOM on VS Code-sized prompts even with healthy headroom. +# 4096 halves the spike; long prompts just prefill in more chunks. +PREFILL_CHUNK="${FT_PREFILL_CHUNK:-4096}" KV_TOKENS="${FT_KV_TOKENS:-131072}" -MOE_CACHE="${FT_MOE_CACHE:-2048}" +# Default max output tokens for requests that omit max_tokens. The reasoning model +# sometimes needs more room to finish its reasoning before answering; the engine +# default is 32k. Copilot sends its own max_tokens, which we cannot override, but +# other clients inherit this. +MAX_OUTPUT="${FT_MAX_OUTPUT:-65536}" +MOE_CACHE="${FT_MOE_CACHE:-auto}" LOG="${FT_LOG:-/tmp/serve_qwen_moe.log}" die() { echo "ERROR: $*" >&2; exit 1; } +# Map FT_MOE_CACHE to argv. Only the auto path carries --kv-reserve-tokens: +# with a fixed size the CLI already validates fit against the pinned KV. +# NOTE: validation must run in THIS shell (no process substitution), or die() +# would only exit a subshell and the launch would continue unvalidated. + # Guard against the exact failure mode VS Code caused before: if this file ever # gets saved with CRLF endings, every argument silently grows a trailing \r. if grep -q $'\r' "${BASH_SOURCE[0]}"; then @@ -70,19 +104,36 @@ start() { local -a SERVE_ARGS=( "--model" "$MODEL" "--moe-backend" "$MOE_BACKEND" + "--cache-type" "$CACHE_TYPE" + $( [ "$MOE_STATS" = 1 ] && echo "--moe-collect-stats" ) + $( [ "$PREFILL_OVERLAP" = 0 ] && echo "--disable-moe-prefill-overlap" ) + $( [ -n "$CPU_LAYERS" ] && echo "--moe-cpu-layers" "$CPU_LAYERS" ) "--attention-backend" "$ATNN" - "--moe-cache-size" "$MOE_CACHE" "--num-tokens" "$KV_TOKENS" + "--memory-ratio" "$MEMORY_RATIO" + "--max-prefill-length" "$PREFILL_CHUNK" + "--max-output-tokens" "$MAX_OUTPUT" "--host" "$HOST" "--port" "$PORT" ) + case "$MOE_CACHE" in + auto) + SERVE_ARGS+=("--moe-cache-auto" "--kv-reserve-tokens" "$KV_TOKENS") + ;; + ''|*[!0-9]*) + die "FT_MOE_CACHE='$MOE_CACHE' is invalid: use 'auto' or a slot count" + ;; + *) + SERVE_ARGS+=("--moe-cache-size" "$MOE_CACHE") + ;; + esac echo "Launching FreeToken server" echo " model : $MODEL" echo " listen : $HOST:$PORT" echo " attn : $ATNN" echo " moe : $MOE_BACKEND" - echo " kv : $KV_TOKENS tokens (gpu moe cache: $MOE_CACHE slots)" + echo " kv : $KV_TOKENS tokens (gpu moe cache: ${MOE_CACHE}${MOE_CACHE:+ }$( [ "$MOE_CACHE" = auto ] && echo "kv-reserve $KV_TOKENS" || echo slots))" echo " python : $PY" echo " log : $LOG" @@ -142,6 +193,13 @@ status() { else echo "NOT RUNNING" fi + # Surface the auto-resolved expert-cache split so users don't have to read + # engine log lines; only present when FT_MOE_CACHE=auto booted the server. + if [ -f "$LOG" ]; then + local resolved + resolved="$(grep -o -- '--moe-cache-auto resolved moe_cache_size=[0-9]* num_pages=[0-9]*' "$LOG" | tail -1)" + [ -n "$resolved" ] && echo "moe cache: ${resolved//--moe-cache-auto resolved /}" + fi [ -f "$LOG" ] && echo "--- last 5 log lines ($LOG) ---" && tail -5 "$LOG" } From 1f76e88e5be3b063de93c282d053d35a32296a24 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Wed, 26 Aug 2026 13:47:34 -0300 Subject: [PATCH 6/6] chore(rocm): drop CI workflow + tinygrad fallback; review cleanup - Remove .github/workflows/rocm.yml (no GH Actions for ROCm) - Remove kernel/tinygrad_fallback.py (dead: nothing imports it) - Remove .vscode/launch.json (hardcoded machine model path) - kernel/utils.py: drop shadowed duplicate _build_stamps definition - kernel/backend.py: drop unused _CUDA_ONLY_PACKAGES - models/qwen3_5_moe/gguf.py: drop unused _q8_0_down_row_bytes - moe/nvfp4_to_mxfp4.py: drop unused _nearest_e2m1_codes/_FP4_SORT_SIGN; assert K is a block multiple instead of silently truncating - moe/expert_banks.py: remove dead dummy-path import in _gguf_banks - models/qwen3_5_moe/gdn.py: rename _fp8 -> _split_proj (also covers GGUF) - engine/engine.py: drop redundant local is_rocm imports; build infos once - utils/graph_gate.py, engine/graph.py, attention/torch.py, tests: strip plan-increment (Inc N) references - docs/install-amd.md, pyproject.toml: drop plan refs; align ROCm torch index - kernel/gguf.py: honor FREETOKEN_KERNEL_CACHE_GFX for the JIT offload-arch - scripts: remove hardcoded /media/smk model paths (require FT_MODEL/arg) - utils.cuh: fix stray [[unlikely]]; statement --- .github/workflows/rocm.yml | 70 -------------- .vscode/launch.json | 28 ------ docs/install-amd.md | 6 +- pyproject.toml | 6 +- python/freetoken/attention/torch.py | 5 +- python/freetoken/engine/engine.py | 7 +- python/freetoken/engine/graph.py | 2 +- python/freetoken/kernel/backend.py | 6 -- .../kernel/csrc/include/freetoken/utils.cuh | 6 +- python/freetoken/kernel/gguf.py | 5 +- python/freetoken/kernel/tinygrad_fallback.py | 92 ------------------- python/freetoken/kernel/utils.py | 6 -- python/freetoken/models/qwen3_5_moe/gdn.py | 7 +- python/freetoken/models/qwen3_5_moe/gguf.py | 4 - python/freetoken/moe/expert_banks.py | 7 +- python/freetoken/moe/nvfp4_to_mxfp4.py | 22 +---- python/freetoken/utils/graph_gate.py | 15 ++- scripts/mirror-hf-configs.sh | 15 +-- scripts/serve-qwen-moe.sh | 3 +- tests/attention/test_torch_backend.py | 2 +- 20 files changed, 42 insertions(+), 272 deletions(-) delete mode 100644 .github/workflows/rocm.yml delete mode 100644 .vscode/launch.json delete mode 100644 python/freetoken/kernel/tinygrad_fallback.py diff --git a/.github/workflows/rocm.yml b/.github/workflows/rocm.yml deleted file mode 100644 index e2f1fe49..00000000 --- a/.github/workflows/rocm.yml +++ /dev/null @@ -1,70 +0,0 @@ -name: ROCm (AMD) correctness smoke - -# ROCm CI job: compiles the AOT kernel cache for the RX 7000 (gfx1100) target on a -# ROCm torch install and runs a torch-free correctness smoke plus the AMD unit tests. -# Gated on a self-hosted runner that has ROCm torch + hipcc. The primary NVIDIA release -# flow is release.yml; this job is additive and must not gate NVIDIA releases. - -on: - workflow_dispatch: - push: - branches: [main] - pull_request: - -jobs: - rocm-smoke: - runs-on: [self-hosted, linux, amd, rocm] - timeout-minutes: 60 - env: - FREETOKEN_DISABLE_JIT: "1" - FREETOKEN_KERNEL_CACHE_GFX: "gfx1100" - steps: - - uses: actions/checkout@v4 - - - name: Check ROCm toolchain - run: | - set -e - command -v hipcc || ls /opt/rocm/bin/hipcc - "${PYTHON:-python3}" -c "import torch.version as v; print('torch hip:', v.hip)" - - - name: Install build deps - run: | - python -m pip install --upgrade pip wheel setuptools - python -m pip install -e "python[rocm]" - - - name: Compile AOT kernel cache for gfx1100 - run: | - FREETOKEN_KERNEL_CACHE_VERBOSE=1 python -m pip wheel ./freetoken-kernel-cache -w dist/cache-rocm - - - name: Install prebuilt kernel cache - run: | - whl="$(find dist/cache-rocm -name 'freetoken_kernel_cache-*.whl' | head -1)" - python -m pip install --force-reinstall "$whl" - - - name: Torch-free AMD unit tests - run: | - python -m pytest \ - tests/utils/test_device_kind.py \ - tests/kernels/test_toolchain_hip.py \ - tests/kernels/test_backend_rocm.py \ - tests/kernels/test_cache_rocm_pairing.py \ - tests/moe/test_nvfp4_to_mxfp4.py \ - -q - - - name: Hardware correctness smoke (serves on RX 7000) - run: | - # Functional path only -- flashinfer/sgl/trtllm are NVIDIA-only and must not - # be selected. AUTO backend must resolve to triton; NVFP4 auto -> triton. - python - <<'PY' - from freetoken.utils.arch import is_rocm, is_gfx_arch_ge - from freetoken.moe.nvfp4_backends import select_nvfp4_backend - import torch - assert is_rocm(), "expected a ROCm torch build" - assert is_gfx_arch_ge(1100), "expected gfx1100-class device (RX 7000)" - print("NVFP4 auto ->", select_nvfp4_backend(torch.device("cuda"), 768, "auto")) - PY - - - name: Serve smoke - run: | - FREETOKEN_DEVICE=cuda python -m freetoken.serve --help >/dev/null \ - && echo "freetoken CLI loads on ROCm" diff --git a/.vscode/launch.json b/.vscode/launch.json deleted file mode 100644 index 64cbc790..00000000 --- a/.vscode/launch.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - // F5 launches the server IN THE FOREGROUND under the ROCm venv (no nohup), so - // stdout/stderr land straight in the Debug Console and breakpoints work. - "version": "0.2.0", - "configurations": [ - { - "name": "FreeToken: serve qwen-moe (debug)", - "type": "debugpy", - "request": "launch", - "module": "freetoken.cli", - "python": "${workspaceFolder}/.venv-rocm/bin/python", - "cwd": "${workspaceFolder}", - "env": { "PYTHONPATH": "${workspaceFolder}/python" }, - "console": "integratedTerminal", - "args": [ - "serve", - "--model", "/media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf", - "--moe-backend", "offload", - "--attention-backend", "triton", - "--moe-cache-size", "2048", - "--num-tokens", "131072", - "--host", "127.0.0.1", - "--port", "1920" - ], - "justMyCode": false - } - ] -} \ No newline at end of file diff --git a/docs/install-amd.md b/docs/install-amd.md index 62ecf2fe..ede4ab80 100644 --- a/docs/install-amd.md +++ b/docs/install-amd.md @@ -6,7 +6,7 @@ recovered via HIP ports where safe. This page covers installing and running on R > 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. See `.plans/amd-gpu-support/plan.md`. +> recovering performance via the HIP kernel ports. ## Requirements @@ -41,10 +41,10 @@ and their backends are rejected with a clean error if requested. | 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 (Inc 3) | +| 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 Inc-1 gate passes | otherwise kernel-launch decode | +| 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 diff --git a/pyproject.toml b/pyproject.toml index 13b52c19..2153c789 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,9 +80,9 @@ fi = ["flashinfer-python[cu13]>=0.6,<0.7"] 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+rocm6.2` from -# https://download.pytorch.org/whl/rocm6.2) so the native extensions build against the HIP -# runtime; this extra pins the rest. See docs/install.md (AMD section, Inc 9). +# 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'", ] diff --git a/python/freetoken/attention/torch.py b/python/freetoken/attention/torch.py index 6ae16c2e..c19a9fcf 100644 --- a/python/freetoken/attention/torch.py +++ b/python/freetoken/attention/torch.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, List +import os import torch from freetoken.core import Batch, get_global_ctx @@ -64,10 +65,8 @@ def __init__(self, config: ModelConfig): self.num_kv_heads = int(getattr(spec, "num_kv_heads", self.num_kv_heads)) break # Debugging: contiguous (per-request) cache instead of the paged pool, to - # isolate cache addressing from the attention compute (Inc 5). + # isolate cache addressing from the attention compute. self._contig: dict[tuple[int, int], list] = {} - import os - self._use_contig = os.environ.get("FT_DEBUG_CONTIG_CACHE") == "1" def _build_metadata(self, batch: Batch) -> TorchMetadata: diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index d86bcf09..74877ce7 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -106,17 +106,14 @@ def _backend_parts_serve(name: str, required: frozenset[AttnType]) -> bool: def _backend_requirements_met(name: str) -> bool: + infos = [attention_backend_info(part) for part in name.split(",")] # On ROCm (AMD) only the portable backends exist: flashinfer/sgl/trtllm (and anything # sm_100-gated) are NVIDIA-only, so short-circuit before probing them at all. - from freetoken.utils.arch import is_rocm - if is_rocm(): return all(not i.requires_flashinfer and not i.requires_sgl_kernel - and not i.requires_sm100 for i in - [attention_backend_info(p) for p in name.split(",")]) + and not i.requires_sm100 for i in infos) # flashinfer first across ALL parts: the sgl probe logs a "falls back to fi" warning, # which would mislead when the candidate is about to fail on flashinfer anyway. - infos = [attention_backend_info(part) for part in name.split(",")] if any(i.requires_flashinfer for i in infos) and not _flashinfer_available(): return False if any(i.requires_sgl_kernel for i in infos) and not _sgl_flash_attn_available(): diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 50bac2f0..276b741e 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -132,7 +132,7 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel # graphs-disabled early return so that config gets the phase too. emit_progress("Capturing CUDA graphs / warming up", 0, 0) self.graph_map: Dict[int, torch.cuda.CUDAGraph] = {} - # Inc-8 parity: on ROCm, honour the Inc-1 graph-gate result. If capture is not + # ROCm parity: honour the graph-capture gate result. If capture is not # viable on this AMD card, skip graphs entirely so decode uses the kernel-launch # path (correct, just not graph-accelerated) rather than erroring mid-capture. from freetoken.utils.arch import is_rocm diff --git a/python/freetoken/kernel/backend.py b/python/freetoken/kernel/backend.py index 7b7629b1..8e685bbc 100644 --- a/python/freetoken/kernel/backend.py +++ b/python/freetoken/kernel/backend.py @@ -13,12 +13,6 @@ from freetoken.utils.arch import is_rocm -# NVIDIA-only optional native packages: even if an importable copy is present on a ROCm -# torch build (e.g. a stray CUDA wheel), they must not be used -- the runtime falls back -# to the portable Triton kernels. Treated as unavailable on ROCm. -_CUDA_ONLY_PACKAGES = frozenset({"flashinfer", "sgl_kernel", "triton_kernels"}) - - def _importable(name: str) -> bool: # find_spec normally returns None when a package is absent, but it can raise # (broken parent package, or a meta_path finder that blocks the name); treat diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index de28877c..1e063374 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -82,8 +82,7 @@ inline auto HIP_CHECK(::hipError_t error, std::source_location location = std::source_location::current()) -> void { - if (error != ::hipSuccess) { - [[unlikely]]; + if (error != ::hipSuccess) [[unlikely]] { ::host::panic(location, "HIP error: ", ::hipGetErrorString(error)); } } @@ -98,8 +97,7 @@ inline auto CUDA_CHECK(::cudaError_t error, std::source_location location = std::source_location::current()) -> void { - if (error != ::cudaSuccess) { - [[unlikely]]; + if (error != ::cudaSuccess) [[unlikely]] { ::host::panic(location, "CUDA error: ", ::cudaGetErrorString(error)); } } diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 13b0ea0c..00d6678d 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -56,9 +56,10 @@ def _module(): # the kernels compile their HIP branches; drop the CUDA-only -ccbin/flag logic. # Explicit --offload-arch (plus PYTORCH_ROCM_ARCH) prevents torch from auto- # emitting ~14 gfx arches, which would multiply build time per arch. - os.environ.setdefault("PYTORCH_ROCM_ARCH", "gfx1100") + gfx = os.getenv("FREETOKEN_KERNEL_CACHE_GFX", "gfx1100") + os.environ.setdefault("PYTORCH_ROCM_ARCH", gfx) extra_cuda_cflags = [ - "-O3", "--offload-arch=gfx1100", "-DUSE_HIP=1", "-DUSE_ROCM=1", + "-O3", f"--offload-arch={gfx}", "-DUSE_HIP=1", "-DUSE_ROCM=1", ] os.environ.pop("CXX", None) os.environ.pop("CC", None) diff --git a/python/freetoken/kernel/tinygrad_fallback.py b/python/freetoken/kernel/tinygrad_fallback.py deleted file mode 100644 index 2b0663fe..00000000 --- a/python/freetoken/kernel/tinygrad_fallback.py +++ /dev/null @@ -1,92 +0,0 @@ -"""tinygrad-JIT fallback for FFI kernels that are not hand-ported to HIP. - -FreeToken's hand-written tvm-ffi kernels (``store`` / ``index`` / ``fast_index_copy`` / -``batch_memcpy``) are CUDA source compiled via nvcc/JIT. The primary AMD port is the -``#if defined(USE_HIP)`` seam in ``device_api.h`` + ``LaunchKernel``/``warp.cuh``. This -module is the **documented fallback** for any kernel that proves intractable to hipify: -tinygrad's JIT compiles one logical kernel to PTX (CUDA) *and* AMDGPU/LLVM (ROCm), so the -same source covers both platforms. - -Constraints (matching the FFI contract): - -* Each fallback takes the same ``tvm.ffi.TensorView`` arguments as the hand-written - kernel and returns the same output tensor(s), so the swap is invisible to callers. -* It runs on the *host* (tinygrad handles GPU dispatch); on ROCm it compiles to AMDGPU. -* It is **never a default**: ``kernel/utils.py`` only routes a kernel to the fallback - when (a) ROCm is active and (b) the hand-HIP AOT/JIT variant is absent/unbuildable. - If tinygrad is not installed, invoking the fallback raises a clear error. - -Because tinygrad is an optional dependency (installed only when the fallback is actually -needed), all imports here are lazy and the module imports with zero third-party deps, so -it is safe to import on the CUDA-only path. -""" - -from __future__ import annotations - -from typing import Callable, Optional - -__all__ = [ - "is_tinygrad_available", - "kernel_fallback_available", - "get_kernel_fallback", -] - -# Kernel names the fallback registry knows how to build (mirrors the FFI kernel set). -_KNOWN_KERNELS = ("store", "index", "fast_index_copy", "batch_memcpy") - -#: Which kernels currently have a *functional* tinygrad reimplementation. As HIP ports -#: land in Inc 7, names are removed from this set (the hand port wins); kernels left here -#: (if any) are the documented fallback set. Default: empty -- the hand-HIP port is the -#: primary path and the fallback is opt-in per kernel. -_FALLBACK_IMPLEMENTED: set[str] = set() - - -def is_tinygrad_available() -> bool: - """True when the ``tinygrad`` package can be imported (JIT-to-ROCm available).""" - try: - import importlib.util # noqa: PLC0415 - - return importlib.util.find_spec("tinygrad") is not None - except Exception: - return False - - -def kernel_fallback_available(kernel: str) -> bool: - """True when a tinygrad fallback for ``kernel`` is both implemented and usable - (tinygrad installed). Always False on the CUDA path unless explicitly enabled, so - the CUDA build never depends on tinygrad.""" - if kernel not in _FALLBACK_IMPLEMENTED: - return False - return is_tinygrad_available() - - -def get_kernel_fallback(kernel: str): - """Return the tinygrad-backed fallback callable for ``kernel``, or raise a clear - error explaining why it is unavailable. Never called on the CUDA path.""" - if kernel not in _FALLBACK_IMPLEMENTED: - raise RuntimeError( - f"FFI kernel {kernel!r} has no tinygrad fallback registered. On ROCm the " - "preferred path is the hand-written HIP port (device_api.h); if you intend " - "to use the tinygrad fallback you must register it in " - "kernel/tinygrad_fallback.py._FALLBACK_IMPLEMENTED and implement the " - "corresponding build function." - ) - if not is_tinygrad_available(): - raise RuntimeError( - f"FFI kernel {kernel!r} requires the tinygrad fallback, but tinygrad is not " - "installed. Install it (`pip install tinygrad`) or provide a hand-written " - "HIP port for this kernel." - ) - from freetoken.kernel import tinygrad_impl # noqa: PLC0415 (lazy; may be None) - - builder = getattr(tinygrad_impl, f"build_{kernel}", None) - if builder is None: - raise RuntimeError( - f"tinygrad fallback for {kernel!r} is registered but has no " - "tinygrad_impl.build_() builder." - ) - return builder - - -def _list_fallbacks() -> list[str]: - return [k for k in _FALLBACK_IMPLEMENTED if kernel_fallback_available(k)] diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index 1beb7398..e04fc396 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -106,12 +106,6 @@ def _version_parts(version: str) -> Tuple[str, List[str]]: return base, local.split(".") if local else [] -def _build_stamps(segments: List[str]) -> set[str]: - """The `g` commit-stamp tokens of a local version segment list - (stamped by scripts/build-release-wheels.sh).""" - return {s for s in segments if re.fullmatch(r"g[0-9a-f]{7,40}", s)} - - def _build_stamps(local_segments: List[str]) -> List[str]: """The `g` commit-stamp tokens of a local version segment list (``["cu130", "g3f01615"]`` -> ``["g3f01615"]``).""" diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index 7b7f227c..9c208e60 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -77,7 +77,10 @@ def __init__( self._block_fp8 = expert_quant == "fp8_block" self._pertensor_fp8 = attn_quant == "fp8_pertensor" self._gguf = expert_quant == "gguf" - self._fp8 = self._block_fp8 or self._pertensor_fp8 or self._gguf + # "Split" projection layout (qkv|z + ba as two GEMMs): used by the fp8 paths + # and by GGUF (native-quant qkv|z + dense bf16 ba). The fused 4-way in_proj is + # only for the plain bf16 case. + self._split_proj = self._block_fp8 or self._pertensor_fp8 or self._gguf self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] if self._block_fp8 or self._pertensor_fp8: @@ -171,7 +174,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: fla = build_fla_metadata(batch, hidden_states.device) batch.fla_metadata = fla - if self._fp8: + if self._split_proj: qkvz = self.in_proj_qkvz.forward(hidden_states) conv_in, z = torch.split(qkvz, [self.conv_dim, self.value_dim], dim=-1) ba = self.in_proj_ba.forward(hidden_states) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index aef81a09..e6ffb04c 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -444,10 +444,6 @@ def iter_gguf_weights( # -------------------------------------------------------------------------------------- -def _q8_0_down_row_bytes(I: int) -> int: - return row_bytes(I, GGML_Q8_0) - - def load_gguf_expert_sources( model_path: str, config: ModelConfig, *, layer_sink=None ) -> dict[str, list[torch.Tensor]]: diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index e0801232..be1fd966 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -259,15 +259,12 @@ def _gguf_banks(model_path, model_config, device, dtype, dummy, parallel=False, "(not safetensors), so the common reader doesn't apply." ) if dummy: - from freetoken.models.weight import dummy_q4_0_moe_expert_sources - raise NotImplementedError("gguf expert banks have no dummy path; load the real GGUF") from freetoken.models.weight import load_gguf_moe_expert_sources - sink = None if dummy else layer_sink - sources = load_gguf_moe_expert_sources(model_path, model_config, layer_sink=sink) + sources = load_gguf_moe_expert_sources(model_path, model_config, layer_sink=layer_sink) return ExpertBanks( - "gguf", {name: sources[name] for name in _BANK_SCHEMAS["gguf"]}, streamed=sink is not None + "gguf", {name: sources[name] for name in _BANK_SCHEMAS["gguf"]}, streamed=layer_sink is not None ) diff --git a/python/freetoken/moe/nvfp4_to_mxfp4.py b/python/freetoken/moe/nvfp4_to_mxfp4.py index f4162a03..a81bf9e9 100644 --- a/python/freetoken/moe/nvfp4_to_mxfp4.py +++ b/python/freetoken/moe/nvfp4_to_mxfp4.py @@ -62,8 +62,6 @@ _FP4_TABLE = _np.asarray(_FP4_CODES, dtype=_np.float32) # Magnitudes sorted ascending for the nearest-code search. _FP4_SORT = _np.asarray(sorted(abs(v) for v in _FP4_CODES[1:8]), dtype=_np.float32) -_FP4_SORT_SIGN = _np.asarray([1.0 if i < 4 else -1.0 for i in range(len(_FP4_SORT))], - dtype=_np.float32) def fp4_e2m1_table() -> Sequence[float]: @@ -71,19 +69,6 @@ def fp4_e2m1_table() -> Sequence[float]: return list(_FP4_CODES) -def _nearest_e2m1_codes(values: _np.ndarray) -> _np.ndarray: - """Nearest e2m1 *code* for each fp32 ``values`` (signed, including 0/NaN).""" - a = _np.abs(values) - diff = _np.abs(a[..., None] - _FP4_SORT) # [..., 7] - idx = diff.argmin(axis=-1) - mag = _FP4_SORT[idx] - neg = _np.signbit(values) - code = (idx + 1).astype(_np.uint8) # _FP4_SORT[i] == table[i+1]; positive codes 1..7 - out = _np.where(neg, code | 0x8, code) - # Magnitudes below the smallest representable value (0.5) round to +0. - return _np.where(mag < 0.25, 0, out) - - def e8m0_scale_and_codes(values: _np.ndarray, block: int = 32) -> tuple[_np.ndarray, _np.ndarray]: """Return ``(scale_codes, fp4_codes)`` for ``values`` shaped ``[..., block]``: an e8m0 ``uint8`` scale per block (the smallest power-of-2 scale covering the block @@ -183,6 +168,10 @@ def convert_nvfp4_to_mxfp4( # Dequantize NVFP4 to fp32, move K to the last axis. K2 = packed.shape[-1] K = K2 * 2 + assert K % block == 0, ( + f"NVFP4 K={K} is not a multiple of the MXFP4 block ({block}); " + "cannot convert without truncating weights" + ) codes = np.stack([packed & 0x0F, (packed >> 4)], axis=-1).reshape( *packed.shape[:-1], K ) @@ -192,9 +181,8 @@ def convert_nvfp4_to_mxfp4( # Requantize to per-`block` e8m0 + e2m1. flat = f32.reshape(-1, K) - # pad to a multiple of block for the reshape (K is a multiple of 32 in practice) n_blocks = K // block - flat_b = flat[:, : n_blocks * block].reshape(-1, block) + flat_b = flat.reshape(-1, block) scale_codes, mxfp4_codes = e8m0_scale_and_codes(flat_b, block=block) out_codes = mxfp4_codes.reshape(flat.shape[0], n_blocks * block) diff --git a/python/freetoken/utils/graph_gate.py b/python/freetoken/utils/graph_gate.py index b528a208..07089ebf 100644 --- a/python/freetoken/utils/graph_gate.py +++ b/python/freetoken/utils/graph_gate.py @@ -1,10 +1,10 @@ """HIP/CUDA graph-capture parity probe. -The Inc-1 hard gate: whether ``torch.cuda.graph`` graph capture works on the target -GPU is the single highest-informational-risk assumption for AMD (ROCm) support. -This module probes it once and records a PASS/FAIL + device result that the rest -of the plan (Inc 8) reads. On CUDA it is expected to PASS; on ROCm it may fail on -some consumer cards, in which case Inc 8 must use the kernel-launch decode path. +Whether ``torch.cuda.graph`` graph capture works on the target GPU is the single +highest-informational-risk assumption for AMD (ROCm) support. This module probes it +once and records a PASS/FAIL + device result that the engine reads when deciding +whether to use CUDA-graph decode. On CUDA it is expected to PASS; on ROCm it may +fail on some consumer cards, in which case decode must use the kernel-launch path. The result is cached to disk under the user cache dir so it survives across runs, and keyed by device kind + device name so a change of GPU invalidates it. @@ -83,7 +83,6 @@ def probe_graph_capture() -> dict: # The child probes both an elementwise op (capturable on both backends) and a GEMM # (hipBLASLt on ROCm), which is what a real decode forward would run. The GEMM is # the discriminating case: on this ROCm build it fatally aborts -> child exit != 0. - import json as _json import subprocess as _subprocess import sys as _sys @@ -105,7 +104,7 @@ def probe_graph_capture() -> dict: "detail": f"fatal during capture: {detail[:240]}", } try: - data = _json.loads(child.stdout) + data = json.loads(child.stdout) except Exception: return { "device_kind": _device_kind(), @@ -177,7 +176,7 @@ def run_graph_gate() -> dict: @lru_cache(maxsize=1) def graph_capture_status() -> str: """Cached graph-capture status: ``"pass"``, ``"fail"``, or ``"unknown"`` (no device / - probe unavailable). Inc 8 reads this to pick HIP-graph vs kernel-launch decode.""" + probe unavailable). The graph runner reads this to pick HIP-graph vs kernel-launch decode.""" try: result = run_graph_gate() if result["ok"]: diff --git a/scripts/mirror-hf-configs.sh b/scripts/mirror-hf-configs.sh index 8cca566b..1ab821e7 100755 --- a/scripts/mirror-hf-configs.sh +++ b/scripts/mirror-hf-configs.sh @@ -6,12 +6,10 @@ # # ./mirror-hf-configs.sh [gguf-file] # -# ./mirror-hf-configs.sh Qwen/Qwen3.6-35B-A3B -# ./mirror-hf-configs.sh Qwen/Qwen3.6-35B-A3B /media/smk/Shared/Models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf +# ./mirror-hf-configs.sh Qwen/Qwen3.6-35B-A3B /path/to/model.gguf # -# With no gguf argument, defaults to $FT_MODEL, else the single *.gguf in -# /media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models or /media/smk/Shared/Models -# matching the repo's basename, else errors. +# With no gguf argument, uses $FT_MODEL if set; otherwise errors. The GGUF must be +# passed explicitly (or via FT_MODEL) -- no implicit model-directory search. # # Files fetched (when they exist upstream): # chat_template.jinja — read by FreeToken's GGUF loader (overrides embedded) @@ -29,12 +27,7 @@ if [ "$#" -ge 1 ]; then GGUF="$1" else GGUF="${FT_MODEL:-}" - if [ -z "$GGUF" ]; then - base="$(basename "${REPO_ID##*:}")" - cand="$(find /media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models /media/smk/Shared/Models -maxdepth 1 -iname "*${base%%-*}*.gguf" 2>/dev/null | head -1)" - [ -n "$cand" ] || { echo "ERROR: no gguf found; pass one explicitly" >&2; exit 1; } - GGUF="$cand" - fi + [ -n "$GGUF" ] || { echo "ERROR: no gguf file given; pass one explicitly or set FT_MODEL" >&2; exit 1; } fi [ -f "$GGUF" ] || { echo "ERROR: not a file: $GGUF" >&2; exit 1; } diff --git a/scripts/serve-qwen-moe.sh b/scripts/serve-qwen-moe.sh index 47de497c..c31cd286 100755 --- a/scripts/serve-qwen-moe.sh +++ b/scripts/serve-qwen-moe.sh @@ -31,7 +31,7 @@ set -euo pipefail REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" PY="${PY:-$REPO/.venv-rocm/bin/python}" -MODEL="${FT_MODEL:-/media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf}" +MODEL="${FT_MODEL:-}" HOST="${FT_HOST:-127.0.0.1}" PORT="${FT_PORT:-1920}" ATNN="${FT_ATTN:-triton}" # triton | torch @@ -94,6 +94,7 @@ server_pids() { } start() { + [ -n "$MODEL" ] || die "no model configured: set FT_MODEL=/path/to/model.gguf" if [ -n "$(server_pids)" ]; then echo "A server is already running (pid $(server_pids | tr '\n' ' ')). Use 'stop' first." exit 1 diff --git a/tests/attention/test_torch_backend.py b/tests/attention/test_torch_backend.py index f655a959..3a132777 100644 --- a/tests/attention/test_torch_backend.py +++ b/tests/attention/test_torch_backend.py @@ -1,4 +1,4 @@ -"""The debug ``"torch"`` attention backend (Inc 4 of fix-attention). +"""The debug ``"torch"`` attention backend. Verifies (a) the backend is registered and selectable, and (b) its pure-PyTorch GQA attention math (with causal masking and the per-head output gate) matches