diff --git a/.gitignore b/.gitignore index bf804e07..edc08af2 100644 --- a/.gitignore +++ b/.gitignore @@ -227,3 +227,13 @@ benchmarks/cross_framework # local e2e/bench artifacts (harnesses may run with repo cwd) /results/ + +# torch.utils.cpp_extension ROCm hipify build artifacts (preprocessed .cu -> .hip, +# and includes rewritten to *_hip.cuh) written next to the GGUF sources. +python/freetoken/kernel/csrc/gguf/*.hip +python/freetoken/kernel/csrc/gguf/*_hip.cuh +python/freetoken/kernel/csrc/gguf/ggml-common_hip.h + +# local session junk +nohup.out +.plans/ 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/docs/install-amd.md b/docs/install-amd.md new file mode 100644 index 00000000..ede4ab80 --- /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. + +## Requirements + +| Component | Requirement | +| --- | --- | +| OS | Linux x86_64 (Windows WDDM pinned-memory is a known edge, not supported yet) | +| GPU | AMD RX 7000 (RDNA 3, `gfx1100`); RX 9000 (`gfx1201`) is future work | +| ROCm | ROCm toolkit with `hipcc` (`/opt/rocm/bin/hipcc` or on `PATH`) | +| torch | ROCm build, e.g. `torch==2.5.1+rocm6.2` | + +The build refuses to mix toolchains: it will **not** silently fall back to `nvcc`/`libcudart` +when only the ROCm toolkit is present, and vice versa. + +## Install + +```bash +# ROCm torch (PyTorch official ROCm wheels) -- must satisfy the repo's torch>=2.11,<2.12 +# build pin, so use the rocm7.2 index (rocm6.2 only carries torch up to 2.5.1). +pip install --index-url https://download.pytorch.org/whl/rocm7.2 \ + "torch==2.11.0+rocm7.2" torchvision triton-rocm==3.6.0 + +# FreeToken with the ROCm extra (builds the native extensions with hipcc) +uv pip install -e ".[rocm]" --no-build-isolation +``` + +`pip install ".[rocm]"` pulls ROCm-compatible `torch`/`triton`; the NVIDIA-only `[accel]` +packages (`flashinfer`, `sgl-kernel`, `triton_kernels`, Marlin) are **not** installed on AMD +and their backends are rejected with a clean error if requested. + +## Verified feature matrix + +| Feature | On AMD | Notes | +| --- | --- | --- | +| Attention | `--attention-backend triton` | flashinfer/fa/trtllm are NVIDIA-only and rejected | +| MoE | `--moe-backend offload / cpu / hybrid` | offload needs pinned host memory | +| Quant | BF16, MXFP4, GGUF (Q4_K/Q8_0), Triton inline-dequant NVFP4 | Marlin INT4 / native NVFP4 SASS unavailable | +| NVFP4 checkpoints with no MXFP4 variant | converted to MXFP4 on load (auto) | `--nvfp4-backend auto` → triton/MXFP4 | +| CUDA graphs (decode) | HIP graph capture **if** the capture probe passes | otherwise kernel-launch decode | +| Multi-GPU (RCCL) | out of scope (single-GPU milestone) | | + +## CLI behavior on AMD + +* `--nvfp4-backend marlin` / `flashinfer` → error (NVIDIA-only). Use `triton` / `auto`. +* `--attention-backend fi` / `fa` / `trtllm` → error (NVIDIA-only). Use `triton` / `auto`. +* `--moe-backend fused` → warning (fused MoE is CUDA-only; falls back to offload/cpu). +* `--nvfp4-backend auto` → resolves to the portable Triton inline-dequant path (or MXFP4 + for a converted checkpoint). + +## Verify + +```bash +ft version # prints an AMD / ROCm banner +ft serve --model Qwen3.6-35B-A3B \ + --moe-backend offload --attention-backend triton --nvfp4-backend auto +``` + +`ldd` of the built `.so` should show `hiprt`/`amdhip64`, not `libcudart`. + +## AOT kernel cache + +Build the prebuilt `+rocm` kernel-cache wheel (no nvcc needed on the target): + +```bash +scripts/build-release-wheels.sh # on a ROCm torch + hipcc box; tags the cache +rocm +``` + +The runtime refuses to pair a `+rocm` cache with a `+cu130` runtime (and vice versa). + +## Notes / limitations + +* `nvtx_annotate` is a no-op on ROCm; roctx profiling is future work. +* FP8 / NVFP4-class formats: BF16 / MXFP4 / GGUF are the supported AMD matrix; performance + parity vs CUDA is not guaranteed for NVFP4-class formats. +* Windows AMD is not yet supported (WDDM zero-copy semantics differ). 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..2153c789 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+rocm7.2` from +# https://download.pytorch.org/whl/rocm7.2) so the native extensions build against the HIP +# runtime; this extra pins the rest. See docs/install-amd.md. +rocm = [ + "triton==3.6.0; platform_system == 'Linux'", +] # NOTE: the Marlin W4A16 NVFP4 expert-GEMM path (sm_80-99) borrows vLLM's AOT wheel # (vllm>=0.14,<0.15), which pins transformers>=4.56,<5 and so is INCOMPATIBLE with the # core transformers>=5.5 requirement. It is therefore not a lockable extra and is left 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..c19a9fcf --- /dev/null +++ b/python/freetoken/attention/torch.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, List + +import os +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. + self._contig: dict[tuple[int, int], list] = {} + 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/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 cd6505d2..e57cbdf1 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -16,7 +16,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 @@ -51,6 +51,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: @@ -102,9 +107,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. + if is_rocm(): + return all(not i.requires_flashinfer and not i.requires_sgl_kernel + 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(): @@ -203,6 +213,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 " @@ -300,6 +317,7 @@ def __init__(self, config: EngineConfig): self.device = bind_assigned_gpu(config.tp_info.rank) _adjust_config(config) + 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) @@ -928,7 +946,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) @@ -1023,6 +1041,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..276b741e 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] = {} + # 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 + 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/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/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..8e685bbc 100644 --- a/python/freetoken/kernel/backend.py +++ b/python/freetoken/kernel/backend.py @@ -10,6 +10,8 @@ import functools import importlib.util +from freetoken.utils.arch import is_rocm + def _importable(name: str) -> bool: # find_spec normally returns None when a package is absent, but it can raise @@ -21,13 +23,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 +60,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..1e063374 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,12 +77,27 @@ 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()) -> void { - if (error != ::cudaSuccess) { - [[unlikely]]; + if (error != ::cudaSuccess) [[unlikely]] { ::host::panic(location, "CUDA error: ", ::cudaGetErrorString(error)); } } @@ -75,7 +107,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 +133,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 +231,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..00d6678d 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -51,16 +51,29 @@ 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. + gfx = os.getenv("FREETOKEN_KERNEL_CACHE_GFX", "gfx1100") + os.environ.setdefault("PYTORCH_ROCM_ARCH", gfx) + extra_cuda_cflags = [ + "-O3", f"--offload-arch={gfx}", "-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 +82,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/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..e04fc396 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] @@ -80,10 +106,18 @@ def _version_parts(version: str) -> Tuple[str, List[str]]: return base, local.split(".") if local else [] -def _build_stamps(segments: List[str]) -> set[str]: +def _build_stamps(local_segments: List[str]) -> List[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)} + (``["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: @@ -93,14 +127,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 +171,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 +246,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 +274,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 +299,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 +336,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..83d8be84 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -10,10 +10,84 @@ 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"} +_TOKENIZER_ARCH = { + "gemma4": "gemma4_text", + "qwen35moe": "qwen3_moe", +} + + +# 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): @@ -46,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/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..9c208e60 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -76,10 +76,14 @@ 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" + # "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._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 +91,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) @@ -161,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 new file mode 100644 index 00000000..e6ffb04c --- /dev/null +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -0,0 +1,508 @@ +"""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 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..df0d41f0 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -109,8 +109,28 @@ 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) + 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/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..be1fd966 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -252,6 +252,22 @@ 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: + raise NotImplementedError("gguf expert banks have no dummy path; load the real GGUF") + from freetoken.models.weight import load_gguf_moe_expert_sources + + 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=layer_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 +317,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..a81bf9e9 --- /dev/null +++ b/python/freetoken/moe/nvfp4_to_mxfp4.py @@ -0,0 +1,227 @@ +"""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) + + +def fp4_e2m1_table() -> Sequence[float]: + """The 16 e2m1 values keyed by 4-bit code (index == code).""" + return list(_FP4_CODES) + + +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 + 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 + ) + 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) + n_blocks = K // 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) + + 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/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 48923e3b..bd35d373 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -414,9 +414,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 a71b6819..afe5cceb 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -150,6 +150,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" @@ -598,6 +604,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", @@ -711,6 +728,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/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/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..07089ebf --- /dev/null +++ b/python/freetoken/utils/graph_gate.py @@ -0,0 +1,188 @@ +"""HIP/CUDA graph-capture parity probe. + +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. +""" +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 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). The graph runner 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/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..1ab821e7 --- /dev/null +++ b/scripts/mirror-hf-configs.sh @@ -0,0 +1,60 @@ +#!/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 /path/to/model.gguf +# +# 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) +# 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:-}" + [ -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; } + +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..c31cd286 --- /dev/null +++ b/scripts/serve-qwen-moe.sh @@ -0,0 +1,213 @@ +#!/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:-}" +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 +# 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}" +# 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 + 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() { + [ -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 + 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" + "--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" + "--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}${MOE_CACHE:+ }$( [ "$MOE_CACHE" = auto ] && echo "kv-reserve $KV_TOKENS" || echo 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 + # 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" +} + +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 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..3a132777 --- /dev/null +++ b/tests/attention/test_torch_backend.py @@ -0,0 +1,62 @@ +"""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 +``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_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 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"