diff --git a/README.md b/README.md index 2a56a086..d2dfaf3e 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ FreeToken is an edge-native Mixture-of-Experts (MoE) serving engine designed for - **Semantic-Aware Caching**: Features semantic anchor checkpoints for recurrent state and KV caches, allowing agentic context edits (e.g., tool calls, thinking blocks) to avoid redundant context recomputation. - **Elastic Memory Management**: Supports dynamic, runtime VRAM re-allocation between expert caches and KV memory without engine restarts or weight reloading. - **Broad MoE & Ecosystem Support**: Supports frontier open-weight MoE models (e.g., DeepSeek-V4-Flash, Qwen3.6-35B-A3B, GLM-5.2) across various parameter scales and quantization formats (e.g., MXFP4, NVFP4, FP8, BF16), with Anthropic/OpenAI-compatible APIs for seamless integration with real-world coding and tool-calling agents (e.g., Codex, Claude Code, OpenCode, OpenClaw, DeepSeek Harness). -- **Diverse Consumer Hardware**: Scales across consumer laptops, gaming desktops, and workstation GPUs, with native support for NVIDIA RTX 30, RTX 40, and RTX 50 series GPUs. +- **Diverse Consumer Hardware**: Scales across consumer laptops, gaming desktops, and workstation GPUs, with native support for NVIDIA RTX 30/40/50 GPUs and experimental ROCm source support for AMD RDNA3/RDNA4 GPUs. ## Getting Started diff --git a/benchmarks/bench_decode_moe.py b/benchmarks/bench_decode_moe.py index 723c2be5..56621792 100644 --- a/benchmarks/bench_decode_moe.py +++ b/benchmarks/bench_decode_moe.py @@ -100,6 +100,8 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: help="hybrid: max PCIe fetches/layer; -1 = auto (benched pcie/cpu bandwidth fraction)", ) p.add_argument("--mem-ratio", type=float, default=0.9, help="target VRAM utilization") + p.add_argument("--gpu", default=None, + help="GPU for the serve: a UUID or nvidia-smi index (as ft serve --gpu)") p.add_argument("--no-graph", action="store_true", help="eager decode instead of CUDA graph") p.add_argument( "--greedy", @@ -183,6 +185,8 @@ def serve_cmd(args: argparse.Namespace, backend: str, port: int) -> list[str]: "--cuda-graph-max-bs", "0" if args.no_graph else "1", "--moe-hybrid-max-fetch", str(args.hybrid_fetch), ] + if args.gpu: + cmd += ["--gpu", args.gpu] if args.cache > 0: cmd += ["--moe-cache-size", str(args.cache)] elif args.cache_rate is not None: diff --git a/benchmarks/bench_load_weight_generic.py b/benchmarks/bench_load_weight_generic.py index 0ef4bd4c..0e1905a0 100644 --- a/benchmarks/bench_load_weight_generic.py +++ b/benchmarks/bench_load_weight_generic.py @@ -140,8 +140,10 @@ def _model_config(model_path: str): if try_get_tp_info() is None: set_tp_info(rank=0, size=1) - torch.cuda.set_device(0) - torch.zeros(1, device="cuda") # init CUDA context (pinning / nvfp4 backend pick) + from freetoken.gpu_select import bind_assigned_gpu + + dev = bind_assigned_gpu() + torch.zeros(1, device=dev) # init CUDA context (pinning / nvfp4 backend pick) cfg = EngineConfig(model_path=model_path, tp_info=DistributedInfo(0, 1), dtype=torch.bfloat16, moe_backend="offload") return cfg.model_config @@ -176,7 +178,7 @@ def _bench_load(mode: str, model_path: str, *, parallel: bool, workers: int, chu s.start() t = time.perf_counter() try: - banks = load_expert_banks(model_path, mc, device=torch.device("cuda:0"), + banks = load_expert_banks(model_path, mc, device=torch.device("cuda", torch.cuda.current_device()), dtype=torch.bfloat16, parallel=parallel, workers=workers, chunk=chunk) except NotImplementedError as e: @@ -220,7 +222,8 @@ def worker_build(ns): s = MemSampler() s.start() t = time.perf_counter() - idx = convert_checkpoint(ns.model, ns.ftw_dir, moe_backend="offload", shard_limit=shard_limit) + dev = f"cuda:{torch.cuda.current_device()}" if ns.gpu else None + idx = convert_checkpoint(ns.model, ns.ftw_dir, moe_backend="offload", shard_limit=shard_limit, device=dev) build_s = time.perf_counter() - t s.stop() print("@@RESULT@@" + json.dumps({ @@ -237,6 +240,8 @@ def _spawn(worker: str, ns): cmd = [sys.executable, os.path.abspath(__file__), "--_worker", worker, "--model", ns.model, "--workers", str(ns.workers), "--chunk-mib", str(ns.chunk_mib), "--shard-gib", str(ns.shard_gib), "--ftw-dir", ns.ftw_dir] + if ns.gpu: + cmd += ["--gpu", ns.gpu] if ns.no_drop_cache: cmd.append("--no-drop-cache") # Capture ONLY stdout (the @@RESULT@@ line); let stderr inherit the terminal so the @@ -260,6 +265,10 @@ def main(): p.add_argument("--no-drop-cache", action="store_true", help="don't evict page cache before each read (warm comparison)") p.add_argument("--keep-ftw", action="store_true", help="keep + reuse the FTW dir across runs") + from freetoken.gpu_select import single_gpu_arg + + p.add_argument("--gpu", type=single_gpu_arg, default=None, + help="GPU UUID or nvidia-smi index (default: the first visible GPU)") p.add_argument("--_worker", default="") ns = p.parse_args() @@ -268,7 +277,17 @@ def main(): if not ns.ftw_dir: ns.ftw_dir = _default_ftw_dir(ns.model) + from freetoken.gpu_select import assign_gpu + + try: + assign_gpu(ns.gpu) + except ValueError as e: + p.error(str(e)) + if ns._worker: + from freetoken.gpu_select import bind_assigned_gpu + + bind_assigned_gpu() return _WORKERS[ns._worker](ns) from freetoken.checkpoint.ftw import is_ftw_checkpoint diff --git a/benchmarks/bench_offload_cache_copy.py b/benchmarks/bench_offload_cache_copy.py index a2434a66..8374510f 100644 --- a/benchmarks/bench_offload_cache_copy.py +++ b/benchmarks/bench_offload_cache_copy.py @@ -16,6 +16,7 @@ import torch +from freetoken.gpu_select import assign_gpu, bind_assigned_gpu, single_gpu_arg from freetoken.moe.offload_cache import _BANK_SCHEMAS, OffloadMoeCache @@ -80,7 +81,8 @@ def expert_bytes(profile: ModelProfile) -> int: def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() - parser.add_argument("--device", type=int, default=0) + parser.add_argument("--gpu", type=single_gpu_arg, default=None, + help="GPU UUID or nvidia-smi index (default: the first visible GPU)") parser.add_argument("--repeat", type=int, default=25) parser.add_argument("--models", type=str, nargs="+", default=list(MODELS), choices=list(MODELS)) parser.add_argument( @@ -228,8 +230,11 @@ def print_table( def main() -> None: args = parse_args() assert torch.cuda.is_available(), "CUDA is required" - torch.cuda.set_device(args.device) - device = torch.device("cuda") + try: + assign_gpu(args.gpu) + device = bind_assigned_gpu() + except (ValueError, RuntimeError) as e: + raise SystemExit(f"error: {e}") from e print("gpu", torch.cuda.get_device_name(device), flush=True) for name in args.models: diff --git a/docs/cli.md b/docs/cli.md index cf4b27a2..ff4af382 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -40,6 +40,7 @@ parsers all resolve automatically from the checkpoint and the GPU. |---|---|---| | `--host` | 127.0.0.1 | Bind address | | `--port` | 1919 | Bind port | +| `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) | | `--max-running-requests` | 4 | Max concurrently running requests | | `--max-output-tokens` | 32768 | Default output budget for requests that omit one | | `--max-seq-len-override` | from checkpoint | Max sequence length | @@ -47,6 +48,21 @@ parsers all resolve automatically from the checkpoint and the GPU. | `--cuda-graph-max-bs`, `--graph` | = max running requests | Max batch size captured as CUDA graphs | | `--decode-log-interval` | 40 | Scheduler status line every N decode steps | +### Choosing a GPU + +For example, a machine with an RTX 5090 and an RTX 3060 Ti: + +```console +$ nvidia-smi -L +GPU 0: NVIDIA GeForce RTX 3060 Ti (UUID: GPU-2f3a9b1c-8d7e-4a05-b6c1-0e5f9a3d7b42) +GPU 1: NVIDIA GeForce RTX 5090 (UUID: GPU-9e8d7c6b-5a49-4f13-8207-c1b0a4e6d3f5) +``` + +```bash +ft serve --model ... --gpu 1 # by nvidia-smi index -- the 5090 +ft serve --model ... --gpu GPU-9e8d7c6b # the same card by UUID (a unique prefix is enough) +``` + ### KV cache & memory | Flag | Default | Meaning | @@ -130,7 +146,7 @@ environment so the agent cannot silently fall back to a paid endpoint. ## ft checkpoint ```bash -ft checkpoint --model --out [--dtype bfloat16] [--moe-backend offload] [--shard-gib 8] [--device cuda:0] +ft checkpoint --model --out [--dtype bfloat16] [--moe-backend offload] [--shard-gib 8] [--gpu ] ``` Converts an HF safetensors checkpoint to FTW, FreeToken's self-contained @@ -142,15 +158,18 @@ keeps them dense for resident serving. See the FTW caveats in ## ft bench bw ```bash -ft bench bw # once per machine +ft bench bw # once per GPU ft bench bw --dtype nvfp4,bf16 # only the formats you serve +ft bench bw --gpu 1 # a specific GPU (UUID or nvidia-smi index, as for ft serve) ``` -Measures host-RAM vs PCIe bandwidth with the real cpu/offload MoE kernels and -writes a profile (`~/.cache/freetoken/benchbw.json`) that `ft serve ---moe-backend auto` and `--moe-hybrid-max-fetch -1` read. Profiles are keyed on -expert format + GPU name, so a profile from different hardware is ignored -rather than misapplied. Selection flags: `--dtype`, `--model`, `--formats`, -`--isa`; decision rule: `--threshold` (default 2.0 — recommend hybrid when CPU -bandwidth > 2× PCIe). +Measures host-RAM vs PCIe bandwidth with the real cpu/offload MoE kernels and writes a +profile that `ft serve --moe-backend auto` and `--moe-hybrid-max-fetch -1` then read. + +- One profile per GPU, at `~/.cache/freetoken/benchbw/.json`. +- Keyed on expert format + GPU, so a profile from other hardware is ignored rather than + misapplied. An older single `benchbw.json` still counts if its GPU name matches. +- What to measure: `--dtype`, `--model`, `--formats`, `--isa`. +- `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe + by that factor. diff --git a/docs/install.md b/docs/install.md index f5205ab3..850558a1 100644 --- a/docs/install.md +++ b/docs/install.md @@ -2,7 +2,9 @@ ## Requirements -- Linux x86_64, NVIDIA GPU, driver r580+ (CUDA 13) +- Linux x86_64 with either: + - NVIDIA GPU, driver r580+ (CUDA 13), or + - AMD RDNA3/RDNA4 GPU (`gfx1100`-`gfx1103`, `gfx1200`, or `gfx1201`) with ROCm 7.14 - Python >= 3.10, with [uv](https://docs.astral.sh/uv/) recommended (plain `pip` + `venv` works too) @@ -15,6 +17,40 @@ uv pip install "freetoken[accel]" CUDA kernels are JIT-compiled on first use, need a CUDA 13 toolkit with `nvcc` on PATH. +### AMD ROCm source install (experimental) + +Use an official ROCm PyTorch image whose PyTorch version satisfies the project's +`torch>=2.11,<2.12` constraint. For RDNA4, the matching ROCm 7.14 image is: + +```bash +VIDEO_GID="$(getent group video | cut -d: -f3)" +RENDER_GID="$(getent group render | cut -d: -f3)" +docker run --rm -it \ + --device=/dev/kfd --device=/dev/dri \ + --group-add="$VIDEO_GID" --group-add="$RENDER_GID" --ipc=host \ + --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ + -e PYTORCH_ROCM_ARCH=gfx1201 -e FREETOKEN_ROCM_ARCH=gfx1201 \ + -v "$PWD:/workspace/FreeToken" -w /workspace/FreeToken \ + rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.11.0 bash +``` + +Inside the container, preserve the ROCm-enabled PyTorch already supplied by the +image and disable build isolation so it is also used to compile the extensions: + +```bash +python -m pip install --no-build-isolation -e . +``` + +Set both architecture variables to `gfx1200` for RX 9060 family GPUs, or to the +actual target reported by `rocminfo`. + +The optional native GGUF kernels also require Thrust headers. Install the generic +headers before the first GGUF kernel JIT build: + +```bash +apt-get update && apt-get install -y --no-install-recommends libthrust-dev +``` + ## Method 2: Install from source ```bash diff --git a/pyproject.toml b/pyproject.toml index 8bd653f8..d22ae67c 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", @@ -57,7 +58,10 @@ dependencies = [ "torch>=2.11,<2.12", "tqdm>=4.66,<5", "transformers>=5.5,<6", - "triton==3.6.0; platform_system == 'Linux'", + # CUDA torch 2.11 resolves Triton 3.6; AMD's ROCm 7.14 image supplies its + # gfx1201-enabled Triton 3.7 build. Keep both supported without replacing the + # runtime-specific wheel selected by the PyTorch distribution. + "triton>=3.6,<3.8; platform_system == 'Linux'", "uvicorn>=0.30,<1", ] diff --git a/python/freetoken/checkpoint/__main__.py b/python/freetoken/checkpoint/__main__.py index cd1f90a5..2786173f 100644 --- a/python/freetoken/checkpoint/__main__.py +++ b/python/freetoken/checkpoint/__main__.py @@ -1,7 +1,7 @@ """CLI: convert an HF safetensors checkpoint to a FreeToken Weight (FTW) checkpoint. ft checkpoint --model --out \ - [--dtype bfloat16] [--moe-backend offload] [--shard-gib 8] + [--dtype bfloat16] [--moe-backend offload] [--shard-gib 8] [--gpu ] The output dir is self-contained: point the server's ``--model`` at it to load via the FTW fast path (auto-detected). @@ -14,6 +14,8 @@ import torch +from freetoken.gpu_select import assign_gpu, bind_assigned_gpu, single_gpu_arg + from .convert import convert_checkpoint _DTYPES = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32} @@ -27,15 +29,24 @@ def main(argv: list[str] | None = None, prog: str = "freetoken.checkpoint") -> i p.add_argument("--moe-backend", default="offload", help="offload (experts -> banks) or e.g. triton (experts stay dense)") p.add_argument("--shard-gib", type=float, default=8.0, help="max shard size in GiB") - p.add_argument("--device", default=None, help="CUDA device for repack (default cuda:0)") + p.add_argument("--gpu", type=single_gpu_arg, default=None, + help="GPU for the repack: a GPU UUID (GPU-xxxx..., as nvidia-smi -L prints) or " + "an nvidia-smi index (default: the first visible GPU)") ns = p.parse_args(argv) + # same as ft serve --gpu: resolve, then bind by UUID at CUDA init + try: + assign_gpu(ns.gpu) + device = f"cuda:{bind_assigned_gpu().index}" + except (ValueError, RuntimeError) as e: + p.error(str(e)) + shard_limit = int(ns.shard_gib * (1 << 30)) shard_limit -= shard_limit % 4096 # keep aligned t = time.perf_counter() index = convert_checkpoint( ns.model, ns.out, dtype=_DTYPES[ns.dtype], - moe_backend=ns.moe_backend, shard_limit=shard_limit, device=ns.device, + moe_backend=ns.moe_backend, shard_limit=shard_limit, device=device, ) dt = time.perf_counter() - t c = index["counts"] diff --git a/python/freetoken/daemon/app.py b/python/freetoken/daemon/app.py index e230ccf3..d7a0a53c 100644 --- a/python/freetoken/daemon/app.py +++ b/python/freetoken/daemon/app.py @@ -58,16 +58,41 @@ class BenchBody(BaseModel): args: list[str] = [] -def _bench_profile_path() -> str: - from freetoken.moe.bench_profile import default_profile_path # torch-free +def _bench_profile_path(gpu_uuid: str | None) -> str | None: + # per-GPU profiles and no torch here: the serve's own card when its --gpu names one, else the newest file + from freetoken.moe.bench_profile import default_profile_path, latest_profile_path # torch-free - return default_profile_path() + if gpu_uuid: + path = default_profile_path(gpu_uuid) + if os.path.isfile(path): + return path + return latest_profile_path() -def _read_bench_profile() -> dict | None: - """The engine host's cached benchbw.json (this is where the serve reads it too), or None.""" +def _serve_gpu_uuid(args: list[str]) -> str | None: + """The full UUID a serve's `--gpu` pins, or None when there is none or it cannot be resolved.""" + for i, a in enumerate(args): + val = a[len("--gpu="):] if a.startswith("--gpu=") else (args[i + 1] if a == "--gpu" and i + 1 < len(args) else None) + if not val: + continue + from freetoken.gpu_select import resolve_gpu_uuids + + try: + resolved = resolve_gpu_uuids([val]) + except ValueError: + return None + if resolved: + return resolved[0] + # no NVML: a UUID value still keys the profile file (canonical prefix), an index cannot + return "GPU-" + val[len("GPU-"):] if val.upper().startswith("GPU-") else None + return None + + +def _read_bench_profile(path: str | None) -> dict | None: + if path is None: + return None try: - with open(_bench_profile_path()) as f: + with open(path) as f: return json.load(f) except (OSError, ValueError): return None @@ -331,19 +356,23 @@ async def gen(): yield _bench_sse("error", {"message": f"failed to spawn bench: {exc}"}) return tail: collections.deque = collections.deque(maxlen=8) # last non-progress lines (errors) + out_path: str | None = None assert proc.stdout is not None async for raw in proc.stdout: line = raw.decode(errors="replace").rstrip() prog = _parse_ftbench(line) if prog is not None: yield _bench_sse("progress", prog) + elif line.startswith("FTBENCH_OUT "): + out_path = line[len("FTBENCH_OUT "):] elif line: tail.append(line) rc = await proc.wait() if rc != 0: yield _bench_sse("error", {"message": "\n".join(tail) or f"bench exited {rc}"}) return - prof = _read_bench_profile() + # the file this run wrote (an older engine prints no FTBENCH_OUT: newest file, as before) + prof = _read_bench_profile(out_path or _bench_profile_path(None)) if prof is None: yield _bench_sse("error", {"message": "bench finished but no profile was written"}) else: @@ -353,7 +382,23 @@ async def gen(): @app.get("/bench/profile", dependencies=auth) async def bench_profile(): - return await run(proxy_pool, _read_bench_profile) + def read() -> dict | None: + return _read_bench_profile(_bench_profile_path(serve_gpu_uuid())) + + def serve_gpu_uuid() -> str | None: + # the running serve reports the full UUID of its card (/v1/stats gpus); a --gpu given as + # a UUID prefix would not match the profile file name + st = manager.status() + if st.get("running"): + try: + gpus = probe.stats(st.get("port") or default_serve_port).get("gpus") or [] + if gpus and gpus[0].get("uuid"): + return gpus[0]["uuid"] + except Exception: # noqa: BLE001 -- the arg below is the fallback + pass + return _serve_gpu_uuid(manager.serve_args()) + + return await run(proxy_pool, read) return app diff --git a/python/freetoken/daemon/serve_manager.py b/python/freetoken/daemon/serve_manager.py index 78cea11f..681bdef4 100644 --- a/python/freetoken/daemon/serve_manager.py +++ b/python/freetoken/daemon/serve_manager.py @@ -862,6 +862,11 @@ def current_pid(self) -> int | None: with self._cond: return self._child.pid if self._child is not None else None + def serve_args(self) -> list[str]: + """Engine args of the running (or last started) serve.""" + with self._cond: + return list(self._args) + def status(self) -> dict: with self._cond: child = self._child diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index c386e28b..7a67f73f 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -10,6 +10,7 @@ from freetoken.attention import AttnType, attention_backend_info, create_attention_backend from freetoken.core import Batch, Context, Req, set_global_ctx from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info +from freetoken.gpu_select import gpu_identity from freetoken.layers import set_rope_device from freetoken.models import create_model, load_weight from freetoken.moe import create_moe_backend, is_offload_moe_backend @@ -294,10 +295,11 @@ def __init__(self, config: EngineConfig): assert not torch.cuda.is_initialized() set_tp_info(rank=config.tp_info.rank, size=config.tp_info.size) _ensure_expandable_segments() # before the first CUDA allocation below - _adjust_config(config) - self.device = torch.device(f"cuda:{config.tp_info.rank}") - torch.cuda.set_device(self.device) + from freetoken.gpu_select import bind_assigned_gpu + + self.device = bind_assigned_gpu(config.tp_info.rank) + _adjust_config(config) torch.manual_seed(42) self.stream = torch.cuda.Stream() torch.cuda.set_stream(self.stream) @@ -423,7 +425,17 @@ def __init__(self, config: EngineConfig): self._warmup_prefill() def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: - if config.tp_info.size == 1 or config.use_pynccl: + use_pynccl = config.use_pynccl + if config.tp_info.size > 1 and use_pynccl: + from freetoken.kernel.backend import is_rocm + + if is_rocm(): + logger.warning_rank0( + "PyNCCL is NVIDIA-only; using PyTorch's ROCm/RCCL process group instead" + ) + use_pynccl = False + + if config.tp_info.size == 1 or use_pynccl: torch.distributed.init_process_group( backend="gloo", rank=config.tp_info.rank, @@ -650,8 +662,10 @@ def _resolve_hybrid_fetch(self, config: EngineConfig, cache) -> None: return # explicit fixed cap from freetoken.moe.bench_profile import load_hybrid_fetch_fraction - gpu_name = torch.cuda.get_device_name(self.device) if torch.cuda.is_available() else None - fraction = load_hybrid_fetch_fraction(cache.quant_format, gpu_name=gpu_name) + gpu_name, gpu_uuid = _profile_gpu(self.device.index) + fraction = load_hybrid_fetch_fraction( + cache.quant_format, gpu_name=gpu_name, gpu_uuid=gpu_uuid + ) if fraction is None: cache.hybrid_max_fetch = 1 logger.warning_rank0( @@ -994,6 +1008,14 @@ def shutdown(self) -> None: destroy_distributed() +def _profile_gpu(index: "int | None" = None) -> Tuple[str | None, str | None]: + """(name, uuid) of visible device ``index`` (default: the current, i.e. bound, device); (None, None) without CUDA.""" + if not torch.cuda.is_available(): + return None, None + ident = gpu_identity(torch.cuda.current_device() if index is None else index) + return ident["name"], ident["uuid"] + + def _ensure_expandable_segments() -> None: """Default the CUDA allocator to expandable segments. @@ -1374,8 +1396,8 @@ def override(attr: str, value: Any): # this is dangerous, use with caution bench_fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16") from freetoken.moe.bench_profile import load_backend_recommendation - gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else None - if load_backend_recommendation(bench_fmt, gpu_name=gpu_name) == "hybrid": + gpu_name, gpu_uuid = _profile_gpu() + if load_backend_recommendation(bench_fmt, gpu_name=gpu_name, gpu_uuid=gpu_uuid) == "hybrid": from freetoken.moe.cpu_executor import compiled_extension_supports _act = getattr(model_config, "hidden_act", "silu") diff --git a/python/freetoken/gpu_select.py b/python/freetoken/gpu_select.py new file mode 100644 index 00000000..2c7d012b --- /dev/null +++ b/python/freetoken/gpu_select.py @@ -0,0 +1,287 @@ +"""--gpu for ft serve / bench bw / checkpoint: resolve entries to GPU UUIDs, bind by UUID at CUDA init. + +Three device-id namespaces, converted explicitly: +- logical: position in the --gpu list == TP rank; each worker takes its own entry by rank and the id ends there. +- physical: NVML / nvidia-smi order, carried as a GPU UUID (_assigned_physical); not affected by CUDA_VISIBLE_DEVICES. +- visible: CUDA ordinal in this process (_assigned_visible), what torch.device("cuda", n) means. + +The parent resolves --gpu entries to full UUIDs via NVML (resolve_gpu_uuids) and fails fast on a typo. +Each worker publishes its own entry (set_assigned_gpu / assign_gpu) and binds it when CUDA comes up (bind_assigned_gpu) by matching the UUID against CUDA's visible devices. +One process runs on one GPU. Binding is unconditional: a process that publishes nothing binds a default ordinal and records it, so assigned_visible_gpu() names that card in every case. +No process mutates CUDA_VISIBLE_DEVICES, and the UUID match holds under any CUDA_DEVICE_ORDER. + +Stdlib only (torch is imported lazily); not under freetoken.utils, which imports transformers. +""" + +from __future__ import annotations + +import argparse +import os +from typing import Sequence + +UUID_PREFIX = "GPU-" + + +def is_gpu_uuid(spec: str) -> bool: + return spec[: len(UUID_PREFIX)].upper() == UUID_PREFIX + + +def is_gpu_index(spec: str) -> bool: + # not str.isdigit(): that also accepts superscripts and other Unicode digits + return spec.isascii() and spec.isdecimal() + + +def _canonical(entry: str) -> str: + """A UUID in the exact form the driver matches (upper-case GPU- prefix), an index as-is.""" + if not (is_gpu_uuid(entry) or is_gpu_index(entry)): + raise ValueError( + f"{entry!r} is neither a GPU UUID (GPU-xxxx..., as `nvidia-smi -L` prints) " + f"nor an nvidia-smi index" + ) + return UUID_PREFIX + entry[len(UUID_PREFIX):] if is_gpu_uuid(entry) else entry + + +def parse_gpu_spec(value: str) -> tuple[str, ...]: + """Split a --gpu value; ValueError on a bad entry, an empty value, or a mix of UUIDs and indices.""" + entries = tuple(_canonical(e.strip()) for e in value.split(",") if e.strip()) + if not entries: + raise ValueError("--gpu needs at least one GPU") + if len({is_gpu_uuid(e) for e in entries}) > 1: + # the driver parses CUDA_VISIBLE_DEVICES as all-UUID or all-index + raise ValueError("--gpu entries must be all UUIDs or all indices") + return entries + + +def gpu_arg(value: str) -> tuple[str, ...]: + """argparse type for a --gpu list.""" + try: + return parse_gpu_spec(value) + except ValueError as exc: + raise argparse.ArgumentTypeError(str(exc)) from exc + + +def single_gpu_arg(value: str) -> str: + """argparse type for a single-GPU --gpu.""" + entries = gpu_arg(value) + if len(entries) != 1: + raise argparse.ArgumentTypeError("takes exactly one GPU") + return entries[0] + + +def _nvml_uuids() -> "list[str] | None": + """Full GPU UUIDs in physical (nvidia-smi) order, or None when NVML is unavailable. + + Own ctypes loader instead of torch's _raw_device_uuid_nvml: that helper only knows the Linux library name, raises (not None) when the library is missing, and is private API. + NVML exports are cdecl on every platform, so CDLL is right on Windows too (same as nvidia-ml-py). + None on any failure -- no library, a stub library without the _v2 symbols, WSL, a dead device -- and callers fall back. + """ + import ctypes + + if os.name == "nt": + candidates = [ + "nvml.dll", + os.path.join(os.environ.get("SystemRoot", r"C:\\Windows"), "System32", "nvml.dll"), + os.path.join(os.environ.get("ProgramFiles", r"C:\\Program Files"), "NVIDIA Corporation", "NVSMI", "nvml.dll"), + ] + else: + candidates = ["libnvidia-ml.so.1"] + try: + for name in candidates: + try: + lib = ctypes.CDLL(name) + break + except OSError: + continue + else: + return None + if lib.nvmlInit() != 0: + return None + try: + count = ctypes.c_int() + if lib.nvmlDeviceGetCount_v2(ctypes.byref(count)) != 0: + return None + uuids = [] + for i in range(count.value): + handle = ctypes.c_void_p() + if lib.nvmlDeviceGetHandleByIndex_v2(i, ctypes.byref(handle)) != 0: + return None + buf = ctypes.create_string_buffer(96) + if lib.nvmlDeviceGetUUID(handle, buf, 96) != 0: + return None + uuids.append(buf.value.decode("ascii", "replace")) + return uuids + finally: + lib.nvmlShutdown() + except (OSError, AttributeError): + return None + + +def _match_uuid(spec: str, uuids: "list[str]", where: str) -> str: + """The unique full UUID that ``spec`` prefixes, else ValueError.""" + hits = [u for u in uuids if u.upper().startswith(spec.upper())] + if len(hits) != 1: + raise ValueError(f"--gpu {spec}: not found or not a unique prefix {where}; run `nvidia-smi -L` to list GPUs") + return hits[0] + + +def resolve_gpu_uuids(specs: Sequence[str]) -> "tuple[str, ...] | None": + """--gpu entries -> full GPU UUIDs, one per TP rank; raises ValueError on a bad entry. + + A preset CUDA_VISIBLE_DEVICES is a quota to stay inside: an index counts within that list, a UUID must name one of its entries. + Returns None when NVML is unavailable -- the worker then interprets the raw entries against CUDA's own enumeration (see bind_assigned_gpu). + """ + specs = parse_gpu_spec(",".join(specs)) + if len({s.upper() for s in specs}) != len(specs): + raise ValueError(f"--gpu {','.join(specs)}: the same GPU appears twice") + uuids = _nvml_uuids() + if uuids is None: + return None + preset_raw = os.environ.get("CUDA_VISIBLE_DEVICES") + preset = None if preset_raw is None else [e.strip() for e in preset_raw.split(",") if e.strip()] + + resolved: list[str] = [] + for spec in specs: + if preset is None: + if is_gpu_uuid(spec): + resolved.append(_match_uuid(spec, uuids, "on this machine")) + elif int(spec) < len(uuids): + resolved.append(uuids[int(spec)]) + else: + raise ValueError(f"--gpu {spec}: only {len(uuids)} GPU(s) on this machine; run `nvidia-smi -L` to list GPUs") + else: + entry = _preset_entry(spec, preset, preset_raw) + # an integer entry is read in physical order, as under CUDA_DEVICE_ORDER=PCI_BUS_ID; a negative or MIG-form entry cannot name a whole GPU + if is_gpu_uuid(entry): + resolved.append(_match_uuid(entry, uuids, f"(from CUDA_VISIBLE_DEVICES={preset_raw!r})")) + elif is_gpu_index(entry) and int(entry) < len(uuids): + resolved.append(uuids[int(entry)]) + else: + raise ValueError( + f"--gpu {spec}: cannot resolve CUDA_VISIBLE_DEVICES entry {entry!r} " + f"({len(uuids)} GPU(s) on this machine)" + ) + if len(set(resolved)) != len(resolved): + raise ValueError(f"--gpu {','.join(specs)}: the same GPU appears twice") + return tuple(resolved) + + +def _preset_entry(spec: str, preset: "list[str]", preset_raw: str) -> str: + """The CUDA_VISIBLE_DEVICES entry ``spec`` selects, else ValueError.""" + if not is_gpu_uuid(spec): + idx = int(spec) + if idx >= len(preset): + raise ValueError( + f"--gpu {spec}: only {len(preset)} GPU(s) are visible through " + f"CUDA_VISIBLE_DEVICES={preset_raw!r} (indices count within that list)" + ) + return preset[idx] + if not all(is_gpu_uuid(p) for p in preset): + raise ValueError( + f"--gpu {spec}: CUDA_VISIBLE_DEVICES={preset_raw!r} lists GPUs by index; " + f"give --gpu as an index into that list" + ) + hits = [p for p in preset if p.upper().startswith(spec.upper()) or spec.upper().startswith(p.upper())] + if len(hits) != 1: + raise ValueError( + f"--gpu {spec}: not one of the GPUs visible through CUDA_VISIBLE_DEVICES={preset_raw!r}" + ) + return hits[0] + + +# The GPU this process was assigned, in whichever namespace it arrived in; bind_assigned_gpu fills in the visible one. +# Process-global on purpose: publishing is torch-free so a worker can do it before heavy imports, and kernel-compat checks (e4m3_native) need the device this process will use. +_assigned_physical: "str | None" = None +_assigned_visible: "int | None" = None + + +def set_assigned_gpu(target: str) -> None: + """Publish this process's GPU before CUDA init; second call must agree. + + A UUID names a physical GPU and is converted at bind time; a bare index is already a visible ordinal (a preset CUDA_VISIBLE_DEVICES has narrowed to it). + """ + global _assigned_physical, _assigned_visible + physical = target if is_gpu_uuid(target) else None + visible = None if physical is not None else int(target) + current = (_assigned_physical, _assigned_visible) + if current not in ((None, None), (physical, visible)): + raise RuntimeError(f"set_assigned_gpu called twice: {current} then {target!r}") + _assigned_physical, _assigned_visible = physical, visible + + +def assign_gpu(spec: "str | None") -> None: + """Resolve one --gpu value and publish it for bind_assigned_gpu; no-op when the flag was not given.""" + if spec is None: + return + resolved = resolve_gpu_uuids([spec]) + set_assigned_gpu(resolved[0] if resolved else parse_gpu_spec(spec)[0]) + + +def _visible_of_physical(uuid: str) -> int: + """CUDA ordinal of the physical GPU ``uuid`` (or unique prefix) among this process's visible devices.""" + import torch + + seen: list[str] = [] + hits: list[int] = [] + for v in range(torch.cuda.device_count()): + u = format_gpu_uuid(getattr(torch.cuda.get_device_properties(v), "uuid", None)) + seen.append(u or "?") + if u is not None and u.upper().startswith(uuid.upper()): + hits.append(v) + if len(hits) == 1: + return hits[0] + if hits: + raise RuntimeError(f"--gpu {uuid}: not a unique prefix (visible: {', '.join(seen)})") + raise RuntimeError( + f"GPU {uuid} is not visible to CUDA in this process " + f"(CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')!r}, " + f"visible: {', '.join(seen) or 'none'})" + ) + + +def bind_assigned_gpu(default: int = 0): + """torch.cuda.set_device this process's GPU and return the device. + + ``default`` is a visible ordinal, used and recorded when nothing was published, so the process always knows which card it runs on. + A published UUID (or prefix) is matched against CUDA's own device list, so the result is right under any CUDA_DEVICE_ORDER. + """ + global _assigned_visible + import torch + + if _assigned_visible is None: + _assigned_visible = default if _assigned_physical is None else _visible_of_physical(_assigned_physical) + if not 0 <= _assigned_visible < torch.cuda.device_count(): + raise RuntimeError( + f"cannot use CUDA device {_assigned_visible}: only {torch.cuda.device_count()} device(s) visible " + f"(CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')!r})" + ) + device = torch.device("cuda", _assigned_visible) + torch.cuda.set_device(device) + return device + + +def assigned_visible_gpu() -> "int | None": + """Visible ordinal this process is pinned to, or None before it publishes or binds a GPU (= the current device). + + Published-but-not-yet-bound still counts: compat checks in the window between publish and bind must judge the assigned card, not whatever the calling thread happens to sit on. + """ + if _assigned_visible is not None: + return _assigned_visible + return None if _assigned_physical is None else _visible_of_physical(_assigned_physical) + + +def format_gpu_uuid(raw) -> str | None: + """nvidia-smi form GPU- from a uuid.UUID.""" + return None if raw is None else f"{UUID_PREFIX}{raw}" + + +def gpu_identity(index: int) -> dict: + """{index, name, uuid, total_bytes} of visible device ``index``.""" + import torch + + props = torch.cuda.get_device_properties(index) + return { + "index": index, + "name": props.name, + "uuid": format_gpu_uuid(getattr(props, "uuid", None)), + "total_bytes": int(props.total_memory), + } diff --git a/python/freetoken/kernel/__main__.py b/python/freetoken/kernel/__main__.py index 7be541a6..5b66d484 100644 --- a/python/freetoken/kernel/__main__.py +++ b/python/freetoken/kernel/__main__.py @@ -6,27 +6,32 @@ def generate_clangd(): import subprocess from freetoken.kernel.utils import DEFAULT_INCLUDE - from freetoken.utils import init_logger + from freetoken.utils import get_rocm_gfx_arch, init_logger, is_rocm from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path logger = init_logger(__name__) logger.info("Generating .clangd file...") include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE - status = subprocess.run( - args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], - capture_output=True, - check=True, - ) - compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] - major, minor = compute_cap.split(".") + + # TODO(ROCm): hiprtc JIT cache should be separate from nvcc JIT cache to avoid stale binaries. + if is_rocm(): + arch_flags = ["-xhip", f"--offload-arch={get_rocm_gfx_arch() or 'gfx1201'}"] + else: + try: + status = subprocess.run( + args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], + capture_output=True, + check=True, + ) + compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] + major, minor = compute_cap.split(".") + except (subprocess.CalledProcessError, FileNotFoundError, ValueError): + import torch + + major, minor = torch.cuda.get_device_capability() + arch_flags = ["-xcuda", f"--cuda-gpu-arch=sm_{major}{minor}"] compile_flags = ",\n ".join( - [ - "-xcuda", - f"--cuda-gpu-arch=sm_{major}{minor}", - "-std=c++20", - "-Wall", - "-Wextra", - ] + arch_flags + ["-std=c++20", "-Wall", "-Wextra"] + [f"-isystem{path}" for path in include_paths] ) clangd_content = f""" diff --git a/python/freetoken/kernel/_toolchain.py b/python/freetoken/kernel/_toolchain.py index b49cebbb..e4065bf8 100644 --- a/python/freetoken/kernel/_toolchain.py +++ b/python/freetoken/kernel/_toolchain.py @@ -1,4 +1,4 @@ -"""CUDA toolchain/torch consistency checks. +"""CUDA/HIP toolchain/torch consistency checks. Standalone on purpose: setup.py and the kernel-cache build backend load this file by path, so it must not import the freetoken package. @@ -14,6 +14,29 @@ ALLOW_MISMATCH_ENV = "FREETOKEN_ALLOW_CUDA_MISMATCH" _TRUE_VALUES = {"1", "true", "yes", "on"} +_AMDHIP64_VERSIONED_RE = re.compile(r"libamdhip64\.so\.(\d+(?:\.\d+)*)$") + + +def select_versioned_rocm_runtime(paths): + """Return the highest numeric libamdhip64 soname from ``paths``. + + Path/string lexical order is not a version order: for example ``7.9`` sorts + after ``7.14``. Keep this helper package-independent so both setup.py and + runtime JIT discovery can use the same selection contract. + """ + candidates = [] + for path in paths: + match = _AMDHIP64_VERSIONED_RE.fullmatch(os.path.basename(os.fspath(path))) + if match is None: + continue + version = tuple(int(part) for part in match.group(1).split(".")) + candidates.append((version, os.fspath(path), path)) + return max(candidates, key=lambda item: (item[0], item[1]))[2] if candidates else None + + +def _is_rocm() -> bool: + import torch + return getattr(torch.version, "hip", None) is not None def _nvcc_path() -> str | None: @@ -49,6 +72,8 @@ def check_nvcc_matches_torch() -> None: nvcc-built binaries link libcudart.so.; at runtime only the torch wheel's own CUDA runtime is guaranteed to be loadable. """ + if _is_rocm(): + return # ROCm uses hipcc, not nvcc if os.getenv(ALLOW_MISMATCH_ENV, "").strip().lower() in _TRUE_VALUES: return torch_major = torch_cuda_major() diff --git a/python/freetoken/kernel/backend.py b/python/freetoken/kernel/backend.py index 3037ad8d..7ddbf9f1 100644 --- a/python/freetoken/kernel/backend.py +++ b/python/freetoken/kernel/backend.py @@ -23,11 +23,15 @@ def _importable(name: str) -> bool: @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,9 +43,29 @@ 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") +@functools.cache +def is_rocm() -> bool: + """True when torch is built for ROCm (AMD GPU).""" + import torch + return getattr(torch.version, "hip", None) is not None + + +@functools.cache +def driver_hip_version() -> int | None: + """ROCm driver version, or None if undetermined.""" + # TODO(ROCm): flashinfer/sgl_kernel have no ROCm builds — Triton fallback is used. + try: + from freetoken.kernel.pinned import _load_pinned_extension + return int(_load_pinned_extension().driver_cuda_version()) or None + except Exception: + return None + + @functools.cache def driver_cuda_version() -> int | None: """Max CUDA version the installed NVIDIA driver supports (``13000`` == CUDA 13.0), @@ -50,6 +74,8 @@ def driver_cuda_version() -> int | None: toolkit version. Resolved through the ``_pinned_tensor`` extension's link-time cudart, so it works wherever the extension builds (including Windows) -- no dlopen by soname.""" + if is_rocm(): + return None try: from freetoken.kernel.pinned import _load_pinned_extension 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 48210b9d..054d7fce 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -29,7 +29,7 @@ #include #include -#include +#include #include #if defined(__linux__) diff --git a/python/freetoken/kernel/csrc/gguf/dispatch.h b/python/freetoken/kernel/csrc/gguf/dispatch.h index f42a2163..bb15096b 100644 --- a/python/freetoken/kernel/csrc/gguf/dispatch.h +++ b/python/freetoken/kernel/csrc/gguf/dispatch.h @@ -5,18 +5,28 @@ #pragma once #include +#include #ifndef WARP_SIZE #define WARP_SIZE 32 #endif -// Warp-shuffle wrappers the donor pulls from sgl-kernel's utils.h (CUDA variants). +// HIP's synchronized shuffle API requires a 64-bit mask even on wave32 targets +// such as gfx1201, while CUDA uses a 32-bit mask. +#if defined(__HIP_PLATFORM_AMD__) +#define SGLANG_SHUFFLE_MASK(mask) static_cast(mask) +#else +#define SGLANG_SHUFFLE_MASK(mask) (mask) +#endif + +// Warp-shuffle wrappers the donor pulls from sgl-kernel's utils.h. #ifndef SGLANG_SHFL_XOR_SYNC -#define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) __shfl_xor_sync((mask), (var), (lane_mask)) +#define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) \ + __shfl_xor_sync(SGLANG_SHUFFLE_MASK(mask), (var), (lane_mask)) #endif #ifndef SGLANG_SHFL_XOR_SYNC_WIDTH #define SGLANG_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \ - __shfl_xor_sync((mask), (var), (lane_mask), (width)) + __shfl_xor_sync(SGLANG_SHUFFLE_MASK(mask), (var), (lane_mask), (width)) #endif #define DISPATCH_CASE_FLOAT_TYPES(...) \ diff --git a/python/freetoken/kernel/csrc/gguf/gguf_dequant_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_dequant_kernel.cu new file mode 100644 index 00000000..3b9d5fc6 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/gguf_dequant_kernel.cu @@ -0,0 +1,40 @@ +// ROCm operation-split binding for the vendored GGUF dequant kernels. +// Kernel implementations remain in the sgl-kernel/llama.cpp-derived headers. +#include +#include +#include +#include + +#include "dispatch.h" +#include "ggml-common.h" +#include "dequantize.cuh" + +torch::Tensor ggml_dequantize( + torch::Tensor W, + int64_t type, + int64_t m, + int64_t n, + std::optional const& dtype) { + const at::cuda::OptionalCUDAGuard 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(); + + DISPATCH_FLOAT_TYPES(DW.scalar_type(), "ggml_dequantize", [&] { + auto to_cuda = ggml_get_to_cuda(type); + TORCH_CHECK( + to_cuda != nullptr, + "ggml_dequantize: unsupported GGUF quant type ", type, + " (dequant kernels exist for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K/IQ2_XXS/" + "IQ2_XS/IQ3_XXS/IQ1_S/IQ4_NL/IQ3_S/IQ2_S/IQ4_XS/IQ1_M)"); + to_cuda((void*)W.data_ptr(), (scalar_t*)DW.data_ptr(), m * n, stream); + }); + + return DW; +} + +#include +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ggml_dequantize", &ggml_dequantize, ""); +} diff --git a/python/freetoken/kernel/csrc/gguf/gguf_mmq_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_mmq_kernel.cu new file mode 100644 index 00000000..666087b0 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/gguf_mmq_kernel.cu @@ -0,0 +1,57 @@ +// ROCm operation-split binding for GGUF large-batch MMQ. +#include +#include +#include +#include + +#include "dispatch.h" +#include "ggml-common.h" +#include "vecdotq.cuh" +#include "mmq.cuh" +#include "quantize_q8_1.cuh" + +torch::Tensor ggml_mul_mat_a8( + torch::Tensor W, + torch::Tensor X, + int64_t type, + int64_t row) { + 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)); + 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(); + 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", [&] { + quantize_row_q8_1_cuda( + (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, batch, stream); + using Fn = void (*)(const void*, const void*, scalar_t*, int, int, int, int, int, cudaStream_t); + Fn fn = nullptr; + switch (type) { + case 2: fn = &ggml_mul_mat_q4_0_q8_1_cuda; break; + case 3: fn = &ggml_mul_mat_q4_1_q8_1_cuda; break; + case 6: fn = &ggml_mul_mat_q5_0_q8_1_cuda; break; + case 7: fn = &ggml_mul_mat_q5_1_q8_1_cuda; break; + case 8: fn = &ggml_mul_mat_q8_0_q8_1_cuda; break; + case 10: fn = &ggml_mul_mat_q2_K_q8_1_cuda; break; + case 11: fn = &ggml_mul_mat_q3_K_q8_1_cuda; break; + case 12: fn = &ggml_mul_mat_q4_K_q8_1_cuda; break; + case 13: fn = &ggml_mul_mat_q5_K_q8_1_cuda; break; + case 14: fn = &ggml_mul_mat_q6_K_q8_1_cuda; break; + default: + TORCH_CHECK(false, "ggml_mul_mat_a8: unsupported GGUF quant type ", type, + " (MMQ kernels exist only for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K; " + "I-quants must route through ggml_dequantize)"); + } + fn(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + col, row, batch, padded, row, stream); + }); + return Y; +} + +#include +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ggml_mul_mat_a8", &ggml_mul_mat_a8, ""); +} diff --git a/python/freetoken/kernel/csrc/gguf/gguf_mmvq_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_mmvq_kernel.cu new file mode 100644 index 00000000..8ea3f818 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/gguf_mmvq_kernel.cu @@ -0,0 +1,102 @@ +// ROCm operation-split binding for GGUF small-batch MMVQ. +#include +#include +#include +#include + +#include "dispatch.h" +#include "ggml-common.h" +#include "vecdotq.cuh" +#include "mmvq.cuh" +#include "quantize_q8_1.cuh" + +torch::Tensor ggml_mul_mat_vec_a8( + torch::Tensor W, + torch::Tensor X, + int64_t type, + int64_t row) { + 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)); + 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(); + 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", [&] { + quantize_row_q8_1_cuda( + (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, vecs, stream); + switch (type) { + case 2: + mul_mat_vec_q4_0_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 3: + mul_mat_vec_q4_1_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 6: + mul_mat_vec_q5_0_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 7: + mul_mat_vec_q5_1_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 8: + mul_mat_vec_q8_0_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 10: + mul_mat_vec_q2_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 11: + mul_mat_vec_q3_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 12: + mul_mat_vec_q4_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 13: + mul_mat_vec_q5_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 14: + mul_mat_vec_q6_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 16: + mul_mat_vec_iq2_xxs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 17: + mul_mat_vec_iq2_xs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 18: + mul_mat_vec_iq3_xxs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 19: + mul_mat_vec_iq1_s_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 20: + mul_mat_vec_iq4_nl_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 21: + mul_mat_vec_iq3_s_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 22: + mul_mat_vec_iq2_s_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 23: + mul_mat_vec_iq4_xs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + case 29: + mul_mat_vec_iq1_m_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); + break; + default: + TORCH_CHECK( + false, + "ggml_mul_mat_vec_a8: unsupported GGUF quant type ", type, + " (MMVQ kernels exist for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K/IQ2_XXS/IQ2_XS/" + "IQ3_XXS/IQ1_S/IQ4_NL/IQ3_S/IQ2_S/IQ4_XS/IQ1_M)"); + } + }); + return Y; +} + +#include +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ggml_mul_mat_vec_a8", &ggml_mul_mat_vec_a8, ""); +} diff --git a/python/freetoken/kernel/csrc/gguf/gguf_moe_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_moe_kernel.cu new file mode 100644 index 00000000..f562f45e --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/gguf_moe_kernel.cu @@ -0,0 +1,87 @@ +// ROCm operation-split binding for GGUF grouped large-batch MoE kernels. +#include +#include +#include +#include + +#include "dispatch.h" +#include "ggml-common.h" +#include "vecdotq.cuh" +#include "mmq.cuh" +#include "moe.cuh" +#include "quantize_q8_1.cuh" + +torch::Tensor ggml_moe_a8( + torch::Tensor X, + torch::Tensor W, + torch::Tensor sorted_token_ids, + torch::Tensor expert_ids, + torch::Tensor num_tokens_post_padded, + int64_t type, + int64_t row, + int64_t top_k, + int64_t tokens) { + int col = X.sizes()[1]; + int padded = (col + 512 - 1) / 512 * 512; + const at::cuda::OptionalCUDAGuard 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(); + 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", [&] { + quantize_row_q8_1_cuda( + (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, tokens, stream); + using Fn = void (*)( + const void*, const void*, scalar_t*, const int*, const int*, const int*, + int, int, int, int, int, int, int, int, cudaStream_t); + Fn fn = nullptr; + switch (type) { + case 2: fn = &ggml_moe_q4_0_q8_1_cuda; break; + case 3: fn = &ggml_moe_q4_1_q8_1_cuda; break; + case 6: fn = &ggml_moe_q5_0_q8_1_cuda; break; + case 7: fn = &ggml_moe_q5_1_q8_1_cuda; break; + case 8: fn = &ggml_moe_q8_0_q8_1_cuda; break; + case 10: fn = &ggml_moe_q2_K_q8_1_cuda; break; + case 11: fn = &ggml_moe_q3_K_q8_1_cuda; break; + case 12: fn = &ggml_moe_q4_K_q8_1_cuda; break; + case 13: fn = &ggml_moe_q5_K_q8_1_cuda; break; + case 14: fn = &ggml_moe_q6_K_q8_1_cuda; break; + default: + TORCH_CHECK(false, "ggml_moe_a8: unsupported GGUF quant type ", type, + " (MMQ kernels exist only for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K; " + "I-quants must route through ggml_dequantize)"); + } + fn(quant_X.data_ptr(), W.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)sorted_token_ids.data_ptr(), (int*)expert_ids.data_ptr(), + (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row, tokens, + padded, row, top_k, sorted_token_ids.sizes()[0], stream); + }); + return Y; +} + +int64_t ggml_moe_get_block_size(int64_t type) { + switch (type) { + case 2: return MOE_X_Q4_0; + case 3: return MOE_X_Q4_1; + case 6: return MOE_X_Q5_0; + case 7: return MOE_X_Q5_1; + case 8: return MOE_X_Q8_0; + case 10: return MOE_X_Q2_K; + case 11: return MOE_X_Q3_K; + case 12: return MOE_X_Q4_K; + case 13: return MOE_X_Q5_K; + case 14: return MOE_X_Q6_K; + default: + TORCH_CHECK(false, "ggml_moe_get_block_size: unsupported GGUF quant type ", type, + " (MMQ kernels exist only for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K; " + "I-quants must route through ggml_dequantize)"); + return 0; + } +} + +#include +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ggml_moe_a8", &ggml_moe_a8, ""); + m.def("ggml_moe_get_block_size", &ggml_moe_get_block_size, ""); +} diff --git a/python/freetoken/kernel/csrc/gguf/gguf_moe_vec_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_moe_vec_kernel.cu new file mode 100644 index 00000000..ddf6fa8f --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/gguf_moe_vec_kernel.cu @@ -0,0 +1,121 @@ +// ROCm operation-split binding for GGUF routed small-batch MoE vector kernels. +#include +#include +#include +#include + +#include "dispatch.h" +#include "ggml-common.h" +#include "vecdotq.cuh" +#include "moe_vec.cuh" +#include "quantize_q8_1.cuh" + +torch::Tensor ggml_moe_a8_vec( + torch::Tensor X, + torch::Tensor W, + torch::Tensor topk_ids, + int64_t top_k, + int64_t type, + int64_t row, + 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)); + 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(); + 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", [&] { + quantize_row_q8_1_cuda( + (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, tokens, stream); + switch (type) { + case 2: + moe_vec_q4_0_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 3: + moe_vec_q4_1_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 6: + moe_vec_q5_0_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 7: + moe_vec_q5_1_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 8: + moe_vec_q8_0_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 10: + moe_vec_q2_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 11: + moe_vec_q3_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 12: + moe_vec_q4_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 13: + moe_vec_q5_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 14: + moe_vec_q6_K_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 16: + moe_vec_iq2_xxs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 17: + moe_vec_iq2_xs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 18: + moe_vec_iq3_xxs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 19: + moe_vec_iq1_s_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 20: + moe_vec_iq4_nl_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 21: + moe_vec_iq3_s_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 22: + moe_vec_iq2_s_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 23: + moe_vec_iq4_xs_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + case 29: + moe_vec_iq1_m_q8_1_cuda(W.data_ptr(), quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, quant_X.stride(0), stream); + break; + default: + TORCH_CHECK(false, "ggml_moe_a8_vec: unsupported GGUF quant type ", type, + " (MMVQ kernels exist for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K/IQ2_XXS/IQ2_XS/" + "IQ3_XXS/IQ1_S/IQ4_NL/IQ3_S/IQ2_S/IQ4_XS/IQ1_M)"); + } + }); + return Y; +} + +#include +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("ggml_moe_a8_vec", &ggml_moe_a8_vec, ""); +} diff --git a/python/freetoken/kernel/csrc/gguf/quantize_q8_1.cuh b/python/freetoken/kernel/csrc/gguf/quantize_q8_1.cuh new file mode 100644 index 00000000..a41ce6cd --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/quantize_q8_1.cuh @@ -0,0 +1,71 @@ +// Extracted from gguf_kernel.cu so ROCm can JIT-compile GGUF operations in +// independent translation units. The implementation is unchanged from the +// vendored sgl-kernel/llama.cpp-derived FreeToken GGUF path. +#pragma once + +#include + +// Q8 activation quantization used by MMVQ/MMQ and grouped MoE kernels. +template +static __global__ void quantize_q8_1( + const scalar_t* __restrict__ x, + void* __restrict__ vy, + const int kx, + const int kx_padded) { + const auto ix = blockDim.x * blockIdx.x + threadIdx.x; + if (ix >= kx_padded) { + return; + } + const auto iy = blockDim.y * blockIdx.y + threadIdx.y; + const int i_padded = iy * kx_padded + ix; + + block_q8_1* y = (block_q8_1*)vy; + + const int ib = i_padded / QK8_1; + const int iqs = i_padded % QK8_1; + + const float xi = ix < kx ? static_cast(x[iy * kx + ix]) : 0.0f; + float amax = fabsf(xi); + float sum = xi; + +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + amax = fmaxf(amax, SGLANG_SHFL_XOR_SYNC_WIDTH(uint32_t(-1), amax, mask, 32)); + sum += SGLANG_SHFL_XOR_SYNC_WIDTH(uint32_t(-1), sum, mask, 32); + } + + const float d = amax / 127; + const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); + + y[ib].qs[iqs] = q; + + if (iqs > 0) { + return; + } + + y[ib].ds.x = __float2half(d); + y[ib].ds.y = __float2half(sum); +} + +template +static void quantize_row_q8_1_cuda( + const scalar_t* x, + void* vy, + const int kx, + const int ky, + cudaStream_t stream) { + const int64_t kx_padded = (kx + 512 - 1) / 512 * 512; + const int block_num_x = + (kx_padded + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE; + constexpr int MAX_BLOCK_SIZE = 65535; + for (int off = 0; off < ky; off += MAX_BLOCK_SIZE) { + const int num_blocks_y = std::min(ky, off + MAX_BLOCK_SIZE) - off; + const dim3 num_blocks(block_num_x, num_blocks_y, 1); + const dim3 block_size(CUDA_DEQUANTIZE_BLOCK_SIZE, 1, 1); + quantize_q8_1<<>>( + &x[off * kx], + (int32_t*)vy + off * (kx_padded / 32 * 9), + kx, + kx_padded); + } +} diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h new file mode 100644 index 00000000..3cb5feaf --- /dev/null +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -0,0 +1,160 @@ +#pragma once + +// HIP compatibility shim: maps CUDA runtime API names to HIP equivalents so +// the same C++ source compiles under both nvcc and hipcc. Include this instead +// of directly when the file needs the runtime API. +// +// On NVIDIA platforms the CUDA headers are included as-is and every macro below +// resolves to the original CUDA symbol, so there is zero overhead. +// +// Supported ROCm targets: +// gfx1100 — RX 7900 XTX / XT +// gfx1101 — RX 7900 GRE +// gfx1102 — RX 7700 / XT +// gfx1103 — RX 7600 / XT +// gfx1200 — RX 9060 family +// gfx1201 — RX 9070 family / Radeon AI PRO R9700 + +#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM) + +#define FREETOKEN_USE_ROCM 1 + +// --- HIP runtime headers --- +#include +#include + +// --- API name mapping (CUDA -> HIP) --- +// HIP already defines most cuda* names as macros that expand to hip* equivalents +// via hip_runtime.h, but a few are missing or differ in signature. Define them +// here so call-sites stay unchanged. + +#ifndef cudaSuccess +#define cudaSuccess hipSuccess +#endif + +#ifndef cudaError_t +#define cudaError_t hipError_t +#endif + +#ifndef cudaGetErrorString +#define cudaGetErrorString hipGetErrorString +#endif + +#ifndef cudaGetLastError +#define cudaGetLastError hipGetLastError +#endif + +#ifndef cudaMallocHost +#define cudaMallocHost hipMallocHost +#endif + +#ifndef cudaFreeHost +#define cudaFreeHost hipFreeHost +#endif + +#ifndef cudaHostAlloc +#define cudaHostAlloc hipHostMalloc +#endif + +#ifndef cudaHostRegister +#define cudaHostRegister hipHostRegister +#endif + +#ifndef cudaHostRegisterPortable +#define cudaHostRegisterPortable hipHostRegisterPortable +#endif + +#ifndef cudaHostRegisterMapped +#define cudaHostRegisterMapped hipHostRegisterMapped +#endif + +#ifndef cudaHostAllocPortable +#define cudaHostAllocPortable hipHostMallocPortable +#endif + +#ifndef cudaHostAllocMapped +#define cudaHostAllocMapped hipHostMallocMapped +#endif + +#ifndef cudaHostGetDevicePointer +#define cudaHostGetDevicePointer hipHostGetDevicePointer +#endif + +#ifndef cudaGetDevice +#define cudaGetDevice hipGetDevice +#endif + +#ifndef cudaDriverGetVersion +#define cudaDriverGetVersion hipDriverGetVersion +#endif + +#ifndef cudaDeviceGetAttribute +#define cudaDeviceGetAttribute hipDeviceGetAttribute +#endif + +#ifndef cudaDevAttrUnifiedAddressing +#define cudaDevAttrUnifiedAddressing hipDeviceAttributeUnifiedAddressing +#endif + +#ifndef cudaDevAttrCanUseHostPointerForRegisteredMem +#define cudaDevAttrCanUseHostPointerForRegisteredMem hipDeviceAttributeCanUseHostPointerForRegisteredMem +#endif + +#ifndef cudaFuncSetAttribute +#define cudaFuncSetAttribute hipFuncSetAttribute +#endif + +#ifndef cudaFuncAttributeMaxDynamicSharedMemorySize +#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize +#endif + +#ifndef cudaLaunchKernelEx +// ROCm 7 exposes the CUDA-compatible extended launch configuration through HIP. +#define cudaLaunchKernelEx hipLaunchKernelEx +#endif + +#ifndef cudaLaunchConfig_t +#define cudaLaunchConfig_t hipLaunchConfig_t +#endif + +#ifndef cudaLaunchAttribute +#define cudaLaunchAttribute hipLaunchAttribute +#endif + +#ifndef cudaLaunchAttributeProgrammaticStreamSerialization +// PDL (Programmatic Dependent Launch) is NVIDIA-specific. +// TODO(ROCm): PDL has no ROCm equivalent — disabled, may affect overlap scheduling latency. +#define cudaLaunchAttributeProgrammaticStreamSerialization 0 +#endif + +#ifndef cudaStream_t +#define cudaStream_t hipStream_t +#endif + +#ifndef cudaStreamSynchronize +#define cudaStreamSynchronize hipStreamSynchronize +#endif + +#ifndef cudaLaunchHostFunc +#define cudaLaunchHostFunc hipLaunchHostFunc +#endif + +#ifndef CUDART_CB +#define CUDART_CB +#endif + +#ifndef __grid_constant__ +#define __grid_constant__ +#endif + +#ifndef dim3 +// HIP already provides dim3; this is a no-op guard. +#endif + +#else // NVIDIA CUDA path + +#define FREETOKEN_USE_ROCM 0 + +#include + +#endif diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 8e917832..f46d2af2 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -44,13 +45,21 @@ namespace PDL { template __always_inline __device__ void wait() { if constexpr (kUsePDL) { +#if FREETOKEN_USE_ROCM + // Programmatic dependent launch is NVIDIA-specific. +#else asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif } } template __always_inline __device__ void launch() { if constexpr (kUsePDL) { +#if FREETOKEN_USE_ROCM + // Programmatic dependent launch is NVIDIA-specific. +#else asm volatile("griddepcontrol.launch_dependents;" :::); +#endif } } @@ -115,6 +124,10 @@ public: } auto with_attr(bool use_pdl) -> LaunchKernel & { +#if FREETOKEN_USE_ROCM + RuntimeCheck(!use_pdl, "Programmatic dependent launch is unavailable on ROCm"); + m_config.numAttrs = 0; +#else if (use_pdl) { m_attr_cache.id = ::cudaLaunchAttributeProgrammaticStreamSerialization; m_attr_cache.val.programmaticStreamSerializationAllowed = 1; @@ -123,6 +136,7 @@ public: } else { m_config.numAttrs = 0; } +#endif return *this; } @@ -138,7 +152,9 @@ private: return config; } cudaLaunchConfig_t m_config; +#if !FREETOKEN_USE_ROCM cudaLaunchAttribute m_attr_cache; +#endif }; } // namespace host diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23e..bf313c52 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 FREETOKEN_USE_ROCM + 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 FREETOKEN_USE_ROCM + 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 FREETOKEN_USE_ROCM + 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 FREETOKEN_USE_ROCM + *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 FREETOKEN_USE_ROCM + *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 FREETOKEN_USE_ROCM + *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,7 @@ __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 !FREETOKEN_USE_ROCM && __CUDA_ARCH__ >= 700 __nanosleep(sleep_ns); #endif sleep_ns = sleep_ns < 2048 ? (sleep_ns << 1) : 2048; @@ -147,7 +171,7 @@ inline bool host_ptr_identity() { } inline void* device_alias(void* ptr, DLDevice dev) { - if (dev.device_type == kDLCUDA || host_ptr_identity()) { + if (dev.device_type == kDLCUDA || dev.device_type == kDLROCM || host_ptr_identity()) { return ptr; } void* mapped = nullptr; @@ -269,7 +293,7 @@ inline auto get_sync_flag_ptr( auto flag_dtype = host::SymbolicDType{}; host::TensorMatcher({1}) .with_dtype(flag_dtype) - .with_device(device) + .with_device(device) .verify(sync_flag); return static_cast(sync_flag.data_ptr()); } @@ -344,17 +368,17 @@ 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}) .with_dtype(indices_dtype) - .with_device(device) + .with_device(device) .verify(src_indices) .verify(dst_indices); @@ -363,7 +387,7 @@ struct FastIndexCopyKernel { const auto num_indices_tensor = num_indices.value(); TensorMatcher({1}) .with_dtype(num_indices_dtype) - .with_device(device) + .with_device(device) .verify(num_indices_tensor); num_indices_data_ptr = static_cast(num_indices_tensor.data_ptr()); @@ -529,14 +553,14 @@ struct MultiIndexCopyKernel { auto indices_dtype = SymbolicDType{}; auto num_indices_dtype = SymbolicDType{}; - TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) + TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) .verify(dst_ptrs).verify(src_ptrs).verify(feat_bytes); - TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) + TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) .verify(dst_indices).verify(src_indices); const int64_t* valid_length = nullptr; if (num_indices.has_value()) { - TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) + TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) .verify(num_indices.value()); valid_length = static_cast(num_indices.value().data_ptr()); } diff --git a/python/freetoken/kernel/csrc/jit/index.cu b/python/freetoken/kernel/csrc/jit/index.cu index ca0e1db2..aca58383 100644 --- a/python/freetoken/kernel/csrc/jit/index.cu +++ b/python/freetoken/kernel/csrc/jit/index.cu @@ -114,15 +114,15 @@ struct IndexKernel { TensorMatcher({-1, D}) // .with_dtype(weights_dtype_) - .with_device(device_) + .with_device(device_) .verify(weights); TensorMatcher({L, D}) // .with_dtype(weights_dtype_) - .with_device(device_) + .with_device(device_) .verify(output); TensorMatcher({L}) // .with_dtype(indices_dtype_) - .with_device(device_) + .with_device(device_) .verify(indices); const auto device = device_.unwrap(); diff --git a/python/freetoken/kernel/csrc/jit/store.cu b/python/freetoken/kernel/csrc/jit/store.cu index 8d84d76e..162dfdfe 100644 --- a/python/freetoken/kernel/csrc/jit/store.cu +++ b/python/freetoken/kernel/csrc/jit/store.cu @@ -72,18 +72,18 @@ struct StoreKernel { TensorMatcher({-1, D}) // .with_strides({X, 1}) - .with_device(device_) + .with_device(device_) .with_dtype(dtype_) .verify(k_cache) .verify(v_cache); TensorMatcher({L, D}) // .with_strides({Y, 1}) - .with_device(device_) + .with_device(device_) .with_dtype(dtype_) .verify(k) .verify(v); TensorMatcher({L}) // - .with_device(device_) + .with_device(device_) .with_dtype(indices_dtype_) .verify(indices); diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index c3947adf..9355f57a 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp @@ -1,5 +1,5 @@ #include -#include +#include #include namespace { diff --git a/python/freetoken/kernel/fla/l2norm.py b/python/freetoken/kernel/fla/l2norm.py index b7153120..f6ebd6cd 100644 --- a/python/freetoken/kernel/fla/l2norm.py +++ b/python/freetoken/kernel/fla/l2norm.py @@ -51,13 +51,12 @@ def l2norm_fwd_kernel1( # ], # key=["D", "NB"], # ) -@triton.jit +@triton.jit(do_not_specialize=["T"]) def l2norm_fwd_kernel( x, y, eps, - NB: tl.constexpr, - T: tl.constexpr, + T, D: tl.constexpr, BT: tl.constexpr, BD: tl.constexpr, @@ -91,7 +90,6 @@ def l2norm_fwd( raise RuntimeError("This layer doesn't support feature dim >= 64KB.") if D <= 512: - NB = triton.cdiv(T, 2048) def grid(meta): return (triton.cdiv(T, meta["BT"]),) @@ -100,7 +98,6 @@ def grid(meta): x, y, eps, - NB=NB, T=T, D=D, BD=BD, diff --git a/python/freetoken/kernel/fla/utils.py b/python/freetoken/kernel/fla/utils.py index 28a95858..e9e0b697 100644 --- a/python/freetoken/kernel/fla/utils.py +++ b/python/freetoken/kernel/fla/utils.py @@ -271,13 +271,13 @@ def _check_platform() -> Literal["nvidia", "amd", "intel", "musa"]: is_nvidia = device_platform == "nvidia" is_intel_alchemist = is_intel and "Intel(R) Arc(TM) A" in torch.xpu.get_device_name(0) is_nvidia_hopper = is_nvidia and ( - "NVIDIA H" in torch.cuda.get_device_name(0) + "NVIDIA H" in torch.cuda.get_device_name(torch.cuda.current_device()) or torch.cuda.get_device_capability()[0] >= 9 ) use_cuda_graph = is_nvidia and os.environ.get("FLA_USE_CUDA_GRAPH", "0") == "1" # Nvidia Ampere or newer, haven't check AMD and intel yet. -is_tf32_supported = is_nvidia and torch.cuda.get_device_capability(0)[0] >= 8 +is_tf32_supported = is_nvidia and torch.cuda.get_device_capability()[0] >= 8 is_gather_supported = hasattr(triton.language, "gather") @@ -309,8 +309,10 @@ def get_shared_memory(cls, arch: str) -> int: @lru_cache(maxsize=None) -def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: +def check_shared_mem(arch: str = "none", tensor_idx: "int | None" = None) -> bool: try: + if tensor_idx is None: + tensor_idx = device_torch_lib.current_device() device_shared_mem_list = get_all_max_shared_mem() max_shared_memory = device_shared_mem_list[tensor_idx] return max_shared_memory >= Backend.get_shared_memory(arch) diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 04a16560..fee3db00 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -1,18 +1,24 @@ -"""Borrowed llama.cpp GGUF dequant/GEMM CUDA kernels, JIT-compiled on first use. +"""Borrowed llama.cpp GGUF dequant/GEMM kernels, JIT-compiled on first use. -The ``.cu``/``.cuh`` under ``csrc/gguf/`` are vendored verbatim from sgl-kernel +The ``.cu``/``.cuh`` under ``csrc/gguf/`` are vendored from sgl-kernel (``csrc/quantization/gguf/``), which are themselves ports of llama.cpp. We compile -them through ``torch.utils.cpp_extension.load`` (the same toolchain sglang/vllm use) -into a torch-op module and expose the handful of ops the GGUF path needs. This is a -separate, torch-native extension that sits alongside FreeToken's tvm-ffi kernels. +them through ``torch.utils.cpp_extension.load`` into torch-op modules and expose +the handful of ops the GGUF path needs. -All ops keep the weight in its native GGUF block layout (packed ``uint8`` rows) and -dequantize *inside* the kernel -- no bf16 copy of the weight is ever materialized. +CUDA keeps the original monolithic translation unit. ROCm uses operation-split +translation units because AMD clang can spend many minutes optimizing the full +all-quant dequant+MMVQ+MMQ+MoE unit even though each operation family compiles in +seconds on RDNA. Splitting by operation preserves all quant coverage without +forcing per-quant JIT modules. + +All ops keep the weight in its native GGUF block layout (packed ``uint8`` rows) +and dequantize inside the kernel -- no bf16 copy of the weight is materialized. """ from __future__ import annotations import functools +import hashlib import os import pathlib import shutil @@ -20,6 +26,46 @@ import torch _CSRC = pathlib.Path(__file__).parent / "csrc" / "gguf" +_ROCM_OPERATION_SOURCES = { + "dequant": "gguf_dequant_kernel.cu", + "mmvq": "gguf_mmvq_kernel.cu", + "mmq": "gguf_mmq_kernel.cu", + "moe_vec": "gguf_moe_vec_kernel.cu", + "moe": "gguf_moe_kernel.cu", +} + + +def _is_rocm() -> bool: + return getattr(torch.version, "hip", None) is not None + + +def _staged_rocm_sources() -> pathlib.Path: + """Copy CUDA sources out of the checkout before PyTorch HIPifies them. + + ``torch.utils.cpp_extension.load`` writes generated ``*_hip`` sources next to + the input file. Keeping the staging directory under the extension cache makes + the source checkout stay clean while still allowing normal incremental builds. + """ + cache_root = pathlib.Path( + os.environ.get( + "TORCH_EXTENSIONS_DIR", + pathlib.Path.home() / ".cache" / "torch_extensions", + ) + ) + digest = hashlib.sha256() + digest.update(f"torch={torch.__version__};hip={torch.version.hip}".encode()) + for source in sorted(_CSRC.iterdir()): + if source.is_file() and "_hip." not in source.name and source.suffix != ".hip": + digest.update(source.name.encode()) + digest.update(source.read_bytes()) + staged = cache_root / f"freetoken_gguf_sources_{digest.hexdigest()[:16]}" + shutil.copytree( + _CSRC, + staged, + dirs_exist_ok=True, + ignore=shutil.ignore_patterns("*_hip.*", "*.hip", "__pycache__"), + ) + return staged def _host_compiler() -> str | None: @@ -47,54 +93,107 @@ def _c_compiler_for(cxx: str) -> str: cc = base.replace("g++", "gcc") return shutil.which(cc) or cc + +@functools.cache +def _rocm_module(operation: str): + """Build one all-quant GGUF operation family on ROCm. + + Keeping quant types together avoids a large fleet of JIT extensions while + keeping AMD clang away from the pathological monolithic translation unit. + """ + if operation not in _ROCM_OPERATION_SOURCES: + raise ValueError(f"unknown ROCm GGUF operation: {operation}") + + from freetoken.kernel.utils import _rocm_link_flags + from torch.utils.cpp_extension import load + + csrc = _staged_rocm_sources() + return load( + name=f"freetoken_gguf_rocm_{operation}_kernels", + sources=[str(csrc / _ROCM_OPERATION_SOURCES[operation])], + extra_include_paths=[str(csrc)], + extra_cuda_cflags=[ + "-O3", + "-DTHRUST_DEVICE_SYSTEM=THRUST_DEVICE_SYSTEM_CPP", + ], + extra_ldflags=_rocm_link_flags(), + verbose=True, + ) + + @functools.cache def _module(): + """Build the original monolithic CUDA extension. + + This remains available on ROCm for compatibility/debugging, but public + wrappers route ROCm calls through ``_rocm_module`` instead. + """ from torch.utils.cpp_extension import load - extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"] - host_cxx = _host_compiler() + is_rocm = _is_rocm() + extra_cuda_cflags = ["-O3"] + extra_ldflags: list[str] = [] + if is_rocm: + from freetoken.kernel.utils import _rocm_link_flags + + extra_ldflags = _rocm_link_flags() + extra_cuda_cflags.append("-DTHRUST_DEVICE_SYSTEM=THRUST_DEVICE_SYSTEM_CPP") + csrc = _staged_rocm_sources() + else: + extra_cuda_cflags.append("--expt-relaxed-constexpr") + csrc = _CSRC + + host_cxx = None if is_rocm else _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. return load( name="freetoken_gguf_kernels", - sources=[str(_CSRC / "gguf_kernel.cu")], - extra_include_paths=[str(_CSRC)], + sources=[str(csrc / "gguf_kernel.cu")], + extra_include_paths=[str(csrc)], extra_cuda_cflags=extra_cuda_cflags, + extra_ldflags=extra_ldflags, verbose=True, ) +def _operation_module(operation: str): + return _rocm_module(operation) if _is_rocm() else _module() + + # ---- thin typed wrappers (signatures mirror sgl_kernel.quantization.gguf) ---- def ggml_dequantize( - weight: torch.Tensor, quant_type: int, m: int, n: int, dtype: torch.dtype | None = None + weight: torch.Tensor, + quant_type: int, + m: int, + n: int, + dtype: torch.dtype | None = None, ) -> torch.Tensor: - """Dequantize a packed GGUF weight ``[m, row_bytes]`` to a dense ``[m, n]`` tensor.""" - return _module().ggml_dequantize(weight, quant_type, m, n, dtype) + """Dequantize a packed GGUF weight ``[m, row_bytes]`` to dense ``[m, n]``.""" + return _operation_module("dequant").ggml_dequantize( + weight, quant_type, m, n, dtype + ) def ggml_mul_mat_vec_a8( weight: torch.Tensor, x: torch.Tensor, quant_type: int, row: int ) -> torch.Tensor: """MMVQ: small-batch GEMV with on-the-fly dequant. ``row`` = output features.""" - return _module().ggml_mul_mat_vec_a8(weight, x, quant_type, row) + return _operation_module("mmvq").ggml_mul_mat_vec_a8( + weight, x, quant_type, row + ) def ggml_mul_mat_a8( weight: torch.Tensor, x: torch.Tensor, quant_type: int, row: int ) -> torch.Tensor: """MMQ: large-batch quantized matmul. ``row`` = output features.""" - return _module().ggml_mul_mat_a8(weight, x, quant_type, row) + return _operation_module("mmq").ggml_mul_mat_a8(weight, x, quant_type, row) def ggml_moe_a8( @@ -109,9 +208,16 @@ def ggml_moe_a8( tokens: int, ) -> torch.Tensor: """MMQ grouped expert matmul over stacked experts ``weight[E, row, *]``.""" - return _module().ggml_moe_a8( - x, weight, sorted_token_ids, expert_ids, num_tokens_post_padded, - quant_type, row, top_k, tokens, + return _operation_module("moe").ggml_moe_a8( + x, + weight, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + quant_type, + row, + top_k, + tokens, ) @@ -125,11 +231,13 @@ def ggml_moe_a8_vec( tokens: int, ) -> torch.Tensor: """MMVQ grouped expert GEMV over stacked experts ``weight[E, row, *]``.""" - return _module().ggml_moe_a8_vec(x, weight, topk_ids, top_k, quant_type, row, tokens) + return _operation_module("moe_vec").ggml_moe_a8_vec( + x, weight, topk_ids, top_k, quant_type, row, tokens + ) def ggml_moe_get_block_size(quant_type: int) -> int: - return _module().ggml_moe_get_block_size(quant_type) + return _operation_module("moe").ggml_moe_get_block_size(quant_type) __all__ = [ diff --git a/python/freetoken/kernel/pynccl.py b/python/freetoken/kernel/pynccl.py index 23ea5735..71bde734 100644 --- a/python/freetoken/kernel/pynccl.py +++ b/python/freetoken/kernel/pynccl.py @@ -27,6 +27,7 @@ def get_buffer(self) -> int: ... @functools.cache def _load_nccl_module() -> Module: + # TODO(ROCm): NCCL -> RCCL migration for multi-GPU tensor parallelism on AMD. return load_aot("pynccl", cuda_files=["pynccl.cu"], extra_ldflags=["-lnccl"]) diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py index 2c38b533..0b7c945c 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/python/freetoken/kernel/triton/activation.py @@ -20,8 +20,9 @@ import triton.language as tl from triton.language.extra import libdevice from triton.language.extra.cuda import gdc_wait, gdc_launch_dependents +from triton.language import target_info -from freetoken.utils.arch import is_sm90_supported +from freetoken.utils.arch import is_rocm, is_sm90_supported SILU = 0 GELU = 1 @@ -48,6 +49,8 @@ def _pdl_supported() -> bool: @triton.jit def _fast_tanh(x): + if target_info.is_hip(): + return libdevice.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], @@ -57,6 +60,8 @@ def _fast_tanh(x): @triton.jit def _fast_ex2(x): + if target_info.is_hip(): + return libdevice.exp2(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], @@ -134,8 +139,9 @@ def _act_and_mul( 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, + o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, BLOCK_D=block_d, num_warps=4, num_stages=num_stages, + **({} if is_rocm() else {"launch_pdl": pdl}), ) return out diff --git a/python/freetoken/kernel/triton/dsv4/sparse_attn.py b/python/freetoken/kernel/triton/dsv4/sparse_attn.py index c22f891e..7dcefe54 100644 --- a/python/freetoken/kernel/triton/dsv4/sparse_attn.py +++ b/python/freetoken/kernel/triton/dsv4/sparse_attn.py @@ -40,9 +40,11 @@ import triton.language as tl BLOCK_H = 16 -# The gather has exactly ONE tl.load site (the pool base is selected per column), so it stages +# CUDA keeps exactly ONE tl.load site (the pool base is selected per column), so it stages # a single [BLOCK_T, D] KV tile -- 67968 B at BLOCK_T=32, num_stages=2, which fits the ~99KB # consumer-Blackwell (sm_120, e.g. RTX 5090) budget. (BLOCK_T=64 would need ~103KB.) +# ROCm uses separately masked pool loads because AMD Triton can assert while canonicalizing +# a pointer-valued select between the two pool bases (triton-lang/triton#9859). BLOCK_T = 32 MAX_SPLITS = 32 MIN_TILES_PER_SPLIT = 4 @@ -63,6 +65,7 @@ def _sparse_attn_paged_kernel( BLOCK_H: tl.constexpr, BLOCK_T: tl.constexpr, HAS_COUNTS: tl.constexpr, + SEPARATE_POOL_LOADS: tl.constexpr, ): pid_m = tl.program_id(0) pid_b = tl.program_id(1) @@ -99,9 +102,19 @@ def _sparse_attn_paged_kernel( # multiple KV tiles in shared memory. Result is bit-identical to the two-pool form -- each # column still reads from the same pool/slot. is_win = offs_t < N_WINDOW - base = tl.where(is_win, win_ptr, cmp_ptr) # [BLOCK_T] per-column pool base pointer - kv_ptrs = base[:, None] + idxs[:, None] * stride_wn + offs_d[None, :] * stride_wd - kv = tl.load(kv_ptrs, mask=valid[:, None], other=0.0).to(tl.float32) # [BLOCK_T, D] + if SEPARATE_POOL_LOADS: + # Avoid a pointer-valued select on ROCm; see triton-lang/triton#9859. + win_ptrs = win_ptr + idxs[:, None] * stride_wn + offs_d[None, :] * stride_wd + cmp_ptrs = cmp_ptr + idxs[:, None] * stride_cn + offs_d[None, :] * stride_cd + win_valid = valid & is_win + cmp_valid = valid & ~is_win + win_kv = tl.load(win_ptrs, mask=win_valid[:, None], other=0.0) + cmp_kv = tl.load(cmp_ptrs, mask=cmp_valid[:, None], other=0.0) + kv = (win_kv + cmp_kv).to(tl.float32) + else: + base = tl.where(is_win, win_ptr, cmp_ptr) # [BLOCK_T] per-column pool base pointer + kv_ptrs = base[:, None] + idxs[:, None] * stride_wn + offs_d[None, :] * stride_wd + kv = tl.load(kv_ptrs, mask=valid[:, None], other=0.0).to(tl.float32) # [BLOCK_T, D] scores = tl.dot(q, tl.trans(kv)) * scale # [BLOCK_H, BLOCK_T] scores = tl.where(valid[None, :], scores, -float("inf")) @@ -140,6 +153,7 @@ def _sparse_attn_paged_splitk_kernel( BLOCK_H: tl.constexpr, BLOCK_T: tl.constexpr, HAS_COUNTS: tl.constexpr, + SEPARATE_POOL_LOADS: tl.constexpr, NUM_SPLITS: tl.constexpr, ): """Stage 1: each program reduces one BLOCK_T-aligned slice of the candidate list and writes @@ -182,9 +196,19 @@ def _sparse_attn_paged_splitk_kernel( idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) valid = idxs >= 0 is_win = offs_t < N_WINDOW - base = tl.where(is_win, win_ptr, cmp_ptr) - kv_ptrs = base[:, None] + idxs[:, None] * stride_wn + offs_d[None, :] * stride_wd - kv = tl.load(kv_ptrs, mask=valid[:, None], other=0.0).to(tl.float32) + if SEPARATE_POOL_LOADS: + # Avoid a pointer-valued select on ROCm; see triton-lang/triton#9859. + win_ptrs = win_ptr + idxs[:, None] * stride_wn + offs_d[None, :] * stride_wd + cmp_ptrs = cmp_ptr + idxs[:, None] * stride_cn + offs_d[None, :] * stride_cd + win_valid = valid & is_win + cmp_valid = valid & ~is_win + win_kv = tl.load(win_ptrs, mask=win_valid[:, None], other=0.0) + cmp_kv = tl.load(cmp_ptrs, mask=cmp_valid[:, None], other=0.0) + kv = (win_kv + cmp_kv).to(tl.float32) + else: + base = tl.where(is_win, win_ptr, cmp_ptr) + kv_ptrs = base[:, None] + idxs[:, None] * stride_wn + offs_d[None, :] * stride_wd + kv = tl.load(kv_ptrs, mask=valid[:, None], other=0.0).to(tl.float32) scores = tl.dot(q, tl.trans(kv)) * scale scores = tl.where(valid[None, :], scores, -float("inf")) @@ -345,6 +369,7 @@ def sparse_attn_paged( BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, HAS_COUNTS=has_counts, + SEPARATE_POOL_LOADS=torch.version.hip is not None, num_warps=8, num_stages=2, ) @@ -374,6 +399,7 @@ def _sparse_attn_paged_splitk( BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, HAS_COUNTS=has_counts, + SEPARATE_POOL_LOADS=torch.version.hip is not None, NUM_SPLITS=n_splits, num_warps=8, num_stages=2, diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 61d3a0e7..6ea0c753 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -59,15 +59,15 @@ def e4m3_native() -> bool: if _native is None: if FORCE_EMU: _native = False + elif torch.version.hip is not None: + # ROCm reports gfx1101 as capability (11, 0), which is not a CUDA + # compute capability and must not select the native fp8e4nv path. + _native = False else: - native = {torch.cuda.get_device_capability(i) >= (8, 9) - for i in range(torch.cuda.device_count())} - if len(native) > 1: - raise NotImplementedError( - "GPUs on both sides of the sm_89 fp8 boundary in one process: " - "the host-side e4m3 convention is process-global" - ) - _native = native.pop() if native else torch.cuda.get_device_capability() >= (8, 9) + from freetoken.gpu_select import assigned_visible_gpu + + # one process runs on one GPU, so its convention is that GPU's; None (-> the current device) only before the process binds + _native = torch.cuda.get_device_capability(assigned_visible_gpu()) >= (8, 9) return _native diff --git a/python/freetoken/kernel/triton/norm.py b/python/freetoken/kernel/triton/norm.py index 3f95c29f..9071e1df 100644 --- a/python/freetoken/kernel/triton/norm.py +++ b/python/freetoken/kernel/triton/norm.py @@ -30,7 +30,7 @@ import triton.language as tl from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait -from freetoken.utils.arch import is_sm90_supported +from freetoken.utils.arch import is_rocm, is_sm90_supported _HEUR = {"BLOCK": lambda a: triton.next_power_of_2(a["H"])} @@ -144,7 +144,8 @@ def _rmsnorm(input, weight, eps, out, gemma: bool): 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, + CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma, + **({} if is_rocm() else {"launch_pdl": pdl}), num_warps=_num_warps(A * B), num_stages=1, ) return out @@ -172,7 +173,8 @@ def _fused_add_rmsnorm(input, residual, weight, eps, gemma: bool): 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, + CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma, + **({} if is_rocm() else {"launch_pdl": pdl}), num_warps=_num_warps(A * B), num_stages=1, ) diff --git a/python/freetoken/kernel/triton/sampling.py b/python/freetoken/kernel/triton/sampling.py index 9d1bd3da..7345d65f 100644 --- a/python/freetoken/kernel/triton/sampling.py +++ b/python/freetoken/kernel/triton/sampling.py @@ -28,7 +28,7 @@ from freetoken.kernel.triton.autotune_cache import autotune_cache_kwargs -_NUM_SM = torch.cuda.get_device_properties(0).multi_processor_count +_NUM_SM = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count _MIN_CHUNK = 4096 # do not split a row finer than this diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index 7a0164b5..533b1222 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -1,9 +1,11 @@ from __future__ import annotations +import hashlib import importlib import os import pathlib import re +from functools import cache from typing import TYPE_CHECKING, List, NamedTuple, Tuple, TypeAlias, Union if TYPE_CHECKING: @@ -19,7 +21,14 @@ DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")] DEFAULT_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"] +DEFAULT_HIP_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_LDFLAGS = [] +DEFAULT_ROCM_ARCHES = ("gfx1100", "gfx1101", "gfx1102", "gfx1103", "gfx1200", "gfx1201") + + +def _is_rocm() -> bool: + import torch + return getattr(torch.version, "hip", None) is not None def _cuda_cflags(extra: List[str]) -> List[str]: @@ -40,6 +49,91 @@ 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 _hip_cflags(extra: List[str]) -> List[str]: + """HIP flags for a kernel build on ROCm.""" + # TODO(ROCm): Triton autotune configs need RDNA-specific tuning (wave count, LDS size). + flags = DEFAULT_HIP_CFLAGS + extra + raw_arches = os.getenv("FREETOKEN_ROCM_ARCH") or os.getenv("PYTORCH_ROCM_ARCH", "") + arches = list(dict.fromkeys(re.findall(r"gfx\d+[a-z]?", raw_arches.lower()))) + if not arches: + from freetoken.utils.arch import get_rocm_gfx_arch + + detected = get_rocm_gfx_arch() + arches = [detected] if detected else list(DEFAULT_ROCM_ARCHES) + return flags + [f"--offload-arch={arch}" for arch in arches] + + +def _select_versioned_rocm_runtime(paths): + """Load the standalone toolchain selector without importing kernel package state.""" + path = pathlib.Path(__file__).with_name("_toolchain.py") + spec = importlib.util.spec_from_file_location("_freetoken_toolchain_runtime", path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load ROCm runtime selector from {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.select_versioned_rocm_runtime(paths) + + +@cache +def _rocm_link_flags() -> List[str]: + """Make ROCm's runtime library discoverable to JIT link commands. + + Traditional ROCm installs provide ``libamdhip64.so`` under ``$ROCM_HOME/lib``. + ROCm 7.14 Python SDK images only provide the versioned soname, while TVM-FFI + still links with ``-lamdhip64``. Supply a cache-local unversioned symlink via + an explicit linker search path without modifying the Python environment. + + The compatibility directory is keyed by the resolved runtime origin so a + long-lived user cache cannot retain a link to a previous ROCm SDK after the + environment or image changes. Different tensor-parallel ranks using the same + runtime still converge on the same cache path. + """ + candidates: list[pathlib.Path] = [] + if os.getenv("ROCM_HOME"): + candidates.append(pathlib.Path(os.environ["ROCM_HOME"])) + try: + from torch.utils.cpp_extension import ROCM_HOME + + if ROCM_HOME: + candidates.append(pathlib.Path(ROCM_HOME)) + except ImportError: + pass + spec = importlib.util.find_spec("_rocm_sdk_core") + if spec and spec.submodule_search_locations: + candidates.append(pathlib.Path(next(iter(spec.submodule_search_locations)))) + candidates.append(pathlib.Path("/opt/rocm")) + + for rocm_home in dict.fromkeys(candidates): + library_dir = rocm_home / "lib" + unversioned = library_dir / "libamdhip64.so" + link_dir = library_dir + if not unversioned.exists(): + versioned = _select_versioned_rocm_runtime(library_dir.glob("libamdhip64.so.*")) + if versioned is None: + continue + runtime_target = versioned.resolve() + cache_key = hashlib.sha256(str(runtime_target).encode("utf-8")).hexdigest()[:16] + link_dir = pathlib.Path.home() / ".cache" / "freetoken" / "rocm-lib" / cache_key + link_dir.mkdir(parents=True, exist_ok=True) + compat_link = link_dir / "libamdhip64.so" + if not compat_link.exists(): + try: + compat_link.symlink_to(runtime_target) + except FileExistsError: + # Multiple tensor-parallel ranks may prepare the same cache. + pass + if not compat_link.exists(): + raise RuntimeError( + f"ROCm runtime compatibility link is unavailable: {compat_link} -> {runtime_target}" + ) + + return [f"-L{link_dir}", f"-Wl,-rpath,{library_dir}"] + + raise RuntimeError("Unable to locate libamdhip64 for ROCm JIT linking") + + CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool] @@ -217,13 +311,20 @@ def load_aot( cpp_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cpp_files] cuda_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cuda_files] + if _is_rocm(): + cuda_cflags = _hip_cflags(extra_cuda_cflags) + runtime_ldflags = _rocm_link_flags() + else: + cuda_cflags = _cuda_cflags(extra_cuda_cflags) + runtime_ldflags = [] + return load( name, cpp_files=cpp_files, cuda_files=cuda_files, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags), - extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, + extra_cuda_cflags=cuda_cflags, + extra_ldflags=DEFAULT_LDFLAGS + runtime_ldflags + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, ) @@ -272,13 +373,20 @@ def load_jit( cuda_sources = [f'#include "{path}"' for path in cuda_paths] cuda_sources += [_make_wrapper(tup) for tup in cuda_wrappers] + if _is_rocm(): + cuda_cflags = _hip_cflags(extra_cuda_cflags) + runtime_ldflags = _rocm_link_flags() + else: + cuda_cflags = _cuda_cflags(extra_cuda_cflags) + runtime_ldflags = [] + return load_inline( name, cpp_sources=cpp_sources, cuda_sources=cuda_sources, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags), - extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, + extra_cuda_cflags=cuda_cflags, + extra_ldflags=DEFAULT_LDFLAGS + runtime_ldflags + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, ) diff --git a/python/freetoken/moe/bench_profile.py b/python/freetoken/moe/bench_profile.py index 46b11d44..d7b72073 100644 --- a/python/freetoken/moe/bench_profile.py +++ b/python/freetoken/moe/bench_profile.py @@ -1,4 +1,4 @@ -"""Torch-free reader for the ``ft bench bw`` hardware profile (``benchbw.json``). +"""Torch-free reader for the ``ft bench bw`` hardware profile (``benchbw/.json``). The engine consults this at MoE-backend *auto* resolution (``engine.py``) to make the offload-vs-hybrid choice hardware-adaptive without importing the (torch-heavy) benchmark @@ -31,10 +31,42 @@ } -def default_profile_path() -> str: - """``$XDG_CACHE_HOME/freetoken/benchbw.json`` (mirrors ``benchbw.default_out_path``).""" +def _cache_dir() -> str: cache = os.environ.get("XDG_CACHE_HOME") or os.path.expanduser("~/.cache") - return os.path.join(cache, "freetoken", "benchbw.json") + return os.path.join(cache, "freetoken") + + +def default_profile_path(gpu_uuid: str | None = None) -> str: + """``$XDG_CACHE_HOME/freetoken/benchbw/.json``, or the legacy ``benchbw.json`` without a uuid. + + One file per GPU: bandwidth differs between slots. + """ + if gpu_uuid: + return os.path.join(_cache_dir(), "benchbw", f"{gpu_uuid}.json") + return os.path.join(_cache_dir(), "benchbw.json") + + +def latest_profile_path() -> str | None: + """Newest ``benchbw/*.json``, else the legacy ``benchbw.json``, else None.""" + per_gpu = os.path.join(_cache_dir(), "benchbw") + newest: tuple[float, str] | None = None + try: + for name in os.listdir(per_gpu): + if not name.endswith(".json"): + continue + path = os.path.join(per_gpu, name) + try: + mtime = os.path.getmtime(path) + except OSError: + continue + if newest is None or mtime > newest[0]: + newest = (mtime, path) + except OSError: + pass + if newest is not None: + return newest[1] + legacy = default_profile_path() + return legacy if os.path.isfile(legacy) else None def _load(path: str) -> dict | None: @@ -45,14 +77,28 @@ def _load(path: str) -> dict | None: return None -def _usable_profile(gpu_name: str | None, path: str | None) -> dict | None: +def _usable_profile( + gpu_name: str | None, path: str | None, gpu_uuid: str | None = None +) -> dict | None: """The cached profile, or ``None`` when there is no file / it was benched on another GPU (bandwidths are hardware-specific, so a mismatch is ignored rather than trusted). - ``path`` overrides the profile location (else ``FREETOKEN_BENCHBW_PATH`` then the default). + Lookup: explicit ``path`` (else ``FREETOKEN_BENCHBW_PATH``) -> ``benchbw/.json`` -> legacy ``benchbw.json``. """ - src = path or os.environ.get("FREETOKEN_BENCHBW_PATH") or default_profile_path() - prof = _load(src) + explicit = path or os.environ.get("FREETOKEN_BENCHBW_PATH") + if explicit: + candidates = [explicit] + else: + candidates = [default_profile_path(gpu_uuid)] if gpu_uuid else [] + candidates.append(default_profile_path()) + prof = None + for src in candidates: + prof = _load(src) + if isinstance(prof, dict): + break + if os.path.exists(src): + # unreadable profile for this card: stay on the safe default, do not borrow the legacy file + return None if not isinstance(prof, dict): return None prof_gpu = (prof.get("gpu") or {}).get("name") @@ -66,7 +112,10 @@ def _usable_profile(gpu_name: str | None, path: str | None) -> dict | None: def load_backend_recommendation( - quant_format: str, gpu_name: str | None = None, path: str | None = None + quant_format: str, + gpu_name: str | None = None, + path: str | None = None, + gpu_uuid: str | None = None, ) -> str | None: """Bench-recommended offload-family backend for ``quant_format`` on this GPU, or ``None``. @@ -77,7 +126,7 @@ def load_backend_recommendation( default (offload) on ``None``. """ fmt = _QUANT_TO_BENCH_FORMAT.get(quant_format, quant_format) - prof = _usable_profile(gpu_name, path) + prof = _usable_profile(gpu_name, path, gpu_uuid) if prof is None: return None @@ -105,7 +154,10 @@ def load_backend_recommendation( def load_hybrid_fetch_fraction( - quant_format: str, gpu_name: str | None = None, path: str | None = None + quant_format: str, + gpu_name: str | None = None, + path: str | None = None, + gpu_uuid: str | None = None, ) -> float | None: """Benched hybrid fetch fraction for ``quant_format``, or ``None``. @@ -119,7 +171,7 @@ def load_hybrid_fetch_fraction( ``None`` = no usable profile; clamped to [0, 1]. """ fmt = _QUANT_TO_BENCH_FORMAT.get(quant_format, quant_format) - prof = _usable_profile(gpu_name, path) + prof = _usable_profile(gpu_name, path, gpu_uuid) if prof is None: return None entries = [(prof.get("dtype_kernels") or {}).get(fmt)] + [ diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py index bab06ba0..f3e5359a 100644 --- a/python/freetoken/moe/benchbw.py +++ b/python/freetoken/moe/benchbw.py @@ -21,13 +21,15 @@ The hybrid-vs-offload choice is dtype-dominated, so the default is a **per-dtype tuning bench** (``--dtype``): one bench per expert format against a canonical geometry -- the minimal set the runtime backend pick matches on. ``--model`` additionally benches specific model geometries for -per-model detail. Results are written to a JSON file (default ``$XDG_CACHE_HOME/freetoken/ -benchbw.json``) so the choice is reproducible. +per-model detail. Results are written to a JSON file per GPU (default ``$XDG_CACHE_HOME/ +freetoken/benchbw/.json``) so the choice is reproducible and a multi-GPU box keeps +one profile per card. ft bench bw # per-dtype tuning bench (default: all formats) ft bench bw --dtype nvfp4,bf16 # only these formats ft bench bw --model qwen3.6-moe # per-model detail instead ft bench bw --dtype all --model all # both + ft bench bw --gpu GPU-2f3a... # bench a specific GPU (UUID or nvidia-smi index) """ from __future__ import annotations @@ -47,6 +49,12 @@ import torch +from freetoken.gpu_select import ( + assign_gpu, + bind_assigned_gpu, + gpu_identity, + single_gpu_arg, +) from freetoken.kernel.pinned import alloc_pinned_tensor from freetoken.moe.cpu_executor import physical_core_cpus, resolve_threads_and_affinity from freetoken.utils import init_logger @@ -134,11 +142,11 @@ class Workload: } -def default_out_path() -> str: +def default_out_path(gpu_uuid: str | None = None) -> str: # Single source of truth with the (torch-free) reader the engine consults. from freetoken.moe.bench_profile import default_profile_path - return default_profile_path() + return default_profile_path(gpu_uuid) def _cgroup_mem_headroom() -> int | None: @@ -696,12 +704,17 @@ def run_benchbw( "benchbw needs a CUDA device to measure PCIe bandwidth (both offload and " "hybrid serve experts to the GPU)." ) + if torch.cuda.device_count() == 0: + raise RuntimeError( + f"no CUDA device visible (CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')!r})." + ) if not 0 <= device_index < torch.cuda.device_count(): raise RuntimeError( - f"--device {device_index} out of range (found {torch.cuda.device_count()} CUDA devices)." + f"device_index {device_index} out of range (found {torch.cuda.device_count()} CUDA devices)." ) device = torch.device("cuda", device_index) torch.cuda.set_device(device) + gpu = gpu_identity(device_index) # Machine-parseable progress on stdout (opt-in), so the daemon/Desktop can stream feedback # while the bench runs -- mirrors ft checkpoint's FTCONVERT lines. `done`/`total` count the @@ -766,7 +779,8 @@ def _prog(label: str) -> None: "timestamp": datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds"), "epoch": int(time.time()), "host": socket.gethostname(), - "gpu": {"index": device_index, "name": torch.cuda.get_device_name(device)}, + # index is the CUDA ordinal the bench ran on; uuid keys the profile file + "gpu": {"index": device_index, "name": gpu["name"], "uuid": gpu["uuid"]}, "cpu": {"physical_cores": len(physical_core_cpus()), "threads_used": cpu["threads"]}, "threshold": threshold, "ceilings": { @@ -780,9 +794,11 @@ def _prog(label: str) -> None: "workloads": workloads_out, } - out_path = os.path.expanduser(out_path or default_out_path()) + out_path = os.path.expanduser(out_path or default_out_path(gpu["uuid"])) _atomic_write_json(out_path, result) result["out_path"] = out_path + if prog_on: + print(f"FTBENCH_OUT {out_path}", flush=True) return result @@ -927,10 +943,12 @@ def main(argv: list[str] | None = None, prog: str = "ft bench bw") -> int: help=f"CPU MoE ISA: 'auto' (default, best), 'all', or a subset of " f"{list(_ISA_TIERS)} to sweep (kernel caps down to hw support)") p.add_argument("-o", "--out", default=None, - help=f"JSON output path (default {default_out_path()})") + help=f"JSON output path (default {default_out_path('')})") p.add_argument("--threshold", type=_positive_float, default=2.0, help="recommend hybrid when CPU BW > threshold x PCIe BW (default 2.0)") - p.add_argument("--device", type=_nonneg_int, default=0, help="CUDA device index (default 0)") + p.add_argument("--gpu", type=single_gpu_arg, default=None, + help="GPU to bench: a GPU UUID (GPU-xxxx..., as nvidia-smi -L prints) or an " + "nvidia-smi index (default: the first visible GPU)") p.add_argument("--cpu-threads", type=_nonneg_int, default=0, help="CPU worker threads (0 = one per physical core)") p.add_argument("--cpu-iters", type=_positive_int, default=8, help="STREAM read passes to time") @@ -942,6 +960,13 @@ def main(argv: list[str] | None = None, prog: str = "ft bench bw") -> int: help="fast_index_copy gather passes to time") ns = p.parse_args(argv) + # same as ft serve --gpu: resolve, then bind by UUID at CUDA init + try: + assign_gpu(ns.gpu) + device_index = bind_assigned_gpu().index + except (ValueError, RuntimeError) as e: + p.error(str(e)) + models = ns.model or () dtypes = ns.dtype # Default (no selection): the full per-dtype tuning bench -- the minimal set the runtime @@ -951,7 +976,7 @@ def main(argv: list[str] | None = None, prog: str = "ft bench bw") -> int: try: result = run_benchbw( - out_path=ns.out, threshold=ns.threshold, device_index=ns.device, models=models, + out_path=ns.out, threshold=ns.threshold, device_index=device_index, models=models, dtypes=dtypes, formats=ns.formats, isas=ns.isas, cpu_threads=ns.cpu_threads, cpu_iters=ns.cpu_iters, pcie_bytes=ns.pcie_mib << 20, pcie_iters=ns.pcie_iters, kernel_cpu_iters=ns.kernel_cpu_iters, kernel_pcie_iters=ns.kernel_pcie_iters, diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py index 5923c823..d7af348a 100644 --- a/python/freetoken/moe/host_banks.py +++ b/python/freetoken/moe/host_banks.py @@ -292,10 +292,14 @@ class PinPipeline: def __init__(self) -> None: self._q: queue.SimpleQueue = queue.SimpleQueue() self._exc: BaseException | None = None + # the current device is thread-local: a fresh thread sits on device 0 and cudaHostRegister would build its context there -- carry the creator's (bound) device into the worker + self._device = torch.cuda.current_device() if torch.cuda.is_available() else None self._thread = threading.Thread(target=self._run, daemon=True) self._thread.start() def _run(self) -> None: + if self._device is not None: + torch.cuda.set_device(self._device) while True: item = self._q.get() if item is None: diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 35541161..48923e3b 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -6,6 +6,7 @@ from freetoken.attention.linear import build_fla_metadata from freetoken.core import Batch, Req from freetoken.env import ENV +from freetoken.gpu_select import gpu_identity from freetoken.message import ( AbortBackendMsg, BaseBackendMsg, @@ -68,6 +69,8 @@ def __init__(self, config: SchedulerConfig): self.stream = torch.cuda.Stream(device=self.device) self.engine_stream_ctx = torch.cuda.stream(self.engine.stream) torch.cuda.set_stream(self.stream) + # sent on the readiness ack for /v1/stats gpus; a list so TP can add one entry per rank + self.gpus = [gpu_identity(self.device.index)] if self.device.type == "cuda" else [] # initialize other managers self.table_manager = TableManager(config.max_running_req, self.engine.page_table) diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 80dbf6ab..3e2acc85 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -174,6 +174,8 @@ class FrontendManager: # "num_mamba_slots"}, from the same ack. Seeds geometry before the first generation reply # (the running snapshot channel) has anything. None until meta arrives. cache_pools: Dict[str, int] | None = None + # one {index, name, uuid, total_bytes} per TP rank, from the same ack; /v1/stats gpus + gpus: List[Dict[str, Any]] = field(default_factory=list) # Backend worker Process handles (TP schedulers + tokenizer/detokenizer), captured from the # BackendHandle after start_backend(). The orderly-shutdown path (lifespan / shell signal # handler) tears these down itself, AFTER setting _SHUTTING_DOWN, so the supervisor observes @@ -1010,6 +1012,7 @@ def _on_meta(meta: dict) -> None: _GLOBAL_STATE.cache_pools = meta.pop("pools", None) _GLOBAL_STATE.swa_full_tokens_ratio = float(meta.pop("swa_full_tokens_ratio", 0.0) or 0.0) _GLOBAL_STATE.cache_budget_bytes = int(meta.pop("cache_budget_bytes", 0) or 0) + _GLOBAL_STATE.gpus = list(meta.pop("gpus", None) or []) _GLOBAL_STATE.unit_bytes = meta # Early-bind: supervise the backend on a daemon thread so uvicorn can bind diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 4954c5f5..a71b6819 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -39,6 +39,10 @@ class ServerArgs(SchedulerConfig): # Comma-separated CORS allow-list for browser/webview clients (e.g. the desktop # app). Empty string disables CORS headers entirely; "*" allows any origin. cors_origins: str = "tauri://localhost,http://tauri.localhost,http://localhost:1420" + # --gpu entries in TP-rank order, empty = not given + gpu: tuple[str, ...] = () + # full UUIDs resolved from --gpu, entry i = TP rank i; None = NVML unavailable, each worker then resolves its raw entry against CUDA's own enumeration + gpu_assigned: "tuple[str, ...] | None" = None @property def share_tokenizer(self) -> bool: @@ -109,6 +113,11 @@ def _positive_int(value: str) -> int: raise argparse.ArgumentTypeError("must be >= 1") return n + def _lazy_gpu_arg(value: str) -> tuple[str, ...]: + from freetoken.gpu_select import gpu_arg + + return gpu_arg(value) + def _infer_tool_call_parser(model_path: str) -> str: try: from freetoken.utils import cached_load_hf_config @@ -221,6 +230,16 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="The tensor parallelism size.", ) + parser.add_argument( + "--gpu", + type=_lazy_gpu_arg, + default=ServerArgs.gpu, + help=( + "GPU(s) to run on, comma-separated; entry i is TP rank i. Each entry is a GPU " + "UUID (GPU-xxxx..., as nvidia-smi -L prints) or an nvidia-smi index" + ), + ) + parser.add_argument( "--max-running-requests", type=int, @@ -611,6 +630,15 @@ def _infer_reasoning_parser(model_path: str) -> str | None: # Parse arguments kwargs = parser.parse_args(args).__dict__.copy() + # reject a too-long list here with a clear reason, not as a dead rank later + if len(kwargs["gpu"]) not in (0, kwargs["tensor_parallel_size"]): + if kwargs["tensor_parallel_size"] == 1 and len(kwargs["gpu"]) > 1: + parser.error("tensor parallelism is not supported yet: --gpu takes one entry") + parser.error( + f"--gpu has {len(kwargs['gpu'])} entries but --tensor-parallel-size is " + f"{kwargs['tensor_parallel_size']}; give one entry per TP rank" + ) + # resolve some arguments run_shell |= kwargs.pop("shell_mode") kwargs["shell_mode"] = run_shell diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py index a5f84388..acef5517 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -59,6 +59,13 @@ def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None: if args.shell_mode: _detach_process_group() + # published (not bound) here: the engine binds it after the allocator setup + from freetoken.gpu_select import set_assigned_gpu + + # resolved UUIDs when we have them, the raw --gpu entries when NVML could not resolve them, else one CUDA ordinal per rank + targets = args.gpu_assigned or args.gpu or tuple(str(r) for r in range(args.tp_info.size)) + set_assigned_gpu(targets[args.tp_info.rank]) + import torch from freetoken.scheduler import Scheduler @@ -90,7 +97,10 @@ def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None: try: from freetoken.kvcache.cache_status import compute_cache_status_meta - ack_queue.put(("meta", compute_cache_status_meta(scheduler.engine))) + meta = compute_cache_status_meta(scheduler.engine) + # the parent must not touch CUDA to learn this + meta["gpus"] = scheduler.gpus + ack_queue.put(("meta", meta)) except Exception: # noqa: BLE001 -- metadata is a nicety; readiness is not pass ack_queue.put("Scheduler is ready") @@ -127,6 +137,19 @@ def launch_server( ) logger = init_logger(__name__, "initializer") + if server_args.gpu: + # resolve here so a typo is one clear error before any worker spawns + from freetoken.gpu_select import resolve_gpu_uuids + + try: + server_args = replace(server_args, gpu_assigned=resolve_gpu_uuids(server_args.gpu)) + except ValueError as exc: + raise SystemExit(f"{prog or 'ft serve'}: error: {exc}") from exc + logger.info( + f"--gpu {','.join(server_args.gpu)} -> " + f"{', '.join(server_args.gpu_assigned) if server_args.gpu_assigned else 'resolved at CUDA init (no NVML)'}" + ) + def start_subprocess() -> "BackendHandle": import multiprocessing as mp diff --git a/python/freetoken/server/stats.py b/python/freetoken/server/stats.py index e89d23a4..76c6c308 100644 --- a/python/freetoken/server/stats.py +++ b/python/freetoken/server/stats.py @@ -130,7 +130,9 @@ def _swa_page_size(config: Any) -> int: def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: """Full /v1/stats doc. throughput is 0 when idle; kv/mamba/swa are null when their total is 0 (owned-KV / non-hybrid / non-SWA). kv and swa share one shape: - pages + the pool's own page_size (tokens = pages x page_size).""" + pages + the pool's own page_size (tokens = pages x page_size). gpus: the engine's GPU as + [{index, name, uuid, total_bytes}] (the primary rank's; a list so TP can extend it), [] + until the readiness meta arrives.""" tr: StatsTracker = state.stats config = state.config ready_at = getattr(state, "ready_at", None) @@ -158,6 +160,7 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: "mamba": mamba, "swa": swa, "vram_bytes": tr.vram_bytes, + "gpus": list(getattr(state, "gpus", None) or []), "throughput": { "decode_tps": round(tr.decode_tps(), 1), "prefill_tps": round(tr.prefill_tps(), 1), diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py index 2e4ad15f..c54579a1 100644 --- a/python/freetoken/utils/__init__.py +++ b/python/freetoken/utils/__init__.py @@ -1,5 +1,9 @@ from .arch import ( is_arch_supported, + is_rocm, + get_rocm_gfx_arch, + is_gfx11xx_family, + is_gfx12xx_family, is_sm90_family, is_sm90_supported, is_sm100_family, @@ -35,6 +39,10 @@ "load_toolcall_anchor_id", "init_logger", "is_arch_supported", + "is_rocm", + "get_rocm_gfx_arch", + "is_gfx11xx_family", + "is_gfx12xx_family", "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..3bf8fb61 100644 --- a/python/freetoken/utils/arch.py +++ b/python/freetoken/utils/arch.py @@ -1,14 +1,77 @@ from __future__ import annotations import functools +import os +import re from typing import Tuple +_GFX_ARCH_RE = re.compile(r"gfx\d+[a-z]?") + + +def _gfx_arch_from(value: object) -> str | None: + match = _GFX_ARCH_RE.search(str(value).lower()) + return match.group(0) if match else None + + +@functools.cache +def is_rocm() -> bool: + """True when torch is built for ROCm (AMD GPU) instead of CUDA.""" + import torch + return getattr(torch.version, "hip", None) is not None + + +@functools.cache +def get_rocm_gfx_arch() -> str | None: + """Return the current AMD GPU target (for example ``gfx1201``). + + Prefer the runtime device because build variables may contain multiple + semicolon-separated targets. Environment variables remain useful for + cross-compilation and systems where no GPU is currently visible. + """ + if not is_rocm(): + return None + + import torch + + if torch.cuda.is_available(): + try: + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + for attr in ("gcnArchName", "arch"): + arch = _gfx_arch_from(getattr(props, attr, "")) + if arch: + return arch + except (AttributeError, RuntimeError): + pass + + for env_var in ("FREETOKEN_ROCM_ARCH", "PYTORCH_ROCM_ARCH", "HCC_AMDGPU_TARGET"): + arch = _gfx_arch_from(os.getenv(env_var, "")) + if arch: + return arch + return None + + +@functools.cache +def is_gfx11xx_family() -> bool: + """True when the current AMD GPU is RDNA3 (gfx110x).""" + arch = get_rocm_gfx_arch() + return arch is not None and arch.startswith("gfx110") + + +@functools.cache +def is_gfx12xx_family() -> bool: + """True when the current AMD GPU is RDNA4 (gfx120x).""" + arch = get_rocm_gfx_arch() + return arch is not None and arch.startswith("gfx120") + + @functools.cache def _get_torch_cuda_version() -> Tuple[int, int] | None: import torch import torch.version + if is_rocm(): + return None if not torch.cuda.is_available() or not torch.version.cuda: return None return torch.cuda.get_device_capability() diff --git a/setup.py b/setup.py index cfe41b7d..98ec6d74 100644 --- a/setup.py +++ b/setup.py @@ -1,21 +1,66 @@ from __future__ import annotations import importlib.util +import os from pathlib import Path from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension +from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension, ROCM_HOME ROOT = Path(__file__).parent +KERNEL_INCLUDE = str(ROOT / "python" / "freetoken" / "kernel" / "csrc" / "include") -def _check_toolchain() -> None: +def _toolchain_module(): 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: + _toolchain_module().check_nvcc_matches_torch() + + +def _is_rocm() -> bool: + import torch + return getattr(torch.version, "hip", None) is not None + + +def _rocm_paths() -> tuple[list[str], list[str], str]: + candidates: list[Path] = [] + if os.getenv("ROCM_HOME"): + candidates.append(Path(os.environ["ROCM_HOME"])) + if ROCM_HOME: + candidates.append(Path(ROCM_HOME)) + + # ROCm 7.14 PyTorch images ship the SDK as a Python package instead of + # installing it at /opt/rocm. + spec = importlib.util.find_spec("_rocm_sdk_core") + if spec and spec.submodule_search_locations: + candidates.append(Path(next(iter(spec.submodule_search_locations)))) + candidates.append(Path("/opt/rocm")) + + for rocm_home in dict.fromkeys(candidates): + include_dir = rocm_home / "include" + library_dir = rocm_home / "lib" + if not (include_dir / "hip" / "hip_runtime.h").exists(): + continue + if (library_dir / "libamdhip64.so").exists(): + return [str(include_dir)], [str(library_dir)], "amdhip64" + versioned = _toolchain_module().select_versioned_rocm_runtime( + library_dir.glob("libamdhip64.so.*") + ) + if versioned is not None: + return [str(include_dir)], [str(library_dir)], f":{versioned.name}" + + searched = ", ".join(str(path) for path in dict.fromkeys(candidates)) + raise RuntimeError( + "A ROCm SDK with HIP headers and libamdhip64 is required to build on ROCm; " + f"searched: {searched}. Set ROCM_HOME to override." + ) def _cuda_runtime_paths() -> tuple[list[str], list[str]]: @@ -31,7 +76,21 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: return [str(cuda_home / "include")], library_dirs -cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths() +IS_ROCM = _is_rocm() + +if IS_ROCM: + runtime_include_dirs, runtime_library_dirs, runtime_lib = _rocm_paths() + runtime_link_args = [f"-Wl,-rpath,{runtime_library_dirs[0]}"] + # These extensions contain host code only. BuildExtension supplies the ROCm + # platform defines to the C++ compiler; offload architecture flags belong on + # HIP device sources and would be rejected by the host compiler here. + extra_compile = ["-O3", "-std=c++17"] +else: + runtime_include_dirs, runtime_library_dirs = _cuda_runtime_paths() + runtime_lib = "cudart" + runtime_link_args = [] + extra_compile = ["-O3", "-std=c++17"] + _check_toolchain() @@ -42,12 +101,13 @@ 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=[KERNEL_INCLUDE, *runtime_include_dirs], + library_dirs=runtime_library_dirs, + libraries=[runtime_lib], + extra_compile_args=extra_compile, + extra_link_args=runtime_link_args, ), - # CPU-compute MoE executor for --moe-backend cpu. Links cudart for the + # CPU-compute MoE executor for --moe-backend cpu. Links cudart/amdhip64 for the # cudaLaunchHostFunc submit/sync graph nodes; the bf16 GEMV microkernels # use per-function target attributes (avx512bf16/avx512f) + a runtime # __builtin_cpu_supports dispatch, so the single binary stays portable @@ -57,10 +117,11 @@ 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=[KERNEL_INCLUDE, *runtime_include_dirs], + library_dirs=runtime_library_dirs, + libraries=[runtime_lib], + extra_compile_args=extra_compile + ["-pthread"], + extra_link_args=runtime_link_args, ), ], cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)}, diff --git a/tests/engine/test_rocm_communication.py b/tests/engine/test_rocm_communication.py new file mode 100644 index 00000000..665eb67c --- /dev/null +++ b/tests/engine/test_rocm_communication.py @@ -0,0 +1,74 @@ +from types import SimpleNamespace + +import torch + +import freetoken.engine.engine as engine_module +import freetoken.kernel.backend as kernel_backend +from freetoken.engine.engine import Engine + + +def _config(*, use_pynccl: bool = True): + return SimpleNamespace( + use_pynccl=use_pynccl, + tp_info=SimpleNamespace(size=2, rank=0), + distributed_timeout=10, + distributed_addr="tcp://127.0.0.1:29500", + max_forward_len=32, + model_config=SimpleNamespace(hidden_size=64), + ) + + +def test_rocm_routes_tensor_parallel_communication_to_rccl(monkeypatch): + calls = [] + cpu_group = object() + + def reject_pynccl(*_args): + raise AssertionError("PyNCCL selected on ROCm") + + monkeypatch.setattr(kernel_backend, "is_rocm", lambda: True) + monkeypatch.setattr( + torch.distributed, + "init_process_group", + lambda **kwargs: calls.append(("init", kwargs)), + ) + monkeypatch.setattr( + torch.distributed, + "new_group", + lambda **kwargs: calls.append(("new", kwargs)) or cpu_group, + ) + monkeypatch.setattr(engine_module, "enable_pynccl_distributed", reject_pynccl) + + result = Engine._init_communication(SimpleNamespace(), _config()) + + assert result is cpu_group + assert calls[0][1]["backend"] == "nccl" + assert calls[1] == ("new", {"backend": "gloo"}) + + +def test_cuda_keeps_custom_pynccl_path(monkeypatch): + calls = [] + world_group = object() + + def reject_new_group(**_kwargs): + raise AssertionError("unexpected RCCL path") + + monkeypatch.setattr(kernel_backend, "is_rocm", lambda: False) + monkeypatch.setattr( + torch.distributed, + "init_process_group", + lambda **kwargs: calls.append(("init", kwargs)), + ) + monkeypatch.setattr(torch.distributed, "group", SimpleNamespace(WORLD=world_group)) + monkeypatch.setattr(torch.distributed, "new_group", reject_new_group) + monkeypatch.setattr( + engine_module, + "enable_pynccl_distributed", + lambda *args: calls.append(("pynccl", args)), + ) + + engine = SimpleNamespace(dtype=torch.float16) + result = Engine._init_communication(engine, _config()) + + assert result is world_group + assert calls[0][1]["backend"] == "gloo" + assert calls[1][0] == "pynccl" diff --git a/tests/kernels/test_backend.py b/tests/kernels/test_backend.py new file mode 100644 index 00000000..eab1e988 --- /dev/null +++ b/tests/kernels/test_backend.py @@ -0,0 +1,38 @@ +import freetoken.kernel.backend as backend + + +_CUDA_ONLY_PROBES = ( + backend.is_flashinfer_installed, + backend.is_sgl_kernel_installed, + backend.is_triton_kernels_installed, +) + + +def _clear_probe_caches() -> None: + for probe in _CUDA_ONLY_PROBES: + probe.cache_clear() + backend.driver_cuda_version.cache_clear() + + +def test_rocm_never_selects_cuda_only_backends(monkeypatch): + def unexpected_probe(_name: str) -> bool: + raise AssertionError("unexpected CUDA-only package probe on ROCm") + + monkeypatch.setattr(backend, "is_rocm", lambda: True) + monkeypatch.setattr(backend, "_importable", unexpected_probe) + _clear_probe_caches() + + assert all(not probe() for probe in _CUDA_ONLY_PROBES) + assert backend.driver_cuda_version() is None + + _clear_probe_caches() + + +def test_cuda_keeps_optional_package_probes(monkeypatch): + monkeypatch.setattr(backend, "is_rocm", lambda: False) + monkeypatch.setattr(backend, "_importable", lambda _name: True) + _clear_probe_caches() + + assert all(probe() for probe in _CUDA_ONLY_PROBES) + + _clear_probe_caches() diff --git a/tests/kernels/test_gguf_rocm.py b/tests/kernels/test_gguf_rocm.py new file mode 100644 index 00000000..52936d32 --- /dev/null +++ b/tests/kernels/test_gguf_rocm.py @@ -0,0 +1,27 @@ +import pytest +import torch + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is None, + reason="a ROCm GPU is required", +) + + +def test_q4_0_dequant_matches_torch_reference(): + from freetoken.kernel.gguf import ggml_dequantize + from freetoken.models.gguf.dequant import GGML_Q4_0, dequantize + + scale = torch.tensor([0.5], dtype=torch.float16).view(torch.uint8) + quants = torch.tensor( + [0x10, 0x32, 0x54, 0x76, 0x98, 0xBA, 0xDC, 0xFE] * 2, + dtype=torch.uint8, + ) + packed_cpu = torch.cat((scale, quants)).reshape(1, 18) + expected = dequantize(packed_cpu, GGML_Q4_0, torch.float32).reshape(1, 32) + + actual = ggml_dequantize( + packed_cpu.to("cuda"), GGML_Q4_0, m=1, n=32, dtype=torch.float32 + ) + + torch.testing.assert_close(actual.cpu(), expected) diff --git a/tests/kernels/test_gguf_rocm_operation_split.py b/tests/kernels/test_gguf_rocm_operation_split.py new file mode 100644 index 00000000..d3d48b09 --- /dev/null +++ b/tests/kernels/test_gguf_rocm_operation_split.py @@ -0,0 +1,96 @@ +"""Routing contract for ROCm GGUF operation-split JIT modules. + +These tests deliberately do not compile a native extension. Physical gfx1101 coverage +lives in the ROCm validation receipts; this file protects the cheap Python dispatch +contract so a future refactor cannot silently route one public GGUF op back through the +pathological monolithic ROCm translation unit. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.kernel import gguf + + +_EXPECTED_SOURCES = { + "dequant": "gguf_dequant_kernel.cu", + "mmvq": "gguf_mmvq_kernel.cu", + "mmq": "gguf_mmq_kernel.cu", + "moe_vec": "gguf_moe_vec_kernel.cu", + "moe": "gguf_moe_kernel.cu", +} + + +def test_rocm_operation_source_contract(): + assert gguf._ROCM_OPERATION_SOURCES == _EXPECTED_SOURCES + + +def test_operation_module_uses_rocm_family(monkeypatch): + seen = [] + marker = object() + monkeypatch.setattr(gguf, "_is_rocm", lambda: True) + monkeypatch.setattr(gguf, "_rocm_module", lambda operation: seen.append(operation) or marker) + monkeypatch.setattr(gguf, "_module", lambda: pytest.fail("ROCm dispatch reached monolithic module")) + + assert gguf._operation_module("mmvq") is marker + assert seen == ["mmvq"] + + +def test_operation_module_keeps_cuda_monolith(monkeypatch): + marker = object() + monkeypatch.setattr(gguf, "_is_rocm", lambda: False) + monkeypatch.setattr(gguf, "_module", lambda: marker) + monkeypatch.setattr(gguf, "_rocm_module", lambda operation: pytest.fail(f"CUDA reached ROCm split {operation}")) + + assert gguf._operation_module("dequant") is marker + + +def test_rocm_loader_rejects_unknown_family_before_build(): + with pytest.raises(ValueError, match="unknown ROCm GGUF operation"): + gguf._rocm_module("not-an-operation") + + +def test_public_wrappers_route_to_owned_operation_family(monkeypatch): + calls = [] + + def op(name): + def invoke(*args, **kwargs): + calls.append((name, args, kwargs)) + return name + return invoke + + modules = { + family: SimpleNamespace( + ggml_dequantize=op("dequant"), + ggml_mul_mat_vec_a8=op("mmvq"), + ggml_mul_mat_a8=op("mmq"), + ggml_moe_a8_vec=op("moe_vec"), + ggml_moe_a8=op("moe"), + ggml_moe_get_block_size=op("moe_block"), + ) + for family in _EXPECTED_SOURCES + } + requested = [] + + def load_family(family): + requested.append(family) + return modules[family] + + monkeypatch.setattr(gguf, "_operation_module", load_family) + + w = torch.empty(1, dtype=torch.uint8) + x = torch.empty(1) + ids = torch.empty(1, dtype=torch.int32) + + assert gguf.ggml_dequantize(w, 2, 1, 32) == "dequant" + assert gguf.ggml_mul_mat_vec_a8(w, x, 2, 1) == "mmvq" + assert gguf.ggml_mul_mat_a8(w, x, 2, 1) == "mmq" + assert gguf.ggml_moe_a8_vec(x, w, ids, 1, 2, 1, 1) == "moe_vec" + assert gguf.ggml_moe_a8(x, w, ids, ids, ids, 2, 1, 1, 1) == "moe" + assert gguf.ggml_moe_get_block_size(2) == "moe_block" + + assert requested == ["dequant", "mmvq", "mmq", "moe_vec", "moe", "moe"] diff --git a/tests/kernels/test_jit_index_store.py b/tests/kernels/test_jit_index_store.py new file mode 100644 index 00000000..7590023a --- /dev/null +++ b/tests/kernels/test_jit_index_store.py @@ -0,0 +1,36 @@ +import pytest +import torch + +from freetoken.kernel import indexing, store_cache + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="a CUDA or ROCm GPU is required" +) + + +def test_indexing_jit_matches_torch_on_cold_and_warm_loads(): + weights = torch.arange(8 * 64, dtype=torch.float32, device="cuda").reshape(8, 64) + + for values in ([7, 2, 0], [1, 6, 3]): + indices = torch.tensor(values, dtype=torch.int32, device="cuda") + actual = indexing(weights, indices) + torch.testing.assert_close(actual, weights[indices.long()]) + + +def test_store_jit_matches_torch_on_cold_and_warm_loads(): + k_cache = torch.zeros((8, 64), dtype=torch.float32, device="cuda") + v_cache = torch.zeros_like(k_cache) + indices = torch.tensor([5, 0, 3], dtype=torch.int64, device="cuda") + + for offset in (0.0, 1000.0): + k = torch.arange(3 * 64, dtype=torch.float32, device="cuda").reshape(3, 64) + k = k + offset + v = k + 500.0 + store_cache(k_cache, v_cache, indices, k, v) + torch.testing.assert_close(k_cache[indices], k) + torch.testing.assert_close(v_cache[indices], v) + + untouched = torch.tensor([1, 2, 4, 6, 7], device="cuda") + torch.testing.assert_close(k_cache[untouched], torch.zeros((5, 64), device="cuda")) + torch.testing.assert_close(v_cache[untouched], torch.zeros((5, 64), device="cuda")) diff --git a/tests/kernels/test_pinned_tensor.py b/tests/kernels/test_pinned_tensor.py index e61108fd..2fee4f25 100644 --- a/tests/kernels/test_pinned_tensor.py +++ b/tests/kernels/test_pinned_tensor.py @@ -122,7 +122,12 @@ def test_host_device_ptr_is_identity_under_uva(): pytest.skip("non-UVA platform: host_device_ptr rejects unregistered memory instead") # Under UVA cudaHostGetDevicePointer degenerates to identity for any host pointer # (no registration validation); rejection of pageable memory only exists on - # non-identity platforms (Windows/WDDM), where the translation is real. + # non-identity CUDA platforms (Windows/WDDM), where the translation is real. + # HIP validates registration even though registered/pinned memory uses the + # identity address on Linux. Calling it with pageable memory also leaves a + # sticky HIP error, so the pinned identity case above is the relevant check. + if torch.version.hip is not None: + return pageable = torch.empty(64, dtype=torch.uint8) ext = _load_pinned_extension() assert ext.host_device_ptr(pageable.data_ptr()) == pageable.data_ptr() diff --git a/tests/kernels/test_rocm_host_pointer_capability.py b/tests/kernels/test_rocm_host_pointer_capability.py new file mode 100644 index 00000000..2cf4a87e --- /dev/null +++ b/tests/kernels/test_rocm_host_pointer_capability.py @@ -0,0 +1,23 @@ +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] +HIP_COMPAT = ROOT / "python" / "freetoken" / "kernel" / "csrc" / "include" / "freetoken" / "hip_compat.h" + + +def test_rocm_host_pointer_capability_is_not_aliased_to_uva(): + text = HIP_COMPAT.read_text() + + assert ( + "#define cudaDevAttrCanUseHostPointerForRegisteredMem " + "hipDeviceAttributeCanUseHostPointerForRegisteredMem" + ) in text + assert ( + "#define cudaDevAttrCanUseHostPointerForRegisteredMem " + "hipDeviceAttributeUnifiedAddressing" + ) not in text + + +if __name__ == "__main__": + test_rocm_host_pointer_capability_is_not_aliased_to_uva() + print("ROCM_HOST_POINTER_SOURCE_CONTRACT=PASS") diff --git a/tests/kernels/test_rocm_runtime_link_cache.py b/tests/kernels/test_rocm_runtime_link_cache.py new file mode 100644 index 00000000..fb1b47cb --- /dev/null +++ b/tests/kernels/test_rocm_runtime_link_cache.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import importlib.util +import os +from pathlib import Path +import sys +import tempfile +from unittest import mock + + +ROOT = Path(__file__).parents[2] +PYTHON_ROOT = ROOT / "python" +UTILS = PYTHON_ROOT / "freetoken" / "kernel" / "utils.py" + +if str(PYTHON_ROOT) not in sys.path: + sys.path.insert(0, str(PYTHON_ROOT)) + + +def _load_utils_module(): + spec = importlib.util.spec_from_file_location("_freetoken_kernel_utils_test", UTILS) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_versioned_rocm_runtime_link_cache_tracks_runtime_origin_and_numeric_version() -> None: + module = _load_utils_module() + + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + home = root / "home" + first_root = root / "rocm-first" + second_root = root / "rocm-second" + first_old = first_root / "lib" / "libamdhip64.so.7.9" + first_new = first_root / "lib" / "libamdhip64.so.7.14" + first_debug = first_root / "lib" / "libamdhip64.so.debug" + second_lib = second_root / "lib" / "libamdhip64.so.7.15" + first_old.parent.mkdir(parents=True) + second_lib.parent.mkdir(parents=True) + home.mkdir() + first_old.write_bytes(b"first-old") + first_new.write_bytes(b"first-new") + first_debug.write_bytes(b"debug") + second_lib.write_bytes(b"second") + + with mock.patch.dict(os.environ, {"HOME": str(home), "ROCM_HOME": str(first_root)}, clear=False): + module._rocm_link_flags.cache_clear() + first_flags = module._rocm_link_flags() + + first_link_dir = Path(next(flag[2:] for flag in first_flags if flag.startswith("-L"))) + first_link = first_link_dir / "libamdhip64.so" + assert first_link.is_symlink() + assert first_link.resolve() == first_new.resolve() + + # Model a long-lived cache surviving a ROCm SDK/image change. A stale + # compat symlink must not pin JIT linking to the vanished runtime, and + # the first selection must be numeric (7.14 > 7.9), not lexical. + first_new.unlink() + + with mock.patch.dict(os.environ, {"HOME": str(home), "ROCM_HOME": str(second_root)}, clear=False): + module._rocm_link_flags.cache_clear() + second_flags = module._rocm_link_flags() + + second_link_dir = Path(next(flag[2:] for flag in second_flags if flag.startswith("-L"))) + second_link = second_link_dir / "libamdhip64.so" + assert second_link.is_symlink() + assert second_link.exists() + assert second_link.resolve() == second_lib.resolve() + assert second_link_dir != first_link_dir + + utils_text = UTILS.read_text() + assert "select_versioned_rocm_runtime" in utils_text + assert 'sorted(library_dir.glob("libamdhip64.so.*"))' not in utils_text + + +if __name__ == "__main__": + test_versioned_rocm_runtime_link_cache_tracks_runtime_origin_and_numeric_version() + print("ROCM_JIT_RUNTIME_RESOLUTION=PASS_NUMERIC_SELECTION_AND_CACHE_LIFETIME") diff --git a/tests/kernels/test_rocm_versioned_runtime_selection.py b/tests/kernels/test_rocm_versioned_runtime_selection.py new file mode 100644 index 00000000..6266766a --- /dev/null +++ b/tests/kernels/test_rocm_versioned_runtime_selection.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] +TOOLCHAIN = ROOT / "python" / "freetoken" / "kernel" / "_toolchain.py" +SETUP = ROOT / "setup.py" + + +def _load_toolchain(): + spec = importlib.util.spec_from_file_location("_freetoken_toolchain_test", TOOLCHAIN) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_rocm_versioned_runtime_selection_is_numeric_not_lexical() -> None: + module = _load_toolchain() + candidates = [ + Path("/sdk/lib/libamdhip64.so.7"), + Path("/sdk/lib/libamdhip64.so.7.9"), + Path("/sdk/lib/libamdhip64.so.7.14"), + Path("/sdk/lib/libamdhip64.so.debug"), + ] + selected = module.select_versioned_rocm_runtime(candidates) + assert selected is not None + assert selected.name == "libamdhip64.so.7.14" + + setup_text = SETUP.read_text() + assert "select_versioned_rocm_runtime" in setup_text + assert 'sorted(library_dir.glob("libamdhip64.so.*"))' not in setup_text + + +if __name__ == "__main__": + test_rocm_versioned_runtime_selection_is_numeric_not_lexical() + print("ROCM_VERSIONED_RUNTIME_SELECTION=PASS_NUMERIC_7_14_OVER_7_9") diff --git a/tests/moe/test_hybrid_fetch.py b/tests/moe/test_hybrid_fetch.py index e080e165..169f8e2b 100644 --- a/tests/moe/test_hybrid_fetch.py +++ b/tests/moe/test_hybrid_fetch.py @@ -6,11 +6,12 @@ """ import json +import os import pytest import torch -from freetoken.moe.bench_profile import load_hybrid_fetch_fraction +from freetoken.moe.bench_profile import default_profile_path, load_backend_recommendation, load_hybrid_fetch_fraction from freetoken.moe.offload_cache import OffloadMoeCache Q = 1 << 16 @@ -65,6 +66,25 @@ def test_load_hybrid_fetch_fraction(tmp_path): assert load_hybrid_fetch_fraction("bf16", gpu_name="OTHER", path=str(path)) is None +def test_profile_lookup_prefers_the_gpu_uuid_file(tmp_path, monkeypatch): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + monkeypatch.delenv("FREETOKEN_BENCHBW_PATH", raising=False) + uuid = "GPU-2f3a9b1c-0000-1111-2222-333344445555" + + def write(path, name, verdict): + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w") as f: + json.dump({"gpu": {"name": name}, "dtypes": {"bf16": verdict}}, f) + + # legacy single file only: used when the name matches, ignored otherwise + write(default_profile_path(), "FAKE GPU", "hybrid") + assert load_backend_recommendation("bf16", gpu_name="FAKE GPU", gpu_uuid=uuid) == "hybrid" + assert load_backend_recommendation("bf16", gpu_name="OTHER", gpu_uuid=uuid) is None + # this card's own file wins over the legacy one + write(default_profile_path(uuid), "FAKE GPU", "offload") + assert load_backend_recommendation("bf16", gpu_name="FAKE GPU", gpu_uuid=uuid) == "offload" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") def test_hybrid_fraction_gpu_matches_cpu_reference(): torch.manual_seed(0) diff --git a/tests/utils/test_rocm_arch.py b/tests/utils/test_rocm_arch.py new file mode 100644 index 00000000..da68a864 --- /dev/null +++ b/tests/utils/test_rocm_arch.py @@ -0,0 +1,91 @@ +import importlib +import pathlib +from types import SimpleNamespace + +import torch + +from freetoken.utils import arch + + +def _clear_arch_caches() -> None: + arch.get_rocm_gfx_arch.cache_clear() + arch.is_gfx11xx_family.cache_clear() + arch.is_gfx12xx_family.cache_clear() + + +def test_rocm_arch_prefers_visible_device_over_multi_arch_build_env(monkeypatch): + monkeypatch.setattr(arch, "is_rocm", lambda: True) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _device: SimpleNamespace(gcnArchName="gfx1201:sramecc-:xnack-"), + ) + monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1200") + _clear_arch_caches() + + assert arch.get_rocm_gfx_arch() == "gfx1201" + assert arch.is_gfx12xx_family() + assert not arch.is_gfx11xx_family() + + _clear_arch_caches() + + +def test_rocm_arch_falls_back_to_cross_compile_env(monkeypatch): + monkeypatch.setattr(arch, "is_rocm", lambda: True) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1200;gfx1201") + _clear_arch_caches() + + assert arch.get_rocm_gfx_arch() == "gfx1200" + + _clear_arch_caches() + + +def test_hip_cflags_emit_one_offload_flag_per_arch(monkeypatch): + from freetoken.kernel.utils import _hip_cflags + + monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1200;gfx1201") + + flags = _hip_cflags(["-Wno-unused-command-line-argument"]) + + assert "--offload-arch=gfx1200" in flags + assert "--offload-arch=gfx1201" in flags + assert not any(";" in flag for flag in flags) + + +def test_rocm_link_flags_support_versioned_modular_sdk(monkeypatch, tmp_path): + import torch.utils.cpp_extension as cpp_extension + + from freetoken.kernel import utils + + sdk = tmp_path / "sdk" + library_dir = sdk / "lib" + library_dir.mkdir(parents=True) + versioned_runtime = library_dir / "libamdhip64.so.7" + versioned_runtime.write_bytes(b"") + real_find_spec = importlib.util.find_spec + + def find_spec(name: str): + if name == "_rocm_sdk_core": + return SimpleNamespace(submodule_search_locations=[str(sdk)]) + return real_find_spec(name) + + monkeypatch.delenv("ROCM_HOME", raising=False) + monkeypatch.setattr(cpp_extension, "ROCM_HOME", None) + monkeypatch.setattr(importlib.util, "find_spec", find_spec) + monkeypatch.setattr(pathlib.Path, "home", lambda: tmp_path) + utils._rocm_link_flags.cache_clear() + + flags = utils._rocm_link_flags() + + compat_root = tmp_path / ".cache" / "freetoken" / "rocm-lib" + compat_dir = pathlib.Path(next(flag[2:] for flag in flags if flag.startswith("-L"))) + compat_link = compat_dir / "libamdhip64.so" + assert compat_dir.parent == compat_root + assert f"-L{compat_dir}" in flags + assert f"-Wl,-rpath,{library_dir}" in flags + assert compat_link.resolve() == versioned_runtime.resolve() + + utils._rocm_link_flags.cache_clear()