From 974164632d5980e09db1432c654f2e816c298066 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Tue, 25 Aug 2026 19:15:39 -0300 Subject: [PATCH 01/17] feat(rocm): AMD ROCm support (gfx1100) + qwen35moe GGUF loader Bring FreeToken up on AMD ROCm (RX 7900 XTX / gfx1100) alongside CUDA via a thin device seam. Includes: - Device detection / architecture gating (is_rocm), build toolchain for tvm-ffi/HIP JIT, --offload-arch=gfx1100 pinning, and HIP backend dispatch. - Pinned-memory + graph-capture gating (ROCm settles to kernel-launch decode). - device_api.h seam, GGUF kernel HIP port, quant mapping, nvfp4->mxfp4. - qwen3_5_moe GGUF loader (config/dense weights/expert offload). - fix(gguf): de-interleave the GDN mrope_interleaved value heads (in_proj_qkvz v/z rows, out_proj cols, in_proj_ba, conv1d v-channels, dt_bias; A_log stored as A=-exp(A_log)) so weights match HF and the model serves correct text. - torch attention backend (ground-truth reference) + ROCm regression tests. - AOT CI (rocm.yml) and docs/install-amd.md. Tests: ROCm suite passes; qwen35voe de-interleave + torch-backend tests added. --- .github/workflows/rocm.yml | 70 +++ .gitignore | 6 + docs/install-amd.md | 83 +++ freetoken-kernel-cache/build_backend.py | 18 +- pyproject.toml | 8 + python/freetoken/attention/__init__.py | 13 + python/freetoken/attention/torch.py | 217 ++++++++ python/freetoken/engine/engine.py | 26 +- python/freetoken/engine/graph.py | 16 +- python/freetoken/kernel/_toolchain.py | 86 +++ python/freetoken/kernel/backend.py | 29 + python/freetoken/kernel/batch_memcpy.py | 7 +- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 32 +- python/freetoken/kernel/csrc/gguf/dispatch.h | 8 + .../freetoken/kernel/csrc/gguf/ggml-common.h | 7 + .../freetoken/kernel/csrc/gguf/gguf_kernel.cu | 38 +- .../csrc/include/freetoken/device_api.h | 99 ++++ .../kernel/csrc/include/freetoken/utils.cuh | 94 ++++ .../kernel/csrc/jit/batch_memcpy.cuh | 64 ++- .../kernel/csrc/jit/fast_index_copy.cuh | 56 +- .../freetoken/kernel/csrc/pinned_tensor.cpp | 70 ++- python/freetoken/kernel/gguf.py | 34 +- python/freetoken/kernel/tinygrad_fallback.py | 92 ++++ python/freetoken/kernel/triton/activation.py | 31 +- python/freetoken/kernel/triton/norm.py | 34 +- python/freetoken/kernel/utils.py | 83 ++- python/freetoken/layers/moe.py | 10 + python/freetoken/models/gguf/config.py | 1 + python/freetoken/models/gguf/dequant.py | 69 +++ python/freetoken/models/gguf/tokenizer.py | 5 +- .../freetoken/models/qwen3_5_moe/__init__.py | 12 + python/freetoken/models/qwen3_5_moe/gdn.py | 14 +- python/freetoken/models/qwen3_5_moe/gguf.py | 512 ++++++++++++++++++ python/freetoken/models/qwen3_5_moe/model.py | 7 + python/freetoken/models/register.py | 8 + python/freetoken/models/weight.py | 14 + python/freetoken/moe/expert_banks.py | 20 + python/freetoken/moe/fused_gguf.py | 52 ++ python/freetoken/moe/nvfp4_backends.py | 12 + python/freetoken/moe/nvfp4_to_mxfp4.py | 239 ++++++++ python/freetoken/moe/offload_cache.py | 6 + python/freetoken/server/args.py | 19 + python/freetoken/utils/__init__.py | 8 + python/freetoken/utils/arch.py | 100 +++- python/freetoken/utils/graph_gate.py | 189 +++++++ python/freetoken/utils/torch_utils.py | 19 + setup.py | 80 ++- tests/attention/test_torch_backend.py | 62 +++ tests/engine/test_attention_backend_rocm.py | 75 +++ tests/kernels/test_backend_rocm.py | 63 +++ tests/kernels/test_cache_rocm_pairing.py | 109 ++++ tests/kernels/test_toolchain_hip.py | 84 +++ tests/models/test_qwen35moe_gguf_deint.py | 71 +++ tests/moe/test_nvfp4_backends_rocm.py | 95 ++++ tests/moe/test_nvfp4_to_mxfp4.py | 111 ++++ tests/utils/test_device_kind.py | 83 +++ 56 files changed, 3379 insertions(+), 91 deletions(-) create mode 100644 .github/workflows/rocm.yml create mode 100644 docs/install-amd.md create mode 100644 python/freetoken/attention/torch.py create mode 100644 python/freetoken/kernel/csrc/include/freetoken/device_api.h create mode 100644 python/freetoken/kernel/tinygrad_fallback.py create mode 100644 python/freetoken/models/qwen3_5_moe/gguf.py create mode 100644 python/freetoken/moe/fused_gguf.py create mode 100644 python/freetoken/moe/nvfp4_to_mxfp4.py create mode 100644 python/freetoken/utils/graph_gate.py create mode 100644 tests/attention/test_torch_backend.py create mode 100644 tests/engine/test_attention_backend_rocm.py create mode 100644 tests/kernels/test_backend_rocm.py create mode 100644 tests/kernels/test_cache_rocm_pairing.py create mode 100644 tests/kernels/test_toolchain_hip.py create mode 100644 tests/models/test_qwen35moe_gguf_deint.py create mode 100644 tests/moe/test_nvfp4_backends_rocm.py create mode 100644 tests/moe/test_nvfp4_to_mxfp4.py create mode 100644 tests/utils/test_device_kind.py diff --git a/.github/workflows/rocm.yml b/.github/workflows/rocm.yml new file mode 100644 index 000000000..e2f1fe49a --- /dev/null +++ b/.github/workflows/rocm.yml @@ -0,0 +1,70 @@ +name: ROCm (AMD) correctness smoke + +# ROCm CI job: compiles the AOT kernel cache for the RX 7000 (gfx1100) target on a +# ROCm torch install and runs a torch-free correctness smoke plus the AMD unit tests. +# Gated on a self-hosted runner that has ROCm torch + hipcc. The primary NVIDIA release +# flow is release.yml; this job is additive and must not gate NVIDIA releases. + +on: + workflow_dispatch: + push: + branches: [main] + pull_request: + +jobs: + rocm-smoke: + runs-on: [self-hosted, linux, amd, rocm] + timeout-minutes: 60 + env: + FREETOKEN_DISABLE_JIT: "1" + FREETOKEN_KERNEL_CACHE_GFX: "gfx1100" + steps: + - uses: actions/checkout@v4 + + - name: Check ROCm toolchain + run: | + set -e + command -v hipcc || ls /opt/rocm/bin/hipcc + "${PYTHON:-python3}" -c "import torch.version as v; print('torch hip:', v.hip)" + + - name: Install build deps + run: | + python -m pip install --upgrade pip wheel setuptools + python -m pip install -e "python[rocm]" + + - name: Compile AOT kernel cache for gfx1100 + run: | + FREETOKEN_KERNEL_CACHE_VERBOSE=1 python -m pip wheel ./freetoken-kernel-cache -w dist/cache-rocm + + - name: Install prebuilt kernel cache + run: | + whl="$(find dist/cache-rocm -name 'freetoken_kernel_cache-*.whl' | head -1)" + python -m pip install --force-reinstall "$whl" + + - name: Torch-free AMD unit tests + run: | + python -m pytest \ + tests/utils/test_device_kind.py \ + tests/kernels/test_toolchain_hip.py \ + tests/kernels/test_backend_rocm.py \ + tests/kernels/test_cache_rocm_pairing.py \ + tests/moe/test_nvfp4_to_mxfp4.py \ + -q + + - name: Hardware correctness smoke (serves on RX 7000) + run: | + # Functional path only -- flashinfer/sgl/trtllm are NVIDIA-only and must not + # be selected. AUTO backend must resolve to triton; NVFP4 auto -> triton. + python - <<'PY' + from freetoken.utils.arch import is_rocm, is_gfx_arch_ge + from freetoken.moe.nvfp4_backends import select_nvfp4_backend + import torch + assert is_rocm(), "expected a ROCm torch build" + assert is_gfx_arch_ge(1100), "expected gfx1100-class device (RX 7000)" + print("NVFP4 auto ->", select_nvfp4_backend(torch.device("cuda"), 768, "auto")) + PY + + - name: Serve smoke + run: | + FREETOKEN_DEVICE=cuda python -m freetoken.serve --help >/dev/null \ + && echo "freetoken CLI loads on ROCm" diff --git a/.gitignore b/.gitignore index bf804e075..95757625b 100644 --- a/.gitignore +++ b/.gitignore @@ -227,3 +227,9 @@ benchmarks/cross_framework # local e2e/bench artifacts (harnesses may run with repo cwd) /results/ + +# torch.utils.cpp_extension ROCm hipify build artifacts (preprocessed .cu -> .hip, +# and includes rewritten to *_hip.cuh) written next to the GGUF sources. +python/freetoken/kernel/csrc/gguf/*.hip +python/freetoken/kernel/csrc/gguf/*_hip.cuh +python/freetoken/kernel/csrc/gguf/ggml-common_hip.h diff --git a/docs/install-amd.md b/docs/install-amd.md new file mode 100644 index 000000000..62ecf2fe7 --- /dev/null +++ b/docs/install-amd.md @@ -0,0 +1,83 @@ +# AMD GPU (ROCm) support + +FreeToken targets Linux + NVIDIA CUDA by default. AMD (ROCm) is a supported, tested +configuration with a **single-GPU** milestone: correct functional path first, performance +recovered via HIP ports where safe. This page covers installing and running on RX 7000. + +> Status: **experimental.** The default and best-tested path remains CUDA. AMD brings up a +> correct functional path (Triton attention + offload/CPU MoE + portable quant) and is +> recovering performance via the HIP kernel ports. See `.plans/amd-gpu-support/plan.md`. + +## Requirements + +| Component | Requirement | +| --- | --- | +| OS | Linux x86_64 (Windows WDDM pinned-memory is a known edge, not supported yet) | +| GPU | AMD RX 7000 (RDNA 3, `gfx1100`); RX 9000 (`gfx1201`) is future work | +| ROCm | ROCm toolkit with `hipcc` (`/opt/rocm/bin/hipcc` or on `PATH`) | +| torch | ROCm build, e.g. `torch==2.5.1+rocm6.2` | + +The build refuses to mix toolchains: it will **not** silently fall back to `nvcc`/`libcudart` +when only the ROCm toolkit is present, and vice versa. + +## Install + +```bash +# ROCm torch (PyTorch official ROCm wheels) -- must satisfy the repo's torch>=2.11,<2.12 +# build pin, so use the rocm7.2 index (rocm6.2 only carries torch up to 2.5.1). +pip install --index-url https://download.pytorch.org/whl/rocm7.2 \ + "torch==2.11.0+rocm7.2" torchvision triton-rocm==3.6.0 + +# FreeToken with the ROCm extra (builds the native extensions with hipcc) +uv pip install -e ".[rocm]" --no-build-isolation +``` + +`pip install ".[rocm]"` pulls ROCm-compatible `torch`/`triton`; the NVIDIA-only `[accel]` +packages (`flashinfer`, `sgl-kernel`, `triton_kernels`, Marlin) are **not** installed on AMD +and their backends are rejected with a clean error if requested. + +## Verified feature matrix + +| Feature | On AMD | Notes | +| --- | --- | --- | +| Attention | `--attention-backend triton` | flashinfer/fa/trtllm are NVIDIA-only and rejected | +| MoE | `--moe-backend offload / cpu / hybrid` | offload needs pinned host memory (Inc 3) | +| Quant | BF16, MXFP4, GGUF (Q4_K/Q8_0), Triton inline-dequant NVFP4 | Marlin INT4 / native NVFP4 SASS unavailable | +| NVFP4 checkpoints with no MXFP4 variant | converted to MXFP4 on load (auto) | `--nvfp4-backend auto` → triton/MXFP4 | +| CUDA graphs (decode) | HIP graph capture **if** the Inc-1 gate passes | otherwise kernel-launch decode | +| Multi-GPU (RCCL) | out of scope (single-GPU milestone) | | + +## CLI behavior on AMD + +* `--nvfp4-backend marlin` / `flashinfer` → error (NVIDIA-only). Use `triton` / `auto`. +* `--attention-backend fi` / `fa` / `trtllm` → error (NVIDIA-only). Use `triton` / `auto`. +* `--moe-backend fused` → warning (fused MoE is CUDA-only; falls back to offload/cpu). +* `--nvfp4-backend auto` → resolves to the portable Triton inline-dequant path (or MXFP4 + for a converted checkpoint). + +## Verify + +```bash +ft version # prints an AMD / ROCm banner +ft serve --model Qwen3.6-35B-A3B \ + --moe-backend offload --attention-backend triton --nvfp4-backend auto +``` + +`ldd` of the built `.so` should show `hiprt`/`amdhip64`, not `libcudart`. + +## AOT kernel cache + +Build the prebuilt `+rocm` kernel-cache wheel (no nvcc needed on the target): + +```bash +scripts/build-release-wheels.sh # on a ROCm torch + hipcc box; tags the cache +rocm +``` + +The runtime refuses to pair a `+rocm` cache with a `+cu130` runtime (and vice versa). + +## Notes / limitations + +* `nvtx_annotate` is a no-op on ROCm; roctx profiling is future work. +* FP8 / NVFP4-class formats: BF16 / MXFP4 / GGUF are the supported AMD matrix; performance + parity vs CUDA is not guaranteed for NVFP4-class formats. +* Windows AMD is not yet supported (WDDM zero-copy semantics differ). diff --git a/freetoken-kernel-cache/build_backend.py b/freetoken-kernel-cache/build_backend.py index 03f17a314..cf255becf 100644 --- a/freetoken-kernel-cache/build_backend.py +++ b/freetoken-kernel-cache/build_backend.py @@ -40,6 +40,11 @@ def _cuda_version_suffix() -> str: return "" cuda_version = getattr(torch.version, "cuda", None) + hip_version = getattr(torch.version, "hip", None) + if hip_version: + # ROCm torch: torch.version.cuda is None; tag the cache with +rocm so it pairs + # only with a ROCm runtime (kernel/utils.py._arch_tags enforces the match). + return "+rocm" if not cuda_version: return "" # The tag advertises torch's CUDA; the cache .so link nvcc's libcudart. @@ -110,7 +115,18 @@ def _build_jit_cache() -> None: # 12.0 -> RTX 50 series, RTX PRO 6000 Blackwell (Blackwell, consumer / workstation) # Override with FREETOKEN_KERNEL_CACHE_ARCHES (space-separated maj.min) or # TVM_FFI_CUDA_ARCH_LIST directly. Needs an nvcc that supports every listed arch. - if "TVM_FFI_CUDA_ARCH_LIST" not in os.environ: + try: + import torch # noqa: PLC0415 + + is_rocm_build = bool(getattr(torch.version, "hip", None)) + except Exception: + is_rocm_build = False + if is_rocm_build: + # ROCm: the CUDA arch-list is meaningless; the gfx arch is passed through + # kernel/utils.py._arch_flags (--offload-arch), defaulting to the RX 7000 + # (gfx1100) target. Env override for other RX 7000 SKUs / future archs. + os.environ.setdefault("FREETOKEN_KERNEL_CACHE_GFX", "gfx1100") + elif "TVM_FFI_CUDA_ARCH_LIST" not in os.environ: os.environ["TVM_FFI_CUDA_ARCH_LIST"] = os.getenv( "FREETOKEN_KERNEL_CACHE_ARCHES", "8.0 8.6 8.9 9.0 10.0 12.0" ) diff --git a/pyproject.toml b/pyproject.toml index 8bd653f87..13b52c19f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ classifiers = [ "Intended Audience :: Science/Research", "Operating System :: POSIX :: Linux", "Environment :: GPU :: NVIDIA CUDA", + "Environment :: GPU :: AMD ROCm", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", @@ -78,6 +79,13 @@ fi = ["flashinfer-python[cu13]>=0.6,<0.7"] # renamed from sgl-kernel at 0.4; still imports as `sgl_kernel`, so never co-install both sgl = ["sglang-kernel==0.4.5"] accel = ["freetoken[fi,sgl]"] +# ROCm (AMD) install: the NVIDIA-only fi/sgl/Marlin packages are NOT pulled in. torch must +# come from the ROCm wheel index (e.g. `pip install torch==2.11.0+rocm6.2` from +# https://download.pytorch.org/whl/rocm6.2) so the native extensions build against the HIP +# runtime; this extra pins the rest. See docs/install.md (AMD section, Inc 9). +rocm = [ + "triton==3.6.0; platform_system == 'Linux'", +] # NOTE: the Marlin W4A16 NVFP4 expert-GEMM path (sm_80-99) borrows vLLM's AOT wheel # (vllm>=0.14,<0.15), which pins transformers>=4.56,<5 and so is INCOMPATIBLE with the # core transformers>=5.5 requirement. It is therefore not a lockable extra and is left diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index 746c04c4b..0dfb21972 100644 --- a/python/freetoken/attention/__init__.py +++ b/python/freetoken/attention/__init__.py @@ -96,6 +96,19 @@ def create_triton_backend(config: ModelConfig): return TritonAttentionBackend(config) +@SUPPORTED_ATTENTION_BACKENDS.register( + "torch", + BackendInfo( + supported_types=frozenset({AttnType.FULL}), + # Debugging/eager ground-truth backend; no package/arch requirements. + ), +) +def create_torch_backend(config: ModelConfig): + from .torch import TorchAttentionBackend + + return TorchAttentionBackend(config) + + @SUPPORTED_ATTENTION_BACKENDS.register( "dsv4_sparse", BackendInfo(supported_types=frozenset({AttnType.DSV4})), diff --git a/python/freetoken/attention/torch.py b/python/freetoken/attention/torch.py new file mode 100644 index 000000000..6ae16c2e3 --- /dev/null +++ b/python/freetoken/attention/torch.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, List + +import torch + +from freetoken.core import Batch, get_global_ctx + +from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata + +if TYPE_CHECKING: + from freetoken.models import ModelConfig + + +@dataclass +class TorchMetadata(BaseAttnMetadata): + """Minimal contiguous-gather metadata for the pure-torch backend. + + ``indices`` maps every logical KV position (across all padded requests, in + ``seqlens_k`` order) to its physical paged-cache slot, exactly like the triton + backend's gather. The torch backend reads the SAME paged cache as triton, so a + triton-vs-torch logit difference isolates the attention *compute* from the + cache addressing. + """ + + indices: torch.Tensor + seqlens_q: List[int] + seqlens_k: List[int] + cached_lens: List[int] + is_decode: bool + cu_seqlens_q: torch.Tensor + + def get_last_indices(self, bs: int) -> torch.Tensor: + return self.cu_seqlens_q[1 : 1 + bs] - 1 + + +class TorchAttentionBackend(BaseAttnBackend): + """Pure-PyTorch full-attention backend (no Triton/CUDA kernels). + + Serves as a numerically-explicit ground truth for debugging the hybrid + qwen35moe model. It stores K/V into the same paged MHAKVCache as + ``TritonAttentionBackend``, gathers the request's full K/V history via the + identical ``indices`` page gather, and computes GQA softmax attention with + PyTorch ops so every intermediate is auditable. + + Intended for correctness debugging / backend A-B comparison, not production + serving. Registered as the ``"torch"`` attention backend (``AttnType.FULL``). + """ + + def __init__(self, config: ModelConfig): + self.config = config + self.kvcache = get_global_ctx().kv_cache + self.device = self.kvcache.device + self.num_q_heads = int(getattr(config, "num_qo_heads", 1)) + self.num_kv_heads = int(getattr(config, "num_kv_heads", 1)) + self.head_dim = int(getattr(config, "head_dim", 1)) + # Prefer the full-attention group spec head_dim (authoritative for kv heads). + specs = getattr(config, "kv_cache_group_specs", lambda: ())() + for spec in specs: + name = getattr(spec, "attn_type", None) + if name is not None and str(name) == "AttnType.FULL": + self.head_dim = int(getattr(spec, "head_dim", self.head_dim)) + self.num_kv_heads = int(getattr(spec, "num_kv_heads", self.num_kv_heads)) + break + # Debugging: contiguous (per-request) cache instead of the paged pool, to + # isolate cache addressing from the attention compute (Inc 5). + self._contig: dict[tuple[int, int], list] = {} + import os + + self._use_contig = os.environ.get("FT_DEBUG_CONTIG_CACHE") == "1" + + def _build_metadata(self, batch: Batch) -> TorchMetadata: + ctx = get_global_ctx() + page_table = ctx.page_table + reqs = batch.padded_reqs + seqlens_q = [req.extend_len for req in reqs] + seqlens_k = [req.device_len for req in reqs] + cached_lens = [req.cached_len for req in reqs] + is_decode = max(seqlens_q) == 1 + indices = torch.cat([page_table[req.table_idx, : req.device_len] for req in reqs]) + if is_decode: + cu_seqlens_q = torch.arange(0, len(reqs) + 1, dtype=torch.int32, device=self.device) + else: + cu_seqlens_q = torch.tensor( + [0] + seqlens_q, dtype=torch.int32, device=self.device + ).cumsum_(0) + return TorchMetadata( + indices=indices, + seqlens_q=seqlens_q, + seqlens_k=seqlens_k, + cached_lens=cached_lens, + is_decode=is_decode, + cu_seqlens_q=cu_seqlens_q, + ) + + def prepare_metadata(self, batch: Batch) -> None: + batch.attn_metadata = self._build_metadata(batch) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer_id: int, + batch: Batch, + attn_spec: AttentionSpec | None = None, + ) -> torch.Tensor: + if self._use_contig: + return self._forward_contig(q, k, v, layer_id, batch, attn_spec) + self.kvcache.store_kv(k, v, batch.out_loc, layer_id) + + k_raw = self.kvcache.k_cache(layer_id) + v_raw = self.kvcache.v_cache(layer_id) + kv_heads, head_dim = k_raw.shape[-2], k_raw.shape[-1] + assert head_dim == q.shape[-1], f"head_dim {head_dim} != {q.shape[-1]}" + k_cache = k_raw.view(-1, kv_heads, head_dim) + v_cache = v_raw.view(-1, kv_heads, head_dim) + + metadata = batch.attn_metadata + assert isinstance(metadata, TorchMetadata) + k_all = k_cache[metadata.indices] # [total_kv, kv_heads, head_dim] + v_all = v_cache[metadata.indices] # [total_kv, kv_heads, head_dim] + + spec = attn_spec or AttentionSpec() + scale = spec.sm_scale if spec.sm_scale is not None else head_dim ** -0.5 + group = self.num_q_heads // kv_heads + + num_q_tokens = q.shape[0] + out = torch.empty((num_q_tokens, self.num_q_heads, head_dim), dtype=q.dtype, device=q.device) + q_off = 0 + k_off = 0 + for lq, lk, cached in zip( + metadata.seqlens_q, metadata.seqlens_k, metadata.cached_lens + ): + qs = q[q_off : q_off + lq] # [lq, num_q, head_dim] + ks = k_all[k_off : k_off + lk] # [lk, kv_heads, head_dim] + vs = v_all[k_off : k_off + lk] + if group > 1: + ks = ks.repeat_interleave(group, dim=1) # [lk, num_q, head_dim] + vs = vs.repeat_interleave(group, dim=1) + # [num_q, lq, lk] + scores = torch.einsum("qhd,khd->hqk", qs.float(), ks.float()) * scale + # causal: query i (global cached+i) attends key col j <= cached+i + if lq > 1 or lk > lq: + rows = torch.arange(lq, device=scores.device) + cols = torch.arange(lk, device=scores.device) + masked = (cols[None, :] > (cached + rows)[:, None]).to(scores.device) + scores = scores.masked_fill(masked[None, :, :], float("-inf")) + probs = torch.softmax(scores, dim=-1) + o = torch.einsum("hqk,khd->qhd", probs, vs.float()).to(q.dtype) # [lq, num_q, head_dim] + out[q_off : q_off + lq] = o + q_off += lq + k_off += lk + + return out + + def _forward_contig(self, q, k, v, layer_id, batch, attn_spec=None): + """Contiguous (non-paged) attention: K/V accumulate per (layer, request) in + a Python list keyed by logical position, so cache addressing is trivially + correct. Isolates the paged-cache addressing from the attention compute.""" + metadata = self._build_metadata(batch) + kv_heads = self.num_kv_heads + head_dim = self.head_dim + spec = attn_spec or AttentionSpec() + scale = spec.sm_scale if spec.sm_scale is not None else head_dim ** -0.5 + group = self.num_q_heads // kv_heads + # Store this forward's K/V rows per request (append in global position order). + q_off = 0 + for i, lq in enumerate(metadata.seqlens_q): + uid = batch.padded_reqs[i].uid + buf = self._contig.setdefault((layer_id, uid), {"k": [], "v": []}) + kseg = k[q_off : q_off + lq].view(lq, kv_heads, head_dim) + vseg = v[q_off : q_off + lq].view(lq, kv_heads, head_dim) + for t in range(lq): + buf["k"].append(kseg[t]) + buf["v"].append(vseg[t]) + q_off += lq + # compute attention from the contiguous cache + q_off = 0 + out = torch.empty( + (q.shape[0], self.num_q_heads, head_dim), dtype=q.dtype, device=q.device + ) + for i, lq in enumerate(metadata.seqlens_q): + uid = batch.padded_reqs[i].uid + buf = self._contig[(layer_id, uid)] + ks = torch.stack(buf["k"]) # [acc, kv_heads, head_dim] + vs = torch.stack(buf["v"]) + acc = ks.shape[0] + cached = acc - lq + qs = q[q_off : q_off + lq] + if group > 1: + ks = ks.repeat_interleave(group, dim=1) + vs = vs.repeat_interleave(group, dim=1) + scores = torch.einsum("qhd,khd->hqk", qs.float(), ks.float()) * scale + rows = torch.arange(lq, device=scores.device) + cols = torch.arange(acc, device=scores.device) + masked = (cols[None, :] > (cached + rows)[:, None]).to(scores.device) + scores = scores.masked_fill(masked[None, :, :], float("-inf")) + probs = torch.softmax(scores, dim=-1) + o = torch.einsum("hqk,khd->qhd", probs, vs.float()).to(q.dtype) + out[q_off : q_off + lq] = o + q_off += lq + return out + + def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + # ROCm/graph capture is disabled for this debugging backend; no-op. + return None + + def prepare_for_capture(self, batch: Batch) -> None: + return None + + def prepare_for_replay(self, batch: Batch) -> None: + return None + + +__all__ = ["TorchAttentionBackend", "TorchMetadata"] diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 22c4a6c7e..5ada3954a 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -15,7 +15,7 @@ from freetoken.moe import create_moe_backend, is_offload_moe_backend from freetoken.moe.expert_banks import load_expert_banks from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache -from freetoken.utils import align_ceil, init_logger, is_sm90_family, is_sm100_family, mem_GB, torch_dtype +from freetoken.utils import align_ceil, device_kind, init_logger, is_rocm, is_sm90_family, is_sm100_family, mem_GB, torch_dtype from .config import EngineConfig from .graph import GraphRunner, get_free_memory @@ -50,6 +50,11 @@ def _flashinfer_available() -> bool: def _sgl_flash_attn_available() -> bool: + from freetoken.utils.arch import is_rocm + + # sgl_kernel is NVIDIA-only; never select it on ROCm even if a stray copy is importable. + if is_rocm(): + return False try: from sgl_kernel.flash_attn import flash_attn_with_kvcache # noqa: F401 except Exception as exc: @@ -101,6 +106,14 @@ def _backend_parts_serve(name: str, required: frozenset[AttnType]) -> bool: def _backend_requirements_met(name: str) -> bool: + # On ROCm (AMD) only the portable backends exist: flashinfer/sgl/trtllm (and anything + # sm_100-gated) are NVIDIA-only, so short-circuit before probing them at all. + from freetoken.utils.arch import is_rocm + + if is_rocm(): + return all(not i.requires_flashinfer and not i.requires_sgl_kernel + and not i.requires_sm100 for i in + [attention_backend_info(p) for p in name.split(",")]) # flashinfer first across ALL parts: the sgl probe logs a "falls back to fi" warning, # which would mislead when the candidate is about to fail on flashinfer anyway. infos = [attention_backend_info(part) for part in name.split(",")] @@ -202,6 +215,13 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att # explicit --attention-backend choices. for part in backend_parts: info = attention_backend_info(part) + from freetoken.utils.arch import is_rocm + + if is_rocm() and (info.requires_flashinfer or info.requires_sgl_kernel or info.requires_sm100): + raise RuntimeError( + f"Attention backend {config.attention_backend!r} is NVIDIA-only and " + f"unavailable on this ROCm (AMD) build; use --attention-backend triton." + ) if info.requires_flashinfer and not _flashinfer_available(): raise RuntimeError( f"Attention backend {config.attention_backend!r} requires flashinfer, which is " @@ -298,6 +318,7 @@ def __init__(self, config: EngineConfig): self.device = torch.device(f"cuda:{config.tp_info.rank}") torch.cuda.set_device(self.device) + logger.info_rank0(f"device_kind={device_kind()} backend={self.device}") torch.manual_seed(42) self.stream = torch.cuda.Stream() torch.cuda.set_stream(self.stream) @@ -1011,6 +1032,9 @@ def _ensure_expandable_segments() -> None: """ if os.environ.get("PYTORCH_ALLOC_CONF") or os.environ.get("PYTORCH_CUDA_ALLOC_CONF"): return + if is_rocm(): + # expandable_segments is a CUDA allocator setting with no ROCm analogue; skip it. + return try: torch.cuda.memory._set_allocator_settings("expandable_segments:True") except Exception as exc: # pragma: no cover - depends on torch build diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 4f2025025..50bac2f0a 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -132,6 +132,18 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel # graphs-disabled early return so that config gets the phase too. emit_progress("Capturing CUDA graphs / warming up", 0, 0) self.graph_map: Dict[int, torch.cuda.CUDAGraph] = {} + # Inc-8 parity: on ROCm, honour the Inc-1 graph-gate result. If capture is not + # viable on this AMD card, skip graphs entirely so decode uses the kernel-launch + # path (correct, just not graph-accelerated) rather than erroring mid-capture. + from freetoken.utils.arch import is_rocm + from freetoken.utils.graph_gate import graph_capture_status + + if is_rocm() and graph_capture_status() == "fail": + logger.info_rank0( + "AMD ROCm build: HIP graph capture gate FAILED on this device; " + "using the kernel-launch decode path (CUDA graphs disabled)." + ) + return None if self.max_graph_bs == 0: return logger.info_rank0("CUDA graph is disabled.") @@ -187,7 +199,9 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel logger.info_rank0(f"Free GPU memory after capturing CUDA graphs: {mem_GB(free_memory)}") def can_use_cuda_graph(self, batch: Batch) -> bool: - return batch.is_decode and batch.size <= self.max_graph_bs + # ``self.graph_map`` is empty when graphs were skipped (ROCm graph-gate fail or + # disabled); decode must then fall back to the kernel-launch path. + return bool(self.graph_map) and batch.is_decode and batch.size <= self.max_graph_bs def replay(self, batch: Batch) -> torch.Tensor: assert self.can_use_cuda_graph(batch) diff --git a/python/freetoken/kernel/_toolchain.py b/python/freetoken/kernel/_toolchain.py index b49cebbb4..eccb96359 100644 --- a/python/freetoken/kernel/_toolchain.py +++ b/python/freetoken/kernel/_toolchain.py @@ -16,6 +16,92 @@ _TRUE_VALUES = {"1", "true", "yes", "on"} +def _hipcc_path() -> str | None: + """Locate hipcc: $ROCM_HOME/bin/hipcc, $HIP_PATH/bin/hipcc, /opt/rocm/bin/hipcc, + then PATH.""" + for env in ("ROCM_HOME", "HIP_PATH"): + root = os.getenv(env) + if root: + candidate = os.path.join(root, "bin", "hipcc") + if os.path.isfile(candidate): + return candidate + default = "/opt/rocm/bin/hipcc" + if os.path.isfile(default): + return default + return shutil.which("hipcc") + + +def hip_hip_version(hipcc: str) -> tuple[int, int] | None: + """HIP toolkit version from ``hipcc --version`` (e.g. (6, 2)), or None.""" + try: + proc = subprocess.run( + [hipcc, "--version"], capture_output=True, text=True, check=True + ) + except (OSError, subprocess.CalledProcessError): + return None + # hipcc --version prints e.g. "HIP version: 6.2.41000" (or a clang version line). + m = re.search(r"HIP version[:\s]+(\d+)\.(\d+)", proc.stdout) + if m: + return int(m.group(1)), int(m.group(2)) + m = re.search(r"(\d+)\.(\d+)\.\d+", proc.stdout) + if m: + return int(m.group(1)), int(m.group(2)) + return None + + +def torch_hip_version() -> str | None: + """The ``torch.version.hip`` string (e.g. "6.2.4100000"), or None on non-ROCm torch.""" + try: + import torch + + return getattr(torch.version, "hip", None) + except Exception: + return None + + +def is_rocm_torch() -> bool: + """True when the installed torch is a ROCm (AMD) build.""" + return bool(torch_hip_version()) + + +def torch_hip_major() -> int | None: + hip = torch_hip_version() + if not hip: + return None + m = re.match(r"(\d+)", hip) + return int(m.group(1)) if m else None + + +def check_hip_matches_torch() -> None: + """Refuse to hipcc-compile kernels across HIP major versions. + + Mirrors check_nvcc_matches_torch: hipcc-built kernels link against the HIP runtime + major they were built with; at runtime only the torch wheel's own HIP runtime is + guaranteed to be loadable. No-op when torch is not ROCm. + """ + if os.getenv(ALLOW_MISMATCH_ENV, "").strip().lower() in _TRUE_VALUES: + return + if not is_rocm_torch(): + return + torch_major = torch_hip_major() + hipcc = _hipcc_path() + if hipcc is None: + raise RuntimeError( + "ROCm torch detected but no hipcc found. Install a ROCm/HIP toolkit " + "(e.g. via /opt/rocm) matching torch's HIP version, or set " + f"{ALLOW_MISMATCH_ENV}=1 to override." + ) + release = hip_hip_version(hipcc) + if release is None: + return + if release[0] != torch_major: + raise RuntimeError( + f"hipcc {release[0]}.{release[1]} would build kernels linking HIP " + f"{release[0]}.x, but torch ships HIP {torch_hip_version()}. Install a " + f"ROCm {torch_major}.x toolkit, or set {ALLOW_MISMATCH_ENV}=1 to override." + ) + + def _nvcc_path() -> str | None: from torch.utils.cpp_extension import CUDA_HOME diff --git a/python/freetoken/kernel/backend.py b/python/freetoken/kernel/backend.py index 3037ad8d7..7b7629b11 100644 --- a/python/freetoken/kernel/backend.py +++ b/python/freetoken/kernel/backend.py @@ -10,6 +10,14 @@ import functools import importlib.util +from freetoken.utils.arch import is_rocm + + +# NVIDIA-only optional native packages: even if an importable copy is present on a ROCm +# torch build (e.g. a stray CUDA wheel), they must not be used -- the runtime falls back +# to the portable Triton kernels. Treated as unavailable on ROCm. +_CUDA_ONLY_PACKAGES = frozenset({"flashinfer", "sgl_kernel", "triton_kernels"}) + def _importable(name: str) -> bool: # find_spec normally returns None when a package is absent, but it can raise @@ -21,13 +29,32 @@ def _importable(name: str) -> bool: return False +def is_native_cuda_available() -> bool: + """True when the current torch build is CUDA and a CUDA-capable device is present. + False on ROCm and CPU builds. Used to gate NVIDIA-native ops/paths.""" + from freetoken.utils.arch import device_kind + + if device_kind() != "cuda": + return False + try: + import torch + + return bool(torch.cuda.is_available()) + except Exception: + return False + + @functools.cache def is_flashinfer_installed() -> bool: + if is_rocm(): + return False return _importable("flashinfer") @functools.cache def is_sgl_kernel_installed() -> bool: + if is_rocm(): + return False return _importable("sgl_kernel") @@ -39,6 +66,8 @@ def is_triton_kernels_installed() -> bool: source tree and has no Windows wheel. It is also not one of the six ops ``freetoken.kernel.triton`` reimplements, so its call-site carries its own fallback. """ + if is_rocm(): + return False return _importable("triton_kernels") diff --git a/python/freetoken/kernel/batch_memcpy.py b/python/freetoken/kernel/batch_memcpy.py index b39e5cde2..0466333c9 100644 --- a/python/freetoken/kernel/batch_memcpy.py +++ b/python/freetoken/kernel/batch_memcpy.py @@ -42,10 +42,15 @@ def _probe(fn) -> None: def load_batch_memcpy(): """Build (once), probe, and return the batch-memcpy entry point, or raise. - The 8-argument cudaMemcpyBatchAsync signature this binding uses is CUDA 13.0's + On ROCm the HIP per-copy grid kernel is used (no CUDA version gate). On CUDA the + 8-argument cudaMemcpyBatchAsync signature this binding uses is CUDA 13.0's (12.8/12.9 had an extra failIdx parameter); gate on the torch runtime version before paying for the JIT build, then verify with a real copy. """ + if torch.version.hip is not None: + fn = _jit_batch_memcpy_module().batch_memcpy + _probe(fn) + return fn cuda = torch.version.cuda if cuda is None or tuple(int(x) for x in cuda.split(".")[:2]) < (13, 0): raise RuntimeError(f"cudaMemcpyBatchAsync binding requires CUDA >= 13.0 (torch built with {cuda})") diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 880e8637a..024bf5a05 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -29,9 +29,31 @@ #include #include +#if defined(USE_HIP) +#include +// HIP host functions carry no special calling convention; define the CUDA host-func +// marker as empty so the static callback signatures below compile unchanged. +#ifndef CUDART_CB +#define CUDART_CB +#endif +#else #include +#endif #include +// Stream-sync / host-func node launch, shared by both backends. +#if defined(USE_HIP) +#define CPU_MOE_STREAM_SYNC(s) \ + hipStreamSynchronize(reinterpret_cast(s)) +#define CPU_MOE_LAUNCH_HOST_FUNC(s, fn, data) \ + hipLaunchHostFunc(reinterpret_cast(s), (fn), (data)) +#else +#define CPU_MOE_STREAM_SYNC(s) \ + cudaStreamSynchronize(reinterpret_cast(s)) +#define CPU_MOE_LAUNCH_HOST_FUNC(s, fn, data) \ + cudaLaunchHostFunc(reinterpret_cast(s), (fn), (data)) +#endif + #if defined(__linux__) #include #include @@ -614,7 +636,7 @@ static bool cumemops_probe(uintptr_t stream, uintptr_t scratch_addr) { auto* s = reinterpret_cast(stream); if (g_cu_write64(s, (unsigned long long)scratch_addr, 7ULL, kCuWriteDefault) != 0) return false; if (g_cu_wait64(s, (unsigned long long)scratch_addr, 7ULL, kCuWaitValueGeq) != 0) return false; - return cudaStreamSynchronize(reinterpret_cast(stream)) == cudaSuccess; + return CPU_MOE_STREAM_SYNC(stream) == 0; } // GPU side of the flag handshake (see the block comment above): enqueued on the @@ -1958,13 +1980,13 @@ struct CpuMoeExecutor { } void submit_with_cuda_stream(uintptr_t stream, uintptr_t task) { - cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::submit_cb, - reinterpret_cast(task)); + CPU_MOE_LAUNCH_HOST_FUNC(stream, &CpuMoeExecutor::submit_cb, + reinterpret_cast(task)); } void sync_with_cuda_stream(uintptr_t stream, uintptr_t task) { - cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::sync_cb, - reinterpret_cast(task)); + CPU_MOE_LAUNCH_HOST_FUNC(stream, &CpuMoeExecutor::sync_cb, + reinterpret_cast(task)); } // Register a (layer, batch-size) slot's task so the coordinator can dispatch it on a diff --git a/python/freetoken/kernel/csrc/gguf/dispatch.h b/python/freetoken/kernel/csrc/gguf/dispatch.h index f42a21633..d46469efb 100644 --- a/python/freetoken/kernel/csrc/gguf/dispatch.h +++ b/python/freetoken/kernel/csrc/gguf/dispatch.h @@ -11,6 +11,13 @@ #endif // Warp-shuffle wrappers the donor pulls from sgl-kernel's utils.h (CUDA variants). +// On ROCm the mask must be 64-bit (HIP static-asserts 32-bit promotion is an error). +#if defined(USE_ROCM) +#define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) \ + __shfl_xor_sync(static_cast(mask), (var), (lane_mask)) +#define SGLANG_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \ + __shfl_xor_sync(static_cast(mask), (var), (lane_mask), (width)) +#else #ifndef SGLANG_SHFL_XOR_SYNC #define SGLANG_SHFL_XOR_SYNC(mask, var, lane_mask) __shfl_xor_sync((mask), (var), (lane_mask)) #endif @@ -18,6 +25,7 @@ #define SGLANG_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \ __shfl_xor_sync((mask), (var), (lane_mask), (width)) #endif +#endif #define DISPATCH_CASE_FLOAT_TYPES(...) \ AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ diff --git a/python/freetoken/kernel/csrc/gguf/ggml-common.h b/python/freetoken/kernel/csrc/gguf/ggml-common.h index 88c21a4ab..5822eb3cb 100644 --- a/python/freetoken/kernel/csrc/gguf/ggml-common.h +++ b/python/freetoken/kernel/csrc/gguf/ggml-common.h @@ -10,6 +10,13 @@ #define GGML_CUDA_DMMV_X 32 #define GGML_CUDA_MMV_Y 1 +#if defined(USE_ROCM) +// ROCm shim: the vendored GGUF launchers (moe.cuh/mmvq.cuh/...) take a CUDA-style +// stream parameter. Map the CUDA stream type to HIP so those signatures compile +// unmodified under USE_ROCM. The including .cu pulls in hip/hip_runtime.h first. +using cudaStream_t = hipStream_t; +#endif + // Data Structures // QK = number of values after dequantization // QR = QK / number of values before dequantization diff --git a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu index d88960d5f..09c83c6f2 100644 --- a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu +++ b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu @@ -1,10 +1,28 @@ // Adatped from // https://github.com/vllm-project/vllm/blob/755ed7b05be4743237d3339c4ff8c22bcaae04f4/csrc/quantization/gguf/gguf_kernel.cu +#if defined(USE_ROCM) +// ROCm torch hipifies these headers into c10::cuda (masquerading-as-CUDA), providing +// c10::cuda::OptionalCUDAGuard / getCurrentCUDAStream backed by HIP. c10/cuda/CUDAGuard.h +// itself is not directly includable on ROCm (missing a generated header). +#include +#include +#include +#include +#else #include #include #include +#endif #include +#if defined(USE_ROCM) +#define GGUF_DEVICE_GUARD(device) c10::cuda::OptionalCUDAGuard device_guard(device) +#define GGUF_CURRENT_STREAM() c10::cuda::getCurrentCUDAStream() +#else +#define GGUF_DEVICE_GUARD(device) at::cuda::OptionalCUDAGuard device_guard(device) +#define GGUF_CURRENT_STREAM() at::cuda::getCurrentCUDAStream() +#endif + // dont use clang-format here, it breaks the include order // clang-format off #include "dispatch.h" @@ -77,11 +95,11 @@ torch::Tensor ggml_dequantize( int64_t m, int64_t n, std::optional const& dtype) { - const at::cuda::OptionalCUDAGuard device_guard(device_of(W)); + const GGUF_DEVICE_GUARD(device_of(W)); auto dtype_ = dtype.value_or(torch::kFloat16); auto options = torch::TensorOptions().dtype(dtype_).device(W.device()); at::Tensor DW = torch::empty({m, n}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); DISPATCH_FLOAT_TYPES(DW.scalar_type(), "ggml_dequantize", [&] { auto to_cuda = ggml_get_to_cuda(type); @@ -99,10 +117,10 @@ torch::Tensor ggml_mul_mat_vec_a8( int col = X.sizes()[1]; int vecs = X.sizes()[0]; const int padded = (col + 512 - 1) / 512 * 512; - const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); + const GGUF_DEVICE_GUARD(device_of(X)); auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); at::Tensor Y = torch::empty({vecs, row}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); at::Tensor quant_X = torch::empty({vecs, padded / 32 * 9}, options); DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_mul_mat_vec_a8", [&] { @@ -197,10 +215,10 @@ torch::Tensor ggml_mul_mat_a8( int col = X.sizes()[1]; int padded = (col + 512 - 1) / 512 * 512; int batch = X.sizes()[0]; - const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); + const GGUF_DEVICE_GUARD(device_of(X)); auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); at::Tensor Y = torch::empty({batch, row}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); at::Tensor quant_X = torch::empty({batch, padded / 32 * 9}, options); DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_mul_mat_a8", [&] { @@ -344,10 +362,10 @@ torch::Tensor ggml_moe_a8( int64_t tokens) { int col = X.sizes()[1]; int padded = (col + 512 - 1) / 512 * 512; - const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); + const GGUF_DEVICE_GUARD(device_of(X)); auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); at::Tensor Y = torch::empty({tokens * top_k, row}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options); DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_a8", [&] { @@ -548,10 +566,10 @@ torch::Tensor ggml_moe_a8_vec( int64_t tokens) { int col = X.sizes()[1]; const int padded = (col + 512 - 1) / 512 * 512; - const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); + const GGUF_DEVICE_GUARD(device_of(X)); auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); at::Tensor Y = torch::zeros({tokens * top_k, row}, options); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options); DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_vec_a8", [&] { diff --git a/python/freetoken/kernel/csrc/include/freetoken/device_api.h b/python/freetoken/kernel/csrc/include/freetoken/device_api.h new file mode 100644 index 000000000..932e2e36f --- /dev/null +++ b/python/freetoken/kernel/csrc/include/freetoken/device_api.h @@ -0,0 +1,99 @@ +// Device API seam for the tvm-ffi JIT/store/index kernels (and future GPU kernels). +// +// FreeToken's hand-written kernels are compiled once, guarded by `#if defined(USE_HIP)` +// (set by kernel/_toolchain.py `_rocm_cflags` and by setup.py for the torch extensions). +// This header maps the small set of runtime calls those kernels use to the active +// backend, so a kernel body written against these macros compiles under both CUDA and +// HIP without `#if` sprinkling at every call site. +// +// Every macro here must expand to a no-op-safe, backend-correct call. Prefer the 1:1 +// HIP counterparts (the HIP runtime is API-compatible at this level). +#pragma once + +#include +#include + +#if defined(USE_HIP) +#include +#else +#include +#include +#endif + +namespace freetoken::device { + +#if defined(USE_HIP) +using Error = hipError_t; +inline constexpr Error kErrorSuccess = hipSuccess; +inline const char* error_string(Error e) { return hipGetErrorString(e); } +using Device = hipDevice_t; +using Stream = hipStream_t; +using DeviceMemPtr = hipDeviceptr_t; +using HostFn = hipHostFn_t; +#else +using Error = cudaError_t; +inline constexpr Error kErrorSuccess = cudaSuccess; +inline const char* error_string(Error e) { return cudaGetErrorString(e); } +using Device = int; +using Stream = cudaStream_t; +using DeviceMemPtr = void*; +using HostFn = void (*)(void*); +#endif + +} // namespace freetoken::device + +// ---- allocation / copy / sync / launch (backend-agnostic call sites) ---- +#if defined(USE_HIP) +#define DEVICE_MALLOC(ptr, n) hipMalloc((void**)(ptr), (n)) +#define DEVICE_FREE(ptr) hipFree(ptr) +#define DEVICE_MEMCPY_ASYNC(dst, src, n, kind, stream) \ + hipMemcpyAsync((dst), (src), (n), (kind), (stream)) +#define DEVICE_MEMCPY_DEVICE_TO_DEVICE hipMemcpyDeviceToDevice +#define DEVICE_MEMCPY_DEVICE_TO_HOST hipMemcpyDeviceToHost +#define DEVICE_MEMCPY_HOST_TO_DEVICE hipMemcpyHostToDevice +#define DEVICE_SYNCTHREADS() __syncthreads() +#define DEVICE_LAUNCH_HOST_FUNC(stream, fn, data) \ + hipLaunchHostFunc((stream), (fn), (data)) +#define DEVICE_STREAM_SYNCHRONIZE(stream) hipStreamSynchronize(stream) +#define DEVICE_ATOMIC_ADD(ptr, val) atomicAdd((ptr), (val)) +#define DEVICE_ATOMIC_MAX(ptr, val) atomicMax((ptr), (val)) +#define DEVICE_ATOMIC_MIN(ptr, val) atomicMin((ptr), (val)) +#define DEVICE_ATOMIC_EXCHANGE(ptr, val) atomicExch((ptr), (val)) +#define DEVICE_THREAD_IDX_X threadIdx.x +#define DEVICE_BLOCK_DIM_X blockDim.x +#define DEVICE_BLOCK_IDX_X blockIdx.x +#define DEVICE_GRID_DIM_X gridDim.x +#else +#define DEVICE_MALLOC(ptr, n) cudaMalloc((void**)(ptr), (n)) +#define DEVICE_FREE(ptr) cudaFree(ptr) +#define DEVICE_MEMCPY_ASYNC(dst, src, n, kind, dir) \ + cudaMemcpyAsync((dst), (src), (n), (kind), (dir)) +#define DEVICE_MEMCPY_DEVICE_TO_DEVICE cudaMemcpyDeviceToDevice +#define DEVICE_MEMCPY_DEVICE_TO_HOST cudaMemcpyDeviceToHost +#define DEVICE_MEMCPY_HOST_TO_DEVICE cudaMemcpyHostToDevice +#define DEVICE_SYNCTHREADS() __syncthreads() +#define DEVICE_LAUNCH_HOST_FUNC(stream, fn, data) \ + cudaLaunchHostFunc((stream), (fn), (data)) +#define DEVICE_STREAM_SYNCHRONIZE(stream) cudaStreamSynchronize(stream) +#define DEVICE_ATOMIC_ADD(ptr, val) atomicAdd((ptr), (val)) +#define DEVICE_ATOMIC_MAX(ptr, val) atomicMax((ptr), (val)) +#define DEVICE_ATOMIC_MIN(ptr, val) atomicMin((ptr), (val)) +#define DEVICE_ATOMIC_EXCHANGE(ptr, val) atomicExch((ptr), (val)) +#define DEVICE_THREAD_IDX_X threadIdx.x +#define DEVICE_BLOCK_DIM_X blockDim.x +#define DEVICE_BLOCK_IDX_X blockIdx.x +#define DEVICE_GRID_DIM_X gridDim.x +#endif + +// The FFI kernels pin tensors to kDLCUDA today; on ROCm the tvm-ffi device type is +// kDLROCM. Kernels that resolve a device must accept both (see tensor.h). This macro +// picks the backend's DLDevice code at call time. +#if defined(USE_HIP) +#define DEVICE_DLDEVICE_ROCM 10 // kDLROCM +#define DEVICE_DLDEVICE_CUDA 2 // kDLCUDA +#define DEVICE_ACTIVE_DLDEVICE DEVICE_DLDEVICE_ROCM +#else +#define DEVICE_DLDEVICE_ROCM 10 +#define DEVICE_DLDEVICE_CUDA 2 +#define DEVICE_ACTIVE_DLDEVICE DEVICE_DLDEVICE_CUDA +#endif diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 8e917832c..de28877ce 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -10,6 +10,16 @@ #include #include +#if defined(USE_HIP) +#include +// HIP has no __grid_constant__ (a CUDA read-only-constant optimization). Define it +// empty so `const __grid_constant__ Params params` compiles as a plain by-value +// parameter, which is correct (just without the CUDA constant-cache hint). +#ifndef __grid_constant__ +#define __grid_constant__ +#endif +#endif + namespace device { inline constexpr auto kWarpThreads = 32u; @@ -42,16 +52,23 @@ __always_inline __device__ auto offset(const T *ptr, U... offset) -> const namespace PDL { +// Programmatic Dependent Launch is a CUDA-only optimization (griddepcontrol). +// HIP has no equivalent; the wait/launch are no-ops there. PDL is optional in the +// kernels (use_pdl defaults to false), so dropping it is purely a perf change. template __always_inline __device__ void wait() { +#if !defined(USE_HIP) if constexpr (kUsePDL) { asm volatile("griddepcontrol.wait;" ::: "memory"); } +#endif } template __always_inline __device__ void launch() { +#if !defined(USE_HIP) if constexpr (kUsePDL) { asm volatile("griddepcontrol.launch_dependents;" :::); } +#endif } } // namespace PDL @@ -60,6 +77,23 @@ template __always_inline __device__ void launch() { namespace host { +#if defined(USE_HIP) +inline auto +HIP_CHECK(::hipError_t error, + std::source_location location = std::source_location::current()) + -> void { + if (error != ::hipSuccess) { + [[unlikely]]; + ::host::panic(location, "HIP error: ", ::hipGetErrorString(error)); + } +} + +inline auto +HIP_CHECK(std::source_location location = std::source_location::current()) + -> void { + return HIP_CHECK(::hipGetLastError(), location); +} +#else inline auto CUDA_CHECK(::cudaError_t error, std::source_location location = std::source_location::current()) @@ -75,7 +109,21 @@ CUDA_CHECK(std::source_location location = std::source_location::current()) -> void { return CUDA_CHECK(::cudaGetLastError(), location); } +#endif +#if defined(USE_HIP) +template inline void set_smem_once(std::size_t smem_size) { + static const auto last_smem_size = [&] { + HIP_CHECK(::hipFuncSetAttribute( + F, ::hipFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + return smem_size; + }(); + RuntimeCheck( + smem_size <= last_smem_size, + "Dynamic shared memory size exceeds the previously set maximum size: ", + last_smem_size, " bytes"); +} +#else template inline void set_smem_once(std::size_t smem_size) { static const auto last_smem_size = [&] { CUDA_CHECK(::cudaFuncSetAttribute( @@ -87,7 +135,52 @@ template inline void set_smem_once(std::size_t smem_size) { "Dynamic shared memory size exceeds the previously set maximum size: ", last_smem_size, " bytes"); } +#endif +#if defined(USE_HIP) +struct LaunchKernel { +public: + explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, DLDevice device, + std::size_t dynamic_shared_mem_bytes = 0) noexcept + : m_grid(grid_dim), m_block(block_dim), + m_stream(resolve_device(device)), m_smem(dynamic_shared_mem_bytes) {} + + explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, hipStream_t stream, + std::size_t dynamic_shared_mem_bytes = 0) noexcept + : m_grid(grid_dim), m_block(block_dim), m_stream(stream), + m_smem(dynamic_shared_mem_bytes) {} + + static auto resolve_device(DLDevice device) -> hipStream_t { + return static_cast( + ::TVMFFIEnvGetStream(device.device_type, device.device_id)); + } + + LaunchKernel(const LaunchKernel &) = delete; + LaunchKernel &operator=(const LaunchKernel &) = delete; + + template + auto operator()(T &&kernel, Args &&...args) const -> void { + // hipLaunchKernel takes a void** args array (pointers to each argument value), + // unlike cudaLaunchKernelEx's variadic form. The array is consumed at launch. + void *arg_array[sizeof...(Args)] = { + const_cast(static_cast(&args))...}; + HIP_CHECK(::hipLaunchKernel(reinterpret_cast(kernel), m_grid, + m_block, arg_array, m_smem, m_stream)); + } + + auto with_attr(bool /*use_pdl*/) -> LaunchKernel & { + // HIP has no programmatic dependent launch / launch attributes; PDL is a + // CUDA-only optimization and is dropped here. + return *this; + } + +private: + dim3 m_grid; + dim3 m_block; + hipStream_t m_stream; + std::size_t m_smem; +}; +#else struct LaunchKernel { public: explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, DLDevice device, @@ -140,5 +233,6 @@ private: cudaLaunchConfig_t m_config; cudaLaunchAttribute m_attr_cache; }; +#endif } // namespace host diff --git a/python/freetoken/kernel/csrc/jit/batch_memcpy.cuh b/python/freetoken/kernel/csrc/jit/batch_memcpy.cuh index 690eff3dd..612b1d82e 100644 --- a/python/freetoken/kernel/csrc/jit/batch_memcpy.cuh +++ b/python/freetoken/kernel/csrc/jit/batch_memcpy.cuh @@ -8,6 +8,26 @@ #include #include +#if defined(USE_HIP) +// HIP has no cudaMemcpyBatchAsync equivalent. Implement a per-copy grid kernel: one +// block per copy, each block cooperatively copying its (src, dst, nbytes) triple. +// The host pointer arrays are staged to device memory for the launch. Copies within +// a batch are unordered (as with cudaMemcpyBatchAsync), so blocks may run in any order. +__global__ void batch_memcpy_kernel(const void* const* srcs, void* const* dsts, + const std::size_t* sizes, std::size_t n) { + const std::size_t i = blockIdx.x; + if (i >= n) { + return; + } + const char* s = static_cast(srcs[i]); + char* d = static_cast(dsts[i]); + const std::size_t sz = sizes[i]; + for (std::size_t j = threadIdx.x; j < sz; j += blockDim.x) { + d[j] = s[j]; + } +} +#endif + // Host wrapper over cudaMemcpyBatchAsync (CUDA >= 13.0, the 8-argument signature; // 12.8/12.9 carried an extra failIdx parameter): enqueue N independent // pointer-to-pointer copies with ONE runtime call, on an explicit (non-legacy) @@ -20,7 +40,49 @@ struct BatchMemcpy { tvm::ffi::TensorView sizes, int64_t stream_handle ) { -#if CUDART_VERSION >= 13000 +#if defined(USE_HIP) + using namespace host; + auto N = SymbolicSize{"batch length"}; + auto ptr_dtype = SymbolicDType{}; + TensorMatcher({N}) + .with_dtype(ptr_dtype) + .with_device() + .verify(dst_ptrs) + .verify(src_ptrs) + .verify(sizes); + auto n = static_cast(N.unwrap()); + if (n == 0) { + return; + } + RuntimeCheck(stream_handle != 0, "batch_memcpy rejects the legacy NULL stream"); + auto stream = reinterpret_cast(stream_handle); + + const void* const* srcs = reinterpret_cast(src_ptrs.data_ptr()); + void* const* dsts = reinterpret_cast(dst_ptrs.data_ptr()); + const std::size_t* sizes_arr = reinterpret_cast(sizes.data_ptr()); + + // Stage the host pointer arrays to device memory for the kernel. Use + // stream-ordered allocation so the free is ordered after the (async) kernel + // launch -- a plain hipFree here could free memory the kernel is still reading. + void* d_srcs = nullptr; + void* d_dsts = nullptr; + void* d_sizes = nullptr; + HIP_CHECK(hipMallocAsync(&d_srcs, n * sizeof(void*), stream)); + HIP_CHECK(hipMallocAsync(&d_dsts, n * sizeof(void*), stream)); + HIP_CHECK(hipMallocAsync(&d_sizes, n * sizeof(std::size_t), stream)); + HIP_CHECK(hipMemcpy(d_srcs, srcs, n * sizeof(void*), hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(d_dsts, dsts, n * sizeof(void*), hipMemcpyHostToDevice)); + HIP_CHECK(hipMemcpy(d_sizes, sizes_arr, n * sizeof(std::size_t), hipMemcpyHostToDevice)); + + void* args[4] = {&d_srcs, &d_dsts, &d_sizes, &n}; + HIP_CHECK(hipLaunchKernel( + reinterpret_cast(batch_memcpy_kernel), dim3(n), dim3(256), args, 0, + stream)); + + HIP_CHECK(hipFreeAsync(d_srcs, stream)); + HIP_CHECK(hipFreeAsync(d_dsts, stream)); + HIP_CHECK(hipFreeAsync(d_sizes, stream)); +#elif CUDART_VERSION >= 13000 using namespace host; auto N = SymbolicSize{"batch length"}; auto ptr_dtype = SymbolicDType{}; diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23ed..2583d1db6 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -34,40 +34,64 @@ inline constexpr auto get_mem_package() { } __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { +#if defined(USE_HIP) + return *src; +#else uint32_t tmp; asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src)); return uint1{tmp}; +#endif } __always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 { +#if defined(USE_HIP) + return *src; +#else uint32_t tmp0, tmp1; asm volatile("ld.global.L1::no_allocate.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src)); return uint2{tmp0, tmp1}; +#endif } __always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 { +#if defined(USE_HIP) + return *src; +#else uint32_t tmp0, tmp1, tmp2, tmp3; asm volatile("ld.global.L1::no_allocate.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src)); return uint4{tmp0, tmp1, tmp2, tmp3}; +#endif } __always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) { +#if defined(USE_HIP) + *dst = value; +#else uint32_t tmp = value.x; asm volatile("st.global.wt.b32 [%0],%1;" ::"l"(dst), "r"(tmp)); +#endif } __always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) { +#if defined(USE_HIP) + *dst = value; +#else uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; asm volatile("st.global.wt.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1)); +#endif } __always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) { +#if defined(USE_HIP) + *dst = value; +#else uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; uint32_t tmp2 = value.z; uint32_t tmp3 = value.w; asm volatile("st.global.wt.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3)); +#endif } __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag_ptr) { @@ -75,7 +99,10 @@ __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag auto* flag = reinterpret_cast(const_cast(flag_ptr)); uint32_t sleep_ns = 128; while (atomicAdd(flag, 0) > 0) { -#if __CUDA_ARCH__ >= 700 +#if defined(USE_HIP) + // HIP has no __nanosleep; busy-wait (functional parity). + (void)sleep_ns; +#elif __CUDA_ARCH__ >= 700 __nanosleep(sleep_ns); #endif sleep_ns = sleep_ns < 2048 ? (sleep_ns << 1) : 2048; @@ -134,6 +161,16 @@ __always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) // process (set at engine launch). inline bool host_ptr_identity() { static const bool identity = [] { +#if defined(USE_HIP) + int device = 0; + if (hipGetDevice(&device) != hipSuccess) { + return false; // fail closed: translate (and surface errors), don't assume identity + } + int uva = 0, reg = 0; + hipDeviceGetAttribute(&uva, hipDeviceAttributeUnifiedAddressing, device); + hipDeviceGetAttribute(®, hipDeviceAttributeCanUseHostPointerForRegisteredMem, device); + return uva == 1 && reg == 1; +#else int device = 0; if (cudaGetDevice(&device) != cudaSuccess) { return false; // fail closed: translate (and surface errors), don't assume identity @@ -142,11 +179,23 @@ inline bool host_ptr_identity() { cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device); cudaDeviceGetAttribute(®, cudaDevAttrCanUseHostPointerForRegisteredMem, device); return uva == 1 && reg == 1; +#endif }(); return identity; } inline void* device_alias(void* ptr, DLDevice dev) { +#if defined(USE_HIP) + if (dev.device_type == kDLROCM || host_ptr_identity()) { + return ptr; + } + void* mapped = nullptr; + const auto err = hipHostGetDevicePointer(&mapped, ptr, 0); + host::RuntimeCheck(err == hipSuccess, + "fast_index_copy: host tensor must be pinned+mapped (hipHostGetDevicePointer: ", + hipGetErrorString(err), ")"); + return mapped; +#else if (dev.device_type == kDLCUDA || host_ptr_identity()) { return ptr; } @@ -156,6 +205,7 @@ inline void* device_alias(void* ptr, DLDevice dev) { "fast_index_copy: host tensor must be pinned+mapped (cudaHostGetDevicePointer: ", cudaGetErrorString(err), ")"); return mapped; +#endif } struct IndexKernelParams { @@ -344,12 +394,12 @@ struct FastIndexCopyKernel { TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(src); TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(dst); TensorMatcher({L}) diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index c3947adfa..da15cdd86 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp @@ -1,12 +1,33 @@ #include +#if defined(USE_HIP) +#include +#else #include +#endif #include +// host alloc flags (mapped+portable) shared by both backends +#if defined(USE_HIP) +#define DEVICE_HOST_ALLOC_FLAGS_MAPPED \ + (hipHostMallocPortable | hipHostMallocMapped) +#define DEVICE_HOST_REGISTER_FLAGS_MAPPED \ + (hipHostRegisterPortable | hipHostRegisterMapped) +#else +#define DEVICE_HOST_ALLOC_FLAGS_MAPPED \ + (cudaHostAllocPortable | cudaHostAllocMapped) +#define DEVICE_HOST_REGISTER_FLAGS_MAPPED \ + (cudaHostRegisterPortable | cudaHostRegisterMapped) +#endif + namespace { void free_pinned(void *ptr) { if (ptr != nullptr) { +#if defined(USE_HIP) + hipHostFree(ptr); +#else cudaFreeHost(ptr); +#endif } } @@ -34,9 +55,16 @@ torch::Tensor create_pinned_tensor_like(torch::Tensor input) { const size_t alloc_nbytes = static_cast(nbytes == 0 ? 1 : nbytes); void *data_ptr = nullptr; +#if defined(USE_HIP) + const hipError_t alloc_err = + hipMallocHost(&data_ptr, alloc_nbytes); + TORCH_CHECK(alloc_err == hipSuccess, + "hipMallocHost failed: ", hipGetErrorString(alloc_err)); +#else const cudaError_t alloc_err = cudaMallocHost(&data_ptr, alloc_nbytes); TORCH_CHECK(alloc_err == cudaSuccess, "cudaMallocHost failed: ", cudaGetErrorString(alloc_err)); +#endif auto options = input.options().device(torch::kCPU).pinned_memory(true); @@ -58,10 +86,17 @@ torch::Tensor alloc_pinned_tensor(std::vector sizes, // Portable + mapped: the offload gather kernel reads these banks straight // from host memory (zero-copy), which requires device-mapped pinned pages. void *data_ptr = nullptr; +#if defined(USE_HIP) + const hipError_t alloc_err = hipHostMalloc( + &data_ptr, alloc_nbytes, DEVICE_HOST_ALLOC_FLAGS_MAPPED); + TORCH_CHECK(alloc_err == hipSuccess, + "hipHostMalloc failed: ", hipGetErrorString(alloc_err)); +#else const cudaError_t alloc_err = cudaHostAlloc( &data_ptr, alloc_nbytes, cudaHostAllocPortable | cudaHostAllocMapped); TORCH_CHECK(alloc_err == cudaSuccess, "cudaHostAlloc failed: ", cudaGetErrorString(alloc_err)); +#endif auto options = torch::TensorOptions() .dtype(dtype) @@ -76,38 +111,69 @@ torch::Tensor alloc_pinned_tensor(std::vector sizes, // device address). Zero-copy consumers resolve bank base addresses through these. bool host_ptr_identity() { int device = 0; +#if defined(USE_HIP) + const hipError_t err = hipGetDevice(&device); + TORCH_CHECK(err == hipSuccess, "hipGetDevice failed: ", hipGetErrorString(err)); + int uva = 0, reg = 0; + hipDeviceGetAttribute(&uva, hipDeviceAttributeUnifiedAddressing, device); + hipDeviceGetAttribute( + ®, hipDeviceAttributeCanUseHostPointerForRegisteredMem, device); +#else const cudaError_t err = cudaGetDevice(&device); TORCH_CHECK(err == cudaSuccess, "cudaGetDevice failed: ", cudaGetErrorString(err)); int uva = 0, reg = 0; cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device); cudaDeviceGetAttribute(®, cudaDevAttrCanUseHostPointerForRegisteredMem, device); +#endif return uva == 1 && reg == 1; } int64_t host_device_ptr(int64_t host_ptr) { void *dev_ptr = nullptr; +#if defined(USE_HIP) + const hipError_t err = hipHostGetDevicePointer( + &dev_ptr, reinterpret_cast(host_ptr), 0); + TORCH_CHECK(err == hipSuccess, + "hipHostGetDevicePointer failed (host memory must be pinned+mapped): ", + hipGetErrorString(err)); +#else const cudaError_t err = cudaHostGetDevicePointer(&dev_ptr, reinterpret_cast(host_ptr), 0); TORCH_CHECK(err == cudaSuccess, "cudaHostGetDevicePointer failed (host memory must be pinned+mapped): ", cudaGetErrorString(err)); +#endif return reinterpret_cast(dev_ptr); } void host_register(int64_t addr, int64_t nbytes) { +#if defined(USE_HIP) + const hipError_t err = + hipHostRegister(reinterpret_cast(addr), static_cast(nbytes), + DEVICE_HOST_REGISTER_FLAGS_MAPPED); + TORCH_CHECK(err == hipSuccess, + "hipHostRegister failed: ", hipGetErrorString(err)); +#else const cudaError_t err = cudaHostRegister(reinterpret_cast(addr), static_cast(nbytes), cudaHostRegisterPortable | cudaHostRegisterMapped); TORCH_CHECK(err == cudaSuccess, "cudaHostRegister failed: ", cudaGetErrorString(err)); +#endif } +// CUDA-only: the maximum CUDA version the NVIDIA driver supports, used to gate +// driver-JIT kernels. On ROCm there is no CUDA driver; return 0 (== "no driver"). int64_t driver_cuda_version() { +#if defined(USE_HIP) + return 0; +#else int version = 0; // stays 0 when no driver is installed const cudaError_t err = cudaDriverGetVersion(&version); TORCH_CHECK(err == cudaSuccess, "cudaDriverGetVersion failed: ", cudaGetErrorString(err)); return version; +#endif } } // namespace @@ -122,7 +188,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("host_device_ptr", &host_device_ptr, "Device-visible alias of a pinned+mapped host address"); m.def("host_register", &host_register, - "cudaHostRegister an existing host range as portable+mapped"); + "cudaHostRegister/hipHostRegister an existing host range as portable+mapped"); m.def("driver_cuda_version", &driver_cuda_version, - "Max CUDA version the installed NVIDIA driver supports (0 if none)"); + "Max CUDA version the installed NVIDIA driver supports (0 if none / ROCm)"); } diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 04a165609..13b0ea0c5 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -51,16 +51,28 @@ def _c_compiler_for(cxx: str) -> str: def _module(): from torch.utils.cpp_extension import load - extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"] - host_cxx = _host_compiler() - if host_cxx is not None: - # Point both nvcc's host pass (-ccbin) and torch's C++ compile (CXX) at a - # libtorch/nvcc-compatible compiler. Force (not setdefault): the system - # default (CXX unset -> g++) can be a gcc too new for the torch headers. - cxx_path = shutil.which(host_cxx) or host_cxx - extra_cuda_cflags += ["-ccbin", cxx_path] - os.environ["CXX"] = cxx_path - os.environ["CC"] = _c_compiler_for(cxx_path) + if torch.version.hip is not None: + # ROCm: hipcc (torch.utils.cpp_extension picks it up), pass the HIP defines so + # the kernels compile their HIP branches; drop the CUDA-only -ccbin/flag logic. + # Explicit --offload-arch (plus PYTORCH_ROCM_ARCH) prevents torch from auto- + # emitting ~14 gfx arches, which would multiply build time per arch. + os.environ.setdefault("PYTORCH_ROCM_ARCH", "gfx1100") + extra_cuda_cflags = [ + "-O3", "--offload-arch=gfx1100", "-DUSE_HIP=1", "-DUSE_ROCM=1", + ] + os.environ.pop("CXX", None) + os.environ.pop("CC", None) + else: + extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"] + host_cxx = _host_compiler() + if host_cxx is not None: + # Point both nvcc's host pass (-ccbin) and torch's C++ compile (CXX) at a + # libtorch/nvcc-compatible compiler. Force (not setdefault): the system + # default (CXX unset -> g++) can be a gcc too new for the torch headers. + cxx_path = shutil.which(host_cxx) or host_cxx + extra_cuda_cflags += ["-ccbin", cxx_path] + os.environ["CXX"] = cxx_path + os.environ["CC"] = _c_compiler_for(cxx_path) # gguf_kernel.cu carries its own PYBIND11_MODULE (appended at the end), so a # plain `load` of the single source compiles + binds the ggml_* ops. @@ -69,7 +81,7 @@ def _module(): sources=[str(_CSRC / "gguf_kernel.cu")], extra_include_paths=[str(_CSRC)], extra_cuda_cflags=extra_cuda_cflags, - verbose=True, + verbose=False, ) diff --git a/python/freetoken/kernel/tinygrad_fallback.py b/python/freetoken/kernel/tinygrad_fallback.py new file mode 100644 index 000000000..2b0663fe3 --- /dev/null +++ b/python/freetoken/kernel/tinygrad_fallback.py @@ -0,0 +1,92 @@ +"""tinygrad-JIT fallback for FFI kernels that are not hand-ported to HIP. + +FreeToken's hand-written tvm-ffi kernels (``store`` / ``index`` / ``fast_index_copy`` / +``batch_memcpy``) are CUDA source compiled via nvcc/JIT. The primary AMD port is the +``#if defined(USE_HIP)`` seam in ``device_api.h`` + ``LaunchKernel``/``warp.cuh``. This +module is the **documented fallback** for any kernel that proves intractable to hipify: +tinygrad's JIT compiles one logical kernel to PTX (CUDA) *and* AMDGPU/LLVM (ROCm), so the +same source covers both platforms. + +Constraints (matching the FFI contract): + +* Each fallback takes the same ``tvm.ffi.TensorView`` arguments as the hand-written + kernel and returns the same output tensor(s), so the swap is invisible to callers. +* It runs on the *host* (tinygrad handles GPU dispatch); on ROCm it compiles to AMDGPU. +* It is **never a default**: ``kernel/utils.py`` only routes a kernel to the fallback + when (a) ROCm is active and (b) the hand-HIP AOT/JIT variant is absent/unbuildable. + If tinygrad is not installed, invoking the fallback raises a clear error. + +Because tinygrad is an optional dependency (installed only when the fallback is actually +needed), all imports here are lazy and the module imports with zero third-party deps, so +it is safe to import on the CUDA-only path. +""" + +from __future__ import annotations + +from typing import Callable, Optional + +__all__ = [ + "is_tinygrad_available", + "kernel_fallback_available", + "get_kernel_fallback", +] + +# Kernel names the fallback registry knows how to build (mirrors the FFI kernel set). +_KNOWN_KERNELS = ("store", "index", "fast_index_copy", "batch_memcpy") + +#: Which kernels currently have a *functional* tinygrad reimplementation. As HIP ports +#: land in Inc 7, names are removed from this set (the hand port wins); kernels left here +#: (if any) are the documented fallback set. Default: empty -- the hand-HIP port is the +#: primary path and the fallback is opt-in per kernel. +_FALLBACK_IMPLEMENTED: set[str] = set() + + +def is_tinygrad_available() -> bool: + """True when the ``tinygrad`` package can be imported (JIT-to-ROCm available).""" + try: + import importlib.util # noqa: PLC0415 + + return importlib.util.find_spec("tinygrad") is not None + except Exception: + return False + + +def kernel_fallback_available(kernel: str) -> bool: + """True when a tinygrad fallback for ``kernel`` is both implemented and usable + (tinygrad installed). Always False on the CUDA path unless explicitly enabled, so + the CUDA build never depends on tinygrad.""" + if kernel not in _FALLBACK_IMPLEMENTED: + return False + return is_tinygrad_available() + + +def get_kernel_fallback(kernel: str): + """Return the tinygrad-backed fallback callable for ``kernel``, or raise a clear + error explaining why it is unavailable. Never called on the CUDA path.""" + if kernel not in _FALLBACK_IMPLEMENTED: + raise RuntimeError( + f"FFI kernel {kernel!r} has no tinygrad fallback registered. On ROCm the " + "preferred path is the hand-written HIP port (device_api.h); if you intend " + "to use the tinygrad fallback you must register it in " + "kernel/tinygrad_fallback.py._FALLBACK_IMPLEMENTED and implement the " + "corresponding build function." + ) + if not is_tinygrad_available(): + raise RuntimeError( + f"FFI kernel {kernel!r} requires the tinygrad fallback, but tinygrad is not " + "installed. Install it (`pip install tinygrad`) or provide a hand-written " + "HIP port for this kernel." + ) + from freetoken.kernel import tinygrad_impl # noqa: PLC0415 (lazy; may be None) + + builder = getattr(tinygrad_impl, f"build_{kernel}", None) + if builder is None: + raise RuntimeError( + f"tinygrad fallback for {kernel!r} is registered but has no " + "tinygrad_impl.build_() builder." + ) + return builder + + +def _list_fallbacks() -> list[str]: + return [k for k in _FALLBACK_IMPLEMENTED if kernel_fallback_available(k)] diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py index 2c38b533e..d57852f56 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/python/freetoken/kernel/triton/activation.py @@ -48,20 +48,15 @@ def _pdl_supported() -> bool: @triton.jit def _fast_tanh(x): - # PTX tanh.approx.f32 — single HW op, matches flashinfer math::tanh. - return tl.inline_asm_elementwise( - "tanh.approx.f32 $0, $1;", "=f,f", [x], - dtype=tl.float32, is_pure=True, pack=1, - ) + # tanh.approx.f32 is a CUDA PTX intrinsic; libdevice.tanh is portable (maps to + # tanhf on both CUDA and AMD). + return libdevice.tanh(x) @triton.jit def _fast_ex2(x): - # PTX ex2.approx.f32 — matches __expf fast path used by flashinfer silu. - return tl.inline_asm_elementwise( - "ex2.approx.f32 $0, $1;", "=f,f", [x], - dtype=tl.float32, is_pure=True, pack=1, - ) + # ex2.approx.f32 is CUDA-only; libdevice.exp2 is portable. + return libdevice.exp2(x) @triton.jit @@ -133,10 +128,18 @@ def _act_and_mul( # 1024/w4/s2 best at rows>=4096). block_d = min(triton.next_power_of_2(d), 1024 if M >= 4096 else 512) num_stages = 2 if block_d == 1024 else 3 - _act_and_mul_kernel[grid]( - o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, launch_pdl=pdl, - BLOCK_D=block_d, num_warps=4, num_stages=num_stages, - ) + # ``launch_pdl`` (Hopper griddepcontrol) is CUDA-only; ROCm triton rejects the + # kwarg, so only pass it on a PDL-capable (sm_90+) device. + if pdl: + _act_and_mul_kernel[grid]( + o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=True, launch_pdl=True, + BLOCK_D=block_d, num_warps=4, num_stages=num_stages, + ) + else: + _act_and_mul_kernel[grid]( + o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=False, + BLOCK_D=block_d, num_warps=4, num_stages=num_stages, + ) return out diff --git a/python/freetoken/kernel/triton/norm.py b/python/freetoken/kernel/triton/norm.py index 3f95c29f0..320f60776 100644 --- a/python/freetoken/kernel/triton/norm.py +++ b/python/freetoken/kernel/triton/norm.py @@ -142,11 +142,18 @@ def _rmsnorm(input, weight, eps, out, gemma: bool): # PDL only on the contiguous (decode-replay) path: on the strided qk-norm's # 32k-CTA prefill grids the per-CTA gdc_wait poll costs more than it hides. pdl = contig and is_sm90_supported() - _rmsnorm_kernel[(A, B)]( - out, input, weight, eps, H, sxa, sxb, soa, sob, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, - num_warps=_num_warps(A * B), num_stages=1, - ) + if pdl: + _rmsnorm_kernel[(A, B)]( + out, input, weight, eps, H, sxa, sxb, soa, sob, + CONTIG=True, ENABLE_PDL=True, launch_pdl=True, GEMMA=gemma, + num_warps=_num_warps(A * B), num_stages=1, + ) + else: + _rmsnorm_kernel[(A, B)]( + out, input, weight, eps, H, sxa, sxb, soa, sob, + CONTIG=contig, ENABLE_PDL=False, GEMMA=gemma, + num_warps=_num_warps(A * B), num_stages=1, + ) return out @@ -170,11 +177,18 @@ def _fused_add_rmsnorm(input, residual, weight, eps, gemma: bool): _, _, sra, srb = _leading(residual) contig = input.ndim == 2 and input.is_contiguous() and residual.is_contiguous() pdl = contig and is_sm90_supported() - _fused_add_rmsnorm_kernel[(A, B)]( - input, residual, weight, eps, H, sxa, sxb, sra, srb, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, - num_warps=_num_warps(A * B), num_stages=1, - ) + if pdl: + _fused_add_rmsnorm_kernel[(A, B)]( + input, residual, weight, eps, H, sxa, sxb, sra, srb, + CONTIG=True, ENABLE_PDL=True, launch_pdl=True, GEMMA=gemma, + num_warps=_num_warps(A * B), num_stages=1, + ) + else: + _fused_add_rmsnorm_kernel[(A, B)]( + input, residual, weight, eps, H, sxa, sxb, sra, srb, + CONTIG=contig, ENABLE_PDL=False, GEMMA=gemma, + num_warps=_num_warps(A * B), num_stages=1, + ) def fused_add_rmsnorm(input, residual, weight, eps: float = 1e-6, enable_pdl: bool = False): diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index 7a0164b59..1beb73982 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -40,6 +40,32 @@ def _rank(a: str) -> int: cc = max(arch_list, key=_rank).rstrip("a").replace(".", "") flags = flags + [f"-gencode=arch=compute_{cc},code=compute_{cc}"] return flags + + +def _rocm_cflags(extra: List[str]) -> List[str]: + """HIP/ROCm flags for a kernel build. Adds the ``USE_HIP``/``USE_ROCM`` defines so the + shared ``.cu``/``.cuh`` sources compile their HIP branches (llama.cpp-style). When + ``TVM_FFI_ROCM_ARCH_LIST`` (e.g. "gfx1100") is set (AOT cache build), we pin + ``--offload-arch``; otherwise hipcc targets the local GPU.""" + flags = DEFAULT_CFLAGS + ["-DUSE_HIP=1", "-DUSE_ROCM=1"] + list(extra) + arch_list = os.getenv("TVM_FFI_ROCM_ARCH_LIST", "").split() + if arch_list: + flags = flags + [f"--offload-arch={arch_list[0]}"] + return flags + + +def _arch_flags(extra: List[str]) -> List[str]: + """GPU kernel build flags for the active backend: HIP flags on ROCm torch, else CUDA.""" + try: + from freetoken.kernel._toolchain import is_rocm_torch + + if is_rocm_torch(): + return _rocm_cflags(extra) + except Exception: + pass + return _cuda_cflags(extra) + + CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool] @@ -86,6 +112,20 @@ def _build_stamps(segments: List[str]) -> set[str]: return {s for s in segments if re.fullmatch(r"g[0-9a-f]{7,40}", s)} +def _build_stamps(local_segments: List[str]) -> List[str]: + """The `g` commit-stamp tokens of a local version segment list + (``["cu130", "g3f01615"]`` -> ``["g3f01615"]``).""" + return [s for s in local_segments if s.startswith("g")] + + +def _arch_tags(local_segments: List[str]) -> List[str]: + """The backend-tag tokens of a local segment list (``cu130`` or ``rocm``). A cache + wheel and runtime wheel must carry the SAME backend tag -- a ``+rocm`` cache is + meaningless to a ``+cu130`` runtime (the fatbin is gfx SASS vs sm SASS) and vice + versa. Returns the tags (usually one, e.g. ``["cu130"]`` or ``["rocm"]``).""" + return [s for s in local_segments if s.startswith("cu") or s.startswith("rocm")] + + def _kernel_cache_version_ok(cache_version: str, runtime_version: str) -> bool: """Same release -- and, when both sides carry a `g` stamp, the same build. @@ -93,14 +133,24 @@ def _kernel_cache_version_ok(cache_version: str, runtime_version: str) -> bool: `.g`), so the old string-prefix test cannot pair a stamped runtime with its cache; and comparing the stamps rejects a runtime/cache pair from two different builds, which bare release numbers (both `0.1.1`) could never detect. Either side - may lack a stamp (dev builds) -- then only the release part is compared.""" + may lack a stamp (dev builds) -- then only the release part is compared. + + The backend arch tag (`cu130` vs `rocm`) must also match when both sides carry one: + a prebuilt kernel-cache fatbin is SASS for a specific backend family, so a CUDA + runtime must never load a ROCm cache (or vice versa).""" cache_base, cache_local = _version_parts(cache_version) runtime_base, runtime_local = _version_parts(runtime_version) if cache_base != runtime_base: return False cache_stamps = _build_stamps(cache_local) runtime_stamps = _build_stamps(runtime_local) - return not (cache_stamps and runtime_stamps and cache_stamps != runtime_stamps) + if cache_stamps and runtime_stamps and cache_stamps != runtime_stamps: + return False + cache_arch = _arch_tags(cache_local) + runtime_arch = _arch_tags(runtime_local) + if cache_arch and runtime_arch and cache_arch != runtime_arch: + return False + return True def _kernel_cache_dir() -> pathlib.Path | None: @@ -127,7 +177,8 @@ def _kernel_cache_dir() -> pathlib.Path | None: f"{package_version!r} does not match freetoken version {runtime_version!r}" ) cache_cuda = re.search(r"\+cu(\d{2,})", package_version) - if cache_cuda is not None: + cache_rocm = re.search(r"\+rocm", package_version) + if cache_cuda is not None and cache_rocm is None: from freetoken.kernel._toolchain import torch_cuda_major cache_major = int(cache_cuda.group(1)[:-1]) @@ -201,9 +252,16 @@ def load_aot( return prebuilt if cuda_files: - from freetoken.kernel._toolchain import check_nvcc_matches_torch + from freetoken.kernel._toolchain import ( + check_hip_matches_torch, + check_nvcc_matches_torch, + is_rocm_torch, + ) - check_nvcc_matches_torch() + if is_rocm_torch(): + check_hip_matches_torch() + else: + check_nvcc_matches_torch() from tvm_ffi.cpp import load @@ -222,7 +280,7 @@ def load_aot( cpp_files=cpp_files, cuda_files=cuda_files, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags), + extra_cuda_cflags=_arch_flags(extra_cuda_cflags), extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, @@ -247,9 +305,16 @@ def load_jit( return prebuilt if cuda_files or cuda_wrappers: - from freetoken.kernel._toolchain import check_nvcc_matches_torch + from freetoken.kernel._toolchain import ( + check_hip_matches_torch, + check_nvcc_matches_torch, + is_rocm_torch, + ) - check_nvcc_matches_torch() + if is_rocm_torch(): + check_hip_matches_torch() + else: + check_nvcc_matches_torch() from tvm_ffi.cpp import load_inline @@ -277,7 +342,7 @@ def load_jit( cpp_sources=cpp_sources, cuda_sources=cuda_sources, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags), + extra_cuda_cflags=_arch_flags(extra_cuda_cflags), extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index d68d8ded5..e16913639 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -531,6 +531,16 @@ def _expert_gemm( return fused_experts_gguf_q4_0( hidden_states, gate_up, down, topk_weights, topk_ids, self.activation ) + if fmt == "gguf": + # Native GGUF qwen3.5-moe experts: gate_up stays Q4_K, down is stored as Q8_0 + # (re-quantized at load; a uniform format the cache can hold). Dequant-in-kernel + # grouped GEMV (MMVQ) over the streamed packed banks. + from freetoken.moe.fused_gguf import fused_experts_gguf + + gate_up, down = views + return fused_experts_gguf( + hidden_states, gate_up, down, topk_weights, topk_ids, self.activation + ) if fmt == "mxfp4_triton": # gpt-oss MXFP4 experts (biased, clamped swiglu): transposed split-K GEMV # decode + grouped `_t` prefill. The swiglu scalars live on the layer diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 63b1a18b9..b1e84c5cc 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -18,6 +18,7 @@ # reuses the model classes but a GGUF parse_config / iter_weights). GGUF_ARCH_TO_REGISTRY: dict[str, str] = { "gemma4": "Gemma4GGUFForCausalLM", + "qwen35moe": "Qwen35moeGGUFForCausalLM", } diff --git a/python/freetoken/models/gguf/dequant.py b/python/freetoken/models/gguf/dequant.py index 77c3ea010..05648f9e4 100644 --- a/python/freetoken/models/gguf/dequant.py +++ b/python/freetoken/models/gguf/dequant.py @@ -23,6 +23,8 @@ GGML_F16 = 1 GGML_Q4_0 = 2 GGML_Q8_0 = 8 +GGML_Q4_K = 12 +GGML_Q5_K = 13 GGML_Q6_K = 14 GGML_BF16 = 30 @@ -33,6 +35,7 @@ GGML_BF16: (1, 2), GGML_Q4_0: (32, 18), GGML_Q8_0: (32, 34), + GGML_Q4_K: (256, 144), GGML_Q6_K: (256, 210), } @@ -42,6 +45,8 @@ GGML_BF16: "BF16", GGML_Q4_0: "Q4_0", GGML_Q8_0: "Q8_0", + GGML_Q4_K: "Q4_K", + GGML_Q5_K: "Q5_K", GGML_Q6_K: "Q6_K", } @@ -115,8 +120,69 @@ def dequant_q6_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: return y.reshape(-1).to(out_dtype) +def quantize_q8_0(w: torch.Tensor) -> torch.Tensor: + """Quantize dense rows to packed Q8_0 blocks (``half d`` + 32 int8). + + ``w``'s last dim must be a multiple of 32 (the Q8_0 block); returns the packed + ``[..., n/32*34]`` uint8 layout the ggml Q8_0 kernels read. Used to re-quantize + K-quant expert banks to a uniform 8-bit type (Q8_0 >= Q5_K/Q6_K precision, so no + quality loss) when the offload cache needs a single per-bank format. + """ + n = w.shape[-1] + assert n % 32 == 0, f"Q8_0 quantize needs last dim % 32 == 0, got {n}" + wq = w.float().view(*w.shape[:-1], n // 32, 32) + d = wq.abs().amax(dim=-1, keepdim=True).clamp(min=1e-9) / 127.0 + q = torch.round(wq / d).to(torch.int8) + dh = d.to(torch.float16).view(torch.uint8) # [..., n//32, 2] + packed = torch.cat([dh, q.view(torch.uint8)], dim=-1) # [..., n//32, 34] + return packed.reshape(*w.shape[:-1], (n // 32) * 34).contiguous() + + +def dequant_q5_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + """Q5_K: 256-elem super-block = half2 dm (dall, dmin), 12B 6-bit scale/min, 32B + qh high-bits, 128B qs low nibbles. Mirrors ggml's dequantize_block_q5_K.""" + raw = raw.reshape(-1, 176) + n = raw.shape[0] + dm = raw[:, 0:4].view(torch.float16).to(torch.float32) # [n,2] + dall, dmin = dm[:, 0], dm[:, 1] + scales = raw[:, 4:16] + qh = raw[:, 16:48].to(torch.int32) + qs = raw[:, 48:176].to(torch.int32) + + def _sm(j): + if j < 4: + d = scales[:, j] & 63 + m = scales[:, j + 4] & 63 + else: + d = (scales[:, j + 4] & 0xF) | ((scales[:, j - 4] >> 6) << 4) + m = (scales[:, j + 4] >> 4) | ((scales[:, j] >> 6) << 4) + return d.to(torch.float32), m.to(torch.float32) + + y = torch.zeros(n, 256, dtype=torch.float32, device=raw.device) + for il in range(4): + s0, m0 = _sm(2 * il) + s1, m1 = _sm(2 * il + 1) + d0, M0 = dall * s0, dmin * m0 + d1, M1 = dall * s1, dmin * m1 + bit0 = 1 << (2 * il) + bit1 = bit0 << 1 + ql = qs[:, 32 * il:32 * il + 32] + ql0, ql1 = ql[:, 0::2], ql[:, 1::2] + h0, h1 = qh[:, 0::2], qh[:, 1::2] + v0 = (ql0 & 0xF) + ((h0 & bit0) != 0).to(torch.float32) * 16 + v1 = (ql1 & 0xF) + ((h1 & bit0) != 0).to(torch.float32) * 16 + even = torch.stack([v0, v1], dim=-1).reshape(n, 32) # [v0[0],v1[0],v0[1],...] + y[:, 64 * il:64 * il + 32] = even * d0.unsqueeze(1) - M0.unsqueeze(1) + w0 = (ql0 >> 4) + ((h0 & bit1) != 0).to(torch.float32) * 16 + w1 = (ql1 >> 4) + ((h1 & bit1) != 0).to(torch.float32) * 16 + odd = torch.stack([w0, w1], dim=-1).reshape(n, 32) + y[:, 64 * il + 32:64 * il + 64] = odd * d1.unsqueeze(1) - M1.unsqueeze(1) + return y.reshape(-1).to(out_dtype) + + _DEQUANT = { GGML_Q4_0: dequant_q4_0, + GGML_Q5_K: dequant_q5_k, GGML_Q6_K: dequant_q6_k, } @@ -143,11 +209,14 @@ def dequantize(raw: torch.Tensor, ggml_type: int, out_dtype: torch.dtype) -> tor "GGML_BF16", "GGML_Q4_0", "GGML_Q8_0", + "GGML_Q4_K", "GGML_Q6_K", "GGML_NAME", "BLOCK_SHAPE", "row_bytes", "dequant_q4_0", + "dequant_q5_k", "dequant_q6_k", + "quantize_q8_0", "dequantize", ] diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index 6d5481c17..c05824190 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -13,7 +13,10 @@ from .reader import gguf_architecture, load_gguf_metadata # GGUF architecture -> transformers GGUF tokenizer-converter key. -_TOKENIZER_ARCH = {"gemma4": "gemma4_text"} +_TOKENIZER_ARCH = { + "gemma4": "gemma4_text", + "qwen35moe": "qwen3_moe", +} def load_gguf_tokenizer(model_path: str): diff --git a/python/freetoken/models/qwen3_5_moe/__init__.py b/python/freetoken/models/qwen3_5_moe/__init__.py index 98936e9f2..cae7dfd34 100644 --- a/python/freetoken/models/qwen3_5_moe/__init__.py +++ b/python/freetoken/models/qwen3_5_moe/__init__.py @@ -1,4 +1,11 @@ from .config import parse_config +from .gguf import ( + convert_qwen35moe_to_gguf, + is_gguf_model, + iter_gguf_weights, + load_gguf_expert_sources, + parse_gguf_config, +) from .model import Qwen3_5MoEForCausalLM from .weight import ( iter_weights, @@ -16,4 +23,9 @@ "load_nvfp4_expert_sources", "load_nvfp4_expert_sources_parallel", "setup_offload_expert_banks", + "parse_gguf_config", + "iter_gguf_weights", + "convert_qwen35moe_to_gguf", + "is_gguf_model", + "load_gguf_expert_sources", ] diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index 2e7320051..7b7f227c0 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -76,10 +76,11 @@ def __init__( # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM). self._block_fp8 = expert_quant == "fp8_block" self._pertensor_fp8 = attn_quant == "fp8_pertensor" - self._fp8 = self._block_fp8 or self._pertensor_fp8 + self._gguf = expert_quant == "gguf" + self._fp8 = self._block_fp8 or self._pertensor_fp8 or self._gguf self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] - if self._fp8: + if self._block_fp8 or self._pertensor_fp8: ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged self.in_proj_qkvz = ColMerged( hidden_size, [self.conv_dim, self.value_dim], has_bias=False @@ -87,6 +88,15 @@ def __init__( self.in_proj_ba = LinearColParallelMerged( hidden_size, [num_v_heads, num_v_heads], has_bias=False ) + elif self._gguf: + # GGUF: qkv|z are native-quant (Q8_0, swapped to GGUFLinear after build), b|a + # stay dense bf16. Same split as the fp8 path (matches the GGUF tensor layout). + self.in_proj_qkvz = LinearColParallelMerged( + hidden_size, [self.conv_dim, self.value_dim], has_bias=False + ) + self.in_proj_ba = LinearColParallelMerged( + hidden_size, [num_v_heads, num_v_heads], has_bias=False + ) else: # Fused input projection (one GEMM instead of four): qkv | z | b | a. self.in_proj = LinearColParallelMerged(hidden_size, self._in_proj_split, has_bias=False) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py new file mode 100644 index 000000000..aef81a099 --- /dev/null +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -0,0 +1,512 @@ +"""Qwen3.5-MoE GGUF adapter: build the FreeToken ``ModelConfig`` from GGUF metadata and +map GGUF tensors to the model's state dict. + +The qwen35moe GGUF arch is a hybrid GatedDeltaNet (linear-attention SSM) + full-attention +MoE (40 layers, every 4th full; 256 routed experts + shared expert). The GGUF geometry +matches the HF qwen3_5_moe model, so ``parse_gguf_config`` produces the *same* +``ModelConfig`` as ``qwen3_5_moe.config.parse_config`` -- only the source is GGUF KV +metadata. ``expert_quant``/``attn_quant``/``dense_quant`` are set to ``"gguf"`` so +``convert_qwen35moe_to_gguf`` can detect the native-quant checkpoint and swap the dense +projections for native GGUF-quant ops. + +Quantized projections (full-attn q/k/v/o, GDN qkv|z and out_proj, shared-expert gate/up/ +down, the token embedding and the lm_head) stay in their native packed block layout (Q8_0 +projections, Q6_K head) and are yielded as ``.qweight`` (uint8); tiny F32 tensors (norms, +router, GDN b/a) dequantize to bf16; GDN conv/A_log/dt_bias stay fp32. Routed experts +(Q4_K gate/up, Q5_K/Q6_K down) go to the offload cache (``load_gguf_expert_sources``). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.models.config import ( + FullAttentionGroupConfig, + LinearGatedDeltaGroupConfig, + ModelConfig, + RotaryConfig, +) +from freetoken.models.gguf.dequant import ( + GGML_F32, + GGML_Q4_K, + GGML_Q8_0, + dequantize, + quantize_q8_0, + row_bytes, +) + +if TYPE_CHECKING: + from freetoken.models.gguf.config import GgufConfigShim + + +def _require_tp1(what: str) -> None: + from freetoken.distributed import get_tp_info + + if get_tp_info().size > 1: + raise NotImplementedError( + f"qwen3.5-moe GGUF {what} currently supports TP=1 only " + "(GGUF quant layers and expert banks are not tensor-parallel sharded)." + ) + + +def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: + m = shim.metadata + + def g(key: str): + val = m.get(f"qwen35moe.{key}") + if val is None: + raise KeyError(f"missing GGUF metadata key qwen35moe.{key}") + return val + + num_layers = int(g("block_count")) + hidden = int(g("embedding_length")) + num_qo_heads = int(g("attention.head_count")) + num_kv_heads = int(g("attention.head_count_kv")) + full_head_dim = int(g("attention.key_length")) + max_pos = int(g("context_length")) + interval = int(g("full_attention_interval")) # every Nth (1-indexed) layer is full + + full_ids = tuple(i for i in range(num_layers) if (i + 1) % interval == 0) + linear_ids = tuple(i for i in range(num_layers) if (i + 1) % interval != 0) + + full_rotary = RotaryConfig( + head_dim=full_head_dim, + rotary_dim=int(g("rope.dimension_count")), + max_position=max_pos, + base=float(g("rope.freq_base")), + scaling=None, + ) + full_group = FullAttentionGroupConfig( + name="full", + layer_ids=full_ids, + num_kv_heads=num_kv_heads, + head_dim=full_head_dim, + rotary_config=full_rotary, + ) + # GDN (SSM) dims. conv_dim = 2*key_dim + value_dim (attn_qkv); value_dim = nv*vhead_dim. + key_head_dim = int(g("ssm.state_size")) + value_head_dim = int(g("ssm.state_size")) + linear_group = LinearGatedDeltaGroupConfig( + name="linear", + layer_ids=linear_ids, + num_key_heads=int(g("ssm.group_count")), + num_value_heads=int(g("ssm.time_step_rank")), + key_head_dim=key_head_dim, + value_head_dim=value_head_dim, + conv_kernel_dim=int(g("ssm.conv_kernel")), + output_gate=True, + ) + groups = tuple(sorted((full_group, linear_group), key=lambda grp: grp.layer_ids[0] or 1 << 30)) + + return ModelConfig( + num_layers=num_layers, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=full_head_dim, + hidden_size=hidden, + vocab_size=int(shim.vocab_size), + intermediate_size=0, # routed MoE; no dense MLP + hidden_act="silu", + rms_norm_eps=float(g("attention.layer_norm_rms_epsilon")), + tie_word_embeddings=bool(shim.tie_word_embeddings), + rotary_config=full_rotary, + num_experts=int(g("expert_count")), + num_experts_per_tok=int(g("expert_used_count")), + moe_intermediate_size=int(g("expert_feed_forward_length")), + shared_expert_intermediate_size=int(g("expert_shared_feed_forward_length")), + norm_topk_prob=True, + model_type="qwen3_5_moe", + architectures=list(shim.architectures), + moe_enabled=True, + use_qk_norm=True, + attention_groups=groups, + expert_quant="gguf", + attn_quant="gguf", + dense_quant="gguf", + lm_head_quant="gguf", + moe_weight_format="gguf", + ) + + +def is_gguf_model(config: ModelConfig) -> bool: + return getattr(config, "moe_weight_format", None) == "gguf" + + +# -------------------------------------------------------------------------------------- +# Model layer swap: dense bf16 Linear -> native GGUF-quant ops. +# -------------------------------------------------------------------------------------- + + +def convert_qwen35moe_to_gguf(model, config: ModelConfig) -> None: + """In place: replace the dense projections + embedding with native GGUF ops. + + Quantized in the checkpoint -> swapped to ``GGUFLinear``/``GGUFEmbedding``: the token + embedding (Q8_0) and the (untied) lm_head (Q6_K), full-attention qkv/o (Q8_0), GDN + in_proj_qkvz + out_proj (Q8_0; in_proj_ba stays dense bf16), and the shared-expert + gate_up/down (Q8_0). Left dense bf16/fp32 (F32 in the GGUF): the norms, the two + routers, and the GDN conv1d/A_log/dt_bias. Routed experts stay on the offload cache. + """ + from freetoken.layers.gguf import GGUFEmbedding, GGUFLinear + from freetoken.models.gguf.dequant import GGML_Q6_K, GGML_Q8_0 + + inner = model.model + inner.embed_tokens = GGUFEmbedding( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + quant_type=GGML_Q8_0, + ) + model.lm_head = GGUFLinear( + config.hidden_size, config.vocab_size, GGML_Q6_K, has_bias=False + ) + shared_I = config.shared_expert_intermediate_size + + for layer in inner.layers.op_list: + if layer._is_linear: + g = layer.linear_attn + g.in_proj_qkvz = GGUFLinear( + config.hidden_size, g.conv_dim + g.value_dim, GGML_Q8_0, has_bias=False + ) + g.out_proj = GGUFLinear( + g.value_dim, config.hidden_size, GGML_Q8_0, has_bias=False + ) + else: + attn = layer.self_attn + attn.qkv_proj = GGUFLinear( + config.hidden_size, + attn.num_q * attn.head_dim * 2 + 2 * attn.kv_attn_dim, + GGML_Q8_0, + has_bias=False, + ) + attn.o_proj = GGUFLinear( + attn.qo_attn_dim, config.hidden_size, GGML_Q8_0, has_bias=False + ) + m = layer.mlp + m.shared_expert.gate_up_proj = GGUFLinear( + config.hidden_size, 2 * shared_I, GGML_Q8_0, has_bias=False + ) + m.shared_expert.down_proj = GGUFLinear( + shared_I, config.hidden_size, GGML_Q8_0, has_bias=False + ) + + +# -------------------------------------------------------------------------------------- +# Weight loading: GGUF tensor names -> FreeToken qwen3_5_moe module params. +# -------------------------------------------------------------------------------------- + +# Gemma-style (1+w) norms get +1 baked in; the GDN gated norm / router / shared-gate are +# standard (no +1). +_GEMMA_NORMS = { + "attn_norm.weight", + "post_attention_norm.weight", + "attn_q_norm.weight", + "attn_k_norm.weight", +} + +_EXPERT_SUFFIXES = ("ffn_gate_exps.weight", "ffn_up_exps.weight", "ffn_down_exps.weight") + + +def _to_bf16(t) -> torch.Tensor: + flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.bfloat16) + return flat.reshape(t.shape) + + +def _to_fp32(t) -> torch.Tensor: + flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.float32) + return flat.reshape(t.shape) + + +def _name_to_key(suffix: str) -> tuple[str, bool]: + """gguf layer suffix -> (module-relative key, is_gemma_norm).""" + if suffix == "attn_norm.weight": + return "input_layernorm.weight", True + if suffix == "post_attention_norm.weight": + return "post_attention_layernorm.weight", True + if suffix == "attn_q_norm.weight": + return "self_attn.q_norm.weight", True + if suffix == "attn_k_norm.weight": + return "self_attn.k_norm.weight", True + if suffix == "ssm_norm.weight": + return "linear_attn.norm.weight", False + if suffix == "ffn_gate_inp.weight": + return "mlp.gate.weight", False + if suffix == "ffn_gate_inp_shexp.weight": + return "mlp.shared_expert_gate.weight", False + if suffix == "ssm_conv1d.weight": + return "linear_attn.conv1d.weight", False # fp32; reshaped below + if suffix == "ssm_a": + return "linear_attn.A_log", False # fp32 + if suffix == "ssm_dt.bias": + return "linear_attn.dt_bias", False # fp32 + return None, False + + +# -------------------------------------------------------------------------------------- +# GDN value-head de-interleaving. +# +# llama.cpp stores the GDN *value* projections with the ``mrope_interleaved`` head order: +# the 32 value heads are split into [even heads, odd heads] (head h lives at GGUF position +# ``(h // 2) + (h % 2) * (num_vheads // 2)``). Full-attention heads are NOT interleaved. +# FreeToken uses the HF contiguous head order, so the value-dim projection weights must be +# de-interleaved when loading. Affected: GDN ``in_proj_qkvz`` (the v and z rows), ``out_proj`` +# (the value input columns), and ``in_proj_ba`` (the per-head b/a rows). +# -------------------------------------------------------------------------------------- + + +def _gdn_head_perm(num_vheads: int) -> list[int]: + """GGUF value-head index of each HF head h (``result[h] = old[perm[h]]``).""" + half = num_vheads // 2 + return [(h // 2) + (h % 2) * half for h in range(num_vheads)] + + +def _deint_q8_rows( + packed: torch.Tensor, num_vheads: int, rows_per_head: int +) -> torch.Tensor: + """De-interleave value heads along the packed rows (output dim).""" + m = packed.reshape(num_vheads, rows_per_head, -1) + return m[_gdn_head_perm(num_vheads)].reshape(packed.shape) + + +def _deint_q8_cols( + packed: torch.Tensor, num_vheads: int, blocks_per_head: int, block_bytes: int = 34 +) -> torch.Tensor: + """De-interleave value heads along the packed columns (input dim, per Q8_0 row).""" + m = packed.reshape(packed.shape[0], num_vheads, blocks_per_head * block_bytes) + return m[:, _gdn_head_perm(num_vheads), :].reshape(packed.shape) + + +def _deint_dense_rows(w: torch.Tensor, num_vheads: int) -> torch.Tensor: + """De-interleave value heads along the leading (head) dim of a dense tensor.""" + m = w.reshape(num_vheads, -1) + return m[_gdn_head_perm(num_vheads)].reshape(w.shape) + + +def iter_gguf_weights( + model_path: str, + device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """Yield (param_name, tensor) for every non-expert qwen3_5_moe param. + + Quantized projections stay packed and are yielded as ``.qweight`` (uint8); the F32 + norms/router/GDN b,a dequantize to bf16; conv1d/A_log/dt_bias stay fp32. Full-attention + q/k/v -> ``self_attn.qkv_proj.qweight``, GDN qkv|z -> ``linear_attn.in_proj_qkvz.qweight`` + (Q8_0, concat along the output dim), GDN b|a -> ``linear_attn.in_proj_ba.weight`` (dense + bf16). Routed experts are skipped (offload cache). + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + from freetoken.utils import cached_load_hf_config + + assert not include_moe_experts, ( + "qwen3.5-moe GGUF stores experts as Q4_K/Q5_K/Q6_K and only supports the offload " + "backend; experts are loaded into the offload cache via load_gguf_expert_sources()." + ) + assert include_non_moe + _require_tp1("weight loading") + + config = parse_gguf_config(cached_load_hf_config(model_path)) + full_layers = set( + next(g.layer_ids for g in config.attention_groups + if isinstance(g, FullAttentionGroupConfig)) + ) + # GDN value-head geometry (for mrope_interleaved de-interleave). + gdn = next(g for g in config.attention_groups + if isinstance(g, LinearGatedDeltaGroupConfig)) + n_vheads = gdn.num_value_heads + vhead_dim = gdn.value_head_dim + key_dim = gdn.num_key_heads * gdn.key_head_dim + q8_blocks_per_head = vhead_dim // 32 # Q8_0 block = 32 + + qkv_buf: dict[int, dict[str, torch.Tensor]] = {} + qkvz_buf: dict[int, dict[str, torch.Tensor]] = {} + ba_buf: dict[int, dict[str, torch.Tensor]] = {} + shexp_buf: dict[int, dict[str, torch.Tensor]] = {} + + for t in iter_gguf_tensors(model_path): + name = t.name + if name == "token_embd.weight": + yield "model.embed_tokens.qweight", t.packed() + continue + if name == "output.weight": + yield "lm_head.qweight", t.packed() + continue + if name == "output_norm.weight": + # GGUF stores Gemma norms as the full (1+w) scale; GemmaRMSNorm multiplies by + # it directly (no +1, unlike the HF safetensors form which stores scale-1). + yield "model.norm.weight", _to_bf16(t) + continue + if not name.startswith("blk."): + continue + if any(name.endswith(sfx) for sfx in _EXPERT_SUFFIXES): + continue # routed experts -> offload banks + + layer = int(name.split(".")[1]) + suffix = name.split(".", 2)[2] + base = f"model.layers.{layer}" + + if suffix == "ssm_conv1d.weight": + # conv channels span [q|k|v]; the v channels are value-head interleaved too. + c = _to_fp32(t) + c = c.clone() + c[key_dim * 2:] = _deint_dense_rows(c[key_dim * 2:], n_vheads) + yield f"{base}.linear_attn.conv1d.weight", c.unsqueeze(1) + continue + if suffix == "ssm_a": + # GGUF stores the GDN decay directly as ``A = -exp(A_log)`` (mamba convention, + # value-head interleaved); recover the log-decay the model consumes. + a = _deint_dense_rows(_to_fp32(t), n_vheads) + yield f"{base}.linear_attn.A_log", torch.log(-a) + continue + if suffix == "ssm_dt.bias": + yield f"{base}.linear_attn.dt_bias", _deint_dense_rows(_to_fp32(t), n_vheads) + continue + if suffix == "ffn_gate_inp_shexp.weight": + yield f"{base}.mlp.shared_expert_gate.weight", _to_bf16(t).unsqueeze(0) + continue + + # The GDN b and a projections fuse into a dense in_proj_ba. + if suffix == "ssm_beta.weight": + ba_buf.setdefault(layer, {})["b"] = _to_bf16(t) + elif suffix == "ssm_alpha.weight": + ba_buf.setdefault(layer, {})["a"] = _to_bf16(t) + else: + key, _gemma = _name_to_key(suffix) + if key is not None: + yield f"{base}.{key}", _to_bf16(t) + continue + + is_full = layer in full_layers + if is_full and suffix in ("attn_q.weight", "attn_k.weight", "attn_v.weight"): + qkv_buf.setdefault(layer, {})[suffix[5]] = t.packed() + elif suffix == "attn_qkv.weight": + qkvz_buf.setdefault(layer, {})["qkv"] = t.packed() + elif suffix == "attn_gate.weight": + qkvz_buf.setdefault(layer, {})["z"] = t.packed() + elif suffix == "attn_output.weight": + yield f"{base}.self_attn.o_proj.qweight", t.packed() + elif suffix == "ssm_out.weight": + yield f"{base}.linear_attn.out_proj.qweight", _deint_q8_cols( + t.packed(), n_vheads, q8_blocks_per_head) + elif suffix == "ffn_gate_shexp.weight": + shexp_buf.setdefault(layer, {})["gate"] = t.packed() + elif suffix == "ffn_up_shexp.weight": + shexp_buf.setdefault(layer, {})["up"] = t.packed() + elif suffix == "ffn_down_shexp.weight": + yield f"{base}.mlp.shared_expert.down_proj.qweight", t.packed() + else: + raise ValueError(f"unmapped qwen3.5-moe GGUF tensor: {name}") + + slots = qkv_buf.get(layer) + if slots is not None and {"q", "k", "v"} <= set(slots): + yield f"{base}.self_attn.qkv_proj.qweight", torch.cat( + [slots["q"], slots["k"], slots["v"]], dim=0) + del qkv_buf[layer] + qz = qkvz_buf.get(layer) + if qz is not None and "qkv" in qz and "z" in qz: + qkv = qz["qkv"] # [2*key_dim + value_dim, cols] + qkv = qkv.clone() + # de-interleave the value rows (last value_dim rows of the qkv projection) + qkv[key_dim * 2:] = _deint_q8_rows( + qkv[key_dim * 2:], n_vheads, vhead_dim) + z = _deint_q8_rows(qz["z"], n_vheads, vhead_dim) + yield f"{base}.linear_attn.in_proj_qkvz.qweight", torch.cat([qkv, z], dim=0) + del qkvz_buf[layer] + ba = ba_buf.get(layer) + if ba is not None and "b" in ba and "a" in ba: + b = _deint_dense_rows(ba["b"], n_vheads) + a = _deint_dense_rows(ba["a"], n_vheads) + yield f"{base}.linear_attn.in_proj_ba.weight", torch.cat([b, a], dim=0) + del ba_buf[layer] + gu = shexp_buf.get(layer) + if gu is not None and "gate" in gu and "up" in gu: + yield f"{base}.mlp.shared_expert.gate_up_proj.qweight", torch.cat( + [gu["gate"], gu["up"]], dim=0) + del shexp_buf[layer] + + assert not qkv_buf, f"incomplete qkv groups: {sorted(qkv_buf)}" + assert not qkvz_buf, f"incomplete GDN qkvz groups: {sorted(qkvz_buf)}" + assert not ba_buf, f"incomplete GDN ba groups: {sorted(ba_buf)}" + assert not shexp_buf, f"incomplete shared-expert gate/up: {sorted(shexp_buf)}" + + +# -------------------------------------------------------------------------------------- +# Routed-expert host banks for the offload cache. +# +# The GGUF stores gate/up as Q4_K on every layer, but ``down`` as Q5_K on 37 layers and +# Q6_K on 3 -- heterogeneous row widths the offload cache cannot hold in one uniform bank +# (``set_bank_sources`` requires every layer to share a shape, and ``ggml_moe_a8_vec`` +# derives the row stride from the quant type). We keep ``gate_up`` native Q4_K and +# re-quantize the ``down`` experts to Q8_0 (8-bit, >= Q5_K/Q6_K precision, so no quality +# loss; a uniform per-bank format that fits the cache machinery). +# -------------------------------------------------------------------------------------- + + +def _q8_0_down_row_bytes(I: int) -> int: + return row_bytes(I, GGML_Q8_0) + + +def load_gguf_expert_sources( + model_path: str, config: ModelConfig, *, layer_sink=None +) -> dict[str, list[torch.Tensor]]: + """Per-layer host banks of the routed experts: ``gate_up`` native Q4_K + ``[E, 2I, row_bytes(H, Q4_K)]`` and ``down`` Q8_0 ``[E, H, row_bytes(I, Q8_0)]``. + + ``ffn_{gate,up}_exps`` are each ``[E, I, row_bytes(H, Q4_K)]`` packed and are fused + along the intermediate dim into ``gate_up``; ``ffn_down_exps`` (Q5_K/Q6_K) is + dequantized and re-quantized to Q8_0. Whole layers complete in two writes + (gate_up + down). ``layer_sink=None`` (serving): pin each layer's banks as they + complete via an internal :class:`PinPipeline`. + """ + from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, alloc_layer_banks + from freetoken.models.gguf.reader import iter_gguf_tensors + + _require_tp1("expert banks") + L, E = config.num_layers, config.num_experts + H, I = config.hidden_size, config.moe_intermediate_size + h_rb = row_bytes(H, GGML_Q4_K) + i_rb = row_bytes(I, GGML_Q8_0) + specs = { + "gate_up": ((E, 2 * I, h_rb), torch.uint8), + "down": ((E, H, i_rb), torch.uint8), + } + hb = alloc_layer_banks(specs, L) + banks = {name: [b.tensor for b in hb[name]] for name in specs} + + def _load(sink) -> None: + tracker = LayerCompletionTracker(2, hb, sink) if sink is not None else None + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + continue + layer = int(t.name.split(".")[1]) + if t.name.endswith("ffn_gate_exps.weight"): + banks["gate_up"][layer][:, :I].copy_(t.packed().reshape(E, I, h_rb)) + elif t.name.endswith("ffn_up_exps.weight"): + banks["gate_up"][layer][:, I:] = t.packed().reshape(E, I, h_rb) + elif t.name.endswith("ffn_down_exps.weight"): + flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.float32) + banks["down"][layer].copy_(quantize_q8_0(flat.reshape(E, H, I))) + else: + continue + if tracker is not None: + tracker.note(layer) + + if layer_sink is not None: + _load(layer_sink) + elif torch.cuda.is_available(): + with PinPipeline() as pins: + _load(pins) + else: + _load(None) + return banks + + +__all__ = [ + "parse_gguf_config", + "iter_gguf_weights", + "convert_qwen35moe_to_gguf", + "is_gguf_model", + "load_gguf_expert_sources", +] diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py index eba7fd24f..32954dc52 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -109,6 +109,13 @@ def __init__(self, config: ModelConfig): ) super().__init__() + # GGUF checkpoints carry native block-quantized weights: swap the dense + # projections + embedding for GGUF-quant ops (experts stay on the offload cache). + from .gguf import convert_qwen35moe_to_gguf, is_gguf_model + + if is_gguf_model(config): + convert_qwen35moe_to_gguf(self, config) + def forward(self) -> torch.Tensor: output = self.model.forward(get_global_ctx().batch.input_ids) return self.lm_head.forward(output) diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 0c033ca01..7cffb20c8 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -58,6 +58,14 @@ class ModelSpec: "freetoken.models.qwen3_5_moe", "Qwen3_5MoEForCausalLM", ), + # GGUF (native Q4_K/Q5_K/Q6_K/Q8_0) qwen3.5-moe: same model classes, GGUF config + + # weight loaders (hybrid GatedDeltaNet + full attention, 256 routed experts). + "Qwen35moeGGUFForCausalLM": ModelSpec( + "freetoken.models.qwen3_5_moe", + "Qwen3_5MoEForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + ), # Dense Qwen3.x (no "Moe" in the arch name, num_experts==0, e.g. Qwen3.6-27B). Shares the # qwen3_5_moe package: the decoder routes its MLP through the dense Qwen3_5DenseMLP and the # loader handles the compressed-tensors NVFP4 layout. diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 6a34f3b9b..636d451ee 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -342,6 +342,20 @@ def load_q4_0_moe_expert_sources( return loader(model_path, model_config, layer_sink=layer_sink) +def load_gguf_moe_expert_sources( + model_path: str, + model_config, + *, + layer_sink=None, +) -> dict: + """Load packed GGUF qwen3.5-moe expert source banks (gate_up native Q4_K, down + re-quantized to Q8_0). ``layer_sink`` (converter) streams each completed layer's + banks.""" + _config, spec = _spec_for_model_path(model_path) + loader = _load_attr(spec.module, "load_gguf_expert_sources") + return loader(model_path, model_config, layer_sink=layer_sink) + + def _num_moe_layers(config) -> int: value = getattr(config, "num_moe_layers", None) if value is not None: diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 8b6116ba8..e0801232e 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -252,6 +252,25 @@ def _q4_0_banks(model_path, model_config, device, dtype, dummy, parallel=False, ) +def _gguf_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks: + if parallel: + raise NotImplementedError( + "parallel reader not implemented for gguf: GGUF is a single packed file " + "(not safetensors), so the common reader doesn't apply." + ) + if dummy: + from freetoken.models.weight import dummy_q4_0_moe_expert_sources + + raise NotImplementedError("gguf expert banks have no dummy path; load the real GGUF") + from freetoken.models.weight import load_gguf_moe_expert_sources + + sink = None if dummy else layer_sink + sources = load_gguf_moe_expert_sources(model_path, model_config, layer_sink=sink) + return ExpertBanks( + "gguf", {name: sources[name] for name in _BANK_SCHEMAS["gguf"]}, streamed=sink is not None + ) + + def _dsfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks: args = model_config.dsv4_args assert args is not None, "ds_fp4 expert banks require dsv4_args on the model config" @@ -301,6 +320,7 @@ def _model_setup_override(model_config): "nvfp4": _nvfp4_banks, "ds_fp4": _dsfp4_banks, "q4_0": _q4_0_banks, + "gguf": _gguf_banks, } diff --git a/python/freetoken/moe/fused_gguf.py b/python/freetoken/moe/fused_gguf.py new file mode 100644 index 000000000..9f8995e94 --- /dev/null +++ b/python/freetoken/moe/fused_gguf.py @@ -0,0 +1,52 @@ +"""Grouped expert GEMM over native GGUF Q4_K gate/up + Q8_0 down banks. + +Ports vLLM/sglang's ``_fused_moe_gguf`` MMVQ path onto FreeToken's offload-cache +interface: experts are streamed to the GPU as packed block bytes and dequantized +*inside* ``ggml_moe_a8_vec`` -- no bf16 expert copy is materialized. ``gate_up`` stays +native Q4_K; ``down`` is stored as Q8_0 (re-quantized at load from the GGUF's +Q5_K/Q6_K -- 8-bit, >= the source precision, so no quality loss) because the offload +cache needs a single uniform per-bank format. We use the MMVQ (vector) kernel for both +prefill and decode, mirroring ``fused_experts_gguf_q4_0``. +""" + +from __future__ import annotations + +import torch + +from freetoken.layers.activation import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul +from freetoken.models.gguf.dequant import GGML_Q4_K, GGML_Q8_0 + +_ACT = {"silu": silu_and_mul, "gelu": gelu_and_mul, "gelu_tanh": gelu_tanh_and_mul} + + +def fused_experts_gguf( + hidden_states: torch.Tensor, + gate_up_q: torch.Tensor, # [num_slots, 2I, row_bytes(H, Q4_K)] uint8 + down_q: torch.Tensor, # [num_slots, H, row_bytes(I, Q8_0)] uint8 + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: str, +) -> torch.Tensor: + from freetoken.kernel.gguf import ggml_moe_a8_vec + + act_fn = _ACT.get(activation) + if act_fn is None: + raise ValueError(f"unsupported MoE activation {activation!r}") + + num_tokens = hidden_states.shape[0] + n2 = gate_up_q.shape[1] # 2 * intermediate + h = down_q.shape[1] # hidden + top_k = topk_ids.shape[1] + + gate_up = ggml_moe_a8_vec( + hidden_states, gate_up_q, topk_ids, top_k, int(GGML_Q4_K), n2, num_tokens + ) + inter = act_fn(gate_up) + out = ggml_moe_a8_vec(inter, down_q, topk_ids, 1, int(GGML_Q8_0), h, num_tokens * top_k) + out = out.reshape(num_tokens, top_k, h) * topk_weights.reshape(num_tokens, top_k, 1).to( + out.dtype + ) + return out.sum(dim=1) + + +__all__ = ["fused_experts_gguf"] diff --git a/python/freetoken/moe/nvfp4_backends.py b/python/freetoken/moe/nvfp4_backends.py index e8a3c6f60..c9f84f782 100644 --- a/python/freetoken/moe/nvfp4_backends.py +++ b/python/freetoken/moe/nvfp4_backends.py @@ -46,6 +46,7 @@ import torch from freetoken.utils import init_logger +from freetoken.utils.arch import is_rocm logger = init_logger(__name__) @@ -213,6 +214,17 @@ def select_nvfp4_backend( raise ValueError( f"bad --nvfp4-backend={requested!r}; expected auto, marlin, flashinfer or triton" ) + # AMD: the marlin (vLLM) and flashinfer b12x fused-MoE kernels are NVIDIA-only + # (NVFP4 SASS / CuTe-DSL sm120). auto -> the portable Triton inline-dequant path; + # a forced NVIDIA-only backend fails loudly, never silently degrades. + if device.type == "cuda" and is_rocm(): + if requested in ("marlin", "flashinfer"): + raise RuntimeError( + f"--nvfp4-backend={requested} is NVIDIA-only and unavailable on this " + "ROCm (AMD) build; use --nvfp4-backend triton (or auto). NVFP4 checkpoints " + "can be converted to MXFP4 (freetoken.moe.nvfp4_to_mxfp4) on load." + ) + return "triton" if requested == "triton": return "triton" if activation != "silu": diff --git a/python/freetoken/moe/nvfp4_to_mxfp4.py b/python/freetoken/moe/nvfp4_to_mxfp4.py new file mode 100644 index 000000000..f4162a032 --- /dev/null +++ b/python/freetoken/moe/nvfp4_to_mxfp4.py @@ -0,0 +1,239 @@ +"""NVFP4 -> MXFP4 (gpt-oss/FreeToken ``mxfp4_triton``) weight converter. + +ModelOpt NVFP4 (the format stored in FreeToken's native ``nvfp4`` banks) packs e2m1 +codes with a *fp8-e4m3* per-16 block scale and a per-output-row fp16 *global* scale. +MXFP4 (FreeToken's ``mxfp4_triton`` banks) packs e2m1 codes with an *e8m0* per-32 +block scale and a per-block bias, and is the native format of the gpt-oss family -- +the AMD-supported quant matrix alongside BF16/GGUF. + +This module converts a checkpoint's NVFP4 expert weights to MXFP4 *on load* (once, +cached per model), so a checkpoint that only ships NVFP4 can still run on AMD via the +portable MXFP4 path (the converter runs on the host, not on the GPU). The two formats +share the e2m1 code-packing (2 codes per byte, low nibble first), so only the scale +granularity (16 -> 32) and scale format (e4m3 -> e8m0) change. + +Layouts handled here (per projection/expert, ``N`` = output rows, ``K`` = input cols): + +* input (NVFP4, native ModelOpt rows): ``packed [N, K//2]`` uint8, ``scale [N, K//16]`` + fp8-e4m3, ``global [N]`` fp16. +* output (MXFP4, ``mxfp4_triton`` bank layout): ``blocks_t [N, K//2]`` uint8, + ``scales_t [N, K//32]`` uint8 e8m0. (FreeToken's MXFP4 stores the projection + transposed, N innermost, matching gpt-oss -- the converter emits that shape.) + +All numeric work is done in numpy so the core is unit-testable without a GPU/torch +runtime; the public entry converts torch tensors to/from numpy on the host. +""" + +from __future__ import annotations + +import math +from typing import Sequence + +try: + import numpy as _np +except ImportError: # pragma: no cover - numpy is a hard dep + _np = None + +__all__ = [ + "convert_nvfp4_to_mxfp4", + "dequantize_nvfp4_block", + "fp4_e2m1_table", + "e8m0_scale_and_codes", +] + +# --------------------------------------------------------------------------- +# e2m1 (fp4) code -> value table, and e8m0 (block scale) encode. +# +# e2m1: 1 sign + 2 exponent + 1 mantissa. With exponent bias 1 the finite set is: +# code value code value +# 0 0.0 8 -0.0 +# 1 0.5 9 -0.5 +# 2 1.0 10 -1.0 +# 3 1.5 11 -1.5 +# 4 2.0 12 -2.0 +# 5 3.0 13 -3.0 +# 6 4.0 14 -4.0 +# 7 6.0 15 -6.0 +# --------------------------------------------------------------------------- +_FP4_CODES = ( + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, +) +_FP4_TABLE = _np.asarray(_FP4_CODES, dtype=_np.float32) +# Magnitudes sorted ascending for the nearest-code search. +_FP4_SORT = _np.asarray(sorted(abs(v) for v in _FP4_CODES[1:8]), dtype=_np.float32) +_FP4_SORT_SIGN = _np.asarray([1.0 if i < 4 else -1.0 for i in range(len(_FP4_SORT))], + dtype=_np.float32) + + +def fp4_e2m1_table() -> Sequence[float]: + """The 16 e2m1 values keyed by 4-bit code (index == code).""" + return list(_FP4_CODES) + + +def _nearest_e2m1_codes(values: _np.ndarray) -> _np.ndarray: + """Nearest e2m1 *code* for each fp32 ``values`` (signed, including 0/NaN).""" + a = _np.abs(values) + diff = _np.abs(a[..., None] - _FP4_SORT) # [..., 7] + idx = diff.argmin(axis=-1) + mag = _FP4_SORT[idx] + neg = _np.signbit(values) + code = (idx + 1).astype(_np.uint8) # _FP4_SORT[i] == table[i+1]; positive codes 1..7 + out = _np.where(neg, code | 0x8, code) + # Magnitudes below the smallest representable value (0.5) round to +0. + return _np.where(mag < 0.25, 0, out) + + +def e8m0_scale_and_codes(values: _np.ndarray, block: int = 32) -> tuple[_np.ndarray, _np.ndarray]: + """Return ``(scale_codes, fp4_codes)`` for ``values`` shaped ``[..., block]``: an + e8m0 ``uint8`` scale per block (the smallest power-of-2 scale covering the block + max-abs, in the MX ``2**(v-127)`` encoding) and the requantized 4-bit codes. + + ``scale_codes`` has shape ``values.shape[:-1]``; ``fp4_codes`` matches ``values``. + """ + v = values.reshape(-1, block) + amax = _np.max(_np.abs(v), axis=-1) + # e2m1 max positive magnitude is 6.0; choose the smallest power-of-2 scale so the + # block max-abs maps near the top of the e2m1 range (best precision): + # s = 2^ceil(log2(max_abs / 6.0)), stored as e8m0 code v with 2**(v-127) == s. + amax_safe = _np.maximum(amax, 1e-38) + exp = _np.ceil(_np.log2(amax_safe / 6.0)).astype(_np.float32) + exp = _np.where(amax == 0.0, 0.0, exp) + scale_codes = (127.0 + exp).astype(_np.uint8) # v-127 == exp + scale_v = (2.0 ** exp).astype(_np.float32) + # Requantize values in the block by its scale, then nearest-e2m1-code. + q = v / scale_v[:, None] + codes = _fp4_quantize(q).astype(_np.uint8) + return scale_codes.reshape(values.shape[:-1]), codes.reshape(values.shape) + + +def _fp4_quantize(values: _np.ndarray) -> _np.ndarray: + a = _np.abs(values) + diff = _np.abs(a[..., None] - _FP4_SORT) + idx = diff.argmin(axis=-1) + mag = _FP4_SORT[idx] + out = _np.where(mag < 0.25, 0, (idx + 1).astype(_np.uint8)) + out = _np.where(_np.signbit(values), out | 8, out) + return out + + +# --------------------------------------------------------------------------- +# NVFP4 -> MXFP4 +# --------------------------------------------------------------------------- + + +def dequantize_nvfp4_block( + packed: _np.ndarray, + scale: _np.ndarray, + global_scale: _np.ndarray, + *, + block: int = 16, +) -> _np.ndarray: + """Dequantize one native NVFP4 projection back to fp32. + + ``packed [N, K//2]`` uint8 (e2m1 pairs), ``scale [N, K//16]`` fp32 (already + converted from fp8-e4m3), ``global_scale [N]`` fp32. Returns ``[N, K]`` fp32. + """ + N, K2 = packed.shape + K = K2 * 2 + lo = (packed & 0x0F).astype(_np.uint8) + hi = (packed >> 4).astype(_np.uint8) + codes = _np.stack([lo, hi], axis=-1).reshape(N, K) # [N, K] + vals = _FP4_TABLE[codes.astype(_np.int64)] # [N, K] + # Per-16 block scale broadcast over K. + bs = _np.repeat(scale, block, axis=-1) # [N, K] + return (vals * bs).astype(_np.float32) * global_scale[:, None].astype(_np.float32) + + +def _pack_codes(codes: _np.ndarray) -> _np.ndarray: + """Pack ``[N, K]`` uint8 4-bit codes -> ``[N, K//2]`` uint8 (low nibble first).""" + N, K = codes.shape + even = codes[..., 0::2] + odd = codes[..., 1::2] + return (even | (odd << 4)).astype(_np.uint8) + + +def convert_nvfp4_to_mxfp4( + packed, + scale, + global_scale, + *, + axis: int = -1, + block: int = 32, +): + """Convert one projection's native NVFP4 expert weights to the MXFP4 layout. + + Args: + packed: ``[..., K//2]`` uint8 e2m1 pairs (low nibble = first code). + scale: ``[..., K//16]`` fp8-e4m3 block scale (fp32/fp16 input accepted). + global_scale: ``[...]`` per-output-row fp16 global scale. + axis: the K (contraction) axis along which blocks are grouped. + + Returns ``(mxfp4_packed [..., K//2] uint8, mxfp4_scales [..., K//32] uint8 e8m0)`` + matching the ``mxfp4_triton`` per-expert bank shape (K innermost). + """ + np = _np + packed = np.asarray(packed) + scale = np.asarray(scale, dtype=np.float32) + global_scale = np.asarray(global_scale, dtype=np.float32) + + if axis not in (-1, packed.ndim - 1): + raise NotImplementedError("converter requires the K axis to be innermost") + + # Dequantize NVFP4 to fp32, move K to the last axis. + K2 = packed.shape[-1] + K = K2 * 2 + codes = np.stack([packed & 0x0F, (packed >> 4)], axis=-1).reshape( + *packed.shape[:-1], K + ) + vals = _FP4_TABLE[codes.astype(np.int64)] + bs = np.repeat(scale, 16, axis=-1) + f32 = (vals * bs).astype(np.float32) * global_scale[..., None].astype(np.float32) + + # Requantize to per-`block` e8m0 + e2m1. + flat = f32.reshape(-1, K) + # pad to a multiple of block for the reshape (K is a multiple of 32 in practice) + n_blocks = K // block + flat_b = flat[:, : n_blocks * block].reshape(-1, block) + scale_codes, mxfp4_codes = e8m0_scale_and_codes(flat_b, block=block) + out_codes = mxfp4_codes.reshape(flat.shape[0], n_blocks * block) + + mxfp4_packed = _pack_codes(out_codes) # [..., K//2] + mxfp4_scales = scale_codes.reshape(*packed.shape[:-1], n_blocks) + return mxfp4_packed, mxfp4_scales + + +# --------------------------------------------------------------------------- +# torch-tensor entrypoint +# --------------------------------------------------------------------------- + + +def _to_numpy(t: object) -> _np.ndarray: + if _np is not None and isinstance(t, _np.ndarray): + return t + try: + return t.detach().cpu().numpy() + except Exception as exc: # pragma: no cover + raise TypeError( + f"convert_nvfp4_to_mxfp4 expects torch tensors or numpy arrays, got {type(t)!r}" + ) from exc + + +def convert_torch_nvfp4_to_mxfp4( + packed, + scale, + global_scale, + *, + block: int = 32, +): + """Torch-tensor variant returning ``(mxfp4_packed, mxfp4_scales)`` torch tensors on + the same device as ``packed`` (host converter: input tensors are pulled to CPU and + the results copied back). Used at load time, cached per model.""" + import torch + + device = packed.device + dtype = packed.dtype + p, s = convert_nvfp4_to_mxfp4( + _to_numpy(packed), _to_numpy(scale), _to_numpy(global_scale), block=block + ) + return torch.from_numpy(p).to(device=device, dtype=dtype), torch.from_numpy(s).to(device=device) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 6ee764061..56264ca54 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -45,6 +45,10 @@ # native GGUF Q4_0 experts: packed block bytes per output row, dequantized inside # the borrowed ggml MoE kernels. gate_up [L*E, 2I, H//32*18], down [L*E, H, I//32*18]. "q4_0": ("gate_up", "down"), + # native GGUF qwen3.5-moe experts: gate_up native Q4_K [L*E, 2I, row_bytes(H,Q4_K)], + # down re-quantized to Q8_0 [L*E, H, row_bytes(I,Q8_0)] (a uniform format the cache + # can hold; the source GGUF down is Q5_K/Q6_K, both >= ... 8-bit >= source precision). + "gguf": ("gate_up", "down"), # native ModelOpt rows for the Triton inline-dequant kernels: packed e2m1 codes + # fp8-e4m3 per-16 block scales + per-output-row fp16 globals (w1/w3 carry distinct # globals, and folding them into the e4m3 block scales would underflow) @@ -83,6 +87,8 @@ "bf16": lambda H, I: 3 * I * H * 2, "fp8_block": lambda H, I: 3 * I * H + ((2 * I // 128) * (H // 128) + (H // 128) * (I // 128)) * 2, "q4_0": lambda H, I: 2 * I * (H // 32) * 18 + H * (I // 32) * 18, + # gate_up Q4_K row_bytes(H, Q4_K)=H//256*144; down Q8_0 row_bytes(I, Q8_0)=I//32*34. + "gguf": lambda H, I: 2 * I * (H // 256) * 144 + H * (I // 32) * 34, "nvfp4": lambda H, I: 2 * I * (H // 2 + H // 16 + 2) + H * (I // 2 + I // 16 + 2), "mxfp4": lambda H, I: 2 * I * (H // 2 + H // 32 + 2) + H * (I // 2 + I // 32 + 2), "ds_fp4": lambda H, I: 2 * I * (H // 2 + H // 32) + H * (I // 2 + I // 32), diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 4954c5f55..2a533be70 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -683,6 +683,25 @@ def _infer_reasoning_parser(model_path: str) -> str | None: kwargs["tp_info"] = DistributedInfo(0, kwargs["tensor_parallel_size"]) del kwargs["tensor_parallel_size"] + # ROCm (AMD) has no NVIDIA-native NVFP4/Marlin path. Reject NVIDIA-only NVFP4 backends + # at parse time with a clean error, and warn for the resident fused MoE backend (the + # offload/cpu/hybrid family is the supported AMD path). + from freetoken.utils.arch import is_rocm + + if is_rocm(): + nvfp4 = kwargs.get("nvfp4_backend") + if nvfp4 in ("marlin", "flashinfer"): + raise SystemExit( + f"--nvfp4-backend {nvfp4} is NVIDIA-only and unavailable on ROCm/AMD; " + f"use --nvfp4-backend triton (inline-dequant) or auto." + ) + if kwargs.get("moe_backend") == "fused": + logger = init_logger(__name__) + logger.warning( + "--moe-backend fused relies on NVIDIA-native fused GEMM; on ROCm/AMD the " + "supported family is offload/hybrid/cpu (the triton/offload path)." + ) + result = ServerArgs(**kwargs) logger = init_logger(__name__) logger.info(f"Parsed arguments:\n{result}") diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py index 2e4ad15f2..01f6c57c0 100644 --- a/python/freetoken/utils/__init__.py +++ b/python/freetoken/utils/__init__.py @@ -1,5 +1,9 @@ from .arch import ( + current_gpu_name, + device_kind, is_arch_supported, + is_cuda, + is_rocm, is_sm90_family, is_sm90_supported, is_sm100_family, @@ -34,7 +38,11 @@ "load_tokenizer", "load_toolcall_anchor_id", "init_logger", + "device_kind", + "current_gpu_name", "is_arch_supported", + "is_cuda", + "is_rocm", "is_sm90_family", "is_sm90_supported", "is_sm100_family", diff --git a/python/freetoken/utils/arch.py b/python/freetoken/utils/arch.py index 8c1c6c3d5..3dd781f41 100644 --- a/python/freetoken/utils/arch.py +++ b/python/freetoken/utils/arch.py @@ -4,14 +4,106 @@ from typing import Tuple +def device_kind() -> str: + """Compute backend the current torch build targets: ``"cuda"`` (NVIDIA), ``"rocm"`` + (AMD/HIP) or ``"cpu"``. Keyed on the *build* (torch.version.hip vs torch.version.cuda), + independent of whether a GPU is present, so feature-gating and graceful-degradation + decisions can be made before any device is available. On ROCm torch, torch.version.cuda + is None and torch.version.hip is set; on CUDA torch the inverse holds.""" + try: + import torch.version + except Exception: + return "cpu" + if getattr(torch.version, "hip", None): + return "rocm" + if getattr(torch.version, "cuda", None): + return "cuda" + return "cpu" + + +def is_rocm() -> bool: + """True when the installed torch is a ROCm (AMD) build.""" + return device_kind() == "rocm" + + +def is_cuda() -> bool: + """True when the installed torch is a CUDA (NVIDIA) build.""" + return device_kind() == "cuda" + + +@functools.cache +def current_gpu_name() -> str | None: + """Device name of the current CUDA-capable device (``torch.cuda.get_device_name``), or + None if torch is unavailable / no device. On ROCm this returns the AMD card name through + the torch.cuda compat layer.""" + try: + import torch + + if not torch.cuda.is_available(): + return None + return torch.cuda.get_device_name(torch.cuda.current_device()) + except Exception: + return None + + @functools.cache def _get_torch_cuda_version() -> Tuple[int, int] | None: - import torch - import torch.version + """Compute capability ``(major, minor)`` of the current CUDA device, or None when it + cannot be determined. Returns None on ROCm torch (gfx archs are not a CUDA compute + capability), when no CUDA device is present, and when torch itself is unavailable -- + so every ``is_sm*``/``is_arch_supported`` gate degrades to the portable path.""" + try: + import torch + import torch.version + + if not torch.cuda.is_available() or not torch.version.cuda: + return None + return torch.cuda.get_device_capability() + except Exception: + return None - if not torch.cuda.is_available() or not torch.version.cuda: + +@functools.cache +def _get_gfx_arch() -> int | None: + """Numeric gfx arch of the current device (e.g. 1100 for ``gfx1100``) on ROCm, or + None when torch is unavailable / not ROCm / no device present. Used by + :func:`is_gfx_arch_ge` for AMD feature gating.""" + try: + import torch + import torch.version + + if not torch.version.hip or not torch.cuda.is_available(): + return None + import re + + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + # ROCm torch exposes the exact gfx string (e.g. ``gfx1100``) as gcnArchName, + # which is the reliable field; the marketing device name is often just + # ``Radeon RX 7900 XTX`` and carries no gfx marker. + gcn = getattr(props, "gcnArchName", None) + if gcn: + m = re.search(r"gfx(\d{3,4})", str(gcn)) + if m: + return int(m.group(1)) + name = torch.cuda.get_device_name(torch.cuda.current_device()) + if name: + m = re.search(r"gfx(\d{3,4})", name) + if m: + return int(m.group(1)) return None - return torch.cuda.get_device_capability() + except Exception: + return None + + +def is_gfx_arch_ge(arch_int: int) -> bool: + """True on ROCm when the current gfx arch number is >= ``arch_int`` (e.g. + ``is_gfx_arch_ge(1100)`` for RDNA 3 / RX 7000). Parses the full gfx string + (``gfx1100`` -> 1100) rather than a CUDA-style ``(major, minor)`` tuple. Returns + False on CUDA and CPU builds, so every gfx gate degrades to the portable path.""" + gfx = _get_gfx_arch() + if gfx is None: + return False + return gfx >= arch_int def is_arch_supported(major: int, minor: int = 0) -> bool: diff --git a/python/freetoken/utils/graph_gate.py b/python/freetoken/utils/graph_gate.py new file mode 100644 index 000000000..b528a208f --- /dev/null +++ b/python/freetoken/utils/graph_gate.py @@ -0,0 +1,189 @@ +"""HIP/CUDA graph-capture parity probe. + +The Inc-1 hard gate: whether ``torch.cuda.graph`` graph capture works on the target +GPU is the single highest-informational-risk assumption for AMD (ROCm) support. +This module probes it once and records a PASS/FAIL + device result that the rest +of the plan (Inc 8) reads. On CUDA it is expected to PASS; on ROCm it may fail on +some consumer cards, in which case Inc 8 must use the kernel-launch decode path. + +The result is cached to disk under the user cache dir so it survives across runs, +and keyed by device kind + device name so a change of GPU invalidates it. +""" +from __future__ import annotations + +import json +import os +from functools import lru_cache + +_CACHE_FILE = "freetoken_graph_gate.json" + + +def _cache_dir() -> str: + base = os.environ.get("XDG_CACHE_HOME") or os.path.join( + os.path.expanduser("~"), ".cache" + ) + path = os.path.join(base, "freetoken") + os.makedirs(path, exist_ok=True) + return path + + +def _cache_path() -> str: + return os.path.join(_cache_dir(), _CACHE_FILE) + + +def _device_kind() -> str: + from freetoken.utils.arch import device_kind + + return device_kind() + + +@lru_cache(maxsize=1) +def _device_name() -> str | None: + """Best-effort current device name via torch.cuda, or None when no device / torch.""" + try: + import torch + + if not torch.cuda.is_available(): + return None + return torch.cuda.get_device_name(torch.cuda.current_device()) + except Exception: + return None + + +def probe_graph_capture() -> dict: + """Run the actual capture probe on the current device. Returns a dict: + ``{"device_kind": ..., "device": ..., "ok": bool, "detail": str}``. + + The GEMM-capture attempt is run in a **fresh subprocess**: on some ROCm builds a + hipBLASLt/capture failure raises an uncatchable fatal HIP error (error 900) that + aborts the whole process, so running it inline would crash the caller (and, worse, + a decode path that attempted graph capture would die). A subprocess lets a fatal + failure surface as a clean ``ok:false`` result instead. + """ + device = _device_name() + try: + import torch + + if not torch.cuda.is_available(): + return { + "device_kind": _device_kind(), + "device": device, + "ok": False, + "detail": "no CUDA-capable device available", + } + except Exception as exc: + detail = next((line.strip() for line in str(exc).splitlines() if line.strip()), "") + return { + "device_kind": _device_kind(), + "device": device, + "ok": False, + "detail": f"torch unavailable: {type(exc).__name__}: {detail}", + } + + # The child probes both an elementwise op (capturable on both backends) and a GEMM + # (hipBLASLt on ROCm), which is what a real decode forward would run. The GEMM is + # the discriminating case: on this ROCm build it fatally aborts -> child exit != 0. + import json as _json + import subprocess as _subprocess + import sys as _sys + + child = _subprocess.run( + [_sys.executable, "-c", _CAPTURE_CHILD], + capture_output=True, + text=True, + timeout=120, + ) + if child.returncode != 0: + detail = next( + (l.strip() for l in child.stderr.splitlines() if l.strip()), + f"graph capture subprocess aborted (rc={child.returncode})", + ) + return { + "device_kind": _device_kind(), + "device": device, + "ok": False, + "detail": f"fatal during capture: {detail[:240]}", + } + try: + data = _json.loads(child.stdout) + except Exception: + return { + "device_kind": _device_kind(), + "device": device, + "ok": False, + "detail": f"unparseable probe output: {child.stdout[:120]}", + } + data.setdefault("device_kind", _device_kind()) + data.setdefault("device", device) + return data + + +#: Child body for the graph-capture probe (see :func:`probe_graph_capture`). Prints a +#: JSON line ``{"ok": true/false, "detail": ...}`` and exits nonzero on a fatal abort. +_CAPTURE_CHILD = r""" +import json, sys +import torch +try: + torch.cuda.synchronize() + s = torch.cuda.Stream() + x = torch.randn(8, 8, device='cuda') + # elementwise (capturable) first, then a GEMM (hipBLASLt on ROCm) + with torch.cuda.stream(s): + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + torch.add(x, x) + s.synchronize() + with torch.cuda.stream(s): + g = torch.cuda.CUDAGraph() + a = torch.randn(64, 64, device='cuda') + with torch.cuda.graph(g): + torch.mm(a, a) + s.synchronize() + torch.cuda.synchronize() + print(json.dumps({"ok": True, "detail": "elementwise+GEMM capture/replay succeeded"})) +except Exception as e: + print(json.dumps({"ok": False, "detail": f"{type(e).__name__}: {str(e)[:160]}"})) +""" + + +def _load_cached() -> dict | None: + try: + with open(_cache_path()) as f: + data = json.load(f) + if ( + data.get("device_kind") == _device_kind() + and data.get("device") == _device_name() + ): + return data + except Exception: + pass + return None + + +def run_graph_gate() -> dict: + """Run (or reuse the cached result for the current device of) the capture probe.""" + cached = _load_cached() + if cached is not None: + return cached + result = probe_graph_capture() + try: + with open(_cache_path(), "w") as f: + json.dump(result, f) + except Exception: + pass + return result + + +@lru_cache(maxsize=1) +def graph_capture_status() -> str: + """Cached graph-capture status: ``"pass"``, ``"fail"``, or ``"unknown"`` (no device / + probe unavailable). Inc 8 reads this to pick HIP-graph vs kernel-launch decode.""" + try: + result = run_graph_gate() + if result["ok"]: + return "pass" + if result.get("device_kind") and result["device_kind"] != "cpu": + return "fail" + return "unknown" + except Exception: + return "unknown" diff --git a/python/freetoken/utils/torch_utils.py b/python/freetoken/utils/torch_utils.py index 9422b9e7d..16223210b 100644 --- a/python/freetoken/utils/torch_utils.py +++ b/python/freetoken/utils/torch_utils.py @@ -21,6 +21,16 @@ def torch_dtype(dtype: torch.dtype): def nvtx_annotate(name: str, layer_id_field: str | None = None): + from freetoken.utils.arch import is_rocm + + # ROCm torch has no torch.cuda.nvtx; mapping to roctx is future work. Under ROCm we + # pass through (no-op decorator) so AMD runs are not coupled to NVIDIA-only tooling. + if is_rocm(): + def passthrough(fn): + return fn + + return passthrough + import torch.cuda.nvtx as nvtx def decorator(fn): @@ -35,3 +45,12 @@ def wrapper(self, *args, **kwargs): return wrapper return decorator + + +def graph_capture(): + """Context manager that captures a CUDA (or, on ROCm, HIP) graph on the current + stream via ``torch.cuda.graph``. Correctness is validated once by the Inc-1 + ``freetoken.utils.graph_gate`` probe; callers rely on that result.""" + import torch + + return torch.cuda.graph() diff --git a/setup.py b/setup.py index cfe41b7d8..c4cf4f1f6 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,7 @@ from __future__ import annotations import importlib.util +import os from pathlib import Path from setuptools import setup @@ -10,15 +11,32 @@ ROOT = Path(__file__).parent -def _check_toolchain() -> None: +def _load_toolchain(): path = ROOT / "python" / "freetoken" / "kernel" / "_toolchain.py" spec = importlib.util.spec_from_file_location("_freetoken_toolchain", path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) - module.check_nvcc_matches_torch() + return module + + +def _check_toolchain() -> None: + module = _load_toolchain() + if module.is_rocm_torch(): + module.check_hip_matches_torch() + else: + module.check_nvcc_matches_torch() + + +def _rocm_home() -> Path | None: + for env in ("ROCM_HOME", "HIP_PATH"): + root = os.getenv(env) + if root and (Path(root) / "include").exists(): + return Path(root) + default = Path("/opt/rocm") + return default if (default / "include").exists() else None -def _cuda_runtime_paths() -> tuple[list[str], list[str]]: +def _cuda_runtime_paths() -> tuple[list[str], list[str], list[str]]: if CUDA_HOME is None: raise RuntimeError( "CUDA_HOME is required to build freetoken.kernel._pinned_tensor " @@ -28,12 +46,43 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: library_dirs = [str(cuda_home / "lib64")] if (cuda_home / "lib").exists(): library_dirs.append(str(cuda_home / "lib")) - return [str(cuda_home / "include")], library_dirs + return [str(cuda_home / "include")], library_dirs, ["cudart"] + + +def _rocm_runtime_paths() -> tuple[list[str], list[str], list[str]]: + home = _rocm_home() + if home is None: + raise RuntimeError( + "ROCm torch detected but no ROCm toolkit found. Install ROCm (e.g. /opt/rocm) " + "matching torch's HIP version to build freetoken.kernel._pinned_tensor " + "(it links the HIP runtime API)." + ) + include_dirs = [str(home / "include")] + library_dirs = [] + for sub in ("lib", "lib64"): + if (home / sub).exists(): + library_dirs.append(str(home / sub)) + # HIP host APIs (hipHostMalloc/hipHostRegister/hipHostGetDevicePointer) and HIP + # graph nodes all live in the HIP runtime, amdhip64. (hiprt is a separate optional + # library not present on all ROCm installs; linking it would break the build.) + libraries = ["amdhip64"] + return include_dirs, library_dirs, libraries + + +def _runtime_paths() -> tuple[list[str], list[str], list[str], list[str]]: + """Returns (include_dirs, library_dirs, libraries, compile_defs) for the active backend.""" + module = _load_toolchain() + if module.is_rocm_torch(): + include_dirs, library_dirs, libraries = _rocm_runtime_paths() + return include_dirs, library_dirs, libraries, ["-DUSE_HIP=1", "-DUSE_ROCM=1"] + include_dirs, library_dirs, libraries = _cuda_runtime_paths() + return include_dirs, library_dirs, libraries, [] -cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths() +include_dirs, library_dirs, libraries, compile_defs = _runtime_paths() _check_toolchain() +_extra_compile_args = ["-O3", "-std=c++17", *compile_defs] setup( ext_modules=[ @@ -42,13 +91,14 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: sources=[ "python/freetoken/kernel/csrc/pinned_tensor.cpp", ], - include_dirs=cuda_include_dirs, - library_dirs=cuda_library_dirs, - libraries=["cudart"], - extra_compile_args=["-O3", "-std=c++17"], + include_dirs=include_dirs, + library_dirs=library_dirs, + libraries=libraries, + extra_compile_args=_extra_compile_args, ), - # CPU-compute MoE executor for --moe-backend cpu. Links cudart for the - # cudaLaunchHostFunc submit/sync graph nodes; the bf16 GEMV microkernels + # CPU-compute MoE executor for --moe-backend cpu. On CUDA it links cudart for the + # cudaLaunchHostFunc submit/sync graph nodes; on ROCm those become HIP graph nodes + # (hipLaunchHostFunc) and we link the HIP runtime instead. The bf16 GEMV microkernels # use per-function target attributes (avx512bf16/avx512f) + a runtime # __builtin_cpu_supports dispatch, so the single binary stays portable # (scalar fallback) -- no global -march is set. @@ -57,10 +107,10 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: sources=[ "python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp", ], - include_dirs=cuda_include_dirs, - library_dirs=cuda_library_dirs, - libraries=["cudart"], - extra_compile_args=["-O3", "-std=c++17", "-pthread"], + include_dirs=include_dirs, + library_dirs=library_dirs, + libraries=libraries, + extra_compile_args=_extra_compile_args + ["-pthread"], ), ], cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)}, diff --git a/tests/attention/test_torch_backend.py b/tests/attention/test_torch_backend.py new file mode 100644 index 000000000..f655a959b --- /dev/null +++ b/tests/attention/test_torch_backend.py @@ -0,0 +1,62 @@ +"""The debug ``"torch"`` attention backend (Inc 4 of fix-attention). + +Verifies (a) the backend is registered and selectable, and (b) its pure-PyTorch +GQA attention math (with causal masking and the per-head output gate) matches +``torch.nn.functional.scaled_dot_product_attention`` on a synthetic case with the +qwen35moe full-attention geometry (num_q=16, num_kv=2, head_dim=256, gate). + +This is the ground-truth backend used to A/B the production triton path. +""" + +import torch + + +def _attention_math(q, ks, vs, lk, scale, group, gate): + """Mirror of TorchAttentionBackend's per-request attention (fp32 intermediates).""" + ks = ks.repeat_interleave(group, dim=1).float() + vs = vs.repeat_interleave(group, dim=1).float() + scores = torch.einsum("qhd,khd->hqk", q.float(), ks) * scale + lq = q.shape[0] + rows = torch.arange(lq) + cols = torch.arange(lk) + masked = (cols[None, :] > (lk - lq + rows)[:, None]).to(scores.device) + scores = scores.masked_fill(masked[None], float("-inf")) + probs = torch.softmax(scores, dim=-1) + o = torch.einsum("hqk,khd->qhd", probs, vs) # [lq, num_q, hd] + lq_, nq, hd_ = o.shape + o = o.reshape(lq_, nq * hd_) * torch.sigmoid(gate.float()) # [lq, num_q*hd] + return o.reshape(lq_, nq, hd_) + + +def test_torch_backend_registered(): + from freetoken.attention import SUPPORTED_ATTENTION_BACKENDS, attention_backend_info, AttnType + + assert "torch" in SUPPORTED_ATTENTION_BACKENDS.supported_names() + info = attention_backend_info("torch") + assert AttnType.FULL in info.supported_types + assert info.hybrid_linear_ok # must coexist with GDN layers + + +def test_torch_attention_matches_sdpa_prefill(): + num_q, num_kv, hd, group = 16, 2, 256, 8 + T = 4 + scale = hd ** -0.5 + q = torch.randn(T, num_q, hd) + k = torch.randn(T, num_kv, hd) + v = torch.randn(T, num_kv, hd) + gate = torch.randn(T, num_q * hd) + + out = _attention_math(q, k, v, T, scale, group, gate) + # reference: sdpa on GQA-expanded heads, then gate + no o_proj here + ke = k.repeat_interleave(group, dim=1) + ve = v.repeat_interleave(group, dim=1) + ref = torch.nn.functional.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0), + ke.transpose(0, 1).unsqueeze(0), + ve.transpose(0, 1).unsqueeze(0), + is_causal=True, + )[0].transpose(0, 1) + ref = ref.reshape(T, num_q * hd) * torch.sigmoid(gate) + ref = ref.reshape(T, num_q, hd) + + assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3) diff --git a/tests/engine/test_attention_backend_rocm.py b/tests/engine/test_attention_backend_rocm.py new file mode 100644 index 000000000..d8ff4e5c9 --- /dev/null +++ b/tests/engine/test_attention_backend_rocm.py @@ -0,0 +1,75 @@ +"""ROCm (AMD) attention-backend resolution: no NVIDIA-native backend may be selected. + +On ROCm, is_sm90/100 gates are False and flashinfer/sgl_kernel are treated as +unavailable, so auto resolution must fall through to the portable Triton backend +for FULL-attention models. +""" + +import pytest + + +def _engine_config(**overrides): + from types import SimpleNamespace + + import torch + + from freetoken.distributed import DistributedInfo + from freetoken.engine.config import EngineConfig + + config = EngineConfig( + model_path="/tmp/freetoken-test-model", + tp_info=DistributedInfo(rank=0, size=1), + dtype=torch.bfloat16, + **overrides, + ) + object.__setattr__( + config, + "model_config", + SimpleNamespace( + has_swa_attention=False, + has_linear_attention=False, + is_moe=False, + num_layers=10, + expert_quant="none", + ), + ) + return config + + +def _patch_rocm(monkeypatch): + from freetoken.engine import engine + from freetoken.kernel import backend + + monkeypatch.setattr(engine, "is_sm100_family", lambda: False) + monkeypatch.setattr(engine, "is_sm90_family", lambda: False) + monkeypatch.setattr(engine, "_flashinfer_available", lambda: False) + monkeypatch.setattr(engine, "_sgl_flash_attn_available", lambda: False) + monkeypatch.setattr(backend, "is_rocm", lambda: True) + monkeypatch.setattr(engine, "is_rocm", lambda: True) + + +def test_rocm_auto_resolves_to_triton(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_rocm(monkeypatch) + config = _engine_config(attention_backend="auto") + _adjust_config(config) + assert config.attention_backend == "triton" + + +def test_rocm_explicit_nvidia_backend_rejected(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_rocm(monkeypatch) + # flashinfer/sgl absent and arch gates false -> trtllm/fa/fi requirements unmet. + for backend in ("fi", "fa", "trtllm"): + config = _engine_config(attention_backend=backend) + with pytest.raises(RuntimeError): + _adjust_config(config) + + +def test_sgl_flash_attn_unavailable_on_rocm(monkeypatch): + from freetoken.engine import engine + + monkeypatch.setattr(engine, "is_rocm", lambda: True) + assert engine._sgl_flash_attn_available() is False diff --git a/tests/kernels/test_backend_rocm.py b/tests/kernels/test_backend_rocm.py new file mode 100644 index 000000000..4c7326351 --- /dev/null +++ b/tests/kernels/test_backend_rocm.py @@ -0,0 +1,63 @@ +"""ROCm-aware optional-package probing in ``kernel/backend.py``. + +Loads ``backend.py`` with a fake ``freetoken.utils.arch`` injected so the ROCm gating +(treat NVIDIA-only native packages as unavailable) is unit-testable without a working +torch or those packages installed. +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +_BACKEND_PATH = ( + Path(__file__).resolve().parents[2] / "python" / "freetoken" / "kernel" / "backend.py" +) + + +def _load_backend(rocm: bool): + # Fake freetoken.utils.arch (backend.py imports is_rocm from it at module load). + arch_mod = types.ModuleType("freetoken.utils.arch") + arch_mod.is_rocm = lambda: rocm + arch_mod.device_kind = lambda: "rocm" if rocm else "cpu" + freetoken = types.ModuleType("freetoken") + freetoken.__path__ = [] + utils = types.ModuleType("freetoken.utils") + utils.__path__ = [] + utils.arch = arch_mod + sys.modules["freetoken"] = freetoken + sys.modules["freetoken.utils"] = utils + sys.modules["freetoken.utils.arch"] = arch_mod + + spec = importlib.util.spec_from_file_location("freetoken.kernel.backend", _BACKEND_PATH) + kernel = types.ModuleType("freetoken.kernel") + kernel.__path__ = [str(_BACKEND_PATH.parent)] + sys.modules["freetoken.kernel"] = kernel + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(autouse=True) +def _clean_sys_modules(): + yield + for name in ("freetoken", "freetoken.utils", "freetoken.utils.arch", "freetoken.kernel"): + sys.modules.pop(name, None) + + +def test_rocm_native_packages_unavailable(): + backend = _load_backend(rocm=True) + assert backend.is_flashinfer_installed() is False + assert backend.is_sgl_kernel_installed() is False + assert backend.is_triton_kernels_installed() is False + assert backend.is_native_cuda_available() is False + + +def test_cuda_native_packages_follow_importability(monkeypatch): + backend = _load_backend(rocm=False) + # Without the packages installed, probes are False; is_native_cuda_available() needs a + # real CUDA-capable torch (returns False here since torch.cuda is unavailable). + assert backend.is_flashinfer_installed() is False + assert backend.is_sgl_kernel_installed() is False diff --git a/tests/kernels/test_cache_rocm_pairing.py b/tests/kernels/test_cache_rocm_pairing.py new file mode 100644 index 000000000..24796632e --- /dev/null +++ b/tests/kernels/test_cache_rocm_pairing.py @@ -0,0 +1,109 @@ +"""ROCm cache/runtime version pairing and gfx-arch gating (torch-free). + +Loads the relevant modules by file path with mocked torch so the logic is unit-testable +without a working torch/GPU install. +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +_ROOT = Path(__file__).resolve().parents[2] +_UTILS_PATH = _ROOT / "python" / "freetoken" / "kernel" / "utils.py" +_ARCH_PATH = _ROOT / "python" / "freetoken" / "utils" / "arch.py" + + +@pytest.fixture(scope="module") +def ku(): + spec = importlib.util.spec_from_file_location("kernel_utils_mod", _UTILS_PATH) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + return m + + +@pytest.fixture(scope="module") +def arch(): + spec = importlib.util.spec_from_file_location("arch_mod", _ARCH_PATH) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + return m + + +def test_cache_version_pairs_rocm(ku): + # Same release + same g-sha + both rocm -> OK. + assert ku._kernel_cache_version_ok( + "0.1.1+rocm.g3f01615", "0.1.1+rocm.g3f01615" + ) + # No g-stamp on either side -> release-only compare -> OK. + assert ku._kernel_cache_version_ok("0.1.1+rocm", "0.1.1+rocm") + + +def test_cache_version_rejects_backend_tag_mismatch(ku): + # CUDA runtime must never pair with a ROCm cache (SASS family differs). + assert not ku._kernel_cache_version_ok( + "0.1.1+rocm.g3f01615", "0.1.1+cu130.g3f01615" + ) + assert not ku._kernel_cache_version_ok( + "0.1.1+cu130.g3f01615", "0.1.1+rocm.g3f01615" + ) + + +def test_cache_version_rejects_different_build(ku): + assert not ku._kernel_cache_version_ok( + "0.1.1+rocm.g3f01615", "0.1.1+rocm.gdeadbee" + ) + assert not ku._kernel_cache_version_ok("0.1.2+rocm", "0.1.1+rocm") + + +def _torch_importable() -> bool: + """True only when torch actually *imports* (a findable-but-broken torch -- e.g. a + CUDA build missing its runtime libs -- must count as absent so the torch-free path + is exercised on such boxes).""" + try: + import torch # noqa: PLC0415 + + return True + except Exception: + return False + + +def test_gfx_arch_ge_degrades_without_torch(arch): + if _torch_importable(): + pytest.skip("torch is importable; no-torch degradation tested on a torch-free box") + assert arch.is_gfx_arch_ge(1100) is False + + +def test_gfx_arch_ge_parses_gfx_string(arch, monkeypatch): + if _torch_importable(): + # Real torch present: is_gfx_arch_ge must reflect the actual device (gfx1100 on + # an RX 7900 XTX). No fake torch injection to avoid cross-test state pollution. + import torch + + assert arch.is_gfx_arch_ge(1100) is True # RX 7000 target + return + arch._get_gfx_arch.cache_clear() + # Fake a ROCm torch whose device name carries the gfx arch string. + tv = types.ModuleType("torch.version") + tv.hip = "6.2.4100000" + tv.cuda = None + t = types.ModuleType("torch") + t.version = tv + _props = types.SimpleNamespace(gcnArchName="gfx1100") + t.cuda = types.SimpleNamespace( + is_available=lambda: True, + current_device=lambda: 0, + get_device_name=lambda dev: "AMD Radeon RX 7900 XTX", + get_device_properties=lambda dev: _props, + ) + monkeypatch.setitem(sys.modules, "torch", t) + monkeypatch.setitem(sys.modules, "torch.version", tv) + assert arch.is_gfx_arch_ge(1100) is True + assert arch.is_gfx_arch_ge(1103) is False + # CUDA builds must return False. + arch._get_gfx_arch.cache_clear() + tv.hip = None + tv.cuda = "13.0" + assert arch.is_gfx_arch_ge(1100) is False diff --git a/tests/kernels/test_toolchain_hip.py b/tests/kernels/test_toolchain_hip.py new file mode 100644 index 000000000..fa8f53cfe --- /dev/null +++ b/tests/kernels/test_toolchain_hip.py @@ -0,0 +1,84 @@ +"""HIP (ROCm) toolchain helper tests. + +These load ``kernel/_toolchain.py`` by file path (as setup.py and the kernel-cache +build backend do) so they need no torch import and no ROCm toolkit to exercise the +pure parsing/detection logic. +""" + +import importlib.util +import os +from pathlib import Path + +import pytest + +_TOOLCHAIN_PATH = ( + Path(__file__).resolve().parents[2] / "python" / "freetoken" / "kernel" / "_toolchain.py" +) + + +def _load_toolchain(): + spec = importlib.util.spec_from_file_location("_freetoken_toolchain", _TOOLCHAIN_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def tc(): + return _load_toolchain() + + +def _write_fake_hipcc(tmp_path, version: str): + bin_dir = tmp_path / "bin" + bin_dir.mkdir(exist_ok=True) + hipcc = bin_dir / "hipcc" + hipcc.write_text( + "#!/bin/sh\n" + f'echo "HIP version: {version}"\n', + encoding="utf-8", + ) + hipcc.chmod(0o755) + return str(hipcc) + + +def test_hip_hip_version(tc, tmp_path): + hipcc = _write_fake_hipcc(tmp_path, "6.2.41000") + assert tc.hip_hip_version(hipcc) == (6, 2) + + +def test_hip_hip_version_missing(tc): + assert tc.hip_hip_version("/nonexistent/hipcc") is None + + +def test_hipcc_path_rocm_home(tc, tmp_path, monkeypatch): + hipcc = _write_fake_hipcc(tmp_path, "6.2.41000") + monkeypatch.setenv("ROCM_HOME", str(tmp_path)) + monkeypatch.delenv("HIP_PATH", raising=False) + assert tc._hipcc_path() == hipcc + + +def test_hipcc_path_default_opt_rocm(tc, monkeypatch): + monkeypatch.delenv("ROCM_HOME", raising=False) + monkeypatch.delenv("HIP_PATH", raising=False) + # If a real /opt/rocm/bin/hipcc exists it wins; otherwise we expect None (no PATH hit + # guaranteed in the test sandbox, so only assert it returns None or a path string). + result = tc._hipcc_path() + if os.path.exists("/opt/rocm/bin/hipcc"): + assert result is not None + else: + assert result is None or result.endswith("hipcc") + + +def test_torch_hip_major_rocm(tc, monkeypatch): + # Simulate ROCm torch via a fake torch.version with hip set. + import sys + import types + + fake_torch = types.ModuleType("torch") + fake_version = types.ModuleType("torch.version") + fake_version.hip = "6.2.4100000" + fake_version.cuda = None + fake_torch.version = fake_version + monkeypatch.setitem(sys.modules, "torch", fake_torch) + assert tc.is_rocm_torch() is True + assert tc.torch_hip_major() == 6 diff --git a/tests/models/test_qwen35moe_gguf_deint.py b/tests/models/test_qwen35moe_gguf_deint.py new file mode 100644 index 000000000..4154fe27e --- /dev/null +++ b/tests/models/test_qwen35moe_gguf_deint.py @@ -0,0 +1,71 @@ +"""qwen35moe GGUF GDN value-head de-interleaving. + +llama.cpp stores the GDN *value* projections with the ``mrope_interleaved`` head order +(even heads 0..nv/2-1 first, then odd heads nv/2..nv-1). FreeToken uses the HF contiguous +head order, so the loader de-interleaves the value projections on load. These tests pin the +permutation and the packed/dense de-interleave helpers (no model weights required). +""" + +import torch + +from freetoken.models.qwen3_5_moe.gguf import ( + _gdn_head_perm, + _deint_dense_rows, + _deint_q8_cols, + _deint_q8_rows, +) + + +def _interleave(x: torch.Tensor, nv: int, rows_per_head: int = 1) -> torch.Tensor: + """Build the GGUF head-interleaved layout from an HF-contiguous tensor. + + GGUF position ``perm[h]`` holds HF head ``h`` (so interleaved[perm[h]] = x[h]). + """ + m = x.reshape(nv, rows_per_head, -1) + perm = _gdn_head_perm(nv) + out = m.clone() + for h in range(nv): + out[perm[h]] = m[h] + return out.reshape(x.shape) + + +def test_head_permutation_is_a_bijection(): + nv = 32 + perm = _gdn_head_perm(nv) + assert sorted(perm) == list(range(nv)) # valid permutation + # GGUF layout: even HF heads occupy positions 0..15, odd heads 16..31. + even = sorted(perm[h] for h in range(0, nv, 2)) + odd = sorted(perm[h] for h in range(1, nv, 2)) + assert even == list(range(nv // 2)) + assert odd == list(range(nv // 2, nv)) + + +def test_deint_dense_rows_recovers_contiguous(): + nv, hd = 32, 128 + x = torch.randn(nv, hd) + inter = _interleave(x, nv, rows_per_head=1) + rec = _deint_dense_rows(inter, nv) + assert torch.allclose(rec, x, atol=1e-6) + + +def test_deint_q8_rows_recovers_contiguous(): + nv, hd = 32, 128 + rph = 64 # rows per value head in the projection output dim + x = torch.randn(nv * rph, 4) # packed rows x (row_bytes mocked) + inter = _interleave(x, nv, rows_per_head=rph) + rec = _deint_q8_rows(inter, nv, rph) + assert torch.allclose(rec, x, atol=1e-6) + + +def test_deint_q8_cols_recovers_contiguous(): + nv = 32 + rows, blocks_per_head, bb = 16, 4, 34 + x = torch.randn(rows, nv * blocks_per_head * bb) + # interleave column head-groups: inter[:, perm[h]*bbp:(perm[h]+1)*bbp] = x[:, h*bbp:(h+1)*bbp] + bbp = blocks_per_head * bb + inter = torch.empty_like(x) + perm = _gdn_head_perm(nv) + for h in range(nv): + inter[:, perm[h] * bbp:(perm[h] + 1) * bbp] = x[:, h * bbp:(h + 1) * bbp] + rec = _deint_q8_cols(inter, nv, blocks_per_head) + assert torch.allclose(rec, x, atol=1e-6) diff --git a/tests/moe/test_nvfp4_backends_rocm.py b/tests/moe/test_nvfp4_backends_rocm.py new file mode 100644 index 000000000..0061d3788 --- /dev/null +++ b/tests/moe/test_nvfp4_backends_rocm.py @@ -0,0 +1,95 @@ +"""ROCm-aware NVFP4 backend selection (torch-free). + +Loads ``nvfp4_backends.py`` with a fake ``freetoken.utils.arch`` so the AMD branch of +``select_nvfp4_backend`` (reject NVIDIA-only marlin/flashinfer, force triton) is +unit-testable without torch/vLLM/flashinfer. +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +_BACKEND_PATH = ( + Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "moe" / "nvfp4_backends.py" +) + + +def _torch_importable() -> bool: + """True only when torch actually *imports* (a findable-but-broken torch -- e.g. a + CUDA build missing its runtime libs -- must count as absent so the torch-free path + is exercised on such boxes).""" + try: + import torch # noqa: PLC0415 + + return True + except Exception: + return False + + +def _load_nvfp4(rocm: bool): + arch = types.ModuleType("freetoken.utils.arch") + arch.is_rocm = lambda: rocm + freetoken = types.ModuleType("freetoken") + freetoken.__path__ = [] + utils = types.ModuleType("freetoken.utils") + utils.__path__ = [] + utils.arch = arch + utils.init_logger = lambda name: types.SimpleNamespace( + info=lambda *a, **k: None, + warning=lambda *a, **k: None, + debug=lambda *a, **k: None, + ) + sys.modules["freetoken"] = freetoken + sys.modules["freetoken.utils"] = utils + sys.modules["freetoken.utils.arch"] = arch + + # nvfp4_backends.py does `import torch` at module scope. On a torch-free box stub it + # so the module imports; on a box with real torch use the real one (never mutate + # sys.modules["torch"], which would corrupt the session's torch/triton state). + if not _torch_importable(): + torch_stub = types.ModuleType("torch") + torch_stub.no_grad = lambda: (lambda f: f) + torch_stub.Tensor = object + sys.modules["torch"] = torch_stub + + spec = importlib.util.spec_from_file_location("nvfp4_backends_mod", _BACKEND_PATH) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + return m + + +@pytest.fixture(autouse=True) +def _clean(): + yield + for name in ("freetoken", "freetoken.utils", "freetoken.utils.arch"): + sys.modules.pop(name, None) + + +def test_rocm_auto_resolves_triton(): + m = _load_nvfp4(rocm=True) + dev = types.SimpleNamespace(type="cuda") + assert m.select_nvfp4_backend(dev, 768, "auto") == "triton" + assert m.select_nvfp4_backend(dev, 768, "triton") == "triton" + + +def test_rocm_rejects_nvidia_backends(): + m = _load_nvfp4(rocm=True) + dev = types.SimpleNamespace(type="cuda") + with pytest.raises(RuntimeError, match="NVIDIA-only"): + m.select_nvfp4_backend(dev, 768, "marlin") + with pytest.raises(RuntimeError, match="NVIDIA-only"): + m.select_nvfp4_backend(dev, 768, "flashinfer") + + +def test_cuda_path_not_short_circuited_by_rocm_guard(): + # On CUDA the ROCm early-return must not be taken. We only exercise the guard + # boundary (is_rocm()==False) without calling torch.cuda (unavailable here); the + # full CUDA auto logic runs on the target box. + m = _load_nvfp4(rocm=False) + assert m.select_nvfp4_backend( + types.SimpleNamespace(type="cpu"), 768, "auto" + ) == "triton" diff --git a/tests/moe/test_nvfp4_to_mxfp4.py b/tests/moe/test_nvfp4_to_mxfp4.py new file mode 100644 index 000000000..194201186 --- /dev/null +++ b/tests/moe/test_nvfp4_to_mxfp4.py @@ -0,0 +1,111 @@ +"""NVFP4 -> MXFP4 converter tests (numpy core, torch-free). + +Validates the error-prone numeric bits: the e2m1 code table, e8m0 scale encoding, and +an end-to-end round trip (dequant NVFP4 -> requant MXFP4 -> dequant back) that must be +close to the reference under a bounded relative error. +""" + +import importlib.util +import math +from pathlib import Path + +import numpy as np +import pytest + +_MODULE_PATH = ( + Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "moe" / "nvfp4_to_mxfp4.py" +) + + +@pytest.fixture(scope="module") +def c(): + spec = importlib.util.spec_from_file_location("nvfp4_to_mxfp4", _MODULE_PATH) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + return m + + +def test_fp4_table_values(c): + table = c.fp4_e2m1_table() + assert len(table) == 16 + # e2m1 canonical values. + assert table[0] == 0.0 + assert table[1] == 0.5 + assert table[4] == 2.0 + assert table[7] == 6.0 + assert table[15] == -6.0 + assert table[9] == -0.5 + + +def test_e8m0_scale_exact_power_of_two(c): + # Block max-abs 3.0 -> scale 2^ceil(log2(3/6)) = 2^-1 = 0.5 -> code 126 (126-127=-1). + values = np.array([[3.0, 1.0, -2.0, 0.5] + [0.0] * 28], dtype=np.float32) + scale_codes, codes = c.e8m0_scale_and_codes(values, block=32) + assert scale_codes.shape == (1,) + assert int(scale_codes[0]) == 126 + # After dividing by 0.5: [6, 2, -4, 1] -> nearest e2m1 codes 7, 4, 14, 2. + assert int(codes[0, 0]) == 7 + assert int(codes[0, 1]) == 4 + assert int(codes[0, 2]) == 14 + assert int(codes[0, 3]) == 2 + + +def test_dequantize_nvfp4_round_trip(c): + rng = np.random.default_rng(0) + N, K = 4, 64 + # Build a plausible NVFP4 weight: choose fp4 codes, block scales, row globals. + codes = rng.integers(0, 16, size=(N, K)).astype(np.uint8) + packed = (codes[:, 0::2] | (codes[:, 1::2] << 4)).astype(np.uint8) + block_scale = (rng.uniform(0.1, 2.0, size=(N, K // 16))).astype(np.float32) + global_scale = (rng.uniform(0.5, 2.0, size=(N,)).astype(np.float32)) + ref = c.dequantize_nvfp4_block(packed, block_scale, global_scale) + assert ref.shape == (N, K) + assert np.isfinite(ref).all() + + +def test_round_trip_matches_reference(c): + rng = np.random.default_rng(1) + N, K = 2, 128 + codes = rng.integers(0, 16, size=(N, K)).astype(np.uint8) + packed = (codes[:, 0::2] | (codes[:, 1::2] << 4)).astype(np.uint8) + block_scale = rng.uniform(0.5, 2.0, size=(N, K // 16)).astype(np.float32) + global_scale = rng.uniform(0.5, 2.0, size=(N,)).astype(np.float32) + ref = c.dequantize_nvfp4_block(packed, block_scale, global_scale) + + mxfp4_packed, mxfp4_scales = c.convert_nvfp4_to_mxfp4( + packed, block_scale, global_scale + ) + assert mxfp4_packed.shape == (N, K // 2) + assert mxfp4_scales.shape == (N, K // 32) + assert mxfp4_scales.dtype == np.uint8 + + # Re-dequantize the MXFP4 output and compare to the NVFP4 reference. + out_codes = np.stack( + [mxfp4_packed & 0x0F, mxfp4_packed >> 4], axis=-1 + ).reshape(N, K) + vals = np.asarray(c.fp4_e2m1_table(), dtype=np.float32)[out_codes.astype(np.int64)] + # e8m0 scale 2**(v-127) + scale_v = (2.0 ** (mxfp4_scales.astype(np.float32) - 127.0))[:, :, None] + mxfp4_vals = (vals.reshape(N, K // 32, 32) * scale_v).reshape(N, K) + # MXFP4 requantizes at a coarser 32-block granularity, so the round trip carries + # fp4 quantization error (inherent to e2m1). Bound it loosely: median < 0.6 and the + # largest magnitude in each block (which defines the scale) must be well-preserved. + rel = np.abs(mxfp4_vals - ref) / (np.abs(ref) + 1e-6) + # Exclude near-zero refs where relative error is meaningless (a code flips +0<->small). + nz = np.abs(ref) > 1e-3 + assert float(np.median(rel[nz])) < 0.6 + assert float(np.percentile(rel[nz], 95)) < 2.0 + # The max-abs element per 32-block reproduces within one e2m1 step of its scale. + block_max_ref = np.abs(ref.reshape(N, K // 32, 32)).max(axis=-1) + block_max_ours = np.abs(mxfp4_vals.reshape(N, K // 32, 32)).max(axis=-1) + scale_ratio = np.abs(block_max_ours / (block_max_ref + 1e-6) - 1.0) + assert float(np.median(scale_ratio)) < 0.5 + + +def test_inner_axis_unsupported(c): + packed = np.zeros((64, 32), dtype=np.uint8) + scale = np.zeros((64, 4), dtype=np.float32) + g = np.ones((64,), dtype=np.float32) + with pytest.raises(NotImplementedError): + c.convert_nvfp4_to_mxfp4(packed.T, scale, g, axis=0) diff --git a/tests/utils/test_device_kind.py b/tests/utils/test_device_kind.py new file mode 100644 index 000000000..0496c4d13 --- /dev/null +++ b/tests/utils/test_device_kind.py @@ -0,0 +1,83 @@ +"""Device-kind (CUDA vs ROCm vs CPU) detection and graceful-degradation gating. + +These run without a working torch install: we load ``utils/arch.py`` by file path and +inject a fake ``torch``/``torch.version`` into ``sys.modules`` so the build-detection +(``torch.version.hip`` vs ``torch.version.cuda``) is unit-testable anywhere. +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +_ARCH_PATH = ( + Path(__file__).resolve().parents[2] / "python" / "freetoken" / "utils" / "arch.py" +) + + +@pytest.fixture(scope="module") +def arch(): + spec = importlib.util.spec_from_file_location("arch_mod", _ARCH_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _set_torch(arch, monkeypatch, hip, cuda): + tv = types.ModuleType("torch.version") + tv.hip = hip + tv.cuda = cuda + t = types.ModuleType("torch") + t.version = tv + t.cuda = types.SimpleNamespace(is_available=lambda: False) + monkeypatch.setitem(sys.modules, "torch", t) + monkeypatch.setitem(sys.modules, "torch.version", tv) + + +def test_device_kind_rocm(arch, monkeypatch): + _set_torch(arch, monkeypatch, hip="6.2.4100000", cuda=None) + assert arch.device_kind() == "rocm" + assert arch.is_rocm() is True + assert arch.is_cuda() is False + + +def test_device_kind_cuda(arch, monkeypatch): + _set_torch(arch, monkeypatch, hip=None, cuda="13.0") + assert arch.device_kind() == "cuda" + assert arch.is_rocm() is False + assert arch.is_cuda() is True + + +def test_device_kind_cpu(arch, monkeypatch): + _set_torch(arch, monkeypatch, hip=None, cuda=None) + assert arch.device_kind() == "cpu" + + +def _torch_importable() -> bool: + """True only when torch actually *imports* (a findable-but-broken torch -- e.g. a + CUDA build missing its runtime libs -- must count as absent so the torch-free path + is exercised on such boxes).""" + try: + import torch # noqa: PLC0415 + + return True + except Exception: + return False + + +def test_arch_gates_degrade_without_torch(arch, monkeypatch): + # Only meaningful where torch is genuinely absent (the CUDA-locked dev box). On a + # box with a working torch, Python re-imports the real torch after the delitem, so + # the assertion target changes -- skip there. + if _torch_importable(): + pytest.skip("torch is importable; no-torch degradation tested on a torch-free box") + # With torch absent, every is_sm* gate is False and device_kind() is "cpu". + monkeypatch.delitem(sys.modules, "torch", raising=False) + monkeypatch.delitem(sys.modules, "torch.version", raising=False) + assert arch.is_sm90_supported() is False + assert arch.is_sm100_supported() is False + assert arch.is_sm90_family() is False + assert arch.is_sm100_family() is False + assert arch.device_kind() == "cpu" From 6f4974eed707ca47a7f017dce7ad4fb507f5a4c9 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Tue, 25 Aug 2026 22:34:03 -0300 Subject: [PATCH 02/17] fix(gguf): register unmerged CONTROL/USER_DEFINED tokens for atomic encoding The GGUF->fast-tokenizer converter relies on BPE merges to keep special strings whole: -style tokens happen to be merge entries and survive, but think tags / tool-call tags are NOT in the merge table, so they silently split into plain pieces (''). The model then receives garbage ids it never saw in training and answers with gibberish + EOS -- observed as 'reasons a little, returns empty content' on Qwen3.6 GGUF checkpoints. _register_control_tokens re-registers every ggml token_type 2/3/4 string that does not already encode atomically, against its EXISTING vocab id: vocab size and id assignments never change, already-atomic tokens are untouched (no-op on healthy setups), and safetensors/HF checkpoints never touch this path. Hardware-independent; safe for CUDA by construction. Also: chat-template resolution now mirrors official HF configs -- chat_template.jinja sidecar next to the .gguf wins, then FT_CHAT_TEMPLATE_REPO (hf_hub_download), then the embedded metadata. CPU-only regression tests included (synthetic BPE tokenizer, no GPU or download needed). --- python/freetoken/models/gguf/tokenizer.py | 77 +++++++++++++++++++- tests/models/test_gguf_tokenizer_specials.py | 77 ++++++++++++++++++++ 2 files changed, 153 insertions(+), 1 deletion(-) create mode 100644 tests/models/test_gguf_tokenizer_specials.py diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index c05824190..83d8be840 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -10,8 +10,12 @@ from typing import Any +from freetoken.utils import init_logger + from .reader import gguf_architecture, load_gguf_metadata +logger = init_logger(__name__) + # GGUF architecture -> transformers GGUF tokenizer-converter key. _TOKENIZER_ARCH = { "gemma4": "gemma4_text", @@ -19,6 +23,73 @@ } +# ggml token_type enum values that mark non-mergeable control strings. +# 1=NORMAL 2=UNKNOWN 3=CONTROL 4=USER_DEFINED 5=UNUSED 6=BYTE +_GGML_SPECIAL_TYPES = (2, 3, 4) + + +def _register_control_tokens(tokenizer, tokens: list[str], types: list[int]) -> None: + """Force every CONTROL/USER_DEFINED/UNKNOWN vocab entry to tokenize atomically. + + The GGUF->fast-tokenizer converter relies on BPE merges for special strings: + ``<|im_start|>`` happens to be a merge-table entry and survives, but ```` + is NOT (it never appears in merged training text), so it silently splits into + plain pieces (''). A model fed those garbage ids answers with + gibberish and EOS -- observed as "thought then returned empty" on Qwen3.6 GGUFs. + Registering each string as an added token makes the AddedVocabulary extract it + before BPE; because the string already exists in the base vocab, the existing id + is reused and the vocab never grows (transformers' own ``add_tokens`` wrapper + no-ops for in-vocab strings, so this goes through the backend directly). + """ + missing = [ + name + for i, (name, ty) in enumerate(zip(tokens, types)) + if int(ty) in _GGML_SPECIAL_TYPES + # Skip BYTE-fallback and unused; skip anything already atomic at its id. + and tokenizer.encode(name, add_special_tokens=False) != [i] + ] + if missing: + tokenizer.backend_tokenizer.add_tokens(missing) + logger.info( + "registered %d unmerged control tokens for atomic encoding (e.g. %s)", + len(missing), + ", ".join(repr(t) for t in missing[:6]), + ) + + +def _resolve_chat_template(meta: dict[str, Any], model_path: str) -> str | None: + """Chat template for GGUF checkpoints: explicit mirrors win, then metadata. + + Priority: a ``chat_template.jinja`` dropped NEXT TO the .gguf file, then + ``FT_CHAT_TEMPLATE_REPO`` ( on the HF Hub, fetched via huggingface_hub), + then the template embedded in the GGUF's ``tokenizer.chat_template`` metadata. + GGUF packagers (llama.cpp/unsloth) sometimes ship modified variants of the + official template — placing the official file beside the checkpoint overrides + it without repacking. The embedded one is the last resort, never wrong-by-default. + """ + import os + + sidecar = os.path.join(os.path.dirname(model_path), "chat_template.jinja") + if os.path.isfile(sidecar): + logger.info("using chat template sidecar %s", sidecar) + with open(sidecar, encoding="utf-8") as fh: + return fh.read() + repo = os.environ.get("FT_CHAT_TEMPLATE_REPO") + if repo: + try: + from huggingface_hub import hf_hub_download + + path = hf_hub_download(repo_id=repo, filename="chat_template.jinja") + with open(path, encoding="utf-8") as fh: + return fh.read() + except Exception as exc: # noqa: BLE001 — offline/bad repo is not fatal + logger.warning("FT_CHAT_TEMPLATE_REPO=%s fetch failed: %s", repo, exc) + embedded = meta.get("tokenizer.chat_template") + if isinstance(embedded, str) and embedded.strip(): + return embedded + return None + + def load_gguf_tokenizer(model_path: str): from transformers import PreTrainedTokenizerFast from transformers.integrations.ggml import convert_gguf_tokenizer @@ -49,7 +120,11 @@ def tok_for(id_key: str, default: str) -> str: unk_token=tok_for("unknown_token_id", ""), pad_token=tok_for("padding_token_id", ""), ) - chat_template = meta.get("tokenizer.chat_template") + # GGUFs are not required to carry per-token types; absent means all-normal. + types = meta.get("tokenizer.ggml.token_type") or [] + if types: + _register_control_tokens(tokenizer, tokens, types) + chat_template = _resolve_chat_template(meta, str(model_path)) if chat_template: tokenizer.chat_template = chat_template return tokenizer diff --git a/tests/models/test_gguf_tokenizer_specials.py b/tests/models/test_gguf_tokenizer_specials.py new file mode 100644 index 000000000..deadb0a4e --- /dev/null +++ b/tests/models/test_gguf_tokenizer_specials.py @@ -0,0 +1,77 @@ +"""CPU-only regression tests for GGUF control-token registration. + +The GGUF->fast-tokenizer converter leaves CONTROL/USER_DEFINED vocab entries +(e.g. Qwen's think tags, tool-call tags) unregistered: unless they happen to be +reachable through BPE merges they silently split into plain pieces, feeding the +model garbage ids ("thought then empty" symptom). ``_register_control_tokens`` +must re-register every such string against its EXISTING id -- vocab size and id +assignments may not change, and already-atomic tokens must be left alone. + +Runs on any machine: no GPU, no model download -- builds synthetic BPE +tokenizers whose vocab mirrors the broken shape (control string in vocab, +unreachable because no merge path leads to it). +""" + +from __future__ import annotations + +from tokenizers import Tokenizer, models, pre_tokenizers + +from freetoken.models.gguf.tokenizer import _register_control_tokens +from transformers import PreTrainedTokenizerFast + + +def _tok(vocab: dict[str, int], merges: list[tuple[str, str]] = []): + backend = Tokenizer(models.BPE(vocab=vocab, merges=list(merges))) + return PreTrainedTokenizerFast(tokenizer_object=backend) + + +def test_unmergeable_control_token_registered_at_existing_id(): + # Mirrors the real Qwen GGUF: '' is IN the vocab (id 11) but no merge + # path produces it, so bare conversion emits per-character pieces. + vocab = { + "<": 1, "t": 2, "h": 3, "i": 4, "n": 5, "k": 6, ">": 7, + "a": 8, "b": 9, "": 11, + } + tok = _tok(vocab) + before = tok.encode("", add_special_tokens=False) + assert before != [11], "sanity: must start out broken (split), like the real bug" + + names_by_id = sorted(vocab, key=vocab.get) + types = [1] * len(names_by_id) + types[names_by_id.index("")] = 4 # USER_DEFINED + _register_control_tokens(tok, names_by_id, types) + + assert tok.encode("", add_special_tokens=False) == [11] + assert tok.vocab_size == len(vocab), "vocab must not grow" + assert tok.convert_tokens_to_ids("") == 11, "id must be preserved" + + +def test_merge_reachable_control_token_left_alone(): + # A CONTROL entry that is ALREADY atomic (reachable via merges) must not be + # re-registered: the filter is encode(name) != [own id]. + vocab = {"a": 1, "b": 2, "ab": 3} + tok = _tok(vocab, merges=[("a", "b")]) + assert tok.encode("ab", add_special_tokens=False) == [3] + + names_by_id = ["a", "b", "ab"] + types = [1, 1, 3] + _register_control_tokens(tok, names_by_id, types) + assert tok.encode("ab", add_special_tokens=False) == [3] + assert tok.vocab_size == 3 + + +def test_normal_tokens_never_registered(): + # NORMAL entries that split must stay split: only CONTROL/UNKNOWN/USER_DEFINED + # types are eligible (BYTE/UNUSED excluded too). + vocab = {"a": 1, "b": 2, "c": 3} + tok = _tok(vocab) + names_by_id = ["a", "b", "c"] + _register_control_tokens(tok, names_by_id, [1, 5, 6]) + assert tok.encode("ab", add_special_tokens=False) == [1, 2] + + +def test_empty_types_is_noop(): + vocab = {"a": 1} + tok = _tok(vocab) + _register_control_tokens(tok, ["a"], []) + assert tok.vocab_size == 1 From bc6232453ca1b2714122c15004a95cdbd31e14c6 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Tue, 25 Aug 2026 22:34:18 -0300 Subject: [PATCH 03/17] fix(server): route Qwen3.6 to the qwen3_coder tool-call parser _infer_tool_call_parser only special-cased qwen3_5/coder names, so a Qwen3.6-* filename fell through to qwen25 (legacy JSON grammar). The 3.5/ 3.6 hybrid family instructs the XML invoke-block format in its chat template ( + blocks -- the Qwen3-Coder grammar), so the model emitted well-formed calls the wrong detector could not see: output swallowed, empty stop. Add qwen3_6/qwen3.6 markers. --- python/freetoken/server/args.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 2a533be70..303eec5dd 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -141,6 +141,12 @@ def _infer_tool_call_parser(model_path: str) -> str: if ( "qwen3_5" in marker or "qwen3.5" in marker + or "qwen3_6" in marker + or "qwen3.6" in marker + # Qwen3-Coder and the 3.5/3.6 hybrid family share the XML invoke-block + # grammar (v); plain "qwen" (2.x) uses the + # older JSON form. A bare "qwen3" marker stays JSON (qwen25) unless it's + # a coder variant. or ("qwen3" in marker and "coder" in marker) ): return "qwen3_coder" From 25d7bd8d999044dd398b83fa26d5502b2496723c Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Tue, 25 Aug 2026 22:34:18 -0300 Subject: [PATCH 04/17] chore(local): ROCm serve scripts + VS Code tasks/debug configs - serve-qwen-moe.sh: single-line array-based launch (backslash-newline continuations get mangled by VS Code shells), setsid detach so a cancelled task cannot kill the server mid-load, CRLF guard, status/log subcommands, 128k KV default (~2.5 GiB on the hybrid arch). - mirror-hf-configs.sh: mirror official HF chat_template/generation_config next to a local GGUF without touching the file. - kill-freetoken.sh: hard-stop helper. - .vscode/: tasks (serve start/stop/status/log, fast tests), debugpy launch config for the server, LF pinning for shell scripts. - .gitignore: nohup.out, .plans/. --- .gitignore | 4 + .vscode/launch.json | 28 +++++++ .vscode/settings.json | 19 +++++ .vscode/tasks.json | 63 ++++++++++++++ scripts/kill-freetoken.sh | 53 ++++++++++++ scripts/mirror-hf-configs.sh | 67 +++++++++++++++ scripts/serve-qwen-moe.sh | 154 +++++++++++++++++++++++++++++++++++ 7 files changed, 388 insertions(+) create mode 100644 .vscode/launch.json create mode 100644 .vscode/settings.json create mode 100644 .vscode/tasks.json create mode 100755 scripts/kill-freetoken.sh create mode 100755 scripts/mirror-hf-configs.sh create mode 100755 scripts/serve-qwen-moe.sh diff --git a/.gitignore b/.gitignore index 95757625b..edc08af25 100644 --- a/.gitignore +++ b/.gitignore @@ -233,3 +233,7 @@ benchmarks/cross_framework python/freetoken/kernel/csrc/gguf/*.hip python/freetoken/kernel/csrc/gguf/*_hip.cuh python/freetoken/kernel/csrc/gguf/ggml-common_hip.h + +# local session junk +nohup.out +.plans/ diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 000000000..64cbc7904 --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,28 @@ +{ + // F5 launches the server IN THE FOREGROUND under the ROCm venv (no nohup), so + // stdout/stderr land straight in the Debug Console and breakpoints work. + "version": "0.2.0", + "configurations": [ + { + "name": "FreeToken: serve qwen-moe (debug)", + "type": "debugpy", + "request": "launch", + "module": "freetoken.cli", + "python": "${workspaceFolder}/.venv-rocm/bin/python", + "cwd": "${workspaceFolder}", + "env": { "PYTHONPATH": "${workspaceFolder}/python" }, + "console": "integratedTerminal", + "args": [ + "serve", + "--model", "/media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf", + "--moe-backend", "offload", + "--attention-backend", "triton", + "--moe-cache-size", "2048", + "--num-tokens", "131072", + "--host", "127.0.0.1", + "--port", "1920" + ], + "justMyCode": false + } + ] +} \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 000000000..7d6299c2a --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,19 @@ +{ + // Prevent VS Code from re-introducing CRLF endings into the serve script — + // CRLF was what broke the backslash-newline continuations before. + "files.eol": "\n", + "files.associations": { + "*.sh": "shellscript" + }, + "[shellscript]": { + "files.eol": "\n", + "editor.tabSize": 4, + "editor.insertSpaces": false + }, + "terminal.integrated.env.linux": { + "PYTHONPATH": "${workspaceFolder}/python" + }, + "python.defaultInterpreterPath": "${workspaceFolder}/.venv-rocm/bin/python", + "python.testing.pytestEnabled": true, + "python.testing.pytestArgs": ["tests"] +} \ No newline at end of file diff --git a/.vscode/tasks.json b/.vscode/tasks.json new file mode 100644 index 000000000..175d2c8f6 --- /dev/null +++ b/.vscode/tasks.json @@ -0,0 +1,63 @@ +{ + // FreeToken server lifecycle tasks. All of them call scripts/serve-qwen-moe.sh, + // which builds its argv as a bash array and launches on one line — safe for the + // VS Code integrated shell (no backslash-newline continuations to mangle). + "version": "2.0.0", + "tasks": [ + { + "label": "FreeToken: serve (start)", + "type": "shell", + "command": "${workspaceFolder}/scripts/serve-qwen-moe.sh", + "args": ["start"], + "options": { + "cwd": "${workspaceFolder}", + "env": {} + }, + "isBackground": true, + "problemMatcher": [], + "presentation": { + "reveal": "always", + "panel": "dedicated", + "clear": true + }, + "detail": "Launch the Qwen MoE server (~3-4 min model load), then wait for readiness." + }, + { + "label": "FreeToken: serve (stop)", + "type": "shell", + "command": "${workspaceFolder}/scripts/serve-qwen-moe.sh", + "args": ["stop"], + "options": { "cwd": "${workspaceFolder}" }, + "problemMatcher": [], + "presentation": { "reveal": "always", "panel": "dedicated" } + }, + { + "label": "FreeToken: serve (status)", + "type": "shell", + "command": "${workspaceFolder}/scripts/serve-qwen-moe.sh", + "args": ["status"], + "options": { "cwd": "${workspaceFolder}" }, + "problemMatcher": [], + "presentation": { "reveal": "always", "panel": "shared" } + }, + { + "label": "FreeToken: serve (follow log)", + "type": "shell", + "command": "${workspaceFolder}/scripts/serve-qwen-moe.sh", + "args": ["log"], + "options": { "cwd": "${workspaceFolder}" }, + "isBackground": true, + "problemMatcher": [], + "presentation": { "reveal": "always", "panel": "dedicated" } + }, + { + "label": "FreeToken: tests (fast)", + "type": "shell", + "command": "${workspaceFolder}/.venv-rocm/bin/python", + "args": ["-m", "pytest", "-q", "-m", "not slow"], + "options": { "cwd": "${workspaceFolder}", "env": { "PYTHONPATH": "${workspaceFolder}/python" } }, + "group": "test", + "problemMatcher": [] + } + ] +} \ No newline at end of file diff --git a/scripts/kill-freetoken.sh b/scripts/kill-freetoken.sh new file mode 100755 index 000000000..18893d54c --- /dev/null +++ b/scripts/kill-freetoken.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +# kill-freetoken.sh +# +# Kill every process related to FreeToken: the API server, backend workers, and the +# multiprocessing spawn/resource-tracker children, then free the listener ports. +# +# Usage: +# ./scripts/kill-freetoken.sh [PORTS...] # default: 1920 1921 + +set -u + +PORTS=("${@:-1920 1921}") +SELF=$$ +SELFCMDLINE="$(ps -p "$SELF" -o args= 2>/dev/null || true)" + +# Kill a pattern, excluding this script's own shell. +kill_matching() { + local pat="$1" + for pid in $(pgrep -f "$pat" 2>/dev/null); do + [ "$pid" = "$SELF" ] && continue + cmdline="$(ps -p "$pid" -o args= 2>/dev/null || true)" + [ -n "$cmdline" ] && [ "$cmdline" = "$SELFCMDLINE" ] && continue + kill -9 "$pid" 2>/dev/null || true + done +} + +# 1) FreeToken CLI / server / backend supervisor / backend workers. +kill_matching "freetoke[n].cli serve" +kill_matching "freetoke[n].cli" +kill_matching "freetoken" +kill_matching "multiprocessing.spawn" +kill_matching "multiprocessing.resource_tracker" +kill_matching "multiprocessing.semaphore" +kill_matching "torch.distributed.launch" + +# 2) Free the listener ports (a worker may still hold one). +for p in "${PORTS[@]}"; do + pid="$(ss -ltnp 2>/dev/null | grep ":$p" | grep -oP 'pid=\K[0-9]+' | head -1)" + [ -n "$pid" ] && kill -9 "$pid" 2>/dev/null || true +done + +sleep 1 + +left="$(pgrep -af "freetoken|multiprocessing.spawn|multiprocessing.resource" 2>/dev/null | grep -v "kill-freetoken.sh" | grep -v "$$" || true)" +if [ -n "$left" ]; then + echo "WARNING: still running (forced KILL):" + echo "$left" + kill_matching "freetoken" + kill_matching "multiprocessing.spawn" + sleep 1 +fi + +echo "FreeToken processes killed; ports ${PORTS[*]} freed." diff --git a/scripts/mirror-hf-configs.sh b/scripts/mirror-hf-configs.sh new file mode 100755 index 000000000..8cca566b9 --- /dev/null +++ b/scripts/mirror-hf-configs.sh @@ -0,0 +1,67 @@ +#!/usr/bin/env bash +# mirror-hf-configs.sh +# +# Mirror a model's official Hugging Face config files next to a local GGUF so +# FreeToken uses them instead of whatever the GGUF packager embedded: +# +# ./mirror-hf-configs.sh [gguf-file] +# +# ./mirror-hf-configs.sh Qwen/Qwen3.6-35B-A3B +# ./mirror-hf-configs.sh Qwen/Qwen3.6-35B-A3B /media/smk/Shared/Models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf +# +# With no gguf argument, defaults to $FT_MODEL, else the single *.gguf in +# /media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models or /media/smk/Shared/Models +# matching the repo's basename, else errors. +# +# Files fetched (when they exist upstream): +# chat_template.jinja — read by FreeToken's GGUF loader (overrides embedded) +# generation_config.json — default sampling params +# tokenizer_config.json — reference only (FreeToken builds its tokenizer from GGUF) +# +# To revert, delete the mirrored files — the GGUF metadata is never modified. + +set -euo pipefail + +REPO_ID="${1:?usage: mirror-hf-configs.sh [gguf-file]}" +shift || true + +if [ "$#" -ge 1 ]; then + GGUF="$1" +else + GGUF="${FT_MODEL:-}" + if [ -z "$GGUF" ]; then + base="$(basename "${REPO_ID##*:}")" + cand="$(find /media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models /media/smk/Shared/Models -maxdepth 1 -iname "*${base%%-*}*.gguf" 2>/dev/null | head -1)" + [ -n "$cand" ] || { echo "ERROR: no gguf found; pass one explicitly" >&2; exit 1; } + GGUF="$cand" + fi +fi +[ -f "$GGUF" ] || { echo "ERROR: not a file: $GGUF" >&2; exit 1; } + +DIR="$(dirname "$GGUF")" +echo "repo : $REPO_ID" +echo "into : $DIR" + +PY="${PY:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)/.venv-rocm/bin/python}" + +for f in chat_template.jinja generation_config.json tokenizer_config.json; do + if "$PY" - "$REPO_ID" "$f" "$DIR" <<'EOF' +import sys +from huggingface_hub import hf_hub_download +repo, fname, dest_dir = sys.argv[1:4] +try: + p = hf_hub_download(repo_id=repo, filename=fname) +except Exception: + sys.exit(1) +import shutil +shutil.copyfile(p, f"{dest_dir}/{fname}") +print(f"{fname}: OK") +EOF + then + : + else + echo "$f: not present upstream, skipped" + fi +done + +echo "done — restart the server to pick them up." diff --git a/scripts/serve-qwen-moe.sh b/scripts/serve-qwen-moe.sh new file mode 100755 index 000000000..6a28444d8 --- /dev/null +++ b/scripts/serve-qwen-moe.sh @@ -0,0 +1,154 @@ +#!/usr/bin/env bash +# serve-qwen-moe.sh +# +# Serve the Qwen3.5/3.6-35B-A3B hybrid MoE GGUF (GatedDeltaNet + full-attention) +# on AMD ROCm (gfx1100 / RX 7900 XTX). Uses the offload MoE backend (experts on the +# CPU/offload cache) and the triton attention backend by default. +# +# Native context window: 262144 (256K) tokens (max_position_embeddings in the model). +# Graph capture is settled as a failure on ROCm, so this runs eager kernel-launch decode. +# +# VS CODE NOTES: +# - The server command is built as a bash array and launched on ONE physical +# line. Backslash-newline continuations get mangled by some VS Code shells / +# task runners (each continuation line then executes as its own command), +# which is exactly how `nohup.out` ended up with bare "--model: command not +# found" errors. Do NOT reintroduce multi-line command strings here. +# - .vscode/settings.json pins files.eol=\n for *.sh; the CRLF guard below +# catches any violation early instead of failing obscurely mid-launch. +# +# Usage: +# ./serve-qwen-moe.sh # launch on 127.0.0.1:1920, triton attention +# FT_ATTN=torch ./serve-qwen-moe.sh # A/B against the pure-torch reference backend +# FT_PORT=1930 ./serve-qwen-moe.sh # pick another port +# ./serve-qwen-moe.sh stop # kill the running server +# ./serve-qwen-moe.sh status # running? + tail of the log +# +# Or from the VS Code Command Palette: "Tasks: Run Task" -> FreeToken: ... + +set -euo pipefail + +REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +PY="${PY:-$REPO/.venv-rocm/bin/python}" + +MODEL="${FT_MODEL:-/media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf}" +HOST="${FT_HOST:-127.0.0.1}" +PORT="${FT_PORT:-1920}" +ATNN="${FT_ATTN:-triton}" # triton | torch +MOE_BACKEND="${FT_MOE:-offload}" # offload required for K-quant experts +# VS Code / Copilot injects a large chat context, so the KV cache must be bigger than +# the tiny 8K the MoE-auto cache leaves. --num-tokens sizes the KV cache in tokens; +# --moe-cache-size limits the GPU expert cache so KV has room (fewer slots = slower decode). +# Hybrid arch: only 10/40 layers are full attention (2 kv heads x 256 dim) -> +# ~20 KiB/token bf16, i.e. 128k ~= 2.7 GiB, full native 256k ~= 5.4 GiB. +KV_TOKENS="${FT_KV_TOKENS:-131072}" +MOE_CACHE="${FT_MOE_CACHE:-2048}" +LOG="${FT_LOG:-/tmp/serve_qwen_moe.log}" + +die() { echo "ERROR: $*" >&2; exit 1; } + +# Guard against the exact failure mode VS Code caused before: if this file ever +# gets saved with CRLF endings, every argument silently grows a trailing \r. +if grep -q $'\r' "${BASH_SOURCE[0]}"; then + die "CRLF line endings detected in $(basename "${BASH_SOURCE[0]}"). Run: sed -i 's/\r$//' ${BASH_SOURCE[0]}" +fi + +[ -x "$PY" ] || die "python not found: $PY (set PY=/path/to/venv-python)" + +server_pids() { + pgrep -f "freetoke[n].cli serve" || true +} + +start() { + if [ -n "$(server_pids)" ]; then + echo "A server is already running (pid $(server_pids | tr '\n' ' ')). Use 'stop' first." + exit 1 + fi + [ -f "$MODEL" ] || die "model not found: $MODEL" + + # One arg per array element; expanded once, single line, no continuations. + local -a SERVE_ARGS=( + "--model" "$MODEL" + "--moe-backend" "$MOE_BACKEND" + "--attention-backend" "$ATNN" + "--moe-cache-size" "$MOE_CACHE" + "--num-tokens" "$KV_TOKENS" + "--host" "$HOST" + "--port" "$PORT" + ) + + echo "Launching FreeToken server" + echo " model : $MODEL" + echo " listen : $HOST:$PORT" + echo " attn : $ATNN" + echo " moe : $MOE_BACKEND" + echo " kv : $KV_TOKENS tokens (gpu moe cache: $MOE_CACHE slots)" + echo " python : $PY" + echo " log : $LOG" + + # setsid: give the server its OWN session/process group. nohup alone only + # ignores SIGHUP — a caller that dies (e.g. a VS Code task cancelled, an agent + # tool timeout) still takes down the whole process group with SIGKILL/SIGTERM, + # which silently killed the server mid model-load once already. + cd "$REPO" + PYTHONPATH="$REPO/python" setsid nohup "$PY" -m freetoken.cli serve "${SERVE_ARGS[@]}" >"$LOG" 2>&1 /dev/null || true + echo "pid=$pid — waiting for readiness (model load takes ~3-4 min)..." +} + +wait_ready() { + for _ in $(seq 1 90); do + if grep -q "API server is ready" "$LOG" 2>/dev/null; then + echo "READY on $HOST:$PORT" + return 0 + fi + if [ -z "$(server_pids)" ]; then + echo "server exited; last log lines:" >&2 + tail -20 "$LOG" >&2 || true + return 1 + fi + sleep 5 + done + echo "timed out waiting for readiness; see $LOG" >&2 + return 1 +} + +stop() { + local pids + pids="$(server_pids)" + if [ -z "$pids" ]; then + echo "no server running" + else + pkill -9 -f "freetoke[n].cli serve" 2>/dev/null || true + echo "stopped (was pid $pids)" + fi + pkill -9 -f "multiprocessing.spawn" 2>/dev/null || true + # free the distributed worker port (server_port+1) + local p pid + for p in "$PORT" "$((PORT + 1))"; do + pid="$(ss -ltnp 2>/dev/null | grep ":$p" | grep -oP 'pid=\K[0-9]+' | head -1 || true)" + if [ -n "$pid" ]; then + kill -9 "$pid" 2>/dev/null || true + fi + done +} + +status() { + local pids + pids="$(server_pids)" + if [ -n "$pids" ]; then + echo "RUNNING (pid $(echo "$pids" | tr '\n' ' ')) on $HOST:$PORT" + else + echo "NOT RUNNING" + fi + [ -f "$LOG" ] && echo "--- last 5 log lines ($LOG) ---" && tail -5 "$LOG" +} + +case "${1:-start}" in + start) start && wait_ready ;; + stop) stop ;; + status) status ;; + log) exec tail -f "$LOG" ;; + *) echo "usage: $0 [start|stop|status|log]" >&2; exit 1 ;; +esac From 729789ae0be705badd7ec3d8b0563a1a4f05c5a3 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Wed, 26 Aug 2026 12:37:47 -0300 Subject: [PATCH 05/17] fix(qwen3.5-moe): GGUF lm_head last-token gather + presence/frequency penalties Root cause of the intermittent empty/loop responses and the non-deterministic 'greedy' decode: the GGUF lm_head (GGUFLinear) returned FULL prefill logits [total, vocab] without gathering each request's final row, while the engine contract is [batch.size, vocab] (the token to sample after each request's prompt). ParallelLMHead and Nvfp4LMHead already gather via attn_metadata.get_last_indices(); the GGUF path did not. Consequences fixed: - The first generated token was sampled from POSITION 0's logits (after '<|im_start|>' -> 'user') instead of the last prompt position (after '...assistant\n thinking\n' -> 'Here'). - Worse, a fresh prefill (142-token batch, row 0 = prompt token 0) and a radix-cache continuation (14-token batch, row 0 = prompt token 128 -> '<|im_end|>' -> immediate stop, EMPTY content) sampled different rows, so the SAME greedy request gave different outputs depending on cache state -- the 'model thinks then returns empty' report. Fix: qwen3_5_moe/model.py gathers last_indices on prefill before the GGUF lm head (matches ParallelLMHead/Nvfp4LMHead; shared code path for both backends). Also in this change: - Presence/frequency penalties were accepted by the API but IGNORED by the sampler. Implemented end-to-end (core.py SamplingParams + Req.prompt_len, generation.py/openai_api.py pass-through, engine/sample.py apply_penalties over generated tokens only, engine.py passes the batch to sample()). Applies to greedy too. Breaks reasoning loops by penalizing repeated tokens. - Serve script: --max-output-tokens 65536 default (FT_MAX_OUTPUT knob) so the reasoning model has room to finish; keeps the Inc-1 diagnostic knobs (moe stats, prefill overlap, cpu layers). Verified: - 28k-token prompt now ANSWERS (finish=stop, content='Four'), no empty, no loop. - Tool calls still return tool_calls with correct args. - TRUE greedy (top_k=1, top_p=1.0) deterministic within a cache state; the first-after-startup fresh run differs from radix hits only at ~1 bf16 ULP (continuation GEMM batch shapes) -- documented residual, both outputs valid. - tests/engine + tests/server: 618 passed, same 15 pre-existing failures. --- python/freetoken/core.py | 15 ++++ python/freetoken/engine/engine.py | 2 +- python/freetoken/engine/sample.py | 50 +++++++++++-- python/freetoken/models/qwen3_5_moe/model.py | 15 +++- python/freetoken/scheduler/scheduler.py | 13 ++++ python/freetoken/scheduler/status.py | 22 ++++++ python/freetoken/server/args.py | 11 +++ python/freetoken/server/generation.py | 4 ++ python/freetoken/server/openai_api.py | 4 ++ scripts/serve-qwen-moe.sh | 74 +++++++++++++++++--- 10 files changed, 196 insertions(+), 14 deletions(-) diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ef0a539cb..82c6e0872 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -25,6 +25,13 @@ class SamplingParams: # Stop strings (OpenAI `stop` / Anthropic `stop_sequences`). Generation finishes when one # appears in the decoded output; the matched substring (and anything after) is trimmed. stop_strs: list[str] = field(default_factory=list) + # OpenAI-style presence/frequency penalties, applied over the tokens this request has + # generated so far (the prompt is excluded): + # logits[t] -= presence_penalty * (t was generated) + frequency_penalty * count(t) + # Positive values push the model away from repeating itself (breaks reasoning loops); + # negative values encourage repetition (coherence). 0.0 = disabled. + presence_penalty: float = 0.0 + frequency_penalty: float = 0.0 @property def is_greedy(self) -> bool: @@ -40,6 +47,9 @@ class Req: uid: int sampling_params: SamplingParams cache_handle: BaseCacheHandle + # Prompt length at creation (see __post_init__); input_ids[prompt_len:] is the + # generated portion used by presence/frequency penalties. + prompt_len: int = 0 # Optional precomputed multimodal soft-token embeddings (GPU, [num_image_tokens, # hidden]) scattered at image-token positions during this request's prefill. mm_embeds: torch.Tensor | None = None @@ -71,6 +81,11 @@ def __post_init__(self) -> None: self.max_device_len = len(self.input_ids) + self.output_len assert 0 <= self.cached_len < self.device_len <= self.max_device_len self._alloc_ids_buf() + # Length of the prompt this request was created with. Generation grows input_ids + # past this point (append_host), so input_ids[prompt_len:] is exactly the tokens + # the model has generated so far -- what presence/frequency penalties are applied + # over. ChunkedReq instances (which never sample) record their partial length. + self.prompt_len = self.device_len def _alloc_ids_buf(self) -> None: self._ids_buf = torch.empty(self.max_device_len, dtype=self.input_ids.dtype) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 5ada3954a..d86bcf09e 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -945,7 +945,7 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: req.complete_one() batch_logits = logits[: batch.size] - next_tokens_gpu = self.sampler.sample(batch_logits, args).to(torch.int32) + next_tokens_gpu = self.sampler.sample(batch_logits, args, batch).to(torch.int32) next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True) copy_done_event = torch.cuda.Event() copy_done_event.record(self.stream) diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 01d14b1aa..be3726b4e 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -7,7 +7,7 @@ from freetoken.utils import is_sm90_supported, nvtx_annotate if TYPE_CHECKING: - from freetoken.core import Batch + from freetoken.core import Batch, Req @dataclass @@ -15,6 +15,36 @@ class BatchSamplingArgs: temperatures: torch.Tensor | None top_k: torch.Tensor | None = None top_p: torch.Tensor | None = None + # True when at least one request carries a presence/frequency penalty; the sampler + # then lowers each request's logits over its generated tokens before sampling. + apply_penalties: bool = False + + +def apply_penalties( + logits: torch.Tensor, + reqs: List["Req"], +) -> None: + """Apply OpenAI presence/frequency penalties to ``logits`` in place (row per req). + + For a token ``t`` the request already generated, its score is lowered by + ``presence_penalty`` plus ``frequency_penalty * count(t)``. The prompt is excluded + (only ``input_ids[req.prompt_len:]`` counts), so the penalty grows with the + generation itself -- positive values push the model away from repeating itself, + which breaks reasoning loops; negative values nudge it toward repetition. + """ + for i, req in enumerate(reqs): + sp = req.sampling_params + pp, fp = sp.presence_penalty, sp.frequency_penalty + if not pp and not fp: + continue + gen = req.input_ids[req.prompt_len :] + if gen.numel() == 0: + continue + uniq, counts = torch.unique(gen, return_counts=True) + vals = torch.full_like(counts, pp, dtype=torch.float32) + fp * counts.to( + torch.float32 + ) + logits[i, uniq.to(logits.device)] -= vals.to(logits.device) def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor: @@ -57,7 +87,10 @@ class Sampler: def prepare(self, batch: Batch) -> BatchSamplingArgs: params = [r.sampling_params for r in batch.reqs] - if all(p.is_greedy for p in params): + apply_penalties = any( + p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 for p in params + ) + if all(p.is_greedy for p in params) and not apply_penalties: return BatchSamplingArgs(temperatures=None) MIN_P = MIN_T = 1e-6 @@ -70,11 +103,20 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: top_k = make_device_tensor(top_ks, torch.int32, self.device) if any(p < 1.0 for p in top_ps): top_p = make_device_tensor(top_ps, torch.float32, self.device) - return BatchSamplingArgs(temperatures, top_k=top_k, top_p=top_p) + return BatchSamplingArgs( + temperatures, + top_k=top_k, + top_p=top_p, + apply_penalties=apply_penalties, + ) @nvtx_annotate("Sampler") - def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: + def sample( + self, logits: torch.Tensor, args: BatchSamplingArgs, batch: Batch + ) -> torch.Tensor: with torch.cuda.nvtx.range("Sampler"): + if args.apply_penalties: + apply_penalties(logits, batch.reqs) if args.temperatures is None: # greedy sampling return torch.argmax(logits, dim=-1) return sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py index 32954dc52..df0d41f0b 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -117,7 +117,20 @@ def __init__(self, config: ModelConfig): convert_qwen35moe_to_gguf(self, config) def forward(self) -> torch.Tensor: - output = self.model.forward(get_global_ctx().batch.input_ids) + ctx = get_global_ctx() + batch = ctx.batch + output = self.model.forward(batch.input_ids) + if batch.is_prefill: + # GGUFLinear (unlike ParallelLMHead / Nvfp4LMHead) does not gather each + # request's final row itself, but the engine contract is [batch.size, vocab] + # (the logits to sample after each request's prompt). Without this gather + # the sampler reads the FIRST prompt position's logits: the first generated + # token comes from position 0, and -- worse -- differs between a fresh + # prefill and a radix-cache continuation of the same prompt (the restored + # continuation has a different first row), which made greedy output flip + # between server states / cache states. + indices = batch.attn_metadata.get_last_indices(batch.size) + output = output[indices] return self.lm_head.forward(output) diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 355411617..0846c940a 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -411,9 +411,22 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: page_size=self.config.page_size, mamba_slots=mamba_slots, swa_tokens=swa_tokens, + moe_stats=self._moe_stats_snapshot(), ) self.send_result(reply) + def _moe_stats_snapshot(self) -> dict | None: + """Per-window MoE cache hit/miss stats for the decode log line, or None when + stats collection is off (the default). Reads device counters once per call; + the status reporter only calls this every decode_log_interval steps.""" + cache = getattr(self.engine, "moe_offload_cache", None) + if cache is None or not getattr(cache, "collect_stats", False): + return None + try: + return cache.decode_miss_stats() + except Exception: + return None + def _match_stop_str(self, req: Req) -> str | None: """First stop string present in this request's generated tail, else None. Decodes only a short suffix (bounded by the longest stop string's char length, so a stop of diff --git a/python/freetoken/scheduler/status.py b/python/freetoken/scheduler/status.py index ca706c506..c7be5e8f1 100644 --- a/python/freetoken/scheduler/status.py +++ b/python/freetoken/scheduler/status.py @@ -34,6 +34,7 @@ def report_batch( page_size: int, mamba_slots: tuple[int, int] | None = None, swa_tokens: tuple[int, int] | None = None, + moe_stats: dict | None = None, ) -> None: if batch.is_prefill: self._report_prefill( @@ -55,6 +56,7 @@ def report_batch( page_size=page_size, mamba_slots=mamba_slots, swa_tokens=swa_tokens, + moe_stats=moe_stats, ) def _report_prefill( @@ -101,6 +103,7 @@ def _report_decode( page_size: int, mamba_slots: tuple[int, int] | None = None, swa_tokens: tuple[int, int] | None = None, + moe_stats: dict | None = None, ) -> None: self._decode_forward_count += 1 self._decode_generated_tokens += len(batch.reqs) @@ -121,6 +124,7 @@ def _report_decode( f"{_mamba_msg(mamba_slots)}" f"gen throughput (token/s): {gen_throughput:.2f}, " f"#queue-req: {queue_reqs}" + f"{_moe_msg(moe_stats)}" ) @@ -128,6 +132,24 @@ def _usage_ratio(used: int, total: int) -> float: return used / total if total > 0 else 0.0 +def _moe_msg(stats: dict | None) -> str: + """MoE cache hit/miss summary for the decode log line (empty when stats are off).""" + if not stats: + return "" + miss = stats.get("miss_rate") + fetch = stats.get("fetch_rate") + cpu = stats.get("cpu_per_layer") + hit = (1.0 - miss) if miss is not None else None + parts = [f"moe hit: {hit:.3f}" if hit is not None else "moe hit: n/a"] + if miss is not None: + parts.append(f"miss: {miss:.3f}") + if fetch is not None: + parts.append(f"fetch: {fetch:.3f}") + if cpu is not None: + parts.append(f"cpu: {cpu:.3f}") + return ", " + ", ".join(parts) + + def _mamba_msg(mamba_slots: tuple[int, int] | None) -> str: """GDN-state (mamba) pool occupancy for hybrid models; empty for the rest.""" if mamba_slots is None: diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 303eec5dd..cb651221b 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -585,6 +585,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--moe-collect-stats", + action="store_true", + default=ServerArgs.moe_collect_stats, + help=( + "Capture MoE cache hit/miss counters into the decode graph and print them " + "in the decode log line (moe hit/miss/fetch/cpu). Off by default; the " + "device-side accumulation is captured into the CUDA graph." + ), + ) + parser.add_argument( "--moe-prefill-hit-d2d", action="store_true", diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index be05d908a..6e2b9ab40 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -163,6 +163,8 @@ def resolve_sampling( ignore_eos: bool, model_sampling: dict[str, Any], stop: str | list[str] | None = None, + presence_penalty: float | None = None, + frequency_penalty: float | None = None, ) -> SamplingParams: """Map a protocol's sampling fields onto the engine's neutral SamplingParams, filling unspecified fields from the checkpoint's recommended defaults.""" @@ -184,6 +186,8 @@ def pick(value, key, framework): top_k=pick(top_k, "top_k", -1), top_p=pick(top_p, "top_p", 1.0), stop_strs=[s for s in stop_list if s], # drop empty strings (would match everything) + presence_penalty=presence_penalty if presence_penalty is not None else 0.0, + frequency_penalty=frequency_penalty if frequency_penalty is not None else 0.0, ) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b4becd263..04822dc7a 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -76,6 +76,8 @@ def chat_request_to_genspec( ignore_eos=req.ignore_eos, model_sampling=model_sampling, stop=req.stop, + presence_penalty=req.presence_penalty, + frequency_penalty=req.frequency_penalty, ), chat_template_kwargs=ctk, template_tools=_tools_for_template(req), @@ -526,6 +528,8 @@ def _resolve_sampling( ignore_eos=req.ignore_eos, model_sampling=model_sampling, stop=req.stop, + presence_penalty=getattr(req, "presence_penalty", None), + frequency_penalty=getattr(req, "frequency_penalty", None), ) diff --git a/scripts/serve-qwen-moe.sh b/scripts/serve-qwen-moe.sh index 6a28444d8..47de497c1 100755 --- a/scripts/serve-qwen-moe.sh +++ b/scripts/serve-qwen-moe.sh @@ -36,17 +36,51 @@ HOST="${FT_HOST:-127.0.0.1}" PORT="${FT_PORT:-1920}" ATNN="${FT_ATTN:-triton}" # triton | torch MOE_BACKEND="${FT_MOE:-offload}" # offload required for K-quant experts -# VS Code / Copilot injects a large chat context, so the KV cache must be bigger than -# the tiny 8K the MoE-auto cache leaves. --num-tokens sizes the KV cache in tokens; -# --moe-cache-size limits the GPU expert cache so KV has room (fewer slots = slower decode). -# Hybrid arch: only 10/40 layers are full attention (2 kv heads x 256 dim) -> -# ~20 KiB/token bf16, i.e. 128k ~= 2.7 GiB, full native 256k ~= 5.4 GiB. +# radix = cross-request GDN-state prefix reuse (default); naive = no prefix reuse. +# Suspect for long-context degradation: radix GDN-state reuse corrupting later requests. +CACHE_TYPE="${FT_CACHE_TYPE:-radix}" +# MoE cache hit/miss stats in the decode log line (--moe-collect-stats). +MOE_STATS="${FT_MOE_STATS:-1}" +# Disable the two-buffer prefill MoE overlap (diagnostic: race check). +PREFILL_OVERLAP="${FT_PREFILL_OVERLAP:-1}" +# MoE layers computed on the CPU executor (diagnostic: '0' = all-GPU, no CPU path). +CPU_LAYERS="${FT_CPU_LAYERS:-}" +# --num-tokens sizes the KV cache in tokens (hybrid arch: only 10/40 layers are full +# attention, 2 kv heads x 256 dim -> ~20 KiB/token bf16; 128k ~= 2.7 GiB). +# The GPU expert-slot cache is sized by FT_MOE_CACHE: +# auto -> --moe-cache-auto: the engine derives slot bytes from the real expert +# tensors and fills all free VRAM AFTER reserving kv-reserve-tokens for KV. +# --kv-reserve-tokens MUST equal --num-tokens here: with an explicit +# --num-tokens the engine skips auto's own KV-half plan (num_page_override +# is set), so without a matching reservation greedy expert fill would eat +# VRAM the pinned KV still needs -> late CUDA OOM. +# -> fixed --moe-cache-size N slots (legacy behavior). +# Headroom: the engine may use FT_MEMORY_RATIO of free VRAM for weights+KV+experts +# combined (default here 0.80, upstream default 0.9). The remainder absorbs prefill +# transients -- the MoE overlap double-buffer alone needs ~3.8 GiB at this model's +# batch shapes; 0.9 left only ~2 GiB and OOM'd mid-prefill under VS Code payloads. +MEMORY_RATIO="${FT_MEMORY_RATIO:-0.80}" +# --max-prefill-length caps chunked-prefill chunk size (engine default 8192). The lm_head +# materializes logits for EVERY chunk token: an 8192-token chunk spikes ~3.8 GiB +# transiently -- enough to OOM on VS Code-sized prompts even with healthy headroom. +# 4096 halves the spike; long prompts just prefill in more chunks. +PREFILL_CHUNK="${FT_PREFILL_CHUNK:-4096}" KV_TOKENS="${FT_KV_TOKENS:-131072}" -MOE_CACHE="${FT_MOE_CACHE:-2048}" +# Default max output tokens for requests that omit max_tokens. The reasoning model +# sometimes needs more room to finish its reasoning before answering; the engine +# default is 32k. Copilot sends its own max_tokens, which we cannot override, but +# other clients inherit this. +MAX_OUTPUT="${FT_MAX_OUTPUT:-65536}" +MOE_CACHE="${FT_MOE_CACHE:-auto}" LOG="${FT_LOG:-/tmp/serve_qwen_moe.log}" die() { echo "ERROR: $*" >&2; exit 1; } +# Map FT_MOE_CACHE to argv. Only the auto path carries --kv-reserve-tokens: +# with a fixed size the CLI already validates fit against the pinned KV. +# NOTE: validation must run in THIS shell (no process substitution), or die() +# would only exit a subshell and the launch would continue unvalidated. + # Guard against the exact failure mode VS Code caused before: if this file ever # gets saved with CRLF endings, every argument silently grows a trailing \r. if grep -q $'\r' "${BASH_SOURCE[0]}"; then @@ -70,19 +104,36 @@ start() { local -a SERVE_ARGS=( "--model" "$MODEL" "--moe-backend" "$MOE_BACKEND" + "--cache-type" "$CACHE_TYPE" + $( [ "$MOE_STATS" = 1 ] && echo "--moe-collect-stats" ) + $( [ "$PREFILL_OVERLAP" = 0 ] && echo "--disable-moe-prefill-overlap" ) + $( [ -n "$CPU_LAYERS" ] && echo "--moe-cpu-layers" "$CPU_LAYERS" ) "--attention-backend" "$ATNN" - "--moe-cache-size" "$MOE_CACHE" "--num-tokens" "$KV_TOKENS" + "--memory-ratio" "$MEMORY_RATIO" + "--max-prefill-length" "$PREFILL_CHUNK" + "--max-output-tokens" "$MAX_OUTPUT" "--host" "$HOST" "--port" "$PORT" ) + case "$MOE_CACHE" in + auto) + SERVE_ARGS+=("--moe-cache-auto" "--kv-reserve-tokens" "$KV_TOKENS") + ;; + ''|*[!0-9]*) + die "FT_MOE_CACHE='$MOE_CACHE' is invalid: use 'auto' or a slot count" + ;; + *) + SERVE_ARGS+=("--moe-cache-size" "$MOE_CACHE") + ;; + esac echo "Launching FreeToken server" echo " model : $MODEL" echo " listen : $HOST:$PORT" echo " attn : $ATNN" echo " moe : $MOE_BACKEND" - echo " kv : $KV_TOKENS tokens (gpu moe cache: $MOE_CACHE slots)" + echo " kv : $KV_TOKENS tokens (gpu moe cache: ${MOE_CACHE}${MOE_CACHE:+ }$( [ "$MOE_CACHE" = auto ] && echo "kv-reserve $KV_TOKENS" || echo slots))" echo " python : $PY" echo " log : $LOG" @@ -142,6 +193,13 @@ status() { else echo "NOT RUNNING" fi + # Surface the auto-resolved expert-cache split so users don't have to read + # engine log lines; only present when FT_MOE_CACHE=auto booted the server. + if [ -f "$LOG" ]; then + local resolved + resolved="$(grep -o -- '--moe-cache-auto resolved moe_cache_size=[0-9]* num_pages=[0-9]*' "$LOG" | tail -1)" + [ -n "$resolved" ] && echo "moe cache: ${resolved//--moe-cache-auto resolved /}" + fi [ -f "$LOG" ] && echo "--- last 5 log lines ($LOG) ---" && tail -5 "$LOG" } From 1f76e88e5be3b063de93c282d053d35a32296a24 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Wed, 26 Aug 2026 13:47:34 -0300 Subject: [PATCH 06/17] chore(rocm): drop CI workflow + tinygrad fallback; review cleanup - Remove .github/workflows/rocm.yml (no GH Actions for ROCm) - Remove kernel/tinygrad_fallback.py (dead: nothing imports it) - Remove .vscode/launch.json (hardcoded machine model path) - kernel/utils.py: drop shadowed duplicate _build_stamps definition - kernel/backend.py: drop unused _CUDA_ONLY_PACKAGES - models/qwen3_5_moe/gguf.py: drop unused _q8_0_down_row_bytes - moe/nvfp4_to_mxfp4.py: drop unused _nearest_e2m1_codes/_FP4_SORT_SIGN; assert K is a block multiple instead of silently truncating - moe/expert_banks.py: remove dead dummy-path import in _gguf_banks - models/qwen3_5_moe/gdn.py: rename _fp8 -> _split_proj (also covers GGUF) - engine/engine.py: drop redundant local is_rocm imports; build infos once - utils/graph_gate.py, engine/graph.py, attention/torch.py, tests: strip plan-increment (Inc N) references - docs/install-amd.md, pyproject.toml: drop plan refs; align ROCm torch index - kernel/gguf.py: honor FREETOKEN_KERNEL_CACHE_GFX for the JIT offload-arch - scripts: remove hardcoded /media/smk model paths (require FT_MODEL/arg) - utils.cuh: fix stray [[unlikely]]; statement --- .github/workflows/rocm.yml | 70 -------------- .vscode/launch.json | 28 ------ docs/install-amd.md | 6 +- pyproject.toml | 6 +- python/freetoken/attention/torch.py | 5 +- python/freetoken/engine/engine.py | 7 +- python/freetoken/engine/graph.py | 2 +- python/freetoken/kernel/backend.py | 6 -- .../kernel/csrc/include/freetoken/utils.cuh | 6 +- python/freetoken/kernel/gguf.py | 5 +- python/freetoken/kernel/tinygrad_fallback.py | 92 ------------------- python/freetoken/kernel/utils.py | 6 -- python/freetoken/models/qwen3_5_moe/gdn.py | 7 +- python/freetoken/models/qwen3_5_moe/gguf.py | 4 - python/freetoken/moe/expert_banks.py | 7 +- python/freetoken/moe/nvfp4_to_mxfp4.py | 22 +---- python/freetoken/utils/graph_gate.py | 15 ++- scripts/mirror-hf-configs.sh | 15 +-- scripts/serve-qwen-moe.sh | 3 +- tests/attention/test_torch_backend.py | 2 +- 20 files changed, 42 insertions(+), 272 deletions(-) delete mode 100644 .github/workflows/rocm.yml delete mode 100644 .vscode/launch.json delete mode 100644 python/freetoken/kernel/tinygrad_fallback.py diff --git a/.github/workflows/rocm.yml b/.github/workflows/rocm.yml deleted file mode 100644 index e2f1fe49a..000000000 --- a/.github/workflows/rocm.yml +++ /dev/null @@ -1,70 +0,0 @@ -name: ROCm (AMD) correctness smoke - -# ROCm CI job: compiles the AOT kernel cache for the RX 7000 (gfx1100) target on a -# ROCm torch install and runs a torch-free correctness smoke plus the AMD unit tests. -# Gated on a self-hosted runner that has ROCm torch + hipcc. The primary NVIDIA release -# flow is release.yml; this job is additive and must not gate NVIDIA releases. - -on: - workflow_dispatch: - push: - branches: [main] - pull_request: - -jobs: - rocm-smoke: - runs-on: [self-hosted, linux, amd, rocm] - timeout-minutes: 60 - env: - FREETOKEN_DISABLE_JIT: "1" - FREETOKEN_KERNEL_CACHE_GFX: "gfx1100" - steps: - - uses: actions/checkout@v4 - - - name: Check ROCm toolchain - run: | - set -e - command -v hipcc || ls /opt/rocm/bin/hipcc - "${PYTHON:-python3}" -c "import torch.version as v; print('torch hip:', v.hip)" - - - name: Install build deps - run: | - python -m pip install --upgrade pip wheel setuptools - python -m pip install -e "python[rocm]" - - - name: Compile AOT kernel cache for gfx1100 - run: | - FREETOKEN_KERNEL_CACHE_VERBOSE=1 python -m pip wheel ./freetoken-kernel-cache -w dist/cache-rocm - - - name: Install prebuilt kernel cache - run: | - whl="$(find dist/cache-rocm -name 'freetoken_kernel_cache-*.whl' | head -1)" - python -m pip install --force-reinstall "$whl" - - - name: Torch-free AMD unit tests - run: | - python -m pytest \ - tests/utils/test_device_kind.py \ - tests/kernels/test_toolchain_hip.py \ - tests/kernels/test_backend_rocm.py \ - tests/kernels/test_cache_rocm_pairing.py \ - tests/moe/test_nvfp4_to_mxfp4.py \ - -q - - - name: Hardware correctness smoke (serves on RX 7000) - run: | - # Functional path only -- flashinfer/sgl/trtllm are NVIDIA-only and must not - # be selected. AUTO backend must resolve to triton; NVFP4 auto -> triton. - python - <<'PY' - from freetoken.utils.arch import is_rocm, is_gfx_arch_ge - from freetoken.moe.nvfp4_backends import select_nvfp4_backend - import torch - assert is_rocm(), "expected a ROCm torch build" - assert is_gfx_arch_ge(1100), "expected gfx1100-class device (RX 7000)" - print("NVFP4 auto ->", select_nvfp4_backend(torch.device("cuda"), 768, "auto")) - PY - - - name: Serve smoke - run: | - FREETOKEN_DEVICE=cuda python -m freetoken.serve --help >/dev/null \ - && echo "freetoken CLI loads on ROCm" diff --git a/.vscode/launch.json b/.vscode/launch.json deleted file mode 100644 index 64cbc7904..000000000 --- a/.vscode/launch.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - // F5 launches the server IN THE FOREGROUND under the ROCm venv (no nohup), so - // stdout/stderr land straight in the Debug Console and breakpoints work. - "version": "0.2.0", - "configurations": [ - { - "name": "FreeToken: serve qwen-moe (debug)", - "type": "debugpy", - "request": "launch", - "module": "freetoken.cli", - "python": "${workspaceFolder}/.venv-rocm/bin/python", - "cwd": "${workspaceFolder}", - "env": { "PYTHONPATH": "${workspaceFolder}/python" }, - "console": "integratedTerminal", - "args": [ - "serve", - "--model", "/media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf", - "--moe-backend", "offload", - "--attention-backend", "triton", - "--moe-cache-size", "2048", - "--num-tokens", "131072", - "--host", "127.0.0.1", - "--port", "1920" - ], - "justMyCode": false - } - ] -} \ No newline at end of file diff --git a/docs/install-amd.md b/docs/install-amd.md index 62ecf2fe7..ede4ab80f 100644 --- a/docs/install-amd.md +++ b/docs/install-amd.md @@ -6,7 +6,7 @@ recovered via HIP ports where safe. This page covers installing and running on R > Status: **experimental.** The default and best-tested path remains CUDA. AMD brings up a > correct functional path (Triton attention + offload/CPU MoE + portable quant) and is -> recovering performance via the HIP kernel ports. See `.plans/amd-gpu-support/plan.md`. +> recovering performance via the HIP kernel ports. ## Requirements @@ -41,10 +41,10 @@ and their backends are rejected with a clean error if requested. | Feature | On AMD | Notes | | --- | --- | --- | | Attention | `--attention-backend triton` | flashinfer/fa/trtllm are NVIDIA-only and rejected | -| MoE | `--moe-backend offload / cpu / hybrid` | offload needs pinned host memory (Inc 3) | +| MoE | `--moe-backend offload / cpu / hybrid` | offload needs pinned host memory | | Quant | BF16, MXFP4, GGUF (Q4_K/Q8_0), Triton inline-dequant NVFP4 | Marlin INT4 / native NVFP4 SASS unavailable | | NVFP4 checkpoints with no MXFP4 variant | converted to MXFP4 on load (auto) | `--nvfp4-backend auto` → triton/MXFP4 | -| CUDA graphs (decode) | HIP graph capture **if** the Inc-1 gate passes | otherwise kernel-launch decode | +| CUDA graphs (decode) | HIP graph capture **if** the capture probe passes | otherwise kernel-launch decode | | Multi-GPU (RCCL) | out of scope (single-GPU milestone) | | ## CLI behavior on AMD diff --git a/pyproject.toml b/pyproject.toml index 13b52c19f..2153c7897 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,9 +80,9 @@ fi = ["flashinfer-python[cu13]>=0.6,<0.7"] sgl = ["sglang-kernel==0.4.5"] accel = ["freetoken[fi,sgl]"] # ROCm (AMD) install: the NVIDIA-only fi/sgl/Marlin packages are NOT pulled in. torch must -# come from the ROCm wheel index (e.g. `pip install torch==2.11.0+rocm6.2` from -# https://download.pytorch.org/whl/rocm6.2) so the native extensions build against the HIP -# runtime; this extra pins the rest. See docs/install.md (AMD section, Inc 9). +# come from the ROCm wheel index (e.g. `pip install torch==2.11.0+rocm7.2` from +# https://download.pytorch.org/whl/rocm7.2) so the native extensions build against the HIP +# runtime; this extra pins the rest. See docs/install-amd.md. rocm = [ "triton==3.6.0; platform_system == 'Linux'", ] diff --git a/python/freetoken/attention/torch.py b/python/freetoken/attention/torch.py index 6ae16c2e3..c19a9fcf9 100644 --- a/python/freetoken/attention/torch.py +++ b/python/freetoken/attention/torch.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, List +import os import torch from freetoken.core import Batch, get_global_ctx @@ -64,10 +65,8 @@ def __init__(self, config: ModelConfig): self.num_kv_heads = int(getattr(spec, "num_kv_heads", self.num_kv_heads)) break # Debugging: contiguous (per-request) cache instead of the paged pool, to - # isolate cache addressing from the attention compute (Inc 5). + # isolate cache addressing from the attention compute. self._contig: dict[tuple[int, int], list] = {} - import os - self._use_contig = os.environ.get("FT_DEBUG_CONTIG_CACHE") == "1" def _build_metadata(self, batch: Batch) -> TorchMetadata: diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index d86bcf09e..74877ce7e 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -106,17 +106,14 @@ def _backend_parts_serve(name: str, required: frozenset[AttnType]) -> bool: def _backend_requirements_met(name: str) -> bool: + infos = [attention_backend_info(part) for part in name.split(",")] # On ROCm (AMD) only the portable backends exist: flashinfer/sgl/trtllm (and anything # sm_100-gated) are NVIDIA-only, so short-circuit before probing them at all. - from freetoken.utils.arch import is_rocm - if is_rocm(): return all(not i.requires_flashinfer and not i.requires_sgl_kernel - and not i.requires_sm100 for i in - [attention_backend_info(p) for p in name.split(",")]) + and not i.requires_sm100 for i in infos) # flashinfer first across ALL parts: the sgl probe logs a "falls back to fi" warning, # which would mislead when the candidate is about to fail on flashinfer anyway. - infos = [attention_backend_info(part) for part in name.split(",")] if any(i.requires_flashinfer for i in infos) and not _flashinfer_available(): return False if any(i.requires_sgl_kernel for i in infos) and not _sgl_flash_attn_available(): diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 50bac2f0a..276b741ef 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -132,7 +132,7 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel # graphs-disabled early return so that config gets the phase too. emit_progress("Capturing CUDA graphs / warming up", 0, 0) self.graph_map: Dict[int, torch.cuda.CUDAGraph] = {} - # Inc-8 parity: on ROCm, honour the Inc-1 graph-gate result. If capture is not + # ROCm parity: honour the graph-capture gate result. If capture is not # viable on this AMD card, skip graphs entirely so decode uses the kernel-launch # path (correct, just not graph-accelerated) rather than erroring mid-capture. from freetoken.utils.arch import is_rocm diff --git a/python/freetoken/kernel/backend.py b/python/freetoken/kernel/backend.py index 7b7629b11..8e685bbc1 100644 --- a/python/freetoken/kernel/backend.py +++ b/python/freetoken/kernel/backend.py @@ -13,12 +13,6 @@ from freetoken.utils.arch import is_rocm -# NVIDIA-only optional native packages: even if an importable copy is present on a ROCm -# torch build (e.g. a stray CUDA wheel), they must not be used -- the runtime falls back -# to the portable Triton kernels. Treated as unavailable on ROCm. -_CUDA_ONLY_PACKAGES = frozenset({"flashinfer", "sgl_kernel", "triton_kernels"}) - - def _importable(name: str) -> bool: # find_spec normally returns None when a package is absent, but it can raise # (broken parent package, or a meta_path finder that blocks the name); treat diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index de28877ce..1e063374c 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -82,8 +82,7 @@ inline auto HIP_CHECK(::hipError_t error, std::source_location location = std::source_location::current()) -> void { - if (error != ::hipSuccess) { - [[unlikely]]; + if (error != ::hipSuccess) [[unlikely]] { ::host::panic(location, "HIP error: ", ::hipGetErrorString(error)); } } @@ -98,8 +97,7 @@ inline auto CUDA_CHECK(::cudaError_t error, std::source_location location = std::source_location::current()) -> void { - if (error != ::cudaSuccess) { - [[unlikely]]; + if (error != ::cudaSuccess) [[unlikely]] { ::host::panic(location, "CUDA error: ", ::cudaGetErrorString(error)); } } diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 13b0ea0c5..00d6678da 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -56,9 +56,10 @@ def _module(): # the kernels compile their HIP branches; drop the CUDA-only -ccbin/flag logic. # Explicit --offload-arch (plus PYTORCH_ROCM_ARCH) prevents torch from auto- # emitting ~14 gfx arches, which would multiply build time per arch. - os.environ.setdefault("PYTORCH_ROCM_ARCH", "gfx1100") + gfx = os.getenv("FREETOKEN_KERNEL_CACHE_GFX", "gfx1100") + os.environ.setdefault("PYTORCH_ROCM_ARCH", gfx) extra_cuda_cflags = [ - "-O3", "--offload-arch=gfx1100", "-DUSE_HIP=1", "-DUSE_ROCM=1", + "-O3", f"--offload-arch={gfx}", "-DUSE_HIP=1", "-DUSE_ROCM=1", ] os.environ.pop("CXX", None) os.environ.pop("CC", None) diff --git a/python/freetoken/kernel/tinygrad_fallback.py b/python/freetoken/kernel/tinygrad_fallback.py deleted file mode 100644 index 2b0663fe3..000000000 --- a/python/freetoken/kernel/tinygrad_fallback.py +++ /dev/null @@ -1,92 +0,0 @@ -"""tinygrad-JIT fallback for FFI kernels that are not hand-ported to HIP. - -FreeToken's hand-written tvm-ffi kernels (``store`` / ``index`` / ``fast_index_copy`` / -``batch_memcpy``) are CUDA source compiled via nvcc/JIT. The primary AMD port is the -``#if defined(USE_HIP)`` seam in ``device_api.h`` + ``LaunchKernel``/``warp.cuh``. This -module is the **documented fallback** for any kernel that proves intractable to hipify: -tinygrad's JIT compiles one logical kernel to PTX (CUDA) *and* AMDGPU/LLVM (ROCm), so the -same source covers both platforms. - -Constraints (matching the FFI contract): - -* Each fallback takes the same ``tvm.ffi.TensorView`` arguments as the hand-written - kernel and returns the same output tensor(s), so the swap is invisible to callers. -* It runs on the *host* (tinygrad handles GPU dispatch); on ROCm it compiles to AMDGPU. -* It is **never a default**: ``kernel/utils.py`` only routes a kernel to the fallback - when (a) ROCm is active and (b) the hand-HIP AOT/JIT variant is absent/unbuildable. - If tinygrad is not installed, invoking the fallback raises a clear error. - -Because tinygrad is an optional dependency (installed only when the fallback is actually -needed), all imports here are lazy and the module imports with zero third-party deps, so -it is safe to import on the CUDA-only path. -""" - -from __future__ import annotations - -from typing import Callable, Optional - -__all__ = [ - "is_tinygrad_available", - "kernel_fallback_available", - "get_kernel_fallback", -] - -# Kernel names the fallback registry knows how to build (mirrors the FFI kernel set). -_KNOWN_KERNELS = ("store", "index", "fast_index_copy", "batch_memcpy") - -#: Which kernels currently have a *functional* tinygrad reimplementation. As HIP ports -#: land in Inc 7, names are removed from this set (the hand port wins); kernels left here -#: (if any) are the documented fallback set. Default: empty -- the hand-HIP port is the -#: primary path and the fallback is opt-in per kernel. -_FALLBACK_IMPLEMENTED: set[str] = set() - - -def is_tinygrad_available() -> bool: - """True when the ``tinygrad`` package can be imported (JIT-to-ROCm available).""" - try: - import importlib.util # noqa: PLC0415 - - return importlib.util.find_spec("tinygrad") is not None - except Exception: - return False - - -def kernel_fallback_available(kernel: str) -> bool: - """True when a tinygrad fallback for ``kernel`` is both implemented and usable - (tinygrad installed). Always False on the CUDA path unless explicitly enabled, so - the CUDA build never depends on tinygrad.""" - if kernel not in _FALLBACK_IMPLEMENTED: - return False - return is_tinygrad_available() - - -def get_kernel_fallback(kernel: str): - """Return the tinygrad-backed fallback callable for ``kernel``, or raise a clear - error explaining why it is unavailable. Never called on the CUDA path.""" - if kernel not in _FALLBACK_IMPLEMENTED: - raise RuntimeError( - f"FFI kernel {kernel!r} has no tinygrad fallback registered. On ROCm the " - "preferred path is the hand-written HIP port (device_api.h); if you intend " - "to use the tinygrad fallback you must register it in " - "kernel/tinygrad_fallback.py._FALLBACK_IMPLEMENTED and implement the " - "corresponding build function." - ) - if not is_tinygrad_available(): - raise RuntimeError( - f"FFI kernel {kernel!r} requires the tinygrad fallback, but tinygrad is not " - "installed. Install it (`pip install tinygrad`) or provide a hand-written " - "HIP port for this kernel." - ) - from freetoken.kernel import tinygrad_impl # noqa: PLC0415 (lazy; may be None) - - builder = getattr(tinygrad_impl, f"build_{kernel}", None) - if builder is None: - raise RuntimeError( - f"tinygrad fallback for {kernel!r} is registered but has no " - "tinygrad_impl.build_() builder." - ) - return builder - - -def _list_fallbacks() -> list[str]: - return [k for k in _FALLBACK_IMPLEMENTED if kernel_fallback_available(k)] diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index 1beb73982..e04fc3965 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -106,12 +106,6 @@ def _version_parts(version: str) -> Tuple[str, List[str]]: return base, local.split(".") if local else [] -def _build_stamps(segments: List[str]) -> set[str]: - """The `g` commit-stamp tokens of a local version segment list - (stamped by scripts/build-release-wheels.sh).""" - return {s for s in segments if re.fullmatch(r"g[0-9a-f]{7,40}", s)} - - def _build_stamps(local_segments: List[str]) -> List[str]: """The `g` commit-stamp tokens of a local version segment list (``["cu130", "g3f01615"]`` -> ``["g3f01615"]``).""" diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index 7b7f227c0..9c208e608 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -77,7 +77,10 @@ def __init__( self._block_fp8 = expert_quant == "fp8_block" self._pertensor_fp8 = attn_quant == "fp8_pertensor" self._gguf = expert_quant == "gguf" - self._fp8 = self._block_fp8 or self._pertensor_fp8 or self._gguf + # "Split" projection layout (qkv|z + ba as two GEMMs): used by the fp8 paths + # and by GGUF (native-quant qkv|z + dense bf16 ba). The fused 4-way in_proj is + # only for the plain bf16 case. + self._split_proj = self._block_fp8 or self._pertensor_fp8 or self._gguf self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] if self._block_fp8 or self._pertensor_fp8: @@ -171,7 +174,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: fla = build_fla_metadata(batch, hidden_states.device) batch.fla_metadata = fla - if self._fp8: + if self._split_proj: qkvz = self.in_proj_qkvz.forward(hidden_states) conv_in, z = torch.split(qkvz, [self.conv_dim, self.value_dim], dim=-1) ba = self.in_proj_ba.forward(hidden_states) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index aef81a099..e6ffb04cf 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -444,10 +444,6 @@ def iter_gguf_weights( # -------------------------------------------------------------------------------------- -def _q8_0_down_row_bytes(I: int) -> int: - return row_bytes(I, GGML_Q8_0) - - def load_gguf_expert_sources( model_path: str, config: ModelConfig, *, layer_sink=None ) -> dict[str, list[torch.Tensor]]: diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index e0801232e..be1fd9668 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -259,15 +259,12 @@ def _gguf_banks(model_path, model_config, device, dtype, dummy, parallel=False, "(not safetensors), so the common reader doesn't apply." ) if dummy: - from freetoken.models.weight import dummy_q4_0_moe_expert_sources - raise NotImplementedError("gguf expert banks have no dummy path; load the real GGUF") from freetoken.models.weight import load_gguf_moe_expert_sources - sink = None if dummy else layer_sink - sources = load_gguf_moe_expert_sources(model_path, model_config, layer_sink=sink) + sources = load_gguf_moe_expert_sources(model_path, model_config, layer_sink=layer_sink) return ExpertBanks( - "gguf", {name: sources[name] for name in _BANK_SCHEMAS["gguf"]}, streamed=sink is not None + "gguf", {name: sources[name] for name in _BANK_SCHEMAS["gguf"]}, streamed=layer_sink is not None ) diff --git a/python/freetoken/moe/nvfp4_to_mxfp4.py b/python/freetoken/moe/nvfp4_to_mxfp4.py index f4162a032..a81bf9e90 100644 --- a/python/freetoken/moe/nvfp4_to_mxfp4.py +++ b/python/freetoken/moe/nvfp4_to_mxfp4.py @@ -62,8 +62,6 @@ _FP4_TABLE = _np.asarray(_FP4_CODES, dtype=_np.float32) # Magnitudes sorted ascending for the nearest-code search. _FP4_SORT = _np.asarray(sorted(abs(v) for v in _FP4_CODES[1:8]), dtype=_np.float32) -_FP4_SORT_SIGN = _np.asarray([1.0 if i < 4 else -1.0 for i in range(len(_FP4_SORT))], - dtype=_np.float32) def fp4_e2m1_table() -> Sequence[float]: @@ -71,19 +69,6 @@ def fp4_e2m1_table() -> Sequence[float]: return list(_FP4_CODES) -def _nearest_e2m1_codes(values: _np.ndarray) -> _np.ndarray: - """Nearest e2m1 *code* for each fp32 ``values`` (signed, including 0/NaN).""" - a = _np.abs(values) - diff = _np.abs(a[..., None] - _FP4_SORT) # [..., 7] - idx = diff.argmin(axis=-1) - mag = _FP4_SORT[idx] - neg = _np.signbit(values) - code = (idx + 1).astype(_np.uint8) # _FP4_SORT[i] == table[i+1]; positive codes 1..7 - out = _np.where(neg, code | 0x8, code) - # Magnitudes below the smallest representable value (0.5) round to +0. - return _np.where(mag < 0.25, 0, out) - - def e8m0_scale_and_codes(values: _np.ndarray, block: int = 32) -> tuple[_np.ndarray, _np.ndarray]: """Return ``(scale_codes, fp4_codes)`` for ``values`` shaped ``[..., block]``: an e8m0 ``uint8`` scale per block (the smallest power-of-2 scale covering the block @@ -183,6 +168,10 @@ def convert_nvfp4_to_mxfp4( # Dequantize NVFP4 to fp32, move K to the last axis. K2 = packed.shape[-1] K = K2 * 2 + assert K % block == 0, ( + f"NVFP4 K={K} is not a multiple of the MXFP4 block ({block}); " + "cannot convert without truncating weights" + ) codes = np.stack([packed & 0x0F, (packed >> 4)], axis=-1).reshape( *packed.shape[:-1], K ) @@ -192,9 +181,8 @@ def convert_nvfp4_to_mxfp4( # Requantize to per-`block` e8m0 + e2m1. flat = f32.reshape(-1, K) - # pad to a multiple of block for the reshape (K is a multiple of 32 in practice) n_blocks = K // block - flat_b = flat[:, : n_blocks * block].reshape(-1, block) + flat_b = flat.reshape(-1, block) scale_codes, mxfp4_codes = e8m0_scale_and_codes(flat_b, block=block) out_codes = mxfp4_codes.reshape(flat.shape[0], n_blocks * block) diff --git a/python/freetoken/utils/graph_gate.py b/python/freetoken/utils/graph_gate.py index b528a208f..07089ebf2 100644 --- a/python/freetoken/utils/graph_gate.py +++ b/python/freetoken/utils/graph_gate.py @@ -1,10 +1,10 @@ """HIP/CUDA graph-capture parity probe. -The Inc-1 hard gate: whether ``torch.cuda.graph`` graph capture works on the target -GPU is the single highest-informational-risk assumption for AMD (ROCm) support. -This module probes it once and records a PASS/FAIL + device result that the rest -of the plan (Inc 8) reads. On CUDA it is expected to PASS; on ROCm it may fail on -some consumer cards, in which case Inc 8 must use the kernel-launch decode path. +Whether ``torch.cuda.graph`` graph capture works on the target GPU is the single +highest-informational-risk assumption for AMD (ROCm) support. This module probes it +once and records a PASS/FAIL + device result that the engine reads when deciding +whether to use CUDA-graph decode. On CUDA it is expected to PASS; on ROCm it may +fail on some consumer cards, in which case decode must use the kernel-launch path. The result is cached to disk under the user cache dir so it survives across runs, and keyed by device kind + device name so a change of GPU invalidates it. @@ -83,7 +83,6 @@ def probe_graph_capture() -> dict: # The child probes both an elementwise op (capturable on both backends) and a GEMM # (hipBLASLt on ROCm), which is what a real decode forward would run. The GEMM is # the discriminating case: on this ROCm build it fatally aborts -> child exit != 0. - import json as _json import subprocess as _subprocess import sys as _sys @@ -105,7 +104,7 @@ def probe_graph_capture() -> dict: "detail": f"fatal during capture: {detail[:240]}", } try: - data = _json.loads(child.stdout) + data = json.loads(child.stdout) except Exception: return { "device_kind": _device_kind(), @@ -177,7 +176,7 @@ def run_graph_gate() -> dict: @lru_cache(maxsize=1) def graph_capture_status() -> str: """Cached graph-capture status: ``"pass"``, ``"fail"``, or ``"unknown"`` (no device / - probe unavailable). Inc 8 reads this to pick HIP-graph vs kernel-launch decode.""" + probe unavailable). The graph runner reads this to pick HIP-graph vs kernel-launch decode.""" try: result = run_graph_gate() if result["ok"]: diff --git a/scripts/mirror-hf-configs.sh b/scripts/mirror-hf-configs.sh index 8cca566b9..1ab821e7a 100755 --- a/scripts/mirror-hf-configs.sh +++ b/scripts/mirror-hf-configs.sh @@ -6,12 +6,10 @@ # # ./mirror-hf-configs.sh [gguf-file] # -# ./mirror-hf-configs.sh Qwen/Qwen3.6-35B-A3B -# ./mirror-hf-configs.sh Qwen/Qwen3.6-35B-A3B /media/smk/Shared/Models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf +# ./mirror-hf-configs.sh Qwen/Qwen3.6-35B-A3B /path/to/model.gguf # -# With no gguf argument, defaults to $FT_MODEL, else the single *.gguf in -# /media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models or /media/smk/Shared/Models -# matching the repo's basename, else errors. +# With no gguf argument, uses $FT_MODEL if set; otherwise errors. The GGUF must be +# passed explicitly (or via FT_MODEL) -- no implicit model-directory search. # # Files fetched (when they exist upstream): # chat_template.jinja — read by FreeToken's GGUF loader (overrides embedded) @@ -29,12 +27,7 @@ if [ "$#" -ge 1 ]; then GGUF="$1" else GGUF="${FT_MODEL:-}" - if [ -z "$GGUF" ]; then - base="$(basename "${REPO_ID##*:}")" - cand="$(find /media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models /media/smk/Shared/Models -maxdepth 1 -iname "*${base%%-*}*.gguf" 2>/dev/null | head -1)" - [ -n "$cand" ] || { echo "ERROR: no gguf found; pass one explicitly" >&2; exit 1; } - GGUF="$cand" - fi + [ -n "$GGUF" ] || { echo "ERROR: no gguf file given; pass one explicitly or set FT_MODEL" >&2; exit 1; } fi [ -f "$GGUF" ] || { echo "ERROR: not a file: $GGUF" >&2; exit 1; } diff --git a/scripts/serve-qwen-moe.sh b/scripts/serve-qwen-moe.sh index 47de497c1..c31cd2865 100755 --- a/scripts/serve-qwen-moe.sh +++ b/scripts/serve-qwen-moe.sh @@ -31,7 +31,7 @@ set -euo pipefail REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" PY="${PY:-$REPO/.venv-rocm/bin/python}" -MODEL="${FT_MODEL:-/media/smk/5fce248d-bbdd-488d-8883-4f000f85cc10/Models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf}" +MODEL="${FT_MODEL:-}" HOST="${FT_HOST:-127.0.0.1}" PORT="${FT_PORT:-1920}" ATNN="${FT_ATTN:-triton}" # triton | torch @@ -94,6 +94,7 @@ server_pids() { } start() { + [ -n "$MODEL" ] || die "no model configured: set FT_MODEL=/path/to/model.gguf" if [ -n "$(server_pids)" ]; then echo "A server is already running (pid $(server_pids | tr '\n' ' ')). Use 'stop' first." exit 1 diff --git a/tests/attention/test_torch_backend.py b/tests/attention/test_torch_backend.py index f655a959b..3a132777d 100644 --- a/tests/attention/test_torch_backend.py +++ b/tests/attention/test_torch_backend.py @@ -1,4 +1,4 @@ -"""The debug ``"torch"`` attention backend (Inc 4 of fix-attention). +"""The debug ``"torch"`` attention backend. Verifies (a) the backend is registered and selectable, and (b) its pure-PyTorch GQA attention math (with causal masking and the per-head output gate) matches From ce11a07f39304922f136bb093f24deb1d8507323 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Sun, 30 Aug 2026 20:23:36 -0300 Subject: [PATCH 07/17] fix(rocm): retain hybrid attention capability flag --- python/freetoken/attention/__init__.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index bc3c3a0d3..61bdaad48 100644 --- a/python/freetoken/attention/__init__.py +++ b/python/freetoken/attention/__init__.py @@ -33,6 +33,9 @@ class BackendInfo: # Whether forward() honors a per-call AttentionSpec (window/sm_scale/sinks). # Non-consumers raise on a non-None spec instead of silently dropping it. consumes_attn_spec: bool = False + # Whether this backend coexists with hybrid-linear (GDN/mamba) models. Linear + # layers bypass attention backend dispatch. + hybrid_linear_ok: bool = True SUPPORTED_ATTENTION_BACKENDS = Registry[BackendCreator]("Attention Backend") From 9e7fe724fe1b119eec4c12d1c8c480aec8340d20 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Sun, 30 Aug 2026 22:07:49 -0300 Subject: [PATCH 08/17] feat(rocm-perf-parity): nvidia-path integrity pass (Inc 1) - suite delta vs main is zero; sys.modules snapshot-fix in fake-module tests; is_rocm pinned in arch-tree tests; unified engine is_rocm reference; strict g-stamp regex restored; qwen3.6 GGUF AOT entry; CUDA+ROCm unit workflows --- .github/workflows/unit-nvidia.yml | 69 +++++++++++++++++++ .github/workflows/unit-rocm.yml | 52 ++++++++++++++ python/freetoken/engine/engine.py | 4 +- python/freetoken/kernel/aot_models.py | 20 ++++++ python/freetoken/kernel/utils.py | 10 ++- tests/engine/test_attention_backend_arch.py | 6 ++ tests/engine/test_attention_backend_matrix.py | 3 + tests/kernels/test_backend_rocm.py | 18 ++++- tests/moe/test_nvfp4_backends_rocm.py | 16 ++++- tests/moe/test_offload.py | 9 +++ .../scheduler/test_abort_inflight_prefill.py | 3 + 11 files changed, 203 insertions(+), 7 deletions(-) create mode 100644 .github/workflows/unit-nvidia.yml create mode 100644 .github/workflows/unit-rocm.yml diff --git a/.github/workflows/unit-nvidia.yml b/.github/workflows/unit-nvidia.yml new file mode 100644 index 000000000..d742c9020 --- /dev/null +++ b/.github/workflows/unit-nvidia.yml @@ -0,0 +1,69 @@ +# NVIDIA unit-tests gate: compile the CUDA sources (catches USE_HIP/USE_ROCM-gated +# refactors that break the CUDA build) and run the test suite on a hosted runner. +# +# The runner has no GPU: cu126 torch installs fine, GPU tests self-skip, and the unit +# suite asserts the CUDA-path *logic* (gating, dispatch, version pairing) which is +# exactly what a HIP port can silently drift. Compile failures surface here. +# +# Trigger policy follows nightly-wheels.yml: only trusted paths for expensive builds. +# fork PRs cannot run the CUDA install (they can: hosted runner, no secrets). Keep +# pull_request restricted to code paths; everything else is workflow_dispatch. + +name: Unit tests (NVIDIA) + +on: + pull_request: + paths: + - "python/**" + - "tests/**" + - "setup.py" + - "pyproject.toml" + - ".github/workflows/unit-nvidia.yml" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: unit-nvidia-${{ github.ref }} + cancel-in-progress: true + +jobs: + cuda-compile-and-unit: + runs-on: ubuntu-latest + timeout-minutes: 60 + env: + # torch's CUDA wheel brings its own cudart headers; nvcc from the same toolkit + # major compiles the _pinned_tensor / gguf extensions against them. + TORCH_CUDA_VERSION: "126" + steps: + - uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install CUDA toolkit (nvcc) + uses: js-data/add-cuda-toolkit@v1.1.0 + with: + cuda-version: "12.6" + + - name: Install pytest (uv is preinstalled on hosted runners) + run: uv pip install --system pytest + + - name: Compile CUDA extensions (nvcc gate) + run: | + uv pip install --system torch --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION} + uv pip install --system -e . --no-build-isolation + + - name: Run unit suite (GPU tests self-skip) + run: | + python -m pytest tests/ -m "not slow" -q -x \ + --ignore=tests/e2e --ignore=tests/dsv4 2>&1 | tail -40 + + - name: Summarize + if: always() + run: echo "CUDA-compile + unit gate finished; see the step log for failures." \ No newline at end of file diff --git a/.github/workflows/unit-rocm.yml b/.github/workflows/unit-rocm.yml new file mode 100644 index 000000000..cb92e9223 --- /dev/null +++ b/.github/workflows/unit-rocm.yml @@ -0,0 +1,52 @@ +# ROCm unit-tests gate. Per plan (rocm-perf-parity/Inc 1): the unit suite must stay +# green on AMD hardware too — the box that serves gfx1100 is the only honest runner. +# +# This job runs on the self-hosted ROCm node, which is an ALWAYS-ON dev box: keep it +# manual (workflow_dispatch) and off-hours so it never contends with local dev or +# benchmark runs. It consumes whatever venv the runner has at .venv-rocm. +# There is deliberately NO pull_request trigger; fork code must never reach the +# self-hosted node (same policy as nightly-wheels.yml). + +name: Unit tests (ROCm, self-hosted) + +on: + workflow_dispatch: + inputs: + run_bench: + description: "Also run the bench_decode_moe smoke after the suite" + type: boolean + default: false + +permissions: + contents: read + +concurrency: + group: unit-rocm + cancel-in-progress: false + +jobs: + rocm-unit: + runs-on: [self-hosted, rocm] + timeout-minutes: 120 + steps: + - uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Show ROCm environment + run: | + .venv-rocm/bin/python -c "import torch; print(torch.__version__, torch.version.hip)" + rocm-smi | head -8 + + - name: Unit suite + run: | + PYTHONPATH=python .venv-rocm/bin/python -m pytest tests/ -m "not slow" -q \ + 2>&1 | tail -60 + + - name: Optional bench smoke + if: ${{ inputs.run_bench }} + run: | + PYTHONPATH=python .venv-rocm/bin/python benchmarks/bench_decode_moe.py \ + --model "${FREETOKEN_BENCH_MODEL:-/home/smk/models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf}" \ + --backend offload --json /tmp/ft-rocm-ci-bench.jsonl 2>&1 | tail -20 + continue-on-error: true \ No newline at end of file diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 113d51db6..eb780129d 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -207,8 +207,10 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att # explicit --attention-backend choices. for part in backend_parts: info = attention_backend_info(part) - from freetoken.utils.arch import is_rocm + # Module-level is_rocm (deliberately not re-imported locally): one reference for + # gating, so patching engine.is_rocm covers config-time validation exactly as it + # covers _backend_requirements_met. if is_rocm() and (info.requires_flashinfer or info.requires_sgl_kernel or info.requires_sm100): raise RuntimeError( f"Attention backend {config.attention_backend!r} is NVIDIA-only and " diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index c9c2fb98e..d75e8b97d 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -89,6 +89,11 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int if fmt == "q4_0": # gemma4/gguf.py _q4_0_expert_specs: GGML Q4_0 rows, 32 elems -> 18 bytes return {"gate_up": 2 * I * (H // 32 * 18), "down": H * (I // 32 * 18)} + if fmt == "gguf": + # qwen3_5_moe/gguf.py iter_gguf_weights: gate_up stays native Q4_K + # (row_bytes(H, Q4_K) = H//256*144), down re-quantized to Q8_0 + # (row_bytes(I, Q8_0) = I//32*34); matches moe/offload_cache.py "gguf". + return {"gate_up": 2 * I * (H // 256 * 144), "down": H * (I // 32 * 34)} if fmt in ("nvfp4", "nvfp4_marlin", "nvfp4_b12x"): # models/nvfp4_banks.py: packed e2m1 pairs + per-16 fp8-e4m3 scales + fp16 # per-row globals; marlin/b12x repacks are byte-identical with the globals @@ -177,6 +182,21 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int moe_intermediate_size=512, expert_formats=("fp8_block",), ), + AotModel( + name="Qwen/Qwen3.6-35B-A3B-GGUF", + # GGUF (native Q4_K/Q5_K/Q6_K/Q8_0) variant of the 3.5/3.6-35B-A3B hybrid: same + # model classes, different register key + config/weight loaders. Q4_K gate_up + + # down re-quantized to Q8_0 (moe/offload_cache.py "gguf" schema). GGUF embeddings + # bypass the index kernel (module docstring), so no embedding indexing spec. + architecture="Qwen3_5MoeForConditionalGeneration", + hidden_size=2048, + kv_groups=((2, 256),), + top_k=8, + moe_intermediate_size=512, + expert_formats=("gguf",), + embed_indexing=False, + arch_aliases=("Qwen35moeGGUFForCausalLM",), + ), AotModel( name="nvidia/Qwen3.6-35B-A3B-NVFP4", architecture="Qwen3_5MoeForConditionalGeneration", diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index e04fc3965..df7583003 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -108,8 +108,14 @@ def _version_parts(version: str) -> Tuple[str, List[str]]: def _build_stamps(local_segments: List[str]) -> List[str]: """The `g` commit-stamp tokens of a local version segment list - (``["cu130", "g3f01615"]`` -> ``["g3f01615"]``).""" - return [s for s in local_segments if s.startswith("g")] + (``["cu130", "g3f01615"]`` -> ``["g3f01615"]``). + + Only genuine build stamps count: ``g`` followed by 7..40 hex chars (the + pre-ROCm regex). A bare ``startswith("g")`` would promote arbitrary local + segments like ``gabcdefgh`` (non-hex letters) to stamps and wrongly reject a + dev build paired with a stamped one (tests/kernels/test_kernel_cache_version.py + pins this case).""" + return [s for s in local_segments if re.fullmatch(r"g[0-9a-f]{7,40}", s)] def _arch_tags(local_segments: List[str]) -> List[str]: diff --git a/tests/engine/test_attention_backend_arch.py b/tests/engine/test_attention_backend_arch.py index 9271bd425..010013142 100644 --- a/tests/engine/test_attention_backend_arch.py +++ b/tests/engine/test_attention_backend_arch.py @@ -64,6 +64,12 @@ def _engine_config(**overrides): def _patch_env(monkeypatch, *, major, flashinfer=True, sgl=True): from freetoken.engine import engine + # These tests assert the NVIDIA arch->backend decision tree in isolation. Pin the + # ROCm gate too: the short-circuit in _backend_requirements_met consults the real + # host (is_rocm() -> True on AMD boxes), which would otherwise leak an environment + # fact into a unit test that fakes the NVIDIA environment. ROCm resolution itself + # is covered by tests/engine/test_attention_backend_rocm.py. + monkeypatch.setattr(engine, "is_rocm", lambda: False) monkeypatch.setattr(engine, "is_sm100_family", lambda: major == 10) monkeypatch.setattr(engine, "is_sm90_family", lambda: major == 9) monkeypatch.setattr(engine, "_flashinfer_available", lambda: flashinfer) diff --git a/tests/engine/test_attention_backend_matrix.py b/tests/engine/test_attention_backend_matrix.py index fc9abfbeb..a737cbf6f 100644 --- a/tests/engine/test_attention_backend_matrix.py +++ b/tests/engine/test_attention_backend_matrix.py @@ -99,6 +99,9 @@ def _config(kind, **overrides): def _patch_env(monkeypatch, *, major=9, flashinfer=True, sgl=True): from freetoken.engine import engine + # Pin the ROCm gate like test_attention_backend_arch.py: these tests assert the + # NVIDIA decision tree; the host's real is_rocm() must not leak in. + monkeypatch.setattr(engine, "is_rocm", lambda: False) monkeypatch.setattr(engine, "is_sm100_family", lambda: major == 10) monkeypatch.setattr(engine, "is_sm90_family", lambda: major == 9) monkeypatch.setattr(engine, "_flashinfer_available", lambda: flashinfer) diff --git a/tests/kernels/test_backend_rocm.py b/tests/kernels/test_backend_rocm.py index 4c7326351..6b62e89e5 100644 --- a/tests/kernels/test_backend_rocm.py +++ b/tests/kernels/test_backend_rocm.py @@ -42,9 +42,23 @@ def _load_backend(rocm: bool): @pytest.fixture(autouse=True) def _clean_sys_modules(): + # Snapshot-restore, NOT pop: when the real freetoken package was already imported + # earlier in the session, popping the 4 fake names leaves the real submodules + # (freetoken.attention, freetoken.kernel.pinned, ...) cached with a dangling parent, + # so every later `getattr(freetoken, "attention")` / `import freetoken.kernel.pinned` + # in unrelated tests raises AttributeError/ImportError. Restoring the exact pre-test + # snapshot (and only creating entries we removed) keeps the session clean. + saved = { + key: mod + for key, mod in sys.modules.items() + if key == "freetoken" or key.startswith("freetoken.") or key == "torch" + } yield - for name in ("freetoken", "freetoken.utils", "freetoken.utils.arch", "freetoken.kernel"): - sys.modules.pop(name, None) + for key in [k for k in sys.modules if k == "freetoken" or k.startswith("freetoken.")]: + if key in saved and saved[key] is not None: + sys.modules[key] = saved[key] + else: + sys.modules.pop(key, None) def test_rocm_native_packages_unavailable(): diff --git a/tests/moe/test_nvfp4_backends_rocm.py b/tests/moe/test_nvfp4_backends_rocm.py index 0061d3788..cc71e4727 100644 --- a/tests/moe/test_nvfp4_backends_rocm.py +++ b/tests/moe/test_nvfp4_backends_rocm.py @@ -64,9 +64,21 @@ def _load_nvfp4(rocm: bool): @pytest.fixture(autouse=True) def _clean(): + # Snapshot-restore, NOT pop: see tests/kernels/test_backend_rocm.py — popping these + # names after replacing the real freetoken package leaves its already-imported + # submodules cached with a dangling parent module, cascading AttributeError/ + # ImportError into unrelated tests for the rest of the session. + saved = { + key: mod + for key, mod in sys.modules.items() + if key == "freetoken" or key.startswith("freetoken.") + } yield - for name in ("freetoken", "freetoken.utils", "freetoken.utils.arch"): - sys.modules.pop(name, None) + for key in [k for k in sys.modules if k == "freetoken" or k.startswith("freetoken.")]: + if key in saved and saved[key] is not None: + sys.modules[key] = saved[key] + else: + sys.modules.pop(key, None) def test_rocm_auto_resolves_triton(): diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 422ca8675..477d215d8 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -447,6 +447,15 @@ def test_graph_capture_reuses_warm_offload_cache_before_capture(monkeypatch): from freetoken.core import Context, Req, get_global_ctx from freetoken.engine.graph import GraphRunner + # This test asserts the capture bookkeeping contract in isolation. The ROCm gate + # (graph_capture_status) reads the real host and would return "fail" on an AMD box, + # skipping capture entirely; capture-path coverage lives with CUDA-graph machinery, + # so pin the gate to pass like the other NVIDIA-env fakes in this suite. + import freetoken.utils.graph_gate as graph_gate + + monkeypatch.setattr(graph_gate, "graph_capture_status", lambda: "pass") + monkeypatch.setattr(graph_gate, "graph_capture_env", lambda: {}, raising=False) + events = [] _init_tp() monkeypatch.setattr(core, "_GLOBAL_CTX", Context(page_size=1)) diff --git a/tests/scheduler/test_abort_inflight_prefill.py b/tests/scheduler/test_abort_inflight_prefill.py index a50d0d323..843b72a78 100644 --- a/tests/scheduler/test_abort_inflight_prefill.py +++ b/tests/scheduler/test_abort_inflight_prefill.py @@ -67,6 +67,9 @@ def _setup(): _mamba_slot_usage=lambda: None, _swa_token_usage=lambda: None, _gpu_mem_bytes=lambda: 0, + # The scheduler now reports per-window MoE stats in the decode log line + # (--moe-collect-stats); the stub has no engine/cache, so stats stay off (None). + _moe_stats_snapshot=lambda: None, _match_stop_str=lambda _req: None, _pending_abort_acks=set(), _last_data=None, From 811415e7ebd675091653a01070866ba63bee7a18 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Sun, 30 Aug 2026 23:41:09 -0300 Subject: [PATCH 09/17] feat(rocm-perf-parity): decode stage profiler (Inc 2) - FREETOKEN_TORCH_PROFILE env hook in forward_batch; moe_router/gate_up/down + attn record_function labels; profile-rocm-decode.sh; stage notes: lm_head+dense 41% of GPU busy, router 7.8%, gather 8.5%, ~43% of step is host/launch gap; stale GGUF JIT lock hygiene in kernel/gguf.py (FileBaton.wait hang) --- python/freetoken/attention/triton.py | 6 +- python/freetoken/engine/engine.py | 42 ++++---- python/freetoken/kernel/gguf.py | 52 +++++++++ python/freetoken/layers/moe.py | 16 +-- python/freetoken/moe/fused_gguf.py | 14 ++- python/freetoken/utils/step_profiler.py | 134 ++++++++++++++++++++++++ scripts/profile-rocm-decode.sh | 97 +++++++++++++++++ 7 files changed, 332 insertions(+), 29 deletions(-) create mode 100644 python/freetoken/utils/step_profiler.py create mode 100755 scripts/profile-rocm-decode.sh diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py index 9eed1e1d2..f397f8491 100644 --- a/python/freetoken/attention/triton.py +++ b/python/freetoken/attention/triton.py @@ -167,7 +167,11 @@ def forward( assert metadata.attn_logits is not None assert metadata.attn_lse is not None assert metadata.num_kv_splits is not None - return decode_paged_attention( + # "attn" record_function label = the profiler-segmented attention stage + # of the decode step (Inc 2, .plans/rocm-perf-parity). The prefill/extend + # branch below stays unlabeled (its stage shows via kernel names). + with torch.profiler.record_function("attn"): + return decode_paged_attention( q=q, k_cache=k_cache, v_cache=v_cache, diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index eb780129d..cbd1fa071 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -937,24 +937,30 @@ def rebuild_runtime_cache( def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: assert torch.cuda.current_stream() == self.stream - with self.ctx.forward_batch(batch): - if self.graph_runner.can_use_cuda_graph(batch): - logits = self.graph_runner.replay(batch) - else: - logits = self.model.forward() - if self.cpu_moe_executor is not None: - # One pinned read: surfaces a fired flag-handshake watchdog (dead coordinator - # -> stale expert outputs) as a loud error instead of silent corruption. - self.cpu_moe_executor.raise_if_unhealthy() - - for req in batch.reqs: - req.complete_one() - - batch_logits = logits[: batch.size] - next_tokens_gpu = self.sampler.sample(batch_logits, args, batch).to(torch.int32) - next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True) - copy_done_event = torch.cuda.Event() - copy_done_event.record(self.stream) + from freetoken.utils.step_profiler import step_profiler + + # Inc 2 instrument of .plans/rocm-perf-parity: stage-time breakdowns via + # FREETOKEN_TORCH_PROFILE (no-op/cached-flag when unset). Wraps the whole + # forward+sample step; range labels live at the MoE/router/attention callsites. + with step_profiler(): + with self.ctx.forward_batch(batch): + if self.graph_runner.can_use_cuda_graph(batch): + logits = self.graph_runner.replay(batch) + else: + logits = self.model.forward() + if self.cpu_moe_executor is not None: + # One pinned read: surfaces a fired flag-handshake watchdog (dead coordinator + # -> stale expert outputs) as a loud error instead of silent corruption. + self.cpu_moe_executor.raise_if_unhealthy() + + for req in batch.reqs: + req.complete_one() + + batch_logits = logits[: batch.size] + next_tokens_gpu = self.sampler.sample(batch_logits, args, batch).to(torch.int32) + next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True) + copy_done_event = torch.cuda.Event() + copy_done_event.record(self.stream) return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event) @torch.inference_mode() diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 00d6678da..661553e73 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -16,6 +16,8 @@ import os import pathlib import shutil +import subprocess +import time import torch @@ -47,6 +49,55 @@ def _c_compiler_for(cxx: str) -> str: cc = base.replace("g++", "gcc") return shutil.which(cc) or cc + +# A JIT-rebuild lock older than this is stale even while its owner lives: no honest +# rebuild of one extension takes hours (the slowest measured gguf_kernel build is +# ~13 min on the gfx1100 box). +_STALE_JIT_LOCK_AGE_S = 3 * 3600 + + +def _clear_stale_jit_lock(module_name: str) -> None: + """Remove a torch-extension ``lock`` whose owner died. + + torch's ``FileBaton`` waits forever when a previous compile was ``kill -9``-ed + mid-rebuild (observed: every later serve hung in warmup at + ``cpp_extension.py _jit_compile -> wait``; the lock file has no owner pid and + no fd is held on it, so there is nothing to poll). Guard on BOTH the age and + the absence of a live freetoken serve/worker process, so a legitimately + concurrent rebuild is never clobbered. + """ + try: + build_dir = pathlib.Path( + torch.utils.cpp_extension._get_build_directory(module_name, False) + ) + lock = build_dir / "lock" + if not lock.exists(): + return + age = time.time() - lock.stat().st_mtime + if age < _STALE_JIT_LOCK_AGE_S: + return + if _freetoken_processes_running(): + # Without the age check a mid-compile server would be clobbered; with it, + # anything left after hours while no freetoken process lives is the corpse + # of a killed run. + return + lock.unlink() + except Exception: # noqa: BLE001 - hygiene must never break the build path + pass + + +def _freetoken_processes_running() -> bool: + try: + out = subprocess.run( + ["pgrep", "-f", r"freetoke[n].cli serve|multiprocessing.s[p]awn"], + capture_output=True, + text=True, + timeout=10, + ) + return bool(out.stdout.strip()) + except Exception: # noqa: BLE001 + return True # cannot tell -> assume a live owner + @functools.cache def _module(): from torch.utils.cpp_extension import load @@ -77,6 +128,7 @@ def _module(): # 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. + _clear_stale_jit_lock("freetoken_gguf_kernels") return load( name="freetoken_gguf_kernels", sources=[str(_CSRC / "gguf_kernel.cu")], diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index e16913639..ac01eed90 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -171,12 +171,16 @@ def forward( if self.weight_format != "bf16": # Quantized resident experts: generic softmax router + format kernel. # The bf16 path below stays on ctx.moe_backend byte-for-byte. - topk_weights, topk_ids = fused_topk( - hidden_states=hidden_states, - gating_output=router_logits, - topk=self.top_k, - renormalize=self.renormalize, - ) + # "moe_router" record_function label = the profiler-segmented router stage + # (Inc 2, .plans/rocm-perf-parity); a range object costs ~1µs and keeps the + # name visible in torch.profiler tables on both backends. + with torch.profiler.record_function("moe_router"): + topk_weights, topk_ids = fused_topk( + hidden_states=hidden_states, + gating_output=router_logits, + topk=self.top_k, + renormalize=self.renormalize, + ) return self._maybe_all_reduce( self._resident_gemm(hidden_states, topk_weights, topk_ids) ) diff --git a/python/freetoken/moe/fused_gguf.py b/python/freetoken/moe/fused_gguf.py index 9f8995e94..db1e386a0 100644 --- a/python/freetoken/moe/fused_gguf.py +++ b/python/freetoken/moe/fused_gguf.py @@ -38,11 +38,17 @@ def fused_experts_gguf( h = down_q.shape[1] # hidden top_k = topk_ids.shape[1] - gate_up = ggml_moe_a8_vec( - hidden_states, gate_up_q, topk_ids, top_k, int(GGML_Q4_K), n2, num_tokens - ) + # "moe_gate_up" / "moe_down" record_function labels = the profiler-segmented + # expert-GEMM halves of the fused MoE forward (Inc 2, .plans/rocm-perf-parity). + with torch.profiler.record_function("moe_gate_up"): + gate_up = ggml_moe_a8_vec( + hidden_states, gate_up_q, topk_ids, top_k, int(GGML_Q4_K), n2, num_tokens + ) inter = act_fn(gate_up) - out = ggml_moe_a8_vec(inter, down_q, topk_ids, 1, int(GGML_Q8_0), h, num_tokens * top_k) + with torch.profiler.record_function("moe_down"): + out = ggml_moe_a8_vec( + inter, down_q, topk_ids, 1, int(GGML_Q8_0), h, num_tokens * top_k + ) out = out.reshape(num_tokens, top_k, h) * topk_weights.reshape(num_tokens, top_k, 1).to( out.dtype ) diff --git a/python/freetoken/utils/step_profiler.py b/python/freetoken/utils/step_profiler.py new file mode 100644 index 000000000..693d0944f --- /dev/null +++ b/python/freetoken/utils/step_profiler.py @@ -0,0 +1,134 @@ +"""Env-gated per-step torch.profiler wrapper (stage-level decode breakdowns). + +``FREETOKEN_TORCH_PROFILE="::"``, unset/empty = no-op (the only +hot-path cost is one cached check). Skip the first ``warm`` call(s) of the wrapped +region, then profile the next ``steps`` calls in one ``torch.profiler`` window and +export on the window's exit: + +- a chrome trace to ```` and +- a top-kernels table to ``-kernels.log``. + +This is the Inc 2 instrument of .plans/rocm-perf-parity: NVTX is a no-op on ROCm, +so the trace segments by the explicit ``torch.profiler.record_function`` range +names (moe_router, moe_gate_up, moe_down, attn, Sampler) that appear as table row +names on both backends. +""" + +from __future__ import annotations + +import os + +_NO_SPEC = object() # "not parsed yet" sentinel + + +def parse_spec(raw: str) -> tuple[int, int, str] | None: + """Parse '::' -> (warm, steps, out); None when empty.""" + raw = raw.strip() + if not raw: + return None + try: + warm_s, steps_s, out = raw.split(":") + warm, steps = int(warm_s), int(steps_s) + if warm < 0 or steps <= 0 or not out: + raise ValueError + except ValueError as exc: + raise ValueError( + "FREETOKEN_TORCH_PROFILE must be '::' " + f"(warm>=0, steps>0), got {raw!r}" + ) from exc + return warm, steps, out + + +def _read_spec() -> tuple[int, int, str] | None: + return parse_spec(os.environ.get("FREETOKEN_TORCH_PROFILE", "")) + + +class _State: + __slots__ = ("spec", "calls", "done") + + def __init__(self) -> None: + self.spec: tuple[int, int, str] | None = _NO_SPEC # type: ignore[assignment] + self.calls = 0 + self.done = False + + +_state = _State() + + +class _NullCtx: + """No-op context manager (profiler disabled or outside the profiled window).""" + + __slots__ = () + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +class _ProfilerCtx: + """One wrapped step inside the profiled window; exports on the last step's exit.""" + + __slots__ = ("prof", "final") + + def __init__(self, prof, final: bool) -> None: + self.prof = prof + self.final = final + if final: + _state.done = True + + def __enter__(self): + self.prof.__enter__() + + def __exit__(self, exc_type=None, exc=None, tb=None) -> bool: # noqa: ANN001 + self.prof.__exit__(exc_type, exc, tb) + if self.final: + warm, steps, out = _state.spec or (0, 0, "") + _export(self.prof, out, steps) + return False + + +_NULL = _NullCtx() + + +def step_profiler(): + """Wrap one scheduler step. No-op unless FREETOKEN_TORCH_PROFILE is set (parsed + on the first call); the window covers `steps` calls after the first `warm`.""" + if _state.done: + return _NULL + if _state.spec is _NO_SPEC: # type: ignore[comparison-overlap] + spec = _read_spec() + if spec is None: + _state.done = True + return _NULL + _state.spec = spec + warm, steps, _ = _state.spec # type: ignore[misc] + _state.calls += 1 + if _state.calls <= warm: + return _NULL + + import torch.profiler + + prof = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=False, + with_stack=False, + ) + return _ProfilerCtx(prof, final=_state.calls >= warm + steps) + + +def _export(prof, out: str, steps: int) -> None: # pragma: no cover - heavy + import os as _os + + _os.makedirs(_os.path.dirname(out) or ".", exist_ok=True) + prof.export_chrome_trace(out) + table = prof.key_averages().table(sort_by="cuda_time_total", row_limit=25) + table_path = out.rsplit(".", 1)[0] + "-kernels.log" + with open(table_path, "w") as f: + f.write(f"torch.profiler key_averages over {steps} profiled step(s)\n") + f.write(table) + f.write("\n") \ No newline at end of file diff --git a/scripts/profile-rocm-decode.sh b/scripts/profile-rocm-decode.sh new file mode 100755 index 000000000..dd3262afe --- /dev/null +++ b/scripts/profile-rocm-decode.sh @@ -0,0 +1,97 @@ +#!/usr/bin/env bash +# profile-rocm-decode.sh +# +# Inc 2 of .plans/rocm-perf-parity: stage-time breakdown of decode on gfx1100 for +# the Qwen3.6-35B-A3B GGUF (matching the P217 baseline conditions). +# +# Runs ft serve with FREETOKEN_TORCH_PROFILE=:: (the env reaches +# the server because serve-qwen-moe.sh execs $PY from this shell), sends one +# AIME-style request, waits for the exported trace, then stops the server. Env: +# PROFILE_OUT chrome-trace base path (default /tmp/ft-rocm-profile/chrome.json) +# PROFILE_WARM decode steps skipped before the profiled window (default 40) +# PROFILE_STEPS steps inside the profiled window (default 40) +# FT_MODEL model path (default the 7900 XTX box's Qwen3.6 GGUF) +# FT_PORT server port (default 1920) + +set -euo pipefail + +REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +PY="${PY:-$REPO/.venv-rocm/bin/python}" + +PROFILE_OUT="${PROFILE_OUT:-/tmp/ft-rocm-profile/chrome.json}" +PROFILE_WARM="${PROFILE_WARM:-40}" +PROFILE_STEPS="${PROFILE_STEPS:-40}" +FT_MODEL="${FT_MODEL:-/home/smk/models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf}" +FT_PORT="${FT_PORT:-1920}" +PORT="$FT_PORT" +export FT_MODEL FT_PORT + +if grep -q $'\r' "${BASH_SOURCE[0]}"; then + echo "ERROR: CRLF line endings in profile-rocm-decode.sh" >&2 + exit 1 +fi +[ -x "$PY" ] || { echo "ERROR: python not found: $PY" >&2; exit 1; } +[ -f "$FT_MODEL" ] || { echo "ERROR: model not found: $FT_MODEL" >&2; exit 1; } + +OUT_DIR="$(dirname "$PROFILE_OUT")" +mkdir -p "$OUT_DIR" +rm -f "$PROFILE_OUT" "${PROFILE_OUT%.json}-kernels.log" + +echo "Profile env: FREETOKEN_TORCH_PROFILE=${PROFILE_WARM}:${PROFILE_STEPS}:${PROFILE_OUT}" + +# serve-qwen-moe.sh refuses a double start; clear any previous instance first. +./scripts/serve-qwen-moe.sh stop >/dev/null 2>&1 || true + +# serve-qwen-moe.sh refuses a double start and returns once ready (~3-4 min load). +FREETOKEN_TORCH_PROFILE="${PROFILE_WARM}:${PROFILE_STEPS}:${PROFILE_OUT}" \ + ./scripts/serve-qwen-moe.sh + +# Send the profiled request (max_tokens > PROFILE_STEPS, so the engine stays +# decoding through the whole profiled window; the trace exports at its last step). +FREETOKEN_TORCH_PROFILE_MAX="${PROFILE_STEPS}" \ +"${PY}" - <<'PYEOF' +import json, os, sys, urllib.request + +port = os.environ["FT_PORT"] +max_tokens = ( + int(os.environ["PROFILE_WARM"]) + int(os.environ["PROFILE_STEPS"]) + 80 +) +prompt = ( + "Every morning Aya goes for a 9 kilometer walk, stops at a coffee shop, then " + "walks back home. She walks at 4 km/h, and the coffee shop detour adds 15 " + "minutes. On a day when she walks at t km/h and the detour still costs her 2 " + "hours total, what is t? Answer with just the number." +) +body = json.dumps({ + "model": os.path.basename(os.environ["FT_MODEL"]), + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_tokens, + "stream": False, +}).encode() +req = urllib.request.Request( + f"http://127.0.0.1:{port}/v1/chat/completions", + data=body, + headers={"Content-Type": "application/json"}, +) +with urllib.request.urlopen(req, timeout=1200) as r: + data = json.loads(r.read()) +print("request completed; completion_tokens:", data.get("usage", {}).get("completion_tokens")) +PYEOF + +# The window's last step exports on exit; decode continues past the window, so poll +# briefly rather than waiting for request completion (server stays decoding). +for i in $(seq 1 10); do + [ -f "$PROFILE_OUT" ] && break + sleep 2 +done + +./scripts/serve-qwen-moe.sh stop || true + +if [ -f "$PROFILE_OUT" ]; then + echo "Trace: $PROFILE_OUT" + echo "Kernels: ${PROFILE_OUT%.json}-kernels.log" +else + echo "ERROR: trace not exported to $PROFILE_OUT; check /tmp/serve_qwen_moe.log" >&2 + tail -20 /tmp/serve_qwen_moe.log >&2 || true + exit 1 +fi \ No newline at end of file From 98eabbd19e016c1f7d72bf73c15f2125746c8bdf Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Sun, 30 Aug 2026 23:52:32 -0300 Subject: [PATCH 10/17] feat(rocm-perf-parity): fused triton router on ROCm (Inc 3) - fused_topk routes ROCm to in-repo kernel/triton/moe_router.fused_topk_softmax (spy-verified, torch-reference-equal); bench 3-run median 37.67 tok/s vs baseline 34.69 (+8.6%), output sha stable per-server (66c6212ab749), no router-fallback warning --- python/freetoken/moe/fused.py | 16 +++++++++++++++- tests/moe/test_fused_moe.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py index 4b9a4875f..e45d5750d 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -45,6 +45,19 @@ def fused_topk( assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" from freetoken.kernel.backend import is_triton_kernels_installed + from freetoken.utils.arch import device_kind, is_rocm + + # ROCm: triton_kernels (OpenAI's, CUDA-oriented) returns False in backend.py, so + # route to the IN-REPO portable triton router instead of the pure-torch chain. + # kernel/triton/moe_router.fused_topk_softmax already handles renormalize + the + # num_token_non_padded device-scalar mask and ties-lowest-id, and is + # torch-reference-tested; this removes ~4 launches + a warpMergeSort topk from + # every decode step on AMD (Inc 3 of .plans/rocm-perf-parity, profiler-measured + # ~0.8 ms/step of GPU busy). + if is_rocm() and device_kind() != "cpu": + from freetoken.kernel.triton.moe_router import fused_topk_softmax + + return fused_topk_softmax(gating_output, topk, renormalize, num_token_non_padded) # triton_kernels ships no Windows wheel, and unlike flashinfer/sgl_kernel it is not one # of the six ops the in-repo triton kernels cover -- so this router needs its own fallback. @@ -58,7 +71,8 @@ def fused_topk( logger.warning_rank0( "fused_topk: triton_kernels is not installed -> pure-torch router fallback " "(numerically equivalent, slower). Expected on Windows (no wheel); on Linux " - "install triton_kernels to restore the fused router." + "install triton_kernels to restore the fused router. On ROCm/AMD the fused " + "in-repo triton router is used instead, so this message means CPU mode." ) return _torch_fused_topk(gating_output, topk, renormalize, num_token_non_padded) diff --git a/tests/moe/test_fused_moe.py b/tests/moe/test_fused_moe.py index 41a75c20e..fb8ff9fe6 100644 --- a/tests/moe/test_fused_moe.py +++ b/tests/moe/test_fused_moe.py @@ -369,3 +369,33 @@ def test_fused_topk_softmax_is_cuda_graph_capturable(): torch.cuda.synchronize() assert (ids[:3] != -1).all() assert (ids[3:] == -1).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="a GPU is required") +@pytest.mark.parametrize("num_tokens,num_experts,topk", [(1, 512, 8), (7, 512, 8), (33, 512, 8)]) +def test_fused_topk_rocm_routes_to_inrepo_triton_router( + num_tokens, num_experts, topk, monkeypatch +): + """On ROCm, fused_topk must route to the in-repo vendored triton router (never the + pure-torch chain -- Inc 3 of .plans/rocm-perf-parity) and match the torch reference.""" + from freetoken.kernel.triton import moe_router as router_mod + from freetoken.moe.fused import _torch_fused_topk, fused_topk + + routed = [] + real = router_mod.fused_topk_softmax + + def spy(*args, **kwargs): + routed.append(1) + return real(*args, **kwargs) + + monkeypatch.setattr(router_mod, "fused_topk_softmax", spy) + gen = torch.Generator(device="cuda").manual_seed(97) + gating = torch.randn(num_tokens, num_experts, generator=gen, device="cuda") + + weights, ids = fused_topk(gating, gating, topk, renormalize=True) + + assert routed, "fused_topk did not route through the in-repo triton router on ROCm" + ref_w, ref_i = _torch_fused_topk(gating, topk, True, None) + assert weights.dtype == torch.float32 and ids.dtype == torch.int32 + assert torch.equal(ids, ref_i) + torch.testing.assert_close(weights, ref_w, rtol=1e-5, atol=1e-6) From 632db88219c678370d7fb8198d32eebb05a60f79 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Sun, 30 Aug 2026 23:53:33 -0300 Subject: [PATCH 11/17] audit(rocm-perf-parity): checkpoint 1 fixes for Inc 1-3 - stale-JIT-lock guard now age-only (pgrep guard was self-defeating: the serve worker matches its own pattern so the stale lock was never cleared in the real hang case); step_profiler import hoisted to module level off the per-decode-step path --- python/freetoken/engine/engine.py | 5 +++-- python/freetoken/kernel/gguf.py | 29 ++++++----------------------- 2 files changed, 9 insertions(+), 25 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cbd1fa071..b4e20ebdb 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -297,6 +297,9 @@ def _materialize_loaded_weight_state_dict( return state_dict +from freetoken.utils.step_profiler import step_profiler + + class ForwardOutput(NamedTuple): next_tokens_gpu: torch.Tensor next_tokens_cpu: torch.Tensor @@ -937,8 +940,6 @@ def rebuild_runtime_cache( def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: assert torch.cuda.current_stream() == self.stream - from freetoken.utils.step_profiler import step_profiler - # Inc 2 instrument of .plans/rocm-perf-parity: stage-time breakdowns via # FREETOKEN_TORCH_PROFILE (no-op/cached-flag when unset). Wraps the whole # forward+sample step; range labels live at the MoE/router/attention callsites. diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 661553e73..22796e4d3 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -16,7 +16,6 @@ import os import pathlib import shutil -import subprocess import time import torch @@ -62,9 +61,12 @@ def _clear_stale_jit_lock(module_name: str) -> None: torch's ``FileBaton`` waits forever when a previous compile was ``kill -9``-ed mid-rebuild (observed: every later serve hung in warmup at ``cpp_extension.py _jit_compile -> wait``; the lock file has no owner pid and - no fd is held on it, so there is nothing to poll). Guard on BOTH the age and - the absence of a live freetoken serve/worker process, so a legitimately - concurrent rebuild is never clobbered. + no fd is held on it, so there is nothing to poll). Only the lock's age can + distinguish stale from live: a live rebuild's lock is minutes old; no honest + build of one extension takes hours (slowest measured gguf_kernel build ~13 min + on the gfx1100 box), so _STALE_JIT_LOCK_AGE_S (3 h) is the staleness bar. A + freshly-created lock is never touched, so a genuinely concurrent rebuild is + not clobbered. """ try: build_dir = pathlib.Path( @@ -76,28 +78,9 @@ def _clear_stale_jit_lock(module_name: str) -> None: age = time.time() - lock.stat().st_mtime if age < _STALE_JIT_LOCK_AGE_S: return - if _freetoken_processes_running(): - # Without the age check a mid-compile server would be clobbered; with it, - # anything left after hours while no freetoken process lives is the corpse - # of a killed run. - return lock.unlink() except Exception: # noqa: BLE001 - hygiene must never break the build path pass - - -def _freetoken_processes_running() -> bool: - try: - out = subprocess.run( - ["pgrep", "-f", r"freetoke[n].cli serve|multiprocessing.s[p]awn"], - capture_output=True, - text=True, - timeout=10, - ) - return bool(out.stdout.strip()) - except Exception: # noqa: BLE001 - return True # cannot tell -> assume a live owner - @functools.cache def _module(): from torch.utils.cpp_extension import load From 78e25488403dc053b6e3f8be7327c696b557f54d Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Mon, 31 Aug 2026 00:40:57 -0300 Subject: [PATCH 12/17] feat(rocm-perf-parity): gguf MMVQ tuning harness + negative result (Inc 4) - overridable GGML_CUDA_MMV_Y via FREETOKEN_GGUF_MMV_Y (JIT cflags-keyed, default 1 preserved); bench_gguf_moe_kernels.py with cross-variant byte-equality gate; sweep 1/2/4/8 at decode shapes: 73-77us/layer-pair within run noise, all variants byte-identical -> keep default, no kernel change --- benchmarks/bench_gguf_moe_kernels.py | 110 ++++++++++++++++++ .../freetoken/kernel/csrc/gguf/ggml-common.h | 4 + python/freetoken/kernel/gguf.py | 13 +++ 3 files changed, 127 insertions(+) create mode 100644 benchmarks/bench_gguf_moe_kernels.py diff --git a/benchmarks/bench_gguf_moe_kernels.py b/benchmarks/bench_gguf_moe_kernels.py new file mode 100644 index 000000000..f77bb41e0 --- /dev/null +++ b/benchmarks/bench_gguf_moe_kernels.py @@ -0,0 +1,110 @@ +"""Micro-bench for the GGUF fused-MoE MMVQ kernel pair (Inc 4, .plans/rocm-perf-parity). + +Times the two ``ggml_moe_a8_vec`` calls the decode hot loop makes per MoE layer at +the Qwen3.6-35B-A3B shapes (H=2048, I=512, top_k=8; Q4_K gate_up + Q8_0 down) on a +stacked-expert bank with ``slots`` resident experts, exactly like the serving path +(``moe/fused_gguf.py``). + +Cross-variant numerics gate: run once at the incumbent config with ``--dump PATH``, +then at any other ``FREETOKEN_GGUF_MMV_Y`` build with ``--check PATH``; the bench +hard-fails if a single output byte differs. + +Usage (each MMV_Y value triggers its own JIT build, ~1-2 min the first time): + PYTHONPATH=python .venv-rocm/bin/python benchmarks/bench_gguf_moe_kernels.py \ + --iters 300 --dump /tmp/moe-ref.pt + FREETOKEN_GGUF_MMV_Y=8 PYTHONPATH=python .venv-rocm/bin/python \ + benchmarks/bench_gguf_moe_kernels.py --iters 300 --check /tmp/moe-ref.pt +""" + +from __future__ import annotations + +import argparse +import hashlib +import os +import statistics +import sys +import time + +import torch +from freetoken.layers.activation import silu_and_mul +from freetoken.kernel.gguf import ggml_moe_a8_vec + +GGML_Q4_K = 12 +GGML_Q8_0 = 8 + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--slots", type=int, default=5873, help="resident expert slot count") + p.add_argument("--hidden", type=int, default=2048, help="model hidden size H") + p.add_argument("--moe-ic", type=int, default=512, help="moe_intermediate_size") + p.add_argument("--topk", type=int, default=8) + p.add_argument("--iters", type=int, default=300) + p.add_argument("--warmup", type=int, default=50) + p.add_argument("--dump", default=None, help="save outputs to this path") + p.add_argument("--check", default=None, help="byte-compare outputs against a saved dump") + args = p.parse_args(argv) + + H, I, topk, slots = args.hidden, args.moe_ic, args.topk, args.slots + n2 = 2 * I + dev = torch.device("cuda") + gen = torch.Generator(device="cpu").manual_seed(0) + + # gate_up bank: [slots, 2I, row_bytes(H, Q4_K)] (row bytes = H/256*144); + # down bank: [slots, H, row_bytes(I, Q8_0)] (I/32*34) — matches moe/fused_gguf.py. + bank_gu = torch.randint( + 0, 255, (slots, n2, H // 256 * 144), dtype=torch.uint8, generator=gen + ).to(dev) + bank_down = torch.randint( + 0, 255, (slots, H, I // 32 * 34), dtype=torch.uint8, generator=gen + ).to(dev) + x = torch.randn(1, H, generator=gen).to(dev).to(torch.bfloat16) + ids = torch.randint(0, slots, (1, topk), generator=gen).to(dev).int() + + def run() -> torch.Tensor: + gate_up = ggml_moe_a8_vec(x, bank_gu, ids, topk, GGML_Q4_K, n2, 1) + inter = silu_and_mul(gate_up.reshape(1 * topk, n2)) + return ggml_moe_a8_vec(inter, bank_down, ids, 1, GGML_Q8_0, H, topk) + + # Warm + reference for the bytes-equality gate (run() is deterministic for a + # fixed seed/config: no atomics, one CTA owns every output row). + out = run() + torch.cuda.synchronize() + torch.cuda.synchronize() + us_pair = bench(run, args.iters) + + sha = hashlib.sha256(out.view(torch.uint8).cpu().numpy().tobytes()).hexdigest()[:16] + mmv_y = os.environ.get("FREETOKEN_GGUF_MMV_Y", "1") + note = "" + if args.dump: + torch.save({"out": out.cpu()}, args.dump) + note = f" (saved reference to {args.dump})" + if args.check: + saved = torch.load(args.check, map_location="cpu", weights_only=True) + if not torch.equal(saved["out"], out.cpu()): + print(f"FAIL: MoE output bytes differ from the saved reference {args.check}", file=sys.stderr) + return 1 + note = f" bytes_equal_to_saved=True" + + print( + f"MMV_Y={mmv_y} slots={slots} H={H} I={I} topk={topk} : " + f"{us_pair:.1f} us / layer-pair out_sha={sha}{note}" + ) + return 0 + + +def bench(fn, iters: int) -> float: + for _ in range(50): + fn() + torch.cuda.synchronize() + ts = [] + for _ in range(iters): + t0 = time.perf_counter() + fn() + torch.cuda.synchronize() + ts.append(time.perf_counter() - t0) + return statistics.median(ts) * 1e6 + + +if __name__ == "__main__": + raise SystemExit(main()) \ No newline at end of file diff --git a/python/freetoken/kernel/csrc/gguf/ggml-common.h b/python/freetoken/kernel/csrc/gguf/ggml-common.h index 5822eb3cb..3f2eaa52b 100644 --- a/python/freetoken/kernel/csrc/gguf/ggml-common.h +++ b/python/freetoken/kernel/csrc/gguf/ggml-common.h @@ -8,7 +8,11 @@ #define CUDA_DEQUANTIZE_BLOCK_SIZE 256 #define CUDA_QUANTIZE_BLOCK_SIZE 256 #define GGML_CUDA_DMMV_X 32 +// Same overridable knob as ggml-common_hip.h (kept byte-identical): rows per warp +// for the MMVQ / MoE-vec launches; tune via -DGGML_CUDA_MMV_Y (FREETOKEN_GGUF_MMV_Y). +#ifndef GGML_CUDA_MMV_Y #define GGML_CUDA_MMV_Y 1 +#endif #if defined(USE_ROCM) // ROCm shim: the vendored GGUF launchers (moe.cuh/mmvq.cuh/...) take a CUDA-style diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 22796e4d3..f05b2def3 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -81,6 +81,8 @@ def _clear_stale_jit_lock(module_name: str) -> None: lock.unlink() except Exception: # noqa: BLE001 - hygiene must never break the build path pass + + @functools.cache def _module(): from torch.utils.cpp_extension import load @@ -109,6 +111,17 @@ def _module(): os.environ["CXX"] = cxx_path os.environ["CC"] = _c_compiler_for(cxx_path) + # Rows-per-warp for the MMVQ/MoE-vec launches (Inc 4, .plans/rocm-perf-parity). + # Overridable for tuning/AB; the JIT cache keys on the cflags, so a changed value + # rebuilds cleanly. CUDA-side the same -D reaches ggml-common.h's #ifndef guard. + mmv_y = os.getenv("FREETOKEN_GGUF_MMV_Y", "").strip() + if mmv_y: + if not mmv_y.isdigit() or int(mmv_y) not in (1, 2, 4, 8): + raise ValueError( + f"FREETOKEN_GGUF_MMV_Y={mmv_y!r}: expected one of 1, 2, 4, 8" + ) + extra_cuda_cflags.append(f"-DGGML_CUDA_MMV_Y={mmv_y}") + # 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. _clear_stale_jit_lock("freetoken_gguf_kernels") From 0680be76a62eba9badd481ba3c17f58c76969088 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Mon, 31 Aug 2026 00:42:26 -0300 Subject: [PATCH 13/17] feat(rocm-perf-parity): decode attention micro-bench + tuning sweep negative result (Inc 5) - incumbent BLOCK_N=32/warps=4 is best of {32/2:1325, 32/4:804, 64/4:1362, 64/8:960}us at kv=8192; SDPA gathered 3618us (4.5x worse) -> triton backend stays the ROCm decode default; reverted config byte-identical to HEAD --- benchmarks/bench_attn_decode.py | 116 ++++++++++++++++++++++++++++++++ 1 file changed, 116 insertions(+) create mode 100644 benchmarks/bench_attn_decode.py diff --git a/benchmarks/bench_attn_decode.py b/benchmarks/bench_attn_decode.py new file mode 100644 index 000000000..289aa1a9e --- /dev/null +++ b/benchmarks/bench_attn_decode.py @@ -0,0 +1,116 @@ +"""Micro-bench for the triton decode paged-attention kernel (Inc 5, rocm-perf-parity). + +Times `freetoken.kernel.triton.attention.decode_paged_attention` at the served +Qwen3.6-35B-A3B decode shape (bs=1, 32 q-heads x 256 dim GQA over 2 kv-heads, +~8k paged KV) and byte-compares each swept (BLOCK_N, num_warps) build variant +against the incumbent config via --dump/--check. + +Usage: + PYTHONPATH=python .venv-rocm/bin/python benchmarks/bench_attn_decode.py \ + --kv 8192 --dump /tmp/attn-ref.pt # incumbent config + # after editing the call-site constexprs: + PYTHONPATH=python .venv-rocm/bin/python benchmarks/bench_attn_decode.py \ + --kv 8192 --check /tmp/attn-ref.pt +""" + +from __future__ import annotations + +import argparse +import hashlib +import statistics +import sys +import time + +import torch +from freetoken.kernel.triton.attention import ( + _MIN_BLOCK_KV, + decode_paged_attention, +) + +MAX_KV_SPLITS = 16 + + +def build(kv_len: int, num_q_heads: int = 32, num_kv_heads: int = 2, head_dim: int = 256): + dev = torch.device("cuda") + gen = torch.Generator(device="cpu").manual_seed(11) + q = torch.randn(1, num_q_heads, head_dim, generator=gen).to(dev).to(torch.bfloat16) * 0.3 + # paged KV of kv_len physical slots; identity page mapping + k_cache = torch.randn(kv_len, num_kv_heads, head_dim, generator=gen).to(dev).to(torch.bfloat16) * 0.3 + v_cache = torch.randn(kv_len, num_kv_heads, head_dim, generator=gen).to(dev).to(torch.bfloat16) * 0.3 + indptr = torch.tensor([0, kv_len], dtype=torch.int32, device=dev) + indices = torch.arange(kv_len, dtype=torch.int32, device=dev) + q_positions = torch.tensor([kv_len - 1], dtype=torch.int32, device=dev) + num_kv_splits = torch.ones(1, dtype=torch.int32, device=dev) + logits = torch.zeros(1, num_q_heads, MAX_KV_SPLITS, head_dim, dtype=torch.float32, device=dev) + lse = torch.zeros(1, num_q_heads, MAX_KV_SPLITS, dtype=torch.float32, device=dev) + return q, k_cache, v_cache, indptr, indices, q_positions, logits, lse, num_kv_splits + + +def run(q, k, v, indptr, indices, q_pos, logits, lse, nks, sm_scale): + return decode_paged_attention( + q=q, + k_cache=k, + v_cache=v, + indptr=indptr, + indices=indices, + q_positions=q_pos, + attn_logits=logits, + attn_lse=lse, + num_kv_splits=nks, + max_kv_splits=MAX_KV_SPLITS, + sm_scale=sm_scale, + ) + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--kv", type=int, default=8192) + p.add_argument("--iters", type=int, default=200) + p.add_argument("--warmup", type=int, default=30) + p.add_argument("--dump", default=None) + p.add_argument("--check", default=None) + args = p.parse_args(argv) + + q, k, v, indptr, indices, q_pos, logits, lse, nks = build(args.kv) + sm_scale = 256**-0.5 + + out = run(q, k, v, indptr, indices, q_pos, logits, lse, nks, sm_scale) + torch.cuda.synchronize() + + def bench() -> float: + for _ in range(args.warmup): + run(q, k, v, indptr, indices, q_pos, logits, lse, nks, sm_scale) + torch.cuda.synchronize() + ts = [] + for _ in range(args.iters): + t0 = time.perf_counter() + run(q, k, v, indptr, indices, q_pos, logits, lse, nks, sm_scale) + torch.cuda.synchronize() + ts.append(time.perf_counter() - t0) + return statistics.median(ts) * 1e6 + + us = bench() + sha = hashlib.sha256(out.float().cpu().numpy().tobytes()).hexdigest()[:16] + note = "" + if args.dump: + torch.save({"out": out.float().cpu()}, args.dump) + note = f" (saved reference to {args.dump})" + if args.check: + saved = torch.load(args.check, map_location="cpu", weights_only=True) + diff = (saved["out"] - out.float().cpu()).abs().max().item() + # Different (BLOCK_N, warps) change the fp32 reduction order -> not bit-equal; + # the correctness gate is max-abs-diff in bf16-output terms (printed, not asserted). + print(f"max_abs_diff vs saved reference: {diff:.3e}") + note = f" max_abs_diff={diff:.3e}" + print( + f"kv={args.kv} q_heads=32 kv_heads=2 dim=256 : {us:.1f} us / call " + f"out_sha={sha}{note}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) + + +_ = _MIN_BLOCK_KV # re-exported constant; keeps the import meaningful for type checkers \ No newline at end of file From ef031bf8299255d3fdc6ad84eeb73b2f1da8322e Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Mon, 31 Aug 2026 01:01:46 -0300 Subject: [PATCH 14/17] feat(rocm-perf-parity): graph-capture variant gate + GRAPH CAPTURE NOW WORKS on gfx1100 (Inc 6) - probe variants default/rocblas/prewarm/rocblas+prewarm; prewarm (one GEMM before capture = hipBLASLt workspace pre-alloc) PASSES, all others fatal-900; env applied at worker spawn in launch.py; v2 cache with variant/env fields; capture bs[1,2,4] succeeds, bench median 44.97 tok/s (+30% vs baseline 34.69), output sha identical to kernel-launch --- python/freetoken/engine/graph.py | 9 +- python/freetoken/server/launch.py | 14 +++ python/freetoken/utils/graph_gate.py | 128 ++++++++++++++++++++------- 3 files changed, 114 insertions(+), 37 deletions(-) diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 276b741ef..24127c235 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -136,12 +136,15 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel # viable on this AMD card, skip graphs entirely so decode uses the kernel-launch # path (correct, just not graph-accelerated) rather than erroring mid-capture. from freetoken.utils.arch import is_rocm - from freetoken.utils.graph_gate import graph_capture_status + from freetoken.utils.graph_gate import graph_capture_status, run_graph_gate if is_rocm() and graph_capture_status() == "fail": + # Variant detail matters: the all-variants record is what closes the thread. + detail = run_graph_gate().get("detail", "") logger.info_rank0( - "AMD ROCm build: HIP graph capture gate FAILED on this device; " - "using the kernel-launch decode path (CUDA graphs disabled)." + "AMD ROCm build: HIP graph capture gate FAILED on this device (all " + f"capture variants: {detail[:160]}); using the kernel-launch decode " + "path (CUDA graphs disabled)." ) return None if self.max_graph_bs == 0: diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py index acef55175..9fcb44856 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -156,6 +156,20 @@ def start_subprocess() -> "BackendHandle": from .supervisor import BackendHandle mp.set_start_method("spawn", force=True) + # Graph-capture variant env (Inc 6, .plans/rocm-perf-parity): the gate probes + # rocBLAS/prewarm variants in subprocesses and may require e.g. + # TORCH_BLAS_PREFER_HIPBLASLT=0 for capture to be viable. Apply it to THIS + # (supervisor) process before the workers are spawned so the spawned engine + # worker inherits it from process start — before any torch import/GEMM; a late + # write at capture time may no-op depending on torch's BLAS-preference caching. + from freetoken.utils.graph_gate import graph_capture_env + + gate_env = graph_capture_env() + if gate_env: + os.environ.update(gate_env) + logger.info( + f"graph-capture gate env applied to workers: {sorted(gate_env)}" + ) detach = server_args.shell_mode # see _detach_process_group world_size = server_args.tp_info.size diff --git a/python/freetoken/utils/graph_gate.py b/python/freetoken/utils/graph_gate.py index 07089ebf2..b521e921b 100644 --- a/python/freetoken/utils/graph_gate.py +++ b/python/freetoken/utils/graph_gate.py @@ -50,11 +50,15 @@ def _device_name() -> str | None: return None -def probe_graph_capture() -> dict: - """Run the actual capture probe on the current device. Returns a dict: - ``{"device_kind": ..., "device": ..., "ok": bool, "detail": str}``. +def probe_graph_capture(variants: tuple[tuple[str, dict[str, str], str], ...] | None = None) -> dict: + """Run the capture probes on the current device across the variants. - The GEMM-capture attempt is run in a **fresh subprocess**: on some ROCm builds a + Returns ``{"device_kind", "device", "ok", "detail", "variant", "env"}`` where + ``variant`` is the first passing entry (or the last failing one) and ``env`` is + the environment the engine worker must be spawned with for capture to work + (empty for the default variant). See :data:`_GRAPH_VARIANTS`. + + Each variant runs in a **fresh subprocess**: on some ROCm builds a hipBLASLt/capture failure raises an uncatchable fatal HIP error (error 900) that aborts the whole process, so running it inline would crash the caller (and, worse, a decode path that attempted graph capture would die). A subprocess lets a fatal @@ -80,52 +84,90 @@ def probe_graph_capture() -> dict: "detail": f"torch unavailable: {type(exc).__name__}: {detail}", } - # The child probes both an elementwise op (capturable on both backends) and a GEMM - # (hipBLASLt on ROCm), which is what a real decode forward would run. The GEMM is - # the discriminating case: on this ROCm build it fatally aborts -> child exit != 0. + # Each variant probes both an elementwise op (capturable on both backends) and a + # GEMM (hipBLASLt on ROCm), which is what a real decode forward would run. The GEMM + # is the discriminating case: on this ROCm build it fatally aborts -> child exit != 0. import subprocess as _subprocess import sys as _sys - child = _subprocess.run( - [_sys.executable, "-c", _CAPTURE_CHILD], - capture_output=True, - text=True, - timeout=120, - ) - if child.returncode != 0: - detail = next( - (l.strip() for l in child.stderr.splitlines() if l.strip()), - f"graph capture subprocess aborted (rc={child.returncode})", + variants = variants if variants is not None else _GRAPH_VARIANTS + last_fail: dict | None = None + for name, extra_env, mode in variants: + env = {**os.environ, **extra_env} + child = _subprocess.run( + [_sys.executable, "-c", _CAPTURE_CHILD, mode], + capture_output=True, + text=True, + timeout=120, + env=env, ) - return { - "device_kind": _device_kind(), - "device": device, - "ok": False, - "detail": f"fatal during capture: {detail[:240]}", - } + if child.returncode != 0: + detail = next( + (l.strip() for l in child.stderr.splitlines() if l.strip()), + f"graph capture subprocess aborted (rc={child.returncode})", + ) + last_fail = { + "device_kind": _device_kind(), + "device": device, + "ok": False, + "detail": f"fatal during capture: {detail[:240]}", + } + continue + try: + data = json.loads(child.stdout) + except Exception: + last_fail = { + "device_kind": _device_kind(), + "device": device, + "ok": False, + "detail": f"unparseable probe output: {child.stdout[:120]}", + } + continue + data.setdefault("device_kind", _device_kind()) + data.setdefault("device", device) + data["variant"] = name + data["env"] = dict(extra_env) + if data.get("ok"): + return data + last_fail = data + assert last_fail is not None + return last_fail + + +@lru_cache(maxsize=1) +def graph_capture_env() -> dict[str, str]: + """Env vars the engine worker must be spawned with for graph capture to work. + + Empty when the default variant passes (or the gate failed / no device); otherwise + the winning variant's extra env (e.g. ``TORCH_BLAS_PREFER_HIPBLASLT=0``). Applied + by the server supervisor BEFORE the engine worker is spawned — never late at + capture time, where a late env write may no-op. + """ try: - data = json.loads(child.stdout) + result = run_graph_gate() + env = result.get("env") or {} + return env if result.get("ok") else {} except Exception: - return { - "device_kind": _device_kind(), - "device": device, - "ok": False, - "detail": f"unparseable probe output: {child.stdout[:120]}", - } - data.setdefault("device_kind", _device_kind()) - data.setdefault("device", device) - return data + return {} #: Child body for the graph-capture probe (see :func:`probe_graph_capture`). Prints a #: JSON line ``{"ok": true/false, "detail": ...}`` and exits nonzero on a fatal abort. +#: Modes: "plain" = capture a GEMM with default BLAS/library state; "prewarm" = run the +#: same-shaped GEMM once before capture (hipBLASLt lazily allocates its workspace on +#: the first GEMM, and a capture that includes that allocation is what aborts on some +#: ROCm builds). _CAPTURE_CHILD = r""" import json, sys +mode = sys.argv[1] if len(sys.argv) > 1 else "plain" import torch try: torch.cuda.synchronize() s = torch.cuda.Stream() x = torch.randn(8, 8, device='cuda') + if mode == "prewarm": + a0 = torch.randn(64, 64, device='cuda') + torch.mm(a0, a0) # pre-allocate the BLAS workspace OUTSIDE any capture # elementwise (capturable) first, then a GEMM (hipBLASLt on ROCm) with torch.cuda.stream(s): g = torch.cuda.CUDAGraph() @@ -139,12 +181,26 @@ def probe_graph_capture() -> dict: torch.mm(a, a) s.synchronize() torch.cuda.synchronize() - print(json.dumps({"ok": True, "detail": "elementwise+GEMM capture/replay succeeded"})) + print(json.dumps({"ok": True, "detail": f"{mode}: elementwise+GEMM capture/replay succeeded"})) except Exception as e: print(json.dumps({"ok": False, "detail": f"{type(e).__name__}: {str(e)[:160]}"})) """ +# Capture-probe variants, tried in order (Inc 6, .plans/rocm-perf-parity). Each entry +# is (name, extra_env, child_mode). "rocblas" forces the rocBLAS GEMM path via +# TORCH_BLAS_PREFER_HIPBLASLT=0 — hipBLASLt is the capture abort's usual suspect on +# gfx1100. "prewarm" pre-runs the GEMM outside capture (workspace pre-alloc). The +# winning variant's env must be applied to the engine worker process at SPAWN time +# (server/launch.py) — never late at capture time, where a late env write may no-op. +_GRAPH_VARIANTS: tuple[tuple[str, dict[str, str], str], ...] = ( + ("default", {}, "plain"), + ("rocblas", {"TORCH_BLAS_PREFER_HIPBLASLT": "0"}, "plain"), + ("prewarm", {}, "prewarm"), + ("rocblas_prewarm", {"TORCH_BLAS_PREFER_HIPBLASLT": "0"}, "prewarm"), +) + + def _load_cached() -> dict | None: try: with open(_cache_path()) as f: @@ -152,6 +208,10 @@ def _load_cached() -> dict | None: if ( data.get("device_kind") == _device_kind() and data.get("device") == _device_name() + # v2 cache format: carries the winning "variant"/"env". An old all-FAIL + # record (no variant) must not mask a possible variant PASS after a + # ROCm/driver upgrade. + and "variant" in data ): return data except Exception: From 2f02e550febc5d763b01adca43c3bebb0efe4254 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Mon, 31 Aug 2026 01:16:35 -0300 Subject: [PATCH 15/17] docs(install-amd): ROCm performance section (Inc 8) - baseline/median table (34.69 -> 45.09 tok/s), enabled features, gate cache + escape hatches (FREETOKEN_TORCH_PROFILE, FREETOKEN_GGUF_MMV_Y), negative tuning results recorded --- docs/install-amd.md | 42 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/docs/install-amd.md b/docs/install-amd.md index ede4ab80f..767632795 100644 --- a/docs/install-amd.md +++ b/docs/install-amd.md @@ -81,3 +81,45 @@ The runtime refuses to pair a `+rocm` cache with a `+cu130` runtime (and vice ve * FP8 / NVFP4-class formats: BF16 / MXFP4 / GGUF are the supported AMD matrix; performance parity vs CUDA is not guaranteed for NVFP4-class formats. * Windows AMD is not yet supported (WDDM zero-copy semantics differ). + +## Performance (gfx1100 / RX 7900 XTX, as of rocm-perf-parity) + +Measured with `benchmarks/bench_decode_moe.py` on Qwen3.6-35B-A3B GGUF (Q4_K_M), +3-run medians, same prompt/sampling protocol throughout (see +`.plans/rocm-perf-parity/` for artifacts and the stage-time profiling notes): + +| state | decode tok/s (median) | +| --- | --- | +| pre-plan baseline (kernel-launch decode, pure-torch router) | 34.69 | +| + in-repo fused triton router | 37.67 | +| + CUDA-graph decode capture (`prewarm` variant) | **45.09** | + +What is enabled on AMD now: + +- **Fused triton router**: `fused_topk` routes ROCm to the in-repo + `kernel/triton/moe_router.fused_topk_softmax` (no `triton_kernels` install needed). + The pure-torch fallback remains for CPU/Windows and is flagged by the + `pure-torch router fallback` log line. +- **CUDA-graph decode**: the capture gate (`utils/graph_gate.py`) now probes capture + *variants* — `default`, `rocblas` (`TORCH_BLAS_PREFER_HIPBLASLT=0`), + `prewarm` (one GEMM before capture; the winner — hipBLASLt's lazy workspace + allocation inside capture was the abort), and `rocblas+prewarm`. The winning + variant's env is applied to the engine worker at spawn. The gate result is cached + at `~/.cache/freetoken/freetoken_graph_gate.json`; delete it to re-probe after a + driver/ROCm upgrade (the cache format invalidates itself automatically). +- Known-negative tuning results (documented, do not redo blindly): + `FREETOKEN_GGUF_MMV_Y` (rows/warp for the GGUF MMVQ launches): 1/2/4/8 within run + noise AND byte-identical outputs — keep 1. Decode attention kernel + (`BLOCK_N=32/warps=4`): best of the swept grid; torch-SDPA gathered attention is + 4.5× slower on gfx1100. + +Escape hatches for runtime bisection: + +- `FREETOKEN_TORCH_PROFILE=::` — torch.profiler trace of a + decode window (`scripts/profile-rocm-decode.sh` wraps this end-to-end). +- `FREETOKEN_GGUF_MMV_Y` — GGUF MMVQ rows-per-warp build knob (see above; JIT cache + is keyed on the value). +- `--no-graph` on `benchmarks/bench_decode_moe.py`, or `--cuda-graph-max-bs 0`, to + fall back to kernel-launch decode; `scripts/serve-qwen-moe.sh stop` clears a wedged + server (a `kill -9` during a JIT rebuild leaves a stale torch-extension lock that + `kernel/gguf.py` now clears automatically once older than 3 h). From 7544d66506ee59765daf1d7e21fe9c24dff4b989 Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Mon, 31 Aug 2026 10:49:22 -0300 Subject: [PATCH 16/17] feat(serve): FT_SERVED_MODEL knob -> --served-model-name (Copilot-fork custom endpoints validate the configured model id against /v1/models; advertise the exact id the client declares, e.g. FT_SERVED_MODEL=qwen3.6) --- scripts/serve-qwen-moe.sh | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/scripts/serve-qwen-moe.sh b/scripts/serve-qwen-moe.sh index c31cd2865..b59df731a 100755 --- a/scripts/serve-qwen-moe.sh +++ b/scripts/serve-qwen-moe.sh @@ -73,6 +73,11 @@ KV_TOKENS="${FT_KV_TOKENS:-131072}" MAX_OUTPUT="${FT_MAX_OUTPUT:-65536}" MOE_CACHE="${FT_MOE_CACHE:-auto}" LOG="${FT_LOG:-/tmp/serve_qwen_moe.log}" +# Advertised model id in /v1/models (--served-model-name). Copilot-fork custom +# endpoints validate the configured model id against this list, so set it to the +# exact id the client config declares (e.g. FT_SERVED_MODEL=qwen3.6). Empty = +# server default (the GGUF filename). +SERVED_MODEL="${FT_SERVED_MODEL:-}" die() { echo "ERROR: $*" >&2; exit 1; } @@ -104,6 +109,7 @@ start() { # One arg per array element; expanded once, single line, no continuations. local -a SERVE_ARGS=( "--model" "$MODEL" + $( [ -n "$SERVED_MODEL" ] && echo "--served-model-name" "$SERVED_MODEL" ) "--moe-backend" "$MOE_BACKEND" "--cache-type" "$CACHE_TYPE" $( [ "$MOE_STATS" = 1 ] && echo "--moe-collect-stats" ) From 93819f7f45db59017c22ead1b2a78125c09650aa Mon Sep 17 00:00:00 2001 From: Samuel Ishida Date: Wed, 2 Sep 2026 01:13:18 -0300 Subject: [PATCH 17/17] feat(rocm): complete Qwen GGUF decode parity path --- .agents/learnings/index.yml | 9 + .../learnings/llamacpp-surpass-rocm-gaps.md | 35 + .agents/learnings/qwen-moe-rocm-base-speed.md | 36 + .../rocm-ollama-gap-implementation.md | 51 + .github/workflows/unit-nvidia.yml | 36 +- .github/workflows/unit-rocm.yml | 167 +- benchmarks/README.md | 37 + benchmarks/__init__.py | 1 + benchmarks/bench_decode_moe.py | 881 ++++++- benchmarks/bench_decode_ollama.py | 425 ++++ benchmarks/bench_decode_replay.py | 927 ++++++++ benchmarks/bench_gguf_linear.py | 127 + benchmarks/bench_gguf_moe_kernels.py | 258 ++- benchmarks/bench_llama_cpp_hip.py | 742 ++++++ benchmarks/check_decode_gate.py | 693 ++++++ benchmarks/probe_qwen_moe_base.py | 263 +++ benchmarks/profile_decode_rocm.py | 98 + docs/install-amd.md | 94 +- python/freetoken/attention/triton.py | 28 + python/freetoken/engine/__init__.py | 6 +- python/freetoken/engine/config.py | 26 + python/freetoken/engine/engine.py | 264 ++- python/freetoken/engine/graph.py | 174 +- python/freetoken/engine/resident_budget.py | 236 ++ python/freetoken/engine/sample.py | 70 +- python/freetoken/kernel/__init__.py | 2 + .../kernel/csrc/gguf/dequantize_hip.cuh | 585 +++++ .../kernel/csrc/gguf/ggml-common_hip.h | 1034 +++++++++ .../kernel/csrc/gguf/gguf_b10434_kernel.cu | 88 + .../freetoken/kernel/csrc/gguf/gguf_kernel.cu | 184 +- .../kernel/csrc/gguf/gguf_moe_gfx1100.cu | 546 +++++ .../kernel/csrc/gguf/llama_b10434/PATCHES.md | 8 + .../kernel/csrc/gguf/llama_b10434/README.md | 19 + .../csrc/gguf/llama_b10434/quantize_q8_1.h | 13 + python/freetoken/kernel/csrc/gguf/mmq_hip.cuh | 883 +++++++ .../freetoken/kernel/csrc/gguf/mmvq_hip.cuh | 354 +++ python/freetoken/kernel/csrc/gguf/moe_hip.cuh | 1381 +++++++++++ python/freetoken/kernel/csrc/gguf/moe_vec.cuh | 76 + .../kernel/csrc/gguf/moe_vec_gfx1100.cuh | 140 ++ .../kernel/csrc/gguf/moe_vec_gfx1100_hip.cuh | 148 ++ .../kernel/csrc/gguf/moe_vec_hip.cuh | 478 ++++ .../kernel/csrc/gguf/vecdotq_hip.cuh | 2039 +++++++++++++++++ python/freetoken/kernel/gguf.py | 637 ++++- python/freetoken/kernel/moe_impl.py | 26 + python/freetoken/kernel/triton/attention.py | 133 ++ python/freetoken/kernel/triton/fused_moe.py | 46 + python/freetoken/kernel/triton/q8_kv.py | 151 ++ python/freetoken/kvcache/__init__.py | 5 + python/freetoken/kvcache/base.py | 133 +- python/freetoken/kvcache/cache_status.py | 17 + python/freetoken/kvcache/mha_pool.py | 143 +- python/freetoken/layers/gguf.py | 51 +- python/freetoken/layers/moe.py | 111 +- python/freetoken/models/config.py | 3 + python/freetoken/models/gguf/dequant.py | 45 +- .../freetoken/models/qwen3_5_moe/__init__.py | 2 + python/freetoken/models/qwen3_5_moe/gguf.py | 178 +- python/freetoken/models/qwen3_5_moe/moe.py | 11 +- python/freetoken/models/weight.py | 12 + python/freetoken/moe/expert_banks.py | 13 +- python/freetoken/moe/fused_gguf.py | 447 +++- python/freetoken/moe/offload_cache.py | 10 +- python/freetoken/scheduler/scheduler.py | 281 ++- python/freetoken/scheduler/status.py | 7 + python/freetoken/server/api_server.py | 6 + python/freetoken/server/args.py | 24 +- python/freetoken/server/launch.py | 14 +- python/freetoken/server/stats.py | 1 + python/freetoken/tokenizer/detokenize.py | 8 + python/freetoken/utils/graph_gate.py | 120 +- python/freetoken/utils/step_profiler.py | 311 ++- scripts/build-llama-reference.sh | 94 + scripts/profile-rocm-decode.sh | 178 +- scripts/serve-qwen-moe.sh | 52 +- tests/benchmarks/__init__.py | 0 tests/benchmarks/test_paired_stats.py | 30 + tests/benchmarks/test_replay_manifest.py | 267 +++ tests/benchmarks/test_rocm_trace.py | 78 + tests/engine/test_cache_budget.py | 6 +- tests/engine/test_gguf_resident.py | 12 + tests/engine/test_graph_blas_policy.py | 114 + tests/engine/test_graph_capture_env.py | 24 + tests/engine/test_resident_budget.py | 28 + tests/engine/test_sample.py | 44 + tests/engine/test_sample_capture.py | 74 + tests/kernels/test_attention_q8.py | 68 + tests/kernels/test_gguf_dispatch.py | 122 + tests/kernels/test_gguf_linear.py | 105 + tests/kernels/test_gguf_moe.py | 271 +++ tests/kvcache/test_kv_storage.py | 40 + tests/kvcache/test_mha_q8.py | 79 + tests/models/test_qwen35moe_gguf_deint.py | 62 + tests/models/test_qwen35moe_moe.py | 113 + tests/moe/test_offload.py | 4 +- tests/scheduler/test_decode_handoff.py | 43 + tests/utils/test_decode_benchmark_metadata.py | 138 ++ tests/utils/test_decode_gate.py | 210 ++ tests/utils/test_step_profiler.py | 25 + 98 files changed, 18465 insertions(+), 424 deletions(-) create mode 100644 .agents/learnings/index.yml create mode 100644 .agents/learnings/llamacpp-surpass-rocm-gaps.md create mode 100644 .agents/learnings/qwen-moe-rocm-base-speed.md create mode 100644 .agents/learnings/rocm-ollama-gap-implementation.md create mode 100644 benchmarks/__init__.py create mode 100644 benchmarks/bench_decode_ollama.py create mode 100644 benchmarks/bench_decode_replay.py create mode 100644 benchmarks/bench_gguf_linear.py create mode 100644 benchmarks/bench_llama_cpp_hip.py create mode 100644 benchmarks/check_decode_gate.py create mode 100644 benchmarks/probe_qwen_moe_base.py create mode 100644 benchmarks/profile_decode_rocm.py create mode 100644 python/freetoken/engine/resident_budget.py create mode 100644 python/freetoken/kernel/csrc/gguf/dequantize_hip.cuh create mode 100644 python/freetoken/kernel/csrc/gguf/ggml-common_hip.h create mode 100644 python/freetoken/kernel/csrc/gguf/gguf_b10434_kernel.cu create mode 100644 python/freetoken/kernel/csrc/gguf/gguf_moe_gfx1100.cu create mode 100644 python/freetoken/kernel/csrc/gguf/llama_b10434/PATCHES.md create mode 100644 python/freetoken/kernel/csrc/gguf/llama_b10434/README.md create mode 100644 python/freetoken/kernel/csrc/gguf/llama_b10434/quantize_q8_1.h create mode 100644 python/freetoken/kernel/csrc/gguf/mmq_hip.cuh create mode 100644 python/freetoken/kernel/csrc/gguf/mmvq_hip.cuh create mode 100644 python/freetoken/kernel/csrc/gguf/moe_hip.cuh create mode 100644 python/freetoken/kernel/csrc/gguf/moe_vec_gfx1100.cuh create mode 100644 python/freetoken/kernel/csrc/gguf/moe_vec_gfx1100_hip.cuh create mode 100644 python/freetoken/kernel/csrc/gguf/moe_vec_hip.cuh create mode 100644 python/freetoken/kernel/csrc/gguf/vecdotq_hip.cuh create mode 100644 python/freetoken/kernel/triton/q8_kv.py create mode 100755 scripts/build-llama-reference.sh create mode 100644 tests/benchmarks/__init__.py create mode 100644 tests/benchmarks/test_paired_stats.py create mode 100644 tests/benchmarks/test_replay_manifest.py create mode 100644 tests/benchmarks/test_rocm_trace.py create mode 100644 tests/engine/test_gguf_resident.py create mode 100644 tests/engine/test_graph_blas_policy.py create mode 100644 tests/engine/test_graph_capture_env.py create mode 100644 tests/engine/test_resident_budget.py create mode 100644 tests/engine/test_sample.py create mode 100644 tests/engine/test_sample_capture.py create mode 100644 tests/kernels/test_attention_q8.py create mode 100644 tests/kernels/test_gguf_dispatch.py create mode 100644 tests/kernels/test_gguf_linear.py create mode 100644 tests/kernels/test_gguf_moe.py create mode 100644 tests/kvcache/test_kv_storage.py create mode 100644 tests/kvcache/test_mha_q8.py create mode 100644 tests/models/test_qwen35moe_moe.py create mode 100644 tests/scheduler/test_decode_handoff.py create mode 100644 tests/utils/test_decode_benchmark_metadata.py create mode 100644 tests/utils/test_decode_gate.py create mode 100644 tests/utils/test_step_profiler.py diff --git a/.agents/learnings/index.yml b/.agents/learnings/index.yml new file mode 100644 index 000000000..30d23345e --- /dev/null +++ b/.agents/learnings/index.yml @@ -0,0 +1,9 @@ +qwen-moe-rocm-base-speed: + file: qwen-moe-rocm-base-speed.md + description: Read before Qwen MoE ROCm base-speed work; records negative kernel/cache/dense results and validation gates. +llamacpp-surpass-rocm-gaps: + file: llamacpp-surpass-rocm-gaps.md + description: Read before follow-up work on native GGUF residency, fused MoE kernels, and ROCm cadence gaps versus llama.cpp. +rocm-ollama-gap-implementation: + file: rocm-ollama-gap-implementation.md + description: Read before q8/GGUF ROCm validation or promotion; records fail-closed dispatch, replay timing semantics, and separate Gate A/B evidence. diff --git a/.agents/learnings/llamacpp-surpass-rocm-gaps.md b/.agents/learnings/llamacpp-surpass-rocm-gaps.md new file mode 100644 index 000000000..b33f08566 --- /dev/null +++ b/.agents/learnings/llamacpp-surpass-rocm-gaps.md @@ -0,0 +1,35 @@ +# llama.cpp Surpass ROCm Gaps + +## Context + +2026-08-31: Qwen3.6-35B-A3B GGUF decode work on RX 7900 XTX/gfx1100 added +native Q4_K/Q5_K/Q6_K paths, residency planning, execution gates, and traces. + +## Hardest decision + +Keep native mixed-K support fail-closed and retain offload defaults: the full +resident budget does not fit this 24-GiB card, while native Q5/Q6 offload only +reached 58.98 tok/s API / 61.19 tok/s scheduler. + +## Alternatives rejected + +- Disable safety reserve to force `fused` residency — risks late OOM and invalid + benchmark evidence. +- Enable grouped GGUF prefill by default — real gfx1100 execution failed during + launch; synthetic ABI success was insufficient. +- Promote gfx1100 rotated-wave or Q5/Q6 path from microbench alone — candidate + and end-to-end evidence did not beat the incumbent or llama.cpp ROCm. +- Treat Torch profiler CPU ranges as additive wall time — synchronization and + overlap inflate attribution; no scheduler edit was accepted. + +## Least confident + +Exact remaining attribution is unresolved because rocprofv3 attach was blocked +by host ptrace policy. Native executor cadence, full resident capacity with a +different context/budget, and a genuinely fused/grouped GGUF kernel remain open. + +## Reuse + +Read before further Qwen ROCm speed work. Preserve model/fixture hashes, +resident execution evidence, zero fetch/remap gates, graph state, exact output +count, and paired same-file llama.cpp measurements. diff --git a/.agents/learnings/qwen-moe-rocm-base-speed.md b/.agents/learnings/qwen-moe-rocm-base-speed.md new file mode 100644 index 000000000..2c618754a --- /dev/null +++ b/.agents/learnings/qwen-moe-rocm-base-speed.md @@ -0,0 +1,36 @@ +# Qwen MoE ROCm Base-Speed Optimization + +## Context + +2026-08-31: Qwen3.6-35B-A3B GGUF Q4_K_M base decode on RX 7900 XTX/gfx1100. +Plan compared FreeToken against Ollama with MTP excluded. + +## Hardest decision + +Keep legacy GGUF MMVQ as default after proving the gfx1100 rotated-wave +candidate 9.7% slower, and close cache/dense/handoff branches when matched +measurements showed no throughput win. Correctness and explicit target failure +outweighed speculative promotion. + +## Alternatives rejected + +- AMD `sdot4` intrinsic rewrite — ROCm target rejects it without `dot1-insts`, + and candidate disassembly had no dot4 instruction. +- Smaller or fixed MoE cache — 4,352 slots fell to 36.875 tok/s; 8,000 slots + did not beat auto; warmed runs had zero fetches. +- Dense MMQ or lm_head rewrite — native MMVQ won every exact bs=1 dense case. +- MTP/speculative path — outside requested base-speed scope. + +## Least confident + +The remaining roughly 13 tok/s gap is attributed to ROCm base execution cadence +and GGUF runtime cost, but profiler CPU totals include overlap attribution and +are not directly additive to API wall time. Revisit with a lower-overhead +timeline or newer ROCm/compiler before changing dispatch again. + +## Reuse + +Read before future Qwen ROCm speed work. Preserve benchmark contract, exact +completion validation, graph/finite-logit gates, cache hit/fetch telemetry, and +forced-only candidate dispatch in `benchmarks/`, `python/freetoken/kernel/gguf.py`, +and `.plans/qwen-moe-speed/`. diff --git a/.agents/learnings/rocm-ollama-gap-implementation.md b/.agents/learnings/rocm-ollama-gap-implementation.md new file mode 100644 index 000000000..8c63008ba --- /dev/null +++ b/.agents/learnings/rocm-ollama-gap-implementation.md @@ -0,0 +1,51 @@ +# ROCm/Ollama Gap Implementation + +## Context + +2026-09-01 implementation of `.plans/rocm-ollama-gap` across KV storage, GGUF dispatch, +profiling, fused MoE, and promotion harnesses. + +## Hardest decision + +Keep q8 KV and gfx1100 GGUF paths explicit and fail-closed: requested CLI settings never count as +observed execution, and candidate code never silently falls back in forced mode. + +## Alternatives rejected + +- Running Ollama/ROCm benchmarks during implementation — user reserved GPU validation for manual review. +- Reusing model dtype for q8 accounting — allocation, packed scales, attention views, and metadata must share one descriptor. +- Combining sampled and teacher-forced results — Gate A absolute throughput and Gate B q8 replay measure different claims. + +## Least confident + +HIP/CUDA compilation and real-model replay remain unverified. The b10434 bridge currently enforces +512-column activation alignment and needs manual ABI, numerical, graph-capture, and performance proof. + +## Reuse + +Read before running `.plans/rocm-ollama-gap` evidence: start with `fixture-manifest.json`, enable +dispatch tracing for per-op records, run q8 with `--kv-type q8_0`, and report Gate A/B separately. + +## Validation closure — 2026-09-01 + +### Hardest decision + +Keep legacy GGUF MMVQ as default after exact 512-token q8 greedy output parity but no speed gain: +candidate `69.040` versus legacy `69.194 tok/s`; fused gate/up also produced degenerate output. + +### Alternatives rejected + +- Promoting gfx1100 MMVQ or fused gate/up — candidate was slower, and fused model output repeated + `We.` despite direct-kernel tolerance. +- Calling b10434 `predicted_ms` decode timing — forced-token server requests bill their forward pass + as one-token `prompt_eval_ms`, so replay accounting now uses that field. + +### Least confident + +Cross-runtime route parity remains unproven because b10434 server has no route instrumentation; +short rocprof q8 graph startup also hit duplicate physical Q8 destinations before request. + +### Reuse + +Read `.plans/rocm-ollama-gap/notes-results.md` before any promotion or kernel follow-up; preserve +legacy default, q8 opt-in, route-parity requirement, and the warm-offload measured bound. diff --git a/.github/workflows/unit-nvidia.yml b/.github/workflows/unit-nvidia.yml index d742c9020..64894ed07 100644 --- a/.github/workflows/unit-nvidia.yml +++ b/.github/workflows/unit-nvidia.yml @@ -59,11 +59,45 @@ jobs: uv pip install --system torch --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION} uv pip install --system -e . --no-build-isolation + - name: Compile/import legacy GGUF CUDA module + run: | + PYTHONPATH=python python - <<'PY' + import torch + from freetoken.kernel import gguf + + assert torch.version.hip is None, "NVIDIA job resolved a ROCm torch wheel" + module = gguf._module() + for symbol in ("ggml_dequantize", "ggml_mul_mat_vec_a8", "ggml_moe_a8_vec"): + assert hasattr(module, symbol), symbol + print("legacy GGUF CUDA module imported:", module.__name__) + PY + + - name: Compile/import pinned single-token GGUF ABI + run: | + PYTHONPATH=python python - <<'PY' + from pathlib import Path + import torch + from torch.utils.cpp_extension import load + + root = Path("python/freetoken/kernel/csrc/gguf") + module = load( + name="freetoken_gguf_b10434_ci", + sources=[str(root / "gguf_moe_gfx1100.cu"), str(root / "gguf_b10434_kernel.cu")], + extra_cuda_cflags=["-O3", "-DFREETOKEN_GGUF_NO_PYBIND=1"], + verbose=False, + ) + for symbol in ("mmvq_bs1_workspace_bytes", "mmvq_bs1", "ggml_moe_mmvq_id"): + assert hasattr(module, symbol), symbol + assert module.mmvq_bs1_workspace_bytes(512, 16, 8) % 256 == 0 + print("pinned b10434 ABI imported:", module.__name__) + PY + - name: Run unit suite (GPU tests self-skip) run: | + set -o pipefail python -m pytest tests/ -m "not slow" -q -x \ --ignore=tests/e2e --ignore=tests/dsv4 2>&1 | tail -40 - name: Summarize if: always() - run: echo "CUDA-compile + unit gate finished; see the step log for failures." \ No newline at end of file + run: echo "CUDA-compile + unit gate finished; see the step log for failures." diff --git a/.github/workflows/unit-rocm.yml b/.github/workflows/unit-rocm.yml index cb92e9223..a4a379dd1 100644 --- a/.github/workflows/unit-rocm.yml +++ b/.github/workflows/unit-rocm.yml @@ -16,6 +16,57 @@ on: description: "Also run the bench_decode_moe smoke after the suite" type: boolean default: false + run_final_gate: + description: "Run the non-optional ten-run promotion gate" + type: boolean + default: false + final_model_path: + description: "Absolute path to exact GGUF on the ROCm runner" + type: string + default: "/home/smk/models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf" + final_model_sha256: + description: "Full SHA-256 of final_model_path" + type: string + default: "" + final_aime_path: + description: "Absolute path to pinned AIME JSONL" + type: string + default: "/tmp/qwen-moe-speed-fixture/test.jsonl" + final_aime_revision: + description: "Pinned AIME dataset revision" + type: string + default: "563bb8404243c5f09de6ec262f2db674fe5bce9b" + final_aime_sha256: + description: "Full SHA-256 of final_aime_path" + type: string + default: "b4e273c02d3e7fe1b74b59eae768fc8230bfb0f79539890cb56f4361caac0331" + final_reference_jsonl: + description: "Optional matched llama.cpp/Ollama JSONL; empty forces gate=false" + type: string + default: "" + final_backend: + description: "MoE backend for final gate" + type: choice + options: [offload, cpu, hybrid] + default: offload + final_candidate_mode: + description: "Explicit GGUF MoE mode for final gate" + type: choice + options: [legacy, auto, gfx1100] + default: legacy + final_graph_mode: + description: "Decode graph mode for final gate" + type: choice + options: [replay, eager] + default: replay + final_decode: + description: "Exact generated token count" + type: string + default: "512" + final_probe_decode: + description: "Finite-logit/parity probe decode count" + type: string + default: "8" permissions: contents: read @@ -38,10 +89,38 @@ jobs: .venv-rocm/bin/python -c "import torch; print(torch.__version__, torch.version.hip)" rocm-smi | head -8 - - name: Unit suite + - name: Compile and smoke-test gfx1100 GGUF MoE candidate + run: | + set -o pipefail + FREETOKEN_GGUF_MOE_IMPL=gfx1100 PYTHONPATH=python \ + .venv-rocm/bin/python benchmarks/bench_gguf_moe_kernels.py \ + --impl gfx1100 --iters 3 --warmup 1 + PYTHONPATH=python .venv-rocm/bin/python -m pytest \ + tests/kernels/test_gguf_moe.py -m slow -q + + - name: Focused parity suite + run: | + PYTHONPATH=python .venv-rocm/bin/python -m pytest -q \ + tests/engine/test_sample_capture.py \ + tests/engine/test_sample.py \ + tests/engine/test_graph_blas_policy.py \ + tests/engine/test_graph_capture_env.py \ + tests/kernels/test_gguf_dispatch.py \ + tests/kernels/test_gguf_linear.py \ + tests/scheduler/test_decode_handoff.py \ + tests/utils/test_decode_benchmark_metadata.py \ + tests/utils/test_decode_gate.py + + - name: Broad suite diagnostic run: | + set -o pipefail PYTHONPATH=python .venv-rocm/bin/python -m pytest tests/ -m "not slow" -q \ - 2>&1 | tail -60 + > /tmp/ft-rocm-broad-suite.log 2>&1 + rc=$? + tail -80 /tmp/ft-rocm-broad-suite.log + echo "broad_suite_exit=$rc" >> /tmp/ft-rocm-broad-suite.log + exit "$rc" + continue-on-error: true - name: Optional bench smoke if: ${{ inputs.run_bench }} @@ -49,4 +128,86 @@ jobs: PYTHONPATH=python .venv-rocm/bin/python benchmarks/bench_decode_moe.py \ --model "${FREETOKEN_BENCH_MODEL:-/home/smk/models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf}" \ --backend offload --json /tmp/ft-rocm-ci-bench.jsonl 2>&1 | tail -20 - continue-on-error: true \ No newline at end of file + continue-on-error: true + + rocm-final-gate: + if: ${{ inputs.run_final_gate }} + needs: rocm-unit + runs-on: [self-hosted, rocm] + timeout-minutes: 360 + env: + FINAL_MODEL_PATH: ${{ inputs.final_model_path }} + FINAL_MODEL_SHA256: ${{ inputs.final_model_sha256 }} + FINAL_AIME_PATH: ${{ inputs.final_aime_path }} + FINAL_AIME_REVISION: ${{ inputs.final_aime_revision }} + FINAL_AIME_SHA256: ${{ inputs.final_aime_sha256 }} + FINAL_REFERENCE_JSONL: ${{ inputs.final_reference_jsonl }} + FINAL_BACKEND: ${{ inputs.final_backend }} + FINAL_CANDIDATE_MODE: ${{ inputs.final_candidate_mode }} + FINAL_GRAPH_MODE: ${{ inputs.final_graph_mode }} + FINAL_DECODE: ${{ inputs.final_decode }} + FINAL_PROBE_DECODE: ${{ inputs.final_probe_decode }} + steps: + - uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Verify pinned final inputs + run: | + test -n "$FINAL_MODEL_SHA256" + test -f "$FINAL_MODEL_PATH" + test "$(sha256sum "$FINAL_MODEL_PATH" | awk '{print $1}')" = "$FINAL_MODEL_SHA256" + test -f "$FINAL_AIME_PATH" + test "$(sha256sum "$FINAL_AIME_PATH" | awk '{print $1}')" = "$FINAL_AIME_SHA256" + + - name: Run exact ten-run FreeToken gate + run: | + graph_arg=() + if [ "$FINAL_GRAPH_MODE" = eager ]; then + graph_arg=(--no-graph) + fi + FREETOKEN_GGUF_MOE_IMPL="$FINAL_CANDIDATE_MODE" \ + PYTHONPATH=python .venv-rocm/bin/python benchmarks/bench_decode_moe.py \ + --model "$FINAL_MODEL_PATH" \ + --aime "$FINAL_AIME_PATH" \ + --aime-revision "$FINAL_AIME_REVISION" \ + --aime-sha256 "$FINAL_AIME_SHA256" \ + --backend "$FINAL_BACKEND" \ + --attention-backend triton \ + --decode "$FINAL_DECODE" --repeats 10 \ + "${graph_arg[@]}" --json /tmp/ft-rocm-final.jsonl + + - name: Run finite-logit and graph parity probe + run: | + FREETOKEN_GGUF_MOE_IMPL="$FINAL_CANDIDATE_MODE" \ + PYTHONPATH=python .venv-rocm/bin/python benchmarks/probe_qwen_moe_base.py \ + --model "$FINAL_MODEL_PATH" \ + --aime "$FINAL_AIME_PATH" \ + --aime-revision "$FINAL_AIME_REVISION" \ + --aime-sha256 "$FINAL_AIME_SHA256" \ + --decode "$FINAL_PROBE_DECODE" \ + --json /tmp/ft-rocm-final-probe.json + + - name: Evaluate promotion gate + run: | + reference_arg=() + if [ -n "$FINAL_REFERENCE_JSONL" ]; then + test -f "$FINAL_REFERENCE_JSONL" + reference_arg=(--reference "$FINAL_REFERENCE_JSONL") + fi + PYTHONPATH=python .venv-rocm/bin/python benchmarks/check_decode_gate.py \ + --freetoken /tmp/ft-rocm-final.jsonl \ + --probe /tmp/ft-rocm-final-probe.json \ + "${reference_arg[@]}" --json /tmp/ft-rocm-final-gate.json + + - name: Upload final gate evidence + if: ${{ always() }} + uses: actions/upload-artifact@v4 + with: + name: rocm-final-gate-${{ github.run_id }} + if-no-files-found: warn + path: | + /tmp/ft-rocm-final.jsonl + /tmp/ft-rocm-final-probe.json + /tmp/ft-rocm-final-gate.json + /tmp/ft-rocm-broad-suite.log diff --git a/benchmarks/README.md b/benchmarks/README.md index 6218903f2..0a3798bc3 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -11,6 +11,43 @@ include the full serving path. AIME-25 prompt, checkpoint-recommended sampling. python benchmarks/bench_decode_moe.py --model /path/to/model --backend offload,cpu,hybrid ``` +For promotion evidence, use one immutable local AIME JSONL (or Hub revision plus SHA-256), +the exact same GGUF, `--decode 512 --context 9216 --batch 512 --ubatch 512 --kv-type q8_0`, +and ten repeats. Validate rows with `benchmarks/check_decode_gate.py`; missing fixture, +runtime KV, graph, thermal, model, or reference identity rejects promotion. Set +`FREETOKEN_GGUF_MOE_IMPL=legacy` for rollback. `FREETOKEN_GRAPH_SAMPLER=1` enables +capture-safe greedy sampler/device-token-chain evidence; dynamic sampling keeps fallback. + +**`bench_decode_replay.py`** — pins one legacy FreeToken greedy continuation and runs a +separate teacher-forced replay lane. Route capture is untimed and never contributes to +replay timing. The golden command must run on an idle GPU: + +```bash +python benchmarks/bench_decode_replay.py golden \ + --model /path/to/model --aime /path/to/test.jsonl \ + --aime-sha256 SHA256 --out .plans/rocm-ollama-gap/replay-manifest.json +python benchmarks/bench_decode_replay.py replay-freetoken \ + --model /path/to/model --manifest .plans/rocm-ollama-gap/replay-manifest.json \ + --aime /path/to/test.jsonl --aime-sha256 SHA256 --routes \ + --json .plans/rocm-ollama-gap/replay-freetoken.jsonl +``` + +`bench_llama_cpp_hip.py --replay-manifest` delegates fixed-token replay to a supplied +`llama-server`; route hashes remain unavailable unless reference binary is instrumented. + +`bench_gguf_moe_kernels.py --paired --impl gfx1100` interleaves legacy/candidate +exact-shape calls and emits raw samples plus paired median recovery and bootstrap CI. +It requires a GPU when run; no result is considered promotion evidence until output +bytes, quant shape, baseline drift, and fallback count pass their gates. + +`profile_decode_rocm.py` parses an existing rocprof artifact or launches only an +explicit command after `--`. It fails closed without clock-correlation records and +reports disjoint token ledgers plus warm-offload event counts/bytes. + +Promotion gate output contains independent `gate_a` (sampled absolute) and `gate_b` +(q8/q8 teacher-forced replay) objects. Gate B requires matched IDs, route hashes, +identities, and paired performance; absence of replay evidence never promotes q8. + **`bench_load_weight_generic.py`** — expert-bank load time: serial vs parallel O_DIRECT vs pre-repacked FTW, each mode in its own subprocess. Linux-only; stages the FTW under `/var/tmp` (`--ftw-dir` overrides; roughly checkpoint-sized). diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 000000000..c4a560a34 --- /dev/null +++ b/benchmarks/__init__.py @@ -0,0 +1 @@ +"""FreeToken benchmark helpers and executable harnesses.""" diff --git a/benchmarks/bench_decode_moe.py b/benchmarks/bench_decode_moe.py index 566217927..867615df0 100644 --- a/benchmarks/bench_decode_moe.py +++ b/benchmarks/bench_decode_moe.py @@ -20,8 +20,8 @@ Prompt: an AIME-25 problem sent as a chat message with thinking enabled -- a real reasoning workload, so expert routing is representative. The server renders the chat template (including checkpoint-shipped encoders like DSV4's ``encoding_dsv4.py``). The -problems come from the ``math-ai/aime25`` dataset on the Hub, downloaded into the usual -HF cache on first run; ``--aime`` points at a local jsonl instead. +problems come from a local immutable jsonl fixture; Hub input is accepted only with an +explicit revision and SHA-256, then cached by the Hub client. Sampling: the checkpoint's recommended params (``generation_config.json``), falling back to temperature 1.0 / top_p 0.95 / top_k 64 for fields the checkpoint does not specify -- @@ -47,6 +47,8 @@ import hashlib import json import os +import platform +import re import signal import socket import subprocess @@ -61,6 +63,13 @@ # Applied for every field the checkpoint's generation_config.json does not specify. FALLBACK_SAMPLING = {"temperature": 1.0, "top_p": 0.95, "top_k": 64} +# Primary comparator contract. Keep these values in metadata even when a local +# runtime cannot prove one of them; an unproven row must be rejected by the gate. +COMPARATOR_CONTEXT = 9216 +COMPARATOR_BATCH = 512 +COMPARATOR_UBATCH = 512 +COMPARATOR_KV_TYPE = "q8_0" + # AIME-25 problems, pulled from the Hub into the usual HF cache on first run. AIME_REPO = "math-ai/aime25" AIME_FILE = "test.jsonl" @@ -77,15 +86,30 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: p.add_argument( "--backend", default="offload", - help="comma list of offload|cpu|hybrid; one server per backend", + help="comma list of fused|offload|cpu|hybrid; one server per backend", ) p.add_argument( "--aime", default=os.environ.get("FREETOKEN_AIME25_JSONL"), - help=f"local jsonl instead of downloading {AIME_REPO}; default $FREETOKEN_AIME25_JSONL", + help=f"local jsonl instead of downloading pinned {AIME_REPO}; default $FREETOKEN_AIME25_JSONL", + ) + p.add_argument( + "--aime-revision", + default=os.environ.get("FREETOKEN_AIME25_REVISION"), + help="Hub dataset commit/revision; required with --aime-sha256 when --aime is absent", + ) + p.add_argument( + "--aime-sha256", + default=os.environ.get("FREETOKEN_AIME25_SHA256"), + help="expected SHA-256 for local or pinned Hub AIME JSONL", ) p.add_argument("--problem", type=int, default=0, help="0-based AIME problem index") - p.add_argument("--decode", type=int, default=256, help="decode tokens to measure (D)") + p.add_argument("--decode", type=int, default=512, help="decode tokens to measure (D)") + p.add_argument("--repeats", type=int, default=10, help="measured requests per server") + p.add_argument("--context", type=int, default=COMPARATOR_CONTEXT) + p.add_argument("--batch", type=int, default=COMPARATOR_BATCH) + p.add_argument("--ubatch", type=int, default=COMPARATOR_UBATCH) + p.add_argument("--kv-type", default=COMPARATOR_KV_TYPE) p.add_argument( "--cache", type=int, @@ -93,6 +117,12 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: help="GPU expert cache slots; 0 = auto-size from free VRAM", ) p.add_argument("--cache-rate", type=float, default=None, help="cache slots as a fraction of L*E") + p.add_argument( + "--kv-reserve-tokens", + type=int, + default=8192, + help="KV tokens reserved while comparing MoE cache sizes", + ) p.add_argument( "--hybrid-fetch", type=int, @@ -100,6 +130,11 @@ 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( + "--attention-backend", + default="triton", + help="attention backend passed to ft serve (default: triton)", + ) 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") @@ -118,18 +153,67 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: return p.parse_args(argv) -def load_problem(path: str | None, index: int) -> tuple[str, str]: - """One AIME-25 (problem, answer). Downloads the dataset unless ``path`` overrides it. +def sha256_file(path: str) -> str: + digest = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() - Accepts both the Hub schema (``problem``) and the pre-formatted jsonl some local copies - use (``prompt``, answer instruction already appended).""" - if not path: + +def _dataset_path( + path: str | None, revision: str | None, expected_sha256: str | None +) -> tuple[str, dict[str, str | None]]: + """Resolve an immutable local or pinned Hub dataset and return provenance.""" + if path: + dataset_path = Path(path).expanduser() + if not dataset_path.is_file(): + sys.exit(f"AIME fixture not found: {dataset_path}") + provenance = { + "source": "local", + "path": str(dataset_path.resolve()), + "revision": revision, + "sha256": sha256_file(dataset_path), + } + else: + if not revision or not expected_sha256: + sys.exit( + "AIME dataset must be local (--aime) or pinned with both " + "--aime-revision and --aime-sha256" + ) from huggingface_hub import hf_hub_download try: - path = hf_hub_download(AIME_REPO, AIME_FILE, repo_type="dataset") + dataset_path = Path( + hf_hub_download(AIME_REPO, AIME_FILE, repo_type="dataset", revision=revision) + ) except Exception as e: # offline, rate-limited, repo moved - sys.exit(f"could not fetch {AIME_REPO}/{AIME_FILE} ({e}); pass --aime ") + sys.exit( + f"could not fetch pinned {AIME_REPO}/{AIME_FILE}@{revision} ({e}); " + "pass --aime " + ) + provenance = { + "source": "huggingface", + "path": str(dataset_path), + "revision": revision, + "sha256": sha256_file(dataset_path), + } + actual = str(provenance["sha256"]) + if expected_sha256 and actual.lower() != expected_sha256.lower(): + sys.exit(f"AIME fixture SHA-256 mismatch: expected {expected_sha256}, got {actual}") + return str(dataset_path), provenance + + +def load_problem_details( + path: str | None, + index: int, + revision: str | None = None, + expected_sha256: str | None = None, +) -> tuple[str, str, dict[str, str | None]]: + """One immutable AIME-25 problem plus answer and dataset provenance.""" + path, provenance = _dataset_path(path, revision, expected_sha256) + if not path: + raise AssertionError("_dataset_path returned an empty path") rows = [json.loads(line) for line in Path(path).read_text().splitlines() if line.strip()] if not 0 <= index < len(rows): sys.exit(f"--problem {index} out of range ({len(rows)} problems available)") @@ -137,7 +221,18 @@ def load_problem(path: str | None, index: int) -> tuple[str, str]: text = row.get("problem") or row["prompt"] if "boxed" not in text: text = f"{text}\n{BOXED_INSTRUCTION}" - return text, str(row.get("answer", "")) + return text, str(row.get("answer", "")), provenance + + +def load_problem( + path: str | None, + index: int, + revision: str | None = None, + expected_sha256: str | None = None, +) -> tuple[str, str]: + """Compatibility wrapper for callers that only need problem text and answer.""" + text, answer, _ = load_problem_details(path, index, revision, expected_sha256) + return text, answer def resolve_sampling(model_path: str, greedy: bool) -> tuple[dict, str]: @@ -179,11 +274,15 @@ def serve_cmd(args: argparse.Namespace, backend: str, port: int) -> list[str]: "--model", args.model, "--host", "127.0.0.1", "--port", str(port), "--moe-backend", backend, + "--attention-backend", args.attention_backend, "--max-running-requests", "1", - "--max-seq-len-override", str(8192 + args.decode), + "--max-seq-len-override", str(args.context), + "--kv-type", args.kv_type, + "--kv-reserve-tokens", str(args.kv_reserve_tokens), "--memory-ratio", str(args.mem_ratio), "--cuda-graph-max-bs", "0" if args.no_graph else "1", "--moe-hybrid-max-fetch", str(args.hybrid_fetch), + "--moe-collect-stats", ] if args.gpu: cmd += ["--gpu", args.gpu] @@ -249,13 +348,532 @@ def stop_server(proc: subprocess.Popen) -> None: time.sleep(3) # let the driver reclaim VRAM before the next backend's server +def _jsonable(value): + if isinstance(value, dict): + return {str(k): _jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(v) for v in value] + if hasattr(value, "item"): + return value.item() + return value + + +def model_fingerprint(model_path: str) -> dict: + """Full content identity plus stable source and GGUF metadata.""" + root = Path(model_path).expanduser() + if not root.exists(): + return {"path": str(root), "error": "model path does not exist"} + if root.is_file(): + stat = root.stat() + sample = 1 << 20 + digest = hashlib.sha256() + with root.open("rb") as f: + digest.update(f.read(sample)) + if stat.st_size > sample: + f.seek(max(0, stat.st_size - sample)) + digest.update(f.read(sample)) + result = { + "path": str(root.resolve()), + "kind": "file", + "size_bytes": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + "sha256": sha256_file(str(root)), + "identity": "full-file-sha256", + "fingerprint": f"head-tail-sha256:{digest.hexdigest()}", + } + if root.suffix == ".gguf": + try: + from freetoken.models.gguf.reader import load_gguf_metadata + + metadata = load_gguf_metadata(str(root)) + result["gguf_metadata"] = _jsonable( + { + key: value + for key, value in metadata.items() + if key.startswith("general.") + or "quantization" in key + or key.endswith(".file_type") + } + ) + except Exception as exc: + result["gguf_metadata_error"] = f"{type(exc).__name__}: {exc}" + return result + files = [] + for path in sorted(root.rglob("*")): + if path.is_file(): + stat = path.stat() + files.append( + { + "path": str(path.relative_to(root)), + "size_bytes": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + "sha256": sha256_file(str(path)), + } + ) + content_manifest = [ + {key: entry[key] for key in ("path", "size_bytes", "sha256")} + for entry in files + ] + digest = hashlib.sha256(json.dumps(content_manifest, sort_keys=True).encode()).hexdigest() + return { + "path": str(root.resolve()), + "kind": "directory", + "files": files, + "sha256": digest, + "identity": "directory-content-sha256", + "fingerprint": f"manifest-sha256:{digest}", + } + + +def runtime_metadata() -> dict: + result = {"python": sys.version, "platform": platform.platform()} + try: + import torch + from freetoken.utils.graph_gate import rocm_blas_report + + result["torch"] = torch.__version__ + result["cuda"] = torch.version.cuda + result["rocm"] = torch.version.hip + # Benchmark metadata must report current worker policy without launching a graph + # probe in the client process; graph gate remains a server-start concern. + result["blas"] = rocm_blas_report(gate={}) + if torch.cuda.is_available(): + index = torch.cuda.current_device() + props = torch.cuda.get_device_properties(index) + result.update( + { + "gpu_index": index, + "gpu_name": props.name, + "gpu_capability": f"{props.major}.{props.minor}", + "gpu_total_memory_bytes": props.total_memory, + "gpu_arch": getattr(props, "gcnArchName", None) + or os.environ.get("FREETOKEN_GPU_ARCH"), + } + ) + except Exception as exc: + result["torch_error"] = f"{type(exc).__name__}: {exc}" + result["gpu_telemetry"] = gpu_telemetry() + # FreeToken currently exposes no server flag for quantized KV storage. Keep + # observation explicit so q8_0 rows cannot be mistaken for BF16/default KV. + result.setdefault("kv_type", None) + result["dirty_diff_hash"] = dirty_diff_hash() + result["env_selector"] = environment_selector_digest() + result["jit_binary"] = jit_binary_sha() + return result + + +def gpu_telemetry() -> dict: + """Best-effort clocks/thermal/power snapshot; unavailable fields stay null.""" + telemetry = { + "core_clock_mhz": None, + "memory_clock_mhz": None, + "hotspot_c": None, + "power_w": None, + "source": None, + } + try: + result = subprocess.run( + ["rocm-smi", "--showclocks", "--showtemp", "--showpower", "--json"], + check=False, + capture_output=True, + text=True, + timeout=5, + ) + except (OSError, subprocess.TimeoutExpired): + return telemetry + if result.returncode != 0: + return telemetry + try: + payload = json.loads(result.stdout) + except json.JSONDecodeError: + payload = {} + values = {} + if isinstance(payload, dict): + values = next((v for v in payload.values() if isinstance(v, dict)), payload) + if isinstance(values, dict): + flat = {str(k).lower().replace(" ", "_"): v for k, v in values.items()} + + def as_number(raw): + """rocm-smi JSON emits strings like "51.0" or "(42Mhz)"; strip decor.""" + if isinstance(raw, (int, float)): + return float(raw) + if isinstance(raw, str): + digits = re.sub(r"[^0-9.\-]", "", raw) + if digits: + try: + return float(digits) + except ValueError: + return None + return None + + for key, aliases in { + "core_clock_mhz": ( + "current_gpu_clk", "gpu_clk", "gpu_clock", + "sclk_clock_speed:", "sclk_clock_speed", "sclk", + ), + "memory_clock_mhz": ( + "current_mem_clk", "mem_clk", "memory_clock", + "mclk_clock_speed:", "mclk_clock_speed", "mclk", + ), + "hotspot_c": ( + "junction_temperature", "hotspot_temperature", + "temperature_(sensor_junction)_(c)", + ), + "power_w": ( + "average_graphics_package_power_(w)", "average_power", + "power_avg", "power", + ), + }.items(): + for alias in aliases: + value = as_number(flat.get(alias)) + if value is not None: + telemetry[key] = value + break + telemetry["source"] = "rocm-smi-json" + return telemetry + + +def run_metadata( + *, + args: argparse.Namespace, + backend: str, + model: dict, + runtime: dict, + graph: dict, + dataset: dict, + sampling: dict, + sampling_src: str, + revision: str, + prompt_sha256: str, + model_id: str, +) -> dict: + """Machine-readable provenance shared by every accepted or rejected row.""" + gguf_metadata = model.get("gguf_metadata") or {} + gguf_type = next( + ( + gguf_metadata[key] + for key in ("general.file_type", "general.quantization_version") + if key in gguf_metadata + ), + None, + ) + return { + "model_sha256": model.get("sha256"), + "model_identity": model.get("identity", "unverified"), + "gguf_type": gguf_type, + "device": runtime.get("gpu_name") or runtime.get("device", "unknown"), + "arch": runtime.get("gpu_arch") or runtime.get("gpu_capability"), + "torch": runtime.get("torch"), + "rocm": runtime.get("rocm"), + "blas": runtime.get("blas"), + "backend": backend, + "context": args.context, + "batch": args.batch, + "ubatch": args.ubatch, + "kv_type": args.kv_type, + "kv_type_source": "requested_cli", + "fixture_sha256": dataset.get("sha256"), + "graph": graph, + "attention": args.attention_backend, + "moe": backend, + "mtp": "off", + "speculative": False, + "decode_batch_size": 1, + "prompt_sha256": prompt_sha256, + "prompt_problem": args.problem, + "prompt_tokens_expected": None, + "completion_tokens_expected": args.decode, + "sampling": sampling, + "sampling_source": sampling_src, + "model_id": model_id, + "dataset": dataset, + "git_revision": revision, + "dirty_diff_hash": runtime.get("dirty_diff_hash"), + "env_selector": runtime.get("env_selector"), + "jit_binary": runtime.get("jit_binary"), + "freetoken_commit": revision, + "lane": "greedy_correctness" if getattr(args, "greedy", False) else "sampled_absolute", + } + + +def execution_metadata( + *, args: argparse.Namespace, backend: str, graph: dict, cache_status: dict +) -> dict: + """Normalize observed runtime facts into one promotion-gate record.""" + geometry = cache_status.get("geometry") if isinstance(cache_status, dict) else None + observed = geometry.get("execution") if isinstance(geometry, dict) else None + observed = observed if isinstance(observed, dict) else {} + kv_storage = observed.get("kv_storage") + kv_storage = kv_storage if isinstance(kv_storage, dict) else {} + return { + "requested_moe_backend": backend, + "effective_moe_backend": observed.get("effective_moe_backend"), + "expert_storage": observed.get("expert_storage"), + "resident_gguf": observed.get("resident_gguf"), + "expert_fetches": observed.get("expert_fetches"), + "expert_remaps": observed.get("expert_remaps"), + "attention_backend": args.attention_backend, + "graph_state": graph.get("state"), + "graph_gate": graph.get("gate"), + "decode_batch_size": 1, + "mtp": "off", + "speculative": False, + "execution_class": observed.get("execution_class"), + "kv_type": observed.get("kv_type", kv_storage.get("storage_type")), + "kv_contract_id": kv_storage.get("contract_id"), + "kv_pointer_generation": kv_storage.get("pointer_generation"), + "kv_unit_bytes": kv_storage.get("unit_bytes"), + "memory_phases": observed.get("memory_phases", []), + } + + +def acceptance_status( + *, + args: argparse.Namespace, + result: dict, + graph: dict, + usage: dict, + model: dict, + runtime: dict | None = None, + dataset: dict | None = None, +) -> dict: + """Return explicit row status; rejected rows never contribute to medians.""" + completion = usage.get("completion_tokens") + context = getattr(args, "context", COMPARATOR_CONTEXT) + batch = getattr(args, "batch", COMPARATOR_BATCH) + ubatch = getattr(args, "ubatch", COMPARATOR_UBATCH) + kv_type = getattr(args, "kv_type", COMPARATOR_KV_TYPE) + stamps = result.get("stamps") or [] + reasons = [] + checks = { + "exact_completion_count": completion == args.decode, + "finite_output": bool(result.get("text")), + "graph_state_known": graph.get("state") != "unknown", + "mtp_off": True, + "speculative_off": True, + "model_sha256": bool(model.get("sha256")), + "fixture_sha256": bool((dataset or {}).get("sha256")), + "kv_type": (runtime or {}).get("kv_type") == COMPARATOR_KV_TYPE, + "thermal_clock_valid": True, + "exact_comparator_config": ( + args.decode == 512 + and context == COMPARATOR_CONTEXT + and batch == COMPARATOR_BATCH + and ubatch == COMPARATOR_UBATCH + and kv_type == COMPARATOR_KV_TYPE + ), + # API streaming does not expose logits. Probe artifacts own this gate. + "finite_logits": "unavailable_api", + } + if len(stamps) < 2: + reasons.append(f"need >=2 token events, got {len(stamps)}") + if completion != args.decode: + reasons.append(f"completion_tokens={completion!r} != --decode {args.decode}") + if not checks["exact_comparator_config"]: + reasons.append( + "comparator requires decode=512 context=9216 batch=512 ubatch=512 kv_type=q8_0" + ) + if not result.get("text"): + reasons.append("server returned no generated text") + if not model.get("sha256"): + reasons.append("model full SHA-256 unavailable") + if dataset is not None and not dataset.get("sha256"): + reasons.append("fixture full SHA-256 unavailable") + if runtime is not None: + if runtime.get("kv_type") != COMPARATOR_KV_TYPE: + reasons.append( + f"observed KV type={runtime.get('kv_type')!r} != {COMPARATOR_KV_TYPE!r}" + ) + telemetry = runtime.get("gpu_telemetry") or {} + required_telemetry = ("core_clock_mhz", "memory_clock_mhz", "hotspot_c") + if not all(isinstance(telemetry.get(key), (int, float)) for key in required_telemetry): + checks["thermal_clock_valid"] = False + reasons.append("GPU clock/thermal telemetry unavailable") + end_telemetry = runtime.get("gpu_telemetry_end") or {} + if all(isinstance(end_telemetry.get(key), (int, float)) for key in required_telemetry): + hotspot_delta = abs( + float(end_telemetry["hotspot_c"]) - float(telemetry["hotspot_c"]) + ) + if hotspot_delta > 8.0: + checks["thermal_clock_valid"] = False + reasons.append(f"GPU hotspot drift={hotspot_delta:.1f}C exceeds 8C") + if result.get("stream_error"): + reasons.append(str(result["stream_error"])) + if graph.get("state") == "unknown": + reasons.append("actual graph state unavailable") + return { + "status": "accepted" if not reasons else "rejected", + "accepted": not reasons, + "checks": checks, + "reasons": reasons, + } + + +def graph_metadata(log_path: str, requested: bool) -> dict[str, str | bool]: + """Infer actual graph path only from engine log evidence.""" + text = Path(log_path).read_text(errors="replace") + if not requested: + return {"requested": False, "state": "disabled", "gate": "not_requested"} + if "HIP graph capture gate FAILED" in text: + return {"requested": True, "state": "eager", "gate": "fail"} + if "CUDA graph is disabled." in text: + return {"requested": True, "state": "eager", "gate": "disabled"} + if "Start capturing CUDA graphs" in text and "Free GPU memory after capturing CUDA graphs" in text: + return {"requested": True, "state": "replay", "gate": "pass"} + return {"requested": True, "state": "unknown", "gate": "unknown"} + + +def decode_cache_stats(log_path: str) -> dict[str, float | str]: + """Read latest engine decode cache counters from the mirrored server log.""" + text = Path(log_path).read_text(errors="replace") + matches = re.findall( + r"Decode batch.*?moe hit: ([0-9.]+), miss: ([0-9.]+), " + r"fetch: ([0-9.]+), cpu: ([0-9.]+)", + text, + ) + if not matches: + return {} + hit, miss, fetch, cpu = (float(value) for value in matches[-1]) + return { + "hit_rate": hit, + "miss_rate": miss, + "fetch_rate": fetch, + "cpu_per_layer": cpu, + "source": "server_log_latest_decode_batch", + } + + +def native_decode_stats(log_path: str) -> dict[str, float | int | str]: + """Read scheduler's device-side decode rate, separate from API arrival timing.""" + text = Path(log_path).read_text(errors="replace") + matches = re.findall( + r"Decode batch.*?gen throughput \(token/s\): ([0-9.]+)", text + ) + if not matches: + return {} + return { + "native_decode_tok_s": float(matches[-1]), + "native_timing_source": "scheduler_decode_log", + "native_timing_samples": len(matches), + } + + +def git_revision() -> str: + try: + return subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=Path(__file__).resolve().parents[1], text=True + ).strip() + except (OSError, subprocess.CalledProcessError): + return "unknown" + + +_RELEVANT_ENV_PREFIXES = ( + "FREETOKEN_", "LLAMA_", "HIP", "ROCR", "ROCM", "HSA_", "GPU_MAX_", + "PYTORCH_ROCM", "TORCH_ROCM", "CUDA_VISIBLE", "ROCM_PATH", "MIOPEN_", +) + + +def dirty_diff_hash() -> str | None: + """Content hash of the exact dirty state the measured code came from. + + A dirty tree is legitimate evidence as long as its content identity is pinned; + the hash is what makes the pinned state reproducible later.""" + repo = Path(__file__).resolve().parents[1] + try: + head = subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=repo, text=True + ).strip() + status = subprocess.check_output( + ["git", "status", "--porcelain=v1", "-uall"], cwd=repo, text=True + ) + except (OSError, subprocess.CalledProcessError): + return None + if not status.strip(): + return None + digest = hashlib.sha256() + digest.update(f"head={head}\n".encode()) + digest.update(status.encode()) + try: + diff = subprocess.check_output( + ["git", "diff", "HEAD"], cwd=repo, text=True, timeout=30 + ) + except (OSError, subprocess.CalledProcessError): + diff = "" + digest.update(f"diff_bytes={len(diff)}\n".encode()) + digest.update(diff.encode("utf-8", errors="replace")) + return digest.hexdigest() + + +def environment_selector_digest(env: dict | None = None) -> dict: + """Digest the runtime-selection environment the spawned server inherits. + + Covers FREETOKEN_* feature selectors plus the HIP/ROCm runtime knobs that can + change kernel selection or code generation. Values are included (they are + config, not secrets); unknown/absent keys contribute nothing.""" + relevant = { + key: value + for key, value in (env if env is not None else os.environ).items() + if key.startswith(("FREETOKEN_", "HIP", "ROCR", "ROCM", "HSA", "PYTORCH_ROCM")) + or key in {"LD_LIBRARY_PATH", "PYTHONPATH"} + } + canonical = "\n".join( + f"{key}={relevant[key]}" for key in sorted(relevant) + ) + return { + "sha256": hashlib.sha256(canonical.encode()).hexdigest(), + "keys": sorted(relevant), + "count": len(relevant), + } + + +def jit_binary_sha() -> dict: + """Best-effort SHA-256 of the JIT-compiled GGUF kernel module binary. + + The server compiles ``freetoken_gguf_kernels`` into torch's extension cache; + the newest ``.so`` in that build directory is what actually loads. Absent or + unreadable cache is reported as null with a reason, never as a fake hash.""" + info: dict = {"sha256": None, "path": None, "reason": None} + try: + from torch.utils.cpp_extension import _get_build_directory + + build_dir = Path(_get_build_directory("freetoken_gguf_kernels", False)) + binaries = sorted(build_dir.glob("*.so"), key=lambda p: p.stat().st_mtime) + if not binaries: + info["reason"] = "no built .so in extension cache (module not compiled yet)" + return info + newest = binaries[-1] + info["path"] = str(newest) + info["sha256"] = sha256_file(str(newest)) + except Exception as exc: # extension cache layout changed / torch absent + info["reason"] = f"{type(exc).__name__}: {exc}" + return info + + +def kernel_implementation_from_log(log_path: str) -> dict: + """Observed kernel implementation/fallback evidence from the server log. + + The engine owns implementation selection; the log is the only client-visible + record of what actually loaded, mirroring how graph state is already read.""" + text = Path(log_path).read_text(errors="replace") + match = re.search(r"Auto-selected MoE backend: (\S+)", text) + fallbacks = re.findall(r"fallback\b", text, re.IGNORECASE) + return { + "auto_selected_moe_backend": match.group(1) if match else None, + "fallback_markers": len(fallbacks), + "source": "server_log", + } + + def stream_generate(origin: str, model_id: str, problem: str, sampling: dict, args: argparse.Namespace) -> dict: """One streamed chat completion; returns per-token arrival stamps, text, and usage.""" + # On current FreeToken, a warmed full-prefix request consumes one output budget slot + # before its first streamed completion ack. Ask for one extra engine slot so strict + # validation still produces exactly D completion tokens instead of accepting D-1. body = { "model": model_id, "messages": [{"role": "user", "content": problem}], - "max_tokens": args.decode, + "max_tokens": args.decode + 1, "ignore_eos": True, "stream": True, "stream_options": {"include_usage": True}, @@ -292,23 +910,37 @@ def stream_generate(origin: str, model_id: str, problem: str, sampling: dict, usage = chunk["usage"] for choice in chunk.get("choices", []): delta = choice.get("delta") or {} - text = delta.get("reasoning_content") or delta.get("content") + text = "".join( + part for part in (delta.get("reasoning_content"), delta.get("content")) if part + ) if text: stamps.append(now) pieces.append(text) - if usage is None: - sys.exit("[bench] stream ended without a usage chunk; is this a FreeToken server?") - return {"t0": t0, "stamps": stamps, "text": "".join(pieces), "usage": usage} + return { + "t0": t0, + "stamps": stamps, + "text": "".join(pieces), + "usage": usage or {}, + "stream_error": None if usage is not None else "missing_usage_chunk", + "prompt_sha256": hashlib.sha256(problem.encode("utf-8")).hexdigest(), + } -def run_one(args: argparse.Namespace, backend: str) -> dict: - problem, answer = load_problem(args.aime, args.problem) +def run_one(args: argparse.Namespace, backend: str) -> list[dict]: + problem, answer, dataset = load_problem_details( + args.aime, args.problem, args.aime_revision, args.aime_sha256 + ) sampling, sampling_src = resolve_sampling(args.model, args.greedy) + model = model_fingerprint(args.model) + runtime = runtime_metadata() + revision = git_revision() port = free_port() origin = f"http://127.0.0.1:{port}" fd, log_path = tempfile.mkstemp(prefix=f"bench-serve-{backend}-", suffix=".log") cmd = serve_cmd(args, backend, port) + if args.repeats < 1: + sys.exit("--repeats must be >= 1") print( f"[bench] model={args.model}\n" f"[bench] backend={backend} cache={args.cache or args.cache_rate or 'auto'} " @@ -332,64 +964,210 @@ def run_one(args: argparse.Namespace, backend: str) -> dict: # Warm the expert cache to a steady-state decode working set. stream_generate(origin, model_id, problem, sampling, args) - r = stream_generate(origin, model_id, problem, sampling, args) - stats = get_json(f"{origin}/v1/stats") + rows = [] + for repeat in range(args.repeats): + # Per-repeat start telemetry: a like-for-like loaded-vs-loaded + # window. The pre-spawn snapshot is idle and would make every + # row fail the hotspot-delta check by construction. + runtime_start = {**runtime, "gpu_telemetry": gpu_telemetry()} + r = stream_generate(origin, model_id, problem, sampling, args) + stats = get_json(f"{origin}/v1/stats") + try: + cache_status = get_json(f"{origin}/v1/cache/status") + except (OSError, ValueError): + cache_status = {} + rows.append( + make_row( + args, + backend, + repeat, + r, + stats, + cache_status, + log_path, + model, + runtime_start, + revision, + dataset, + sampling, + sampling_src, + model_id, + ) + ) finally: stop_server(proc) pump.join(timeout=10) - stamps, usage = r["stamps"], r["usage"] - if len(stamps) < 2: - sys.exit(f"[bench] need >=2 token events to measure decode, got {len(stamps)}") - completion = usage["completion_tokens"] - if completion != args.decode: - print(f"[bench] WARNING: completion_tokens={completion} != --decode {args.decode}", flush=True) - steps = completion - 1 - decode_time = stamps[-1] - stamps[0] + for row in rows: + print_row(row, args, answer, rows) + return rows + + +def make_row( + args: argparse.Namespace, + backend: str, + repeat: int, + result: dict, + stats: dict, + cache_status: dict, + log_path: str, + model: dict, + runtime: dict, + revision: str, + dataset: dict, + sampling: dict, + sampling_src: str, + model_id: str, +) -> dict: + stamps, usage = result["stamps"], result["usage"] + runtime = dict(runtime) + runtime["gpu_telemetry_end"] = gpu_telemetry() + completion = usage.get("completion_tokens") + steps = completion - 1 if isinstance(completion, int) and completion > 0 else 0 + decode_time = stamps[-1] - stamps[0] if len(stamps) >= 2 else 0.0 gaps = sorted((b - a) * 1e3 for a, b in zip(stamps, stamps[1:])) + graph = graph_metadata(log_path, not args.no_graph) + execution = execution_metadata( + args=args, backend=backend, graph=graph, cache_status=cache_status + ) + # CLI request is not observation. The readiness/status response is the first + # server-owned source proving physical KV allocation and contract. + observed_kv = execution.get("kv_type") + if observed_kv is not None: + runtime["kv_type"] = observed_kv + runtime["kv_storage"] = { + "contract_id": execution.get("kv_contract_id"), + "pointer_generation": execution.get("kv_pointer_generation"), + "unit_bytes": execution.get("kv_unit_bytes"), + } + acceptance = acceptance_status( + args=args, + result=result, + graph=graph, + usage=usage, + model=model, + runtime=runtime, + dataset=dataset, + ) + metadata = run_metadata( + args=args, + backend=backend, + model=model, + runtime=runtime, + graph=graph, + dataset=dataset, + sampling=sampling, + sampling_src=sampling_src, + revision=revision, + prompt_sha256=result.get("prompt_sha256", "unknown"), + model_id=model_id, + ) + metadata["execution"] = execution + metadata["prompt_tokens_expected"] = usage.get("prompt_tokens") + metadata["kernel_observed"] = kernel_implementation_from_log(log_path) + metadata["observed_kv_type"] = execution.get("kv_type") + metadata["kv_type_source"] = "engine_cache_metadata" if execution.get("kv_type") else "unobserved" row = { + "schema": "freetoken-base-decode-v2", + "status": acceptance["status"], + "acceptance": acceptance, + "metadata": metadata, + "git_revision": revision, + "lane": metadata["lane"], "model": args.model, + "model_id": model_id, + "model_fingerprint": model, "backend": backend, + "repeat": repeat, "problem": args.problem, - "prompt_tokens": usage["prompt_tokens"], + "dataset": dataset, + "prompt_tokens": usage.get("prompt_tokens"), "decode_steps": steps, - "decode_tok_s": steps / decode_time if decode_time > 0 else 0.0, - "ms_per_token": decode_time / steps * 1e3 if steps > 0 else 0.0, - "event_ms_p50": gaps[len(gaps) // 2], - "event_ms_p99": gaps[min(len(gaps) - 1, int(len(gaps) * 0.99))], - "ttft_ms": (stamps[0] - r["t0"]) * 1e3, + "decode_tok_s": steps / decode_time if decode_time > 0 and steps > 0 else None, + "ms_per_token": decode_time / steps * 1e3 if decode_time > 0 and steps > 0 else None, + "event_ms_p50": gaps[len(gaps) // 2] if gaps else None, + "event_ms_p99": gaps[min(len(gaps) - 1, int(len(gaps) * 0.99))] if gaps else None, + "ttft_ms": (stamps[0] - result["t0"]) * 1e3 if stamps else None, "events": len(stamps), "completion_tokens": completion, + "decode_requested": args.decode, + "context": args.context, + "batch": args.batch, + "ubatch": args.ubatch, + "kv_type": args.kv_type, + "fixture_sha256": dataset.get("sha256"), "vram_gib": stats.get("vram_bytes", 0) / 2**30, "sampling": sampling, - "output_sha1": hashlib.sha1(r["text"].encode()).hexdigest()[:12], + "sampling_source": sampling_src, + "mtp": "off", + "speculative": False, + "decode_batch_size": 1, + "attention_backend": args.attention_backend, + "graph": graph, + "cache_request": { + "slots": args.cache, + "rate": args.cache_rate, + "auto": args.cache <= 0 and args.cache_rate is None, + "memory_ratio": args.mem_ratio, + }, + "cache_status": cache_status, + "cache_decode_stats": decode_cache_stats(log_path), + "runtime": runtime, + "execution": execution, + "decode_window_s": decode_time, + "timing_window": "first_generated_token_to_last_generated_token", + **native_decode_stats(log_path), + "output_sha1": hashlib.sha1(result.get("text", "").encode()).hexdigest()[:12], + "output_sha256": hashlib.sha256(result.get("text", "").encode()).hexdigest(), + "output_sample": result.get("text", "")[:240], + "stream_error": result.get("stream_error"), "server_log": log_path, + "stage_summary": { + "status": "unavailable", + "reason": "launch timeline is collected by scripts/profile-rocm-decode.sh", + }, } + return row - print(f"\n==== decode bs=1 [{backend}] via /v1/chat/completions ====", flush=True) - print(f" decode throughput : {row['decode_tok_s']:8.2f} tok/s ({row['ms_per_token']:.3f} ms/token)") - print(f" TTFT (warm) : {row['ttft_ms']:8.1f} ms (prompt {row['prompt_tokens']} tok)") - print(f" decode measured : {steps} steps in {decode_time:.3f} s " - f"(event p50 {row['event_ms_p50']:.3f} / p99 {row['event_ms_p99']:.3f} ms, " - f"{len(stamps)} events)") + +def print_row(row: dict, args: argparse.Namespace, answer: str, rows: list[dict]) -> None: + throughput = row["decode_tok_s"] + ms_per_token = row["ms_per_token"] + ttft = row["ttft_ms"] + p50 = row["event_ms_p50"] + p99 = row["event_ms_p99"] + throughput_text = "rejected" if throughput is None else f"{throughput:8.2f} tok/s" + ms_text = "n/a" if ms_per_token is None else f"{ms_per_token:.3f} ms/token" + ttft_text = "n/a" if ttft is None else f"{ttft:8.1f} ms" + p50_text = "n/a" if p50 is None else f"{p50:.3f}" + p99_text = "n/a" if p99 is None else f"{p99:.3f}" + print(f"\n==== decode bs=1 [{row['backend']}] via /v1/chat/completions ====", flush=True) + print(f" repeat : {row['repeat'] + 1}/{len(rows)}") + print(f" decode throughput : {throughput_text} ({ms_text})") + print(f" TTFT (warm) : {ttft_text} (prompt {row['prompt_tokens']} tok)") + print(f" decode measured : {row['decode_steps']} steps in {row['decode_window_s']:.3f} s " + f"(event p50 {p50_text} / p99 {p99_text} ms, " + f"{row['events']} events)") print(f" vram (server) : {row['vram_gib']:8.2f} GiB") sha_note = "greedy" if args.greedy else "sampled, per-server deterministic" print(f" output sha1 : {row['output_sha1']} ({sha_note}; compare across backends)") - print(f" output sample : {r['text'][:240]!r}") - return row + print(f" graph : {row['graph']['state']} (gate={row['graph']['gate']})") + print(f" status : {row['status']} ({'; '.join(row['acceptance']['reasons']) or 'ok'})") + print(f" output sample : {row['output_sample']!r}") def main(argv: list[str] | None = None) -> int: args = parse_args(argv) backends = [b.strip() for b in args.backend.split(",") if b.strip()] - unknown = [b for b in backends if b not in ("offload", "cpu", "hybrid")] + unknown = [b for b in backends if b not in ("fused", "offload", "cpu", "hybrid")] if unknown: sys.exit(f"unknown backend(s): {unknown}") failed = [] + rejected = [] for backend in backends: try: - row = run_one(args, backend) + rows = run_one(args, backend) # SystemExit inherits BaseException, not Exception, so name both: a mid-decode # connection drop (server crash) must not abort the remaining backends either. except (SystemExit, Exception) as e: @@ -400,10 +1178,15 @@ def main(argv: list[str] | None = None) -> int: continue if args.json_out: with open(args.json_out, "a") as f: - f.write(json.dumps(row) + "\n") + for row in rows: + f.write(json.dumps(row, sort_keys=True) + "\n") + rejected.extend(row for row in rows if row["status"] != "accepted") if failed: print(f"\n[bench] backends that failed: {failed}", flush=True) return 1 + if rejected: + print(f"\n[bench] rejected measured rows: {len(rejected)}", flush=True) + return 1 return 0 diff --git a/benchmarks/bench_decode_ollama.py b/benchmarks/bench_decode_ollama.py new file mode 100644 index 000000000..b908f8497 --- /dev/null +++ b/benchmarks/bench_decode_ollama.py @@ -0,0 +1,425 @@ +"""Ollama base-decode adapter for the FreeToken Qwen MoE comparison. + +The adapter uses Ollama's native ``/api/chat`` NDJSON stream and keeps two metrics +separate: + +* ``client_arrival_tok_s`` matches FreeToken's SSE arrival window and is the only + cross-runtime metric; +* ``native_decode_tok_s`` comes from Ollama's ``eval_count/eval_duration`` fields and + is diagnostic only. + +Ollama's public API exposes draft-token configuration through ``/api/show``'s generated +Modelfile. Rows carry ``mtp=off`` only when ``draft_num_predict=0`` is explicit; otherwise +they remain unknown and must not be presented as the user's base-mode reference. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import platform +import re +import subprocess +import sys +import time +import urllib.error +import urllib.request +from pathlib import Path + +from bench_decode_moe import ( + git_revision, + load_problem_details, + runtime_metadata, + sha256_file, +) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--origin", default="http://127.0.0.1:11434") + p.add_argument("--model", required=True, help="Ollama model name/tag") + p.add_argument("--decode", type=int, default=512, help="num_predict and exact eval_count") + p.add_argument("--repeats", type=int, default=3) + p.add_argument("--aime", default=os.environ.get("FREETOKEN_AIME25_JSONL")) + p.add_argument("--aime-revision", default=os.environ.get("FREETOKEN_AIME25_REVISION")) + p.add_argument("--aime-sha256", default=os.environ.get("FREETOKEN_AIME25_SHA256")) + p.add_argument("--problem", type=int, default=0) + p.add_argument("--greedy", action="store_true") + p.add_argument( + "--ollama-gguf", + default=os.environ.get("OLLAMA_GGUF_PATH"), + help="local GGUF file verified as Ollama's loaded model blob", + ) + p.add_argument( + "--reference-gguf", + default=os.environ.get("FREETOKEN_REFERENCE_GGUF"), + help="FreeToken GGUF file used for byte-identity comparison", + ) + p.add_argument( + "--spawn-hip", + action="store_true", + help="spawn one ollama serve with a HIP-forced environment for one " + "directional row (never edits the installed unit); failure blocks nothing", + ) + p.add_argument("--json", dest="json_out", default=None) + return p.parse_args(argv) + + +def worker_backend_evidence(origin: str) -> dict: + """Best-effort worker backend evidence; unknown stays unknown, never a guess.""" + evidence: dict = {"backend": None, "source": None, "gpu_share": None} + try: + ps = request_json(origin, "/api/ps") + except Exception: + return evidence + models = ps.get("models") or [] + if not models: + return evidence + entry = models[0] + size_vram = entry.get("size_vram") or 0 + size = entry.get("size") or 0 + evidence["source"] = "api_ps" + evidence["gpu_share"] = (size_vram / size) if size else None + # Ollama does not expose its loaded library over the HTTP API; the backend + # stays unknown unless the row was captured from an explicitly HIP-forced + # spawned instance. Unknown backend downgrades the row to directional-only. + evidence["worker_backend"] = "unknown" + return evidence + + +def spawn_hip_forced_server(port: int, timeout: float = 120): + """Spawn ``ollama serve`` with a HIP-forced library override. + + Directional evidence only, per the plan: the installed unit is never edited, + and failure to obtain the row blocks nothing.""" + env = dict(os.environ) + env["OLLAMA_LLM_LIBRARY"] = "rocm" + env.setdefault("OLLAMA_HOST", f"127.0.0.1:{port}") + env["OLLAMA_HOST"] = f"127.0.0.1:{port}" + proc = subprocess.Popen( + ["ollama", "serve"], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if proc.poll() is not None: + return None + try: + request_json(f"http://127.0.0.1:{port}", "/api/version") + return proc + except (OSError, RuntimeError, ValueError, urllib.error.URLError): + time.sleep(1) + return None + + +def request_json(origin: str, path: str, body: dict | None = None) -> dict: + data = None if body is None else json.dumps(body).encode() + headers = {"Content-Type": "application/json"} if data is not None else {} + req = urllib.request.Request(f"{origin.rstrip('/')}{path}", data=data, headers=headers) + try: + with urllib.request.urlopen(req, timeout=30) as resp: + return json.load(resp) + except urllib.error.HTTPError as exc: + detail = exc.read()[:500] + raise RuntimeError(f"Ollama {path} failed: HTTP {exc.code}: {detail!r}") from exc + + +def ollama_model_info(origin: str, model: str) -> tuple[dict, dict, str | None]: + version = request_json(origin, "/api/version") + shown = request_json(origin, "/api/show", {"name": model}) + digest = next( + ( + row.get("digest") + for row in request_json(origin, "/api/tags").get("models", []) + if row.get("name") == model or row.get("model") == model + ), + None, + ) + return version, shown, digest + + +def sampling(greedy: bool) -> dict: + if greedy: + return {"temperature": 0.0, "top_p": 1.0, "top_k": -1} + return {"temperature": 1.0, "top_p": 0.95, "top_k": 64} + + +def _mtp_status(model_info: dict) -> str: + """Return proof state from Ollama's explicit draft-token runtime parameter.""" + parameters = str(model_info.get("parameters", "")) + draft = re.search(r"(?m)^\s*draft_num_predict\s+(\d+)\s*$", parameters) + if draft and int(draft.group(1)) == 0: + return "off" + if draft: + return "unknown" + return "unknown" + + +def ollama_blob_identity( + model_digest: str | None, + *, + ollama_gguf: str | None = None, + reference_gguf: str | None = None, +) -> dict: + """Separate Ollama manifest digest from verified GGUF byte identity.""" + identity = { + "manifest_digest": model_digest, + "manifest_digest_kind": "ollama-model-manifest" if model_digest else "missing", + "same_blob": False, + "status": "unproven", + "ollama_gguf": None, + "reference_gguf": None, + } + if ollama_gguf: + path = Path(ollama_gguf).expanduser() + if not path.is_file(): + identity["status"] = "invalid_ollama_gguf_path" + return identity + identity["ollama_gguf"] = { + "path": str(path.resolve()), + "size_bytes": path.stat().st_size, + "sha256": sha256_file(str(path)), + } + if reference_gguf: + path = Path(reference_gguf).expanduser() + if not path.is_file(): + identity["status"] = "invalid_reference_gguf_path" + return identity + identity["reference_gguf"] = { + "path": str(path.resolve()), + "size_bytes": path.stat().st_size, + "sha256": sha256_file(str(path)), + } + if identity["ollama_gguf"] and identity["reference_gguf"]: + ollama_sha = identity["ollama_gguf"]["sha256"] + reference_sha = identity["reference_gguf"]["sha256"] + identity["same_blob"] = ollama_sha == reference_sha + identity["status"] = "verified" if identity["same_blob"] else "mismatch" + elif identity["ollama_gguf"] or identity["reference_gguf"]: + identity["status"] = "one_side_only" + return identity + + +def stream_chat( + origin: str, + model: str, + problem: str, + options: dict, + decode: int, +) -> dict: + body = { + "model": model, + "messages": [{"role": "user", "content": problem}], + "stream": True, + "think": True, + "options": {**options, "num_predict": decode}, + } + req = urllib.request.Request( + f"{origin.rstrip('/')}/api/chat", + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, + ) + stamps: list[float] = [] + pieces: list[str] = [] + done: dict | None = None + t0 = time.perf_counter() + try: + resp = urllib.request.urlopen(req, timeout=1800) + except urllib.error.HTTPError as exc: + raise RuntimeError(f"Ollama /api/chat failed: HTTP {exc.code}: {exc.read()[:500]!r}") from exc + with resp: + for raw in resp: + line = raw.strip() + if not line: + continue + chunk = json.loads(line) + message = chunk.get("message") or {} + fragments = [message.get("thinking"), message.get("content")] + text = "".join(fragment for fragment in fragments if fragment) + if text: + stamps.append(time.perf_counter()) + pieces.append(text) + if chunk.get("done"): + done = chunk + break + if done is None: + raise RuntimeError("Ollama stream ended without done=true") + return { + "t0": t0, + "stamps": stamps, + "text": "".join(pieces), + "done": done, + "prompt": problem, + } + + +def make_row( + args: argparse.Namespace, + repeat: int, + result: dict, + version: dict, + model_info: dict, + model_digest: str | None, + dataset: dict, + mtp: str, + options: dict, + blob_identity: dict, +) -> dict: + done = result["done"] + eval_count = done.get("eval_count") + if eval_count is None: + raise RuntimeError("rejected Ollama run: done message missing eval_count") + if eval_count != args.decode: + raise RuntimeError( + f"rejected Ollama run: eval_count={eval_count} != --decode {args.decode}" + ) + stamps = result["stamps"] + if len(stamps) < 2: + raise RuntimeError(f"rejected Ollama run: need >=2 output events, got {len(stamps)}") + if not result["text"]: + raise RuntimeError("rejected Ollama run: server returned no generated text") + client_window = stamps[-1] - stamps[0] + native_duration_ns = done.get("eval_duration") + if not isinstance(native_duration_ns, (int, float)) or native_duration_ns <= 0: + raise RuntimeError("rejected Ollama run: done message missing positive eval_duration") + gaps = sorted((b - a) * 1e3 for a, b in zip(stamps, stamps[1:])) + return { + "schema": "ollama-base-decode-v1", + "status": "accepted", + "acceptance": { + "status": "accepted", + "accepted": True, + "checks": { + "exact_completion_count": True, + "finite_output": True, + "finite_logits": "unavailable_api", + "mtp_off": mtp == "off", + "same_blob": blob_identity["same_blob"], + }, + "reasons": [], + }, + "origin": args.origin, + "model": args.model, + "model_sha256": ( + blob_identity.get("ollama_gguf") or {} + ).get("sha256"), + "reference_identity": blob_identity, + "repeat": repeat, + "git_revision": git_revision(), + "dataset": dataset, + "prompt_sha256": hashlib.sha256(result["prompt"].encode("utf-8")).hexdigest(), + "decode_requested": args.decode, + "eval_count": eval_count, + "prompt_eval_count": done.get("prompt_eval_count"), + "client_arrival_tok_s": (eval_count - 1) / client_window if client_window > 0 else 0.0, + "client_arrival_window_s": client_window, + "native_decode_tok_s": eval_count / (native_duration_ns / 1e9), + "eval_duration_ns": native_duration_ns, + "prompt_eval_duration_ns": done.get("prompt_eval_duration"), + "total_duration_ns": done.get("total_duration"), + "load_duration_ns": done.get("load_duration"), + "event_ms_p50": gaps[len(gaps) // 2], + "event_ms_p99": gaps[min(len(gaps) - 1, int(len(gaps) * 0.99))], + "ttft_ms": (stamps[0] - result["t0"]) * 1e3, + "events": len(stamps), + "output_sha1": hashlib.sha1(result["text"].encode()).hexdigest()[:12], + "output_sample": result["text"][:240], + "options": options, + "think": True, + "speculative": "off" if mtp == "off" else "unknown", + "mtp": mtp, + "comparable_to_freetoken_base": mtp in ("off", "unsupported"), + "ollama_version": version, + "model_info": { + "digest": model_digest, + **{ + key: model_info.get(key) + for key in ("modified_at", "details", "parameters", "template") + if key in model_info + }, + }, + "runtime": {"python": sys.version, "platform": platform.platform(), **runtime_metadata()}, + } + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + if args.decode < 2: + raise SystemExit("--decode must be >= 2") + if args.repeats < 1: + raise SystemExit("--repeats must be >= 1") + hip_proc = None + hip_forced = False + if args.spawn_hip: + hip_port = 11500 + (int(time.time()) % 4000) + hip_proc = spawn_hip_forced_server(hip_port) + if hip_proc is None: + print( + "WARNING: HIP-forced ollama spawn failed; keeping the running " + "service row directional-only (plan: failure blocks nothing)", + file=sys.stderr, + ) + else: + args.origin = f"http://127.0.0.1:{hip_port}" + hip_forced = True + problem, _, dataset = load_problem_details( + args.aime, args.problem, args.aime_revision, args.aime_sha256 + ) + version, model_info, model_digest = ollama_model_info(args.origin, args.model) + blob_identity = ollama_blob_identity( + model_digest, + ollama_gguf=args.ollama_gguf, + reference_gguf=args.reference_gguf, + ) + options = sampling(args.greedy) + mtp = _mtp_status(model_info) + backend = worker_backend_evidence(args.origin) + if hip_forced: + backend = {**backend, "worker_backend": "hip-forced", "requested_backend": "hip"} + # Warm Ollama's loaded model before collecting measured rows. + stream_chat(args.origin, args.model, problem, options, args.decode) + rows = [] + for repeat in range(args.repeats): + row = make_row( + args, + repeat, + stream_chat(args.origin, args.model, problem, options, args.decode), + version, + model_info, + model_digest, + dataset, + mtp, + options, + blob_identity, + ) + row["worker_backend"] = backend + row["backend_evidence"] = backend + row["directional_only"] = not hip_forced + rows.append(row) + print( + f"repeat {repeat + 1}/{args.repeats}: client {row['client_arrival_tok_s']:.2f} " + f"tok/s, native {row['native_decode_tok_s']:.2f} tok/s, mtp={mtp}", + flush=True, + ) + if mtp == "unknown": + print("WARNING: Ollama MTP status unknown; rows are non-comparable base reference", file=sys.stderr) + if args.json_out: + with open(args.json_out, "a") as f: + for row in rows: + f.write(json.dumps(row, sort_keys=True) + "\n") + if hip_proc is not None: + try: + hip_proc.terminate() + hip_proc.wait(timeout=30) + except (ProcessLookupError, subprocess.TimeoutExpired): + hip_proc.kill() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/bench_decode_replay.py b/benchmarks/bench_decode_replay.py new file mode 100644 index 000000000..2f798a4f5 --- /dev/null +++ b/benchmarks/bench_decode_replay.py @@ -0,0 +1,927 @@ +"""Teacher-forced replay lanes and golden-output pinning (rocm-ollama-gap Inc 0). + +This module owns the replay contract shared by every runtime comparator: + +* **Golden capture** -- one in-process FreeToken greedy run pins the legacy + 512-token ID sequence, per-run logit finiteness, and the boxed answer. The + golden must reproduce from fresh processes; it is the correctness anchor for + every later fused-kernel increment. +* **Teacher-forced replay** -- every runtime is fed the *same* token sequence. + FreeToken is driven in-process with the sampler overridden to emit the pinned + next token, so logits are timed without sampling feedback and routes cannot + diverge after the first logit difference. llama.cpp is driven through + ``llama-server`` with ``cache_prompt`` and ``n_predict=1`` so each step + evaluates exactly one forced token per request over identical inputs. +* **Correctness replay** -- a separate untimed pass that captures per-(token, + layer) top-k route-ID hashes. Route capture adds host synchronization, so it + must never run inside timed passes; performance replay is primary only when + route hashes match exactly across runtimes. llama.cpp route capture requires + an instrumented build and is recorded as unavailable until then. + +Lane discipline: every row produced here carries ``lane=teacher_forced_replay`` +and is disjoint from ``sampled_absolute`` and ``greedy_correctness`` rows. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import shlex +import socket +import statistics +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +from pathlib import Path + +from bench_decode_moe import ( + git_revision, + load_problem_details, + model_fingerprint, + runtime_metadata, +) + +MANIFEST_SCHEMA = "freetoken-replay-manifest-v1" +LANE = "teacher_forced_replay" + + +# --------------------------------------------------------------------------- +# Manifest schema +# --------------------------------------------------------------------------- + + +def _sha256_bytes(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def route_hash(topk_ids) -> str: + """Stable hash of one (token, layer) top-k route-ID tuple.""" + return _sha256_bytes(json.dumps([int(v) for v in topk_ids]).encode())[:16] + + +def continuation_hash(ids: list[int]) -> str: + return _sha256_bytes(json.dumps([int(v) for v in ids]).encode()) + + +def validate_manifest(manifest: object) -> list[str]: + """Structural + contract validation; empty list means acceptable.""" + problems: list[str] = [] + if not isinstance(manifest, dict): + return ["manifest is not an object"] + if manifest.get("schema") != MANIFEST_SCHEMA: + problems.append(f"schema={manifest.get('schema')!r} != {MANIFEST_SCHEMA!r}") + if manifest.get("lane") != LANE: + problems.append(f"lane={manifest.get('lane')!r} != {LANE!r}") + prompt_ids = manifest.get("prompt_ids") + continuation = manifest.get("continuation_ids") + if not isinstance(prompt_ids, list) or not prompt_ids: + problems.append("prompt_ids missing or empty") + elif any(not isinstance(v, int) or v < 0 for v in prompt_ids): + problems.append("prompt_ids must be non-negative token IDs") + if not isinstance(continuation, list) or len(continuation) < 2: + problems.append("continuation_ids missing or shorter than two pinned IDs") + else: + if any(not isinstance(v, int) or v < 0 for v in continuation): + problems.append("continuation_ids must be non-negative token IDs") + if manifest.get("measured_tokens") != len(continuation): + problems.append("measured_tokens != len(continuation_ids)") + warmup = manifest.get("warmup_tokens", 0) + measured = manifest.get("measured_tokens", 0) + if not isinstance(warmup, int) or warmup < 0: + problems.append("warmup_tokens must be a non-negative integer") + if not isinstance(measured, int) or measured < 1: + problems.append("measured_tokens must be a positive integer") + for key in ("model_sha256", "fixture_sha256", "tokenizer_sha256"): + value = manifest.get(key) + if not (isinstance(value, str) and len(value) == 64 and all(c in "0123456789abcdefABCDEF" for c in value)): + problems.append(f"{key} missing or not a full SHA-256") + if not isinstance(manifest.get("route_top_k"), int) or manifest["route_top_k"] < 1: + problems.append("route_top_k missing or not positive") + golden = manifest.get("golden") + if not isinstance(golden, dict) or not ( + isinstance(golden.get("ids_sha256"), str) + and len(golden["ids_sha256"]) == 64 + and all(c in "0123456789abcdefABCDEF" for c in golden["ids_sha256"]) + ): + problems.append("golden.ids_sha256 missing") + prompt_text = manifest.get("prompt_text") + prompt_text_sha = manifest.get("prompt_text_sha256") + if prompt_text is not None: + if not isinstance(prompt_text, str): + problems.append("prompt_text must be a string when present") + elif not ( + isinstance(prompt_text_sha, str) + and len(prompt_text_sha) == 64 + and all(c in "0123456789abcdefABCDEF" for c in prompt_text_sha) + ): + problems.append("prompt_text_sha256 missing or not a full SHA-256") + elif _sha256_bytes(prompt_text.encode()) != prompt_text_sha: + problems.append("prompt_text_sha256 does not match prompt_text") + return problems + + +def load_manifest(path: str) -> dict: + manifest = json.loads(Path(path).read_text()) + problems = validate_manifest(manifest) + if problems: + raise ValueError(f"replay manifest rejected ({path}): {'; '.join(problems)}") + return manifest + + +def summarize_steps(steps_ms: list[float], warmup_steps: int) -> dict: + """Warmup-excluded step statistics; the raw per-step list is always kept.""" + warmup_steps = max(0, min(int(warmup_steps), len(steps_ms))) + measured = list(steps_ms[warmup_steps:]) + if not measured: + return { + "steps": 0, + "warmup_steps": warmup_steps, + "ms_per_token_median": None, + "raw_steps_ms": list(steps_ms), + } + return { + "steps": len(measured), + "warmup_steps": len(steps_ms) - len(measured), + "ms_per_token_median": statistics.median(measured), + "ms_per_token_mean": statistics.fmean(measured), + "ms_per_token_min": min(measured), + "ms_per_token_max": max(measured), + "raw_steps_ms": list(steps_ms), + } + + +# --------------------------------------------------------------------------- +# FreeToken in-process adapter +# --------------------------------------------------------------------------- + + +def _build_llm(args: argparse.Namespace): + import torch + + from freetoken.llm import LLM + + kwargs = { + "attention_backend": args.attention_backend, + "max_running_req": 1, + "max_extend_tokens": 8192, + "max_seq_len_override": args.context, + "moe_backend": args.moe_backend, + "memory_ratio": args.memory_ratio, + "moe_cache_auto": args.cache == 0, + "cuda_graph_max_bs": 1 if args.graph else 0, + "kv_storage_type": args.kv_type, + } + if args.cache > 0: + kwargs["moe_cache_size"] = args.cache + return LLM(args.model, dtype=torch.bfloat16, **kwargs) + + +def _tokenizer_sha(llm) -> str: + tok = llm.tokenizer + payload = { + "class": type(tok).__name__, + "vocab_size": getattr(tok, "vocab_size", None), + "eos_token_id": getattr(tok, "eos_token_id", None), + "bos_token_id": getattr(tok, "bos_token_id", None), + "pad_token_id": getattr(tok, "pad_token_id", None), + } + return _sha256_bytes(json.dumps(payload, sort_keys=True).encode()) + + +class _RouteCapture: + """Record per-(token, layer) top-k route-ID hashes from the router logits. + + Routes are re-derived with ``torch.topk`` over the float router logits -- + the same decision the layer's own router makes (softmax is order-preserving, + ``torch.topk`` tie-breaking is index-stable). Every call appends host-side + data, so capture belongs ONLY in the untimed correctness replay.""" + + def __init__(self, llm): + import torch + + self._torch = torch + self.records: list[dict] = [] + self._restores: list = [] + self._offsets: dict[int, int] = {} + for index, block in _routed_blocks(llm): + top_k = int(getattr(block, "top_k", 8)) + original = block.experts.forward + + def make_wrapper(fn, layer_index, k): + def wrapped(*call_args, **call_kwargs): + out = fn(*call_args, **call_kwargs) + try: + logits = call_kwargs.get("router_logits") + if logits is None and len(call_args) >= 2: + candidate = call_args[1] + if hasattr(candidate, "dim") and candidate.dim() == 2: + logits = candidate + if logits is not None: + ids = self._torch.topk( + logits.detach().float(), k=k, dim=-1 + ).indices + start = self._offsets.get(layer_index, 0) + hashes = [route_hash(row) for row in ids.tolist()] + self.records.append( + { + "layer": layer_index, + "start": start, + "hashes": hashes, + } + ) + self._offsets[layer_index] = start + len(hashes) + except Exception: # capture must never break the run + pass + return out + + return wrapped + + block.experts.forward = make_wrapper(original, index, top_k) + self._restores.append((block.experts, original)) + + def restore(self) -> None: + for obj, original in self._restores: + try: + obj.forward = original + except Exception: + pass + + +def _routed_blocks(llm) -> list[tuple[int, object]]: + """(model-layer-index, MoE block) for every layer with routed experts.""" + layers = getattr(llm.engine.model, "model", None) + layers = getattr(layers, "layers", None) or [] + # FreeToken's execution model stores transformer blocks in OPList.op_list; + # ordinary ModuleList containers remain directly iterable. + layers = getattr(layers, "op_list", layers) + found = [] + for index, layer in enumerate(layers): + for attr in ("moe", "mlp"): + block = getattr(layer, attr, None) + if block is not None and hasattr(block, "experts"): + found.append((index, block)) + break + return found + + +def build_manifest(args: argparse.Namespace) -> dict: + """Golden run + tokenizer/model/fixture identity, ready for replay lanes. + + Golden means the current default FreeToken path -- no candidate kernels, no + forced switches. This output is the fusion golden for Inc 9/12.""" + import torch + + from freetoken.core import SamplingParams + + problem, answer, dataset = load_problem_details( + args.aime, args.problem, args.aime_revision, args.aime_sha256 + ) + llm = _build_llm(args) + captured: list = [] + sampler = llm.engine.sampler + original_sample = sampler.sample + original_sample_into_device = sampler.sample_into_device + + def capture_sample(logits, sample_args, batch): + captured.append(logits.detach().float().cpu()) + return original_sample(logits, sample_args, batch) + + def capture_sample_into_device(logits, sample_args, batch, out, scratch): + captured.append(logits.detach().float().cpu()) + return original_sample_into_device(logits, sample_args, batch, out, scratch) + + try: + encoded = llm.tokenizer.encode(problem) + if hasattr(encoded, "tolist"): + encoded = encoded.tolist() + prompt_ids = [int(v) for v in encoded] + tokenizer_sha = _tokenizer_sha(llm) + sampler.sample = capture_sample + sampler.sample_into_device = capture_sample_into_device + try: + output = llm.generate( + [problem], + SamplingParams( + temperature=0.0, top_p=1.0, top_k=-1, + # Offline scheduler currently consumes one fewer output + # slot at exact max length than HTTP generation. Request + # one sentinel step, then pin only requested decode IDs. + max_tokens=args.decode + 1, ignore_eos=True, + ), + )[0] + finally: + sampler.sample = original_sample + sampler.sample_into_device = original_sample_into_device + finally: + llm.shutdown() + + token_ids = [int(v) for v in output["token_ids"]][: args.decode] + if len(token_ids) != args.decode: + raise RuntimeError(f"golden produced {len(token_ids)} usable ids, expected {args.decode}") + rows = torch.cat(captured, dim=0)[: args.decode] if captured else torch.empty(0) + if rows.shape[0] != args.decode: + raise RuntimeError(f"golden captured {rows.shape[0]} usable logit rows, expected {args.decode}") + if not bool(torch.isfinite(rows).all()): + raise RuntimeError("golden run produced non-finite logits") + text = output["text"] + boxed = text.split("\\boxed{")[-1].split("}")[0] if "\\boxed{" in text else None + model = model_fingerprint(args.model) + runtime = runtime_metadata() + manifest = { + "schema": MANIFEST_SCHEMA, + "lane": LANE, + "created": time.strftime("%Y-%m-%dT%H:%M:%S%z"), + "model_path": str(Path(args.model).expanduser().resolve()), + "model_sha256": model.get("sha256"), + "model_size_bytes": model.get("size_bytes"), + "model_identity": model.get("identity", "unverified"), + "fixture_sha256": dataset.get("sha256"), + "fixture_revision": dataset.get("revision"), + "tokenizer_sha256": tokenizer_sha, + "prompt_text": problem, + "prompt_text_sha256": _sha256_bytes(problem.encode()), + "prompt_ids": prompt_ids, + "continuation_ids": token_ids, + "warmup_tokens": args.warmup, + "measured_tokens": len(token_ids), + "route_top_k": args.route_top_k, + "golden": { + "source": "freetoken-legacy-greedy", + "ids_sha256": continuation_hash(token_ids), + "text_sha256": _sha256_bytes(text.encode()), + "answer": boxed, + "expected_answer": answer, + "answer_match": boxed == answer, + "finite_logits": True, + "logit_rows": int(rows.shape[0]), + "logit_sha256": _sha256_bytes(rows.numpy().tobytes()), + "decode": args.decode, + }, + "runtime": { + key: runtime.get(key) + for key in ("torch", "rocm", "gpu_capability", "gpu_arch", "gpu_name") + }, + "git_revision": git_revision(), + } + problems = validate_manifest(manifest) + if problems: + raise RuntimeError(f"built manifest is invalid: {'; '.join(problems)}") + return manifest + + +def replay_freetoken( + args: argparse.Namespace, manifest: dict, *, capture_routes: bool = False +) -> dict: + """Teacher-forced replay against FreeToken on the pinned token sequence. + + The measured window forces ``measured_tokens`` pinned IDs and reports host + submission gaps plus a synchronized end-to-end window. With + ``capture_routes`` the run becomes the untimed correctness replay that + additionally hashes per-(token, layer) top-k route IDs.""" + import torch + + from freetoken.core import SamplingParams + + prompt_ids = list(manifest["prompt_ids"]) + pinned = list(manifest["continuation_ids"]) + warmup = int(manifest.get("warmup_tokens", 0)) + measured = int(manifest["measured_tokens"]) + + llm = _build_llm(args) + stamps: list[tuple[float, int]] = [] + timed_ids: list[int] = [] + sampler = None + original_sample = None + original_sample_into_device = None + route_digest = None + route_hash_status = "not_requested" + try: + # Warmup pass: identical forced inputs, discarded (JIT/graph amortized here). + if warmup: + llm.generate( + [prompt_ids], + SamplingParams( + temperature=0.0, top_p=1.0, top_k=-1, + max_tokens=warmup, ignore_eos=True, + ), + ) + sampler = llm.engine.sampler + original_sample = sampler.sample + original_sample_into_device = sampler.sample_into_device + + def forced_sample(logits, sample_args, batch): + out = original_sample(logits, sample_args, batch) + if len(timed_ids) < measured: + index = len(timed_ids) + stamps.append((time.perf_counter(), int(pinned[index]))) + timed_ids.append(int(pinned[index])) + out.copy_(torch.full_like(out, pinned[index])) + return out + + def forced_sample_into_device(logits, sample_args, batch, out, scratch): + sampled = original_sample_into_device(logits, sample_args, batch, out, scratch) + if len(timed_ids) < measured: + index = len(timed_ids) + stamps.append((time.perf_counter(), int(pinned[index]))) + timed_ids.append(int(pinned[index])) + out.copy_(torch.full_like(out, pinned[index])) + return out + return sampled + + sampler.sample = forced_sample + sampler.sample_into_device = forced_sample_into_device + torch.cuda.synchronize() + window_start = time.perf_counter() + result = llm.generate( + [prompt_ids], + SamplingParams( + temperature=0.0, top_p=1.0, top_k=-1, + max_tokens=measured + 1, ignore_eos=True, + ), + )[0] + torch.cuda.synchronize() + window_s = time.perf_counter() - window_start + sampler.sample = original_sample + sampler.sample_into_device = original_sample_into_device + if capture_routes: + route_digest, route_hash_status = _capture_freetoken_routes_on_llm( + args, manifest, llm + ) + finally: + if sampler is not None and original_sample is not None: + sampler.sample = original_sample + if sampler is not None and original_sample_into_device is not None: + sampler.sample_into_device = original_sample_into_device + llm.shutdown() + + token_ids = [int(v) for v in result.get("token_ids") or []] + forced_ids = token_ids[:measured] + expected_ids = pinned[:measured] + steps_ms = [ + (b - ta) * 1e3 for (ta, _), (b, _) in zip(stamps, stamps[1:]) + ][:measured] + return { + "schema": "freetoken-replay-v1", + "lane": LANE, + "runtime": "freetoken", + "status": "accepted" if forced_ids == expected_ids and len(token_ids) == measured else "rejected", + "execution": "graph_replay" if args.graph else "eager", + "repeat": getattr(args, "repeat", None), + "model_sha256": manifest["model_sha256"], + "fixture_sha256": manifest["fixture_sha256"], + "tokenizer_sha256": manifest["tokenizer_sha256"], + "manifest_ids_sha256": manifest["golden"]["ids_sha256"], + "forced": True, + "ids_match": forced_ids == expected_ids and len(token_ids) == measured, + "forced_ids_sha256": continuation_hash(forced_ids), + "steps": summarize_steps(steps_ms, 0), + "window_ms": window_s * 1e3, + "ms_per_token_window": window_s * 1e3 / measured if measured else None, + "timing_domain": "host_window_synchronized", + "route_digest": route_digest, + "route_hash_status": route_hash_status, + "route_capture_timing": "untimed_separate_pass" if capture_routes else None, + "mtp": "off", + "speculative": False, + "decode_batch_size": 1, + "context": args.context, + "batch": args.batch, + "ubatch": args.ubatch, + "kv_type": args.kv_type, + "decode_ms_per_token_median": window_s * 1e3 / measured if measured else None, + "acceptance": { + "accepted": forced_ids == expected_ids and len(token_ids) == measured, + "ids_match": forced_ids == expected_ids and len(token_ids) == measured, + }, + } + + +def _route_digest(records: list[dict]) -> dict[str, str]: + """Combine per-layer route hashes into deterministic per-token digests.""" + per_step: dict[int, list[tuple[int, str]]] = {} + for record in records: + layer = int(record["layer"]) + start = int(record.get("start", 0)) + for offset, digest in enumerate(record.get("hashes", [])): + per_step.setdefault(start + offset, []).append((layer, digest)) + return { + str(step): hashlib.sha256( + json.dumps(sorted(values), separators=(",", ":")).encode() + ).hexdigest() + for step, values in sorted(per_step.items()) + } + + +def _capture_freetoken_routes_on_llm( + args: argparse.Namespace, manifest: dict, llm +) -> tuple[dict | None, str]: + """Run untimed route capture on already initialized FreeToken state.""" + import torch + + from freetoken.core import SamplingParams + + prompt_ids = list(manifest["prompt_ids"]) + pinned = list(manifest["continuation_ids"]) + measured = int(manifest["measured_tokens"]) + expected_layers = {index for index, _ in _routed_blocks(llm)} + capture = _RouteCapture(llm) + sampler = llm.engine.sampler + original_sample = sampler.sample + original_sample_into_device = sampler.sample_into_device + forced: list[int] = [] + ids_match = False + try: + def forced_sample(logits, sample_args, batch): + out = original_sample(logits, sample_args, batch) + if len(forced) < measured: + token = int(pinned[len(forced)]) + forced.append(token) + out.copy_(torch.full_like(out, token)) + return out + + def forced_sample_into_device(logits, sample_args, batch, out, scratch): + sampled = original_sample_into_device(logits, sample_args, batch, out, scratch) + if len(forced) < measured: + token = int(pinned[len(forced)]) + forced.append(token) + out.copy_(torch.full_like(out, token)) + return out + return sampled + + sampler.sample = forced_sample + sampler.sample_into_device = forced_sample_into_device + result = llm.generate( + [prompt_ids], + SamplingParams( + temperature=0.0, top_p=1.0, top_k=-1, + max_tokens=measured + 1, ignore_eos=True, + ), + )[0] + token_ids = [int(v) for v in result.get("token_ids") or []] + ids_match = token_ids == pinned[:measured] + records = [dict(record) for record in capture.records] + finally: + sampler.sample = original_sample + sampler.sample_into_device = original_sample_into_device + capture.restore() + if not ids_match or len(forced) != measured: + return None, "ids_mismatch" + if not records: + return None, "unavailable" + observed_layers = {int(record["layer"]) for record in records} + if observed_layers != expected_layers: + return None, "incomplete_layer_capture" + return _route_digest(records), "captured" + + +def capture_freetoken_routes( + args: argparse.Namespace, manifest: dict +) -> tuple[dict | None, str]: + """Run untimed forced replay solely for route capture. + + Host synchronization and route copies stay outside measured replay timing. + """ + llm = _build_llm(args) + try: + return _capture_freetoken_routes_on_llm(args, manifest, llm) + finally: + llm.shutdown() + + +# --------------------------------------------------------------------------- +# llama.cpp adapter +# --------------------------------------------------------------------------- + + +def _server_json(origin: str, path: str, body: dict | None = None, timeout: float = 30): + data = None if body is None else json.dumps(body).encode() + request = urllib.request.Request( + f"{origin.rstrip('/')}{path}", + data=data, + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=timeout) as response: + return json.load(response) + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _stop_llama_server(proc: subprocess.Popen) -> None: + if proc.poll() is not None: + return + try: + proc.terminate() + proc.wait(timeout=30) + except (ProcessLookupError, subprocess.TimeoutExpired): + try: + proc.kill() + proc.wait(timeout=10) + except (ProcessLookupError, subprocess.TimeoutExpired): + pass + + +def replay_llama(args: argparse.Namespace, manifest: dict, repeat: int) -> dict: + """Teacher-forced replay against a pinned ``llama-server``. + + Each step sends ``prompt = prompt_ids + pinned[:k]`` with ``cache_prompt`` + and ``n_predict=1``. The sampled output is discarded. Because each forced + token is appended to cached prompt, llama-server reports its forward pass + under ``timings.prompt_ms``; ``predicted_ms`` is only an additional sampled + token and is retained as a diagnostic. + The input context stays byte-identical to every other runtime's.""" + binary = Path(args.server).expanduser() + if not binary.is_file(): + raise SystemExit(f"llama-server binary not found: {binary}") + port = _free_port() + origin = f"http://127.0.0.1:{port}" + command = [ + str(binary), + "-m", str(Path(args.model).expanduser()), + "--host", "127.0.0.1", "--port", str(port), + "-c", str(getattr(args, "context", 9216)), + "-b", str(getattr(args, "batch", 512)), + "-ub", str(getattr(args, "ubatch", 512)), "-np", "1", + "-ngl", "99", "-fa", "on", + "-ctk", getattr(args, "kv_type", "q8_0"), + "-ctv", getattr(args, "kv_type", "q8_0"), + ] + extra = os.environ.get("LLAMA_SERVER_EXTRA_ARGS", "").strip() + if extra: + command = command + shlex.split(extra) + with tempfile.NamedTemporaryFile(prefix="replay-llama-", suffix=".log", delete=False) as log: + log_path = log.name + log_handle = open(log_path, "wb") + proc = subprocess.Popen( + command, stdout=log_handle, stderr=subprocess.STDOUT, start_new_session=True + ) + try: + deadline = time.monotonic() + args.timeout + while True: + if proc.poll() is not None: + raise RuntimeError(f"llama-server exited {proc.returncode}; log={log_path}") + try: + if _server_json(origin, "/health", timeout=5).get("status") == "ok": + break + except (OSError, ValueError, urllib.error.HTTPError): + pass + if time.monotonic() > deadline: + raise RuntimeError(f"llama-server not ready after {args.timeout:.0f}s") + time.sleep(1) + tokens = _server_json(origin, "/tokenize", {"content": prompt_text_of(manifest)}) + prompt_ids_match = tokens_list(tokens) == manifest["prompt_ids"] + pinned = manifest["continuation_ids"] + warmup = int(manifest.get("warmup_tokens", 0)) + measured = int(manifest["measured_tokens"]) + steps_ms: list[float] = [] + prompt_costs: list[float] = [] + decode_costs: list[float] = [] + base = manifest["prompt_ids"] + for _ in range(warmup): + _server_json( + origin, + "/completion", + { + "prompt": base, + "n_predict": 1, + "cache_prompt": True, + "temperature": 0.0, + "top_k": 1, + "seed": 0, + }, + timeout=600, + ) + for step in range(measured): + count = step + 1 + body = { + "prompt": base + pinned[:count], + "n_predict": 1, + "cache_prompt": True, + "temperature": 0.0, + "top_p": 1.0, + "top_k": 1, + "seed": 0, + } + response = _server_json(origin, "/completion", body, timeout=600) + timings = response.get("timings") or {} + prompt_cost = float(timings.get("prompt_ms") or 0.0) + decode_cost = float(timings.get("predicted_ms") or 0.0) + prompt_costs.append(prompt_cost) + decode_costs.append(decode_cost) + # The forced token is appended to the cached prompt. llama-server + # therefore accounts its forward pass as one prompt-eval token; + # predicted_ms covers only an additional sampled token and is near + # zero for this request shape. + steps_ms.append(prompt_cost) + finally: + _stop_llama_server(proc) + log_handle.close() + timing_reasons = [] + if not prompt_ids_match: + timing_reasons.append("llama-server tokenizer did not reproduce pinned prompt IDs") + if len(steps_ms) != measured: + timing_reasons.append( + f"forced-token timing count={len(steps_ms)} != measured={measured}" + ) + if any(value <= 0 for value in steps_ms): + timing_reasons.append("llama-server returned missing/non-positive forced-token timing") + return { + "schema": "freetoken-replay-v1", + "lane": LANE, + "runtime": "llama-cpp-hip", + "backend": getattr(args, "backend", "hip"), + "repeat": repeat, + "status": "accepted" if not timing_reasons else "rejected", + "acceptance": { + "accepted": not timing_reasons, + "reasons": timing_reasons, + "prompt_ids_match": prompt_ids_match, + "decode_timing_complete": len(steps_ms) == measured, + }, + "execution": "kernel", + "model_sha256": manifest["model_sha256"], + "fixture_sha256": manifest["fixture_sha256"], + "tokenizer_sha256": manifest["tokenizer_sha256"], + "manifest_ids_sha256": manifest["golden"]["ids_sha256"], + "forced": True, + "ids_match": prompt_ids_match, + "prompt_ids_match": prompt_ids_match, + "steps": summarize_steps(steps_ms, 0), + "prompt_ms_per_token_median": ( + statistics.median(prompt_costs) if prompt_costs else None + ), + "decode_ms_per_token_median": ( + statistics.median(steps_ms) if steps_ms else None + ), + "route_digest": None, + "route_hash_status": "unavailable_without_instrumentation", + "timing_domain": "llama_server_prompt_eval_ms_for_forced_token", + "mtp": "off", + "speculative": False, + "decode_batch_size": 1, + "context": getattr(args, "context", 9216), + "batch": getattr(args, "batch", 512), + "ubatch": getattr(args, "ubatch", 512), + "kv_type": getattr(args, "kv_type", "q8_0"), + "server_log": log_path, + } + + +def prompt_text_of(manifest: dict) -> str: + """Return raw prompt text matching pinned FreeToken prompt IDs. + + The golden manifest stores raw prompt text and its hash. This avoids trying to + load a Transformers tokenizer from a GGUF file, which is not a valid tokenizer + source. Legacy manifests may provide ``tokenizer_path`` explicitly.""" + prompt = manifest.get("prompt_text") + if isinstance(prompt, str): + expected = manifest.get("prompt_text_sha256") + if expected and _sha256_bytes(prompt.encode()) != expected: + raise ValueError("prompt_text does not match prompt_text_sha256") + return prompt + tokenizer_path = manifest.get("tokenizer_path") + if not isinstance(tokenizer_path, str): + raise ValueError("manifest lacks prompt_text; regenerate golden manifest with current tool") + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True) + return tokenizer.decode(manifest["prompt_ids"]) + + +def tokens_list(tokens) -> list[int] | None: + if isinstance(tokens, dict): + tokens = tokens.get("tokens") + if not isinstance(tokens, list): + return None + try: + return [int(v) for v in tokens] + except (TypeError, ValueError): + return None + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + sub = parser.add_subparsers(dest="command", required=True) + + def common(p: argparse.ArgumentParser) -> None: + p.add_argument("--model", required=True) + p.add_argument("--aime", default=os.environ.get("FREETOKEN_AIME25_JSONL")) + p.add_argument("--aime-revision", default=os.environ.get("FREETOKEN_AIME25_REVISION")) + p.add_argument("--aime-sha256", default=os.environ.get("FREETOKEN_AIME25_SHA256")) + p.add_argument("--problem", type=int, default=0) + p.add_argument("--context", type=int, default=9216) + p.add_argument("--batch", type=int, default=512) + p.add_argument("--ubatch", type=int, default=512) + p.add_argument("--attention-backend", default="triton") + p.add_argument("--moe-backend", default="offload") + p.add_argument("--memory-ratio", type=float, default=0.9) + p.add_argument("--cache", type=int, default=0) + p.add_argument("--graph", action="store_true", help="capture/replay graphs") + p.add_argument( + "--kv-type", default="q8_0", + help="llama.cpp cache type flags for the replay server (llama side only)", + ) + + golden = sub.add_parser("golden", help="pin legacy greedy golden IDs/hash") + common(golden) + golden.add_argument("--decode", type=int, default=512) + golden.add_argument( + "--warmup", type=int, default=8, help="untimed forced steps before the measured window" + ) + golden.add_argument("--route-top-k", type=int, default=8) + golden.add_argument("--out", required=True, help="manifest JSON to create") + + replay = sub.add_parser("replay-freetoken", help="timed teacher-forced replay") + common(replay) + replay.add_argument("--manifest", required=True) + replay.add_argument("--routes", action="store_true", help="add the untimed route pass") + replay.add_argument("--repeats", type=int, default=10) + replay.add_argument("--json", dest="json_out", default=None) + + rllama = sub.add_parser("replay-llama", help="timed replay against llama-server") + common(rllama) + rllama.add_argument("--manifest", required=True) + rllama.add_argument("--server", required=True, help="llama-server HIP binary") + rllama.add_argument("--repeats", type=int, default=1, help="server restarts (each reruns all steps)") + rllama.add_argument("--timeout", type=float, default=1800) + rllama.add_argument("--json", dest="json_out", default=None) + return parser.parse_args(argv) + + +def _write_rows(rows: list[dict], json_out: str | None) -> None: + if not json_out: + return + with open(json_out, "a") as handle: + for row in rows: + handle.write(json.dumps(row, sort_keys=True) + "\n") + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + if args.command == "golden": + manifest = build_manifest(args) + Path(args.out).write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") + golden = manifest["golden"] + print( + json.dumps( + { + "ids_sha256": golden["ids_sha256"], + "answer": golden["answer"], + "answer_match": golden["answer_match"], + "finite_logits": golden["finite_logits"], + "logit_sha256": golden["logit_sha256"], + }, + indent=2, + ) + ) + return 0 + if args.command == "replay-freetoken": + manifest = load_manifest(args.manifest) + rows = [] + for repeat in range(args.repeats): + args.repeat = repeat + row = replay_freetoken(args, manifest, capture_routes=args.routes) + row["repeat"] = repeat + rows.append(row) + stats = row["steps"] + median = stats.get("ms_per_token_median") if stats else None + median_text = "n/a" if median is None else f"{median:.3f} ms/token" + print( + f"repeat {repeat + 1}/{args.repeats}: {median_text}, " + f"ids_match={row['ids_match']}", + flush=True, + ) + _write_rows(rows, args.json_out) + return 0 if rows and all(row["ids_match"] for row in rows) else 1 + if args.command == "replay-llama": + manifest = load_manifest(args.manifest) + rows = [replay_llama(args, manifest, repeat) for repeat in range(args.repeats)] + _write_rows(rows, args.json_out) + for row in rows: + median = (row["steps"] or {}).get("ms_per_token_median") + print( + f"{row['runtime']}: {median if median is None else f'{median:.3f} ms/token'} " + f"prompt_ids_match={row['prompt_ids_match']}", + flush=True, + ) + return 0 if rows and all(row["status"] == "accepted" for row in rows) else 1 + raise RuntimeError(f"unhandled command {args.command!r}") + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/bench_gguf_linear.py b/benchmarks/bench_gguf_linear.py new file mode 100644 index 000000000..7f39d83f3 --- /dev/null +++ b/benchmarks/bench_gguf_linear.py @@ -0,0 +1,127 @@ +"""Exact-shape GGUF dense GEMV benchmark for Qwen3.6-MoE decode.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import time + +import torch + +from freetoken.kernel.gguf import ggml_mul_mat_a8, ggml_mul_mat_vec_a8 +from freetoken.models.gguf.dequant import GGML_Q6_K, GGML_Q8_0, row_bytes +from freetoken.utils.graph_gate import rocm_blas_report + + +CASES = { + "attn_qkv": (8192, 2048, GGML_Q8_0), + "attn_gate": (4096, 2048, GGML_Q8_0), + "attn_output": (2048, 4096, GGML_Q8_0), + "gdn_out": (2048, 4096, GGML_Q8_0), + "shared_gate_up": (512, 2048, GGML_Q8_0), + "shared_down": (2048, 512, GGML_Q8_0), + "lm_head": (248320, 2048, GGML_Q6_K), +} + + +def _packed_weight(out_features: int, in_features: int, quant_type: int) -> torch.Tensor: + if quant_type == GGML_Q8_0: + blocks = torch.zeros( + (out_features, in_features // 32, 34), dtype=torch.uint8, device="cuda" + ) + # Valid positive fp16 scale, followed by bounded signed int8 payload. + blocks[..., :2] = torch.tensor([128, 63], dtype=torch.uint8, device="cuda") + blocks[..., 2:] = torch.randint( + 0, 255, blocks[..., 2:].shape, dtype=torch.uint8, device="cuda" + ) + return blocks.reshape(out_features, row_bytes(in_features, quant_type)) + if quant_type == GGML_Q6_K: + blocks = torch.zeros( + (out_features, in_features // 256, 210), dtype=torch.uint8, device="cuda" + ) + # Q6_K layout ends with fp16 d; scales and quants may be arbitrary bytes. + blocks[..., 192:208] = torch.randint( + 1, 255, blocks[..., 192:208].shape, dtype=torch.uint8, device="cuda" + ) + blocks[..., 208:210] = torch.tensor([128, 63], dtype=torch.uint8, device="cuda") + return blocks.reshape(out_features, row_bytes(in_features, quant_type)) + raise ValueError(f"unsupported quant type {quant_type}") + + +def _time_call(fn, warmup: int, iters: int) -> tuple[float, torch.Tensor]: + for _ in range(warmup): + out = fn() + torch.cuda.synchronize() + # Host timing avoids event allocation in the measured loop; synchronize once + # after each op so async failures cannot contaminate later samples. + samples = [] + for i in range(iters): + start = time.perf_counter() + out = fn() + torch.cuda.synchronize() + samples.append((time.perf_counter() - start) * 1e6) + if not torch.isfinite(out).all(): + raise RuntimeError("GGUF linear benchmark produced non-finite output") + return statistics.median(samples), out + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--case", choices=[*CASES, "all"], default="all") + parser.add_argument("--iters", type=int, default=30) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--batch", type=int, default=1, help="exact dense batch shape") + parser.add_argument("--json") + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("CUDA/ROCm device required") + if args.batch < 1: + raise SystemExit("--batch must be >= 1") + torch.manual_seed(20260831) + blas = rocm_blas_report(gate={}) + if blas["verification"] == "mismatch": + raise SystemExit(f"requested BLAS policy not effective: {blas}") + + names = list(CASES) if args.case == "all" else [args.case] + rows = [] + generator = torch.Generator(device="cpu").manual_seed(20260831) + for name in names: + out_features, in_features, quant_type = CASES[name] + weight = _packed_weight(out_features, in_features, quant_type) + x = torch.randn( + args.batch, in_features, generator=generator, dtype=torch.bfloat16, device="cpu" + ).to("cuda") + calls = { + "mmvq": lambda: ggml_mul_mat_vec_a8(weight, x, quant_type, out_features), + "mmq": lambda: ggml_mul_mat_a8(weight, x, quant_type, out_features), + } + for impl, fn in calls.items(): + median_us, out = _time_call(fn, args.warmup, args.iters) + rows.append( + { + "case": name, + "impl": impl, + "out_features": out_features, + "in_features": in_features, + "quant_type": quant_type, + "batch": args.batch, + "blas": blas, + "median_us": round(median_us, 3), + "finite": bool(torch.isfinite(out).all().item()), + "output_sha256": hashlib.sha256( + out.detach().cpu().float().numpy().tobytes() + ).hexdigest(), + } + ) + print(json.dumps(rows[-1], sort_keys=True)) + del weight, x + torch.cuda.empty_cache() + if args.json: + with open(args.json, "w", encoding="utf-8") as handle: + json.dump(rows, handle, indent=2) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_gguf_moe_kernels.py b/benchmarks/bench_gguf_moe_kernels.py index f77bb41e0..6edbea21c 100644 --- a/benchmarks/bench_gguf_moe_kernels.py +++ b/benchmarks/bench_gguf_moe_kernels.py @@ -1,4 +1,4 @@ -"""Micro-bench for the GGUF fused-MoE MMVQ kernel pair (Inc 4, .plans/rocm-perf-parity). +"""Micro-bench for the GGUF fused-MoE MMVQ kernel pair (Inc 1, .plans/qwen-moe-speed). Times the two ``ggml_moe_a8_vec`` calls the decode hot loop makes per MoE layer at the Qwen3.6-35B-A3B shapes (H=2048, I=512, top_k=8; Q4_K gate_up + Q8_0 down) on a @@ -27,7 +27,17 @@ import torch from freetoken.layers.activation import silu_and_mul -from freetoken.kernel.gguf import ggml_moe_a8_vec +from freetoken.kernel.gguf import ( + ggml_moe_a8_vec, + ggml_moe_a8_vec_workspace, + ggml_moe_mmvq_id, + ggml_moe_mmvdq_id, + ggml_moe_gate_up_swiglu_id, +) +try: + from benchmarks.lib.paired_stats import paired_summary +except ModuleNotFoundError: # direct ``python benchmarks/bench_gguf_moe_kernels.py`` + from lib.paired_stats import paired_summary GGML_Q4_K = 12 GGML_Q8_0 = 8 @@ -41,9 +51,33 @@ def main(argv: list[str] | None = None) -> int: p.add_argument("--topk", type=int, default=8) p.add_argument("--iters", type=int, default=300) p.add_argument("--warmup", type=int, default=50) + p.add_argument( + "--trace", + default=None, + help="export a profiler trace and print quantize/MMVQ/activation/reduce stages", + ) + p.add_argument("--impl", default=os.environ.get("FREETOKEN_GGUF_MOE_IMPL", "legacy")) + p.add_argument("--down-type", choices=("q8_0", "q6_k"), default="q8_0") + p.add_argument("--fuse-gate-up", action="store_true") p.add_argument("--dump", default=None, help="save outputs to this path") p.add_argument("--check", default=None, help="byte-compare outputs against a saved dump") + p.add_argument( + "--paired", action="store_true", + help="interleave legacy and selected candidate, then emit paired statistics", + ) + p.add_argument("--json", dest="json_out", default=None, help="write result artifact") args = p.parse_args(argv) + if args.impl not in {"legacy", "auto", "gfx1100", "rdna3_mmid", "rdna3_mmvdq"}: + p.error("--impl must be legacy, auto, gfx1100, rdna3_mmid, or rdna3_mmvdq") + if args.impl == "rdna3_mmvdq" and args.down_type != "q6_k": + p.error("rdna3_mmvdq benchmark requires --down-type q6_k") + if args.fuse_gate_up and args.impl != "rdna3_mmid": + p.error("--fuse-gate-up requires --impl rdna3_mmid") + if args.paired and args.impl in {"legacy", "auto"}: + p.error("--paired requires an explicit candidate --impl") + if args.paired and args.fuse_gate_up: + p.error("--paired does not support --fuse-gate-up; compare same unfused pair") + os.environ["FREETOKEN_GGUF_MOE_IMPL"] = args.impl H, I, topk, slots = args.hidden, args.moe_ic, args.topk, args.slots n2 = 2 * I @@ -52,26 +86,116 @@ def main(argv: list[str] | None = None) -> int: # gate_up bank: [slots, 2I, row_bytes(H, Q4_K)] (row bytes = H/256*144); # down bank: [slots, H, row_bytes(I, Q8_0)] (I/32*34) — matches moe/fused_gguf.py. - bank_gu = torch.randint( - 0, 255, (slots, n2, H // 256 * 144), dtype=torch.uint8, generator=gen - ).to(dev) - bank_down = torch.randint( - 0, 255, (slots, H, I // 32 * 34), dtype=torch.uint8, generator=gen - ).to(dev) + gu_blocks = torch.zeros((slots, n2, H // 256, 144), dtype=torch.uint8, device=dev) + gu_blocks[..., :4] = torch.tensor([128, 63, 128, 63], dtype=torch.uint8, device=dev) + # Keep synthetic Q4_K metadata/quants bounded. Arbitrary metadata bytes can + # produce BF16 intermediates large enough to overflow Q8_1's half scale and + # make a byte-comparison gate observe NaN payload differences. + gu_blocks[..., 4:16] = 1 + gu_blocks[..., 16:] = 1 + bank_gu = gu_blocks.reshape(slots, n2, H // 256 * 144) + if args.down_type == "q8_0": + down_blocks = torch.zeros((slots, H, I // 32, 34), dtype=torch.uint8, device=dev) + down_blocks[..., :2] = torch.tensor([128, 63], dtype=torch.uint8, device=dev) + down_blocks[..., 2:] = torch.randint( + 0, 255, (slots, H, I // 32, 32), dtype=torch.uint8, generator=gen + ).to(dev) + down_type = GGML_Q8_0 + else: + down_blocks = torch.randint( + 0, 255, (slots, H, I // 256, 210), dtype=torch.uint8, generator=gen + ).to(dev) + down_blocks[..., -2:] = torch.tensor([128, 63], dtype=torch.uint8, device=dev) + down_type = 14 + down_row_bytes = I // 32 * 34 if args.down_type == "q8_0" else I // 256 * 210 + bank_down = down_blocks.reshape(slots, H, down_row_bytes) x = torch.randn(1, H, generator=gen).to(dev).to(torch.bfloat16) ids = torch.randint(0, slots, (1, topk), generator=gen).to(dev).int() + weights = torch.rand(1, topk, generator=gen).to(dev) + + def q8_1_shape(tokens: int, cols: int) -> tuple[int, int]: + padded = (cols + 512 - 1) // 512 * 512 + return tokens, padded // 32 * 9 + + # Match serving path: both variants reuse fixed graph-address output and + # Q8_1 scratch buffers instead of timing allocator churn. + gate_output = torch.empty((topk, n2), dtype=x.dtype, device=dev) + gate_quant_x = torch.empty(q8_1_shape(1, H), dtype=torch.int32, device=dev) + down_output = torch.empty((topk, H), dtype=x.dtype, device=dev) + down_quant_x = torch.empty(q8_1_shape(topk, I), dtype=torch.int32, device=dev) - def run() -> torch.Tensor: - gate_up = ggml_moe_a8_vec(x, bank_gu, ids, topk, GGML_Q4_K, n2, 1) - inter = silu_and_mul(gate_up.reshape(1 * topk, n2)) - return ggml_moe_a8_vec(inter, bank_down, ids, 1, GGML_Q8_0, H, topk) + def run(impl: str | None = None) -> torch.Tensor: + impl = impl or args.impl + if args.fuse_gate_up and impl == args.impl: + inter = ggml_moe_gate_up_swiglu_id( + x, bank_gu, ids, topk, I, 1, + int(bank_gu.stride(0)), int(bank_gu.stride(1)), "slot", + ) + elif impl == "rdna3_mmvdq": + gate_up = ggml_moe_mmvdq_id( + x, bank_gu, ids, topk, GGML_Q4_K, n2, 1, + int(bank_gu.stride(0)), int(bank_gu.stride(1)), "slot", + ) + elif impl in {"rdna3_mmid", "gfx1100"}: + gate_up = ggml_moe_mmvq_id( + x, bank_gu, ids, topk, GGML_Q4_K, n2, 1, + int(bank_gu.stride(0)), int(bank_gu.stride(1)), "slot", + gate_output, gate_quant_x, + ) + else: + gate_up = ggml_moe_a8_vec_workspace( + x, bank_gu, ids, topk, GGML_Q4_K, n2, 1, + gate_output, gate_quant_x, + ) + if not (args.fuse_gate_up and impl == args.impl): + with torch.profiler.record_function("moe_activation"): + inter = silu_and_mul(gate_up.reshape(1 * topk, n2)) + route_ids = ids.reshape(-1, 1) if impl in {"rdna3_mmid", "rdna3_mmvdq", "gfx1100"} else ids + if impl == "rdna3_mmvdq": + down = ggml_moe_mmvdq_id( + inter, bank_down, route_ids, 1, down_type, H, topk, + int(bank_down.stride(0)), int(bank_down.stride(1)), "slot", + ) + elif impl in {"rdna3_mmid", "gfx1100"}: + down = ggml_moe_mmvq_id( + inter, bank_down, route_ids, 1, down_type, H, topk, + int(bank_down.stride(0)), int(bank_down.stride(1)), "slot", + down_output, down_quant_x, + ) + else: + down = ggml_moe_a8_vec_workspace( + inter, bank_down, ids, 1, down_type, H, topk, + down_output, down_quant_x, + ) + with torch.profiler.record_function("moe_route_reduce"): + return (down.reshape(1, topk, H) * weights.reshape(1, topk, 1)).sum(dim=1) # Warm + reference for the bytes-equality gate (run() is deterministic for a # fixed seed/config: no atomics, one CTA owns every output row). - out = run() - torch.cuda.synchronize() + if args.paired: + os.environ["FREETOKEN_GGUF_MOE_IMPL"] = "legacy" + for _ in range(args.warmup): + run("legacy") + os.environ["FREETOKEN_GGUF_MOE_IMPL"] = args.impl + for _ in range(args.warmup): + run(args.impl) + legacy_samples: list[float] = [] + candidate_samples: list[float] = [] + for _ in range(args.iters): + os.environ["FREETOKEN_GGUF_MOE_IMPL"] = "legacy" + legacy_samples.append(bench_once(lambda: run("legacy"))) + os.environ["FREETOKEN_GGUF_MOE_IMPL"] = args.impl + candidate_samples.append(bench_once(lambda: run(args.impl))) + paired = paired_summary(legacy_samples, candidate_samples) + os.environ["FREETOKEN_GGUF_MOE_IMPL"] = args.impl + out = run(args.impl) + else: + out = run() + torch.cuda.synchronize() + us_pair = bench(run, args.iters, args.warmup) + paired = None torch.cuda.synchronize() - us_pair = bench(run, args.iters) + stages = profile_stages(run, args.trace, args.warmup) if args.trace else {} sha = hashlib.sha256(out.view(torch.uint8).cpu().numpy().tobytes()).hexdigest()[:16] mmv_y = os.environ.get("FREETOKEN_GGUF_MMV_Y", "1") @@ -86,25 +210,109 @@ def run() -> torch.Tensor: return 1 note = f" bytes_equal_to_saved=True" - print( - f"MMV_Y={mmv_y} slots={slots} H={H} I={I} topk={topk} : " - f"{us_pair:.1f} us / layer-pair out_sha={sha}{note}" - ) + result = { + "schema": "freetoken-gguf-moe-microbench-v2", + "impl": args.impl, + "mmv_y": mmv_y, + "slots": slots, + "hidden": H, + "moe_ic": I, + "topk": topk, + "down_type": args.down_type, + "iters": args.iters, + "warmup": args.warmup, + "output_sha256": sha, + "paired": paired, + } + if paired is None: + result["median_us_layer_pair"] = us_pair + print( + f"impl={args.impl} MMV_Y={mmv_y} slots={slots} H={H} I={I} topk={topk} : " + f"{us_pair:.1f} us / layer-pair out_sha={sha}{note}" + ) + else: + print( + f"legacy_vs_{args.impl} pairs={paired['pairs']} " + f"legacy={paired['median_legacy_us']:.1f} us candidate={paired['median_candidate_us']:.1f} us " + f"recovery={paired['median_recovery_us']:.1f} us " + f"CI=[{paired['recovery_p02_5_us']:.1f},{paired['recovery_p97_5_us']:.1f}] us out_sha={sha}{note}" + ) + if stages: + print(" profiler stages: " + ", ".join(f"{k}={v:.2f} us" for k, v in stages.items())) + if args.json_out: + with open(args.json_out, "w") as handle: + import json + json.dump(result, handle, indent=2, sort_keys=True) + handle.write("\n") return 0 -def bench(fn, iters: int) -> float: - for _ in range(50): +def bench_once(fn) -> float: + """Time one already-warmed device call with a fresh event pair.""" + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1e3 + + +def bench(fn, iters: int, warmup: int) -> float: + for _ in range(warmup): fn() torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) ts = [] for _ in range(iters): - t0 = time.perf_counter() + start.record() fn() + end.record() + end.synchronize() + ts.append(start.elapsed_time(end)) + return statistics.median(ts) * 1e3 + + +def profile_stages(fn, trace_path: str, warmup: int) -> dict[str, float]: + """Profile actual device kernels; never infer stages from one fused-call wall timer.""" + activities = [torch.profiler.ProfilerActivity.CPU] + if hasattr(torch.profiler.ProfilerActivity, "CUDA"): + activities.append(torch.profiler.ProfilerActivity.CUDA) + with torch.profiler.profile(activities=activities, record_shapes=True, acc_events=True) as prof: + for _ in range(max(1, warmup // 2)): + fn() + torch.cuda.synchronize() + for _ in range(10): + fn() torch.cuda.synchronize() - ts.append(time.perf_counter() - t0) - return statistics.median(ts) * 1e6 + prof.export_chrome_trace(trace_path) + stage_events: dict[str, list[float]] = { + "q8_quantize": [], + "mmvq": [], + "activation": [], + "route_reduce": [], + } + def device_us(event) -> float: + value = getattr(event, "device_time_total", None) + if value is None: + value = getattr(event, "self_device_time_total", 0.0) + return float(value) + + for event in prof.key_averages(): + name = event.key.lower() + if "moe_activation" in name: + stage_events["activation"].append(event.self_cpu_time_total / max(event.count, 1)) + elif "moe_route_reduce" in name: + stage_events["route_reduce"].append(event.self_cpu_time_total / max(event.count, 1)) + elif "quantize_q8_1" in name: + stage_events["q8_quantize"].append(device_us(event) / max(event.count, 1)) + elif "moe_vec" in name: + stage_events["mmvq"].append(device_us(event) / max(event.count, 1)) + # Profiler time units are microseconds. Multiple symbols can be emitted for one + # stage, so report their sum and retain the trace for exact attribution. + return {name: sum(values) for name, values in stage_events.items() if values} if __name__ == "__main__": - raise SystemExit(main()) \ No newline at end of file + raise SystemExit(main()) diff --git a/benchmarks/bench_llama_cpp_hip.py b/benchmarks/bench_llama_cpp_hip.py new file mode 100644 index 000000000..3209e4d02 --- /dev/null +++ b/benchmarks/bench_llama_cpp_hip.py @@ -0,0 +1,742 @@ +"""Run separate llama.cpp HIP reference lanes for decode parity diagnostics. + +``llama-cli`` owns raw-prompt/output/sampling evidence. ``llama-bench`` owns native +decode throughput only; its rows never close FreeToken's serving parity gate. +Missing binaries are recorded as ``unavailable`` rows when ``--json`` is supplied. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import shlex +import socket +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +from pathlib import Path + +from bench_decode_moe import load_problem_details, model_fingerprint + + +PINNED_COMMIT = "7e4c0a968" +COMPARATOR_CONTEXT = 9216 +COMPARATOR_BATCH = 512 +COMPARATOR_UBATCH = 512 +COMPARATOR_KV_TYPE = "q8_0" + + +_EVAL_RE = re.compile( + r"eval\s+time\s*=.*?/\s*(?P\d+)\s+runs?\s*\([^)]*?" + r"(?P[0-9]+(?:\.[0-9]+)?)\s+tokens?\s+per\s+second", + re.IGNORECASE | re.DOTALL, +) +_JSON_TPS_KEYS = ( + "eval_tokens_per_second", + "tokens_per_second", + "eval_t_s", + "eval_t/s", + "t/s", +) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True, help="exact local GGUF file") + parser.add_argument( + "--cli", + default=os.environ.get("LLAMA_CLI"), + help="llama-cli HIP binary; absent = unavailable lane", + ) + parser.add_argument( + "--bench", + default=os.environ.get("LLAMA_BENCH"), + help="llama-bench HIP binary; absent = unavailable lane", + ) + parser.add_argument( + "--server-rocm", + default=os.environ.get("LLAMA_SERVER_ROCM"), + help="llama-server ROCm binary; absent = unavailable lane", + ) + parser.add_argument( + "--server-vulkan", + default=os.environ.get("LLAMA_SERVER_VULKAN"), + help="llama-server Vulkan binary; absent = unavailable lane", + ) + parser.add_argument("--server-context", type=int, default=9216) + parser.add_argument("--server-batch", type=int, default=512) + parser.add_argument("--server-ubatch", type=int, default=512) + parser.add_argument("--aime", default=os.environ.get("FREETOKEN_AIME25_JSONL")) + parser.add_argument("--aime-revision", default=os.environ.get("FREETOKEN_AIME25_REVISION")) + parser.add_argument("--aime-sha256", default=os.environ.get("FREETOKEN_AIME25_SHA256")) + parser.add_argument("--problem", type=int, default=0) + parser.add_argument("--decode", type=int, default=512) + parser.add_argument("--repeats", type=int, default=10) + parser.add_argument( + "--commit", + default=os.environ.get("LLAMA_CPP_COMMIT", PINNED_COMMIT), + help="expected llama.cpp source commit; reference builder defaults to b10434", + ) + parser.add_argument("--greedy", action="store_true") + parser.add_argument( + "--replay-manifest", + default=None, + help="run the fixed-token teacher-forced replay lane from this manifest " + "against the supplied ROCm server binary instead of the decode lanes", + ) + parser.add_argument("--timeout", type=float, default=1800) + parser.add_argument("--json", dest="json_out", default=None) + return parser.parse_args(argv) + + +def _version(binary: str) -> str | None: + try: + result = subprocess.run( + [binary, "--version"], + check=False, + capture_output=True, + text=True, + timeout=30, + ) + except (OSError, subprocess.TimeoutExpired): + return None + text = (result.stdout + result.stderr).strip() + return text[:1000] or None + + +def _binary_commit(version: str | None) -> str | None: + if not version: + return os.environ.get("LLAMA_CPP_COMMIT") + match = re.search(r"\b(?:commit|rev(?:ision)?)\s*[:=]?\s*([0-9a-f]{7,40})\b", version, re.I) + return match.group(1) if match else os.environ.get("LLAMA_CPP_COMMIT") + + +def _base_row( + *, + lane: str, + args: argparse.Namespace, + model_sha256: str | None, + prompt_sha256: str, + dataset: dict, + binary: str | None, + binary_version: str | None, + backend: str = "hip", + schema: str = "llama-cpp-hip-decode-v1", +) -> dict: + return { + "schema": schema, + "status": "unavailable", + "lane": lane, + "binary": binary, + "binary_version": binary_version, + "binary_commit": _binary_commit(binary_version), + "expected_commit": args.commit, + "backend": backend, + "model": str(Path(args.model).expanduser().resolve()), + "model_sha256": model_sha256, + "model_identity": "full-file-sha256" if model_sha256 else "unverified", + "dataset": dataset, + "prompt_sha256": prompt_sha256, + "prompt_protocol": "raw_user_prompt", + "sampling": { + "temperature": 0.0 if args.greedy else 1.0, + "top_p": 1.0 if args.greedy else 0.95, + "top_k": -1 if args.greedy else 64, + "seed": 0, + }, + "decode_requested": args.decode, + "context": args.server_context, + "batch": args.server_batch, + "ubatch": args.server_ubatch, + "kv_type": COMPARATOR_KV_TYPE, + "fixture_sha256": dataset.get("sha256"), + "placement": { + "ngl": 99, + "flash_attention": True, + "extra_args": os.environ.get("LLAMA_SERVER_EXTRA_ARGS", "").strip(), + }, + "mtp": "off", + "speculative": False, + "eval_count": None, + "native_decode_tok_s": None, + "client_arrival_tok_s": None, + "output_sha256": None, + "reference_status": "reference_only", + "acceptance": { + "status": "unavailable", + "accepted": False, + "checks": { + "exact_completion_count": False, + "model_sha256": bool(model_sha256), + "finite_output": False, + "prompt_options_match": False, + "fixture_sha256": bool(dataset.get("sha256")), + "exact_comparator_config": ( + args.decode == 512 + and args.server_context == COMPARATOR_CONTEXT + and args.server_batch == COMPARATOR_BATCH + and args.server_ubatch == COMPARATOR_UBATCH + ), + "mtp_off": True, + "speculative_off": True, + "native_only_for_llama_bench": lane == "llama-bench", + }, + "reasons": [], + }, + } + + +def _parse_cli_timing(output: str) -> tuple[int, float] | None: + matches = list(_EVAL_RE.finditer(output)) + if not matches: + return None + match = matches[-1] + return int(match.group("count")), float(match.group("tps")) + + +def _run_cli( + binary: str, + *, + args: argparse.Namespace, + prompt: str, + model_sha256: str | None, + prompt_sha256: str, + dataset: dict, + repeat: int, +) -> dict: + row = _base_row( + lane="llama-cli", + args=args, + model_sha256=model_sha256, + prompt_sha256=prompt_sha256, + dataset=dataset, + binary=binary, + binary_version=_version(binary), + ) + command = [ + binary, + "-m", + args.model, + "-p", + prompt, + "-n", + str(args.decode), + "--temp", + str(row["sampling"]["temperature"]), + "--top-p", + str(row["sampling"]["top_p"]), + "--top-k", + str(row["sampling"]["top_k"]), + "--seed", + str(row["sampling"]["seed"]), + "--no-display-prompt", + "--no-conversation", + "-c", str(args.server_context), + "-b", str(args.server_batch), + "-ub", str(args.server_ubatch), + "-ngl", "99", + "-fa", "on", + "-ctk", "q8_0", + "-ctv", "q8_0", + ] + try: + result = subprocess.run( + command, + check=False, + capture_output=True, + text=True, + timeout=args.timeout, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + row["status"] = "rejected" + row["acceptance"]["status"] = "rejected" + row["acceptance"]["reasons"].append(f"binary execution failed: {type(exc).__name__}") + return row + output = result.stdout + result.stderr + timing = _parse_cli_timing(output) + row["repeat"] = repeat + row["command"] = command + row["returncode"] = result.returncode + row["output_sha256"] = hashlib.sha256(result.stdout.encode()).hexdigest() + if result.returncode != 0: + row["status"] = "rejected" + row["acceptance"]["status"] = "rejected" + row["acceptance"]["reasons"].append(f"binary exited {result.returncode}") + return row + if timing is None: + row["status"] = "rejected" + row["acceptance"]["status"] = "rejected" + row["acceptance"]["reasons"].append("llama-cli eval timing/count not found") + return row + eval_count, tok_s = timing + row["eval_count"] = eval_count + row["native_decode_tok_s"] = tok_s + checks = row["acceptance"]["checks"] + checks["exact_completion_count"] = eval_count == args.decode + checks["model_sha256"] = bool(model_sha256) + checks["finite_output"] = bool(result.stdout) + checks["prompt_options_match"] = True + reasons = row["acceptance"]["reasons"] + if eval_count != args.decode: + reasons.append(f"eval_count={eval_count} != --decode {args.decode}") + if not model_sha256: + reasons.append("model full SHA-256 unavailable") + if dataset.get("sha256") is None: + reasons.append("fixture full SHA-256 unavailable") + if not checks["exact_comparator_config"]: + reasons.append("comparator requires decode=512 context=9216 batch=512 ubatch=512") + if row["binary_commit"] != args.commit: + reasons.append( + f"binary commit={row['binary_commit']!r} != expected {args.commit!r}" + ) + if not result.stdout: + reasons.append("llama-cli emitted no stdout") + row["status"] = "accepted" if not reasons else "rejected" + row["acceptance"]["status"] = row["status"] + row["acceptance"]["accepted"] = not reasons + return row + + +def _find_bench_result(stdout: str) -> tuple[int | None, float | None]: + for line in stdout.splitlines(): + try: + value = json.loads(line) + except json.JSONDecodeError: + continue + if not isinstance(value, dict): + continue + count = next( + (value[key] for key in ("n_gen", "eval_count", "tokens") if key in value), + None, + ) + tps = next((value[key] for key in _JSON_TPS_KEYS if key in value), None) + if isinstance(tps, (int, float)): + return int(count) if isinstance(count, (int, float)) else None, float(tps) + match = re.search( + r"(?P\d+)\s+tokens?\s+per\s+second.*?(?P[0-9]+(?:\.[0-9]+)?)", + stdout, + re.IGNORECASE, + ) + return ( + (int(match.group("count")) if match else None), + (float(match.group("tps")) if match else None), + ) + + +def _run_bench( + binary: str, + *, + args: argparse.Namespace, + model_sha256: str | None, + prompt_sha256: str, + dataset: dict, +) -> dict: + row = _base_row( + lane="llama-bench", + args=args, + model_sha256=model_sha256, + prompt_sha256=prompt_sha256, + dataset=dataset, + binary=binary, + binary_version=_version(binary), + ) + command = [binary, "-m", args.model, "-n", str(args.decode), "-r", str(args.repeats), "-o", "json"] + try: + result = subprocess.run( + command, + check=False, + capture_output=True, + text=True, + timeout=args.timeout, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + row["status"] = "rejected" + row["acceptance"]["status"] = "rejected" + row["acceptance"]["reasons"].append(f"binary execution failed: {type(exc).__name__}") + return row + count, tok_s = _find_bench_result(result.stdout + result.stderr) + row["command"] = command + row["returncode"] = result.returncode + row["eval_count"] = count + row["native_decode_tok_s"] = tok_s + checks = row["acceptance"]["checks"] + checks["exact_completion_count"] = count == args.decode + checks["model_sha256"] = bool(model_sha256) + checks["finite_output"] = True + checks["prompt_options_match"] = False + reasons = row["acceptance"]["reasons"] + if result.returncode != 0: + reasons.append(f"binary exited {result.returncode}") + if count != args.decode: + reasons.append(f"eval_count={count!r} != --decode {args.decode}") + if tok_s is None: + reasons.append("llama-bench JSON eval throughput not found") + if not model_sha256: + reasons.append("model full SHA-256 unavailable") + if dataset.get("sha256") is None: + reasons.append("fixture full SHA-256 unavailable") + if not checks["exact_comparator_config"]: + reasons.append("comparator requires decode=512 context=9216 batch=512 ubatch=512") + if row["binary_commit"] != args.commit: + reasons.append( + f"binary commit={row['binary_commit']!r} != expected {args.commit!r}" + ) + row["status"] = "accepted" if not reasons else "rejected" + row["acceptance"]["status"] = row["status"] + row["acceptance"]["accepted"] = not reasons + row["acceptance"]["reasons"] = reasons + return row + + +def _unavailable(lane: str, args: argparse.Namespace, model_sha256: str | None, prompt_sha256: str, dataset: dict) -> dict: + row = _base_row( + lane=lane, + args=args, + model_sha256=model_sha256, + prompt_sha256=prompt_sha256, + dataset=dataset, + binary=None, + binary_version=None, + ) + row["acceptance"]["reasons"].append("binary not supplied") + return row + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _server_command(binary: str, args: argparse.Namespace, port: int) -> list[str]: + command = [ + binary, + "-m", args.model, + "--host", "127.0.0.1", + "--port", str(port), + "-c", str(args.server_context), + "-b", str(args.server_batch), + "-ub", str(args.server_ubatch), + "-np", "1", + "-ngl", "99", + "-fa", "on", + "-ctk", "q8_0", + "-ctv", "q8_0", + ] + extra = os.environ.get("LLAMA_SERVER_EXTRA_ARGS", "").strip() + if extra: + command.extend(shlex.split(extra)) + return command + + +def _server_json(origin: str, path: str, *, timeout: float = 10) -> dict: + with urllib.request.urlopen(f"{origin}{path}", timeout=timeout) as response: + return json.load(response) + + +def _wait_server(origin: str, proc: subprocess.Popen, log_path: str, timeout: float) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if proc.poll() is not None: + raise RuntimeError(f"llama-server exited {proc.returncode}; log={log_path}") + try: + health = _server_json(origin, "/health", timeout=5) + if health.get("status") in ("ok", "ready"): + return + except (OSError, ValueError, urllib.error.HTTPError): + pass + time.sleep(1) + raise RuntimeError(f"llama-server not ready after {timeout:.0f}s; log={log_path}") + + +def _stop_server(proc: subprocess.Popen) -> None: + if proc.poll() is not None: + return + try: + proc.terminate() + proc.wait(timeout=30) + except (ProcessLookupError, subprocess.TimeoutExpired): + try: + proc.kill() + proc.wait(timeout=10) + except (ProcessLookupError, subprocess.TimeoutExpired): + pass + + +def _stream_server(origin: str, model: str, prompt: str, args: argparse.Namespace) -> dict: + body = { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": args.decode, + "ignore_eos": True, + "stream": True, + "stream_options": {"include_usage": True}, + "temperature": 0.0 if args.greedy else 1.0, + "top_p": 1.0 if args.greedy else 0.95, + "top_k": -1 if args.greedy else 64, + "seed": 0, + } + request = urllib.request.Request( + f"{origin}/v1/chat/completions", + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, + ) + stamps: list[float] = [] + pieces: list[str] = [] + usage: dict = {} + started = time.perf_counter() + with urllib.request.urlopen(request, timeout=args.timeout) as response: + for raw in response: + line = raw.strip() + if not line.startswith(b"data:"): + continue + payload = line[len(b"data:"):].strip() + if payload == b"[DONE]": + break + chunk = json.loads(payload) + if isinstance(chunk.get("usage"), dict): + usage = chunk["usage"] + for choice in chunk.get("choices", []): + delta = choice.get("delta") or {} + text = "".join( + part for part in (delta.get("reasoning_content"), delta.get("content")) if part + ) + if text: + stamps.append(time.perf_counter()) + pieces.append(text) + return { + "t0": started, + "stamps": stamps, + "text": "".join(pieces), + "usage": usage, + } + + +def _server_native_timing(log_path: str) -> tuple[int, float] | None: + text = Path(log_path).read_text(errors="replace") + return _parse_cli_timing(text) + + +def _run_server( + binary: str, + *, + backend: str, + args: argparse.Namespace, + prompt: str, + model_sha256: str | None, + prompt_sha256: str, + dataset: dict, +) -> list[dict]: + port = _free_port() + origin = f"http://127.0.0.1:{port}" + lane = f"llama-server-{backend}" + version = _version(binary) + rows: list[dict] = [] + with tempfile.NamedTemporaryFile(prefix=f"{lane}-", suffix=".log", delete=False) as log: + log_path = log.name + command = _server_command(binary, args, port) + log_handle = open(log_path, "wb") + proc = subprocess.Popen( + command, + stdout=log_handle, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + try: + _wait_server(origin, proc, log_path, args.timeout) + model_id = _server_json(origin, "/v1/models").get("data", [{}])[0].get("id", args.model) + _stream_server(origin, model_id, prompt, args) + for repeat in range(args.repeats): + row = _base_row( + lane=lane, + args=args, + model_sha256=model_sha256, + prompt_sha256=prompt_sha256, + dataset=dataset, + binary=binary, + binary_version=version, + backend=backend, + schema="llama-cpp-server-decode-v1", + ) + row["repeat"] = repeat + row["command"] = command + try: + result = _stream_server(origin, model_id, prompt, args) + except (OSError, ValueError, urllib.error.HTTPError, json.JSONDecodeError) as exc: + row["status"] = "rejected" + row["acceptance"]["status"] = "rejected" + row["acceptance"]["reasons"].append(f"server request failed: {type(exc).__name__}") + rows.append(row) + continue + usage = result["usage"] + count = usage.get("completion_tokens") + stamps = result["stamps"] + client_window = stamps[-1] - stamps[0] if len(stamps) >= 2 else 0.0 + row["eval_count"] = count + row["client_arrival_tok_s"] = ( + (count - 1) / client_window if isinstance(count, int) and client_window > 0 else None + ) + row["client_arrival_window_s"] = client_window + row["output_sha256"] = hashlib.sha256(result["text"].encode()).hexdigest() + native = _server_native_timing(log_path) + if native: + row["native_eval_count"], row["native_decode_tok_s"] = native + checks = row["acceptance"]["checks"] + checks["exact_completion_count"] = count == args.decode + checks["model_sha256"] = bool(model_sha256) + checks["finite_output"] = bool(result["text"]) + checks["prompt_options_match"] = True + reasons = row["acceptance"]["reasons"] + if count != args.decode: + reasons.append(f"completion_tokens={count!r} != --decode {args.decode}") + if len(stamps) < 2: + reasons.append(f"need >=2 output events, got {len(stamps)}") + if not result["text"]: + reasons.append("llama-server emitted no output") + if not model_sha256: + reasons.append("model full SHA-256 unavailable") + if dataset.get("sha256") is None: + reasons.append("fixture full SHA-256 unavailable") + if not checks["exact_comparator_config"]: + reasons.append("comparator requires decode=512 context=9216 batch=512 ubatch=512") + if row["binary_commit"] != args.commit: + reasons.append( + f"binary commit={row['binary_commit']!r} != expected {args.commit!r}" + ) + row["status"] = "accepted" if not reasons else "rejected" + row["acceptance"]["status"] = row["status"] + row["acceptance"]["accepted"] = not reasons + rows.append(row) + except (OSError, ValueError, RuntimeError, urllib.error.HTTPError) as exc: + row = _unavailable(lane, args, model_sha256, prompt_sha256, dataset) + row["binary"] = binary + row["binary_version"] = version + row["backend"] = backend + row["status"] = "rejected" + row["acceptance"]["status"] = "rejected" + row["acceptance"]["reasons"] = [f"server startup failed: {type(exc).__name__}: {exc}"] + rows.append(row) + finally: + _stop_server(proc) + log_handle.close() + for row in rows: + row["server_log"] = log_path + return rows + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + if args.decode < 2 or args.repeats < 1: + raise SystemExit("--decode must be >= 2 and --repeats must be >= 1") + if args.replay_manifest: + # Delegated fixed-token replay lane: one common token sequence per + # runtime, timing logits without sampling (rocm-ollama-gap Inc 0). + from bench_decode_replay import load_manifest, replay_llama + + manifest = load_manifest(args.replay_manifest) + rows = [] + for backend, binary in (("rocm", args.server_rocm), ("vulkan", args.server_vulkan)): + if not binary: + continue + replay_args = argparse.Namespace(**vars(args)) + replay_args.server = binary + replay_args.backend = backend + replay_args.context = args.server_context + replay_args.batch = args.server_batch + replay_args.ubatch = args.server_ubatch + replay_args.kv_type = COMPARATOR_KV_TYPE + rows.extend( + replay_llama(replay_args, manifest, repeat) + for repeat in range(args.repeats) + ) + if args.json_out: + with open(args.json_out, "a") as handle: + for row in rows: + handle.write(json.dumps(row, sort_keys=True) + "\n") + for row in rows: + print( + f"{row['runtime']}: steps={row['steps'].get('steps')} " + f"prompt_ids_match={row['prompt_ids_match']}", + flush=True, + ) + return 0 if rows and all( + row["prompt_ids_match"] + and row["steps"].get("steps") == manifest["measured_tokens"] + for row in rows + ) else 1 + model = model_fingerprint(args.model) + model_sha256 = model.get("sha256") + prompt, _, dataset = load_problem_details( + args.aime, args.problem, args.aime_revision, args.aime_sha256 + ) + prompt_sha256 = hashlib.sha256(prompt.encode("utf-8")).hexdigest() + rows = [] + if args.cli: + rows.extend( + _run_cli( + args.cli, + args=args, + prompt=prompt, + model_sha256=model_sha256, + prompt_sha256=prompt_sha256, + dataset=dataset, + repeat=repeat, + ) + for repeat in range(args.repeats) + ) + else: + rows.append(_unavailable("llama-cli", args, model_sha256, prompt_sha256, dataset)) + if args.bench: + rows.append( + _run_bench( + args.bench, + args=args, + model_sha256=model_sha256, + prompt_sha256=prompt_sha256, + dataset=dataset, + ) + ) + else: + rows.append(_unavailable("llama-bench", args, model_sha256, prompt_sha256, dataset)) + for backend, binary in (("rocm", args.server_rocm), ("vulkan", args.server_vulkan)): + if binary: + rows.extend( + _run_server( + binary, + backend=backend, + args=args, + prompt=prompt, + model_sha256=model_sha256, + prompt_sha256=prompt_sha256, + dataset=dataset, + ) + ) + else: + rows.append(_unavailable( + f"llama-server-{backend}", args, model_sha256, prompt_sha256, dataset + )) + for row in rows: + print( + f"{row['lane']}: {row['status']} native={row['native_decode_tok_s']} " + f"count={row['eval_count']}", + flush=True, + ) + if args.json_out: + with open(args.json_out, "a") as handle: + for row in rows: + handle.write(json.dumps(row, sort_keys=True) + "\n") + if any(row["status"] == "accepted" for row in rows): + return 0 + if all(row["status"] == "unavailable" for row in rows): + return 0 + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/check_decode_gate.py b/benchmarks/check_decode_gate.py new file mode 100644 index 000000000..a27bc9979 --- /dev/null +++ b/benchmarks/check_decode_gate.py @@ -0,0 +1,693 @@ +"""Evaluate final FreeToken base-decode promotion evidence without heuristics.""" + +from __future__ import annotations + +import argparse +import json +import random +import re +import statistics +from pathlib import Path +from typing import Iterable + + +COMPARATOR_CONTEXT = 9216 +COMPARATOR_BATCH = 512 +COMPARATOR_UBATCH = 512 +COMPARATOR_KV_TYPE = "q8_0" + +# Disjoint measurement lanes (rocm-ollama-gap Inc 0). A row belongs to exactly +# one lane; mixing lanes in one promotion claim is always a rejection. +LANE_SAMPLED = "sampled_absolute" +LANE_GREEDY = "greedy_correctness" +LANE_REPLAY = "teacher_forced_replay" +LANES = (LANE_SAMPLED, LANE_GREEDY, LANE_REPLAY) + + +def _row_lane(row: dict) -> str | None: + """The lane a row belongs to, or None when the row is unclassifiable.""" + lane = row.get("lane") + if lane in LANES: + return str(lane) + metadata = _mapping(row.get("metadata")) + meta_lane = metadata.get("lane") + if isinstance(meta_lane, str) and meta_lane in LANES: + return meta_lane + # Legacy freetoken-base rows predate explicit lanes: classify by sampling. + if str(row.get("schema", "")).startswith("freetoken-base-"): + sampling = _sampling(row) + if sampling is not None and sampling.get("temperature") == 0: + return LANE_GREEDY + return LANE_SAMPLED + # Older gate fixtures had no schema/lane field but did carry the sampled + # throughput shape. Preserve their meaning without accepting arbitrary + # unlabeled rows. + sampling = _sampling(row) + if isinstance(row.get("decode_tok_s"), (int, float)) and sampling is not None: + return LANE_GREEDY if sampling.get("temperature") == 0 else LANE_SAMPLED + return None + + +def _lane_reasons(rows: list[dict]) -> list[str]: + """Reject mixed or unclassifiable lanes before any promotion arithmetic.""" + reasons: list[str] = [] + lanes: set[str] = set() + for index, row in enumerate(rows): + lane = _row_lane(row) + if lane is None: + reasons.append(f"accepted row {index} has unknown/mismatched lane identity") + classified = {lane for lane in (_row_lane(row) for row in rows) if lane} + if len(classified) > 1: + reasons.append( + f"rows mix disjoint lanes: {sorted(classified)}; lanes cannot be combined" + ) + return reasons + + +def load_jsonl(path: str) -> list[dict]: + rows = [] + for line_no, line in enumerate(Path(path).read_text().splitlines(), 1): + if not line.strip(): + continue + value = json.loads(line) + if not isinstance(value, dict): + raise ValueError(f"{path}:{line_no}: expected JSON object") + rows.append(value) + return rows + + +def _accepted(rows: Iterable[dict], *, lane: str | None = None) -> list[dict]: + return [ + row for row in rows + if row.get("status") == "accepted" + and (lane is None or row.get("lane") == lane) + ] + + +def _mapping(value: object) -> dict: + return value if isinstance(value, dict) else {} + + +def _sha256(value: object) -> str | None: + if isinstance(value, str) and re.fullmatch(r"[0-9a-fA-F]{64}", value): + return value.lower() + return None + + +def _model_sha(row: dict) -> str | None: + metadata = _mapping(row.get("metadata")) + fingerprint = _mapping(row.get("model_fingerprint")) + for value in (row.get("model_sha256"), fingerprint.get("sha256"), metadata.get("model_sha256")): + digest = _sha256(value) + if digest: + return digest + return None + + +def _prompt_sha(row: dict) -> str | None: + metadata = _mapping(row.get("metadata")) + return _sha256(row.get("prompt_sha256")) or _sha256(metadata.get("prompt_sha256")) + + +def _comparator_value(row: dict, key: str): + metadata = _mapping(row.get("metadata")) + return row.get(key, metadata.get(key)) + + +def _comparator_reasons(rows: list[dict], *, label: str) -> tuple[list[str], str | None]: + """Require one exact benchmark contract before promotion arithmetic.""" + reasons: list[str] = [] + fixture_shas: set[str] = set() + expected = { + "context": COMPARATOR_CONTEXT, + "batch": COMPARATOR_BATCH, + "ubatch": COMPARATOR_UBATCH, + "kv_type": COMPARATOR_KV_TYPE, + } + for index, row in enumerate(rows): + for key, wanted in expected.items(): + value = _comparator_value(row, key) + if value != wanted: + reasons.append( + f"{label} row {index} {key}={value!r} != comparator {wanted!r}" + ) + fixture = _sha256(_comparator_value(row, "fixture_sha256")) + if fixture is None: + reasons.append(f"{label} row {index} missing full fixture SHA-256") + else: + fixture_shas.add(fixture) + if len(fixture_shas) != 1: + reasons.append(f"{label} rows do not carry one full fixture SHA-256") + return reasons, next(iter(fixture_shas), None) + + +def _sampling(row: dict) -> dict | None: + metadata = _mapping(row.get("metadata")) + for value in (row.get("sampling"), row.get("options"), metadata.get("sampling")): + if isinstance(value, dict): + return value + return None + + +def _execution(row: dict) -> dict | None: + metadata = _mapping(row.get("metadata")) + value = row.get("execution", metadata.get("execution")) + return value if isinstance(value, dict) else None + + +def _execution_reasons(rows: list[dict]) -> list[str]: + """Require one observed eligible execution class, never requested CLI flags.""" + reasons: list[str] = [] + required = ( + "effective_moe_backend", "expert_storage", "resident_gguf", "expert_fetches", + "expert_remaps", "attention_backend", "graph_state", "decode_batch_size", + "mtp", "speculative", + ) + fingerprints = set() + for index, row in enumerate(rows): + execution = _execution(row) + if execution is None: + reasons.append(f"accepted row {index} missing execution-mode evidence") + continue + missing = [key for key in required if key not in execution] + if missing: + reasons.append(f"accepted row {index} execution evidence missing: {','.join(missing)}") + continue + execution_class = execution.get("execution_class") + # Older accepted fixtures predate the explicit class. Their complete + # resident evidence maps unambiguously to resident_fused. + if execution_class is None and execution["effective_moe_backend"] == "fused": + execution_class = "resident_fused" + if execution_class == "resident_fused": + if execution["effective_moe_backend"] != "fused": + reasons.append(f"accepted row {index} resident_fused backend is not fused") + if execution["expert_storage"] != "resident_gguf" or execution["resident_gguf"] is not True: + reasons.append(f"accepted row {index} does not prove resident native GGUF experts") + if execution["expert_fetches"] != 0 or execution["expert_remaps"] != 0: + reasons.append(f"accepted row {index} reports expert fetch/remap activity") + elif execution_class == "offload_warm": + if execution["effective_moe_backend"] not in {"offload", "hybrid"}: + reasons.append(f"accepted row {index} offload_warm backend is invalid") + for key in ("expert_fetch_bytes", "expert_copy_bytes"): + if execution.get(key) != 0: + reasons.append(f"accepted row {index} reports nonzero warm {key}") + else: + reasons.append(f"accepted row {index} execution_class={execution_class!r} is not eligible") + if execution["graph_state"] != "replay": + reasons.append(f"accepted row {index} graph_state={execution['graph_state']!r} != 'replay'") + if execution["decode_batch_size"] != 1: + reasons.append(f"accepted row {index} decode_batch_size is not 1") + if execution["mtp"] != "off" or execution["speculative"] is not False: + reasons.append(f"accepted row {index} enables MTP/speculative execution") + fingerprints.add( + (execution_class, *(execution.get(key) for key in required)) + ) + if len(fingerprints) > 1: + reasons.append("accepted rows mix execution modes") + return reasons + + +def _same_sampling(left: dict, right: dict) -> bool: + keys = ("temperature", "top_p", "top_k") + return isinstance(left, dict) and isinstance(right, dict) and all( + left.get(key) == right.get(key) for key in keys + ) + + +def _bootstrap_p02_5(values: list[float], *, samples: int = 10_000, seed: int = 0) -> float: + if not values: + return float("nan") + rng = random.Random(seed) + medians = [statistics.median(rng.choices(values, k=len(values))) for _ in range(samples)] + medians.sort() + return medians[min(len(medians) - 1, int(samples * 0.025))] + + +def _reference( + rows: list[dict], + model_sha: str, + prompt_sha: str | None, + sampling: dict, + fixture_sha: str | None, +) -> tuple[dict | None, list[str]]: + reasons = [] + # Ollama only becomes eligible when both local files were hashed and matched. A + # manifest digest alone is deliberately insufficient. + ollama = [ + row for row in _accepted(rows) + if row.get("schema", "").startswith("ollama-") + and row.get("mtp") == "off" + and row.get("speculative") in (False, "off") + and _mapping(row.get("acceptance")).get("accepted") is True + and _mapping(_mapping(row.get("acceptance")).get("checks")).get("mtp_off") is True + and _mapping(row.get("reference_identity")).get("status") == "verified" + and _mapping(row.get("reference_identity")).get("same_blob") is True + and row.get("prompt_sha256") == prompt_sha + and _sha256(_comparator_value(row, "fixture_sha256")) == fixture_sha + and not _comparator_reasons([row], label="reference")[0] + and isinstance(row.get("client_arrival_tok_s"), (int, float)) + ] + if ollama: + identity = ollama[0]["reference_identity"] + reference_gguf = _mapping(identity).get("reference_gguf") + if _sha256(_mapping(reference_gguf).get("sha256")) != model_sha: + reasons.append("verified Ollama identity does not match FreeToken model SHA-256") + elif not all(_sampling(row) and _same_sampling(_sampling(row), sampling) for row in ollama): + reasons.append("Ollama sampling options do not match FreeToken") + else: + values = [float(row["client_arrival_tok_s"]) for row in ollama] + return { + "kind": "ollama-client-arrival", + "sha256": model_sha, + "rows": len(values), + "median_tok_s": statistics.median(values), + "identity": identity, + "by_repeat": { + row.get("repeat"): float(row["client_arrival_tok_s"]) + for row in ollama if row.get("repeat") is not None + }, + }, reasons + + cli = [ + row for row in _accepted(rows, lane="llama-cli") + if row.get("model_sha256") == model_sha + and row.get("prompt_sha256") == prompt_sha + and _sha256(_comparator_value(row, "fixture_sha256")) == fixture_sha + and not _comparator_reasons([row], label="reference")[0] + and row.get("mtp") == "off" + and row.get("speculative") is False + and _mapping(row.get("acceptance")).get("accepted") is True + and _mapping(_mapping(row.get("acceptance")).get("checks")).get("mtp_off") is True + and row.get("eval_count") == row.get("decode_requested") + and isinstance(row.get("native_decode_tok_s"), (int, float)) + and _sampling(row) is not None + and _same_sampling(_sampling(row), sampling) + ] + if cli: + values = [float(row["native_decode_tok_s"]) for row in cli] + return { + "kind": "llama-cli-native", + "sha256": model_sha, + "rows": len(values), + "median_tok_s": statistics.median(values), + "identity": "full-file-sha256", + "by_repeat": { + row.get("repeat"): float(row["native_decode_tok_s"]) + for row in cli if row.get("repeat") is not None + }, + }, reasons + reasons.append("no matched full-SHA llama-cli or verified same-blob Ollama reference") + return None, reasons + + +def _probe_reasons(probe: object, model_sha: str | None, prompt_sha: str | None) -> list[str]: + """Validate finite-logit and eager/graph evidence from the offline probe.""" + if not isinstance(probe, dict): + return ["finite-logit/parity probe unavailable"] + reasons: list[str] = [] + probe_model = _sha256(_mapping(probe.get("model")).get("sha256")) + if not model_sha or probe_model != model_sha: + reasons.append("finite-logit probe model SHA-256 does not match FreeToken rows") + probe_prompt = _sha256(probe.get("prompt_sha256")) + if not prompt_sha or probe_prompt != prompt_sha: + reasons.append("finite-logit probe prompt SHA-256 does not match FreeToken rows") + if probe.get("mtp") != "off" or probe.get("speculative") is not False: + reasons.append("finite-logit probe did not prove MTP/speculative mode is off") + + lanes = [] + for name in ("eager", "graph"): + lane = probe.get(name) + if lane is None: + continue + if not isinstance(lane, dict): + reasons.append(f"{name} finite-logit probe record is malformed") + continue + lanes.append(name) + if lane.get("finite_logits") is not True: + reasons.append(f"{name} finite-logit probe failed") + if not isinstance(lane.get("decode_rows"), int) or lane["decode_rows"] < 1: + reasons.append(f"{name} finite-logit probe has no decode rows") + if not lanes: + reasons.append("finite-logit probe has no eager or graph lane") + comparison = _mapping(probe.get("comparison")) + if probe.get("graph") is not None and comparison.get("token_ids_equal") is not True: + reasons.append("eager/graph greedy token parity failed or is unavailable") + return reasons + + +def _paired_deltas(rows: list[dict], reference: dict | None) -> list[float]: + if not reference: + return [] + by_repeat = reference.get("by_repeat") or {} + deltas = [] + for row in rows: + key = row.get("repeat") + value = row.get("decode_tok_s") + ref_value = by_repeat.get(key) + if isinstance(value, (int, float)) and isinstance(ref_value, (int, float)): + deltas.append(float(value) - float(ref_value)) + return deltas + + +def _bootstrap_interval( + values: list[float], *, samples: int = 10_000, seed: int = 0 +) -> tuple[float, float]: + if not values: + return float("nan"), float("nan") + rng = random.Random(seed) + medians = [statistics.median(rng.choices(values, k=len(values))) for _ in range(samples)] + medians.sort() + low = medians[min(len(medians) - 1, int(samples * 0.025))] + high = medians[min(len(medians) - 1, int(samples * 0.975))] + return low, high + + +def _evaluate_sampled_gate( + freetoken_rows: list[dict], + reference_rows: list[dict] | None = None, + probe: dict | None = None, + *, + min_runs: int = 10, + directional_threshold: float = 80.225, +) -> dict: + reference_rows = reference_rows or [] + rejected = [row for row in freetoken_rows if row.get("status") != "accepted"] + accepted = _accepted(freetoken_rows) + reasons: list[str] = [] + if rejected: + reasons.append(f"{len(rejected)} FreeToken rows rejected") + if len(accepted) < min_runs: + reasons.append(f"need >= {min_runs} accepted FreeToken rows, got {len(accepted)}") + + model_shas = {_model_sha(row) for row in accepted} + model_shas.discard(None) + model_sha = next(iter(model_shas), None) if len(model_shas) == 1 else None + if len(model_shas) != 1: + reasons.append("FreeToken rows do not carry one full model SHA-256") + prompt_shas = {_prompt_sha(row) for row in accepted} + if len(prompt_shas) != 1 or None in prompt_shas: + reasons.append("FreeToken rows do not carry one prompt SHA-256") + sampling_rows = [_sampling(row) for row in accepted] + sampling = sampling_rows[0] if sampling_rows else None + if sampling is None or not all(value is not None and _same_sampling(value, sampling) for value in sampling_rows): + reasons.append("FreeToken sampling options are missing or mixed") + comparator_reasons, fixture_sha = _comparator_reasons(accepted, label="FreeToken") + reasons.extend(comparator_reasons) + exact = all( + row.get("completion_tokens") == row.get("decode_requested") + and _mapping(row.get("acceptance")).get("accepted") is True + for row in accepted + ) + if not exact: + reasons.append("accepted rows contain exact-completion or acceptance mismatch") + reasons.extend(_execution_reasons(accepted)) + reasons.extend(_lane_reasons(accepted)) + for index, row in enumerate(accepted): + kernel_observed = _mapping(_mapping(row.get("metadata")).get("kernel_observed")) + fallback_markers = kernel_observed.get("fallback_markers") + if isinstance(fallback_markers, int) and fallback_markers > 0: + reasons.append( + f"accepted row {index} carries candidate-eligible fallback evidence " + f"(fallback_markers={fallback_markers})" + ) + if any(_mapping(row.get("metadata")).get("mtp") != "off" for row in accepted): + reasons.append("MTP/speculative mode is not off") + values = [float(row["decode_tok_s"]) for row in accepted if isinstance(row.get("decode_tok_s"), (int, float))] + if len(values) != len(accepted): + reasons.append("accepted rows missing decode throughput") + median = statistics.median(values) if values else None + p02_5 = _bootstrap_p02_5(values) if values else None + minimum = min(values) if values else None + if values and minimum < 70: + reasons.append(f"run below 70 tok/s: {minimum:.3f}") + if values and not p02_5 > 75: + reasons.append(f"bootstrap p02.5={p02_5:.3f} is not >75 tok/s") + + prompt_sha = next(iter(prompt_shas), None) if len(prompt_shas) == 1 else None + probe_reasons = _probe_reasons(probe, model_sha, prompt_sha) + reasons.extend(probe_reasons) + + reference = None + if model_sha and prompt_sha and sampling is not None: + reference, reference_reasons = _reference( + reference_rows, model_sha, prompt_sha, sampling, fixture_sha + ) + reasons.extend(reference_reasons) + else: + reasons.append("reference matching blocked by FreeToken identity/options failure") + if reference is None: + reasons.append("matched reference unavailable; directional Ollama threshold cannot promote parity") + threshold = directional_threshold + threshold_source = "directional-ollama-unproven" + else: + threshold = float(reference["median_tok_s"]) + threshold_source = reference["kind"] + if not values or not median > threshold: + reasons.append( + f"FreeToken median {median if median is not None else 'n/a'} " + f"does not beat reference {threshold:.3f}" + ) + + paired_deltas = _paired_deltas(accepted, reference) + paired_low, paired_high = _bootstrap_interval(paired_deltas, seed=20260831) + if reference is not None: + if len(paired_deltas) < min_runs: + reasons.append( + f"need >= {min_runs} complete FreeToken/reference pairs, got {len(paired_deltas)}" + ) + elif not paired_low > 0: + reasons.append( + f"paired median delta bootstrap interval [{paired_low:.3f}, {paired_high:.3f}] crosses zero" + ) + + # Greedy rows must not silently vary output, when this gate is run on greedy evidence. + output_hashes = {row.get("output_sha1") or row.get("output_sha256") for row in accepted} + if sampling and sampling.get("temperature") == 0 and output_hashes and None not in output_hashes and len(output_hashes) > 1: + reasons.append("greedy/output hashes are not stable") + gate = not reasons + return { + "reference_identity": reference, + "probe": { + "provided": isinstance(probe, dict), + "valid": not probe_reasons, + "reasons": probe_reasons, + }, + "threshold_source": threshold_source, + "threshold_tok_s": threshold, + "runs": { + "accepted": len(accepted), + "required": min_runs, + "median_tok_s": median, + "p02_5_bootstrap_tok_s": p02_5, + "min_tok_s": minimum, + "max_tok_s": max(values) if values else None, + "paired": len(paired_deltas), + "paired_median_delta_tok_s": statistics.median(paired_deltas) if paired_deltas else None, + "paired_delta_p02_5_bootstrap_tok_s": paired_low if paired_deltas else None, + "paired_delta_p97_5_bootstrap_tok_s": paired_high if paired_deltas else None, + }, + "rejected": len(rejected), + "gate": gate, + "reasons": reasons, + } + + +def _replay_rate(row: dict) -> float | None: + """Normalize replay timing to tok/s; missing timing is never zero.""" + for key in ("decode_tok_s", "replay_tok_s", "native_decode_tok_s"): + value = row.get(key) + if isinstance(value, (int, float)) and value > 0: + return float(value) + steps = _mapping(row.get("steps")) + milliseconds = steps.get("ms_per_token_median") + if not isinstance(milliseconds, (int, float)): + milliseconds = row.get("decode_ms_per_token_median") + if isinstance(milliseconds, (int, float)) and milliseconds > 0: + return 1000.0 / float(milliseconds) + return None + + +def _replay_identity_reasons(rows: list[dict], label: str) -> list[str]: + reasons: list[str] = [] + for index, row in enumerate(rows): + if _row_lane(row) != LANE_REPLAY: + reasons.append(f"{label} row {index} is not teacher_forced_replay") + if row.get("forced") is not True: + reasons.append(f"{label} row {index} is not forced replay") + if row.get("ids_match") is not True and row.get("prompt_ids_match") is not True: + reasons.append(f"{label} row {index} input/token IDs are not proven identical") + if row.get("route_hash_status") != "matched": + reasons.append(f"{label} row {index} route hashes are not matched") + for key in ("model_sha256", "fixture_sha256", "tokenizer_sha256", "manifest_ids_sha256"): + if _sha256(row.get(key)) is None: + reasons.append(f"{label} row {index} missing full {key}") + if row.get("mtp") != "off" or row.get("speculative") is not False: + reasons.append(f"{label} row {index} enables MTP/speculative execution") + if row.get("decode_batch_size") != 1: + reasons.append(f"{label} row {index} decode batch size is not 1") + if _replay_rate(row) is None: + reasons.append(f"{label} row {index} missing positive replay timing") + return reasons + + +def _evaluate_replay_gate( + rows: list[dict], + reference_rows: list[dict], + *, + min_runs: int, +) -> dict: + """Gate B: q8/q8 forced replay, independent from sampled absolute speed.""" + accepted = _accepted(rows) + references = _accepted(reference_rows) + reasons = _replay_identity_reasons(accepted, "FreeToken replay") + reasons.extend(_replay_identity_reasons(references, "reference replay")) + if len(accepted) < min_runs: + reasons.append(f"Gate B needs >= {min_runs} accepted replay rows, got {len(accepted)}") + if len(references) < min_runs: + reasons.append(f"Gate B needs >= {min_runs} accepted reference replay rows, got {len(references)}") + for label, candidate in (("FreeToken replay", accepted), ("reference replay", references)): + comparator_reasons, _ = _comparator_reasons(candidate, label=label) + reasons.extend(comparator_reasons) + + by_repeat = {} + for index, row in enumerate(references): + key = row.get("repeat", index) + if key in by_repeat: + reasons.append(f"reference replay has duplicate repeat={key!r}") + by_repeat[key] = row + candidate_values: list[float] = [] + reference_values: list[float] = [] + deltas: list[float] = [] + for index, row in enumerate(accepted): + value = _replay_rate(row) + if value is not None and value < 70: + reasons.append(f"Gate B replay run below 70 tok/s: {value:.3f}") + key = row.get("repeat", index) + ref = by_repeat.get(key) + if ref is None: + reasons.append(f"Gate B missing paired reference for repeat={key!r}") + continue + if not isinstance(row.get("route_digest"), dict) or row.get("route_digest") != ref.get("route_digest"): + reasons.append(f"Gate B route hash mismatch at repeat={key!r}") + for identity in ("model_sha256", "fixture_sha256", "tokenizer_sha256", "manifest_ids_sha256"): + if row.get(identity) != ref.get(identity): + reasons.append(f"Gate B {identity} mismatch at repeat={key!r}") + reference_value = _replay_rate(ref) + if value is not None and reference_value is not None: + candidate_values.append(value) + reference_values.append(reference_value) + deltas.append(value - reference_value) + low, high = _bootstrap_interval(deltas, seed=20260831) + if len(deltas) < min_runs: + reasons.append(f"Gate B needs >= {min_runs} complete replay pairs, got {len(deltas)}") + elif not low > 0: + reasons.append(f"Gate B paired median delta bootstrap interval [{low:.3f}, {high:.3f}] does not exceed zero") + return { + "gate": not reasons, + "lane": LANE_REPLAY, + "runs": { + "accepted": len(accepted), + "reference_accepted": len(references), + "paired": len(deltas), + "required": min_runs, + "median_tok_s": statistics.median(candidate_values) if candidate_values else None, + "reference_median_tok_s": statistics.median(reference_values) if reference_values else None, + "paired_median_delta_tok_s": statistics.median(deltas) if deltas else None, + "paired_delta_p02_5_bootstrap_tok_s": low if deltas else None, + "paired_delta_p97_5_bootstrap_tok_s": high if deltas else None, + }, + "reasons": reasons, + } + + +def evaluate_gate( + freetoken_rows: list[dict], + reference_rows: list[dict] | None = None, + probe: dict | None = None, + *, + min_runs: int = 10, + directional_threshold: float = 80.225, +) -> dict: + """Evaluate Gate A sampled and Gate B replay lanes independently.""" + reference_rows = reference_rows or [] + sampled_rows = [row for row in freetoken_rows if _row_lane(row) == LANE_SAMPLED] + replay_rows = [row for row in freetoken_rows if _row_lane(row) == LANE_REPLAY] + gate_a = _evaluate_sampled_gate( + sampled_rows, reference_rows, probe, + min_runs=min_runs, directional_threshold=directional_threshold, + ) + gate_b = None + if replay_rows: + gate_b = _evaluate_replay_gate( + replay_rows, + [row for row in reference_rows if _row_lane(row) == LANE_REPLAY], + min_runs=min_runs, + ) + result = dict(gate_a) + result["gate_a"] = { + "gate": gate_a["gate"], + "runs": gate_a["runs"], + "reasons": gate_a["reasons"], + } + result["gate_b"] = gate_b + result["gates"] = { + "gate_a": gate_a["gate"], + "gate_b": None if gate_b is None else gate_b["gate"], + } + result["lane_counts"] = { + LANE_SAMPLED: len(sampled_rows), + LANE_GREEDY: sum(_row_lane(row) == LANE_GREEDY for row in freetoken_rows), + LANE_REPLAY: len(replay_rows), + } + if gate_b is not None: + result["gate"] = gate_a["gate"] and gate_b["gate"] + result["reasons"] = gate_a["reasons"] + [ + f"Gate B: {reason}" for reason in gate_b["reasons"] + ] + return result + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--freetoken", required=True) + parser.add_argument("--reference", action="append", default=[]) + parser.add_argument("--probe", help="finite-logit/eager-graph parity probe JSON") + parser.add_argument("--min-runs", type=int, default=10) + parser.add_argument("--json", dest="json_out") + args = parser.parse_args(argv) + try: + reference_rows = [row for path in args.reference for row in load_jsonl(path)] + probe = json.loads(Path(args.probe).read_text()) if args.probe else None + freetoken_rows = load_jsonl(args.freetoken) + result = evaluate_gate(freetoken_rows, reference_rows, probe, min_runs=args.min_runs) + except (OSError, TypeError, ValueError, json.JSONDecodeError) as exc: + result = { + "reference_identity": None, + "probe": {"provided": bool(args.probe), "valid": False, "reasons": []}, + "threshold_source": "input-error", + "threshold_tok_s": None, + "runs": { + "accepted": 0, + "required": args.min_runs, + "median_tok_s": None, + "p02_5_bootstrap_tok_s": None, + "min_tok_s": None, + "max_tok_s": None, + }, + "rejected": 0, + "gate": False, + "gate_a": {"gate": False, "runs": {}, "reasons": []}, + "gate_b": None, + "gates": {"gate_a": False, "gate_b": None}, + "lane_counts": {LANE_SAMPLED: 0, LANE_GREEDY: 0, LANE_REPLAY: 0}, + "reasons": [f"invalid gate input: {type(exc).__name__}: {exc}"], + } + encoded = json.dumps(result, indent=2, sort_keys=True, allow_nan=False) + print(encoded) + if args.json_out: + Path(args.json_out).write_text(encoded + "\n") + return 0 if result["gate"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/probe_qwen_moe_base.py b/benchmarks/probe_qwen_moe_base.py new file mode 100644 index 000000000..304296100 --- /dev/null +++ b/benchmarks/probe_qwen_moe_base.py @@ -0,0 +1,263 @@ +"""Finite-logit and eager/graph parity probe for real Qwen MoE checkpoints. + +This is intentionally separate from the throughput benchmark. It drives FreeToken's +offline scheduler, captures logits immediately before sampling, and then checks the +greedy token IDs. HTTP readiness or generated text alone cannot detect NaN logits or +captured-graph corruption. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from pathlib import Path + +import torch + +from bench_decode_moe import load_problem_details, model_fingerprint, runtime_metadata + + +@dataclass +class ProbeResult: + requested: str + actual: str + token_ids: list[int] + logits: list[torch.Tensor] + + def artifact(self) -> dict: + rows = torch.cat(self.logits, dim=0) if self.logits else torch.empty(0) + digest = hashlib.sha256(rows.numpy().tobytes()).hexdigest()[:16] + finite = bool(torch.isfinite(rows).all()) + return { + "requested": self.requested, + "actual": self.actual, + "decode_rows": len(self.logits), + "logit_shape": list(rows.shape), + "finite_logits": finite, + "logit_sha256": digest, + "token_ids": self.token_ids, + } + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--model", required=True) + p.add_argument("--aime", default=os.environ.get("FREETOKEN_AIME25_JSONL")) + p.add_argument("--aime-revision", default=os.environ.get("FREETOKEN_AIME25_REVISION")) + p.add_argument("--aime-sha256", default=os.environ.get("FREETOKEN_AIME25_SHA256")) + p.add_argument("--problem", type=int, default=0) + p.add_argument("--decode", type=int, default=8) + p.add_argument("--memory-ratio", type=float, default=0.9) + p.add_argument("--cache", type=int, default=0) + p.add_argument("--json", dest="json_out", default=None) + p.add_argument("--_worker", action="store_true", help=argparse.SUPPRESS) + p.add_argument("--_mode", choices=("eager", "graph"), help=argparse.SUPPRESS) + p.add_argument("--_graph-gate", default="unknown", help=argparse.SUPPRESS) + p.add_argument("--_prompt-file", help=argparse.SUPPRESS) + p.add_argument("--_artifact-file", help=argparse.SUPPRESS) + p.add_argument("--_logits-file", help=argparse.SUPPRESS) + return p.parse_args(argv) + + +def run_probe( + args: argparse.Namespace, + prompt: str, + requested: str, + graph_gate: str, +) -> ProbeResult: + from freetoken.core import SamplingParams + from freetoken.llm import LLM + + graph_requested = requested == "graph" + kwargs = { + "attention_backend": "triton", + "max_running_req": 1, + "max_extend_tokens": max(8192, args.decode + 128), + "max_seq_len_override": 8192 + args.decode, + "moe_backend": "offload", + "memory_ratio": args.memory_ratio, + "moe_cache_auto": args.cache == 0, + "cuda_graph_max_bs": 1 if graph_requested else 0, + } + if args.cache > 0: + kwargs["moe_cache_size"] = args.cache + llm = LLM(args.model, dtype=torch.bfloat16, **kwargs) + captured: list[torch.Tensor] = [] + original_sample = llm.engine.sampler.sample + + def capture_sample(logits, sample_args, batch): + if batch.is_decode: + captured.append(logits.detach().float().cpu()) + return original_sample(logits, sample_args, batch) + + llm.engine.sampler.sample = capture_sample + try: + actual = "replay" if llm.engine.graph_runner.graph_map else "eager" + if graph_requested and graph_gate == "pass" and actual != "replay": + raise RuntimeError("graph gate passed but no graph was captured") + output = llm.generate( + [prompt], + SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + # Match the serving path: current full-prefix scheduling consumes one + # engine output budget slot before the first returned completion token. + max_tokens=args.decode + 1, + ignore_eos=True, + ), + )[0] + token_ids = list(output["token_ids"]) + finally: + llm.shutdown() + if len(token_ids) != args.decode: + raise RuntimeError( + f"{requested} probe generated {len(token_ids)} tokens, expected {args.decode}" + ) + result = ProbeResult(requested, actual, token_ids, captured) + artifact = result.artifact() + if not artifact["finite_logits"]: + raise RuntimeError(f"{requested} probe found non-finite logits: {artifact}") + if artifact["decode_rows"] != args.decode: + raise RuntimeError( + f"{requested} probe captured {artifact['decode_rows']} decode logit rows, " + f"expected {args.decode}" + ) + return result + + +def run_probe_worker(args: argparse.Namespace) -> int: + """Run one LLM lifetime in a pristine child process. + + Engine construction intentionally rejects an already initialized CUDA runtime, so eager + and graph probes cannot share this interpreter. Keep this boundary explicit rather than + weakening that runtime invariant. + """ + prompt = Path(args._prompt_file).read_text() + result = run_probe(args, prompt, args._mode, args._graph_gate) + artifact = result.artifact() + torch.save(torch.cat(result.logits, dim=0), args._logits_file) + Path(args._artifact_file).write_text(json.dumps(artifact, sort_keys=True)) + return 0 + + +def run_probe_child( + args: argparse.Namespace, + prompt: str, + requested: str, + graph_gate: str, + workdir: str, +) -> ProbeResult: + prompt_file = Path(workdir) / f"{requested}-prompt.txt" + artifact_file = Path(workdir) / f"{requested}-artifact.json" + logits_file = Path(workdir) / f"{requested}-logits.pt" + prompt_file.write_text(prompt) + cmd = [ + sys.executable, + __file__, + "--_worker", + "--model", args.model, + "--problem", str(args.problem), + "--decode", str(args.decode), + "--memory-ratio", str(args.memory_ratio), + "--cache", str(args.cache), + "--_mode", requested, + "--_graph-gate", graph_gate, + "--_prompt-file", str(prompt_file), + "--_artifact-file", str(artifact_file), + "--_logits-file", str(logits_file), + ] + if args.aime: + cmd += ["--aime", args.aime] + if args.aime_revision: + cmd += ["--aime-revision", args.aime_revision] + if args.aime_sha256: + cmd += ["--aime-sha256", args.aime_sha256] + child = subprocess.run(cmd, text=True, capture_output=True, check=False) + if child.returncode != 0: + tail = "\n".join((child.stdout + child.stderr).splitlines()[-30:]) + raise RuntimeError(f"{requested} probe worker failed (rc={child.returncode}):\n{tail}") + if not artifact_file.is_file() or not logits_file.is_file(): + raise RuntimeError(f"{requested} probe worker produced no artifacts") + artifact = json.loads(artifact_file.read_text()) + rows = torch.load(logits_file, map_location="cpu", weights_only=True) + if rows.ndim != 2 or rows.shape[0] != args.decode: + raise RuntimeError(f"{requested} probe worker returned bad logits shape {tuple(rows.shape)}") + return ProbeResult( + requested=artifact["requested"], + actual=artifact["actual"], + token_ids=artifact["token_ids"], + logits=list(rows.split(1, dim=0)), + ) + + +def compare(eager: ProbeResult, graph: ProbeResult) -> dict: + if eager.token_ids != graph.token_ids: + raise RuntimeError("eager/graph greedy token IDs differ") + if len(eager.logits) != len(graph.logits): + raise RuntimeError("eager/graph decode logit row counts differ") + max_abs = 0.0 + max_rel = 0.0 + for left, right in zip(eager.logits, graph.logits): + if left.shape != right.shape: + raise RuntimeError(f"eager/graph logit shapes differ: {left.shape} vs {right.shape}") + delta = (left - right).abs() + max_abs = max(max_abs, float(delta.max())) + max_rel = max(max_rel, float((delta / right.abs().clamp_min(1e-6)).max())) + torch.testing.assert_close(left, right, rtol=2e-2, atol=2e-2) + return {"token_ids_equal": True, "max_abs": max_abs, "max_rel": max_rel} + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + if args._worker: + if not all((args._mode, args._prompt_file, args._artifact_file, args._logits_file)): + raise SystemExit("probe worker arguments incomplete") + return run_probe_worker(args) + if args.decode < 1: + raise SystemExit("--decode must be >= 1") + prompt, answer, dataset = load_problem_details( + args.aime, args.problem, args.aime_revision, args.aime_sha256 + ) + from freetoken.utils.graph_gate import graph_capture_status + + gate = graph_capture_status() + with tempfile.TemporaryDirectory(prefix="freetoken-qwen-probe-") as workdir: + eager = run_probe_child(args, prompt, "eager", gate, workdir) + graph = None + comparison = {"status": "skipped", "reason": f"graph gate={gate}"} + if gate == "pass": + graph = run_probe_child(args, prompt, "graph", gate, workdir) + comparison = compare(eager, graph) + artifact = { + "schema": "qwen-moe-base-probe-v2", + "model": model_fingerprint(args.model), + "dataset": dataset, + "problem": args.problem, + "prompt_sha256": hashlib.sha256(prompt.encode("utf-8")).hexdigest(), + "answer": answer, + "decode": args.decode, + "mtp": "off", + "speculative": False, + "graph_gate": gate, + "eager": eager.artifact(), + "graph": graph.artifact() if graph else None, + "comparison": comparison, + "runtime": runtime_metadata(), + } + print(json.dumps(artifact, indent=2, sort_keys=True)) + if args.json_out: + with open(args.json_out, "w") as f: + json.dump(artifact, f, indent=2, sort_keys=True) + f.write("\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/profile_decode_rocm.py b/benchmarks/profile_decode_rocm.py new file mode 100644 index 000000000..dbdf9a7dd --- /dev/null +++ b/benchmarks/profile_decode_rocm.py @@ -0,0 +1,98 @@ +"""Launch-mode ROCm trace wrapper and ledger converter. + +No command is executed unless the caller supplies ``--`` followed by an explicit +runtime command. The implementation is safe to import and unit-test on CPU-only hosts. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +from pathlib import Path + +try: + from benchmarks.lib.rocm_trace import ( + calibrate_clocks, correlate_events, token_ledger, warm_offload_summary, + ) +except ModuleNotFoundError: # direct ``python benchmarks/profile_decode_rocm.py`` + from lib.rocm_trace import calibrate_clocks, correlate_events, token_ledger, warm_offload_summary + + +def parse_args(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--trace", type=Path, help="existing rocprof JSON artifact") + parser.add_argument("--out", type=Path, required=True, help="ledger JSON output") + parser.add_argument("--rocprof", default="rocprofv3") + parser.add_argument("--profile-output", type=Path) + parser.add_argument("command", nargs=argparse.REMAINDER) + return parser.parse_args(argv) + + +def load_trace(path: Path) -> dict: + payload = json.loads(path.read_text()) + if isinstance(payload, list): + return {"events": payload} + if not isinstance(payload, dict): + raise ValueError("rocprof artifact must be an object or event list") + return payload + + +def build_report(payload: dict) -> dict: + records = payload.get("clock_correlations") + if not records: + raise ValueError("clock calibration records unavailable") + calibration = calibrate_clocks(records) + if not calibration.accepted: + raise ValueError(f"clock calibration residual exceeds 10us: {calibration.max_residual_ns:.0f}ns") + events = payload.get("events", []) + hip = payload.get("hip_api", [event for event in events if event.get("kind") == "hip_api"]) + kernels = payload.get("kernels", [event for event in events if event.get("kind") == "kernel"]) + copies = payload.get("copies", [event for event in events if event.get("kind") == "copy"]) + correlated = correlate_events(hip, kernels, copies) + return { + "schema": "freetoken-rocm-ledger-v1", + "clock": { + "scale": calibration.scale, + "offset_ns": calibration.offset_ns, + "max_residual_ns": calibration.max_residual_ns, + "accepted": calibration.accepted, + }, + "ledgers": token_ledger(payload.get("tokens", payload.get("token_ranges", [])), correlated, payload.get("host_ranges", [])), + "events": correlated, + "warm_offload": warm_offload_summary(correlated), + "missing": payload.get("missing", {}), + } + + +def main(argv=None) -> int: + args = parse_args(argv) + trace = args.trace + command = [item for item in args.command if item != "--"] + if trace is None: + if not command: + raise SystemExit("provide --trace or an explicit command after --") + profile_output = args.profile_output or args.out.with_suffix(".rocprof.json") + run = [ + args.rocprof, + "--hip-trace", + "--marker-trace", + "--kernel-trace", + "--memory-copy-trace", + "--output-file", + str(profile_output), + "--output-format", + "json", + "--", + *command, + ] + subprocess.run(run, check=True) + trace = profile_output + report = build_report(load_trace(trace)) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/install-amd.md b/docs/install-amd.md index 767632795..423826dd4 100644 --- a/docs/install-amd.md +++ b/docs/install-amd.md @@ -41,8 +41,8 @@ and their backends are rejected with a clean error if requested. | Feature | On AMD | Notes | | --- | --- | --- | | Attention | `--attention-backend triton` | flashinfer/fa/trtllm are NVIDIA-only and rejected | -| MoE | `--moe-backend offload / cpu / hybrid` | offload needs pinned host memory | -| Quant | BF16, MXFP4, GGUF (Q4_K/Q8_0), Triton inline-dequant NVFP4 | Marlin INT4 / native NVFP4 SASS unavailable | +| MoE | `--moe-backend fused / offload / cpu / hybrid` | GGUF `fused` is native resident only when preflight fits; offload needs pinned host memory | +| Quant | BF16, MXFP4, GGUF (Q4_K/Q5_K/Q6_K/Q8_0), Triton inline-dequant NVFP4 | Marlin INT4 / native NVFP4 SASS unavailable | | NVFP4 checkpoints with no MXFP4 variant | converted to MXFP4 on load (auto) | `--nvfp4-backend auto` → triton/MXFP4 | | CUDA graphs (decode) | HIP graph capture **if** the capture probe passes | otherwise kernel-launch decode | | Multi-GPU (RCCL) | out of scope (single-GPU milestone) | | @@ -51,7 +51,9 @@ and their backends are rejected with a clean error if requested. * `--nvfp4-backend marlin` / `flashinfer` → error (NVIDIA-only). Use `triton` / `auto`. * `--attention-backend fi` / `fa` / `trtllm` → error (NVIDIA-only). Use `triton` / `auto`. -* `--moe-backend fused` → warning (fused MoE is CUDA-only; falls back to offload/cpu). +* `--moe-backend fused` → native GGUF resident path when allocation-free fit check passes; + explicit fit failure is an error. `--moe-backend auto` selects native residency only when it + fits, otherwise existing offload. * `--nvfp4-backend auto` → resolves to the portable Triton inline-dequant path (or MXFP4 for a converted checkpoint). @@ -77,16 +79,18 @@ The runtime refuses to pair a `+rocm` cache with a `+cu130` runtime (and vice ve ## Notes / limitations -* `nvtx_annotate` is a no-op on ROCm; roctx profiling is future work. +* `nvtx_annotate` is a no-op on ROCm; `FREETOKEN_ROCTX_MARKERS=1` enables low-overhead + ROCTX ranges, and `scripts/profile-rocm-decode.sh PROFILE_MODE=rocprofv3` captures + runtime, marker, and kernel traces when `rocprofv3` is installed. * FP8 / NVFP4-class formats: BF16 / MXFP4 / GGUF are the supported AMD matrix; performance parity vs CUDA is not guaranteed for NVFP4-class formats. * Windows AMD is not yet supported (WDDM zero-copy semantics differ). -## Performance (gfx1100 / RX 7900 XTX, as of rocm-perf-parity) +## Performance (gfx1100 / RX 7900 XTX, as of qwen-moe-speed) Measured with `benchmarks/bench_decode_moe.py` on Qwen3.6-35B-A3B GGUF (Q4_K_M), 3-run medians, same prompt/sampling protocol throughout (see -`.plans/rocm-perf-parity/` for artifacts and the stage-time profiling notes): +`.plans/rocm-perf-parity/` for historical artifacts and the stage-time profiling notes): | state | decode tok/s (median) | | --- | --- | @@ -94,6 +98,84 @@ Measured with `benchmarks/bench_decode_moe.py` on Qwen3.6-35B-A3B GGUF (Q4_K_M), | + in-repo fused triton router | 37.67 | | + CUDA-graph decode capture (`prewarm` variant) | **45.09** | +Current Qwen3.6 base-speed gate on current HEAD, using the exact FreeToken +`Qwen3.6-35B-A3B-UD-Q4_K_M.gguf` file, one request, 512 generated tokens, +offload MoE, Triton attention, and graph replay: + +| lane | measured decode tok/s (median) | +| --- | --- | +| FreeToken sampled, 10 runs | 62.295 | +| FreeToken greedy, 3 runs | 59.650 | +| FreeToken graph-off control, Inc2 baseline | 38.611 | +| Ollama sampled, 3 runs, client arrival | 80.225 | + +Inc 4 native GPU-offload diagnostic (one exact 128-token graph run, not a promotion +median): **58.98 tok/s** API arrival, **61.19 tok/s** scheduler timing. Startup fell +from roughly 3.5 minutes for legacy Q8 conversion to roughly 19 seconds; decode did +not materially improve. This rejects loader conversion as sole cause and leaves +attention/dense/MoE dispatch and backend differences as the remaining speed gap. + +This gate did not clear the 75 tok/s external floor or the 80 tok/s engineering +target. Ollama was configured without MTP, but its Q4_K_M model blob differs +from FreeToken's GGUF file; treat 80.225 tok/s as directional, not an exact +same-file claim. Full provenance and JSONL artifacts are in +`.plans/qwen-moe-speed/notes-baseline.md` and `notes-profile.md`. + +### ROCm execution policy + +`FREETOKEN_ROCM_BLAS=auto|hipblas|hipblaslt|rocblas` selects the ROCm BLAS +preference before workers start. `rocblas` is an alias for hipBLAS; explicit +requests fail when the installed PyTorch API cannot honor them. Startup logs and +benchmark JSON report requested and effective policy. Changing BLAS variables +after worker startup has no effect. + +KV storage is explicit: `--kv-type bf16` (default), `fp16`, or opt-in `q8_0`. +Q8 uses the pinned 32-value row contract with FP16 scales and currently accepts +plain full-attention MHA groups only. Allocation, store, decode, prefill, and +pointer-generation metadata must all report q8 before a q8 benchmark row is +eligible. Roll back with `--kv-type bf16`. + +ROCm GGUF JIT stages its `.cu` entrypoint into the torch extension cache before +PyTorch HIPify runs. This keeps generated HIP intermediates and tracked checkout +headers separate; delete only the relevant torch-extension cache after toolchain +upgrades if a rebuild is needed. + +GGUF Q4_K/Q5_K/Q6_K/Q8_0 MoE uses MMVQ for decode. Native Q5_K/Q6_K rows retain +their source type; Q5_K cache rows carry zeroed Q6_K-sized tails and kernels receive +explicit expert and row strides. A grouped-MMQ ABI probe exists, but real-model route +validation exposed an illegal HIP access, so it is opt-in only via +`FREETOKEN_GGUF_GROUPED_PREFILL=1`; default prefill stays on proven vector dispatch. +The `gfx1100` MoE kernel is available only as an explicit, forced candidate; `legacy` +remains default because measured candidate throughput was slower. + +Greedy/no-penalty decode may opt into fixed-address sampler graph capture: + +```bash +FREETOKEN_GRAPH_SAMPLER=1 ft serve --model /path/to/model.gguf \ + --moe-backend offload --attention-backend triton +``` + +Dynamic temperature/top-k/top-p, penalties, and unsupported sampling modes use +the regular sampler path. Unset `FREETOKEN_GRAPH_SAMPLER` to retain default +behavior. Use `FREETOKEN_GGUF_MOE_IMPL=legacy` for explicit rollback: + +```bash +export FREETOKEN_GGUF_MOE_IMPL=legacy +unset FREETOKEN_GRAPH_SAMPLER +ft serve --model /path/to/model.gguf --moe-backend offload --attention-backend triton +``` + +Final promotion requires ten accepted runs, exact 512-token completion, stable +sampling/output checks, a matching finite-logit/eager-graph probe, p02.5 bootstrap +lower bound above 75 tok/s, no run below 70 tok/s, and a matched full-file-SHA-256 +llama.cpp/Ollama reference. Run `benchmarks/probe_qwen_moe_base.py` and pass its +JSON to `benchmarks/check_decode_gate.py --probe`; missing probe evidence fails +closed. The current 80.225 tok/s Ollama result lacks byte identity and is +directional only. Gate A covers sampled absolute throughput. Gate B is separate: +q8/q8 teacher-forced replay, matched IDs and route hashes, paired bootstrap delta. +Local CUDA compilation was unavailable (`nvcc` absent); NVIDIA compilation remains +a CI gate. + What is enabled on AMD now: - **Fused triton router**: `fused_topk` routes ROCm to the in-repo diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py index f397f8491..3e4b82490 100644 --- a/python/freetoken/attention/triton.py +++ b/python/freetoken/attention/triton.py @@ -149,6 +149,26 @@ def forward( assert isinstance(metadata, TritonMetadata) self.kvcache.store_kv(k, v, batch.out_loc, layer_id) + if getattr(self.kvcache, "is_quantized", False): + from freetoken.kernel.triton.attention import q8_paged_attention + + k_view = self.kvcache.k_cache_view(layer_id) + v_view = self.kvcache.v_cache_view(layer_id) + scale = spec.sm_scale if (spec := (attn_spec or AttentionSpec())).sm_scale is not None else q.shape[-1] ** -0.5 + return q8_paged_attention( + q=q, + k_payload=k_view.payload, + v_payload=v_view.payload, + k_scales=k_view.scales, + v_scales=v_view.scales, + indptr=metadata.indptr, + indices=metadata.indices, + q_to_req=metadata.q_to_req, + q_positions=metadata.q_positions, + sm_scale=scale, + sliding_window=spec.sliding_window, + ) + k_raw = self.kvcache.k_cache(layer_id) v_raw = self.kvcache.v_cache(layer_id) kv_heads, head_dim = k_raw.shape[-2], k_raw.shape[-1] @@ -242,6 +262,14 @@ def prepare_metadata(self, batch: Batch) -> None: [0] + seqlens_q, dtype=torch.int32, device=device ).cumsum_(0) indices = torch.cat([page_table[req.table_idx, : req.device_len] for req in reqs]) + if getattr(self.kvcache, "is_quantized", False): + # Duplicate physical destinations would make the captured Q8 store + # last-writer-wins and pair payload/scales nondeterministically. + from freetoken.kernel.triton.q8_kv import validate_unique_destinations + + destinations = getattr(batch, "out_loc", None) + if destinations is not None: + validate_unique_destinations(destinations) swa_indices = None if getattr(self.kvcache, "swa_paged", False): # Global-paged SWA (naive + radix): the swa-layer gather reads swa-pool slots = full->swa diff --git a/python/freetoken/engine/__init__.py b/python/freetoken/engine/__init__.py index cf74975da..332fe80e3 100644 --- a/python/freetoken/engine/__init__.py +++ b/python/freetoken/engine/__init__.py @@ -1,5 +1,5 @@ -from .config import EngineConfig -from .engine import Engine, ForwardOutput +from .config import EngineConfig, KVStorageType +from .engine import DeviceTokenChain, Engine, ForwardOutput from .sample import BatchSamplingArgs -__all__ = ["Engine", "EngineConfig", "ForwardOutput", "BatchSamplingArgs"] +__all__ = ["DeviceTokenChain", "Engine", "EngineConfig", "KVStorageType", "ForwardOutput", "BatchSamplingArgs"] diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 543012f39..07fd0a805 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass +from enum import Enum from functools import cached_property from typing import TYPE_CHECKING, List @@ -9,6 +10,26 @@ from freetoken.models.register import _load_attr, get_model_spec from freetoken.utils import cached_load_hf_config + +class KVStorageType(str, Enum): + """Physical KV-cache format, independent from model compute dtype.""" + + BF16 = "bf16" + FP16 = "fp16" + Q8_0 = "q8_0" + + @classmethod + def parse(cls, value: "KVStorageType | str") -> "KVStorageType": + if isinstance(value, cls): + return value + try: + return cls(str(value).lower()) + except ValueError as exc: + raise ValueError( + f"unsupported KV storage {value!r}; expected one of " + f"{', '.join(item.value for item in cls)}" + ) from exc + if TYPE_CHECKING: from freetoken.models import ModelConfig @@ -20,6 +41,7 @@ class EngineConfig: dtype: torch.dtype max_running_req: int = 4 attention_backend: str = "auto" + kv_storage_type: KVStorageType | None = None moe_backend: str = "auto" # NVFP4 routed-expert GEMM backend (--nvfp4-backend): auto|marlin|flashinfer|triton. nvfp4_backend: str = "triton" @@ -81,6 +103,10 @@ class EngineConfig: # is final. Mutually exclusive with num_page_override. num_token_override: int | None = None + def __post_init__(self) -> None: + if self.kv_storage_type is not None: + object.__setattr__(self, "kv_storage_type", KVStorageType.parse(self.kv_storage_type)) + @cached_property def hf_config(self): return cached_load_hf_config(self.model_path) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index b4e20ebdb..87409f11e 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -17,12 +17,13 @@ from freetoken.moe.expert_banks import load_expert_banks from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache from freetoken.utils import align_ceil, device_kind, init_logger, is_rocm, is_sm90_family, is_sm100_family, mem_GB, torch_dtype +from freetoken.utils.step_profiler import profiler_phase from .config import EngineConfig from .graph import GraphRunner, get_free_memory from .sample import BatchSamplingArgs, Sampler from freetoken.kvcache import create_kv_pool, resolve_pool_class -from freetoken.kvcache.base import CacheRebuildRejected +from freetoken.kvcache.base import CacheRebuildRejected, validate_kv_storage_config from freetoken.kvcache.cache_status import _supports_swa_ratio from freetoken.kvcache.linear_state_pool import ( _linear_pool_min_slots, _linear_pool_num_slots, state_pool_bytes, @@ -44,6 +45,29 @@ def _require_offload_cache_size(cache_size: int, num_experts: int) -> None: ) +def _resolve_gguf_resident_backend(config: EngineConfig, free_bytes: int) -> None: + """Resolve GGUF native residency before model tensors or caches are allocated.""" + model_config = config.model_config + if getattr(model_config, "moe_weight_format", None) != "gguf": + return + if not getattr(model_config, "is_moe", False) or config.moe_backend not in ("auto", "fused"): + return + from freetoken.engine.resident_budget import estimate_gguf_resident_budget, resolve_gguf_moe_backend + + budget = estimate_gguf_resident_budget(config.model_path, config, free_bytes) + selected = resolve_gguf_moe_backend(config, free_bytes) + if config.moe_backend == "fused" and selected != "fused": + raise RuntimeError( + "--moe-backend fused cannot fit native GGUF residency before allocation: " + f"required={budget.required_bytes} free={budget.free_bytes} breakdown={budget.as_dict()}" + ) + if config.moe_backend == "auto": + object.__setattr__(config, "moe_backend", selected) + logger.info_rank0( + f"GGUF resident preflight selected moe_backend={selected!r}: {budget.as_dict()}" + ) + + def _flashinfer_available() -> bool: from freetoken.kernel.backend import is_flashinfer_installed @@ -297,15 +321,93 @@ def _materialize_loaded_weight_state_dict( return state_dict -from freetoken.utils.step_profiler import step_profiler - - class ForwardOutput(NamedTuple): next_tokens_gpu: torch.Tensor next_tokens_cpu: torch.Tensor copy_done_event: torch.cuda.Event +class TokenStaging: + """Double-buffer sampled tokens across engine and scheduler streams. + + Overlap scheduling drains previous CPU tokens after launching current forward, so one + pinned-CPU staging slot remains in use while the other is written. Reusing either slot is + safe only after its prior copy event was drained; the scheduler's one-step overlap provides + that invariant. GPU sampled output remains caller-owned, avoiding an extra device copy. + """ + + def __init__(self, device: torch.device, capacity: int) -> None: + if device.type != "cuda": + raise ValueError("TokenStaging requires CUDA device") + self.capacity = max(1, capacity) + self._next_slot = 0 + self._cpu = [ + torch.empty(self.capacity, dtype=torch.int32, pin_memory=True) for _ in range(2) + ] + self._events = [torch.cuda.Event(), torch.cuda.Event()] + # Keep D2H launch off engine stream. Next decode can consume device IDs while CPU + # response observation waits on copy completion independently. + self.copy_stream = torch.cuda.Stream(device=device) + + def stage( + self, tokens: torch.Tensor, stream: torch.cuda.Stream + ) -> tuple[torch.Tensor, torch.Tensor, torch.cuda.Event]: + count = tokens.numel() + if count > self.capacity: + raise ValueError(f"token batch {count} exceeds staging capacity {self.capacity}") + slot = self._next_slot + cpu = self._cpu[slot][:count] + gpu = tokens.reshape(-1) + if gpu.dtype != torch.int32: + gpu = gpu.to(torch.int32) + event = self._events[slot] + with torch.cuda.stream(self.copy_stream): + self.copy_stream.wait_stream(stream) + cpu.copy_(gpu, non_blocking=True) + gpu.record_stream(self.copy_stream) + event.record(self.copy_stream) + self._next_slot = (slot + 1) % 2 + return gpu, cpu, event + + +class DeviceTokenChain: + """Stable graph IDs plus two-slot device handoff and CPU observation rings.""" + + def __init__(self, device: torch.device, capacity: int) -> None: + if device.type != "cuda": + raise ValueError("DeviceTokenChain requires CUDA device") + self.capacity = max(1, capacity) + # Graph captures need one fixed output address. Runtime handoff needs two addresses so + # copy-stream D2H can overlap next decode without reading a buffer being rewritten. + self.device_tokens = torch.empty(self.capacity, dtype=torch.int32, device=device) + self.sampled_indices = torch.empty(self.capacity, dtype=torch.int64, device=device) + self._handoff_tokens = torch.empty( + (2, self.capacity), dtype=torch.int32, device=device + ) + self._next_handoff = 0 + self.staging = TokenStaging(device, self.capacity) + + def next_device_tokens(self) -> torch.Tensor: + slot = self._next_handoff + self._next_handoff = (slot + 1) % 2 + return self._handoff_tokens[slot] + + def publish(self, tokens: torch.Tensor, stream: torch.cuda.Stream) -> torch.Tensor: + """Snapshot graph output into handoff slot before asynchronous observation.""" + del stream # caller's engine stream orders this copy before ``stage`` wait + target = self.next_device_tokens() + flat = tokens.reshape(-1) + if flat.dtype != torch.int32: + flat = flat.to(torch.int32) + target[: flat.numel()].copy_(flat) + return target[: flat.numel()] + + def stage( + self, tokens: torch.Tensor, stream: torch.cuda.Stream + ) -> tuple[torch.Tensor, torch.Tensor, torch.cuda.Event]: + return self.staging.stage(tokens, stream) + + class Engine: def __init__(self, config: EngineConfig): assert not torch.cuda.is_initialized() @@ -315,13 +417,24 @@ def __init__(self, config: EngineConfig): from freetoken.gpu_select import bind_assigned_gpu self.device = bind_assigned_gpu(config.tp_info.rank) + _resolve_gguf_resident_backend(config, get_free_memory(self.device)) _adjust_config(config) logger.info_rank0(f"device_kind={device_kind()} backend={self.device}") + from freetoken.utils.graph_gate import rocm_blas_report + + self.blas_policy = rocm_blas_report() + if self.blas_policy["verification"] == "mismatch": + raise RuntimeError( + "requested ROCm BLAS policy was not effective: " + f"{self.blas_policy}" + ) + logger.info_rank0(f"BLAS policy effective={self.blas_policy}") torch.manual_seed(42) self.stream = torch.cuda.Stream() torch.cuda.set_stream(self.stream) self.dtype = config.dtype self.config = config # retained for runtime cache rebuild (rebuild_runtime_cache) + self.memory_phases = [] # KV pool family fixed at construction from the model config: its classmethods own the # page-token geometry and cost arithmetic the engine needs BEFORE the pool exists # (num_pages sizing, --moe-cache-auto); the instance owns rebuild/validation after. @@ -339,6 +452,7 @@ def __init__(self, config: EngineConfig): set_rope_device(self.device) with torch.device("meta"), torch_dtype(config.dtype): self.model = create_model(config.model_config) + self._memory_phase_start("load") self.model.load_state_dict(self._load_weight_state_dict(config)) post_weights_free = self._sync_get_memory()[0] self._weights_bytes = self._baseline_free - post_weights_free @@ -360,6 +474,7 @@ def __init__(self, config: EngineConfig): self._init_offload_moe_cache(config) if hasattr(self.model, "prepare_for_runtime"): self.model.prepare_for_runtime() + self._memory_phase_end("load") # ======================= KV cache initialization ======================== new_free = self._sync_get_memory()[1] @@ -372,6 +487,13 @@ def __init__(self, config: EngineConfig): self.ctx.kv_cache = self.kv_cache = create_kv_pool( config, self.num_pages, device=self.device, dtype=self.dtype ) + descriptor = getattr(self.kv_cache, "storage_descriptor", None) + self.kv_storage_metadata = { + "storage_type": getattr(getattr(descriptor, "storage_type", None), "value", None), + "contract_id": getattr(descriptor, "contract_id", None), + "pointer_generation": getattr(self.kv_cache, "pointer_generation", None), + "unit_bytes": list(getattr(self.kv_cache, "unit_bytes", lambda: (None, None))()), + } # ======================= Linear (GatedDeltaNet) state initialization ======================== linear_group = config.model_config.linear_attention_group() @@ -403,6 +525,10 @@ def __init__(self, config: EngineConfig): # re-point here (and again on any table realloc). The graph-input snapshot that reads # through them belongs to the attention backend, built later in init_capture_graph. self.kv_cache.attach_page_table(self.page_table) + if hasattr(self, "kv_storage_metadata"): + self.kv_storage_metadata["pointer_generation"] = getattr( + self.kv_cache, "pointer_generation", None + ) # ======================= Attention & MoE backend initialization ======================== self.ctx.attn_backend = self.attn_backend = create_attention_backend( @@ -413,6 +539,13 @@ def __init__(self, config: EngineConfig): # ======================= Sampler initialization ======================== self.sampler = Sampler(self.device, config.model_config.vocab_size) + graph_capacity = max(config.cuda_graph_bs or [0]) + self.token_chain = DeviceTokenChain( + self.device, + max(config.max_running_req, config.cuda_graph_max_bs or 0, graph_capacity, 1), + ) + # Compatibility alias for diagnostics and existing callers. + self.token_staging = self.token_chain.staging post_free_memory = self._sync_get_memory()[0] logger.info_rank0(f"Free memory after initialization: {mem_GB(post_free_memory)}") @@ -431,6 +564,7 @@ def __init__(self, config: EngineConfig): if self.linear_state_pool is not None: self.dummy_req.linear_slot_idx = self.linear_state_pool.padding_slot self.page_table[self.dummy_req.table_idx].fill_(num_tokens) # point to dummy page + self._memory_phase_start("capture") self.graph_runner = GraphRunner( stream=self.stream, device=self.device, @@ -443,10 +577,53 @@ def __init__(self, config: EngineConfig): vocab_size=config.model_config.vocab_size, dummy_req=self.dummy_req, moe_offload_cache=self.moe_offload_cache, + sampler=self.sampler, + token_chain=self.token_chain, ) if config.attention_backend.split(",")[0] == "triton": # Prefill runs on the first comma part; warm its autotune cache. self._warmup_prefill() + self._memory_phase_end("capture") + + def _memory_phase_start(self, name: str) -> None: + """Reset allocator peaks and begin a driver-memory phase sample.""" + if not torch.cuda.is_available(): + return + torch.cuda.reset_peak_memory_stats(self.device) + free, total = torch.cuda.mem_get_info(self.device) + self._active_memory_phase = (name, int(free), int(total), int(free)) + + def _memory_phase_poll(self) -> None: + active = getattr(self, "_active_memory_phase", None) + if active is None or not torch.cuda.is_available(): + return + name, start_free, total, minimum_free = active + free, _ = torch.cuda.mem_get_info(self.device) + self._active_memory_phase = (name, start_free, total, min(minimum_free, int(free))) + + def _memory_phase_end(self, name: str) -> None: + """Synchronize and retain allocator plus driver high-water observations.""" + active = getattr(self, "_active_memory_phase", None) + if active is None or not torch.cuda.is_available(): + return + torch.cuda.synchronize(self.device) + self._memory_phase_poll() + _name, start_free, total, minimum_free = self._active_memory_phase + free, _ = torch.cuda.mem_get_info(self.device) + from freetoken.engine.resident_budget import phase_memory + + self.memory_phases.append( + phase_memory( + name, + start_free_bytes=start_free, + end_free_bytes=int(free), + allocator_peak_allocated_bytes=int(torch.cuda.max_memory_allocated(self.device)), + allocator_peak_reserved_bytes=int(torch.cuda.max_memory_reserved(self.device)), + minimum_driver_free_bytes=minimum_free, + total_driver_bytes=total, + ) + ) + self._active_memory_phase = None def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: if config.tp_info.size == 1 or config.use_pynccl: @@ -775,6 +952,10 @@ def _resize_kv_pool(self, config, num_pages: int, num_swa_pages: int | None) -> self.model.mark_for_rebind() self.kv_cache.rebuild_from_config(config, num_pages, num_swa_pages=num_swa_pages) self.num_pages = num_pages + if hasattr(self, "kv_storage_metadata"): + self.kv_storage_metadata["pointer_generation"] = getattr( + self.kv_cache, "pointer_generation", None + ) def _refresh_seq_state(self, config) -> None: num_tokens = self.num_pages * config.page_size @@ -936,32 +1117,57 @@ def rebuild_runtime_cache( vocab_size=config.model_config.vocab_size, dummy_req=self.dummy_req, moe_offload_cache=self.moe_offload_cache, + sampler=self.sampler, + token_chain=self.token_chain, ) def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: assert torch.cuda.current_stream() == self.stream - # Inc 2 instrument of .plans/rocm-perf-parity: stage-time breakdowns via - # FREETOKEN_TORCH_PROFILE (no-op/cached-flag when unset). Wraps the whole - # forward+sample step; range labels live at the MoE/router/attention callsites. - with step_profiler(): - with self.ctx.forward_batch(batch): - if self.graph_runner.can_use_cuda_graph(batch): - logits = self.graph_runner.replay(batch) - else: + captured_tokens = None + # Scheduler owns the outer Inc2 step-profiler scope so its trace includes scheduling, + # graph input copies, engine forward/sample, and result drain in one window. + with self.ctx.forward_batch(batch): + with profiler_phase("model_forward_setup"): + use_graph = self.graph_runner.can_use_cuda_graph(batch) + if use_graph: + with profiler_phase("graph_replay_submission"): + logits, captured_tokens = self.graph_runner.replay(batch, args) + else: + with profiler_phase("model_forward"): logits = self.model.forward() - if self.cpu_moe_executor is not None: - # One pinned read: surfaces a fired flag-handshake watchdog (dead coordinator - # -> stale expert outputs) as a loud error instead of silent corruption. - self.cpu_moe_executor.raise_if_unhealthy() - - for req in batch.reqs: - req.complete_one() - - batch_logits = logits[: batch.size] - next_tokens_gpu = self.sampler.sample(batch_logits, args, batch).to(torch.int32) - next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True) - copy_done_event = torch.cuda.Event() - copy_done_event.record(self.stream) + if self.cpu_moe_executor is not None: + # One pinned read: surfaces a fired flag-handshake watchdog (dead coordinator + # -> stale expert outputs) as a loud error instead of silent corruption. + self.cpu_moe_executor.raise_if_unhealthy() + + for req in batch.reqs: + req.complete_one() + + batch_logits = logits[: batch.size] + singleton_device_chain = ( + batch.is_decode + and batch.size == 1 + and self.sampler.capture_safe(args) + ) + if captured_tokens is None and singleton_device_chain: + with profiler_phase("sampler"): + sampled_tokens = self.sampler.sample_into_device( + batch_logits, + args, + batch, + self.token_chain.next_device_tokens()[: batch.size], + self.token_chain.sampled_indices[: batch.size], + ) + elif captured_tokens is None: + with profiler_phase("sampler"): + sampled_tokens = self.sampler.sample(batch_logits, args, batch) + else: + with profiler_phase("token_handoff"): + sampled_tokens = self.token_chain.publish(captured_tokens, self.stream) + with profiler_phase("token_d2h_copy"): + next_tokens_gpu, next_tokens_cpu, copy_done_event = self.token_chain.stage( + sampled_tokens, self.stream + ) return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event) @torch.inference_mode() @@ -1254,6 +1460,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution object.__setattr__(config, attr, value) model_config = config.model_config + validate_kv_storage_config(config) single_stream_only = getattr(model_config, "single_stream_only", False) is_dsv4 = getattr(model_config, "dsv4_args", None) is not None has_swa_attention = getattr(model_config, "has_swa_attention", False) @@ -1357,6 +1564,11 @@ def override(attr: str, value: Any): # this is dangerous, use with caution ) logger.info_rank0(f"Auto-selected attention backend: {config.attention_backend}") _validate_attention_backend_choice(config, override, required_attn_types) + if getattr(config, "kv_storage_type", None) is not None and str(config.kv_storage_type) in { + "KVStorageType.Q8_0", "q8_0" + }: + if any(part.strip() != "triton" for part in config.attention_backend.split(",")): + raise ValueError("q8_0 KV storage requires --attention-backend triton") if config.moe_cache_rate is not None: total_experts = config.model_config.num_moe_layers * config.model_config.num_experts @@ -1495,7 +1707,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution if ( is_moe - and expert_quant not in ("none", "fp8_block") + and expert_quant not in ("none", "fp8_block", "gguf") and not is_offload_moe_backend(config.moe_backend) ): raise ValueError( diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 24127c235..b8139cb2e 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -1,41 +1,89 @@ from __future__ import annotations import gc +import os from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List +from typing import TYPE_CHECKING, Any, Dict, List import torch from freetoken.core import Batch, Req, get_global_ctx from freetoken.distributed import get_tp_info from freetoken.utils import init_logger, mem_GB from freetoken.utils.progress import emit_progress +from freetoken.utils.step_profiler import profiler_phase from tqdm import tqdm if TYPE_CHECKING: from freetoken.attention import BaseAttnBackend from freetoken.models import BaseLLMModel from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.engine.sample import BatchSamplingArgs, Sampler logger = init_logger(__name__) +def _has_weight_format(value, wanted: str, seen: set[int] | None = None) -> bool: + """Walk FreeToken's BaseOP tree without assuming torch.nn.Module APIs.""" + if seen is None: + seen = set() + if value is None or id(value) in seen or isinstance(value, torch.Tensor): + return False + seen.add(id(value)) + if getattr(value, "weight_format", None) == wanted: + return True + if isinstance(value, (list, tuple)): + return any(_has_weight_format(item, wanted, seen) for item in value) + attrs = getattr(value, "__dict__", None) + return bool(attrs) and any(_has_weight_format(item, wanted, seen) for item in attrs.values()) + + @dataclass class GraphCaptureBuffer: input_ids: torch.Tensor out_loc: torch.Tensor positions: torch.Tensor logits: torch.Tensor + sampled_tokens: torch.Tensor + sampled_indices: torch.Tensor table_idx: torch.Tensor # per-request slot id for GatedDeltaNet state gather/scatter # Decode GDN query indptr = arange(bs+1); a constant per captured bs, filled once. fla_cu_seqlens: torch.Tensor @classmethod - def init(cls, bs: int, vocab_size: int, device: torch.device) -> GraphCaptureBuffer: + def init( + cls, + bs: int, + vocab_size: int, + device: torch.device, + *, + sampled_tokens: torch.Tensor | None = None, + sampled_indices: torch.Tensor | None = None, + ) -> GraphCaptureBuffer: + if sampled_tokens is None: + sampled_tokens = torch.empty(bs, dtype=torch.int32, device=device) + if sampled_indices is None: + sampled_indices = torch.empty(bs, dtype=torch.int64, device=device) + if ( + sampled_tokens.ndim != 1 + or sampled_tokens.shape[0] < bs + or sampled_tokens.dtype != torch.int32 + or sampled_tokens.device != device + ): + raise ValueError("sampled token chain must be a device int32 vector with graph capacity") + if ( + sampled_indices.ndim != 1 + or sampled_indices.shape[0] < bs + or sampled_indices.dtype != torch.int64 + or sampled_indices.device != device + ): + raise ValueError("sampled index scratch must be a device int64 vector with graph capacity") return GraphCaptureBuffer( input_ids=torch.zeros(bs, dtype=torch.int32, device=device), out_loc=torch.zeros(bs, dtype=torch.int32, device=device), positions=torch.zeros(bs, dtype=torch.int32, device=device), logits=torch.empty(bs, vocab_size, dtype=torch.float32, device=device), + sampled_tokens=sampled_tokens[:bs], + sampled_indices=sampled_indices[:bs], table_idx=torch.zeros(bs, dtype=torch.int32, device=device), fla_cu_seqlens=torch.arange(bs + 1, dtype=torch.int32, device=device), ) @@ -56,13 +104,14 @@ def set_batch(self, batch: Batch) -> None: ) def copy_from(self, batch: Batch) -> None: - _slice = slice(batch.padded_size) - self.input_ids[_slice] = batch.input_ids - if batch.out_loc is not None: - self.out_loc[_slice] = batch.out_loc - self.positions[_slice] = batch.positions - if batch.linear_table_idx is not None: - self.table_idx[_slice] = batch.linear_table_idx + with profiler_phase("graph_input_copy"): + _slice = slice(batch.padded_size) + self.input_ids[_slice] = batch.input_ids + if batch.out_loc is not None: + self.out_loc[_slice] = batch.out_loc + self.positions[_slice] = batch.positions + if batch.linear_table_idx is not None: + self.table_idx[_slice] = batch.linear_table_idx def _determine_cuda_graph_bs( @@ -105,6 +154,8 @@ def __init__( vocab_size: int, dummy_req: Req, moe_offload_cache: OffloadMoeCache | None = None, + sampler: "Sampler | None" = None, + token_chain: Any | None = None, ) -> None: cuda_graph_bs = _determine_cuda_graph_bs( cuda_graph_bs=cuda_graph_bs, @@ -116,10 +167,41 @@ def __init__( self.graph_bs_list = sorted(cuda_graph_bs) self.dummy_req = dummy_req self.moe_offload_cache = moe_offload_cache + self.sampler = sampler + self.token_chain = token_chain + self.resident_gguf = _has_weight_format(model, "gguf") + self.graph_telemetry = { + "expert_storage": "resident_gguf" if self.resident_gguf else "offload_or_dense", + "expert_fetches": 0, + "expert_remaps": 0, + } + self.capture_sampler = sampler is not None and os.environ.get( + "FREETOKEN_GRAPH_SAMPLER", "0" + ).strip().lower() in {"1", "true", "yes", "on"} self.stream = stream self.device = device self._capture_graphs(max_seq_len, vocab_size, model) + def runtime_telemetry(self) -> dict: + """Return graph/storage facts without synchronizing or touching model state.""" + try: + from freetoken.kernel.gguf import gguf_dispatch_report + + dispatch = gguf_dispatch_report() + except Exception: # noqa: BLE001 -- optional GGUF telemetry + dispatch = [] + result = { + **self.graph_telemetry, + "resident_gguf": self.resident_gguf, + "graph_batches": sorted(self.graph_map), + "sampler_graph_batches": sorted(self.sampler_graph_map), + } + if dispatch and os.environ.get("FREETOKEN_GGUF_DISPATCH_TRACE", "").lower() in { + "1", "true", "yes", "on" + }: + result["gguf_dispatch"] = dispatch + return result + def _reset_moe_offload_cache(self) -> None: if self.moe_offload_cache is not None: self.moe_offload_cache.reset() @@ -132,11 +214,15 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel # graphs-disabled early return so that config gets the phase too. emit_progress("Capturing CUDA graphs / warming up", 0, 0) self.graph_map: Dict[int, torch.cuda.CUDAGraph] = {} + self.sampler_graph_map: Dict[int, torch.cuda.CUDAGraph] = {} # ROCm parity: honour the graph-capture gate result. If capture is not # viable on this AMD card, skip graphs entirely so decode uses the kernel-launch # path (correct, just not graph-accelerated) rather than erroring mid-capture. from freetoken.utils.arch import is_rocm - from freetoken.utils.graph_gate import graph_capture_status, run_graph_gate + from freetoken.utils.graph_gate import graph_capture_status, rocm_blas_report, run_graph_gate + + if is_rocm(): + logger.info_rank0(f"graph capture BLAS policy: {rocm_blas_report()}") if is_rocm() and graph_capture_status() == "fail": # Variant detail matters: the all-variants record is what closes the thread. @@ -150,6 +236,13 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel if self.max_graph_bs == 0: return logger.info_rank0("CUDA graph is disabled.") + # A forced/selected GGUF candidate must compile, execute both measured + # quant paths, and synchronize before any graph captures its addresses. + # Auto mode falls back to legacy inside this gate; forced gfx1100 errors. + from freetoken.kernel.gguf import ensure_gguf_moe_candidate_ready + if ensure_gguf_moe_candidate_ready(): + logger.info_rank0("gfx1100 GGUF MoE candidate compile/self-test passed before graph capture") + self.attn_backend.init_capture_graph(max_seq_len=max_seq_len, bs_list=self.graph_bs_list) torch.cuda.synchronize(self.device) @@ -160,7 +253,18 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel free_memory = get_free_memory(self.device) logger.info_rank0(f"Free GPU memory before capturing CUDA graphs: {mem_GB(free_memory)}") - self.buffer = GraphCaptureBuffer.init(self.max_graph_bs, vocab_size, self.device) + sampled_tokens = None + sampled_indices = None + if self.token_chain is not None: + sampled_tokens = self.token_chain.device_tokens + sampled_indices = self.token_chain.sampled_indices + self.buffer = GraphCaptureBuffer.init( + self.max_graph_bs, + vocab_size, + self.device, + sampled_tokens=sampled_tokens, + sampled_indices=sampled_indices, + ) self._reset_moe_offload_cache() pbar = tqdm( @@ -196,23 +300,58 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel if pool is None: pool = graph.pool() # reuse cuda graph handle to reduce memory self.graph_map[bs] = graph + if self.capture_sampler: + from freetoken.engine.sample import BatchSamplingArgs + + sampler_graph = torch.cuda.CUDAGraph() + try: + with torch.cuda.graph(sampler_graph, pool=pool, stream=self.stream): + self.sampler.sample_into_device( + self.buffer.logits[:bs], + BatchSamplingArgs(temperatures=None), + batch, + self.buffer.sampled_tokens[:bs], + self.buffer.sampled_indices[:bs], + ) + except Exception as exc: # graph stage is optional; model graph remains valid + logger.warning_rank0( + f"greedy sampler graph disabled for bs={bs}; using fallback sampler " + f"({type(exc).__name__}: {str(exc)[:160]})" + ) + else: + self.sampler_graph_map[bs] = sampler_graph self._reset_moe_offload_cache() free_memory = get_free_memory(self.device) logger.info_rank0(f"Free GPU memory after capturing CUDA graphs: {mem_GB(free_memory)}") + logger.info_rank0(f"GGUF graph telemetry: {self.runtime_telemetry()}") def can_use_cuda_graph(self, batch: Batch) -> bool: # ``self.graph_map`` is empty when graphs were skipped (ROCm graph-gate fail or # disabled); decode must then fall back to the kernel-launch path. return bool(self.graph_map) and batch.is_decode and batch.size <= self.max_graph_bs - def replay(self, batch: Batch) -> torch.Tensor: + def replay( + self, batch: Batch, sample_args: "BatchSamplingArgs | None" = None + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]: assert self.can_use_cuda_graph(batch) - self.buffer.copy_from(batch) - g = self.graph_map[batch.padded_size] - self.attn_backend.prepare_for_replay(batch) - g.replay() - return self.buffer.logits[: batch.size] + with profiler_phase("graph_replay"): + self.buffer.copy_from(batch) + g = self.graph_map[batch.padded_size] + self.attn_backend.prepare_for_replay(batch) + g.replay() + logits = self.buffer.logits[: batch.size] + if sample_args is None: + return logits + sampler_graph = self.sampler_graph_map.get(batch.padded_size) + sampled = ( + self.buffer.sampled_tokens[: batch.size] + if sampler_graph is not None and self.sampler.capture_safe(sample_args) + else None + ) + if sampled is not None: + sampler_graph.replay() + return logits, sampled def pad_batch(self, batch: Batch) -> None: padded_size = ( # choose the first available batch size @@ -230,5 +369,6 @@ def destroy_cuda_graphs(self) -> None: # free-before-alloc cannot reclaim this GPU memory. empty_cache() is left to the # caller / next capture (GraphRunner._capture_graphs already runs it). self.graph_map = {} + self.sampler_graph_map = {} self.buffer = None gc.collect() diff --git a/python/freetoken/engine/resident_budget.py b/python/freetoken/engine/resident_budget.py new file mode 100644 index 000000000..c6a646dff --- /dev/null +++ b/python/freetoken/engine/resident_budget.py @@ -0,0 +1,236 @@ +"""Allocation-free fit planning for native GGUF resident MoE execution.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Iterable, Literal + +_MIB = 1 << 20 +_GIB = 1 << 30 +GRAPH_RESERVE_BYTES = 768 * _MIB +MIN_LOAD_SCRATCH_BYTES = 512 * _MIB +MIN_SAFETY_BYTES = 1_500 * _MIB + + +@dataclass(frozen=True) +class PhaseMemory: + """One synchronized allocator/driver high-water observation.""" + + name: str + start_free_bytes: int + end_free_bytes: int + allocator_peak_allocated_bytes: int + allocator_peak_reserved_bytes: int + minimum_driver_free_bytes: int + total_driver_bytes: int + + @property + def driver_used_high_water_bytes(self) -> int: + return self.total_driver_bytes - self.minimum_driver_free_bytes + + @property + def non_torch_bytes(self) -> int: + # Driver usage already includes allocator-reserved memory. This is a diagnostic + # remainder, never an additive term in required_bytes. + return max(0, self.driver_used_high_water_bytes - self.allocator_peak_reserved_bytes) + + @property + def required_bytes(self) -> int: + return max(self.driver_used_high_water_bytes, self.allocator_peak_reserved_bytes) + + def as_dict(self) -> dict[str, int | str]: + return { + "name": self.name, + "start_free_bytes": self.start_free_bytes, + "end_free_bytes": self.end_free_bytes, + "allocator_peak_allocated_bytes": self.allocator_peak_allocated_bytes, + "allocator_peak_reserved_bytes": self.allocator_peak_reserved_bytes, + "minimum_driver_free_bytes": self.minimum_driver_free_bytes, + "total_driver_bytes": self.total_driver_bytes, + "driver_used_high_water_bytes": self.driver_used_high_water_bytes, + "non_torch_bytes": self.non_torch_bytes, + "required_bytes": self.required_bytes, + } + + +def phase_memory( + name: str, + *, + start_free_bytes: int, + end_free_bytes: int, + allocator_peak_allocated_bytes: int, + allocator_peak_reserved_bytes: int, + minimum_driver_free_bytes: int, + total_driver_bytes: int, +) -> PhaseMemory: + """Build a phase observation; caller supplies synchronized device counters.""" + if minimum_driver_free_bytes < 0 or minimum_driver_free_bytes > total_driver_bytes: + raise ValueError("minimum driver free must be within total device memory") + if allocator_peak_allocated_bytes < 0 or allocator_peak_reserved_bytes < 0: + raise ValueError("allocator peaks must be non-negative") + return PhaseMemory( + name=name, + start_free_bytes=int(start_free_bytes), + end_free_bytes=int(end_free_bytes), + allocator_peak_allocated_bytes=int(allocator_peak_allocated_bytes), + allocator_peak_reserved_bytes=int(allocator_peak_reserved_bytes), + minimum_driver_free_bytes=int(minimum_driver_free_bytes), + total_driver_bytes=int(total_driver_bytes), + ) + + +@dataclass(frozen=True) +class ResidentBudget: + free_bytes: int + total_vram_bytes: int + packed_model_bytes: int + kv_bytes: int + gdn_state_bytes: int + page_table_bytes: int + graph_reserve_bytes: int + peak_load_scratch_bytes: int + safety_bytes: int + phases: tuple[PhaseMemory, ...] = () + + @property + def required_bytes(self) -> int: + static = sum( + ( + self.packed_model_bytes, + self.kv_bytes, + self.gdn_state_bytes, + self.page_table_bytes, + self.graph_reserve_bytes, + self.peak_load_scratch_bytes, + self.safety_bytes, + ) + ) + observed = max((phase.required_bytes for phase in self.phases), default=0) + return max(static, observed) + + @property + def fits(self) -> bool: + return self.required_bytes <= self.free_bytes + + def as_dict(self) -> dict[str, int | bool]: + return { + "free_bytes": self.free_bytes, + "total_vram_bytes": self.total_vram_bytes, + "packed_model_bytes": self.packed_model_bytes, + "kv_bytes": self.kv_bytes, + "gdn_state_bytes": self.gdn_state_bytes, + "page_table_bytes": self.page_table_bytes, + "graph_reserve_bytes": self.graph_reserve_bytes, + "peak_load_scratch_bytes": self.peak_load_scratch_bytes, + "safety_bytes": self.safety_bytes, + "required_bytes": self.required_bytes, + "fits": self.fits, + "phases": [phase.as_dict() for phase in self.phases], + "observed_required_bytes": max((phase.required_bytes for phase in self.phases), default=0), + } + + +def required_phase_bytes(phases: Iterable[PhaseMemory], safety_bytes: int = 0) -> int: + """Worst observed driver high-water plus explicit safety, without double counting.""" + return max((phase.required_bytes for phase in phases), default=0) + int(safety_bytes) + + +def _gguf_payload_bytes(model_path: str) -> tuple[int, int]: + """Return (all packed tensor bytes, largest tensor bytes) from GGUF headers.""" + from freetoken.models.gguf.reader import _reader + + reader = _reader(model_path) + total = 0 + largest = 0 + import gguf + + for tensor in reader.tensors: + shape = [int(dim) for dim in tensor.shape] + block, type_size = gguf.GGML_QUANT_SIZES[tensor.tensor_type] + fastest = shape[0] + if fastest % block: + raise ValueError( + f"{tensor.name}: fastest dimension {fastest} is not a multiple of {block}" + ) + size = math.prod(shape) // block * type_size + total += size + largest = max(largest, size) + return total, largest + + +def _total_vram_bytes(free_bytes: int) -> int: + try: + import torch + + if torch.cuda.is_available(): + return int(torch.cuda.get_device_properties(torch.cuda.current_device()).total_memory) + except Exception: + pass + return free_bytes + + +def _page_table_width(max_seq_len: int, page_size: int) -> int: + page_aligned = ((max_seq_len + page_size - 1) // page_size) * page_size + return ((page_aligned + 31) // 32) * 32 + + +def estimate_gguf_resident_budget(model_path, config, free_bytes: int) -> ResidentBudget: + """Estimate complete native-resident startup footprint without CUDA allocation.""" + packed_model_bytes, largest_tensor_bytes = _gguf_payload_bytes(model_path) + + from freetoken.kvcache import resolve_pool_class + from freetoken.kvcache.linear_state_pool import _linear_pool_num_slots, state_pool_bytes + + pool_cls = resolve_pool_class(config.model_config) + per_page, fixed, page_tokens, _ = pool_cls.kv_cost(config) + pages = max(1, (int(config.max_seq_len) + page_tokens - 1) // page_tokens) + kv_bytes = pages * per_page + fixed + # _adjust_config changes linear models' default radix cache to hybrid_radix after this + # preflight. Price that final slot geometry without mutating the frozen config. + linear = config.model_config.linear_attention_group() + cache_type = getattr(config, "cache_type", "radix") + if linear is not None and cache_type != "naive": + max_req = int(config.max_running_req) + cache_slots = max(4, int(getattr(config, "linear_state_cache_ratio", 2.0) * max_req)) + state_slots = 4 * max_req + cache_slots + 1 + else: + state_slots = _linear_pool_num_slots(config) + gdn_state = state_pool_bytes(config, num_slots=state_slots) + page_table = ( + (int(config.max_running_req) + 1) + * _page_table_width(int(config.max_seq_len), int(config.page_size)) + * 4 + ) + peak_scratch = max(MIN_LOAD_SCRATCH_BYTES, largest_tensor_bytes) + total_vram = _total_vram_bytes(int(free_bytes)) + safety = max(MIN_SAFETY_BYTES, int(total_vram * 0.08)) + return ResidentBudget( + free_bytes=int(free_bytes), + total_vram_bytes=total_vram, + packed_model_bytes=packed_model_bytes, + kv_bytes=kv_bytes, + gdn_state_bytes=int(gdn_state), + page_table_bytes=page_table, + graph_reserve_bytes=GRAPH_RESERVE_BYTES, + peak_load_scratch_bytes=peak_scratch, + safety_bytes=safety, + ) + + +def resolve_gguf_moe_backend(config, free_bytes: int) -> Literal["fused", "offload"]: + """Resolve GGUF auto mode; callers handle explicit fused failure details.""" + if getattr(config.model_config, "moe_weight_format", None) != "gguf": + return "offload" + budget = estimate_gguf_resident_budget(config.model_path, config, int(free_bytes)) + return "fused" if budget.fits else "offload" + + +__all__ = [ + "PhaseMemory", + "ResidentBudget", + "estimate_gguf_resident_budget", + "phase_memory", + "required_phase_bytes", + "resolve_gguf_moe_backend", +] diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index be3726b4e..0519c2fae 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -1,6 +1,6 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, List import torch @@ -84,6 +84,7 @@ def sample_impl( class Sampler: device: torch.device vocab_size: int + _args_cache: dict[tuple, BatchSamplingArgs] = field(default_factory=dict, init=False, repr=False) def prepare(self, batch: Batch) -> BatchSamplingArgs: params = [r.sampling_params for r in batch.reqs] @@ -93,6 +94,20 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: if all(p.is_greedy for p in params) and not apply_penalties: return BatchSamplingArgs(temperatures=None) + key = tuple( + ( + p.temperature, + p.top_k, + p.top_p, + p.presence_penalty, + p.frequency_penalty, + ) + for p in params + ) + cached = self._args_cache.get(key) + if cached is not None: + return cached + MIN_P = MIN_T = 1e-6 ts = [max(0.0 if p.is_greedy else p.temperature, MIN_T) for p in params] top_ks = [p.top_k if p.top_k >= 1 else self.vocab_size for p in params] @@ -103,20 +118,69 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: top_k = make_device_tensor(top_ks, torch.int32, self.device) if any(p < 1.0 for p in top_ps): top_p = make_device_tensor(top_ps, torch.float32, self.device) - return BatchSamplingArgs( + args = BatchSamplingArgs( temperatures, top_k=top_k, top_p=top_p, apply_penalties=apply_penalties, ) + # Request sampling combinations are normally stable for a decode stream. Keep cache + # bounded so varied multi-request traffic cannot retain unbounded device tensors. + if len(self._args_cache) >= 128: + self._args_cache.clear() + self._args_cache[key] = args + return args @nvtx_annotate("Sampler") def sample( self, logits: torch.Tensor, args: BatchSamplingArgs, batch: Batch ) -> torch.Tensor: - with torch.cuda.nvtx.range("Sampler"): + with torch.profiler.record_function("sampler_impl"): if args.apply_penalties: apply_penalties(logits, batch.reqs) if args.temperatures is None: # greedy sampling return torch.argmax(logits, dim=-1) return sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) + + @staticmethod + def capture_safe(args: BatchSamplingArgs) -> bool: + """Whether sampling can run inside a graph without changing semantics.""" + # Graph capture owns no RNG state. Dynamic sampling and penalties must retain the + # existing post-forward path so temperature/top-k/top-p/penalty behavior survives. + return args.temperatures is None and not args.apply_penalties + + def sample_into( + self, + logits: torch.Tensor, + args: BatchSamplingArgs, + batch: Batch | None, + out: torch.Tensor, + scratch: torch.Tensor | None = None, + ) -> torch.Tensor: + """Capture-safe fixed-address greedy sample into ``out``. + + ``batch`` stays in the signature so callers can share the sampler contract; it is + intentionally unused for this mode. Unsupported modes fail loudly instead of + silently becoming greedy. + """ + del batch + if not self.capture_safe(args): + raise ValueError("sample_into supports greedy sampling without penalties only") + if out.ndim != 1 or out.shape[0] < logits.shape[0] or out.dtype != torch.int32: + raise ValueError("sample_into output must be reusable int32 vector") + if scratch is None or scratch.ndim != 1 or scratch.shape[0] < logits.shape[0] or scratch.dtype != torch.int64: + raise ValueError("sample_into scratch must be reusable int64 vector") + torch.argmax(logits, dim=-1, out=scratch[: logits.shape[0]]) + out[: logits.shape[0]].copy_(scratch[: logits.shape[0]]) + return out[: logits.shape[0]] + + def sample_into_device( + self, + logits: torch.Tensor, + args: BatchSamplingArgs, + batch: Batch | None, + out: torch.Tensor, + scratch: torch.Tensor, + ) -> torch.Tensor: + """Write capture-safe sampled IDs into caller-owned device storage.""" + return self.sample_into(logits, args, batch, out, scratch) diff --git a/python/freetoken/kernel/__init__.py b/python/freetoken/kernel/__init__.py index 5dc515711..be3be6e9b 100644 --- a/python/freetoken/kernel/__init__.py +++ b/python/freetoken/kernel/__init__.py @@ -8,6 +8,7 @@ get_fp4_lut, moe_align_block_size_triton, moe_sum_reduce_triton, + moe_weighted_sum_reduce_triton, mxfp4_fused_moe_kernel_t_triton, mxfp4_splitk_gemv_triton, ) @@ -35,6 +36,7 @@ "gpt_oss_swiglu_triton", "moe_align_block_size_triton", "moe_sum_reduce_triton", + "moe_weighted_sum_reduce_triton", "create_pinned_tensor_like", "copy_to_pinned_tensor", ] diff --git a/python/freetoken/kernel/csrc/gguf/dequantize_hip.cuh b/python/freetoken/kernel/csrc/gguf/dequantize_hip.cuh new file mode 100644 index 000000000..005842f12 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/dequantize_hip.cuh @@ -0,0 +1,585 @@ +// !!! This is a file automatically generated by hipify!!! +#include "hip/hip_runtime.h" +// copied from +// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/dequantize.cuh +// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/convert.cu +// Dequant functions +static __device__ __forceinline__ void dequantize_q4_0(const void* vx, const int ib, const int iqs, dfloat2& v) { + const block_q4_0* x = (const block_q4_0*)vx; + + const dfloat d = x[ib].d; + + const int vui = x[ib].qs[iqs]; + + v.x = __int2half_rn(vui & 0xF); + v.y = __int2half_rn(vui >> 4); + + v = __hsub2(v, __floats2half2_rn(8.0f, 8.0f)); + v = __hmul2(v, {d, d}); +} + +static __device__ __forceinline__ void dequantize_q4_1(const void* vx, const int ib, const int iqs, dfloat2& v) { + const block_q4_1* x = (const block_q4_1*)vx; + + const dfloat d = __low2half(x[ib].dm); + const dfloat m = __high2half(x[ib].dm); + + const int vui = x[ib].qs[iqs]; + + v.x = __int2half_rn(vui & 0xF); + v.y = __int2half_rn(vui >> 4); + + v = __hmul2(v, {d, d}); + v = __hadd2(v, {m, m}); +} + +static __device__ __forceinline__ void dequantize_q5_0(const void* vx, const int ib, const int iqs, dfloat2& v) { + const block_q5_0* x = (const block_q5_0*)vx; + + const dfloat d = x[ib].d; + + uint32_t qh; + memcpy(&qh, x[ib].qh, sizeof(qh)); + + const int xh_0 = ((qh >> (iqs + 0)) << 4) & 0x10; + const int xh_1 = ((qh >> (iqs + 12))) & 0x10; + + v.x = __int2half_rn((x[ib].qs[iqs] & 0xf) | xh_0); + v.y = __int2half_rn((x[ib].qs[iqs] >> 4) | xh_1); + + v = __hsub2(v, __floats2half2_rn(16.0f, 16.0f)); + v = __hmul2(v, {d, d}); +} + +static __device__ __forceinline__ void dequantize_q5_1(const void* vx, const int ib, const int iqs, dfloat2& v) { + const block_q5_1* x = (const block_q5_1*)vx; + + const dfloat d = __low2half(x[ib].dm); + const dfloat m = __high2half(x[ib].dm); + + uint32_t qh; + memcpy(&qh, x[ib].qh, sizeof(qh)); + + const int xh_0 = ((qh >> (iqs + 0)) << 4) & 0x10; + const int xh_1 = ((qh >> (iqs + 12))) & 0x10; + + v.x = __int2half_rn((x[ib].qs[iqs] & 0xf) | xh_0); + v.y = __int2half_rn((x[ib].qs[iqs] >> 4) | xh_1); + + v = __hmul2(v, {d, d}); + v = __hadd2(v, {m, m}); +} + +static __device__ __forceinline__ void dequantize_q8_0(const void* vx, const int ib, const int iqs, dfloat2& v) { + const block_q8_0* x = (const block_q8_0*)vx; + + const dfloat d = x[ib].d; + + v.x = __int2half_rn(x[ib].qs[iqs + 0]); + v.y = __int2half_rn(x[ib].qs[iqs + 1]); + + v = __hmul2(v, {d, d}); +} + +template +static __global__ void dequantize_block(const void* __restrict__ vx, dst_t* __restrict__ y, const int k) { + const int i = 2 * (blockDim.x * blockIdx.x + threadIdx.x); + + if (i >= k) { + return; + } + + const int ib = i / qk; // block index + const int iqs = (i % qk) / qr; // quant index + const int iybs = i - i % qk; // y block start index + const int y_offset = qr == 1 ? 1 : qk / 2; + + // dequantize + dfloat2 v; + dequantize_kernel(vx, ib, iqs, v); + + y[iybs + iqs + 0] = convert_from_half(v.x); + y[iybs + iqs + y_offset] = convert_from_half(v.y); +} + +template +static __global__ void dequantize_block_q2_K(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const auto i = blockIdx.x; + const block_q2_K* x = (const block_q2_K*)vx; + + const auto tid = threadIdx.x; + const int n = tid / 32; + const int l = tid - 32 * n; + const int is = 8 * n + l / 16; + + const uint8_t q = x[i].qs[32 * n + l]; + dst_t* y = yy + i * QK_K + 128 * n; + + half dall = __low2half(x[i].dm); + half dmin = __high2half(x[i].dm); + y[l + 0] = convert_from_half(__hsub( + __hmul(dall, __int2half_rn((x[i].scales[is + 0] & 0xF) * ((q >> 0) & 3))), + __hmul(dmin, __int2half_rn(x[i].scales[is + 0] >> 4)))); + y[l + 32] = convert_from_half(__hsub( + __hmul(dall, __int2half_rn((x[i].scales[is + 2] & 0xF) * ((q >> 2) & 3))), + __hmul(dmin, __int2half_rn(x[i].scales[is + 2] >> 4)))); + y[l + 64] = convert_from_half(__hsub( + __hmul(dall, __int2half_rn((x[i].scales[is + 4] & 0xF) * ((q >> 4) & 3))), + __hmul(dmin, __int2half_rn(x[i].scales[is + 4] >> 4)))); + y[l + 96] = convert_from_half(__hsub( + __hmul(dall, __int2half_rn((x[i].scales[is + 6] & 0xF) * ((q >> 6) & 3))), + __hmul(dmin, __int2half_rn(x[i].scales[is + 6] >> 4)))); +} + +template +static __global__ void dequantize_block_q3_K(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const auto i = blockIdx.x; + const block_q3_K* x = (const block_q3_K*)vx; + + const auto r = threadIdx.x / 4; + const int tid = r / 2; + const int is0 = r % 2; + const int l0 = 16 * is0 + 4 * (threadIdx.x % 4); + const int n = tid / 4; + const int j = tid - 4 * n; + + uint8_t m = 1 << (4 * n + j); + int is = 8 * n + 2 * j + is0; + int shift = 2 * j; + + int8_t us = is < 4 ? (x[i].scales[is - 0] & 0xF) | (((x[i].scales[is + 8] >> 0) & 3) << 4) + : is < 8 ? (x[i].scales[is - 0] & 0xF) | (((x[i].scales[is + 4] >> 2) & 3) << 4) + : is < 12 ? (x[i].scales[is - 8] >> 4) | (((x[i].scales[is + 0] >> 4) & 3) << 4) + : (x[i].scales[is - 8] >> 4) | (((x[i].scales[is - 4] >> 6) & 3) << 4); + half d_all = x[i].d; + half dl = __hmul(d_all, __int2half_rn(us - 32)); + + dst_t* y = yy + i * QK_K + 128 * n + 32 * j; + const uint8_t* q = x[i].qs + 32 * n; + const uint8_t* hm = x[i].hmask; + + for (int l = l0; l < l0 + 4; ++l) { + y[l] = convert_from_half(__hmul(dl, __int2half_rn((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)))); + } +} + +static inline __device__ void get_scale_min_k4(int j, const uint8_t* q, uint8_t& d, uint8_t& m) { + if (j < 4) { + d = q[j] & 63; + m = q[j + 4] & 63; + } else { + d = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4); + m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4); + } +} + +template +static __global__ void dequantize_block_q4_K(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const block_q4_K* x = (const block_q4_K*)vx; + + const auto i = blockIdx.x; + + // assume 32 threads + const auto tid = threadIdx.x; + const int il = tid / 8; + const int ir = tid % 8; + const int is = 2 * il; + const int n = 4; + + dst_t* y = yy + i * QK_K + 64 * il + n * ir; + + const half dall = __low2half(x[i].dm); + const half dmin = __high2half(x[i].dm); + + const uint8_t* q = x[i].qs + 32 * il + n * ir; + + uint8_t sc, m; + get_scale_min_k4(is + 0, x[i].scales, sc, m); + const half d1 = __hmul(dall, __int2half_rn(sc)); + const half m1 = __hmul(dmin, __int2half_rn(m)); + get_scale_min_k4(is + 1, x[i].scales, sc, m); + const half d2 = __hmul(dall, __int2half_rn(sc)); + const half m2 = __hmul(dmin, __int2half_rn(m)); + for (int l = 0; l < n; ++l) { + y[l + 0] = convert_from_half(__hsub(__hmul(d1, __int2half_rn(q[l] & 0xF)), m1)); + y[l + 32] = convert_from_half(__hsub(__hmul(d2, __int2half_rn(q[l] >> 4)), m2)); + } +} + +template +static __global__ void dequantize_block_q5_K(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const block_q5_K* x = (const block_q5_K*)vx; + + const auto i = blockIdx.x; + + // assume 64 threads - this is very slightly better than the one below + const auto tid = threadIdx.x; + const int il = tid / 16; // il is in 0...3 + const int ir = tid % 16; // ir is in 0...15 + const int is = 2 * il; // is is in 0...6 + + dst_t* y = yy + i * QK_K + 64 * il + 2 * ir; + + const half dall = __low2half(x[i].dm); + const half dmin = __high2half(x[i].dm); + + const uint8_t* ql = x[i].qs + 32 * il + 2 * ir; + const uint8_t* qh = x[i].qh + 2 * ir; + + uint8_t sc, m; + get_scale_min_k4(is + 0, x[i].scales, sc, m); + const half d1 = __hmul(dall, __int2half_rn(sc)); + const half m1 = __hmul(dmin, __int2half_rn(m)); + get_scale_min_k4(is + 1, x[i].scales, sc, m); + const half d2 = __hmul(dall, __int2half_rn(sc)); + const half m2 = __hmul(dmin, __int2half_rn(m)); + + uint8_t hm = 1 << (2 * il); + y[0] = convert_from_half(__hsub(__hmul(d1, __int2half_rn((ql[0] & 0xF) + (qh[0] & hm ? 16 : 0))), m1)); + y[1] = convert_from_half(__hsub(__hmul(d1, __int2half_rn((ql[1] & 0xF) + (qh[1] & hm ? 16 : 0))), m1)); + hm <<= 1; + y[32] = convert_from_half(__hsub(__hmul(d2, __int2half_rn((ql[0] >> 4) + (qh[0] & hm ? 16 : 0))), m2)); + y[33] = convert_from_half(__hsub(__hmul(d2, __int2half_rn((ql[1] >> 4) + (qh[1] & hm ? 16 : 0))), m2)); +} + +template +static __global__ void dequantize_block_q6_K(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const block_q6_K* x = (const block_q6_K*)vx; + + const auto i = blockIdx.x; + + // assume 64 threads - this is very slightly better than the one below + const auto tid = threadIdx.x; + const int ip = tid / 32; // ip is 0 or 1 + const int il = tid - 32 * ip; // 0...32 + const int is = 8 * ip + il / 16; + + dst_t* y = yy + i * QK_K + 128 * ip + il; + + const half d = x[i].d; + + const uint8_t* ql = x[i].ql + 64 * ip + il; + const uint8_t qh = x[i].qh[32 * ip + il]; + const int8_t* sc = x[i].scales + is; + + y[0] = convert_from_half( + __hmul(d, __int2half_rn(sc[0] * ((int8_t)((ql[0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32)))); + y[32] = convert_from_half( + __hmul(d, __int2half_rn(sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32)))); + y[64] = convert_from_half( + __hmul(d, __int2half_rn(sc[4] * ((int8_t)((ql[0] >> 4) | (((qh >> 4) & 3) << 4)) - 32)))); + y[96] = convert_from_half( + __hmul(d, __int2half_rn(sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32)))); +} + +template +static __global__ void dequantize_block_iq2_xxs(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const auto i = blockIdx.x; + const block_iq2_xxs* x = (const block_iq2_xxs*)vx; + + const auto tid = threadIdx.x; + const int il = tid / 8; // 0...3 + const int ib = tid % 8; // 0...7 + dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; + const uint16_t* q2 = x[i].qs + 4 * ib; + const uint8_t* aux8 = (const uint8_t*)q2; + const uint8_t* grid = (const uint8_t*)(iq2xxs_grid + aux8[il]); + const uint32_t aux32 = q2[2] | (q2[3] << 16); + const float d = __half2float(x[i].d) * (0.5f + (aux32 >> 28)) * 0.25f; + const uint8_t signs = ksigns_iq2xs[(aux32 >> 7 * il) & 127]; + for (int j = 0; j < 8; ++j) + y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); +} + +template +static __global__ void dequantize_block_iq2_xs(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const auto i = blockIdx.x; + const block_iq2_xs* x = (const block_iq2_xs*)vx; + + const auto tid = threadIdx.x; + const int il = tid / 8; // 0...3 + const int ib = tid % 8; // 0...7 + dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; + const uint16_t* q2 = x[i].qs + 4 * ib; + const uint8_t* grid = (const uint8_t*)(iq2xs_grid + (q2[il] & 511)); + const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib] >> 4 * (il / 2)) & 0xf)) * 0.25f; + const uint8_t signs = ksigns_iq2xs[q2[il] >> 9]; + for (int j = 0; j < 8; ++j) + y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); +} + +template +static __global__ void dequantize_block_iq2_s(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const auto i = blockIdx.x; + const block_iq2_s* x = (const block_iq2_s*)vx; + + const auto tid = threadIdx.x; + const int il = tid / 8; // 0...3 + const int ib = tid % 8; // 0...7 + dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; + const uint8_t* grid = (const uint8_t*)(iq2s_grid + (x[i].qs[4 * ib + il] | ((x[i].qh[ib] << (8 - 2 * il)) & 0x300))); + const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib] >> 4 * (il / 2)) & 0xf)) * 0.25f; + const uint8_t signs = x[i].qs[QK_K / 8 + 4 * ib + il]; + for (int j = 0; j < 8; ++j) + y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); +} + +template +static __global__ void dequantize_block_iq3_xxs(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const auto i = blockIdx.x; + const block_iq3_xxs* x = (const block_iq3_xxs*)vx; + + const auto tid = threadIdx.x; + const int il = tid / 8; // 0...3 + const int ib = tid % 8; // 0...7 + dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; + const uint8_t* q3 = x[i].qs + 8 * ib; + const uint16_t* gas = (const uint16_t*)(x[i].qs + QK_K / 4) + 2 * ib; + const uint8_t* grid1 = (const uint8_t*)(iq3xxs_grid + q3[2 * il + 0]); + const uint8_t* grid2 = (const uint8_t*)(iq3xxs_grid + q3[2 * il + 1]); + const uint32_t aux32 = gas[0] | (gas[1] << 16); + const float d = __half2float(x[i].d) * (0.5f + (aux32 >> 28)) * 0.5f; + const uint8_t signs = ksigns_iq2xs[(aux32 >> 7 * il) & 127]; + for (int j = 0; j < 4; ++j) { + y[j + 0] = d * grid1[j] * (signs & kmask_iq2xs[j + 0] ? -1.f : 1.f); + y[j + 4] = d * grid2[j] * (signs & kmask_iq2xs[j + 4] ? -1.f : 1.f); + } +} + +template +static __global__ void dequantize_block_iq3_s(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const auto i = blockIdx.x; + const block_iq3_s* x = (const block_iq3_s*)vx; + + const auto tid = threadIdx.x; + const int il = tid / 8; // 0...3 + const int ib = tid % 8; // 0...7 + dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; + const uint8_t* qs = x[i].qs + 8 * ib; + const uint8_t* grid1 = (const uint8_t*)(iq3xs_grid + (qs[2 * il + 0] | ((x[i].qh[ib] << (8 - 2 * il)) & 256))); + const uint8_t* grid2 = (const uint8_t*)(iq3xs_grid + (qs[2 * il + 1] | ((x[i].qh[ib] << (7 - 2 * il)) & 256))); + const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib / 2] >> 4 * (ib % 2)) & 0xf)) * 0.5f; + const uint8_t signs = x[i].signs[4 * ib + il]; + for (int j = 0; j < 4; ++j) { + y[j + 0] = d * grid1[j] * (signs & kmask_iq2xs[j + 0] ? -1.f : 1.f); + y[j + 4] = d * grid2[j] * (signs & kmask_iq2xs[j + 4] ? -1.f : 1.f); + } +} + +template +static __global__ void dequantize_block_iq1_s(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const int64_t i = blockIdx.x; + const block_iq1_s* x = (const block_iq1_s*)vx; + + const int64_t tid = threadIdx.x; + const int64_t il = tid / 8; // 0...3 + const int64_t ib = tid % 8; // 0...7 + dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; + const float delta = x[i].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA; + const float d = __half2float(x[i].d) * (2 * ((x[i].qh[ib] >> 12) & 7) + 1); + uint32_t grid32[2]; + const int8_t* q = (const int8_t*)grid32; + grid32[0] = iq1s_grid_gpu[x[i].qs[4 * ib + il] | (((x[i].qh[ib] >> 3 * il) & 7) << 8)]; + grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f; + grid32[0] &= 0x0f0f0f0f; + for (int j = 0; j < 8; ++j) { + y[j] = d * (q[j] + delta); + } +} + +template +static __global__ void dequantize_block_iq1_m(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const int64_t i = blockIdx.x; + const block_iq1_m* x = (const block_iq1_m*)vx; + + const int64_t tid = threadIdx.x; + const int64_t il = tid / 8; // 0...3 + const int64_t ib = tid % 8; // 0...7 + dst_t* y = yy + i * QK_K + 32 * ib + 8 * il; + const uint16_t* sc = (const uint16_t*)x[i].scales; + iq1m_scale_t scale; + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); + const int64_t ib16 = 2 * ib + il / 2; // sc[ib16/4] >> 3*(ib16%4) -> sc[ib/2] >> 3*((2*ib+il/2)%4); + const float d = __half2float(scale.f16) * (2 * ((sc[ib16 / 4] >> 3 * (ib16 % 4)) & 0x7) + 1); + const float delta = x[i].qh[2 * ib + il / 2] & (0x08 << 4 * (il % 2)) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA; + uint32_t grid32[2]; + const int8_t* q = (const int8_t*)grid32; + grid32[0] = iq1s_grid_gpu[x[i].qs[4 * ib + il] | (((x[i].qh[2 * ib + il / 2] >> 4 * (il % 2)) & 7) << 8)]; + grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f; + grid32[0] &= 0x0f0f0f0f; + for (int j = 0; j < 8; ++j) { + y[j] = d * (q[j] + delta); + } +} + +template +static __global__ void dequantize_block_iq4_nl(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const auto i = blockIdx.x; + const block_iq4_nl* x = (const block_iq4_nl*)vx + i * (QK_K / QK4_NL); + + const auto tid = threadIdx.x; + const int il = tid / 8; // 0...3 + const int ib = tid % 8; // 0...7 + dst_t* y = yy + i * QK_K + 32 * ib + 4 * il; + const uint8_t* q4 = x[ib].qs + 4 * il; + const float d = __half2float(x[ib].d); + for (int j = 0; j < 4; ++j) { + y[j + 0] = d * kvalues_iq4nl[q4[j] & 0xf]; + y[j + 16] = d * kvalues_iq4nl[q4[j] >> 4]; + } +} + +template +static __global__ void dequantize_block_iq4_xs(const void* __restrict__ vx, dst_t* __restrict__ yy) { + const auto i = blockIdx.x; + const block_iq4_xs* x = (const block_iq4_xs*)vx; + + const auto tid = threadIdx.x; + const int il = tid / 8; // 0...3 + const int ib = tid % 8; // 0...7 + dst_t* y = yy + i * QK_K + 32 * ib + 4 * il; + const uint8_t* q4 = x[i].qs + 16 * ib + 4 * il; + const float d = __half2float(x[i].d) * + ((((x[i].scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((x[i].scales_h >> 2 * ib) & 3) << 4)) - 32); + for (int j = 0; j < 4; ++j) { + y[j + 0] = d * kvalues_iq4nl[q4[j] & 0xf]; + y[j + 16] = d * kvalues_iq4nl[q4[j] >> 4]; + } +} + +template +static void +dequantize_block_cuda(const void* __restrict__ vx, dst_t* __restrict__ y, const int k, hipStream_t stream) { + const int num_blocks = (k + 2 * CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / (2 * CUDA_DEQUANTIZE_BLOCK_SIZE); + hipLaunchKernelGGL(( dequantize_block), dim3(num_blocks), dim3(CUDA_DEQUANTIZE_BLOCK_SIZE), 0, stream, vx, y, k); +} + +template +static void dequantize_row_q2_K_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = k / QK_K; + hipLaunchKernelGGL(( dequantize_block_q2_K), dim3(nb), dim3(64), 0, stream, vx, y); +} + +template +static void dequantize_row_q3_K_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = k / QK_K; + hipLaunchKernelGGL(( dequantize_block_q3_K), dim3(nb), dim3(64), 0, stream, vx, y); +} + +template +static void dequantize_row_q4_K_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = k / QK_K; + hipLaunchKernelGGL(( dequantize_block_q4_K), dim3(nb), dim3(32), 0, stream, vx, y); +} + +template +static void dequantize_row_q5_K_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = k / QK_K; + hipLaunchKernelGGL(( dequantize_block_q5_K), dim3(nb), dim3(64), 0, stream, vx, y); +} + +template +static void dequantize_row_q6_K_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = k / QK_K; + hipLaunchKernelGGL(( dequantize_block_q6_K), dim3(nb), dim3(64), 0, stream, vx, y); +} + +template +static void dequantize_row_iq2_xxs_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = k / QK_K; + hipLaunchKernelGGL(( dequantize_block_iq2_xxs), dim3(nb), dim3(32), 0, stream, vx, y); +} + +template +static void dequantize_row_iq2_xs_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = k / QK_K; + hipLaunchKernelGGL(( dequantize_block_iq2_xs), dim3(nb), dim3(32), 0, stream, vx, y); +} + +template +static void dequantize_row_iq2_s_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = k / QK_K; + hipLaunchKernelGGL(( dequantize_block_iq2_s), dim3(nb), dim3(32), 0, stream, vx, y); +} + +template +static void dequantize_row_iq3_xxs_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = k / QK_K; + hipLaunchKernelGGL(( dequantize_block_iq3_xxs), dim3(nb), dim3(32), 0, stream, vx, y); +} + +template +static void dequantize_row_iq3_s_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = k / QK_K; + hipLaunchKernelGGL(( dequantize_block_iq3_s), dim3(nb), dim3(32), 0, stream, vx, y); +} + +template +static void dequantize_row_iq1_s_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = k / QK_K; + hipLaunchKernelGGL(( dequantize_block_iq1_s), dim3(nb), dim3(32), 0, stream, vx, y); +} + +template +static void dequantize_row_iq1_m_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = k / QK_K; + hipLaunchKernelGGL(( dequantize_block_iq1_m), dim3(nb), dim3(32), 0, stream, vx, y); +} + +template +static void dequantize_row_iq4_nl_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = (k + QK_K - 1) / QK_K; + hipLaunchKernelGGL(( dequantize_block_iq4_nl), dim3(nb), dim3(32), 0, stream, vx, y); +} + +template +static void dequantize_row_iq4_xs_cuda(const void* vx, dst_t* y, const int k, hipStream_t stream) { + const int nb = (k + QK_K - 1) / QK_K; + hipLaunchKernelGGL(( dequantize_block_iq4_xs), dim3(nb), dim3(32), 0, stream, vx, y); +} + +template +static to_cuda_ggml_t ggml_get_to_cuda(int64_t type) { + switch (type) { + case 2: + return dequantize_block_cuda; + case 3: + return dequantize_block_cuda; + case 6: + return dequantize_block_cuda; + case 7: + return dequantize_block_cuda; + case 8: + return dequantize_block_cuda; + case 10: + return dequantize_row_q2_K_cuda; + case 11: + return dequantize_row_q3_K_cuda; + case 12: + return dequantize_row_q4_K_cuda; + case 13: + return dequantize_row_q5_K_cuda; + case 14: + return dequantize_row_q6_K_cuda; + case 16: + return dequantize_row_iq2_xxs_cuda; + case 17: + return dequantize_row_iq2_xs_cuda; + case 18: + return dequantize_row_iq3_xxs_cuda; + case 19: + return dequantize_row_iq1_s_cuda; + case 20: + return dequantize_row_iq4_nl_cuda; + case 21: + return dequantize_row_iq3_s_cuda; + case 22: + return dequantize_row_iq2_s_cuda; + case 23: + return dequantize_row_iq4_xs_cuda; + case 29: + return dequantize_row_iq1_m_cuda; + default: + return nullptr; + } +} diff --git a/python/freetoken/kernel/csrc/gguf/ggml-common_hip.h b/python/freetoken/kernel/csrc/gguf/ggml-common_hip.h new file mode 100644 index 000000000..5579c87b4 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/ggml-common_hip.h @@ -0,0 +1,1034 @@ +// !!! This is a file automatically generated by hipify!!! +// adapted from +// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/ggml-common.h +// copied from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-common.h +#define QK_K 256 +#define K_QUANTS_PER_ITERATION 2 +#define WARP_SIZE_GGUF 32 +#define K_SCALE_SIZE 12 +#define CUDA_DEQUANTIZE_BLOCK_SIZE 256 +#define CUDA_QUANTIZE_BLOCK_SIZE 256 +#define GGML_CUDA_DMMV_X 32 +// Same overridable knob as ggml-common_hip.h (kept byte-identical): rows per warp +// for the MMVQ / MoE-vec launches; tune via -DGGML_CUDA_MMV_Y (FREETOKEN_GGUF_MMV_Y). +#ifndef GGML_CUDA_MMV_Y +#define GGML_CUDA_MMV_Y 1 +#endif + +// Data Structures +// QK = number of values after dequantization +// QR = QK / number of values before dequantization +// QI = number of 32 bit integers before dequantization + +#define QK4_0 32 +#define QR4_0 2 +#define QI4_0 (QK4_0 / (4 * QR4_0)) +typedef struct { + half d; // delta + uint8_t qs[QK4_0 / 2]; // nibbles / quants +} block_q4_0; + +#define QK4_1 32 +#define QR4_1 2 +#define QI4_1 (QK4_1 / (4 * QR4_1)) +typedef struct { + half2 dm; // dm.x = delta, dm.y = min + uint8_t qs[QK4_1 / 2]; // nibbles / quants +} block_q4_1; + +#define QK5_0 32 +#define QR5_0 2 +#define QI5_0 (QK5_0 / (4 * QR5_0)) +typedef struct { + half d; // delta + uint8_t qh[4]; // 5-th bit of quants + uint8_t qs[QK5_0 / 2]; // nibbles / quants +} block_q5_0; + +#define QK5_1 32 +#define QR5_1 2 +#define QI5_1 (QK5_1 / (4 * QR5_1)) +typedef struct { + half2 dm; // dm.x = delta, dm.y = min + uint8_t qh[4]; // 5-th bit of quants + uint8_t qs[QK5_1 / 2]; // nibbles / quants +} block_q5_1; + +#define QK8_0 32 +#define QR8_0 1 +#define QI8_0 (QK8_0 / (4 * QR8_0)) +typedef struct { + half d; // delta + int8_t qs[QK8_0]; // quants +} block_q8_0; + +#define QK8_1 32 +#define QR8_1 1 +#define QI8_1 (QK8_1 / (4 * QR8_1)) +typedef struct { + half2 ds; // ds.x = delta, ds.y = sum + int8_t qs[QK8_0]; // quants +} block_q8_1; + +#define QR2_K 4 +#define QI2_K (QK_K / (4 * QR2_K)) +typedef struct { + uint8_t scales[QK_K / 16]; // scales and mins, quantized with 4 bits + uint8_t qs[QK_K / 4]; // quants + half2 dm; // super-block scale for quantized scales/mins +} block_q2_K; + +#define QR3_K 4 +#define QI3_K (QK_K / (4 * QR3_K)) +typedef struct { + uint8_t hmask[QK_K / 8]; // quants - high bit + uint8_t qs[QK_K / 4]; // quants - low 2 bits + uint8_t scales[K_SCALE_SIZE]; // scales, quantized with 6 bits + half d; // super-block scale +} block_q3_K; + +#define QR4_K 2 +#define QI4_K (QK_K / (4 * QR4_K)) +typedef struct { + half2 dm; // super-block scale for quantized scales/mins + uint8_t scales[3 * QK_K / 64]; // scales, quantized with 6 bits + uint8_t qs[QK_K / 2]; // 4--bit quants +} block_q4_K; + +#define QR5_K 2 +#define QI5_K (QK_K / (4 * QR5_K)) +typedef struct { + half2 dm; // super-block scale for quantized scales/mins + uint8_t scales[K_SCALE_SIZE]; // scales and mins, quantized with 6 bits + uint8_t qh[QK_K / 8]; // quants, high bit + uint8_t qs[QK_K / 2]; // quants, low 4 bits +} block_q5_K; + +#define QR6_K 2 +#define QI6_K (QK_K / (4 * QR6_K)) +typedef struct { + uint8_t ql[QK_K / 2]; // quants, lower 4 bits + uint8_t qh[QK_K / 4]; // quants, upper 2 bits + int8_t scales[QK_K / 16]; // scales + half d; // delta +} block_q6_K; + +#define QR2_XXS 8 +#define QI2_XXS (QK_K / (4 * QR2_XXS)) +typedef struct { + half d; + uint16_t qs[QK_K / 8]; +} block_iq2_xxs; + +#define QR2_XS 8 +#define QI2_XS (QK_K / (4 * QR2_XS)) +typedef struct { + half d; + uint16_t qs[QK_K / 8]; + uint8_t scales[QK_K / 32]; +} block_iq2_xs; + +#define QR2_S 8 +#define QI2_S (QK_K / (4 * QR2_S)) +typedef struct { + half d; + uint8_t qs[QK_K / 4]; + uint8_t qh[QK_K / 32]; + uint8_t scales[QK_K / 32]; +} block_iq2_s; + +#define QR3_XXS 8 +#define QI3_XXS (QK_K / (4 * QR3_XXS)) +typedef struct { + half d; + uint8_t qs[3 * (QK_K / 8)]; +} block_iq3_xxs; + +#define QR3_XS 8 +#define QI3_XS (QK_K / (4 * QR3_XS)) +#define IQ3S_N_SCALE QK_K / 64 +typedef struct { + half d; + uint8_t qs[QK_K / 4]; + uint8_t qh[QK_K / 32]; + uint8_t signs[QK_K / 8]; + uint8_t scales[IQ3S_N_SCALE]; +} block_iq3_s; + +// 1.5625 bpw +#define QR1_S 8 +#define QI1_S (QK_K / (4 * QR1_S)) +typedef struct { + half d; + uint8_t qs[QK_K / 8]; + uint16_t qh[QK_K / 32]; +} block_iq1_s; + +// 1.75 bpw +#define QR1_M 8 +#define QI1_M (QK_K / (4 * QR1_M)) +typedef struct { + uint8_t qs[QK_K / 8]; // grid index, low 8 bits + uint8_t qh[QK_K / 16]; // grid index, high 3 bits + grid shift bit (for two groups of 8) + uint8_t scales[QK_K / 32]; // 3-bit block scales (4-bit if QK_K == 64) +} block_iq1_m; + +// Used by IQ1_M quants +typedef union { + half f16; + uint16_t u16; +} iq1m_scale_t; + +#define QK4_NL 32 +#define QR4_NL 2 +#define QI4_NL (QK4_NL / (4 * QR4_NL)) +typedef struct { + half d; + uint8_t qs[QK4_NL / 2]; +} block_iq4_nl; + +#define QR4_XS 8 +#define QI4_XS (QK_K / (4 * QR4_XS)) +typedef struct { + half d; + uint16_t scales_h; + uint8_t scales_l[QK_K / 64]; + uint8_t qs[QK_K / 2]; +} block_iq4_xs; + +static const __device__ uint64_t iq2xxs_grid[256] = { + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, + 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, 0x08080808082b082b, 0x08080808082b2b08, + 0x08080808082b2b2b, 0x0808080819080819, 0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, + 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b, + 0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808, 0x0808081908191919, + 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08, 0x0808082b08080808, 0x0808082b0808082b, + 0x0808082b082b082b, 0x0808082b2b08082b, 0x0808190808080819, 0x0808190808081908, 0x0808190808190808, + 0x08081908082b0819, 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08, + 0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, 0x080819082b2b1908, + 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08, 0x08081919082b0808, 0x080819191908192b, + 0x08081919192b2b19, 0x080819192b080808, 0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, + 0x0808192b19080808, 0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919, + 0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819, 0x08082b0819081908, + 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08, 0x08082b1908081908, 0x08082b1919080808, + 0x08082b2b0808082b, 0x08082b2b08191908, 0x0819080808080819, 0x0819080808081908, 0x0819080808190808, + 0x08190808082b0819, 0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808, + 0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808, 0x0819081919190808, + 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908, 0x0819082b19081919, 0x0819190808080808, + 0x0819190808082b08, 0x08191908082b0808, 0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, + 0x0819191908192b08, 0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819, + 0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819, 0x08192b1908080808, + 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819, 0x082b080808080808, 0x082b08080808082b, + 0x082b080808082b2b, 0x082b080819081908, 0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, + 0x082b0819082b2b19, 0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819, + 0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b, 0x082b191908080808, + 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808, 0x082b2b0808082b08, 0x082b2b08082b0808, + 0x082b2b082b191908, 0x082b2b2b19081908, 0x1908080808080819, 0x1908080808081908, 0x1908080808190808, + 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08, + 0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908, 0x190808082b190808, + 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819, 0x190808192b080808, 0x190808192b081919, + 0x1908082b08080819, 0x1908082b08190808, 0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, + 0x1908190808080808, 0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19, + 0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819, 0x19082b0808081908, + 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b1908080808, 0x19082b1919192b08, + 0x19082b19192b0819, 0x19082b192b08082b, 0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, + 0x1919080808082b08, 0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808, + 0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908, 0x1919082b2b190819, + 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b, 0x1919192b08080819, 0x1919192b19191908, + 0x19192b0808080808, 0x19192b0808190819, 0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, + 0x19192b2b08082b08, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08, + 0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808, 0x192b190808080808, + 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b, 0x192b19192b081908, 0x192b2b081908082b, + 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, + 0x2b08081908081908, 0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819, + 0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808, + 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908, 0x2b08192b08082b19, 0x2b08192b19080808, + 0x2b08192b192b0808, 0x2b082b080808082b, 0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, + 0x2b19080808190808, 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b, + 0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b, 0x2b19190819081908, + 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808, 0x2b2b08080808082b, 0x2b2b080819190808, + 0x2b2b08082b081919, 0x2b2b081908082b19, 0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, + 0x2b2b2b1908081908, +}; + +static const __device__ uint64_t iq2xs_grid[512] = { + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, + 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, 0x0808080808192b19, 0x08080808082b0808, + 0x08080808082b082b, 0x08080808082b1919, 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, + 0x080808081908192b, 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, + 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, + 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819, 0x080808082b191908, 0x080808082b192b19, + 0x080808082b2b0808, 0x0808081908080819, 0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, + 0x0808081908190808, 0x080808190819082b, 0x0808081908191919, 0x0808081908192b08, 0x0808081908192b2b, + 0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, 0x080808191908082b, 0x0808081919081919, + 0x0808081919082b08, 0x0808081919190819, 0x0808081919191908, 0x08080819192b0808, 0x08080819192b2b08, + 0x080808192b080819, 0x080808192b081908, 0x080808192b190808, 0x0808082b08080808, 0x0808082b0808082b, + 0x0808082b08081919, 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808, + 0x0808082b19080819, 0x0808082b19081908, 0x0808082b19190808, 0x0808082b19191919, 0x0808082b2b080808, + 0x0808082b2b082b2b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b, 0x0808190808082b19, + 0x0808190808190808, 0x080819080819082b, 0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819, + 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819081919, 0x0808190819082b08, + 0x0808190819190819, 0x0808190819191908, 0x080819081919192b, 0x08081908192b0808, 0x080819082b080819, + 0x080819082b081908, 0x080819082b190808, 0x0808191908080808, 0x080819190808082b, 0x0808191908081919, + 0x0808191908082b08, 0x0808191908190819, 0x0808191908191908, 0x08081919082b0808, 0x0808191919080819, + 0x0808191919081908, 0x0808191919190808, 0x08081919192b0819, 0x080819192b080808, 0x0808192b08080819, + 0x0808192b08081908, 0x0808192b08190808, 0x0808192b082b192b, 0x0808192b19080808, 0x0808192b1908082b, + 0x0808192b2b081908, 0x08082b0808080808, 0x08082b080808082b, 0x08082b0808081919, 0x08082b0808082b08, + 0x08082b0808082b2b, 0x08082b0808190819, 0x08082b0808191908, 0x08082b08082b0808, 0x08082b08082b1919, + 0x08082b0819080819, 0x08082b0819081908, 0x08082b0819190808, 0x08082b0819192b08, 0x08082b082b080808, + 0x08082b082b2b0808, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908, 0x08082b1908190808, + 0x08082b1919080808, 0x08082b192b080819, 0x08082b192b082b19, 0x08082b2b08080808, 0x08082b2b082b0808, + 0x08082b2b082b2b08, 0x08082b2b2b19192b, 0x08082b2b2b2b0808, 0x0819080808080819, 0x0819080808081908, + 0x081908080808192b, 0x0819080808082b19, 0x0819080808190808, 0x081908080819082b, 0x0819080808191919, + 0x0819080808192b08, 0x08190808082b0819, 0x08190808082b1908, 0x0819080819080808, 0x081908081908082b, + 0x0819080819081919, 0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x08190808192b0808, + 0x08190808192b2b2b, 0x081908082b080819, 0x081908082b081908, 0x081908082b190808, 0x0819081908080808, + 0x081908190808082b, 0x0819081908081919, 0x0819081908082b08, 0x0819081908190819, 0x0819081908191908, + 0x08190819082b0808, 0x0819081919080819, 0x0819081919081908, 0x0819081919190808, 0x081908192b080808, + 0x081908192b191908, 0x081908192b19192b, 0x0819082b08080819, 0x0819082b08081908, 0x0819082b0808192b, + 0x0819082b08190808, 0x0819082b19080808, 0x0819082b192b0808, 0x0819190808080808, 0x081919080808082b, + 0x0819190808081919, 0x0819190808082b08, 0x0819190808190819, 0x0819190808191908, 0x08191908082b0808, + 0x0819190819080819, 0x0819190819081908, 0x0819190819082b19, 0x0819190819190808, 0x08191908192b1908, + 0x081919082b080808, 0x0819191908080819, 0x0819191908081908, 0x0819191908190808, 0x0819191919080808, + 0x0819192b08080808, 0x0819192b08191908, 0x0819192b19082b19, 0x08192b0808080819, 0x08192b0808081908, + 0x08192b0808190808, 0x08192b080819082b, 0x08192b0819080808, 0x08192b0819191908, 0x08192b082b08192b, + 0x08192b1908080808, 0x08192b1908081919, 0x08192b19192b192b, 0x08192b2b19190819, 0x08192b2b2b2b2b19, + 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, 0x082b080808082b08, 0x082b080808082b2b, + 0x082b080808190819, 0x082b080808191908, 0x082b0808082b0808, 0x082b080819080819, 0x082b080819081908, + 0x082b080819190808, 0x082b08082b080808, 0x082b08082b2b0808, 0x082b081908080819, 0x082b081908081908, + 0x082b081908190808, 0x082b081919080808, 0x082b081919082b08, 0x082b0819192b1919, 0x082b082b08080808, + 0x082b082b082b082b, 0x082b082b2b080808, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908, + 0x082b190808190808, 0x082b1908082b2b19, 0x082b190819080808, 0x082b191908080808, 0x082b191919080819, + 0x082b19191919082b, 0x082b19192b192b19, 0x082b192b08080819, 0x082b192b08192b2b, 0x082b192b2b2b192b, + 0x082b2b0808080808, 0x082b2b0808082b08, 0x082b2b0808082b2b, 0x082b2b08082b0808, 0x082b2b0819191919, + 0x082b2b082b082b08, 0x082b2b082b2b082b, 0x082b2b19192b2b08, 0x082b2b192b190808, 0x082b2b2b08082b08, + 0x082b2b2b082b0808, 0x082b2b2b2b08082b, 0x082b2b2b2b082b08, 0x082b2b2b2b082b2b, 0x1908080808080819, + 0x1908080808081908, 0x190808080808192b, 0x1908080808082b19, 0x1908080808190808, 0x190808080819082b, + 0x1908080808191919, 0x1908080808192b08, 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, + 0x190808081908082b, 0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b, 0x1908080819190819, + 0x1908080819191908, 0x19080808192b0808, 0x19080808192b1919, 0x190808082b080819, 0x190808082b081908, + 0x190808082b190808, 0x1908081908080808, 0x190808190808082b, 0x1908081908081919, 0x1908081908082b08, + 0x1908081908190819, 0x1908081908191908, 0x19080819082b0808, 0x1908081919080819, 0x1908081919081908, + 0x1908081919190808, 0x190808192b080808, 0x190808192b081919, 0x190808192b2b082b, 0x1908082b08080819, + 0x1908082b08081908, 0x1908082b08190808, 0x1908082b0819082b, 0x1908082b082b2b19, 0x1908082b19080808, + 0x1908190808080808, 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808190819, + 0x1908190808191908, 0x1908190808192b19, 0x19081908082b0808, 0x1908190819080819, 0x1908190819081908, + 0x1908190819190808, 0x190819082b080808, 0x190819082b191908, 0x1908191908080819, 0x1908191908081908, + 0x1908191908190808, 0x19081919082b1908, 0x1908191919080808, 0x190819192b192b2b, 0x1908192b08080808, + 0x1908192b08082b2b, 0x1908192b19081908, 0x1908192b19190808, 0x19082b0808080819, 0x19082b0808081908, + 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, 0x19082b0819191908, 0x19082b08192b082b, + 0x19082b1908080808, 0x19082b1908190819, 0x19082b1919081908, 0x19082b1919190808, 0x19082b19192b2b19, + 0x19082b2b08081908, 0x1919080808080808, 0x191908080808082b, 0x1919080808081919, 0x1919080808082b08, + 0x1919080808190819, 0x1919080808191908, 0x19190808082b0808, 0x19190808082b2b08, 0x1919080819080819, + 0x1919080819081908, 0x1919080819190808, 0x191908082b080808, 0x1919081908080819, 0x1919081908081908, + 0x1919081908190808, 0x1919081908191919, 0x1919081919080808, 0x191908191908082b, 0x1919082b08080808, + 0x1919082b19081908, 0x1919082b2b2b2b2b, 0x1919190808080819, 0x1919190808081908, 0x1919190808190808, + 0x19191908082b0819, 0x1919190819080808, 0x19191908192b0808, 0x191919082b080819, 0x191919082b2b0819, + 0x1919191908080808, 0x1919191908082b08, 0x191919192b080808, 0x191919192b082b08, 0x1919192b082b0819, + 0x1919192b192b2b08, 0x1919192b2b2b0819, 0x19192b0808080808, 0x19192b0808191908, 0x19192b0819080819, + 0x19192b0819190808, 0x19192b082b192b19, 0x19192b1908192b2b, 0x19192b1919080808, 0x19192b191908082b, + 0x19192b2b2b081919, 0x192b080808080819, 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, + 0x192b080819191908, 0x192b0808192b082b, 0x192b08082b08192b, 0x192b08082b2b2b19, 0x192b081908080808, + 0x192b082b082b1908, 0x192b082b19082b2b, 0x192b082b2b19082b, 0x192b190808080808, 0x192b19080819192b, + 0x192b191908190808, 0x192b191919080808, 0x192b191919081919, 0x192b19192b2b1908, 0x192b2b0808080819, + 0x192b2b08192b2b2b, 0x192b2b19082b1919, 0x192b2b2b0808192b, 0x192b2b2b19191908, 0x192b2b2b192b082b, + 0x2b08080808080808, 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819, + 0x2b08080808191908, 0x2b080808082b0808, 0x2b080808082b2b2b, 0x2b08080819080819, 0x2b08080819081908, + 0x2b08080819190808, 0x2b0808082b080808, 0x2b0808082b08082b, 0x2b0808082b2b2b08, 0x2b0808082b2b2b2b, + 0x2b08081908080819, 0x2b08081908081908, 0x2b0808190808192b, 0x2b08081908190808, 0x2b08081919080808, + 0x2b08081919190819, 0x2b08081919192b19, 0x2b08082b08080808, 0x2b08082b082b0808, 0x2b08082b2b080808, + 0x2b08082b2b08082b, 0x2b08082b2b2b0808, 0x2b08082b2b2b2b08, 0x2b08190808080819, 0x2b08190808081908, + 0x2b08190808190808, 0x2b0819080819082b, 0x2b08190808191919, 0x2b08190819080808, 0x2b081908192b0808, + 0x2b0819082b082b19, 0x2b08191908080808, 0x2b08191919081908, 0x2b0819192b2b1919, 0x2b08192b08192b08, + 0x2b08192b192b2b2b, 0x2b082b0808080808, 0x2b082b0808082b08, 0x2b082b08082b1919, 0x2b082b0819192b2b, + 0x2b082b082b080808, 0x2b082b082b08082b, 0x2b082b082b2b2b08, 0x2b082b190808192b, 0x2b082b2b082b082b, + 0x2b082b2b2b080808, 0x2b082b2b2b082b08, 0x2b082b2b2b19192b, 0x2b082b2b2b2b2b08, 0x2b19080808080819, + 0x2b19080808081908, 0x2b19080808190808, 0x2b19080819080808, 0x2b1908081919192b, 0x2b1908082b081908, + 0x2b19081908080808, 0x2b190819082b082b, 0x2b190819192b1908, 0x2b19082b1919192b, 0x2b19082b2b082b19, + 0x2b19190808080808, 0x2b19190808081919, 0x2b19190819081908, 0x2b19190819190808, 0x2b19190819192b08, + 0x2b191919082b2b19, 0x2b1919192b190808, 0x2b1919192b19082b, 0x2b19192b19080819, 0x2b192b0819190819, + 0x2b192b082b2b192b, 0x2b192b1919082b19, 0x2b192b2b08191919, 0x2b192b2b192b0808, 0x2b2b080808080808, + 0x2b2b08080808082b, 0x2b2b080808082b08, 0x2b2b080808082b2b, 0x2b2b0808082b0808, 0x2b2b0808082b2b2b, + 0x2b2b08082b2b0808, 0x2b2b081919190819, 0x2b2b081919192b19, 0x2b2b08192b2b192b, 0x2b2b082b08080808, + 0x2b2b082b0808082b, 0x2b2b082b08082b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b080808, 0x2b2b082b2b2b0808, + 0x2b2b190819080808, 0x2b2b19082b191919, 0x2b2b192b192b1919, 0x2b2b192b2b192b08, 0x2b2b2b0808082b2b, + 0x2b2b2b08082b0808, 0x2b2b2b08082b082b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b0808, 0x2b2b2b082b2b2b08, + 0x2b2b2b1908081908, 0x2b2b2b192b081908, 0x2b2b2b192b08192b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b082b2b2b, + 0x2b2b2b2b2b190819, 0x2b2b2b2b2b2b2b2b, +}; + +static const __device__ uint64_t iq2s_grid[1024] = { + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, 0x0808080808082b2b, + 0x0808080808190819, 0x0808080808191908, 0x080808080819192b, 0x0808080808192b19, 0x08080808082b0808, + 0x08080808082b082b, 0x08080808082b1919, 0x08080808082b2b08, 0x0808080819080819, 0x0808080819081908, + 0x080808081908192b, 0x0808080819082b19, 0x0808080819190808, 0x080808081919082b, 0x0808080819191919, + 0x0808080819192b08, 0x08080808192b0819, 0x08080808192b1908, 0x08080808192b192b, 0x08080808192b2b19, + 0x080808082b080808, 0x080808082b08082b, 0x080808082b081919, 0x080808082b082b08, 0x080808082b190819, + 0x080808082b191908, 0x080808082b2b0808, 0x080808082b2b1919, 0x080808082b2b2b2b, 0x0808081908080819, + 0x0808081908081908, 0x080808190808192b, 0x0808081908082b19, 0x0808081908190808, 0x080808190819082b, + 0x0808081908191919, 0x0808081908192b08, 0x08080819082b0819, 0x08080819082b1908, 0x0808081919080808, + 0x080808191908082b, 0x0808081919081919, 0x0808081919082b08, 0x0808081919190819, 0x0808081919191908, + 0x080808191919192b, 0x0808081919192b19, 0x08080819192b0808, 0x08080819192b1919, 0x08080819192b2b08, + 0x080808192b080819, 0x080808192b081908, 0x080808192b190808, 0x080808192b19082b, 0x080808192b191919, + 0x080808192b2b0819, 0x080808192b2b1908, 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b08081919, + 0x0808082b08082b08, 0x0808082b08190819, 0x0808082b08191908, 0x0808082b082b0808, 0x0808082b082b2b2b, + 0x0808082b19080819, 0x0808082b19081908, 0x0808082b1908192b, 0x0808082b19082b19, 0x0808082b19190808, + 0x0808082b19191919, 0x0808082b2b080808, 0x0808082b2b081919, 0x0808082b2b082b2b, 0x0808082b2b191908, + 0x0808082b2b2b082b, 0x0808190808080819, 0x0808190808081908, 0x080819080808192b, 0x0808190808082b19, + 0x0808190808190808, 0x080819080819082b, 0x0808190808191919, 0x0808190808192b08, 0x08081908082b0819, + 0x08081908082b1908, 0x08081908082b192b, 0x08081908082b2b19, 0x0808190819080808, 0x080819081908082b, + 0x0808190819081919, 0x0808190819082b08, 0x0808190819082b2b, 0x0808190819190819, 0x0808190819191908, + 0x080819081919192b, 0x0808190819192b19, 0x08081908192b0808, 0x08081908192b082b, 0x08081908192b1919, + 0x080819082b080819, 0x080819082b081908, 0x080819082b08192b, 0x080819082b082b19, 0x080819082b190808, + 0x080819082b191919, 0x080819082b192b08, 0x080819082b2b0819, 0x080819082b2b1908, 0x0808191908080808, + 0x080819190808082b, 0x0808191908081919, 0x0808191908082b08, 0x0808191908082b2b, 0x0808191908190819, + 0x0808191908191908, 0x080819190819192b, 0x0808191908192b19, 0x08081919082b0808, 0x08081919082b1919, + 0x08081919082b2b08, 0x0808191919080819, 0x0808191919081908, 0x080819191908192b, 0x0808191919082b19, + 0x0808191919190808, 0x080819191919082b, 0x0808191919191919, 0x0808191919192b08, 0x08081919192b0819, + 0x08081919192b1908, 0x080819192b080808, 0x080819192b08082b, 0x080819192b081919, 0x080819192b082b08, + 0x080819192b190819, 0x080819192b191908, 0x080819192b2b0808, 0x0808192b08080819, 0x0808192b08081908, + 0x0808192b0808192b, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b08191919, 0x0808192b19080808, + 0x0808192b19081919, 0x0808192b19082b08, 0x0808192b19190819, 0x0808192b19191908, 0x0808192b192b0808, + 0x0808192b2b080819, 0x0808192b2b081908, 0x0808192b2b190808, 0x08082b0808080808, 0x08082b080808082b, + 0x08082b0808081919, 0x08082b0808082b08, 0x08082b0808190819, 0x08082b0808191908, 0x08082b080819192b, + 0x08082b0808192b19, 0x08082b08082b0808, 0x08082b08082b1919, 0x08082b08082b2b2b, 0x08082b0819080819, + 0x08082b0819081908, 0x08082b081908192b, 0x08082b0819082b19, 0x08082b0819190808, 0x08082b081919082b, + 0x08082b0819191919, 0x08082b0819192b08, 0x08082b08192b0819, 0x08082b08192b1908, 0x08082b082b080808, + 0x08082b082b081919, 0x08082b082b191908, 0x08082b082b2b2b2b, 0x08082b1908080819, 0x08082b1908081908, + 0x08082b1908190808, 0x08082b190819082b, 0x08082b1908191919, 0x08082b1908192b08, 0x08082b19082b0819, + 0x08082b1919080808, 0x08082b1919081919, 0x08082b1919082b08, 0x08082b1919190819, 0x08082b1919191908, + 0x08082b19192b0808, 0x08082b192b080819, 0x08082b192b190808, 0x08082b2b08080808, 0x08082b2b08190819, + 0x08082b2b08191908, 0x08082b2b082b082b, 0x08082b2b082b2b08, 0x08082b2b082b2b2b, 0x08082b2b19190808, + 0x08082b2b2b192b19, 0x0819080808080819, 0x0819080808081908, 0x081908080808192b, 0x0819080808082b19, + 0x0819080808190808, 0x081908080819082b, 0x0819080808191919, 0x0819080808192b08, 0x08190808082b0819, + 0x08190808082b1908, 0x08190808082b192b, 0x0819080819080808, 0x081908081908082b, 0x0819080819081919, + 0x0819080819082b08, 0x0819080819190819, 0x0819080819191908, 0x081908081919192b, 0x0819080819192b19, + 0x08190808192b0808, 0x08190808192b082b, 0x08190808192b1919, 0x08190808192b2b08, 0x081908082b080819, + 0x081908082b081908, 0x081908082b08192b, 0x081908082b190808, 0x081908082b191919, 0x081908082b192b08, + 0x081908082b2b0819, 0x081908082b2b1908, 0x0819081908080808, 0x081908190808082b, 0x0819081908081919, + 0x0819081908082b08, 0x0819081908082b2b, 0x0819081908190819, 0x0819081908191908, 0x081908190819192b, + 0x0819081908192b19, 0x08190819082b0808, 0x08190819082b082b, 0x08190819082b1919, 0x08190819082b2b08, + 0x0819081919080819, 0x0819081919081908, 0x081908191908192b, 0x0819081919082b19, 0x0819081919190808, + 0x081908191919082b, 0x0819081919191919, 0x0819081919192b08, 0x08190819192b0819, 0x08190819192b1908, + 0x081908192b080808, 0x081908192b08082b, 0x081908192b081919, 0x081908192b082b08, 0x081908192b190819, + 0x081908192b191908, 0x0819082b08080819, 0x0819082b08081908, 0x0819082b08082b19, 0x0819082b08190808, + 0x0819082b08191919, 0x0819082b082b0819, 0x0819082b082b1908, 0x0819082b19080808, 0x0819082b19081919, + 0x0819082b19190819, 0x0819082b19191908, 0x0819082b2b080819, 0x0819082b2b081908, 0x0819082b2b190808, + 0x0819190808080808, 0x081919080808082b, 0x0819190808081919, 0x0819190808082b08, 0x0819190808190819, + 0x0819190808191908, 0x081919080819192b, 0x0819190808192b19, 0x08191908082b0808, 0x08191908082b1919, + 0x08191908082b2b08, 0x0819190819080819, 0x0819190819081908, 0x081919081908192b, 0x0819190819082b19, + 0x0819190819190808, 0x081919081919082b, 0x0819190819191919, 0x0819190819192b08, 0x08191908192b0819, + 0x08191908192b1908, 0x081919082b080808, 0x081919082b08082b, 0x081919082b081919, 0x081919082b082b08, + 0x081919082b190819, 0x081919082b191908, 0x081919082b2b0808, 0x0819191908080819, 0x0819191908081908, + 0x081919190808192b, 0x0819191908082b19, 0x0819191908190808, 0x081919190819082b, 0x0819191908191919, + 0x0819191908192b08, 0x08191919082b0819, 0x08191919082b1908, 0x0819191919080808, 0x081919191908082b, + 0x0819191919081919, 0x0819191919082b08, 0x0819191919190819, 0x0819191919191908, 0x08191919192b0808, + 0x081919192b080819, 0x081919192b081908, 0x081919192b190808, 0x0819192b08080808, 0x0819192b08081919, + 0x0819192b08082b08, 0x0819192b08190819, 0x0819192b08191908, 0x0819192b082b0808, 0x0819192b19080819, + 0x0819192b19081908, 0x0819192b19190808, 0x0819192b2b080808, 0x0819192b2b2b2b2b, 0x08192b0808080819, + 0x08192b0808081908, 0x08192b080808192b, 0x08192b0808082b19, 0x08192b0808190808, 0x08192b0808191919, + 0x08192b0808192b08, 0x08192b08082b0819, 0x08192b0819080808, 0x08192b081908082b, 0x08192b0819081919, + 0x08192b0819082b08, 0x08192b0819190819, 0x08192b0819191908, 0x08192b08192b0808, 0x08192b082b080819, + 0x08192b082b081908, 0x08192b1908080808, 0x08192b190808082b, 0x08192b1908081919, 0x08192b1908082b08, + 0x08192b1908190819, 0x08192b1908191908, 0x08192b19082b0808, 0x08192b1919080819, 0x08192b1919081908, + 0x08192b1919190808, 0x08192b19192b2b19, 0x08192b192b2b082b, 0x08192b2b08081908, 0x08192b2b08190808, + 0x08192b2b19080808, 0x08192b2b1919192b, 0x082b080808080808, 0x082b08080808082b, 0x082b080808081919, + 0x082b080808082b08, 0x082b080808190819, 0x082b080808191908, 0x082b08080819192b, 0x082b080808192b19, + 0x082b0808082b0808, 0x082b0808082b1919, 0x082b0808082b2b2b, 0x082b080819080819, 0x082b080819081908, + 0x082b080819190808, 0x082b08081919082b, 0x082b080819191919, 0x082b0808192b1908, 0x082b08082b080808, + 0x082b08082b082b2b, 0x082b08082b191908, 0x082b08082b2b2b2b, 0x082b081908080819, 0x082b081908081908, + 0x082b081908190808, 0x082b08190819082b, 0x082b081908191919, 0x082b0819082b0819, 0x082b081919080808, + 0x082b08191908082b, 0x082b081919081919, 0x082b081919190819, 0x082b081919191908, 0x082b0819192b0808, + 0x082b08192b080819, 0x082b08192b081908, 0x082b08192b190808, 0x082b082b08080808, 0x082b082b08082b2b, + 0x082b082b082b082b, 0x082b082b082b2b08, 0x082b082b082b2b2b, 0x082b082b19081908, 0x082b082b19190808, + 0x082b082b2b082b08, 0x082b082b2b082b2b, 0x082b082b2b2b2b08, 0x082b190808080819, 0x082b190808081908, + 0x082b19080808192b, 0x082b190808082b19, 0x082b190808190808, 0x082b190808191919, 0x082b190808192b08, + 0x082b1908082b0819, 0x082b1908082b1908, 0x082b190819080808, 0x082b19081908082b, 0x082b190819081919, + 0x082b190819082b08, 0x082b190819190819, 0x082b190819191908, 0x082b1908192b0808, 0x082b19082b080819, + 0x082b19082b081908, 0x082b19082b190808, 0x082b191908080808, 0x082b191908081919, 0x082b191908082b08, + 0x082b191908190819, 0x082b191908191908, 0x082b1919082b0808, 0x082b191919080819, 0x082b191919081908, + 0x082b191919190808, 0x082b1919192b192b, 0x082b19192b080808, 0x082b192b08080819, 0x082b192b08081908, + 0x082b192b08190808, 0x082b192b19080808, 0x082b192b19192b19, 0x082b2b0808080808, 0x082b2b0808081919, + 0x082b2b0808190819, 0x082b2b0808191908, 0x082b2b0819080819, 0x082b2b0819081908, 0x082b2b0819190808, + 0x082b2b082b082b2b, 0x082b2b082b2b2b2b, 0x082b2b1908080819, 0x082b2b1908081908, 0x082b2b1908190808, + 0x082b2b192b191919, 0x082b2b2b08082b2b, 0x082b2b2b082b082b, 0x082b2b2b192b1908, 0x082b2b2b2b082b08, + 0x082b2b2b2b082b2b, 0x1908080808080819, 0x1908080808081908, 0x190808080808192b, 0x1908080808082b19, + 0x1908080808190808, 0x190808080819082b, 0x1908080808191919, 0x1908080808192b08, 0x1908080808192b2b, + 0x19080808082b0819, 0x19080808082b1908, 0x19080808082b192b, 0x1908080819080808, 0x190808081908082b, + 0x1908080819081919, 0x1908080819082b08, 0x1908080819082b2b, 0x1908080819190819, 0x1908080819191908, + 0x190808081919192b, 0x1908080819192b19, 0x19080808192b0808, 0x19080808192b082b, 0x19080808192b1919, + 0x190808082b080819, 0x190808082b081908, 0x190808082b190808, 0x190808082b191919, 0x190808082b192b08, + 0x190808082b2b0819, 0x190808082b2b1908, 0x1908081908080808, 0x190808190808082b, 0x1908081908081919, + 0x1908081908082b08, 0x1908081908190819, 0x1908081908191908, 0x190808190819192b, 0x1908081908192b19, + 0x19080819082b0808, 0x19080819082b082b, 0x19080819082b1919, 0x1908081919080819, 0x1908081919081908, + 0x190808191908192b, 0x1908081919082b19, 0x1908081919190808, 0x190808191919082b, 0x1908081919191919, + 0x1908081919192b08, 0x19080819192b0819, 0x19080819192b1908, 0x190808192b080808, 0x190808192b08082b, + 0x190808192b081919, 0x190808192b082b08, 0x190808192b190819, 0x190808192b191908, 0x190808192b2b0808, + 0x1908082b08080819, 0x1908082b08081908, 0x1908082b08190808, 0x1908082b0819082b, 0x1908082b08191919, + 0x1908082b08192b08, 0x1908082b082b1908, 0x1908082b19080808, 0x1908082b19081919, 0x1908082b19082b08, + 0x1908082b19190819, 0x1908082b19191908, 0x1908082b192b0808, 0x1908082b2b080819, 0x1908082b2b081908, + 0x1908190808080808, 0x190819080808082b, 0x1908190808081919, 0x1908190808082b08, 0x1908190808082b2b, + 0x1908190808190819, 0x1908190808191908, 0x190819080819192b, 0x1908190808192b19, 0x19081908082b0808, + 0x19081908082b082b, 0x19081908082b1919, 0x19081908082b2b08, 0x1908190819080819, 0x1908190819081908, + 0x190819081908192b, 0x1908190819082b19, 0x1908190819190808, 0x190819081919082b, 0x1908190819191919, + 0x1908190819192b08, 0x19081908192b0819, 0x19081908192b1908, 0x190819082b080808, 0x190819082b08082b, + 0x190819082b081919, 0x190819082b082b08, 0x190819082b190819, 0x190819082b191908, 0x190819082b2b0808, + 0x1908191908080819, 0x1908191908081908, 0x190819190808192b, 0x1908191908082b19, 0x1908191908190808, + 0x190819190819082b, 0x1908191908191919, 0x1908191908192b08, 0x19081919082b0819, 0x19081919082b1908, + 0x1908191919080808, 0x190819191908082b, 0x1908191919081919, 0x1908191919082b08, 0x1908191919190819, + 0x1908191919191908, 0x19081919192b0808, 0x19081919192b2b2b, 0x190819192b080819, 0x190819192b081908, + 0x190819192b190808, 0x1908192b08080808, 0x1908192b0808082b, 0x1908192b08081919, 0x1908192b08082b08, + 0x1908192b08190819, 0x1908192b08191908, 0x1908192b082b0808, 0x1908192b19080819, 0x1908192b19081908, + 0x1908192b19190808, 0x1908192b2b080808, 0x1908192b2b2b1919, 0x19082b0808080819, 0x19082b0808081908, + 0x19082b0808082b19, 0x19082b0808190808, 0x19082b080819082b, 0x19082b0808191919, 0x19082b0808192b08, + 0x19082b08082b0819, 0x19082b08082b1908, 0x19082b0819080808, 0x19082b081908082b, 0x19082b0819081919, + 0x19082b0819082b08, 0x19082b0819190819, 0x19082b0819191908, 0x19082b08192b0808, 0x19082b082b081908, + 0x19082b082b190808, 0x19082b1908080808, 0x19082b190808082b, 0x19082b1908081919, 0x19082b1908082b08, + 0x19082b1908190819, 0x19082b1908191908, 0x19082b19082b0808, 0x19082b1919080819, 0x19082b1919081908, + 0x19082b1919190808, 0x19082b192b080808, 0x19082b192b19192b, 0x19082b2b08080819, 0x19082b2b08081908, + 0x19082b2b08190808, 0x19082b2b19080808, 0x1919080808080808, 0x191908080808082b, 0x1919080808081919, + 0x1919080808082b08, 0x1919080808190819, 0x1919080808191908, 0x191908080819192b, 0x1919080808192b19, + 0x19190808082b0808, 0x19190808082b082b, 0x19190808082b1919, 0x19190808082b2b08, 0x1919080819080819, + 0x1919080819081908, 0x191908081908192b, 0x1919080819082b19, 0x1919080819190808, 0x191908081919082b, + 0x1919080819191919, 0x1919080819192b08, 0x19190808192b0819, 0x19190808192b1908, 0x191908082b080808, + 0x191908082b08082b, 0x191908082b081919, 0x191908082b082b08, 0x191908082b190819, 0x191908082b191908, + 0x1919081908080819, 0x1919081908081908, 0x191908190808192b, 0x1919081908082b19, 0x1919081908190808, + 0x191908190819082b, 0x1919081908191919, 0x1919081908192b08, 0x19190819082b0819, 0x19190819082b1908, + 0x1919081919080808, 0x191908191908082b, 0x1919081919081919, 0x1919081919082b08, 0x1919081919190819, + 0x1919081919191908, 0x19190819192b0808, 0x191908192b080819, 0x191908192b081908, 0x191908192b190808, + 0x1919082b08080808, 0x1919082b08081919, 0x1919082b08082b08, 0x1919082b08190819, 0x1919082b08191908, + 0x1919082b082b0808, 0x1919082b19080819, 0x1919082b19081908, 0x1919082b19190808, 0x1919082b192b2b19, + 0x1919082b2b080808, 0x1919190808080819, 0x1919190808081908, 0x191919080808192b, 0x1919190808082b19, + 0x1919190808190808, 0x191919080819082b, 0x1919190808191919, 0x1919190808192b08, 0x19191908082b0819, + 0x19191908082b1908, 0x1919190819080808, 0x191919081908082b, 0x1919190819081919, 0x1919190819082b08, + 0x1919190819190819, 0x1919190819191908, 0x19191908192b0808, 0x191919082b080819, 0x191919082b081908, + 0x191919082b190808, 0x1919191908080808, 0x191919190808082b, 0x1919191908081919, 0x1919191908082b08, + 0x1919191908190819, 0x1919191908191908, 0x19191919082b0808, 0x1919191919080819, 0x1919191919081908, + 0x1919191919190808, 0x191919192b080808, 0x1919192b08080819, 0x1919192b08081908, 0x1919192b08190808, + 0x1919192b082b192b, 0x1919192b19080808, 0x19192b0808080808, 0x19192b080808082b, 0x19192b0808081919, + 0x19192b0808082b08, 0x19192b0808190819, 0x19192b0808191908, 0x19192b08082b0808, 0x19192b0819080819, + 0x19192b0819081908, 0x19192b0819190808, 0x19192b0819192b2b, 0x19192b082b080808, 0x19192b1908080819, + 0x19192b1908081908, 0x19192b1908190808, 0x19192b1919080808, 0x19192b2b08080808, 0x19192b2b08192b19, + 0x19192b2b2b081919, 0x19192b2b2b2b2b08, 0x192b080808080819, 0x192b080808081908, 0x192b08080808192b, + 0x192b080808190808, 0x192b08080819082b, 0x192b080808191919, 0x192b080808192b08, 0x192b0808082b0819, + 0x192b0808082b1908, 0x192b080819080808, 0x192b080819081919, 0x192b080819082b08, 0x192b080819190819, + 0x192b080819191908, 0x192b0808192b0808, 0x192b08082b081908, 0x192b08082b190808, 0x192b081908080808, + 0x192b08190808082b, 0x192b081908081919, 0x192b081908082b08, 0x192b081908190819, 0x192b081908191908, + 0x192b0819082b0808, 0x192b081919080819, 0x192b081919081908, 0x192b081919190808, 0x192b08192b080808, + 0x192b08192b192b19, 0x192b082b08081908, 0x192b082b08190808, 0x192b082b19080808, 0x192b082b1919192b, + 0x192b082b2b2b0819, 0x192b190808080808, 0x192b190808081919, 0x192b190808082b08, 0x192b190808190819, + 0x192b190808191908, 0x192b1908082b0808, 0x192b190819080819, 0x192b190819081908, 0x192b190819190808, + 0x192b19082b080808, 0x192b191908080819, 0x192b191908081908, 0x192b191908190808, 0x192b191919080808, + 0x192b191919082b2b, 0x192b1919192b2b08, 0x192b19192b19082b, 0x192b192b08080808, 0x192b192b2b191908, + 0x192b2b0808080819, 0x192b2b0808081908, 0x192b2b0808190808, 0x192b2b08192b1919, 0x192b2b082b192b08, + 0x192b2b1908080808, 0x192b2b19082b2b2b, 0x192b2b2b1908082b, 0x192b2b2b2b2b0819, 0x2b08080808080808, + 0x2b0808080808082b, 0x2b08080808081919, 0x2b08080808082b08, 0x2b08080808190819, 0x2b08080808191908, + 0x2b08080808192b19, 0x2b080808082b0808, 0x2b080808082b1919, 0x2b08080819080819, 0x2b08080819081908, + 0x2b08080819190808, 0x2b0808081919082b, 0x2b08080819191919, 0x2b08080819192b08, 0x2b080808192b0819, + 0x2b0808082b080808, 0x2b0808082b081919, 0x2b0808082b190819, 0x2b0808082b191908, 0x2b08081908080819, + 0x2b08081908081908, 0x2b08081908082b19, 0x2b08081908190808, 0x2b0808190819082b, 0x2b08081908191919, + 0x2b08081908192b08, 0x2b080819082b0819, 0x2b080819082b1908, 0x2b08081919080808, 0x2b0808191908082b, + 0x2b08081919081919, 0x2b08081919082b08, 0x2b08081919190819, 0x2b08081919191908, 0x2b0808192b080819, + 0x2b0808192b081908, 0x2b0808192b190808, 0x2b0808192b2b2b19, 0x2b08082b08080808, 0x2b08082b08081919, + 0x2b08082b08082b2b, 0x2b08082b08190819, 0x2b08082b08191908, 0x2b08082b19080819, 0x2b08082b19081908, + 0x2b08082b19190808, 0x2b08190808080819, 0x2b08190808081908, 0x2b0819080808192b, 0x2b08190808082b19, + 0x2b08190808190808, 0x2b0819080819082b, 0x2b08190808191919, 0x2b08190808192b08, 0x2b081908082b0819, + 0x2b08190819080808, 0x2b0819081908082b, 0x2b08190819081919, 0x2b08190819082b08, 0x2b08190819190819, + 0x2b08190819191908, 0x2b081908192b0808, 0x2b0819082b080819, 0x2b0819082b081908, 0x2b0819082b190808, + 0x2b08191908080808, 0x2b0819190808082b, 0x2b08191908081919, 0x2b08191908082b08, 0x2b08191908190819, + 0x2b08191908191908, 0x2b081919082b0808, 0x2b08191919080819, 0x2b08191919081908, 0x2b08191919190808, + 0x2b0819192b080808, 0x2b0819192b082b2b, 0x2b08192b08080819, 0x2b08192b08081908, 0x2b08192b08190808, + 0x2b08192b082b2b19, 0x2b08192b19080808, 0x2b082b0808080808, 0x2b082b0808081919, 0x2b082b0808190819, + 0x2b082b0808191908, 0x2b082b0819080819, 0x2b082b0819081908, 0x2b082b0819190808, 0x2b082b082b2b082b, + 0x2b082b1908080819, 0x2b082b1908081908, 0x2b082b1919080808, 0x2b082b19192b1919, 0x2b082b2b082b082b, + 0x2b082b2b19192b08, 0x2b082b2b19192b2b, 0x2b082b2b2b08082b, 0x2b082b2b2b2b082b, 0x2b19080808080819, + 0x2b19080808081908, 0x2b19080808082b19, 0x2b19080808190808, 0x2b1908080819082b, 0x2b19080808191919, + 0x2b19080808192b08, 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908081908082b, 0x2b19080819081919, + 0x2b19080819082b08, 0x2b19080819190819, 0x2b19080819191908, 0x2b190808192b0808, 0x2b1908082b080819, + 0x2b1908082b081908, 0x2b1908082b190808, 0x2b19081908080808, 0x2b19081908081919, 0x2b19081908190819, + 0x2b19081908191908, 0x2b19081919080819, 0x2b19081919081908, 0x2b19081919190808, 0x2b19081919192b2b, + 0x2b19082b08080819, 0x2b19082b08081908, 0x2b19082b08190808, 0x2b19082b19080808, 0x2b19082b2b2b192b, + 0x2b19190808080808, 0x2b1919080808082b, 0x2b19190808081919, 0x2b19190808082b08, 0x2b19190808190819, + 0x2b19190808191908, 0x2b191908082b0808, 0x2b19190819080819, 0x2b19190819081908, 0x2b19190819190808, + 0x2b1919082b080808, 0x2b1919082b19192b, 0x2b19191908080819, 0x2b19191908081908, 0x2b19191908190808, + 0x2b19191919080808, 0x2b1919192b192b08, 0x2b1919192b2b0819, 0x2b19192b08080808, 0x2b19192b1908192b, + 0x2b19192b192b1908, 0x2b192b0808080819, 0x2b192b0808081908, 0x2b192b0808190808, 0x2b192b08082b192b, + 0x2b192b0819080808, 0x2b192b082b2b2b19, 0x2b192b1908080808, 0x2b192b1919082b19, 0x2b192b191919082b, + 0x2b192b2b2b190808, 0x2b2b080808080808, 0x2b2b080808081919, 0x2b2b080808082b2b, 0x2b2b080808191908, + 0x2b2b0808082b082b, 0x2b2b0808082b2b2b, 0x2b2b080819080819, 0x2b2b080819081908, 0x2b2b080819190808, + 0x2b2b08082b2b082b, 0x2b2b08082b2b2b2b, 0x2b2b081919080808, 0x2b2b0819192b1919, 0x2b2b082b0808082b, + 0x2b2b082b08082b2b, 0x2b2b082b082b082b, 0x2b2b082b082b2b08, 0x2b2b082b082b2b2b, 0x2b2b082b2b08082b, + 0x2b2b082b2b082b08, 0x2b2b082b2b082b2b, 0x2b2b082b2b2b2b08, 0x2b2b190808080819, 0x2b2b190808081908, + 0x2b2b190808190808, 0x2b2b190819080808, 0x2b2b19082b082b19, 0x2b2b19082b2b1908, 0x2b2b191908080808, + 0x2b2b191908192b19, 0x2b2b192b19190819, 0x2b2b2b0808082b2b, 0x2b2b2b08082b2b08, 0x2b2b2b082b2b082b, + 0x2b2b2b1919191908, 0x2b2b2b192b08192b, 0x2b2b2b2b08082b08, 0x2b2b2b2b08082b2b, 0x2b2b2b2b082b0808, + 0x2b2b2b2b082b082b, 0x2b2b2b2b082b2b08, 0x2b2b2b2b2b082b08, 0x2b2b2b2b2b2b2b2b, +}; + +static const __device__ uint32_t iq3xxs_grid[256] = { + 0x04040404, 0x04040414, 0x04040424, 0x04040c0c, 0x04040c1c, 0x04040c3e, 0x04041404, 0x04041414, 0x04041c0c, + 0x04042414, 0x04043e1c, 0x04043e2c, 0x040c040c, 0x040c041c, 0x040c0c04, 0x040c0c14, 0x040c140c, 0x040c142c, + 0x040c1c04, 0x040c1c14, 0x040c240c, 0x040c2c24, 0x040c3e04, 0x04140404, 0x04140414, 0x04140424, 0x04140c0c, + 0x04141404, 0x04141414, 0x04141c0c, 0x04141c1c, 0x04141c3e, 0x04142c0c, 0x04142c3e, 0x04143e2c, 0x041c040c, + 0x041c043e, 0x041c0c04, 0x041c0c14, 0x041c142c, 0x041c3e04, 0x04240c1c, 0x04241c3e, 0x04242424, 0x04242c3e, + 0x04243e1c, 0x04243e2c, 0x042c040c, 0x042c043e, 0x042c1c14, 0x042c2c14, 0x04341c2c, 0x04343424, 0x043e0c04, + 0x043e0c24, 0x043e0c34, 0x043e241c, 0x043e340c, 0x0c04040c, 0x0c04041c, 0x0c040c04, 0x0c040c14, 0x0c04140c, + 0x0c04141c, 0x0c041c04, 0x0c041c14, 0x0c041c24, 0x0c04243e, 0x0c042c04, 0x0c0c0404, 0x0c0c0414, 0x0c0c0c0c, + 0x0c0c1404, 0x0c0c1414, 0x0c14040c, 0x0c14041c, 0x0c140c04, 0x0c140c14, 0x0c14140c, 0x0c141c04, 0x0c143e14, + 0x0c1c0404, 0x0c1c0414, 0x0c1c1404, 0x0c1c1c0c, 0x0c1c2434, 0x0c1c3434, 0x0c24040c, 0x0c24042c, 0x0c242c04, + 0x0c2c1404, 0x0c2c1424, 0x0c2c2434, 0x0c2c3e0c, 0x0c34042c, 0x0c3e1414, 0x0c3e2404, 0x14040404, 0x14040414, + 0x14040c0c, 0x14040c1c, 0x14041404, 0x14041414, 0x14041434, 0x14041c0c, 0x14042414, 0x140c040c, 0x140c041c, + 0x140c042c, 0x140c0c04, 0x140c0c14, 0x140c140c, 0x140c1c04, 0x140c341c, 0x140c343e, 0x140c3e04, 0x14140404, + 0x14140414, 0x14140c0c, 0x14140c3e, 0x14141404, 0x14141414, 0x14141c3e, 0x14142404, 0x14142c2c, 0x141c040c, + 0x141c0c04, 0x141c0c24, 0x141c3e04, 0x141c3e24, 0x14241c2c, 0x14242c1c, 0x142c041c, 0x142c143e, 0x142c240c, + 0x142c3e24, 0x143e040c, 0x143e041c, 0x143e0c34, 0x143e242c, 0x1c04040c, 0x1c040c04, 0x1c040c14, 0x1c04140c, + 0x1c04141c, 0x1c042c04, 0x1c04342c, 0x1c043e14, 0x1c0c0404, 0x1c0c0414, 0x1c0c1404, 0x1c0c1c0c, 0x1c0c2424, + 0x1c0c2434, 0x1c14040c, 0x1c14041c, 0x1c140c04, 0x1c14142c, 0x1c142c14, 0x1c143e14, 0x1c1c0c0c, 0x1c1c1c1c, + 0x1c241c04, 0x1c24243e, 0x1c243e14, 0x1c2c0404, 0x1c2c0434, 0x1c2c1414, 0x1c2c2c2c, 0x1c340c24, 0x1c341c34, + 0x1c34341c, 0x1c3e1c1c, 0x1c3e3404, 0x24040424, 0x24040c3e, 0x24041c2c, 0x24041c3e, 0x24042c1c, 0x24042c3e, + 0x240c3e24, 0x24141404, 0x24141c3e, 0x24142404, 0x24143404, 0x24143434, 0x241c043e, 0x241c242c, 0x24240424, + 0x24242c0c, 0x24243424, 0x242c142c, 0x242c241c, 0x242c3e04, 0x243e042c, 0x243e0c04, 0x243e0c14, 0x243e1c04, + 0x2c040c14, 0x2c04240c, 0x2c043e04, 0x2c0c0404, 0x2c0c0434, 0x2c0c1434, 0x2c0c2c2c, 0x2c140c24, 0x2c141c14, + 0x2c143e14, 0x2c1c0414, 0x2c1c2c1c, 0x2c240c04, 0x2c24141c, 0x2c24143e, 0x2c243e14, 0x2c2c0414, 0x2c2c1c0c, + 0x2c342c04, 0x2c3e1424, 0x2c3e2414, 0x34041424, 0x34042424, 0x34042434, 0x34043424, 0x340c140c, 0x340c340c, + 0x34140c3e, 0x34143424, 0x341c1c04, 0x341c1c34, 0x34242424, 0x342c042c, 0x342c2c14, 0x34341c1c, 0x343e041c, + 0x343e140c, 0x3e04041c, 0x3e04042c, 0x3e04043e, 0x3e040c04, 0x3e041c14, 0x3e042c14, 0x3e0c1434, 0x3e0c2404, + 0x3e140c14, 0x3e14242c, 0x3e142c14, 0x3e1c0404, 0x3e1c0c2c, 0x3e1c1c1c, 0x3e1c3404, 0x3e24140c, 0x3e24240c, + 0x3e2c0404, 0x3e2c0414, 0x3e2c1424, 0x3e341c04, +}; + +static const __device__ uint32_t iq3xs_grid[512] = { + 0x04040404, 0x0404040c, 0x04040414, 0x0404042c, 0x0404043e, 0x04040c04, 0x04040c0c, 0x04040c14, 0x04040c24, + 0x04040c34, 0x04041404, 0x0404140c, 0x0404142c, 0x04041c1c, 0x04042404, 0x04042414, 0x0404242c, 0x0404243e, + 0x04042c0c, 0x04042c1c, 0x04043404, 0x04043414, 0x04043e0c, 0x04043e24, 0x04043e3e, 0x040c0404, 0x040c040c, + 0x040c0414, 0x040c0424, 0x040c0c04, 0x040c0c0c, 0x040c0c2c, 0x040c1404, 0x040c141c, 0x040c143e, 0x040c1c0c, + 0x040c1c2c, 0x040c2424, 0x040c340c, 0x040c342c, 0x040c3e14, 0x04140404, 0x0414040c, 0x0414042c, 0x0414043e, + 0x04140c04, 0x04140c1c, 0x04140c34, 0x0414140c, 0x0414142c, 0x04141c04, 0x04141c24, 0x04142414, 0x0414242c, + 0x0414243e, 0x04142c0c, 0x04142c1c, 0x04143e04, 0x04143e1c, 0x041c041c, 0x041c0c0c, 0x041c0c2c, 0x041c1404, + 0x041c1414, 0x041c1c0c, 0x041c1c1c, 0x041c1c34, 0x041c2424, 0x041c2c04, 0x041c2c14, 0x041c343e, 0x041c3e0c, + 0x041c3e2c, 0x04240404, 0x04240c1c, 0x04240c3e, 0x0424140c, 0x04241424, 0x04241c14, 0x04242404, 0x0424241c, + 0x04242c0c, 0x04243e04, 0x042c0414, 0x042c0424, 0x042c1404, 0x042c1414, 0x042c1434, 0x042c1c1c, 0x042c240c, + 0x042c242c, 0x042c243e, 0x042c3434, 0x042c3e1c, 0x04340434, 0x04340c0c, 0x04340c1c, 0x04341c0c, 0x04342c14, + 0x04343e0c, 0x043e0404, 0x043e0414, 0x043e0424, 0x043e1404, 0x043e1414, 0x043e1434, 0x043e1c1c, 0x043e2c04, + 0x043e2c24, 0x0c040404, 0x0c04040c, 0x0c040414, 0x0c040424, 0x0c040c04, 0x0c040c0c, 0x0c040c1c, 0x0c040c2c, + 0x0c040c3e, 0x0c041404, 0x0c041414, 0x0c041c0c, 0x0c041c24, 0x0c041c34, 0x0c042c24, 0x0c042c34, 0x0c04340c, + 0x0c043e14, 0x0c0c0404, 0x0c0c040c, 0x0c0c041c, 0x0c0c0434, 0x0c0c0c04, 0x0c0c0c24, 0x0c0c140c, 0x0c0c1c04, + 0x0c0c1c1c, 0x0c0c240c, 0x0c0c2c04, 0x0c0c2c14, 0x0c0c3e04, 0x0c0c3e34, 0x0c140404, 0x0c140c14, 0x0c140c2c, + 0x0c140c3e, 0x0c141404, 0x0c141424, 0x0c141c14, 0x0c142404, 0x0c14241c, 0x0c142c2c, 0x0c143404, 0x0c143e14, + 0x0c1c040c, 0x0c1c0424, 0x0c1c043e, 0x0c1c0c04, 0x0c1c0c1c, 0x0c1c140c, 0x0c1c143e, 0x0c1c1c04, 0x0c1c1c24, + 0x0c1c240c, 0x0c1c3414, 0x0c1c3e04, 0x0c24041c, 0x0c24042c, 0x0c240c14, 0x0c240c24, 0x0c241c0c, 0x0c241c1c, + 0x0c242414, 0x0c242434, 0x0c242c04, 0x0c242c24, 0x0c2c040c, 0x0c2c0c04, 0x0c2c0c1c, 0x0c2c140c, 0x0c2c1c04, + 0x0c2c1c14, 0x0c2c2c0c, 0x0c341404, 0x0c341424, 0x0c34143e, 0x0c342424, 0x0c342434, 0x0c3e040c, 0x0c3e041c, + 0x0c3e0c04, 0x0c3e0c14, 0x0c3e140c, 0x0c3e1c2c, 0x0c3e240c, 0x0c3e3414, 0x0c3e3e04, 0x14040404, 0x1404040c, + 0x1404041c, 0x1404042c, 0x1404043e, 0x14040c04, 0x14040c14, 0x14040c24, 0x14040c34, 0x1404140c, 0x1404141c, + 0x1404143e, 0x14041c04, 0x14041c14, 0x1404240c, 0x1404241c, 0x1404242c, 0x14042c04, 0x14042c14, 0x1404343e, + 0x14043e04, 0x14043e1c, 0x14043e2c, 0x140c0404, 0x140c0414, 0x140c0c04, 0x140c0c1c, 0x140c0c3e, 0x140c1414, + 0x140c142c, 0x140c1c0c, 0x140c1c24, 0x140c2414, 0x140c2c0c, 0x1414040c, 0x14140424, 0x1414043e, 0x1414140c, + 0x1414141c, 0x14141c04, 0x14141c3e, 0x1414240c, 0x14142c1c, 0x14142c3e, 0x14143e0c, 0x14143e24, 0x141c0404, + 0x141c0414, 0x141c042c, 0x141c0c0c, 0x141c1414, 0x141c1424, 0x141c1c0c, 0x141c1c1c, 0x141c2414, 0x141c2c04, + 0x141c3434, 0x1424040c, 0x1424043e, 0x14241404, 0x1424141c, 0x14241c14, 0x14241c2c, 0x1424240c, 0x14243e14, + 0x14243e2c, 0x142c0424, 0x142c0c0c, 0x142c1414, 0x142c1c3e, 0x142c2404, 0x142c2c1c, 0x142c3e04, 0x14340404, + 0x14340414, 0x1434043e, 0x1434140c, 0x14342c2c, 0x1434340c, 0x143e042c, 0x143e0c0c, 0x143e1434, 0x143e1c04, + 0x143e241c, 0x143e2c04, 0x1c040414, 0x1c040c0c, 0x1c040c1c, 0x1c040c2c, 0x1c040c3e, 0x1c041414, 0x1c041c0c, + 0x1c041c1c, 0x1c041c2c, 0x1c042414, 0x1c042424, 0x1c04243e, 0x1c042c0c, 0x1c04341c, 0x1c043e0c, 0x1c0c040c, + 0x1c0c041c, 0x1c0c042c, 0x1c0c0c24, 0x1c0c140c, 0x1c0c141c, 0x1c0c2404, 0x1c0c3404, 0x1c0c3e14, 0x1c0c3e34, + 0x1c140404, 0x1c140c14, 0x1c141404, 0x1c141c14, 0x1c141c24, 0x1c142c04, 0x1c1c040c, 0x1c1c0c04, 0x1c1c0c24, + 0x1c1c140c, 0x1c1c141c, 0x1c1c143e, 0x1c1c1c04, 0x1c1c240c, 0x1c1c241c, 0x1c1c243e, 0x1c1c2c2c, 0x1c1c3e1c, + 0x1c24041c, 0x1c240c0c, 0x1c240c34, 0x1c241414, 0x1c241c0c, 0x1c242c14, 0x1c243404, 0x1c243424, 0x1c2c040c, + 0x1c2c0c04, 0x1c2c0c14, 0x1c2c142c, 0x1c2c1c14, 0x1c2c2424, 0x1c2c2c34, 0x1c2c3e1c, 0x1c340c34, 0x1c34240c, + 0x1c3e040c, 0x1c3e041c, 0x1c3e1404, 0x1c3e1414, 0x1c3e1c2c, 0x24040404, 0x24040424, 0x24040c14, 0x24041404, + 0x24041424, 0x2404143e, 0x24041c14, 0x2404240c, 0x24042c04, 0x24043e04, 0x240c0414, 0x240c043e, 0x240c0c0c, + 0x240c0c1c, 0x240c1414, 0x240c1c04, 0x240c1c2c, 0x240c241c, 0x240c2c0c, 0x240c2c2c, 0x2414040c, 0x2414041c, + 0x24140c04, 0x24140c2c, 0x2414140c, 0x24141c1c, 0x24142404, 0x24142c3e, 0x24143414, 0x24143e04, 0x241c0424, + 0x241c0c0c, 0x241c0c1c, 0x241c1404, 0x241c1414, 0x241c1c0c, 0x241c1c2c, 0x24240404, 0x24240414, 0x24241424, + 0x24241c3e, 0x24242404, 0x24243e0c, 0x242c042c, 0x242c043e, 0x242c140c, 0x242c3414, 0x24340c1c, 0x24341c24, + 0x24343404, 0x243e0c04, 0x243e0c2c, 0x243e1c04, 0x243e241c, 0x243e2c0c, 0x2c040414, 0x2c040c04, 0x2c040c24, + 0x2c041414, 0x2c042404, 0x2c042424, 0x2c04243e, 0x2c042c14, 0x2c043434, 0x2c043e24, 0x2c0c040c, 0x2c0c041c, + 0x2c0c042c, 0x2c0c0c14, 0x2c0c140c, 0x2c0c1c14, 0x2c0c3e14, 0x2c140404, 0x2c140c0c, 0x2c14141c, 0x2c141c04, + 0x2c141c34, 0x2c142c1c, 0x2c1c0414, 0x2c1c043e, 0x2c1c0c04, 0x2c1c143e, 0x2c1c2424, 0x2c1c2c0c, 0x2c1c342c, + 0x2c1c3e1c, 0x2c24040c, 0x2c240424, 0x2c241404, 0x2c241c14, 0x2c242434, 0x2c2c0c14, 0x2c2c1434, 0x2c2c2c0c, + 0x2c2c2c1c, 0x2c342414, 0x2c3e0414, 0x2c3e0424, 0x2c3e1414, 0x34040c0c, 0x34040c1c, 0x34040c2c, 0x34041c0c, + 0x34041c1c, 0x34043404, 0x340c0404, 0x340c1404, 0x340c143e, 0x340c3424, 0x34140c14, 0x34141c24, 0x34142414, + 0x34142c2c, 0x34143414, 0x34143e04, 0x341c0404, 0x341c0c24, 0x341c140c, 0x341c2404, 0x3424142c, 0x3424241c, + 0x34243414, 0x342c0404, 0x342c041c, 0x342c1c24, 0x342c3404, 0x3434042c, 0x34342404, 0x343e0c0c, 0x343e0c1c, + 0x3e040404, 0x3e040424, 0x3e04043e, 0x3e041404, 0x3e041414, 0x3e041c34, 0x3e042404, 0x3e042c24, 0x3e043414, + 0x3e0c0414, 0x3e0c0c0c, 0x3e0c1424, 0x3e0c241c, 0x3e0c242c, 0x3e14040c, 0x3e140424, 0x3e140c04, 0x3e140c34, + 0x3e14140c, 0x3e141c04, 0x3e142c0c, 0x3e1c0414, 0x3e1c1c14, 0x3e1c1c2c, 0x3e1c2c1c, 0x3e24040c, 0x3e24042c, + 0x3e240c1c, 0x3e241404, 0x3e242c04, 0x3e2c1414, 0x3e2c2414, 0x3e340414, 0x3e341c0c, 0x3e3e0404, +}; + +#define IQ1S_DELTA 0.125f +#define IQ1M_DELTA 0.125f +static const __device__ uint64_t iq1s_grid_gpu[2048] = { + 0x00000000, 0x00000002, 0x00000101, 0x00000200, 0x00000202, 0x00010001, 0x00010101, 0x00020000, 0x00020002, + 0x00020200, 0x00020202, 0x01000101, 0x01010001, 0x01010100, 0x01010102, 0x01020101, 0x02000000, 0x02000002, + 0x02000200, 0x02000202, 0x02010101, 0x02020000, 0x02020002, 0x02020200, 0x02020202, 0x00000110, 0x00000111, + 0x00010011, 0x00010110, 0x00010112, 0x00010211, 0x00010212, 0x00020111, 0x01000011, 0x01000112, 0x01000211, + 0x01010012, 0x01010111, 0x01010212, 0x01020011, 0x01020110, 0x01020112, 0x01020210, 0x02000111, 0x02010011, + 0x02010110, 0x02010112, 0x02020111, 0x00000020, 0x00000022, 0x00000220, 0x00000222, 0x00010121, 0x00020020, + 0x00020022, 0x00020220, 0x00020222, 0x01000121, 0x01010021, 0x01010221, 0x01020120, 0x01020221, 0x02000020, + 0x02000022, 0x02000220, 0x02000222, 0x02010021, 0x02010121, 0x02010221, 0x02020020, 0x02020022, 0x02020220, + 0x02020222, 0x00011001, 0x00011100, 0x00011102, 0x00021101, 0x01001001, 0x01001201, 0x01011101, 0x01011202, + 0x01021100, 0x01021101, 0x02011001, 0x02011201, 0x02021101, 0x00001011, 0x00001110, 0x00001111, 0x00001112, + 0x00011111, 0x00011210, 0x00011212, 0x00021211, 0x01001010, 0x01001111, 0x01001212, 0x01011010, 0x01011011, + 0x01011110, 0x01011111, 0x01011112, 0x01011211, 0x01021010, 0x01021012, 0x01021111, 0x01021210, 0x01021212, + 0x02001011, 0x02011011, 0x02011111, 0x02011210, 0x02011212, 0x02021011, 0x02021110, 0x02021111, 0x02021112, + 0x02021211, 0x00011120, 0x00011221, 0x01001021, 0x01001120, 0x01011020, 0x01011022, 0x01011121, 0x01011220, + 0x01021020, 0x01021021, 0x01021122, 0x01021221, 0x02001121, 0x02011021, 0x02011120, 0x02011221, 0x00002000, + 0x00002002, 0x00002200, 0x00002202, 0x00012101, 0x00022000, 0x00022002, 0x00022200, 0x00022202, 0x01002101, + 0x01012001, 0x01012102, 0x01022101, 0x02002000, 0x02002002, 0x02002200, 0x02002202, 0x02012101, 0x02022000, + 0x02022002, 0x02022200, 0x02022202, 0x00002111, 0x00012011, 0x00012110, 0x00012211, 0x00022110, 0x00022111, + 0x01002011, 0x01012010, 0x01012011, 0x01012111, 0x01022011, 0x01022110, 0x01022211, 0x02012011, 0x02012110, + 0x02012112, 0x02012211, 0x02022111, 0x00002020, 0x00002022, 0x00002220, 0x00002222, 0x00012121, 0x00022020, + 0x00022022, 0x00022220, 0x00022222, 0x01002121, 0x01012021, 0x01012221, 0x01022021, 0x01022121, 0x02002020, + 0x02002022, 0x02002121, 0x02002220, 0x02002222, 0x02012121, 0x02022020, 0x02022022, 0x02022220, 0x02022222, + 0x00110000, 0x00110001, 0x00110100, 0x00110201, 0x00120100, 0x00120101, 0x01100001, 0x01100100, 0x01110000, + 0x01110101, 0x01110200, 0x01120001, 0x01120100, 0x01120101, 0x01120201, 0x02110001, 0x02110100, 0x02110102, + 0x02120001, 0x02120101, 0x00100011, 0x00100110, 0x00100112, 0x00100211, 0x00110010, 0x00110012, 0x00110111, + 0x00110210, 0x00120011, 0x00120110, 0x00120211, 0x01100111, 0x01100212, 0x01110010, 0x01110011, 0x01110012, + 0x01110110, 0x01110111, 0x01110112, 0x01110211, 0x01120010, 0x01120111, 0x02100110, 0x02110012, 0x02110111, + 0x02120011, 0x02120110, 0x00110021, 0x00110120, 0x00110122, 0x00120121, 0x01100020, 0x01100122, 0x01100221, + 0x01110022, 0x01110121, 0x01110220, 0x01110222, 0x01120120, 0x01120122, 0x02100121, 0x02110021, 0x02110120, + 0x02110122, 0x02120121, 0x00101001, 0x00101102, 0x00101201, 0x00111100, 0x00111101, 0x00111200, 0x00111201, + 0x00121001, 0x00121102, 0x01101001, 0x01101101, 0x01101102, 0x01101200, 0x01101202, 0x01111001, 0x01111100, + 0x01111101, 0x01111102, 0x01111201, 0x01121002, 0x01121101, 0x01121200, 0x02101100, 0x02101201, 0x02111000, + 0x02111100, 0x02111101, 0x02111200, 0x02111201, 0x02111202, 0x02121001, 0x02121100, 0x02121101, 0x02121201, + 0x00101012, 0x00101111, 0x00101212, 0x00111011, 0x00111110, 0x00111111, 0x00111112, 0x00111211, 0x00121010, + 0x00121012, 0x00121111, 0x00121210, 0x00121212, 0x01101011, 0x01101110, 0x01101111, 0x01101112, 0x01111011, + 0x01111012, 0x01111110, 0x01111111, 0x01111112, 0x01111211, 0x01111212, 0x01121011, 0x01121110, 0x01121111, + 0x01121112, 0x01121211, 0x02101010, 0x02101012, 0x02101110, 0x02101111, 0x02101210, 0x02101212, 0x02111010, + 0x02111011, 0x02111110, 0x02111111, 0x02111112, 0x02111211, 0x02111212, 0x02121010, 0x02121012, 0x02121111, + 0x00101021, 0x00101120, 0x00101121, 0x00101122, 0x00111121, 0x00111122, 0x00111220, 0x00111222, 0x00121021, + 0x00121122, 0x01101020, 0x01101022, 0x01101120, 0x01101121, 0x01101220, 0x01101222, 0x01111021, 0x01111121, + 0x01111122, 0x01111220, 0x01111221, 0x01121021, 0x01121120, 0x01121121, 0x01121220, 0x01121221, 0x01121222, + 0x02101122, 0x02101222, 0x02111022, 0x02111121, 0x02121120, 0x02121221, 0x00112001, 0x00112102, 0x00122101, + 0x01102001, 0x01102100, 0x01102102, 0x01102201, 0x01112000, 0x01112101, 0x01112200, 0x01112202, 0x01122000, + 0x01122001, 0x01122100, 0x01122102, 0x01122201, 0x02102101, 0x02112001, 0x02112100, 0x02122101, 0x00112010, + 0x00112012, 0x00112111, 0x00112212, 0x00122011, 0x00122111, 0x01102012, 0x01102110, 0x01102111, 0x01102210, + 0x01112011, 0x01112110, 0x01112111, 0x01112112, 0x01112211, 0x01112212, 0x01122010, 0x01122111, 0x01122212, + 0x02102211, 0x02112011, 0x02112012, 0x02112111, 0x02112210, 0x02122011, 0x02122112, 0x02122211, 0x00102221, + 0x00112122, 0x00122120, 0x00122122, 0x01102120, 0x01102122, 0x01102221, 0x01112020, 0x01112022, 0x01112121, + 0x01112220, 0x01122021, 0x01122122, 0x01122221, 0x02102121, 0x02112021, 0x02112122, 0x02112222, 0x00200000, + 0x00200002, 0x00200200, 0x00200202, 0x00210101, 0x00220000, 0x00220002, 0x00220101, 0x00220200, 0x00220202, + 0x01200101, 0x01210001, 0x01210201, 0x01220001, 0x01220101, 0x02200000, 0x02200002, 0x02200200, 0x02200202, + 0x02210101, 0x02220000, 0x02220002, 0x02220101, 0x02220200, 0x02220202, 0x00200111, 0x00210011, 0x00210110, + 0x00210211, 0x00220111, 0x01200012, 0x01200110, 0x01200211, 0x01210111, 0x01210210, 0x01210212, 0x01220011, + 0x01220110, 0x01220111, 0x01220112, 0x02200111, 0x02210010, 0x02210112, 0x02210211, 0x02220111, 0x00200021, + 0x00200220, 0x00200222, 0x00210021, 0x00210121, 0x00220020, 0x00220022, 0x00220220, 0x00220222, 0x01200121, + 0x01210021, 0x01210122, 0x01210221, 0x01220121, 0x02200021, 0x02200220, 0x02200222, 0x02210021, 0x02210121, + 0x02220020, 0x02220022, 0x02220220, 0x02220222, 0x00201101, 0x00211100, 0x00211102, 0x00211201, 0x00221101, + 0x01201100, 0x01201101, 0x01201102, 0x01201201, 0x01211002, 0x01211101, 0x01211200, 0x01211202, 0x01221102, + 0x02201101, 0x02211001, 0x02211100, 0x02211201, 0x02221001, 0x02221101, 0x00201211, 0x00211111, 0x00221011, + 0x00221211, 0x01201010, 0x01201111, 0x01201210, 0x01211011, 0x01211110, 0x01211111, 0x01211211, 0x01221012, + 0x01221111, 0x01221210, 0x02201211, 0x02211010, 0x02211110, 0x02211111, 0x02211210, 0x02211212, 0x02221011, + 0x02221110, 0x02221112, 0x02221211, 0x00201121, 0x00211020, 0x00211022, 0x00211221, 0x00221121, 0x01201021, + 0x01201221, 0x01211121, 0x01221020, 0x01221021, 0x01221221, 0x02201120, 0x02201122, 0x02211020, 0x02211222, + 0x00202000, 0x00202002, 0x00202200, 0x00202202, 0x00212101, 0x00222000, 0x00222002, 0x00222200, 0x00222202, + 0x01202101, 0x01212001, 0x01212100, 0x01222101, 0x02202000, 0x02202002, 0x02202200, 0x02202202, 0x02222000, + 0x02222002, 0x02222200, 0x02222202, 0x00202211, 0x00212011, 0x00212110, 0x00212211, 0x00222111, 0x01202112, + 0x01202211, 0x01212012, 0x01212111, 0x01222011, 0x01222110, 0x01222112, 0x01222211, 0x02202111, 0x02212010, + 0x02212112, 0x02212211, 0x02222110, 0x02222111, 0x00202020, 0x00202022, 0x00202220, 0x00202222, 0x00222020, + 0x00222022, 0x00222220, 0x00222222, 0x01202121, 0x01212021, 0x01212122, 0x01212221, 0x01222121, 0x02202020, + 0x02202022, 0x02202220, 0x02202222, 0x02212121, 0x02222020, 0x02222022, 0x02222220, 0x02222222, 0x10000101, + 0x10010001, 0x10010102, 0x10020101, 0x11000201, 0x11010002, 0x11010101, 0x11010200, 0x11010202, 0x11020001, + 0x11020100, 0x11020102, 0x12010100, 0x12010201, 0x12020001, 0x12020102, 0x10000010, 0x10000011, 0x10000110, + 0x10000112, 0x10000211, 0x10010012, 0x10010111, 0x10010112, 0x10010210, 0x10010212, 0x10020011, 0x10020112, + 0x10020211, 0x11000111, 0x11000210, 0x11000212, 0x11010011, 0x11010110, 0x11010111, 0x11010112, 0x11010211, + 0x11010212, 0x11020111, 0x11020210, 0x11020212, 0x12000011, 0x12000110, 0x12000112, 0x12010010, 0x12010012, + 0x12010111, 0x12020010, 0x12020011, 0x12020012, 0x10000121, 0x10010021, 0x10010120, 0x10010122, 0x10020121, + 0x11000021, 0x11010022, 0x11010121, 0x11010222, 0x11020120, 0x11020221, 0x12000221, 0x12010120, 0x12020121, + 0x10001001, 0x10011101, 0x10011201, 0x10021201, 0x11001101, 0x11001200, 0x11001202, 0x11011001, 0x11011100, + 0x11011101, 0x11011102, 0x11021001, 0x11021002, 0x11021101, 0x11021200, 0x11021202, 0x12001001, 0x12001102, + 0x12001201, 0x12011000, 0x12011002, 0x12011101, 0x12021000, 0x12021001, 0x12021201, 0x10001011, 0x10001012, + 0x10001111, 0x10001212, 0x10011011, 0x10011110, 0x10011111, 0x10011112, 0x10011211, 0x10021010, 0x10021111, + 0x10021212, 0x11001011, 0x11001110, 0x11001111, 0x11001112, 0x11001211, 0x11011010, 0x11011011, 0x11011110, + 0x11011111, 0x11011112, 0x11011210, 0x11011211, 0x11021011, 0x11021110, 0x11021111, 0x11021112, 0x11021211, + 0x12001012, 0x12001110, 0x12001111, 0x12001210, 0x12011011, 0x12011110, 0x12011111, 0x12011112, 0x12011211, + 0x12011212, 0x12021111, 0x12021210, 0x12021212, 0x10001021, 0x10001121, 0x10001221, 0x10011120, 0x10011121, + 0x10011220, 0x10011222, 0x10021021, 0x10021120, 0x10021221, 0x11001020, 0x11001022, 0x11001121, 0x11001220, + 0x11011020, 0x11011021, 0x11011022, 0x11011121, 0x11011122, 0x11011221, 0x11021022, 0x11021121, 0x11021220, + 0x12001021, 0x12001121, 0x12001222, 0x12011120, 0x12011121, 0x12021021, 0x12021120, 0x12021122, 0x10002101, + 0x10012001, 0x10012101, 0x10012202, 0x10022101, 0x11002002, 0x11002201, 0x11012000, 0x11012101, 0x11012200, + 0x11022001, 0x11022100, 0x11022102, 0x11022201, 0x12002101, 0x12012001, 0x12012100, 0x12012102, 0x12012201, + 0x12022101, 0x10002011, 0x10002111, 0x10002112, 0x10002212, 0x10012010, 0x10012110, 0x10012111, 0x10012210, + 0x10022011, 0x10022110, 0x10022112, 0x11002010, 0x11002111, 0x11002212, 0x11012011, 0x11012012, 0x11012110, + 0x11012111, 0x11012112, 0x11012211, 0x11022010, 0x11022012, 0x11022111, 0x11022112, 0x11022212, 0x12002112, + 0x12002211, 0x12012012, 0x12012111, 0x12012112, 0x12012210, 0x12022011, 0x12022110, 0x12022112, 0x12022211, + 0x10012122, 0x11002120, 0x11002122, 0x11002221, 0x11012121, 0x11012220, 0x11012222, 0x11022120, 0x11022221, + 0x12012120, 0x12022121, 0x10100001, 0x10100100, 0x10100101, 0x10100102, 0x10100201, 0x10110002, 0x10110101, + 0x10110202, 0x10120001, 0x10120100, 0x10120201, 0x11100000, 0x11100101, 0x11100200, 0x11110001, 0x11110100, + 0x11110101, 0x11110102, 0x11110201, 0x11120101, 0x11120200, 0x12100102, 0x12100201, 0x12110101, 0x12110200, + 0x12120000, 0x12120001, 0x12120102, 0x12120201, 0x10100111, 0x10100210, 0x10100211, 0x10100212, 0x10110011, + 0x10110110, 0x10110111, 0x10110112, 0x10110210, 0x10110211, 0x10120010, 0x10120111, 0x10120112, 0x10120210, + 0x10120212, 0x11100011, 0x11100110, 0x11100111, 0x11100112, 0x11100211, 0x11110010, 0x11110011, 0x11110012, + 0x11110110, 0x11110111, 0x11110112, 0x11110210, 0x11110211, 0x11110212, 0x11120011, 0x11120110, 0x11120111, + 0x11120112, 0x11120211, 0x12100012, 0x12100111, 0x12110011, 0x12110110, 0x12110111, 0x12110112, 0x12110211, + 0x12120010, 0x12120111, 0x12120212, 0x10100021, 0x10100122, 0x10110022, 0x10110121, 0x10110222, 0x10120021, + 0x10120120, 0x11100022, 0x11100121, 0x11100222, 0x11110021, 0x11110120, 0x11110121, 0x11110122, 0x11110221, + 0x11120022, 0x11120121, 0x12100121, 0x12110020, 0x12110022, 0x12110121, 0x12110221, 0x12110222, 0x12120120, + 0x10101100, 0x10101101, 0x10111001, 0x10111100, 0x10111101, 0x10111102, 0x10111200, 0x10111201, 0x10121001, + 0x10121101, 0x10121200, 0x10121202, 0x11101001, 0x11101100, 0x11101101, 0x11101102, 0x11101201, 0x11101202, + 0x11111000, 0x11111001, 0x11111100, 0x11111101, 0x11111102, 0x11111200, 0x11111201, 0x11111202, 0x11121001, + 0x11121002, 0x11121100, 0x11121101, 0x11121102, 0x11121201, 0x12101000, 0x12101200, 0x12101202, 0x12111001, + 0x12111100, 0x12111101, 0x12111102, 0x12111201, 0x12121001, 0x12121100, 0x12121101, 0x12121202, 0x10101011, + 0x10101012, 0x10101110, 0x10101111, 0x10101112, 0x10101211, 0x10111010, 0x10111011, 0x10111012, 0x10111110, + 0x10111111, 0x10111112, 0x10111211, 0x10111212, 0x10121011, 0x10121110, 0x10121111, 0x10121112, 0x10121211, + 0x11101010, 0x11101011, 0x11101012, 0x11101110, 0x11101111, 0x11101112, 0x11101210, 0x11101211, 0x11111010, + 0x11111011, 0x11111012, 0x11111110, 0x11111111, 0x11111112, 0x11111210, 0x11111211, 0x11111212, 0x11121010, + 0x11121011, 0x11121110, 0x11121111, 0x11121112, 0x11121210, 0x11121211, 0x11121212, 0x12101011, 0x12101110, + 0x12101111, 0x12101211, 0x12101212, 0x12111010, 0x12111011, 0x12111110, 0x12111111, 0x12111112, 0x12111210, + 0x12111211, 0x12121011, 0x12121110, 0x12121111, 0x12121112, 0x12121211, 0x10101020, 0x10101021, 0x10101022, + 0x10101120, 0x10101122, 0x10101220, 0x10101221, 0x10111021, 0x10111120, 0x10111121, 0x10111220, 0x10111221, + 0x10121020, 0x10121021, 0x10121022, 0x10121120, 0x10121121, 0x10121122, 0x10121220, 0x10121221, 0x11101021, + 0x11101121, 0x11101122, 0x11101220, 0x11101221, 0x11101222, 0x11111020, 0x11111021, 0x11111022, 0x11111120, + 0x11111121, 0x11111122, 0x11111220, 0x11111221, 0x11111222, 0x11121021, 0x11121120, 0x11121121, 0x11121221, + 0x12101022, 0x12101121, 0x12101122, 0x12101220, 0x12101221, 0x12101222, 0x12111021, 0x12111121, 0x12111222, + 0x12121022, 0x12121121, 0x12121122, 0x12121220, 0x12121221, 0x10102100, 0x10102101, 0x10102102, 0x10102201, + 0x10112000, 0x10112101, 0x10112200, 0x10122001, 0x10122202, 0x11102101, 0x11102200, 0x11102202, 0x11112001, + 0x11112100, 0x11112101, 0x11112102, 0x11112200, 0x11112201, 0x11122000, 0x11122002, 0x11122100, 0x11122101, + 0x12102002, 0x12102201, 0x12112000, 0x12112002, 0x12112101, 0x12112200, 0x12122001, 0x12122201, 0x10102011, + 0x10102012, 0x10102111, 0x10102212, 0x10112011, 0x10112110, 0x10112111, 0x10112112, 0x10112211, 0x10122111, + 0x11102011, 0x11102110, 0x11102111, 0x11102112, 0x11102211, 0x11112010, 0x11112011, 0x11112012, 0x11112110, + 0x11112111, 0x11112112, 0x11112210, 0x11112211, 0x11112212, 0x11122011, 0x11122110, 0x11122111, 0x11122112, + 0x11122211, 0x12102011, 0x12102111, 0x12102211, 0x12112011, 0x12112110, 0x12112111, 0x12112112, 0x12112210, + 0x12112211, 0x12122111, 0x10102120, 0x10102220, 0x10112121, 0x10112222, 0x10122020, 0x10122121, 0x10122122, + 0x10122221, 0x11102121, 0x11102220, 0x11102221, 0x11112021, 0x11112121, 0x11112122, 0x11112220, 0x11112221, + 0x11122022, 0x11122121, 0x11122220, 0x11122222, 0x12102021, 0x12102222, 0x12112022, 0x12112121, 0x12112122, + 0x12112220, 0x12112222, 0x12122021, 0x10200101, 0x10210100, 0x10210102, 0x10210201, 0x10220101, 0x11200100, + 0x11210000, 0x11210101, 0x11210102, 0x11210200, 0x11210202, 0x11220001, 0x11220100, 0x11220102, 0x11220201, + 0x12200001, 0x12210102, 0x12220101, 0x10200011, 0x10200110, 0x10200112, 0x10200211, 0x10210012, 0x10210111, + 0x10220011, 0x10220012, 0x10220112, 0x10220211, 0x11200111, 0x11200211, 0x11210011, 0x11210111, 0x11210112, + 0x11210211, 0x11220111, 0x11220112, 0x11220212, 0x12200110, 0x12200212, 0x12210012, 0x12210111, 0x12220011, + 0x12220112, 0x12220211, 0x10210021, 0x10210122, 0x10210221, 0x11200020, 0x11200021, 0x11200122, 0x11210121, + 0x11210122, 0x11210220, 0x11220020, 0x12200121, 0x12210021, 0x12210122, 0x12220121, 0x10211001, 0x10211002, + 0x10211101, 0x10211102, 0x10211202, 0x10221001, 0x10221102, 0x10221201, 0x11201000, 0x11201002, 0x11201101, + 0x11201200, 0x11201202, 0x11211001, 0x11211100, 0x11211101, 0x11211102, 0x11211201, 0x11211202, 0x11221000, + 0x11221002, 0x11221101, 0x12201100, 0x12201101, 0x12201201, 0x12211000, 0x12211002, 0x12211100, 0x12211101, + 0x12211102, 0x12211200, 0x12211202, 0x12221001, 0x12221100, 0x12221201, 0x10201111, 0x10201210, 0x10201212, + 0x10211011, 0x10211111, 0x10211112, 0x10211211, 0x11201110, 0x11201111, 0x11201112, 0x11201211, 0x11211010, + 0x11211011, 0x11211110, 0x11211111, 0x11211112, 0x11211211, 0x11221011, 0x11221110, 0x11221111, 0x11221112, + 0x11221211, 0x12201112, 0x12201211, 0x12201212, 0x12211011, 0x12211111, 0x12211112, 0x12211211, 0x12211212, + 0x12221012, 0x12221111, 0x12221112, 0x12221210, 0x10201022, 0x10201221, 0x10211121, 0x10221020, 0x10221122, + 0x10221220, 0x10221221, 0x11201020, 0x11201121, 0x11201220, 0x11201222, 0x11211021, 0x11211120, 0x11211121, + 0x11211122, 0x11211220, 0x11211222, 0x11221020, 0x11221121, 0x11221220, 0x12201020, 0x12201022, 0x12201121, + 0x12201222, 0x12211120, 0x12211122, 0x12211220, 0x12211221, 0x12221020, 0x12221120, 0x12221122, 0x12221222, + 0x10212102, 0x10212201, 0x10222101, 0x11202001, 0x11212002, 0x11212101, 0x11212202, 0x11222001, 0x11222201, + 0x12202101, 0x12212001, 0x12212200, 0x12222102, 0x10202011, 0x10202110, 0x10212010, 0x10212111, 0x10222011, + 0x10222110, 0x10222112, 0x10222211, 0x11202010, 0x11202011, 0x11202111, 0x11202112, 0x11202210, 0x11212011, + 0x11212110, 0x11212111, 0x11212112, 0x11212211, 0x11222010, 0x11222111, 0x11222212, 0x12202012, 0x12202110, + 0x12202212, 0x12212111, 0x12222011, 0x12222110, 0x12222111, 0x12222211, 0x10212021, 0x10212122, 0x10212220, + 0x11202021, 0x11202120, 0x11202221, 0x11212020, 0x11212121, 0x11212220, 0x11212222, 0x11222120, 0x11222121, + 0x11222221, 0x12202122, 0x12212120, 0x12212220, 0x12212222, 0x12222122, 0x20000000, 0x20000002, 0x20000200, + 0x20000202, 0x20020000, 0x20020002, 0x20020200, 0x20020202, 0x21000101, 0x21010000, 0x21010001, 0x21010100, + 0x21010102, 0x21010201, 0x21020101, 0x22000000, 0x22000002, 0x22000200, 0x22000202, 0x22010101, 0x22020000, + 0x22020002, 0x22020200, 0x22020202, 0x20000111, 0x20010011, 0x20010110, 0x20010112, 0x20010211, 0x20020111, + 0x21000011, 0x21000110, 0x21000211, 0x21010010, 0x21010012, 0x21010111, 0x21010112, 0x21010210, 0x21010211, + 0x21020110, 0x21020112, 0x21020211, 0x22000111, 0x22000211, 0x22010110, 0x22010112, 0x22010211, 0x22020111, + 0x20000020, 0x20000022, 0x20000220, 0x20000222, 0x20010121, 0x20020020, 0x20020022, 0x20020220, 0x20020222, + 0x21010021, 0x21010120, 0x21010221, 0x21020121, 0x22000020, 0x22000022, 0x22000220, 0x22000222, 0x22010121, + 0x22020020, 0x22020022, 0x22020220, 0x22020222, 0x20011100, 0x20011201, 0x21001001, 0x21001100, 0x21011001, + 0x21011101, 0x21011202, 0x21021001, 0x21021100, 0x21021201, 0x22011100, 0x22011201, 0x20001011, 0x20001211, + 0x20011012, 0x20011111, 0x20011212, 0x20021112, 0x20021211, 0x21001010, 0x21001011, 0x21001111, 0x21001210, + 0x21011011, 0x21011110, 0x21011111, 0x21011112, 0x21011211, 0x21011212, 0x21021111, 0x21021112, 0x21021210, + 0x21021212, 0x22001011, 0x22001110, 0x22001112, 0x22001211, 0x22011010, 0x22011012, 0x22011111, 0x22011210, + 0x22021112, 0x20011021, 0x20011122, 0x20011221, 0x20021121, 0x21001021, 0x21001120, 0x21001221, 0x21001222, + 0x21011020, 0x21011121, 0x21011221, 0x21011222, 0x21021021, 0x21021122, 0x21021222, 0x22001121, 0x22011021, + 0x22011222, 0x22021120, 0x20002000, 0x20002002, 0x20002200, 0x20002202, 0x20012101, 0x20022000, 0x20022002, + 0x20022200, 0x20022202, 0x21002001, 0x21002101, 0x21012001, 0x21012100, 0x21012201, 0x21022101, 0x21022201, + 0x22002000, 0x22002002, 0x22002200, 0x22002202, 0x22012101, 0x22022000, 0x22022002, 0x22022200, 0x22022202, + 0x20002111, 0x20002112, 0x20012011, 0x20012110, 0x20012112, 0x20022111, 0x21002011, 0x21002110, 0x21002112, + 0x21002211, 0x21012010, 0x21012012, 0x21012111, 0x21012212, 0x21022011, 0x21022110, 0x22002111, 0x22012112, + 0x22012211, 0x22022111, 0x20002020, 0x20002022, 0x20002220, 0x20002222, 0x20012121, 0x20022020, 0x20022022, + 0x20022220, 0x20022222, 0x21002121, 0x21012021, 0x21012120, 0x21012122, 0x22002020, 0x22002022, 0x22002220, + 0x22002222, 0x22012121, 0x22022020, 0x22022022, 0x22022220, 0x22022222, 0x20100101, 0x20110001, 0x20110102, + 0x20110200, 0x20110201, 0x20120101, 0x21100001, 0x21100102, 0x21100201, 0x21110101, 0x21110200, 0x21110202, + 0x21120201, 0x21120202, 0x22100101, 0x22110001, 0x22110100, 0x22110102, 0x22110201, 0x22120101, 0x20100011, + 0x20100110, 0x20100112, 0x20100211, 0x20110010, 0x20110111, 0x20110210, 0x20110212, 0x20120011, 0x20120110, + 0x20120112, 0x20120211, 0x21100010, 0x21100111, 0x21110010, 0x21110011, 0x21110110, 0x21110111, 0x21110112, + 0x21110211, 0x21120012, 0x21120111, 0x22100110, 0x22100112, 0x22110012, 0x22110111, 0x22110210, 0x22120011, + 0x22120110, 0x22120112, 0x22120211, 0x20100121, 0x20110021, 0x20110120, 0x20110221, 0x20120121, 0x21100120, + 0x21100122, 0x21100221, 0x21110020, 0x21110022, 0x21110121, 0x21110220, 0x21120122, 0x21120221, 0x22100121, + 0x22110120, 0x22110122, 0x22120221, 0x20101001, 0x20101100, 0x20101102, 0x20111000, 0x20111101, 0x20111200, + 0x20121102, 0x21101000, 0x21101202, 0x21111001, 0x21111100, 0x21111101, 0x21111102, 0x21111200, 0x21111201, + 0x21121000, 0x21121001, 0x21121002, 0x21121101, 0x22101100, 0x22101102, 0x22111002, 0x22111100, 0x22111101, + 0x22111200, 0x22121001, 0x22121201, 0x20101010, 0x20101111, 0x20101210, 0x20101212, 0x20111010, 0x20111011, + 0x20111110, 0x20111111, 0x20111112, 0x20111211, 0x20121011, 0x20121111, 0x20121211, 0x20121212, 0x21101011, + 0x21101110, 0x21101111, 0x21101112, 0x21101211, 0x21111010, 0x21111011, 0x21111012, 0x21111110, 0x21111111, + 0x21111112, 0x21111210, 0x21111211, 0x21111212, 0x21121011, 0x21121110, 0x21121111, 0x21121112, 0x21121211, + 0x22101011, 0x22101111, 0x22101210, 0x22111011, 0x22111012, 0x22111110, 0x22111111, 0x22111112, 0x22111211, + 0x22111212, 0x22121010, 0x22121012, 0x22121111, 0x22121210, 0x22121212, 0x20101021, 0x20101120, 0x20111020, + 0x20111121, 0x20111221, 0x20121020, 0x20121122, 0x20121221, 0x21101121, 0x21101220, 0x21101221, 0x21111021, + 0x21111022, 0x21111121, 0x21111122, 0x21111221, 0x21121121, 0x21121220, 0x22101022, 0x22101120, 0x22101221, + 0x22101222, 0x22111022, 0x22111120, 0x22111121, 0x22121120, 0x22121122, 0x22121221, 0x20102101, 0x20112102, + 0x20112201, 0x20122101, 0x21102001, 0x21102102, 0x21112000, 0x21112002, 0x21112101, 0x21112102, 0x21112202, + 0x21122100, 0x21122101, 0x22102101, 0x22112001, 0x22112102, 0x22112201, 0x22122101, 0x20102110, 0x20102112, + 0x20102211, 0x20112010, 0x20112012, 0x20112111, 0x20112210, 0x20112212, 0x20122010, 0x20122011, 0x20122110, + 0x20122112, 0x21102010, 0x21102012, 0x21102111, 0x21102210, 0x21102212, 0x21112011, 0x21112110, 0x21112111, + 0x21112112, 0x21112211, 0x21122012, 0x21122111, 0x21122112, 0x21122212, 0x22102011, 0x22102110, 0x22112010, + 0x22112012, 0x22112111, 0x22112212, 0x22122011, 0x22122112, 0x20102121, 0x20112121, 0x20122121, 0x21102120, + 0x21102122, 0x21102221, 0x21112020, 0x21112121, 0x21112220, 0x21122021, 0x22102121, 0x22112021, 0x22112120, + 0x22112121, 0x22112122, 0x20200000, 0x20200002, 0x20200200, 0x20200202, 0x20210101, 0x20220000, 0x20220002, + 0x20220200, 0x20220202, 0x21200101, 0x21210001, 0x21210100, 0x21210102, 0x21210201, 0x22200000, 0x22200002, + 0x22200200, 0x22200202, 0x22210101, 0x22220000, 0x22220002, 0x22220200, 0x22220202, 0x20200111, 0x20200211, + 0x20210011, 0x20210110, 0x20210112, 0x20210211, 0x20210212, 0x21200112, 0x21200211, 0x21210011, 0x21210111, + 0x21210210, 0x21210212, 0x21220011, 0x21220110, 0x22200111, 0x22210010, 0x22210012, 0x22210112, 0x22210211, + 0x20200022, 0x20200220, 0x20200222, 0x20210020, 0x20210221, 0x20220022, 0x20220220, 0x20220222, 0x21200121, + 0x21210021, 0x21210122, 0x21210221, 0x21220121, 0x22200020, 0x22200022, 0x22200220, 0x22200222, 0x22210121, + 0x22220020, 0x22220022, 0x22220220, 0x22220222, 0x20211201, 0x20221101, 0x21201001, 0x21201100, 0x21211000, + 0x21211100, 0x21211101, 0x21211200, 0x21211202, 0x21221001, 0x21221101, 0x21221102, 0x21221200, 0x21221201, + 0x22201101, 0x20201112, 0x20201211, 0x20211010, 0x20211012, 0x20211111, 0x20211210, 0x20221112, 0x20221211, + 0x21201012, 0x21201111, 0x21211011, 0x21211110, 0x21211111, 0x21211112, 0x21211211, 0x21221111, 0x21221212, + 0x22201011, 0x22201110, 0x22201111, 0x22201112, 0x22201211, 0x22211012, 0x22211111, 0x22211210, 0x20201121, + 0x20211021, 0x20211122, 0x20211222, 0x20221021, 0x20221121, 0x21201120, 0x21201122, 0x21201222, 0x21211022, + 0x21211121, 0x21211122, 0x21211220, 0x21221020, 0x21221022, 0x22201122, 0x22211020, 0x22211121, 0x22211122, + 0x22211221, 0x22221021, 0x22221120, 0x22221122, 0x20202000, 0x20202002, 0x20202200, 0x20202202, 0x20222000, + 0x20222002, 0x20222200, 0x20222202, 0x21212001, 0x21212100, 0x21212102, 0x21212201, 0x22202000, 0x22202002, + 0x22202200, 0x22202202, 0x22212101, 0x22222000, 0x22222002, 0x22222200, 0x22222202, 0x20202111, 0x20212110, + 0x20212211, 0x20222011, 0x20222111, 0x21202011, 0x21212010, 0x21212111, 0x21212212, 0x21222011, 0x21222112, + 0x21222211, 0x22212010, 0x22212112, 0x20202020, 0x20202022, 0x20202220, 0x20202222, 0x20222020, 0x20222022, + 0x20222220, 0x20222222, 0x21212021, 0x21212120, 0x21212122, 0x22202020, 0x22202022, 0x22202220, 0x22202222, + 0x22212121, 0x22222020, 0x22222022, 0x22222220, 0x22222222, +}; + +static const __device__ uint8_t ksigns_iq2xs[128] = { + 0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15, 144, 17, 18, 147, 20, 149, + 150, 23, 24, 153, 154, 27, 156, 29, 30, 159, 160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, + 172, 45, 46, 175, 48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63, 192, 65, + 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207, 80, 209, 210, 83, 212, 85, 86, 215, + 216, 89, 90, 219, 92, 221, 222, 95, 96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, + 238, 111, 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255, +}; + +static const __device__ uint64_t ksigns64[128] = { + 0x0000000000000000, 0xff000000000000ff, 0xff0000000000ff00, 0x000000000000ffff, 0xff00000000ff0000, + 0x0000000000ff00ff, 0x0000000000ffff00, 0xff00000000ffffff, 0xff000000ff000000, 0x00000000ff0000ff, + 0x00000000ff00ff00, 0xff000000ff00ffff, 0x00000000ffff0000, 0xff000000ffff00ff, 0xff000000ffffff00, + 0x00000000ffffffff, 0xff0000ff00000000, 0x000000ff000000ff, 0x000000ff0000ff00, 0xff0000ff0000ffff, + 0x000000ff00ff0000, 0xff0000ff00ff00ff, 0xff0000ff00ffff00, 0x000000ff00ffffff, 0x000000ffff000000, + 0xff0000ffff0000ff, 0xff0000ffff00ff00, 0x000000ffff00ffff, 0xff0000ffffff0000, 0x000000ffffff00ff, + 0x000000ffffffff00, 0xff0000ffffffffff, 0xff00ff0000000000, 0x0000ff00000000ff, 0x0000ff000000ff00, + 0xff00ff000000ffff, 0x0000ff0000ff0000, 0xff00ff0000ff00ff, 0xff00ff0000ffff00, 0x0000ff0000ffffff, + 0x0000ff00ff000000, 0xff00ff00ff0000ff, 0xff00ff00ff00ff00, 0x0000ff00ff00ffff, 0xff00ff00ffff0000, + 0x0000ff00ffff00ff, 0x0000ff00ffffff00, 0xff00ff00ffffffff, 0x0000ffff00000000, 0xff00ffff000000ff, + 0xff00ffff0000ff00, 0x0000ffff0000ffff, 0xff00ffff00ff0000, 0x0000ffff00ff00ff, 0x0000ffff00ffff00, + 0xff00ffff00ffffff, 0xff00ffffff000000, 0x0000ffffff0000ff, 0x0000ffffff00ff00, 0xff00ffffff00ffff, + 0x0000ffffffff0000, 0xff00ffffffff00ff, 0xff00ffffffffff00, 0x0000ffffffffffff, 0xffff000000000000, + 0x00ff0000000000ff, 0x00ff00000000ff00, 0xffff00000000ffff, 0x00ff000000ff0000, 0xffff000000ff00ff, + 0xffff000000ffff00, 0x00ff000000ffffff, 0x00ff0000ff000000, 0xffff0000ff0000ff, 0xffff0000ff00ff00, + 0x00ff0000ff00ffff, 0xffff0000ffff0000, 0x00ff0000ffff00ff, 0x00ff0000ffffff00, 0xffff0000ffffffff, + 0x00ff00ff00000000, 0xffff00ff000000ff, 0xffff00ff0000ff00, 0x00ff00ff0000ffff, 0xffff00ff00ff0000, + 0x00ff00ff00ff00ff, 0x00ff00ff00ffff00, 0xffff00ff00ffffff, 0xffff00ffff000000, 0x00ff00ffff0000ff, + 0x00ff00ffff00ff00, 0xffff00ffff00ffff, 0x00ff00ffffff0000, 0xffff00ffffff00ff, 0xffff00ffffffff00, + 0x00ff00ffffffffff, 0x00ffff0000000000, 0xffffff00000000ff, 0xffffff000000ff00, 0x00ffff000000ffff, + 0xffffff0000ff0000, 0x00ffff0000ff00ff, 0x00ffff0000ffff00, 0xffffff0000ffffff, 0xffffff00ff000000, + 0x00ffff00ff0000ff, 0x00ffff00ff00ff00, 0xffffff00ff00ffff, 0x00ffff00ffff0000, 0xffffff00ffff00ff, + 0xffffff00ffffff00, 0x00ffff00ffffffff, 0xffffffff00000000, 0x00ffffff000000ff, 0x00ffffff0000ff00, + 0xffffffff0000ffff, 0x00ffffff00ff0000, 0xffffffff00ff00ff, 0xffffffff00ffff00, 0x00ffffff00ffffff, + 0x00ffffffff000000, 0xffffffffff0000ff, 0xffffffffff00ff00, 0x00ffffffff00ffff, 0xffffffffffff0000, + 0x00ffffffffff00ff, 0x00ffffffffffff00, 0xffffffffffffffff, +}; + +static const __device__ uint8_t kmask_iq2xs[8] = {1, 2, 4, 8, 16, 32, 64, 128}; +static const __device__ int8_t kvalues_iq4nl[16] = { + -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; + +typedef half dfloat; // dequantize float +typedef half2 dfloat2; +typedef void (*dequantize_kernel_t)(const void* vx, const int ib, const int iqs, dfloat2& v); +template +using to_cuda_ggml_t = void (*)(const void* __restrict__ x, dst_t* __restrict__ y, int k, hipStream_t stream); +typedef float (*vec_dot_q_cuda_t)(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs); +typedef void (*allocate_tiles_cuda_t)(int** x_ql, half2** x_dm, int** x_qh, int** x_sc); +typedef void (*load_tiles_cuda_t)( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row); +typedef float (*vec_dot_q_mul_mat_cuda_t)( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ms, + const int& i, + const int& j, + const int& k); + +// Utility function + +template +static __device__ __forceinline__ dst_t convert_from_half(half val) { + return val; +} + +template <> +__device__ __forceinline__ c10::BFloat16 convert_from_half(half val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 || defined(USE_MUSA) + return __float2bfloat16(__half2float(val)); +#else + return __half2float(val); +#endif // defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 || defined(USE_MUSA) +} + +template <> +__device__ __forceinline__ float convert_from_half(half val) { + return __half2float(val); +} + +#if defined(USE_ROCM) + +#ifndef __has_builtin +#define __has_builtin(x) 0 +#endif + +typedef int8_t int8x4_t __attribute__((ext_vector_type(4))); +static __device__ __forceinline__ int __vsubss4(const int a, const int b) { + const int8x4_t va = reinterpret_cast(a); + const int8x4_t vb = reinterpret_cast(b); +#if __has_builtin(__builtin_elementwise_sub_sat) + const int8x4_t c = __builtin_elementwise_sub_sat(va, vb); + return reinterpret_cast(c); +#else + int8x4_t c; + int16_t tmp; +#pragma unroll + for (int i = 0; i < 4; i++) { + tmp = va[i] - vb[i]; + if (tmp > std::numeric_limits::max()) tmp = std::numeric_limits::max(); + if (tmp < std::numeric_limits::min()) tmp = std::numeric_limits::min(); + c[i] = tmp; + } + return reinterpret_cast(c); +#endif // __has_builtin(__builtin_elementwise_sub_sat) +} + +static __device__ __forceinline__ int __dp4a(const int a, const int b, int c) { +#if __has_builtin(__builtin_amdgcn_sdot4) + c = __builtin_amdgcn_sdot4(a, b, c, false); +#else + const int8x4_t va = reinterpret_cast(a); + const int8x4_t vb = reinterpret_cast(b); + c += va[0] * vb[0] + va[1] * vb[1] + va[2] * vb[2] + va[3] * vb[3]; +#endif + return c; +} + +static __device__ __forceinline__ uint32_t __vcmpeq4(const uint32_t a, const uint32_t b) { + uint32_t neq = a ^ b; + return !(neq & 0xff000000) * 0xff000000 | !(neq & 0x00ff0000) * 0x00ff0000 | !(neq & 0x0000ff00) * 0x0000ff00 | + !(neq & 0x000000ff) * 0x000000ff; +} + +static __device__ __forceinline__ uint32_t __vsub4(const uint32_t a, const uint32_t b) { + return (static_cast(((a & 0xff000000) >> 24) - ((b & 0xff000000) >> 24)) << 24) + + (static_cast(((a & 0x00ff0000) >> 16) - ((b & 0x00ff0000) >> 16)) << 16) + + (static_cast(((a & 0x0000ff00) >> 8) - ((b & 0x0000ff00) >> 8)) << 8) + + (static_cast(((a & 0x000000ff) >> 0) - ((b & 0x000000ff) >> 0)) << 0); +} +#endif // defined(USE_ROCM) diff --git a/python/freetoken/kernel/csrc/gguf/gguf_b10434_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_b10434_kernel.cu new file mode 100644 index 000000000..f188162e6 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/gguf_b10434_kernel.cu @@ -0,0 +1,88 @@ +// Narrow caller-owned workspace ABI for pinned llama.cpp b10434 single-token work. +// The candidate source supplies the quantized GEMV implementation; this translation unit +// owns ABI validation, stable Q8_1 scratch slicing, and the public binding. Building both +// translation units into one extension keeps one pybind module and one stream contract. +#include +#include +#include + +void bind_gguf_moe_gfx1100(pybind11::module_&); + +torch::Tensor ggml_moe_mmvq_id( + torch::Tensor X, torch::Tensor W, torch::Tensor topk_ids, + int64_t top_k, int64_t type, int64_t row, int64_t tokens, + int64_t expert_stride_bytes, int64_t row_stride_bytes, + const std::string& id_space, torch::Tensor output, torch::Tensor quant_X_input); + +namespace { +constexpr int64_t kAlign = 256; +constexpr int64_t align256(int64_t value) { return (value + kAlign - 1) / kAlign * kAlign; } +} + +int64_t mmvq_bs1_workspace_bytes(int64_t hidden, int64_t rows, int64_t channels) { + TORCH_CHECK(hidden > 0 && rows > 0 && channels > 0, "MMVQ shape must be positive"); + TORCH_CHECK(hidden % 32 == 0, "MMVQ b10434 ABI requires hidden dimension divisible by 32"); + const int64_t activation = align256((hidden / 32) * 36); + const int64_t output = align256(channels * rows * static_cast(sizeof(float))); + return activation + output; +} + +torch::Tensor mmvq_bs1( + torch::Tensor x, torch::Tensor weight, torch::Tensor output, torch::Tensor workspace, + int64_t quant_type, int64_t rows, int64_t channels, + torch::Tensor route_ids = torch::Tensor()) { + TORCH_CHECK(x.is_cuda() && weight.is_cuda() && output.is_cuda() && workspace.is_cuda(), + "MMVQ b10434 ABI requires CUDA/HIP tensors"); + TORCH_CHECK(x.device() == weight.device() && output.device() == weight.device() && + workspace.device() == weight.device(), + "MMVQ b10434 tensors must share device"); + TORCH_CHECK(x.is_contiguous() && weight.stride(2) == 1 && output.is_contiguous() && + workspace.is_contiguous(), + "MMVQ b10434 ABI requires contiguous caller-owned buffers"); + TORCH_CHECK(x.dim() == 2 && x.size(0) == 1, "MMVQ b10434 accepts [1,H] activation"); + TORCH_CHECK(output.scalar_type() == torch::kFloat32 && output.sizes() == + torch::IntArrayRef({channels, rows}), "MMVQ output must be FP32 [channels, rows]"); + TORCH_CHECK(workspace.numel() * workspace.element_size() >= + mmvq_bs1_workspace_bytes(x.size(1), rows, channels), "MMVQ workspace too small"); + TORCH_CHECK(workspace.scalar_type() == torch::kUInt8 && workspace.dim() == 1, + "MMVQ workspace must be a contiguous uint8 byte buffer"); + TORCH_CHECK(quant_type == 2 || quant_type == 8 || quant_type == 12 || + quant_type == 13 || quant_type == 14, "unsupported GGUF quant type"); + TORCH_CHECK(weight.dim() == 3 && weight.size(1) == rows && weight.size(0) >= channels, + "MMVQ weight must be [experts, rows, row_bytes]"); + // The linked b10434 candidate currently uses the upstream 512-column Q8_1 tile. + // Qwen H=2048 is exact; rejecting other geometry avoids hidden padding allocation. + TORCH_CHECK(x.size(1) % 512 == 0, + "MMVQ b10434 ABI requires hidden dimension divisible by 512"); + torch::Tensor ids = route_ids; + if (!ids.defined()) { + // Convenience path for eager ABI smoke only. Captured callers pass route_ids so + // no device allocation occurs after graph capture begins. + ids = torch::arange(channels, torch::TensorOptions().dtype(torch::kInt32).device(x.device())) + .view({1, channels}); + } + TORCH_CHECK(ids.is_cuda() && ids.device() == weight.device() && ids.is_contiguous() && + ids.scalar_type() == torch::kInt32 && ids.sizes() == torch::IntArrayRef({1, channels}), + "MMVQ route_ids must be contiguous CUDA int32 [1, channels]"); + TORCH_CHECK(x.scalar_type() == torch::kBFloat16, + "MMVQ b10434 ABI currently accepts BF16 activation"); + const int64_t activation_bytes = (x.size(1) / 32) * 36; + const int64_t activation_region = align256(activation_bytes); + auto quant_x = workspace.narrow(0, 0, activation_bytes).view(torch::kInt32) + .view({1, x.size(1) / 32 * 9}); + // Candidate GEMV preserves activation dtype. Use the caller-owned aligned output + // region as BF16 scratch, then cast into the ABI's FP32 output without allocation. + auto candidate_output = workspace.narrow(0, activation_region, channels * rows * 2) + .view(torch::kBFloat16).view({channels, rows}); + ggml_moe_mmvq_id( + x, weight, ids, channels, quant_type, rows, 1, + weight.stride(0), weight.stride(1), "raw", candidate_output, quant_x); + output.copy_(candidate_output); + return output; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + bind_gguf_moe_gfx1100(module); + module.def("mmvq_bs1_workspace_bytes", &mmvq_bs1_workspace_bytes); + module.def("mmvq_bs1", &mmvq_bs1); +} diff --git a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu index 09c83c6f2..d04c26460 100644 --- a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu +++ b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu @@ -1,5 +1,9 @@ -// Adatped from +// Adapted from // https://github.com/vllm-project/vllm/blob/755ed7b05be4743237d3339c4ff8c22bcaae04f4/csrc/quantization/gguf/gguf_kernel.cu +// Algorithm cross-check: llama.cpp 7e4c0a968 (b10434), +// ggml/src/ggml-cuda/mmvq.cu and ggml/src/ggml-cuda/ggml-cuda.cu. +// Local HIP wrappers below preserve the GGUF packed Q4_K/Q5_K/Q6_K/Q8_0 +// contracts; native strided Q5_K/Q6_K is FreeToken-specific cache glue. #if defined(USE_ROCM) // ROCm torch hipifies these headers into c10::cuda (masquerading-as-CUDA), providing // c10::cuda::OptionalCUDAGuard / getCurrentCUDAStream backed by HIP. c10/cuda/CUDAGuard.h @@ -8,6 +12,7 @@ #include #include #include +using cudaStream_t = hipStream_t; #else #include #include @@ -27,6 +32,17 @@ // clang-format off #include "dispatch.h" +#if defined(USE_ROCM) +// These are checked-in HIP translations. Do not rely on ignored/generated *.hip +// sidecars; this selector is the active JIT source contract. +#include "ggml-common_hip.h" +#include "vecdotq_hip.cuh" +#include "dequantize_hip.cuh" +#include "mmvq_hip.cuh" +#include "mmq_hip.cuh" +#include "moe_hip.cuh" +#include "moe_vec_hip.cuh" +#else #include "ggml-common.h" #include "vecdotq.cuh" #include "dequantize.cuh" @@ -34,6 +50,7 @@ #include "mmq.cuh" #include "moe.cuh" #include "moe_vec.cuh" +#endif // clang-format off // Q8 gemv @@ -563,15 +580,43 @@ torch::Tensor ggml_moe_a8_vec( int64_t top_k, int64_t type, int64_t row, - int64_t tokens) { + int64_t tokens, + torch::Tensor output = torch::Tensor(), + torch::Tensor quant_X_input = torch::Tensor()) { int col = X.sizes()[1]; const int padded = (col + 512 - 1) / 512 * 512; const GGUF_DEVICE_GUARD(device_of(X)); auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); - at::Tensor Y = torch::zeros({tokens * top_k, row}, options); + at::Tensor Y; + if (output.defined()) { + TORCH_CHECK(output.is_cuda() && output.is_contiguous(), + "ggml_moe_a8_vec output must be contiguous CUDA/HIP tensor"); + TORCH_CHECK(output.device() == W.device() && output.dtype() == X.dtype(), + "ggml_moe_a8_vec output device/dtype must match input/output"); + TORCH_CHECK(output.sizes() == torch::IntArrayRef({tokens * top_k, row}), + "ggml_moe_a8_vec output shape mismatch"); + Y = output; + // The legacy vector kernels skip invalid/padded routes. Preserve their old zero-fill + // contract while allowing callers to reuse fixed graph-address output storage. + Y.zero_(); + } else { + Y = torch::zeros({tokens * top_k, row}, options); + } cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); - at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options); + at::Tensor quant_X; + if (quant_X_input.defined()) { + TORCH_CHECK(quant_X_input.is_cuda() && quant_X_input.is_contiguous(), + "ggml_moe_a8_vec quant_X must be contiguous CUDA/HIP tensor"); + TORCH_CHECK(quant_X_input.device() == W.device() && + quant_X_input.scalar_type() == torch::kInt32, + "ggml_moe_a8_vec quant_X device/dtype mismatch"); + TORCH_CHECK(quant_X_input.sizes() == torch::IntArrayRef({tokens, padded / 32 * 9}), + "ggml_moe_a8_vec quant_X shape mismatch"); + quant_X = quant_X_input; + } else { + 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) { @@ -827,6 +872,74 @@ torch::Tensor ggml_moe_a8_vec( return Y; } +torch::Tensor ggml_moe_a8_vec_strided( + torch::Tensor X, torch::Tensor W, torch::Tensor topk_ids, + int64_t top_k, int64_t type, int64_t row, int64_t tokens, + int64_t expert_stride_bytes, int64_t row_stride_bytes, + torch::Tensor output = torch::Tensor(), + torch::Tensor quant_X_input = torch::Tensor()) { + TORCH_CHECK(X.is_cuda() && W.is_cuda() && topk_ids.is_cuda(), + "ggml_moe_a8_vec_strided requires CUDA/HIP tensors"); + TORCH_CHECK(X.is_contiguous() && W.is_contiguous() && topk_ids.is_contiguous(), + "ggml_moe_a8_vec_strided requires contiguous tensors"); + TORCH_CHECK(type == 13 || type == 14, + "ggml_moe_a8_vec_strided supports Q5_K (13) and Q6_K (14), got ", type); + TORCH_CHECK(X.dim() == 2 && W.dim() == 3 && topk_ids.dim() == 2, + "invalid ggml_moe_a8_vec_strided tensor ranks"); + TORCH_CHECK(tokens == X.size(0) && top_k == topk_ids.size(1), + "ggml_moe_a8_vec_strided token/top-k shape mismatch"); + TORCH_CHECK(expert_stride_bytes == W.stride(0) && row_stride_bytes == W.stride(1), + "weight strides must match contiguous uint8 tensor"); + const int col = X.size(1); + const int padded = (col + 512 - 1) / 512 * 512; + const GGUF_DEVICE_GUARD(device_of(X)); + auto output_options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); + at::Tensor Y; + if (output.defined()) { + TORCH_CHECK(output.is_cuda() && output.is_contiguous(), + "ggml_moe_a8_vec_strided output must be contiguous CUDA/HIP tensor"); + TORCH_CHECK(output.device() == W.device() && output.dtype() == X.dtype(), + "ggml_moe_a8_vec_strided output device/dtype mismatch"); + TORCH_CHECK(output.sizes() == torch::IntArrayRef({tokens * top_k, row}), + "ggml_moe_a8_vec_strided output shape mismatch"); + Y = output; + Y.zero_(); + } else { + Y = torch::zeros({tokens * top_k, row}, output_options); + } + auto quant_options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); + at::Tensor quant_X; + if (quant_X_input.defined()) { + TORCH_CHECK(quant_X_input.is_cuda() && quant_X_input.is_contiguous(), + "ggml_moe_a8_vec_strided quant_X must be contiguous CUDA/HIP tensor"); + TORCH_CHECK(quant_X_input.device() == W.device() && + quant_X_input.scalar_type() == torch::kInt32, + "ggml_moe_a8_vec_strided quant_X device/dtype mismatch"); + TORCH_CHECK(quant_X_input.sizes() == torch::IntArrayRef({tokens, padded / 32 * 9}), + "ggml_moe_a8_vec_strided quant_X shape mismatch"); + quant_X = quant_X_input; + } else { + quant_X = torch::empty({tokens, padded / 32 * 9}, quant_options); + } + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); + DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_vec_a8_strided", [&] { + quantize_row_q8_1_cuda( + (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, tokens, stream); + if (type == 13) { + moe_vec_q5_K_q8_1_strided_cuda( + (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, + quant_X.stride(0), expert_stride_bytes, row_stride_bytes, stream); + } else { + moe_vec_q6_K_q8_1_strided_cuda( + (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), + (int*)topk_ids.data_ptr(), top_k, tokens, col, row, + quant_X.stride(0), expert_stride_bytes, row_stride_bytes, stream); + } + }); + return Y; +} + int64_t ggml_moe_get_block_size(int64_t type) { switch (type) { case 2: @@ -861,6 +974,67 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("ggml_mul_mat_vec_a8", &ggml_mul_mat_vec_a8, ""); m.def("ggml_mul_mat_a8", &ggml_mul_mat_a8, ""); m.def("ggml_moe_a8", &ggml_moe_a8, ""); - m.def("ggml_moe_a8_vec", &ggml_moe_a8_vec, ""); + m.def("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) { + return ggml_moe_a8_vec(X, W, topk_ids, top_k, type, row, tokens); + }, "", + py::arg("X"), py::arg("W"), py::arg("topk_ids"), py::arg("top_k"), + py::arg("type"), py::arg("row"), py::arg("tokens")); + m.def("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, + torch::Tensor output) { + return ggml_moe_a8_vec(X, W, topk_ids, top_k, type, row, tokens, output); + }, "", + py::arg("X"), py::arg("W"), py::arg("topk_ids"), py::arg("top_k"), + py::arg("type"), py::arg("row"), py::arg("tokens"), py::arg("output")); + m.def("ggml_moe_a8_vec_workspace", + [](torch::Tensor X, torch::Tensor W, torch::Tensor topk_ids, + int64_t top_k, int64_t type, int64_t row, int64_t tokens, + torch::Tensor output, torch::Tensor quant_X) { + return ggml_moe_a8_vec( + X, W, topk_ids, top_k, type, row, tokens, output, quant_X); + }, "", + py::arg("X"), py::arg("W"), py::arg("topk_ids"), py::arg("top_k"), + py::arg("type"), py::arg("row"), py::arg("tokens"), py::arg("output"), + py::arg("quant_X")); + m.def("ggml_moe_a8_vec_strided", + [](torch::Tensor X, torch::Tensor W, torch::Tensor topk_ids, + int64_t top_k, int64_t type, int64_t row, int64_t tokens, + int64_t expert_stride_bytes, int64_t row_stride_bytes) { + return ggml_moe_a8_vec_strided( + X, W, topk_ids, top_k, type, row, tokens, + expert_stride_bytes, row_stride_bytes); + }, "", + py::arg("X"), py::arg("W"), py::arg("topk_ids"), py::arg("top_k"), + py::arg("type"), py::arg("row"), py::arg("tokens"), + py::arg("expert_stride_bytes"), py::arg("row_stride_bytes")); + m.def("ggml_moe_a8_vec_strided", + [](torch::Tensor X, torch::Tensor W, torch::Tensor topk_ids, + int64_t top_k, int64_t type, int64_t row, int64_t tokens, + int64_t expert_stride_bytes, int64_t row_stride_bytes, + torch::Tensor output) { + return ggml_moe_a8_vec_strided( + X, W, topk_ids, top_k, type, row, tokens, + expert_stride_bytes, row_stride_bytes, output); + }, "", + py::arg("X"), py::arg("W"), py::arg("topk_ids"), py::arg("top_k"), + py::arg("type"), py::arg("row"), py::arg("tokens"), + py::arg("expert_stride_bytes"), py::arg("row_stride_bytes"), + py::arg("output")); + m.def("ggml_moe_a8_vec_strided_workspace", + [](torch::Tensor X, torch::Tensor W, torch::Tensor topk_ids, + int64_t top_k, int64_t type, int64_t row, int64_t tokens, + int64_t expert_stride_bytes, int64_t row_stride_bytes, + torch::Tensor output, torch::Tensor quant_X) { + return ggml_moe_a8_vec_strided( + X, W, topk_ids, top_k, type, row, tokens, + expert_stride_bytes, row_stride_bytes, output, quant_X); + }, "", + py::arg("X"), py::arg("W"), py::arg("topk_ids"), py::arg("top_k"), + py::arg("type"), py::arg("row"), py::arg("tokens"), + py::arg("expert_stride_bytes"), py::arg("row_stride_bytes"), + py::arg("output"), py::arg("quant_X")); m.def("ggml_moe_get_block_size", &ggml_moe_get_block_size, ""); } diff --git a/python/freetoken/kernel/csrc/gguf/gguf_moe_gfx1100.cu b/python/freetoken/kernel/csrc/gguf/gguf_moe_gfx1100.cu new file mode 100644 index 000000000..1c92a35ea --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/gguf_moe_gfx1100.cu @@ -0,0 +1,546 @@ +// Separate gfx1100 GGUF MoE candidate module. Legacy gguf_kernel.cu remains +// independently loadable so candidate compile failure cannot strand fallback. +#if defined(USE_ROCM) +#include +#include +#include +#include +using cudaStream_t = hipStream_t; +#else +#include +#include +#include +#endif +#include +#include +#include + +#if defined(USE_ROCM) +#define GGUF_DEVICE_GUARD(device) c10::cuda::OptionalCUDAGuard device_guard(device) +#define GGUF_CURRENT_STREAM() c10::cuda::getCurrentCUDAStream() +#else +#define GGUF_DEVICE_GUARD(device) at::cuda::OptionalCUDAGuard device_guard(device) +#define GGUF_CURRENT_STREAM() at::cuda::getCurrentCUDAStream() +#endif + +// Keep include order aligned with gguf_kernel.cu. These are the active JIT +// headers; ignored/generated *.hip sidecars are not dependencies. +#include "dispatch.h" +#if defined(USE_ROCM) +#include "ggml-common_hip.h" +#include "vecdotq_hip.cuh" +#include "moe_vec_gfx1100_hip.cuh" +#else +#include "ggml-common.h" +#include "vecdotq.cuh" +#include "moe_vec_gfx1100.cuh" +#endif + +template +static __global__ void quantize_q8_1_gfx1100( + const scalar_t* __restrict__ x, void* __restrict__ vy, + const int kx, const int kx_padded) { + const int ix = blockDim.x * blockIdx.x + threadIdx.x; + if (ix >= kx_padded) { + return; + } + const int iy = blockDim.y * blockIdx.y + threadIdx.y; + const int i_padded = iy * kx_padded + ix; + block_q8_1* y = static_cast(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; + y[ib].qs[iqs] = amax == 0.0f ? 0 : static_cast(roundf(xi / d)); + if (iqs == 0) { + y[ib].ds.x = __float2half(d); + y[ib].ds.y = __float2half(sum); + } +} + +template +static void quantize_row_q8_1_gfx1100( + 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; + quantize_q8_1_gfx1100<<< + dim3(block_num_x, num_blocks_y, 1), dim3(CUDA_DEQUANTIZE_BLOCK_SIZE, 1, 1), + 0, stream>>>( + &x[off * kx], static_cast(vy) + off * (kx_padded / 32 * 9), + kx, kx_padded); + } +} + +torch::Tensor ggml_moe_a8_vec_gfx1100( + torch::Tensor X, torch::Tensor W, torch::Tensor topk_ids, + int64_t top_k, int64_t type, int64_t row, int64_t tokens) { + TORCH_CHECK(X.is_cuda() && W.is_cuda() && topk_ids.is_cuda(), + "gfx1100 GGUF MoE candidate requires CUDA/HIP tensors"); + TORCH_CHECK(X.is_contiguous() && W.is_contiguous() && topk_ids.is_contiguous(), + "gfx1100 GGUF MoE candidate requires contiguous tensors"); + TORCH_CHECK(type == 8 || type == 12, + "gfx1100 candidate supports Q8_0 (8) and Q4_K (12), got ", type); + TORCH_CHECK(X.dim() == 2 && W.dim() == 3 && topk_ids.dim() == 2, + "invalid gfx1100 GGUF MoE tensor ranks"); + const int col = X.sizes()[1]; + const int padded = (col + 512 - 1) / 512 * 512; + const GGUF_DEVICE_GUARD(device_of(X)); + auto output_options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); + at::Tensor Y = torch::zeros({tokens * top_k, row}, output_options); + auto quant_options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); + at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, quant_options); + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); + + DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_a8_vec_gfx1100", [&] { + quantize_row_q8_1_gfx1100( + static_cast(X.data_ptr()), quant_X.data_ptr(), col, tokens, stream); + if (type == 12) { + moe_vec_q4_K_q8_1_gfx1100( + W.data_ptr(), quant_X.data_ptr(), static_cast(Y.data_ptr()), + static_cast(topk_ids.data_ptr()), top_k, tokens, col, row, + quant_X.stride(0), stream); + } else { + moe_vec_q8_0_q8_1_gfx1100( + W.data_ptr(), quant_X.data_ptr(), static_cast(Y.data_ptr()), + static_cast(topk_ids.data_ptr()), top_k, tokens, col, row, + quant_X.stride(0), stream); + } + }); + return Y; +} + +torch::Tensor ggml_moe_mmvq_id( + torch::Tensor X, torch::Tensor W, torch::Tensor topk_ids, + int64_t top_k, int64_t type, int64_t row, int64_t tokens, + int64_t expert_stride_bytes, int64_t row_stride_bytes, + const std::string& id_space, + torch::Tensor output = torch::Tensor(), + torch::Tensor quant_X_input = torch::Tensor()) { + TORCH_CHECK(X.is_cuda() && W.is_cuda() && topk_ids.is_cuda(), + "ggml_moe_mmvq_id requires CUDA/HIP tensors"); + TORCH_CHECK(X.is_contiguous() && topk_ids.is_contiguous(), + "ggml_moe_mmvq_id requires contiguous activation/ID tensors"); + TORCH_CHECK(W.dim() == 3 && W.scalar_type() == torch::kUInt8 && W.stride(2) == 1, + "ggml_moe_mmvq_id requires packed uint8 [E,rows,bytes] weights"); + TORCH_CHECK(X.dim() == 2 && topk_ids.dim() == 2 && topk_ids.scalar_type() == torch::kInt, + "invalid ggml_moe_mmvq_id tensor ranks or ID dtype"); + TORCH_CHECK(tokens == X.size(0) && top_k == topk_ids.size(1), + "ggml_moe_mmvq_id token/top-k shape mismatch"); + TORCH_CHECK(type == 8 || type == 12 || type == 13 || type == 14, + "ggml_moe_mmvq_id supports Q4_K/Q5_K/Q6_K/Q8_0, got ", type); + TORCH_CHECK(id_space == "raw" || id_space == "slot", + "ggml_moe_mmvq_id id_space must be raw or slot"); + TORCH_CHECK(expert_stride_bytes == W.stride(0) && row_stride_bytes == W.stride(1), + "GGUF expert/row strides must match the supplied packed bank"); + const int col = X.size(1); + const int padded = (col + 512 - 1) / 512 * 512; + const GGUF_DEVICE_GUARD(device_of(X)); + auto output_options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); + at::Tensor Y; + if (output.defined()) { + TORCH_CHECK(output.is_cuda() && output.is_contiguous() && output.device() == W.device(), + "ggml_moe_mmvq_id output must be contiguous on bank device"); + TORCH_CHECK(output.scalar_type() == X.scalar_type() && + output.sizes() == torch::IntArrayRef({tokens * top_k, row}), + "ggml_moe_mmvq_id output shape/dtype mismatch"); + Y = output; + Y.zero_(); + } else { + Y = torch::zeros({tokens * top_k, row}, output_options); + } + auto quant_options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); + at::Tensor quant_X; + if (quant_X_input.defined()) { + TORCH_CHECK(quant_X_input.is_cuda() && quant_X_input.is_contiguous() && + quant_X_input.device() == W.device() && + quant_X_input.scalar_type() == torch::kInt32 && + quant_X_input.sizes() == torch::IntArrayRef({tokens, padded / 32 * 9}), + "ggml_moe_mmvq_id quant_X shape/device/dtype mismatch"); + quant_X = quant_X_input; + } else { + quant_X = torch::empty({tokens, padded / 32 * 9}, quant_options); + } + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); + DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_mmvq_id", [&] { + quantize_row_q8_1_gfx1100( + static_cast(X.data_ptr()), quant_X.data_ptr(), col, tokens, stream); + if (type == 12) { + moe_vec_q4_K_q8_1_gfx1100_id( + W.data_ptr(), quant_X.data_ptr(), static_cast(Y.data_ptr()), + static_cast(topk_ids.data_ptr()), top_k, tokens, col, row, + quant_X.stride(0), expert_stride_bytes, row_stride_bytes, stream); + } else if (type == 13) { + moe_vec_q5_K_q8_1_gfx1100_id( + W.data_ptr(), quant_X.data_ptr(), static_cast(Y.data_ptr()), + static_cast(topk_ids.data_ptr()), top_k, tokens, col, row, + quant_X.stride(0), expert_stride_bytes, row_stride_bytes, stream); + } else if (type == 14) { + moe_vec_q6_K_q8_1_gfx1100_id( + W.data_ptr(), quant_X.data_ptr(), static_cast(Y.data_ptr()), + static_cast(topk_ids.data_ptr()), top_k, tokens, col, row, + quant_X.stride(0), expert_stride_bytes, row_stride_bytes, stream); + } else { + moe_vec_q8_0_q8_1_gfx1100_id( + W.data_ptr(), quant_X.data_ptr(), static_cast(Y.data_ptr()), + static_cast(topk_ids.data_ptr()), top_k, tokens, col, row, + quant_X.stride(0), expert_stride_bytes, row_stride_bytes, stream); + } + }); + return Y; +} + +template +static __global__ void moe_gate_up_swiglu_id_kernel( + const uint8_t* __restrict__ weights, + const int32_t* __restrict__ quant_x, + scalar_t* __restrict__ output, + const int* __restrict__ ids, + const int top_k, + const int ncols, + const int nrows, + const int quant_x_stride, + const int64_t expert_stride_bytes, + const int64_t row_stride_bytes, + const int experts) { + const int row = blockIdx.x * blockDim.y + threadIdx.y; + const int route = blockIdx.y; + const int expert = ids[route]; + if (row >= nrows || expert < 0 || expert >= experts) return; + + const int token = route / top_k; + const int blocks_per_row = ncols / QK_K; + const int lanes_per_chunk = QI4_K / VDR_Q4_K_Q8_1_MMVQ; + const int blocks_per_wave = VDR_Q4_K_Q8_1_MMVQ * WARP_SIZE / QI4_K; + const uint8_t* expert_base = weights + expert * expert_stride_bytes; + const uint8_t* gate_base = expert_base + row * row_stride_bytes; + const uint8_t* up_base = expert_base + (row + nrows) * row_stride_bytes; + const block_q8_1* x = reinterpret_cast( + quant_x + token * quant_x_stride); + float gate = 0.0f; + float up = 0.0f; + for (int i = threadIdx.x / lanes_per_chunk; i < blocks_per_row; + i += blocks_per_wave) { + const int iqs = VDR_Q4_K_Q8_1_MMVQ * (threadIdx.x % lanes_per_chunk); + gate += vec_dot_q4_K_q8_1(gate_base + i * sizeof(block_q4_K), &x[i], iqs); + up += vec_dot_q4_K_q8_1(up_base + i * sizeof(block_q4_K), &x[i], iqs); + } + for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { + gate += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), gate, mask); + up += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), up, mask); + } + if (threadIdx.x == 0) { + const float silu = gate / (1.0f + expf(-gate)); + output[route * nrows + row] = static_cast(silu * up); + } +} + +torch::Tensor ggml_moe_gate_up_swiglu_id( + torch::Tensor X, torch::Tensor W, torch::Tensor topk_ids, + int64_t top_k, int64_t nrows, int64_t tokens, + int64_t expert_stride_bytes, int64_t row_stride_bytes, + const std::string& id_space, torch::Tensor output = torch::Tensor(), + torch::Tensor quant_X_input = torch::Tensor()) { + TORCH_CHECK(X.is_cuda() && W.is_cuda() && topk_ids.is_cuda(), + "ggml_moe_gate_up_swiglu_id requires CUDA/HIP tensors"); + TORCH_CHECK(X.is_contiguous() && topk_ids.is_contiguous(), + "ggml_moe_gate_up_swiglu_id requires contiguous activation/ID tensors"); + TORCH_CHECK(W.dim() == 3 && W.scalar_type() == torch::kUInt8 && W.stride(2) == 1, + "ggml_moe_gate_up_swiglu_id requires packed uint8 [E,2I,bytes] weights"); + TORCH_CHECK(X.dim() == 2 && topk_ids.dim() == 2 && topk_ids.scalar_type() == torch::kInt, + "invalid ggml_moe_gate_up_swiglu_id tensor ranks or ID dtype"); + TORCH_CHECK(tokens == X.size(0) && top_k == topk_ids.size(1) && + nrows * 2 == W.size(1), + "ggml_moe_gate_up_swiglu_id shape mismatch"); + TORCH_CHECK(id_space == "raw" || id_space == "slot", + "ggml_moe_gate_up_swiglu_id id_space must be raw or slot"); + const int col = X.size(1); + TORCH_CHECK(col > 0 && col % QK_K == 0 && row_stride_bytes >= + (col / QK_K) * static_cast(sizeof(block_q4_K)) && + expert_stride_bytes >= W.size(1) * row_stride_bytes, + "ggml_moe_gate_up_swiglu_id requires aligned columns and valid strides"); + const GGUF_DEVICE_GUARD(device_of(X)); + auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); + at::Tensor Y; + if (output.defined()) { + TORCH_CHECK(output.is_cuda() && output.is_contiguous() && output.device() == W.device() && + output.scalar_type() == X.scalar_type() && + output.sizes() == torch::IntArrayRef({tokens * top_k, nrows}), + "ggml_moe_gate_up_swiglu_id output shape/device/dtype mismatch"); + Y = output; + Y.zero_(); + } else { + Y = torch::zeros({tokens * top_k, nrows}, options); + } + const int padded = (col + 512 - 1) / 512 * 512; + auto quant_options = torch::TensorOptions().dtype(torch::kInt32).device(W.device()); + at::Tensor quant_X; + if (quant_X_input.defined()) { + TORCH_CHECK(quant_X_input.is_cuda() && quant_X_input.is_contiguous() && + quant_X_input.device() == W.device() && + quant_X_input.scalar_type() == torch::kInt32 && + quant_X_input.sizes() == torch::IntArrayRef({tokens, padded / 32 * 9}), + "ggml_moe_gate_up_swiglu_id quant_X shape/device/dtype mismatch"); + quant_X = quant_X_input; + } else { + quant_X = torch::empty({tokens, padded / 32 * 9}, quant_options); + } + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); + DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_gate_up_swiglu_id", [&] { + quantize_row_q8_1_gfx1100( + static_cast(X.data_ptr()), quant_X.data_ptr(), col, tokens, stream); + moe_gate_up_swiglu_id_kernel<<< + dim3((nrows + 3) / 4, tokens * top_k, 1), dim3(WARP_SIZE, 4, 1), 0, stream>>>( + W.data_ptr(), quant_X.data_ptr(), + static_cast(Y.data_ptr()), static_cast(topk_ids.data_ptr()), + top_k, col, nrows, quant_X.stride(0), expert_stride_bytes, row_stride_bytes, + W.size(0)); + }); + return Y; +} + +static __device__ __forceinline__ void mmvdq_get_scale_min_k4( + int j, const uint8_t* q, uint8_t& d, uint8_t& m) { + if (j < 4) { + d = q[j] & 63; + m = q[j + 4] & 63; + } else { + d = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4); + m = (q[j + 4] >> 4) | ((q[j] >> 6) << 4); + } +} + +template +static __device__ __forceinline__ float mmvdq_value( + const uint8_t* block, int index); + +template <> +static __device__ __forceinline__ float mmvdq_value<12>( + const uint8_t* block, int index) { + const block_q4_K* q = reinterpret_cast(block); + const int il = index / 64; + const int local = index & 63; + const bool high = local >= 32; + const int in_group = high ? local - 32 : local; + const int ir = in_group / 4; + const int l = in_group & 3; + const int scale_index = 2 * il + (high ? 1 : 0); + uint8_t scale, minimum; + mmvdq_get_scale_min_k4(scale_index, q->scales, scale, minimum); + const half d = __hmul(__low2half(q->dm), __int2half_rn(scale)); + const half m = __hmul(__high2half(q->dm), __int2half_rn(minimum)); + const uint8_t packed = q->qs[32 * il + 4 * ir + l]; + const int value = high ? packed >> 4 : packed & 0xF; + return __half2float(__hsub(__hmul(d, __int2half_rn(value)), m)); +} + +template <> +static __device__ __forceinline__ float mmvdq_value<13>( + const uint8_t* block, int index) { + const block_q5_K* q = reinterpret_cast(block); + const int il = index / 64; + const int local = index & 63; + const bool high = local >= 32; + const int in_group = high ? local - 32 : local; + const int ir = in_group / 2; + const int l = in_group & 1; + const int scale_index = 2 * il + (high ? 1 : 0); + uint8_t scale, minimum; + mmvdq_get_scale_min_k4(scale_index, q->scales, scale, minimum); + const half d = __hmul(__low2half(q->dm), __int2half_rn(scale)); + const half m = __hmul(__high2half(q->dm), __int2half_rn(minimum)); + const int high_bit = high ? 1 : 0; + const uint8_t packed = q->qs[32 * il + 2 * ir + l]; + const uint8_t high_bits = q->qh[2 * ir + l]; + const int value = (high ? packed >> 4 : packed & 0xF) + + ((high_bits & (1u << (2 * il + high_bit))) ? 16 : 0); + return __half2float(__hsub(__hmul(d, __int2half_rn(value)), m)); +} + +template <> +static __device__ __forceinline__ float mmvdq_value<14>( + const uint8_t* block, int index) { + const block_q6_K* q = reinterpret_cast(block); + const int ip = index / 128; + const int local = index & 127; + const int group = local / 32; + const int pos = local & 31; + const uint8_t ql = q->ql[64 * ip + (group & 1) * 32 + pos]; + const uint8_t qh = q->qh[32 * ip + pos]; + const int low = group < 2 ? ql & 0xF : ql >> 4; + const int value = low | (((qh >> (2 * group)) & 3) << 4); + const int scale_index = 8 * ip + (pos / 16) + 2 * group; + return __half2float(__hmul( + q->d, __int2half_rn(q->scales[scale_index] * (value - 32)))); +} + +template +static __global__ void moe_mmvdq_id_kernel( + const uint8_t* __restrict__ weights, + const scalar_t* __restrict__ x, + scalar_t* __restrict__ output, + const int* __restrict__ ids, + const int top_k, + const int ncols, + const int nrows, + const int64_t expert_stride_bytes, + const int64_t row_stride_bytes, + const int experts) { + const int row = blockIdx.x * blockDim.y + threadIdx.y; + const int route = blockIdx.y; + const int expert = ids[route]; + if (row >= nrows || expert < 0 || expert >= experts) return; + + const int token = route / top_k; + const uint8_t* row_data = weights + expert * expert_stride_bytes + + row * row_stride_bytes; + float sum = 0.0f; + for (int k = threadIdx.x; k < ncols; k += blockDim.x) { + const uint8_t* block = row_data + (k / QK_K) * block_bytes; + sum += mmvdq_value(block, k & (QK_K - 1)) * + static_cast(x[token * ncols + k]); + } + + const int lane = threadIdx.x; + for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { + sum += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), sum, mask); + } + if (lane == 0) output[route * nrows + row] = static_cast(sum); +} + +torch::Tensor ggml_moe_mmvdq_id( + torch::Tensor X, torch::Tensor W, torch::Tensor topk_ids, + int64_t top_k, int64_t type, int64_t row, int64_t tokens, + int64_t expert_stride_bytes, int64_t row_stride_bytes, + const std::string& id_space, torch::Tensor output = torch::Tensor()) { + TORCH_CHECK(X.is_cuda() && W.is_cuda() && topk_ids.is_cuda(), + "ggml_moe_mmvdq_id requires CUDA/HIP tensors"); + TORCH_CHECK(X.is_contiguous() && topk_ids.is_contiguous(), + "ggml_moe_mmvdq_id requires contiguous activation/ID tensors"); + TORCH_CHECK(W.dim() == 3 && W.scalar_type() == torch::kUInt8 && W.stride(2) == 1, + "ggml_moe_mmvdq_id requires packed uint8 [E,rows,bytes] weights"); + TORCH_CHECK(X.dim() == 2 && topk_ids.dim() == 2 && topk_ids.scalar_type() == torch::kInt, + "invalid ggml_moe_mmvdq_id tensor ranks or ID dtype"); + TORCH_CHECK(tokens == X.size(0) && top_k == topk_ids.size(1), + "ggml_moe_mmvdq_id token/top-k shape mismatch"); + TORCH_CHECK(type == 12 || type == 13 || type == 14, + "ggml_moe_mmvdq_id supports Q4_K/Q5_K/Q6_K, got ", type); + TORCH_CHECK(id_space == "raw" || id_space == "slot", + "ggml_moe_mmvdq_id id_space must be raw or slot"); + TORCH_CHECK(X.scalar_type() == torch::kFloat || X.scalar_type() == torch::kHalf || + X.scalar_type() == torch::kBFloat16, + "ggml_moe_mmvdq_id supports F32/F16/BF16 activations"); + const int col = X.size(1); + TORCH_CHECK(col > 0 && col % QK_K == 0 && row == W.size(1), + "ggml_moe_mmvdq_id requires QK_K-aligned columns and matching rows"); + const int block_bytes = type == 12 ? sizeof(block_q4_K) : + (type == 13 ? sizeof(block_q5_K) : sizeof(block_q6_K)); + TORCH_CHECK(row_stride_bytes >= (col / QK_K) * block_bytes && + expert_stride_bytes >= row * row_stride_bytes, + "GGUF MMVDQ strides are smaller than packed row extent"); + const GGUF_DEVICE_GUARD(device_of(X)); + auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device()); + at::Tensor Y; + if (output.defined()) { + TORCH_CHECK(output.is_cuda() && output.is_contiguous() && output.device() == W.device() && + output.scalar_type() == X.scalar_type() && + output.sizes() == torch::IntArrayRef({tokens * top_k, row}), + "ggml_moe_mmvdq_id output shape/device/dtype mismatch"); + Y = output; + Y.zero_(); + } else { + Y = torch::zeros({tokens * top_k, row}, options); + } + cudaStream_t stream = GGUF_CURRENT_STREAM().stream(); + DISPATCH_FLOAT_TYPES(X.scalar_type(), "ggml_moe_mmvdq_id", [&] { + const scalar_t* x = static_cast(X.data_ptr()); + scalar_t* y = static_cast(Y.data_ptr()); + const int* ids = static_cast(topk_ids.data_ptr()); + if (type == 12) { + moe_mmvdq_id_kernel<<< + dim3((row + 3) / 4, tokens * top_k, 1), dim3(WARP_SIZE, 4, 1), 0, stream>>>( + W.data_ptr(), x, y, ids, top_k, col, row, + expert_stride_bytes, row_stride_bytes, W.size(0)); + } else if (type == 13) { + moe_mmvdq_id_kernel<<< + dim3((row + 3) / 4, tokens * top_k, 1), dim3(WARP_SIZE, 4, 1), 0, stream>>>( + W.data_ptr(), x, y, ids, top_k, col, row, + expert_stride_bytes, row_stride_bytes, W.size(0)); + } else { + moe_mmvdq_id_kernel<<< + dim3((row + 3) / 4, tokens * top_k, 1), dim3(WARP_SIZE, 4, 1), 0, stream>>>( + W.data_ptr(), x, y, ids, top_k, col, row, + expert_stride_bytes, row_stride_bytes, W.size(0)); + } + }); + return Y; +} + +void bind_gguf_moe_gfx1100(pybind11::module_& m) { + m.def("ggml_moe_a8_vec_gfx1100", &ggml_moe_a8_vec_gfx1100, ""); + m.def("ggml_moe_mmvq_id", + [](torch::Tensor X, torch::Tensor W, torch::Tensor ids, + int64_t top_k, int64_t type, int64_t row, int64_t tokens, + int64_t expert_stride, int64_t row_stride, const std::string& id_space) { + return ggml_moe_mmvq_id(X, W, ids, top_k, type, row, tokens, + expert_stride, row_stride, id_space); + }, ""); + m.def("ggml_moe_mmvq_id_workspace", + [](torch::Tensor X, torch::Tensor W, torch::Tensor ids, + int64_t top_k, int64_t type, int64_t row, int64_t tokens, + int64_t expert_stride, int64_t row_stride, const std::string& id_space, + torch::Tensor output, torch::Tensor quant_X) { + return ggml_moe_mmvq_id(X, W, ids, top_k, type, row, tokens, + expert_stride, row_stride, id_space, output, quant_X); + }, ""); + m.def("ggml_moe_mmvdq_id", + [](torch::Tensor X, torch::Tensor W, torch::Tensor ids, + int64_t top_k, int64_t type, int64_t row, int64_t tokens, + int64_t expert_stride, int64_t row_stride, const std::string& id_space) { + return ggml_moe_mmvdq_id(X, W, ids, top_k, type, row, tokens, + expert_stride, row_stride, id_space); + }, ""); + m.def("ggml_moe_mmvdq_id_workspace", + [](torch::Tensor X, torch::Tensor W, torch::Tensor ids, + int64_t top_k, int64_t type, int64_t row, int64_t tokens, + int64_t expert_stride, int64_t row_stride, const std::string& id_space, + torch::Tensor output) { + return ggml_moe_mmvdq_id(X, W, ids, top_k, type, row, tokens, + expert_stride, row_stride, id_space, output); + }, ""); + m.def("ggml_moe_gate_up_swiglu_id", + [](torch::Tensor X, torch::Tensor W, torch::Tensor ids, + int64_t top_k, int64_t nrows, int64_t tokens, + int64_t expert_stride, int64_t row_stride, const std::string& id_space) { + return ggml_moe_gate_up_swiglu_id(X, W, ids, top_k, nrows, tokens, + expert_stride, row_stride, id_space); + }, ""); + m.def("ggml_moe_gate_up_swiglu_id_workspace", + [](torch::Tensor X, torch::Tensor W, torch::Tensor ids, + int64_t top_k, int64_t nrows, int64_t tokens, + int64_t expert_stride, int64_t row_stride, const std::string& id_space, + torch::Tensor output, torch::Tensor quant_X) { + return ggml_moe_gate_up_swiglu_id(X, W, ids, top_k, nrows, tokens, + expert_stride, row_stride, id_space, + output, quant_X); + }, ""); +} + +#ifndef FREETOKEN_GGUF_NO_PYBIND +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + bind_gguf_moe_gfx1100(m); +} +#endif diff --git a/python/freetoken/kernel/csrc/gguf/llama_b10434/PATCHES.md b/python/freetoken/kernel/csrc/gguf/llama_b10434/PATCHES.md new file mode 100644 index 000000000..99ea62f00 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/llama_b10434/PATCHES.md @@ -0,0 +1,8 @@ +# Local b10434 delta + +- HIP portability wrappers use `USE_ROCM` and the existing FreeToken GGUF dispatch headers. +- `ggml_cuda_pdl_sync`, `ggml_cuda_pdl_lc`, and launch-attribute helpers are stubbed to no-ops; + stream-ordered launches make this safe for the caller-owned workspace. No PDL mapping is + claimed until graph capture/replay evidence exists. +- `small_k` template variant is omitted. RDNA3 dispatch forces `use_small_k=false`. +- No `mul_mat_vec_q_moe` or MMID compaction is ported. diff --git a/python/freetoken/kernel/csrc/gguf/llama_b10434/README.md b/python/freetoken/kernel/csrc/gguf/llama_b10434/README.md new file mode 100644 index 000000000..eb3c5514a --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/llama_b10434/README.md @@ -0,0 +1,19 @@ +# llama.cpp b10434 single-token source contract + +Pinned source commit: + +```text +7e4c0a96880dae4fc4268ad441f8a6446bd5460a +``` + +This directory records the narrow single-token ABI used by FreeToken. Active quant block and +vec-dot declarations live in the sibling `ggml-common*`, `vecdotq*`, and `moe_vec*` headers so +CUDA/HIP builds share one source surface. Multi-token MMID compaction and `small_k` are excluded. + +ABI choices: + +- Q8_1 activation rows: `ceil(H / 32) * 36` bytes, unpadded contract requires `H % 32 == 0`. +- Output: caller-owned FP32 `[channels, Nrows]`. +- Every workspace region is aligned to 256 bytes. +- PDL calls are no-ops; stream order already serializes the caller-owned workspace during graph + replay. Native PDL mapping requires a later capture/replay proof. diff --git a/python/freetoken/kernel/csrc/gguf/llama_b10434/quantize_q8_1.h b/python/freetoken/kernel/csrc/gguf/llama_b10434/quantize_q8_1.h new file mode 100644 index 000000000..a073a8cd0 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/llama_b10434/quantize_q8_1.h @@ -0,0 +1,13 @@ +#pragma once + +#include + +// Layout consumed by the b10434 single-token ABI. ``qs`` is signed Q8_1 payload; +// ``d`` is the scale and ``s`` is the activation sum. +struct freetoken_block_q8_1 { + uint16_t d; + uint16_t s; + int8_t qs[32]; +}; + +static_assert(sizeof(freetoken_block_q8_1) == 36, "Q8_1 ABI changed"); diff --git a/python/freetoken/kernel/csrc/gguf/mmq_hip.cuh b/python/freetoken/kernel/csrc/gguf/mmq_hip.cuh new file mode 100644 index 000000000..094307691 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/mmq_hip.cuh @@ -0,0 +1,883 @@ +// !!! This is a file automatically generated by hipify!!! +#include "hip/hip_runtime.h" +// copied from +// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/mmq.cuh +// copied from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmq.cu +template < + typename scalar_t, + int qk, + int qr, + int qi, + bool need_sum, + typename block_q_t, + int mmq_x, + int mmq_y, + int nwarps, + allocate_tiles_cuda_t allocate_tiles, + load_tiles_cuda_t load_tiles, + int vdr, + vec_dot_q_mul_mat_cuda_t vec_dot> +static __device__ __forceinline__ void mul_mat_q( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const block_q_t* x = (const block_q_t*)vx; + const block_q8_1* y = (const block_q8_1*)vy; + + const int blocks_per_row_x = ncols_x / qk; + const int blocks_per_col_y = nrows_y / QK8_1; + const int blocks_per_warp = WARP_SIZE_GGUF / qi; + + const int& ncols_dst = ncols_y; + + const auto row_dst_0 = blockIdx.x * mmq_y; + const int& row_x_0 = row_dst_0; + + const auto col_dst_0 = blockIdx.y * mmq_x; + const int& col_y_0 = col_dst_0; + + int* tile_x_ql = nullptr; + half2* tile_x_dm = nullptr; + int* tile_x_qh = nullptr; + int* tile_x_sc = nullptr; + + allocate_tiles(&tile_x_ql, &tile_x_dm, &tile_x_qh, &tile_x_sc); + + __shared__ int tile_y_qs[mmq_x * WARP_SIZE_GGUF]; + __shared__ half2 tile_y_ds[mmq_x * WARP_SIZE_GGUF / QI8_1]; + + float sum[mmq_y / WARP_SIZE_GGUF][mmq_x / nwarps] = {{0.0f}}; + + for (int ib0 = 0; ib0 < blocks_per_row_x; ib0 += blocks_per_warp) { + load_tiles( + x + row_x_0 * blocks_per_row_x + ib0, + tile_x_ql, + tile_x_dm, + tile_x_qh, + tile_x_sc, + threadIdx.y, + nrows_x - row_x_0 - 1, + threadIdx.x, + blocks_per_row_x); + +#pragma unroll + for (int ir = 0; ir < qr && ib0 + ir * blocks_per_warp / qr < blocks_per_row_x; ++ir) { + const auto kqs = ir * WARP_SIZE_GGUF + threadIdx.x; + const int kbxd = kqs / QI8_1; + +#pragma unroll + for (int i = 0; i < mmq_x; i += nwarps) { + const int col_y_eff = min(col_y_0 + threadIdx.y + i, ncols_y - 1); // to prevent out-of-bounds memory accesses + const block_q8_1* by0 = &y[col_y_eff * blocks_per_col_y + ib0 * (qk / QK8_1) + kbxd]; + const int index_y = (threadIdx.y + i) * WARP_SIZE_GGUF + kqs % WARP_SIZE_GGUF; + tile_y_qs[index_y] = get_int_from_int8_aligned(by0->qs, threadIdx.x % QI8_1); + } + +#pragma unroll + for (int ids0 = 0; ids0 < mmq_x; ids0 += nwarps * QI8_1) { + const int ids = (ids0 + threadIdx.y * QI8_1 + threadIdx.x / (WARP_SIZE_GGUF / QI8_1)) % mmq_x; + const auto kby = threadIdx.x % (WARP_SIZE_GGUF / QI8_1); + const int col_y_eff = min(col_y_0 + ids, ncols_y - 1); + + // if the sum is not needed it's faster to transform the scale to f32 ahead of time + const half2* dsi_src = + &y[col_y_eff * blocks_per_col_y + ib0 * (qk / QK8_1) + ir * (WARP_SIZE_GGUF / QI8_1) + kby].ds; + half2* dsi_dst = &tile_y_ds[ids * (WARP_SIZE_GGUF / QI8_1) + kby]; + if (need_sum) { + *dsi_dst = *dsi_src; + } else { + float* dfi_dst = (float*)dsi_dst; + *dfi_dst = __low2float(*dsi_src); + } + } + + __syncthreads(); + + // #pragma unroll // unrolling this loop causes too much register pressure + for (int k = ir * WARP_SIZE_GGUF / qr; k < (ir + 1) * WARP_SIZE_GGUF / qr; k += vdr) { +#pragma unroll + for (int j = 0; j < mmq_x; j += nwarps) { +#pragma unroll + for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) { + sum[i / WARP_SIZE_GGUF][j / nwarps] += vec_dot( + tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc, tile_y_qs, tile_y_ds, threadIdx.x + i, threadIdx.y + j, k); + } + } + } + __syncthreads(); + } + } + +#pragma unroll + for (int j = 0; j < mmq_x; j += nwarps) { + const auto col_dst = col_dst_0 + j + threadIdx.y; + if (col_dst >= ncols_dst) { + return; + } + +#pragma unroll + for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) { + const auto row_dst = row_dst_0 + threadIdx.x + i; + if (row_dst >= nrows_dst) { + continue; + } + dst[col_dst * nrows_dst + row_dst] = sum[i / WARP_SIZE_GGUF][j / nwarps]; + } + } +} + +#if defined(USE_ROCM) +#define MMQ_X_Q4_0 64 +#define MMQ_Y_Q4_0 128 +#define NWARPS_Q4_0 8 +#else +#define MMQ_X_Q4_0 4 +#define MMQ_Y_Q4_0 32 +#define NWARPS_Q4_0 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_0, 2) +#endif + mul_mat_q4_0( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_Q4_0; + const int mmq_y = MMQ_Y_Q4_0; + const int nwarps = NWARPS_Q4_0; + + mul_mat_q< + scalar_t, + QK4_0, + QR4_0, + QI4_0, + true, + block_q4_0, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q4_0, + load_tiles_q4_0, + VDR_Q4_0_Q8_1_MMQ, + vec_dot_q4_0_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_q4_0_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + hipStream_t stream) { + int mmq_x = MMQ_X_Q4_0; + int mmq_y = MMQ_Y_Q4_0; + int nwarps = NWARPS_Q4_0; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + hipLaunchKernelGGL(( mul_mat_q4_0) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + hipLaunchKernelGGL(( mul_mat_q4_0) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_Q4_1 64 +#define MMQ_Y_Q4_1 128 +#define NWARPS_Q4_1 8 +#else +#define MMQ_X_Q4_1 4 +#define MMQ_Y_Q4_1 32 +#define NWARPS_Q4_1 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_1, 2) +#endif + mul_mat_q4_1( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_Q4_1; + const int mmq_y = MMQ_Y_Q4_1; + const int nwarps = NWARPS_Q4_1; + + mul_mat_q< + scalar_t, + QK4_1, + QR4_1, + QI4_1, + true, + block_q4_1, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q4_1, + load_tiles_q4_1, + VDR_Q4_1_Q8_1_MMQ, + vec_dot_q4_1_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_q4_1_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + hipStream_t stream) { + int mmq_x = MMQ_X_Q4_1; + int mmq_y = MMQ_Y_Q4_1; + int nwarps = NWARPS_Q4_1; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + hipLaunchKernelGGL(( mul_mat_q4_1) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + hipLaunchKernelGGL(( mul_mat_q4_1) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_Q5_0 64 +#define MMQ_Y_Q5_0 128 +#define NWARPS_Q5_0 8 +#else +#define MMQ_X_Q5_0 4 +#define MMQ_Y_Q5_0 32 +#define NWARPS_Q5_0 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_0, 2) +#endif + mul_mat_q5_0( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_Q5_0; + const int mmq_y = MMQ_Y_Q5_0; + const int nwarps = NWARPS_Q5_0; + + mul_mat_q< + scalar_t, + QK5_0, + QR5_0, + QI5_0, + false, + block_q5_0, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q5_0, + load_tiles_q5_0, + VDR_Q5_0_Q8_1_MMQ, + vec_dot_q5_0_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_q5_0_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + hipStream_t stream) { + const int mmq_x = MMQ_X_Q5_0; + const int mmq_y = MMQ_Y_Q5_0; + const int nwarps = NWARPS_Q5_0; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + hipLaunchKernelGGL(( mul_mat_q5_0) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + hipLaunchKernelGGL(( mul_mat_q5_0) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_Q5_1 64 +#define MMQ_Y_Q5_1 128 +#define NWARPS_Q5_1 8 +#else +#define MMQ_X_Q5_1 4 +#define MMQ_Y_Q5_1 32 +#define NWARPS_Q5_1 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_1, 2) +#endif + mul_mat_q5_1( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_Q5_1; + const int mmq_y = MMQ_Y_Q5_1; + const int nwarps = NWARPS_Q5_1; + + mul_mat_q< + scalar_t, + QK5_1, + QR5_1, + QI5_1, + true, + block_q5_1, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q5_1, + load_tiles_q5_1, + VDR_Q5_1_Q8_1_MMQ, + vec_dot_q5_1_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_q5_1_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + hipStream_t stream) { + const int mmq_x = MMQ_X_Q5_1; + const int mmq_y = MMQ_Y_Q5_1; + const int nwarps = NWARPS_Q5_1; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + hipLaunchKernelGGL(( mul_mat_q5_1) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + hipLaunchKernelGGL(( mul_mat_q5_1) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_Q8_0 64 +#define MMQ_Y_Q8_0 128 +#define NWARPS_Q8_0 8 +#else +#define MMQ_X_Q8_0 4 +#define MMQ_Y_Q8_0 32 +#define NWARPS_Q8_0 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q8_0, 2) +#endif + mul_mat_q8_0( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_Q8_0; + const int mmq_y = MMQ_Y_Q8_0; + const int nwarps = NWARPS_Q8_0; + + mul_mat_q< + scalar_t, + QK8_0, + QR8_0, + QI8_0, + false, + block_q8_0, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q8_0, + load_tiles_q8_0, + VDR_Q8_0_Q8_1_MMQ, + vec_dot_q8_0_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_q8_0_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + hipStream_t stream) { + const int mmq_x = MMQ_X_Q8_0; + const int mmq_y = MMQ_Y_Q8_0; + const int nwarps = NWARPS_Q8_0; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + hipLaunchKernelGGL(( mul_mat_q8_0) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + hipLaunchKernelGGL(( mul_mat_q8_0) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_Q2_K 64 +#define MMQ_Y_Q2_K 128 +#define NWARPS_Q2_K 8 +#else +#define MMQ_X_Q2_K 4 +#define MMQ_Y_Q2_K 32 +#define NWARPS_Q2_K 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q2_K, 2) +#endif + mul_mat_q2_K( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_Q2_K; + const int mmq_y = MMQ_Y_Q2_K; + const int nwarps = NWARPS_Q2_K; + + mul_mat_q< + scalar_t, + QK_K, + QR2_K, + QI2_K, + false, + block_q2_K, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q2_K, + load_tiles_q2_K, + VDR_Q2_K_Q8_1_MMQ, + vec_dot_q2_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_q2_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + hipStream_t stream) { + const int mmq_x = MMQ_X_Q2_K; + const int mmq_y = MMQ_Y_Q2_K; + const int nwarps = NWARPS_Q2_K; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + hipLaunchKernelGGL(( mul_mat_q2_K) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + hipLaunchKernelGGL(( mul_mat_q2_K) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_Q3_K 64 +#define MMQ_Y_Q3_K 128 +#define NWARPS_Q3_K 8 +#else +#define MMQ_X_Q3_K 4 +#define MMQ_Y_Q3_K 32 +#define NWARPS_Q3_K 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q3_K, 2) +#endif + mul_mat_q3_K( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + + const int mmq_x = MMQ_X_Q3_K; + const int mmq_y = MMQ_Y_Q3_K; + const int nwarps = NWARPS_Q3_K; + + mul_mat_q< + scalar_t, + QK_K, + QR3_K, + QI3_K, + false, + block_q3_K, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q3_K, + load_tiles_q3_K, + VDR_Q3_K_Q8_1_MMQ, + vec_dot_q3_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_q3_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + hipStream_t stream) { + const int mmq_x = MMQ_X_Q3_K; + const int mmq_y = MMQ_Y_Q3_K; + const int nwarps = NWARPS_Q3_K; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + hipLaunchKernelGGL(( mul_mat_q3_K) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + hipLaunchKernelGGL(( mul_mat_q3_K) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_Q4_K 64 +#define MMQ_Y_Q4_K 128 +#define NWARPS_Q4_K 8 +#else +#define MMQ_X_Q4_K 4 +#define MMQ_Y_Q4_K 32 +#define NWARPS_Q4_K 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_K, 2) +#endif + mul_mat_q4_K( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_Q4_K; + const int mmq_y = MMQ_Y_Q4_K; + const int nwarps = NWARPS_Q4_K; + + mul_mat_q< + scalar_t, + QK_K, + QR4_K, + QI4_K, + true, + block_q4_K, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q4_K, + load_tiles_q4_K, + VDR_Q4_K_Q8_1_MMQ, + vec_dot_q4_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_q4_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + hipStream_t stream) { + const int mmq_x = MMQ_X_Q4_K; + const int mmq_y = MMQ_Y_Q4_K; + const int nwarps = NWARPS_Q4_K; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + hipLaunchKernelGGL(( mul_mat_q4_K) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + hipLaunchKernelGGL(( mul_mat_q4_K) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_Q5_K 64 +#define MMQ_Y_Q5_K 128 +#define NWARPS_Q5_K 8 +#else +#define MMQ_X_Q5_K 4 +#define MMQ_Y_Q5_K 32 +#define NWARPS_Q5_K 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_K, 2) +#endif + mul_mat_q5_K( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_Q5_K; + const int mmq_y = MMQ_Y_Q5_K; + const int nwarps = NWARPS_Q5_K; + + mul_mat_q< + scalar_t, + QK_K, + QR5_K, + QI5_K, + true, + block_q5_K, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q5_K, + load_tiles_q5_K, + VDR_Q5_K_Q8_1_MMQ, + vec_dot_q5_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_q5_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + hipStream_t stream) { + const int mmq_x = MMQ_X_Q5_K; + const int mmq_y = MMQ_Y_Q5_K; + const int nwarps = NWARPS_Q5_K; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + hipLaunchKernelGGL(( mul_mat_q5_K) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + hipLaunchKernelGGL(( mul_mat_q5_K) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} + +#if defined(USE_ROCM) +#define MMQ_X_Q6_K 64 +#define MMQ_Y_Q6_K 128 +#define NWARPS_Q6_K 8 +#else +#define MMQ_X_Q6_K 4 +#define MMQ_Y_Q6_K 32 +#define NWARPS_Q6_K 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q6_K, 2) +#endif + mul_mat_q6_K( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst) { + const int mmq_x = MMQ_X_Q6_K; + const int mmq_y = MMQ_Y_Q6_K; + const int nwarps = NWARPS_Q6_K; + + mul_mat_q< + scalar_t, + QK_K, + QR6_K, + QI6_K, + false, + block_q6_K, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q6_K, + load_tiles_q6_K, + VDR_Q6_K_Q8_1_MMQ, + vec_dot_q6_K_q8_1_mul_mat>(vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +template +static void ggml_mul_mat_q6_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + hipStream_t stream) { + const int mmq_x = MMQ_X_Q6_K; + const int mmq_y = MMQ_Y_Q6_K; + const int nwarps = NWARPS_Q6_K; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (ncols_y + mmq_x - 1) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + const bool need_check = false; + hipLaunchKernelGGL(( mul_mat_q6_K) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } else { + const bool need_check = true; + hipLaunchKernelGGL(( mul_mat_q6_K) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); + } +} diff --git a/python/freetoken/kernel/csrc/gguf/mmvq_hip.cuh b/python/freetoken/kernel/csrc/gguf/mmvq_hip.cuh new file mode 100644 index 000000000..d8c74e410 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/mmvq_hip.cuh @@ -0,0 +1,354 @@ +// !!! This is a file automatically generated by hipify!!! +#include "hip/hip_runtime.h" +// copied from +// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/mmvq.cuh +// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu +template +static __global__ void mul_mat_vec_q( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int ncols, + const int nrows, + const int nvecs) { + const auto row = blockIdx.x * blockDim.y + threadIdx.y; + const auto vec = blockIdx.y; + + if (row >= nrows || vec >= nvecs) { + return; + } + + const int blocks_per_row = ncols / qk; + const int blocks_per_warp = vdr * WARP_SIZE / qi; + const int nrows_y = (ncols + 512 - 1) / 512 * 512; + + // partial sum for each thread + float tmp = 0.0f; + + const block_q_t* x = (const block_q_t*)vx; + const block_q8_1* y = (const block_q8_1*)vy; + + for (auto i = threadIdx.x / (qi / vdr); i < blocks_per_row; i += blocks_per_warp) { + const int ibx = row * blocks_per_row + i; // x block index + + const int iby = vec * (nrows_y / QK8_1) + i * (qk / QK8_1); // y block index that aligns with ibx + + const int iqs = vdr * (threadIdx.x % (qi / vdr)); // x block quant index when casting the quants to int + + tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs); + } + + // sum up partial sums and write back result +#pragma unroll + for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { + tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask); + } + + if (threadIdx.x == 0) { + dst[vec * nrows + row] = tmp; + } +} + +template +static void mul_mat_vec_q4_0_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_q4_1_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_q5_0_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_q5_1_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_q8_0_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_q2_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_q3_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_q4_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_q5_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_q6_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_iq2_xxs_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_iq2_xs_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_iq2_s_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_iq3_xxs_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_iq1_s_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_iq1_m_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_iq4_nl_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_iq4_xs_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} + +template +static void mul_mat_vec_iq3_s_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int ncols, + const int nrows, + const int nvecs, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, nvecs, 1); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( mul_mat_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, ncols, nrows, nvecs); +} diff --git a/python/freetoken/kernel/csrc/gguf/moe_hip.cuh b/python/freetoken/kernel/csrc/gguf/moe_hip.cuh new file mode 100644 index 000000000..0db30c76c --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/moe_hip.cuh @@ -0,0 +1,1381 @@ +// !!! This is a file automatically generated by hipify!!! +#include "hip/hip_runtime.h" +// copied from +// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/moe.cuh +#include + +/* Adapted from ./csrc/quantization/gguf/mmq.cuh + */ +template < + typename scalar_t, + int qk, + int qr, + int qi, + bool need_sum, + typename block_q_t, + int mmq_x, + int mmq_y, + int nwarps, + allocate_tiles_cuda_t allocate_tiles, + load_tiles_cuda_t load_tiles, + int vdr, + vec_dot_q_mul_mat_cuda_t vec_dot> +static __device__ __forceinline__ void moe_q( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* __restrict__ sorted_token_ids, + const int* __restrict__ expert_ids, + const int* __restrict__ num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int blocks_per_row_x = ncols_x / qk; + const int blocks_per_col_y = nrows_y / QK8_1; + const int blocks_per_warp = WARP_SIZE_GGUF / qi; + + const int ncols_dst = ncols_y * top_k; + + const auto row_dst_0 = blockIdx.x * mmq_y; + const int& row_x_0 = row_dst_0; + + const auto col_dst_0 = blockIdx.y * mmq_x; + + int token_offs[mmq_x / nwarps]; + for (int i = 0; i < mmq_x; i += nwarps) { + token_offs[i / nwarps] = sorted_token_ids[col_dst_0 + threadIdx.y + i]; + } + + const int exp_idx = expert_ids[blockIdx.y]; + if (exp_idx > 255 || exp_idx < 0) return; + if (blockIdx.y * mmq_x > num_tokens_post_padded[0]) return; + + const block_q_t* x = (const block_q_t*)((char*)vx + exp_idx * exp_stride); + const block_q8_1* y = (const block_q8_1*)(vy); + + int* tile_x_ql = nullptr; + half2* tile_x_dm = nullptr; + int* tile_x_qh = nullptr; + int* tile_x_sc = nullptr; + + allocate_tiles(&tile_x_ql, &tile_x_dm, &tile_x_qh, &tile_x_sc); + + __shared__ int tile_y_qs[mmq_x * WARP_SIZE_GGUF]; + __shared__ half2 tile_y_ds[mmq_x * WARP_SIZE_GGUF / QI8_1]; + + float sum[mmq_y / WARP_SIZE_GGUF][mmq_x / nwarps] = {{0.0f}}; + + for (int ib0 = 0; ib0 < blocks_per_row_x; ib0 += blocks_per_warp) { + load_tiles( + x + row_x_0 * blocks_per_row_x + ib0, + tile_x_ql, + tile_x_dm, + tile_x_qh, + tile_x_sc, + threadIdx.y, + nrows_x - row_x_0 - 1, + threadIdx.x, + blocks_per_row_x); + + const int n_per_r = ((qk * blocks_per_warp) / qr); +#pragma unroll + for (int ir = 0; ir < qr && ib0 * qk + ir * n_per_r < ncols_x; ++ir) { + const auto kqs = ir * WARP_SIZE_GGUF + threadIdx.x; + const int kbxd = kqs / QI8_1; + +#pragma unroll + for (int i = 0; i < mmq_x; i += nwarps) { + const int col_y_eff = token_offs[i / nwarps] / top_k; + const int block_x = ib0 * (qk / QK8_1) + kbxd; + if (col_y_eff < ncols_y && block_x < blocks_per_col_y) { + const block_q8_1* by0 = &y[col_y_eff * blocks_per_col_y + block_x]; + const int index_y = (threadIdx.y + i) * WARP_SIZE_GGUF + kqs % WARP_SIZE_GGUF; + tile_y_qs[index_y] = get_int_from_int8_aligned(by0->qs, threadIdx.x % QI8_1); + } + } + + if (threadIdx.x < n_per_r / QK8_1) { + const auto kby = threadIdx.x % (WARP_SIZE_GGUF / QI8_1); + const int col_y_eff = token_offs[threadIdx.y] / top_k; + const int block_x = ib0 * (qk / QK8_1) + ir * (WARP_SIZE_GGUF / QI8_1) + kby; + + if (col_y_eff < ncols_y && block_x < blocks_per_col_y) { + const half2* dsi_src = &y[col_y_eff * blocks_per_col_y + block_x].ds; + half2* dsi_dst = &tile_y_ds[threadIdx.y * (WARP_SIZE_GGUF / QI8_1) + kby]; + + if (need_sum) { + *dsi_dst = *dsi_src; + } else { + float* dfi_dst = (float*)dsi_dst; + *dfi_dst = __low2float(*dsi_src); + } + } + } + __syncthreads(); + + // #pragma unroll // unrolling this loop causes too much register pressure + for (int k = ir * WARP_SIZE_GGUF / qr; k < (ir + 1) * WARP_SIZE_GGUF / qr; k += vdr) { +#pragma unroll + for (int j = 0; j < mmq_x; j += nwarps) { +#pragma unroll + for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) { + sum[i / WARP_SIZE_GGUF][j / nwarps] += vec_dot( + tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc, tile_y_qs, tile_y_ds, threadIdx.x + i, threadIdx.y + j, k); + } + } + } + __syncthreads(); + } + } + +#pragma unroll + for (int j = 0; j < mmq_x; j += nwarps) { + const int col_dst = token_offs[j / nwarps]; + if (col_dst >= ncols_dst) { + return; + } + +#pragma unroll + for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) { + const auto row_dst = row_dst_0 + threadIdx.x + i; + if (row_dst >= nrows_dst) { + continue; + } + dst[col_dst * nrows_dst + row_dst] = sum[i / WARP_SIZE_GGUF][j / nwarps]; + } + } +} + +#if defined(USE_ROCM) +#define MOE_X_Q4_0 8 +#define MOE_Y_Q4_0 128 +#define NWARPS_Q4_0 8 +#else +#define MOE_X_Q4_0 4 +#define MOE_Y_Q4_0 32 +#define NWARPS_Q4_0 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_0, 2) +#endif + moe_q4_0( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_Q4_0; + const int mmq_y = MOE_Y_Q4_0; + const int nwarps = NWARPS_Q4_0; + + moe_q< + scalar_t, + QK4_0, + QR4_0, + QI4_0, + true, + block_q4_0, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q4_0, + load_tiles_q4_0, + VDR_Q4_0_Q8_1_MMQ, + vec_dot_q4_0_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_q4_0_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + hipStream_t stream) { + int mmq_x = MOE_X_Q4_0; + int mmq_y = MOE_Y_Q4_0; + int nwarps = NWARPS_Q4_0; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (tokens_post_padded) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + hipLaunchKernelGGL(( moe_q4_0), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + hipLaunchKernelGGL(( moe_q4_0), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_Q4_1 8 +#define MOE_Y_Q4_1 128 +#define NWARPS_Q4_1 8 +#else +#define MOE_X_Q4_1 4 +#define MOE_Y_Q4_1 32 +#define NWARPS_Q4_1 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_1, 2) +#endif + moe_q4_1( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_Q4_1; + const int mmq_y = MOE_Y_Q4_1; + const int nwarps = NWARPS_Q4_1; + + moe_q< + scalar_t, + QK4_1, + QR4_1, + QI4_1, + true, + block_q4_1, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q4_1, + load_tiles_q4_1, + VDR_Q4_1_Q8_1_MMQ, + vec_dot_q4_1_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_q4_1_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + hipStream_t stream) { + int mmq_x = MOE_X_Q4_1; + int mmq_y = MOE_Y_Q4_1; + int nwarps = NWARPS_Q4_1; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (tokens_post_padded) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + hipLaunchKernelGGL(( moe_q4_1), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + hipLaunchKernelGGL(( moe_q4_1), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_Q5_0 8 +#define MOE_Y_Q5_0 128 +#define NWARPS_Q5_0 8 +#else +#define MOE_X_Q5_0 4 +#define MOE_Y_Q5_0 32 +#define NWARPS_Q5_0 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_0, 2) +#endif + moe_q5_0( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_Q5_0; + const int mmq_y = MOE_Y_Q5_0; + const int nwarps = NWARPS_Q5_0; + + moe_q< + scalar_t, + QK5_0, + QR5_0, + QI5_0, + false, + block_q5_0, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q5_0, + load_tiles_q5_0, + VDR_Q5_0_Q8_1_MMQ, + vec_dot_q5_0_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_q5_0_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + hipStream_t stream) { + const int mmq_x = MOE_X_Q5_0; + const int mmq_y = MOE_Y_Q5_0; + const int nwarps = NWARPS_Q5_0; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (tokens_post_padded) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + hipLaunchKernelGGL(( moe_q5_0), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + hipLaunchKernelGGL(( moe_q5_0), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_Q5_1 8 +#define MOE_Y_Q5_1 128 +#define NWARPS_Q5_1 8 +#else +#define MOE_X_Q5_1 4 +#define MOE_Y_Q5_1 32 +#define NWARPS_Q5_1 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_1, 2) +#endif + moe_q5_1( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_Q5_1; + const int mmq_y = MOE_Y_Q5_1; + const int nwarps = NWARPS_Q5_1; + + moe_q< + scalar_t, + QK5_1, + QR5_1, + QI5_1, + true, + block_q5_1, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q5_1, + load_tiles_q5_1, + VDR_Q5_1_Q8_1_MMQ, + vec_dot_q5_1_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_q5_1_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + hipStream_t stream) { + const int mmq_x = MOE_X_Q5_1; + const int mmq_y = MOE_Y_Q5_1; + const int nwarps = NWARPS_Q5_1; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (tokens_post_padded) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + hipLaunchKernelGGL(( moe_q5_1), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + hipLaunchKernelGGL(( moe_q5_1), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_Q8_0 8 +#define MOE_Y_Q8_0 128 +#define NWARPS_Q8_0 8 +#else +#define MOE_X_Q8_0 4 +#define MOE_Y_Q8_0 32 +#define NWARPS_Q8_0 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q8_0, 2) +#endif + moe_q8_0( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_Q8_0; + const int mmq_y = MOE_Y_Q8_0; + const int nwarps = NWARPS_Q8_0; + + moe_q< + scalar_t, + QK8_0, + QR8_0, + QI8_0, + false, + block_q8_0, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q8_0, + load_tiles_q8_0, + VDR_Q8_0_Q8_1_MMQ, + vec_dot_q8_0_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_q8_0_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + hipStream_t stream) { + const int mmq_x = MOE_X_Q8_0; + const int mmq_y = MOE_Y_Q8_0; + const int nwarps = NWARPS_Q8_0; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (tokens_post_padded) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + hipLaunchKernelGGL(( moe_q8_0), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + hipLaunchKernelGGL(( moe_q8_0), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_Q2_K 8 +#define MOE_Y_Q2_K 128 +#define NWARPS_Q2_K 8 +#else +#define MOE_X_Q2_K 4 +#define MOE_Y_Q2_K 32 +#define NWARPS_Q2_K 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q2_K, 2) +#endif + moe_q2_K( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_Q2_K; + const int mmq_y = MOE_Y_Q2_K; + const int nwarps = NWARPS_Q2_K; + + moe_q< + scalar_t, + QK_K, + QR2_K, + QI2_K, + false, + block_q2_K, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q2_K, + load_tiles_q2_K, + VDR_Q2_K_Q8_1_MMQ, + vec_dot_q2_K_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_q2_K_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + hipStream_t stream) { + const int mmq_x = MOE_X_Q2_K; + const int mmq_y = MOE_Y_Q2_K; + const int nwarps = NWARPS_Q2_K; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (tokens_post_padded) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + hipLaunchKernelGGL(( moe_q2_K), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + hipLaunchKernelGGL(( moe_q2_K), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_Q3_K 8 +#define MOE_Y_Q3_K 128 +#define NWARPS_Q3_K 8 +#else +#define MOE_X_Q3_K 4 +#define MOE_Y_Q3_K 32 +#define NWARPS_Q3_K 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q3_K, 2) +#endif + moe_q3_K( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + + const int mmq_x = MOE_X_Q3_K; + const int mmq_y = MOE_Y_Q3_K; + const int nwarps = NWARPS_Q3_K; + + moe_q< + scalar_t, + QK_K, + QR3_K, + QI3_K, + false, + block_q3_K, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q3_K, + load_tiles_q3_K, + VDR_Q3_K_Q8_1_MMQ, + vec_dot_q3_K_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} +template +static void ggml_moe_q3_K_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + hipStream_t stream) { + const int mmq_x = MOE_X_Q3_K; + const int mmq_y = MOE_Y_Q3_K; + const int nwarps = NWARPS_Q3_K; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (tokens_post_padded) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + hipLaunchKernelGGL(( moe_q3_K), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + hipLaunchKernelGGL(( moe_q3_K), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_Q4_K 8 +#define MOE_Y_Q4_K 128 +#define NWARPS_Q4_K 8 +#else +#define MOE_X_Q4_K 4 +#define MOE_Y_Q4_K 32 +#define NWARPS_Q4_K 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_K, 2) +#endif + moe_q4_K( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_Q4_K; + const int mmq_y = MOE_Y_Q4_K; + const int nwarps = NWARPS_Q4_K; + + moe_q< + scalar_t, + QK_K, + QR4_K, + QI4_K, + true, + block_q4_K, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q4_K, + load_tiles_q4_K, + VDR_Q4_K_Q8_1_MMQ, + vec_dot_q4_K_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_q4_K_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + hipStream_t stream) { + const int mmq_x = MOE_X_Q4_K; + const int mmq_y = MOE_Y_Q4_K; + const int nwarps = NWARPS_Q4_K; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (tokens_post_padded) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + hipLaunchKernelGGL(( moe_q4_K), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + hipLaunchKernelGGL(( moe_q4_K), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_Q5_K 8 +#define MOE_Y_Q5_K 128 +#define NWARPS_Q5_K 8 +#else +#define MOE_X_Q5_K 4 +#define MOE_Y_Q5_K 32 +#define NWARPS_Q5_K 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_K, 2) +#endif + moe_q5_K( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_Q5_K; + const int mmq_y = MOE_Y_Q5_K; + const int nwarps = NWARPS_Q5_K; + + moe_q< + scalar_t, + QK_K, + QR5_K, + QI5_K, + true, + block_q5_K, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q5_K, + load_tiles_q5_K, + VDR_Q5_K_Q8_1_MMQ, + vec_dot_q5_K_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_q5_K_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + hipStream_t stream) { + const int mmq_x = MOE_X_Q5_K; + const int mmq_y = MOE_Y_Q5_K; + const int nwarps = NWARPS_Q5_K; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (tokens_post_padded) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + hipLaunchKernelGGL(( moe_q5_K), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + hipLaunchKernelGGL(( moe_q5_K), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} + +#if defined(USE_ROCM) +#define MOE_X_Q6_K 8 +#define MOE_Y_Q6_K 128 +#define NWARPS_Q6_K 8 +#else +#define MOE_X_Q6_K 4 +#define MOE_Y_Q6_K 32 +#define NWARPS_Q6_K 4 +#endif + +template +static __global__ void +#if defined(USE_ROCM) +__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q6_K, 2) +#endif + moe_q6_K( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k) { + const int mmq_x = MOE_X_Q6_K; + const int mmq_y = MOE_Y_Q6_K; + const int nwarps = NWARPS_Q6_K; + + moe_q< + scalar_t, + QK_K, + QR6_K, + QI6_K, + false, + block_q6_K, + mmq_x, + mmq_y, + nwarps, + allocate_tiles_q6_K, + load_tiles_q6_K, + VDR_Q6_K_Q8_1_MMQ, + vec_dot_q6_K_q8_1_mul_mat>( + vx, + vy, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); +} + +template +static void ggml_moe_q6_K_q8_1_cuda( + const void* inp, + const void* w, + scalar_t* dst, + const int* sorted_token_ids, + const int* expert_ids, + const int* num_tokens_post_padded, + const int exp_stride, + const int ncols_x, + const int nrows_x, + const int ncols_y, + const int nrows_y, + const int nrows_dst, + const int top_k, + const int tokens_post_padded, + hipStream_t stream) { + const int mmq_x = MOE_X_Q6_K; + const int mmq_y = MOE_Y_Q6_K; + const int nwarps = NWARPS_Q6_K; + + const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y; + const int block_num_y = (tokens_post_padded) / mmq_x; + const dim3 block_nums(block_num_x, block_num_y, 1); + const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1); + + if (nrows_x % mmq_y == 0) { + constexpr bool need_check = false; + hipLaunchKernelGGL(( moe_q6_K), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } else { + constexpr bool need_check = true; + hipLaunchKernelGGL(( moe_q6_K), dim3(block_nums), dim3(block_dims), 0, stream, + w, + inp, + dst, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + exp_stride, + ncols_x, + nrows_x, + ncols_y, + nrows_y, + nrows_dst, + top_k); + } +} diff --git a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh index 8cef9e080..a282571d0 100644 --- a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh +++ b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh @@ -2,6 +2,9 @@ // https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/moe_vec.cuh // copied and adapted from // https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu +// Port target cross-check: llama.cpp 7e4c0a968 (b10434), +// ggml/src/ggml-cuda/mmvq.cu; ``moe_vec_q`` preserves Q4_K/Q5_K/Q6_K +// packed rows and ``moe_vec_q_strided`` adds explicit cache row strides. template static __global__ void moe_vec_q( const void* __restrict__ vx, @@ -241,6 +244,79 @@ static void moe_vec_q6_K_q8_1_cuda( <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); } +// Native mixed-Q5_K/Q6_K offload stores every down row at the Q6_K stride so +// cache banks remain shape-uniform. Q5_K rows occupy their native prefix and +// the stride-aware kernel skips the zero padding between rows/experts. +template +static __global__ void moe_vec_q_strided( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* topk_ids, + const int topk, + const int ncols, + const int nrows, + const int token_stride, + const int64_t expert_stride_bytes, + const int64_t row_stride_bytes) { + const auto row = blockIdx.x * blockDim.y + threadIdx.y; + const auto token = blockIdx.z / topk; + const auto expert = topk_ids[blockIdx.z]; + if (row >= nrows) return; + + const int blocks_per_row = ncols / qk; + const int blocks_per_warp = vdr * WARP_SIZE / qi; + float tmp = 0.0f; + const auto* expert_base = static_cast(vx) + + static_cast(expert) * expert_stride_bytes; + const block_q_t* x = reinterpret_cast( + expert_base + row * row_stride_bytes); + const block_q8_1* y = reinterpret_cast( + static_cast(vy) + token * token_stride); + for (auto i = threadIdx.x / (qi / vdr); i < blocks_per_row; i += blocks_per_warp) { + const int iby = i * (qk / QK8_1); + const int iqs = vdr * (threadIdx.x % (qi / vdr)); + tmp += vec_dot_q_cuda(&x[i], &y[iby], iqs); + } +#pragma unroll + for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) + tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask); + if (threadIdx.x == 0) dst[blockIdx.z * nrows + row] = tmp; +} + +template +static void moe_vec_q5_K_q8_1_strided_cuda( + const void* vx, const void* vy, scalar_t* dst, const int* topk_ids, + const int top_k, const int tokens, const int ncols, const int nrows, + const int token_stride, const int64_t expert_stride_bytes, + const int64_t row_stride_bytes, cudaStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + moe_vec_q_strided + <<>>( + vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, + expert_stride_bytes, row_stride_bytes); +} + +template +static void moe_vec_q6_K_q8_1_strided_cuda( + const void* vx, const void* vy, scalar_t* dst, const int* topk_ids, + const int top_k, const int tokens, const int ncols, const int nrows, + const int token_stride, const int64_t expert_stride_bytes, + const int64_t row_stride_bytes, cudaStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + moe_vec_q_strided + <<>>( + vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, + expert_stride_bytes, row_stride_bytes); +} + template static void moe_vec_iq2_xxs_q8_1_cuda( const void* vx, diff --git a/python/freetoken/kernel/csrc/gguf/moe_vec_gfx1100.cuh b/python/freetoken/kernel/csrc/gguf/moe_vec_gfx1100.cuh new file mode 100644 index 000000000..7ec641c62 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/moe_vec_gfx1100.cuh @@ -0,0 +1,140 @@ +// gfx1100 decode candidate for packed GGUF Q4_K/Q8_0 MoE vectors. +// +// This is intentionally separate from moe_vec.cuh. The incumbent assigns +// consecutive lanes to consecutive quant chunks. The candidate rotates the +// wave32 lane map by eight lanes and processes a tunable number of output rows per wave. It +// still consumes native packed rows and calls the shared quant math, so the +// comparison isolates layout/occupancy from a changed quantization contract. +#pragma once + +template +static __global__ void moe_vec_gfx1100( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* topk_ids, + const int topk, + const int ncols, + const int nrows, + const int token_stride) { + constexpr int rows_per_wave = GGML_CUDA_MMV_Y; + const int row = blockIdx.x * rows_per_wave + threadIdx.y; + const int token = blockIdx.z / topk; + const int expert = topk_ids[blockIdx.z]; + + if (row >= nrows) { + return; + } + + const int lane = static_cast(threadIdx.x) & (WARP_SIZE - 1); + const int lanes_per_chunk = qi / vdr; + const int blocks_per_row = ncols / qk; + const int blocks_per_wave = vdr * WARP_SIZE / qi; + const int block_lane = lane / lanes_per_chunk; + const int iqs = vdr * (lane % lanes_per_chunk); + const block_q_t* x = static_cast(vx) + + expert * nrows * blocks_per_row; + const block_q8_1* y = static_cast( + static_cast(static_cast(vy) + token * token_stride)); + + float tmp = 0.0f; + for (int i = block_lane; i < blocks_per_row; i += blocks_per_wave) { + const int ibx = row * blocks_per_row + i; + const int iby = i * (qk / QK8_1); + tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs); + } + +#pragma unroll + for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { + tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask); + } + + if (threadIdx.x == 0) { + dst[blockIdx.z * nrows + row] = tmp; + } +} + +template +static void moe_vec_q4_K_q8_1_gfx1100( + const void* vx, const void* vy, scalar_t* dst, const int* topk_ids, + const int top_k, const int tokens, const int ncols, const int nrows, + const int token_stride, cudaStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + moe_vec_gfx1100 + <<>>( + vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_q8_0_q8_1_gfx1100( + const void* vx, const void* vy, scalar_t* dst, const int* topk_ids, + const int top_k, const int tokens, const int ncols, const int nrows, + const int token_stride, cudaStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + moe_vec_gfx1100 + <<>>( + vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static __global__ void moe_vec_id_gfx1100( + const void* __restrict__ vx, const void* __restrict__ vy, + scalar_t* __restrict__ dst, const int* topk_ids, const int topk, + const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, const int64_t row_stride_bytes) { + const int row = blockIdx.x * GGML_CUDA_MMV_Y + threadIdx.y; + const int route = blockIdx.z; + const int token = route / topk; + const int expert = topk_ids[route]; + if (row >= nrows || expert < 0) return; + const int blocks_per_row = ncols / qk; + const int blocks_per_wave = vdr * WARP_SIZE / qi; + const int lanes_per_chunk = qi / vdr; + const int block_lane = threadIdx.x / lanes_per_chunk; + const int iqs = vdr * (threadIdx.x % lanes_per_chunk); + const char* expert_base = static_cast(vx) + + static_cast(expert) * expert_stride_bytes; + const block_q_t* x = reinterpret_cast( + expert_base + static_cast(row) * row_stride_bytes); + const block_q8_1* y = reinterpret_cast( + static_cast(vy) + static_cast(token) * token_stride); + float tmp = 0.0f; + for (int i = block_lane; i < blocks_per_row; i += blocks_per_wave) + tmp += vec_dot_q_cuda(&x[i], &y[i * (qk / QK8_1)], iqs); +#pragma unroll + for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) + tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask); + if (threadIdx.x == 0) dst[route * nrows + row] = tmp; +} + +#define GGUF_GFX1100_ID_LAUNCH(NAME, QK, QI, BLOCK, VDR, DOT) \ +template \ +static void NAME(const void* vx, const void* vy, scalar_t* dst, const int* ids, \ + const int topk, const int tokens, const int ncols, const int nrows, \ + const int token_stride, const int64_t expert_stride, const int64_t row_stride, \ + cudaStream_t stream) { \ + const dim3 grid((nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y, 1, tokens * topk); \ + const dim3 block(WARP_SIZE, GGML_CUDA_MMV_Y, 1); \ + moe_vec_id_gfx1100 \ + <<>>(vx, vy, dst, ids, topk, ncols, nrows, token_stride, \ + expert_stride, row_stride); \ +} + +GGUF_GFX1100_ID_LAUNCH(moe_vec_q4_K_q8_1_gfx1100_id, QK_K, QI4_K, block_q4_K, + VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1) +GGUF_GFX1100_ID_LAUNCH(moe_vec_q5_K_q8_1_gfx1100_id, QK_K, QI5_K, block_q5_K, + VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1) +GGUF_GFX1100_ID_LAUNCH(moe_vec_q6_K_q8_1_gfx1100_id, QK_K, QI6_K, block_q6_K, + VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1) +GGUF_GFX1100_ID_LAUNCH(moe_vec_q8_0_q8_1_gfx1100_id, QK8_0, QI8_0, block_q8_0, + VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1) + +#undef GGUF_GFX1100_ID_LAUNCH diff --git a/python/freetoken/kernel/csrc/gguf/moe_vec_gfx1100_hip.cuh b/python/freetoken/kernel/csrc/gguf/moe_vec_gfx1100_hip.cuh new file mode 100644 index 000000000..b0021a578 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/moe_vec_gfx1100_hip.cuh @@ -0,0 +1,148 @@ +// !!! This is a file automatically generated by hipify!!! +#include "hip/hip_runtime.h" +// gfx1100 decode candidate for packed GGUF Q4_K/Q8_0 MoE vectors. +// +// This is intentionally separate from moe_vec.cuh. The incumbent assigns +// consecutive lanes to consecutive quant chunks. The candidate keeps the +// wave32 lane map explicit and processes a tunable number of output rows per wave. It +// still consumes native packed rows and calls the shared quant math, so the +// comparison isolates layout/occupancy from a changed quantization contract. +#pragma once + +template +static __global__ void moe_vec_gfx1100( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* topk_ids, + const int topk, + const int ncols, + const int nrows, + const int token_stride) { + constexpr int rows_per_wave = GGML_CUDA_MMV_Y; + const int row = blockIdx.x * rows_per_wave + threadIdx.y; + const int token = blockIdx.z / topk; + const int expert = topk_ids[blockIdx.z]; + + if (row >= nrows) { + return; + } + + const int lane = static_cast(threadIdx.x) & (WARP_SIZE - 1); + const int lanes_per_chunk = qi / vdr; + const int blocks_per_row = ncols / qk; + const int blocks_per_wave = vdr * WARP_SIZE / qi; + const int block_lane = lane / lanes_per_chunk; + const int iqs = vdr * (lane % lanes_per_chunk); + const block_q_t* x = static_cast(vx) + + expert * nrows * blocks_per_row; + const block_q8_1* y = static_cast( + static_cast(static_cast(vy) + token * token_stride)); + + float tmp = 0.0f; + for (int i = block_lane; i < blocks_per_row; i += blocks_per_wave) { + const int ibx = row * blocks_per_row + i; + const int iby = i * (qk / QK8_1); + tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs); + } + +#pragma unroll + for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { + tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask); + } + + if (threadIdx.x == 0) { + dst[blockIdx.z * nrows + row] = tmp; + } +} + +template +static void moe_vec_q4_K_q8_1_gfx1100( + const void* vx, const void* vy, scalar_t* dst, const int* topk_ids, + const int top_k, const int tokens, const int ncols, const int nrows, + const int token_stride, hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_gfx1100) + , dim3(block_nums), dim3(block_dims), 0, stream, + vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_q8_0_q8_1_gfx1100( + const void* vx, const void* vy, scalar_t* dst, const int* topk_ids, + const int top_k, const int tokens, const int ncols, const int nrows, + const int token_stride, hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_gfx1100) + , dim3(block_nums), dim3(block_dims), 0, stream, + vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +// ID-aware RDNA3 path. Unlike the historical candidate above, this kernel +// carries physical expert and row strides explicitly, so raw model IDs and +// offload slot IDs use the same ABI without assuming compact row packing. +template +static __global__ void moe_vec_id_gfx1100( + const void* __restrict__ vx, const void* __restrict__ vy, + scalar_t* __restrict__ dst, const int* topk_ids, const int topk, + const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, const int64_t row_stride_bytes) { + const int row = blockIdx.x * GGML_CUDA_MMV_Y + threadIdx.y; + const int route = blockIdx.z; + const int token = route / topk; + const int expert = topk_ids[route]; + if (row >= nrows || expert < 0) return; + + const int blocks_per_row = ncols / qk; + const int blocks_per_wave = vdr * WARP_SIZE / qi; + const int lanes_per_chunk = qi / vdr; + const int block_lane = threadIdx.x / lanes_per_chunk; + const int iqs = vdr * (threadIdx.x % lanes_per_chunk); + const char* expert_base = static_cast(vx) + + static_cast(expert) * expert_stride_bytes; + const block_q_t* x = reinterpret_cast( + expert_base + static_cast(row) * row_stride_bytes); + const block_q8_1* y = reinterpret_cast( + static_cast(vy) + static_cast(token) * token_stride); + float tmp = 0.0f; + for (int i = block_lane; i < blocks_per_row; i += blocks_per_wave) { + const int iby = i * (qk / QK8_1); + tmp += vec_dot_q_cuda(&x[i], &y[iby], iqs); + } +#pragma unroll + for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) + tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask); + if (threadIdx.x == 0) dst[route * nrows + row] = tmp; +} + +#define GGUF_GFX1100_ID_LAUNCH(NAME, QK, QI, BLOCK, VDR, DOT) \ +template \ +static void NAME(const void* vx, const void* vy, scalar_t* dst, const int* ids, \ + const int topk, const int tokens, const int ncols, const int nrows, \ + const int token_stride, const int64_t expert_stride, const int64_t row_stride, \ + hipStream_t stream) { \ + const dim3 grid((nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y, 1, tokens * topk); \ + const dim3 block(WARP_SIZE, GGML_CUDA_MMV_Y, 1); \ + hipLaunchKernelGGL((moe_vec_id_gfx1100), \ + grid, block, 0, stream, vx, vy, dst, ids, topk, ncols, nrows, token_stride, \ + expert_stride, row_stride); \ +} + +GGUF_GFX1100_ID_LAUNCH(moe_vec_q4_K_q8_1_gfx1100_id, QK_K, QI4_K, block_q4_K, + VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1) +GGUF_GFX1100_ID_LAUNCH(moe_vec_q5_K_q8_1_gfx1100_id, QK_K, QI5_K, block_q5_K, + VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1) +GGUF_GFX1100_ID_LAUNCH(moe_vec_q6_K_q8_1_gfx1100_id, QK_K, QI6_K, block_q6_K, + VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1) +GGUF_GFX1100_ID_LAUNCH(moe_vec_q8_0_q8_1_gfx1100_id, QK8_0, QI8_0, block_q8_0, + VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1) + +#undef GGUF_GFX1100_ID_LAUNCH diff --git a/python/freetoken/kernel/csrc/gguf/moe_vec_hip.cuh b/python/freetoken/kernel/csrc/gguf/moe_vec_hip.cuh new file mode 100644 index 000000000..2d4f046a9 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/moe_vec_hip.cuh @@ -0,0 +1,478 @@ +// !!! This is a file automatically generated by hipify!!! +#include "hip/hip_runtime.h" +// copied from +// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/moe_vec.cuh +// copied and adapted from +// https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu +// Port target cross-check: llama.cpp 7e4c0a968 (b10434), +// ggml/src/ggml-cuda/mmvq.cu; this file is the local HIP translation. +template +static __global__ void moe_vec_q( + const void* __restrict__ vx, + const void* __restrict__ vy, + scalar_t* __restrict__ dst, + const int* topk_ids, + const int topk, + const int ncols, + const int nrows, + const int token_stride) { + const auto row = blockIdx.x * blockDim.y + threadIdx.y; + + const auto token = blockIdx.z / topk; + const auto expert = (topk_ids)[blockIdx.z]; + + if (row >= nrows) { + return; + } + + const int blocks_per_row = ncols / qk; + const int blocks_per_warp = vdr * WARP_SIZE / qi; + + // partial sum for each thread + float tmp = 0.0f; + + const block_q_t* x = ((const block_q_t*)vx) + expert * nrows * blocks_per_row; + const block_q8_1* y = (const block_q8_1*)(((const int*)vy) + token * token_stride); + + for (auto i = threadIdx.x / (qi / vdr); i < blocks_per_row; i += blocks_per_warp) { + const int ibx = row * blocks_per_row + i; // x block index + + const int iby = i * (qk / QK8_1); // y block index that aligns with ibx + + const int iqs = vdr * (threadIdx.x % (qi / vdr)); // x block quant index when casting the quants to int + + tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs); + } + + // sum up partial sums and write back result +#pragma unroll + for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { + tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask); + } + + if (threadIdx.x == 0) { + dst[blockIdx.z * nrows + row] = tmp; + } +} + +template +static void moe_vec_q4_0_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_q4_1_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_q5_0_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_q5_1_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_q8_0_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_q2_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_q3_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_q4_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_q5_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_q6_K_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +// Native mixed-Q5_K/Q6_K offload uses a uniform Q6_K row stride. Q5_K data +// occupies its native prefix; this kernel skips padding between rows/experts. +template +static __global__ void moe_vec_q_strided( + const void* __restrict__ vx, const void* __restrict__ vy, + scalar_t* __restrict__ dst, const int* topk_ids, const int topk, + const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, const int64_t row_stride_bytes) { + const auto row = blockIdx.x * blockDim.y + threadIdx.y; + const auto token = blockIdx.z / topk; + const auto expert = topk_ids[blockIdx.z]; + if (row >= nrows) return; + const int blocks_per_row = ncols / qk; + const int blocks_per_warp = vdr * WARP_SIZE / qi; + float tmp = 0.0f; + const auto* expert_base = static_cast(vx) + + static_cast(expert) * expert_stride_bytes; + const block_q_t* x = reinterpret_cast( + expert_base + row * row_stride_bytes); + const block_q8_1* y = reinterpret_cast( + static_cast(vy) + token * token_stride); + for (auto i = threadIdx.x / (qi / vdr); i < blocks_per_row; i += blocks_per_warp) { + const int iby = i * (qk / QK8_1); + const int iqs = vdr * (threadIdx.x % (qi / vdr)); + tmp += vec_dot_q_cuda(&x[i], &y[iby], iqs); + } +#pragma unroll + for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) + tmp += SGLANG_SHFL_XOR_SYNC(uint32_t(-1), tmp, mask); + if (threadIdx.x == 0) dst[blockIdx.z * nrows + row] = tmp; +} + +template +static void moe_vec_q5_K_q8_1_strided_cuda( + const void* vx, const void* vy, scalar_t* dst, const int* topk_ids, + const int top_k, const int tokens, const int ncols, const int nrows, + const int token_stride, const int64_t expert_stride_bytes, + const int64_t row_stride_bytes, hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q_strided), + dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, + top_k, ncols, nrows, token_stride, expert_stride_bytes, row_stride_bytes); +} + +template +static void moe_vec_q6_K_q8_1_strided_cuda( + const void* vx, const void* vy, scalar_t* dst, const int* topk_ids, + const int top_k, const int tokens, const int ncols, const int nrows, + const int token_stride, const int64_t expert_stride_bytes, + const int64_t row_stride_bytes, hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q_strided), + dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, + top_k, ncols, nrows, token_stride, expert_stride_bytes, row_stride_bytes); +} + +template +static void moe_vec_iq2_xxs_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_iq2_xs_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_iq2_s_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_iq3_xxs_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_iq1_s_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_iq1_m_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_iq4_nl_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_iq4_xs_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} + +template +static void moe_vec_iq3_s_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + hipStream_t stream) { + const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; + const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); + hipLaunchKernelGGL(( moe_vec_q) + , dim3(block_nums), dim3(block_dims), 0, stream, vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); +} diff --git a/python/freetoken/kernel/csrc/gguf/vecdotq_hip.cuh b/python/freetoken/kernel/csrc/gguf/vecdotq_hip.cuh new file mode 100644 index 000000000..d838f5779 --- /dev/null +++ b/python/freetoken/kernel/csrc/gguf/vecdotq_hip.cuh @@ -0,0 +1,2039 @@ +// !!! This is a file automatically generated by hipify!!! +#include "hip/hip_runtime.h" +// copied from +// https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/vecdotq.cuh +// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/vecdotq.cuh +// and https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmq.cu +static __device__ __forceinline__ int get_int_b2(const void* x, const int& i32) { + const uint16_t* x16 = (const uint16_t*)x; // assume at least 2 byte alignment + + int x32 = x16[2 * i32 + 0] << 0; + x32 |= x16[2 * i32 + 1] << 16; + + return x32; +} + +static __device__ __forceinline__ int get_int_b4(const void* x, const int& i32) { + return ((const int*)x)[i32]; // assume at least 4 byte alignment +} + +static __device__ __forceinline__ int get_int_from_int8(const int8_t* x8, const int& i32) { + const uint16_t* x16 = (const uint16_t*)(x8 + sizeof(int) * i32); // assume at least 2 byte alignment + int x32 = 0; + x32 |= x16[0] << 0; + x32 |= x16[1] << 16; + return x32; +} + +static __device__ __forceinline__ int get_int_from_uint8(const uint8_t* x8, const int& i32) { + const uint16_t* x16 = (const uint16_t*)(x8 + sizeof(int) * i32); // assume at least 2 byte alignment + int x32 = 0; + x32 |= x16[0] << 0; + x32 |= x16[1] << 16; + return x32; +} + +static __device__ __forceinline__ int get_int_from_int8_aligned(const int8_t* x8, const int& i32) { + return *((const int*)(x8 + sizeof(int) * i32)); // assume at least 4 byte alignment +} + +static __device__ __forceinline__ int get_int_from_uint8_aligned(const uint8_t* x8, const int& i32) { + return *((const int*)(x8 + sizeof(int) * i32)); // assume at least 4 byte alignment +} + +// VDR = vec dot ratio, how many contiguous integers each thread processes when the vec dot kernel is called +// MMVQ = mul_mat_vec_q, MMQ = mul_mat_q + +#define VDR_Q4_0_Q8_1_MMVQ 2 +#define VDR_Q4_0_Q8_1_MMQ 4 + +template +static __device__ __forceinline__ float +vec_dot_q4_0_q8_1_impl(const int* v, const int* u, const float& d4, const half2& ds8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; + const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; + + // SIMD dot product of quantized values + sumi = __dp4a(vi0, u[2 * i + 0], sumi); + sumi = __dp4a(vi1, u[2 * i + 1], sumi); + } + + const float2 ds8f = __half22float2(ds8); + + // second part effectively subtracts 8 from each quant value + return d4 * (sumi * ds8f.x - (8 * vdr / QI4_0) * ds8f.y); +#endif +} + +#define VDR_Q4_1_Q8_1_MMVQ 2 +#define VDR_Q4_1_Q8_1_MMQ 4 + +template +static __device__ __forceinline__ float +vec_dot_q4_1_q8_1_impl(const int* v, const int* u, const half2& dm4, const half2& ds8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; + const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; + + // SIMD dot product of quantized values + sumi = __dp4a(vi0, u[2 * i + 0], sumi); + sumi = __dp4a(vi1, u[2 * i + 1], sumi); + } + + const float2 tmp = __half22float2(__hmul2(dm4, ds8)); + const float d4d8 = tmp.x; + const float m4s8 = tmp.y; + + // scale second part of sum by QI8_1/(vdr * QR4_1) to compensate for multiple threads adding it + return sumi * d4d8 + m4s8 / (QI8_1 / (vdr * QR4_1)); +#endif +} + +#define VDR_Q5_0_Q8_1_MMVQ 2 +#define VDR_Q5_0_Q8_1_MMQ 4 + +template +static __device__ __forceinline__ float +vec_dot_q5_0_q8_1_impl(const int* vl, const int* vh, const int* u, const float& d5, const half2& ds8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits + vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 + vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 + vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 + vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 + sumi = __dp4a(vi0, u[2 * i + 0], sumi); // SIMD dot product of quantized values + + int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits + vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 + vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 + vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 + vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 + sumi = __dp4a(vi1, u[2 * i + 1], sumi); // SIMD dot product of quantized values + } + + const float2 ds8f = __half22float2(ds8); + + // second part effectively subtracts 16 from each quant value + return d5 * (sumi * ds8f.x - (16 * vdr / QI5_0) * ds8f.y); +#endif +} + +#define VDR_Q5_1_Q8_1_MMVQ 2 +#define VDR_Q5_1_Q8_1_MMQ 4 + +template +static __device__ __forceinline__ float +vec_dot_q5_1_q8_1_impl(const int* vl, const int* vh, const int* u, const half2& dm5, const half2& ds8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits + vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 + vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 + vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 + vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 + sumi = __dp4a(vi0, u[2 * i + 0], sumi); // SIMD dot product of quantized values + + int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits + vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 + vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 + vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 + vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 + sumi = __dp4a(vi1, u[2 * i + 1], sumi); // SIMD dot product of quantized values + } + + const float2 tmp = __half22float2(__hmul2(dm5, ds8)); + const float d5d8 = tmp.x; + const float m5s8 = tmp.y; + + // scale second part of sum by QI5_1 / vdr to compensate for multiple threads adding it + return sumi * d5d8 + m5s8 / (QI5_1 / vdr); +#endif +} + +#define VDR_Q8_0_Q8_1_MMVQ 2 +#define VDR_Q8_0_Q8_1_MMQ 8 + +template +static __device__ __forceinline__ float +vec_dot_q8_0_q8_1_impl(const int* v, const int* u, const float& d8_0, const float& d8_1) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + // SIMD dot product of quantized values + sumi = __dp4a(v[i], u[i], sumi); + } + return d8_0 * d8_1 * sumi; +#endif +} + +template +static __device__ __forceinline__ float +vec_dot_q8_1_q8_1_impl(const int* v, const int* u, const half2& dm8, const half2& ds8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + // SIMD dot product of quantized values + sumi = __dp4a(v[i], u[i], sumi); + } + + const float2 tmp = __half22float2(__hmul2(dm8, ds8)); + const float d8d8 = tmp.x; + const float m8s8 = tmp.y; + + // scale second part of sum by QI8_1/ vdr to compensate for multiple threads adding it + return sumi * d8d8 + m8s8 / (QI8_1 / vdr); +#endif +} + +#define VDR_Q2_K_Q8_1_MMVQ 1 +#define VDR_Q2_K_Q8_1_MMQ 2 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq( + const int& v, + const int* __restrict__ u, + const uint8_t* __restrict__ scales, + const half2& dm2, + const float* __restrict__ d8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR2_K; ++i) { + const int sc = scales[2 * i]; + + const int vi = (v >> (2 * i)) & 0x03030303; + + sumf_d += d8[i] * (__dp4a(vi, u[i], 0) * (sc & 0xF)); // SIMD dot product + + // fill int with 4x m + int m = sc >> 4; + m |= m << 8; + m |= m << 16; + sumf_m += d8[i] * __dp4a(m, u[i], 0); // multiply constant q2_K part with sum of q8_1 values + } + + const float2 dm2f = __half22float2(dm2); + + return dm2f.x * sumf_d - dm2f.y * sumf_m; +#endif +} + +static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmq( + const int* __restrict__ v, + const int* __restrict__ u, + const uint8_t* __restrict__ scales, + const half2& dm2, + const float& d8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + int sumi_d = 0; + int sumi_m = 0; + +#pragma unroll + for (int i0 = 0; i0 < QI8_1; i0 += QI8_1 / 2) { + int sumi_d_sc = 0; + + const int sc = scales[i0 / (QI8_1 / 2)]; + + // fill int with 4x m + int m = sc >> 4; + m |= m << 8; + m |= m << 16; + +#pragma unroll + for (int i = i0; i < i0 + QI8_1 / 2; ++i) { + sumi_d_sc = __dp4a(v[i], u[i], sumi_d_sc); // SIMD dot product + sumi_m = __dp4a(m, u[i], sumi_m); // multiply sum of q8_1 values with m + } + + sumi_d += sumi_d_sc * (sc & 0xF); + } + + const float2 dm2f = __half22float2(dm2); + + return d8 * (dm2f.x * sumi_d - dm2f.y * sumi_m); +#endif +} + +#define VDR_Q3_K_Q8_1_MMVQ 1 +#define VDR_Q3_K_Q8_1_MMQ 2 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmvq( + const int& vl, + const int& vh, + const int* __restrict__ u, + const uint8_t* __restrict__ scales, + const int& scale_offset, + const float& d3, + const float* __restrict__ d8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + + float sumf = 0.0f; + +#pragma unroll + for (int i = 0; i < QR3_K; ++i) { + const int isc = scale_offset + 2 * i; + + const int isc_low = isc % (QK_K / 32); + const int sc_shift_low = 4 * (isc / (QK_K / 32)); + const int sc_low = (scales[isc_low] >> sc_shift_low) & 0xF; + + const int isc_high = isc % (QK_K / 64); + const int sc_shift_high = 2 * (isc / (QK_K / 64)); + const int sc_high = ((scales[(QK_K / 32) + isc_high] >> sc_shift_high) & 3) << 4; + + const int sc = (sc_low | sc_high) - 32; + + const int vil = (vl >> (2 * i)) & 0x03030303; + + const int vih = ((vh >> i) << 2) & 0x04040404; + + const int vi = __vsubss4(vil, vih); + + sumf += d8[i] * (__dp4a(vi, u[i], 0) * sc); // SIMD dot product + } + + return d3 * sumf; +#endif +} + +static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmq( + const int* __restrict__ v, + const int* __restrict__ u, + const int8_t* __restrict__ scales, + const float& d3, + const float& d8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + int sumi = 0; + +#pragma unroll + for (int i0 = 0; i0 < QR3_K * VDR_Q3_K_Q8_1_MMQ; i0 += QI8_1 / 2) { + int sumi_sc = 0; + + for (int i = i0; i < i0 + QI8_1 / 2; ++i) { + sumi_sc = __dp4a(v[i], u[i], sumi_sc); // SIMD dot product + } + + sumi += sumi_sc * scales[i0 / (QI8_1 / 2)]; + } + + return d3 * d8 * sumi; +#endif +} + +#define VDR_Q4_K_Q8_1_MMVQ 2 +#define VDR_Q4_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_vmmq( + const int* __restrict__ v, + const int* __restrict__ u, + const uint8_t* __restrict__ sc, + const uint8_t* __restrict__ m, + const half2& dm4, + const float* __restrict__ d8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR4_K; ++i) { + const int v0i = (v[0] >> (4 * i)) & 0x0F0F0F0F; + const int v1i = (v[1] >> (4 * i)) & 0x0F0F0F0F; + + const int dot1 = __dp4a(v1i, u[2 * i + 1], __dp4a(v0i, u[2 * i + 0], 0)); // SIMD dot product + const int dot2 = __dp4a(0x01010101, u[2 * i + 1], __dp4a(0x01010101, u[2 * i + 0], 0)); // sum of u + + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); // multiply constant part of q4_K with sum of q8_1 values + } + + const float2 dm4f = __half22float2(dm4); + return dm4f.x * sumf_d - dm4f.y * sumf_m; +#endif +} + +static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_mmq( + const int* __restrict__ v, + const int* __restrict__ u, + const uint8_t* __restrict__ sc, + const uint8_t* __restrict__ m, + const half2& dm4, + const half2* __restrict__ ds8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR4_K * VDR_Q4_K_Q8_1_MMQ / QI8_1; ++i) { + int sumi_d = 0; + +#pragma unroll + for (int j = 0; j < QI8_1; ++j) { + sumi_d = __dp4a((v[j] >> (4 * i)) & 0x0F0F0F0F, u[i * QI8_1 + j], sumi_d); // SIMD dot product + } + + const float2 ds8f = __half22float2(ds8[i]); + + sumf_d += ds8f.x * (sc[i] * sumi_d); + sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x * sumf_d - dm4f.y * sumf_m; +#endif +} + +#define VDR_Q5_K_Q8_1_MMVQ 2 +#define VDR_Q5_K_Q8_1_MMQ 8 + +static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_vmmq( + const int* __restrict__ vl, + const int* __restrict__ vh, + const int* __restrict__ u, + const uint8_t* __restrict__ sc, + const uint8_t* __restrict__ m, + const half2& dm5, + const float* __restrict__ d8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const int vl0i = (vl[0] >> (4 * i)) & 0x0F0F0F0F; + const int vl1i = (vl[1] >> (4 * i)) & 0x0F0F0F0F; + + const int vh0i = ((vh[0] >> i) << 4) & 0x10101010; + const int vh1i = ((vh[1] >> i) << 4) & 0x10101010; + + const int v0i = vl0i | vh0i; + const int v1i = vl1i | vh1i; + + const int dot1 = __dp4a(v0i, u[2 * i + 0], __dp4a(v1i, u[2 * i + 1], 0)); // SIMD dot product + const int dot2 = __dp4a(0x01010101, u[2 * i + 0], __dp4a(0x01010101, u[2 * i + 1], 0)); // sum of u + + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); + } + + const float2 dm5f = __half22float2(dm5); + return dm5f.x * sumf_d - dm5f.y * sumf_m; +#endif +} + +static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_mmq( + const int* __restrict__ v, + const int* __restrict__ u, + const uint8_t* __restrict__ sc, + const uint8_t* __restrict__ m, + const half2& dm4, + const half2* __restrict__ ds8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR5_K * VDR_Q5_K_Q8_1_MMQ / QI8_1; ++i) { + int sumi_d = 0; + +#pragma unroll + for (int j = 0; j < QI8_1; ++j) { + sumi_d = __dp4a(v[i * QI8_1 + j], u[i * QI8_1 + j], sumi_d); // SIMD dot product + } + + const float2 ds8f = __half22float2(ds8[i]); + + sumf_d += ds8f.x * (sc[i] * sumi_d); + sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x * sumf_d - dm4f.y * sumf_m; +#endif +} + +#define VDR_Q6_K_Q8_1_MMVQ 1 +#define VDR_Q6_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq( + const int& vl, + const int& vh, + const int* __restrict__ u, + const int8_t* __restrict__ scales, + const float& d, + const float* __restrict__ d8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + float sumf = 0.0f; + +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + const int sc = scales[4 * i]; + const int vil = (vl >> (4 * i)) & 0x0F0F0F0F; + const int vih = ((vh >> (4 * i)) << 4) & 0x30303030; + const int vi = __vsubss4((vil | vih), 0x20202020); // vi = (vil | vih) - 32 + + sumf += d8[i] * (__dp4a(vi, u[i], 0) * sc); // SIMD dot product + } + + return d * sumf; +#endif +} + +static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmq( + const int* __restrict__ v, + const int* __restrict__ u, + const int8_t* __restrict__ sc, + const float& d6, + const float* __restrict__ d8) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + float sumf_d = 0.0f; + +#pragma unroll + for (int i0 = 0; i0 < VDR_Q6_K_Q8_1_MMQ; i0 += 4) { + int2 sumi_d = {0, 0}; // 2 q6_K scales per q8_1 scale + +#pragma unroll + for (int i = i0; i < i0 + 2; ++i) { + sumi_d.x = __dp4a(v[2 * i + 0], u[2 * i + 0], sumi_d.x); // SIMD dot product + sumi_d.x = __dp4a(v[2 * i + 1], u[2 * i + 1], sumi_d.x); // SIMD dot product + + sumi_d.y = __dp4a(v[2 * i + 4], u[2 * i + 4], sumi_d.y); // SIMD dot product + sumi_d.y = __dp4a(v[2 * i + 5], u[2 * i + 5], sumi_d.y); // SIMD dot product + } + + sumf_d += d8[i0 / 4] * (sc[i0 / 2 + 0] * sumi_d.x + sc[i0 / 2 + 1] * sumi_d.y); + } + + return d6 * sumf_d; +#endif +} + +static __device__ __forceinline__ float +vec_dot_q4_0_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { + const block_q4_0* bq4_0 = (const block_q4_0*)vbq; + + int v[VDR_Q4_0_Q8_1_MMVQ]; + int u[2 * VDR_Q4_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q4_0_Q8_1_MMVQ; ++i) { + v[i] = get_int_from_uint8(bq4_0->qs, iqs + i); + u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_0); + } + + return vec_dot_q4_0_q8_1_impl(v, u, __half2float(bq4_0->d), bq8_1->ds); +} + +template +static __device__ __forceinline__ void allocate_tiles_q4_0(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_qs[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI4_0) + mmq_y / QI4_0]; + *x_ql = tile_x_qs; + *x_dm = (half2*)tile_x_d; +} + +template +static __device__ __forceinline__ void load_tiles_q4_0( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI4_0; + const int kqsx = k % QI4_0; + + const block_q4_0* bx0 = (const block_q4_0*)vx; + float* x_dmf = (float*)x_dm; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + if (need_check) { + i = min(i, i_max); + } + const block_q4_0* bxi = bx0 + i * blocks_per_row + kbx; + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs, kqsx); + // x_dmf[i * (WARP_SIZE_GGUF/QI4_0) + i / QI4_0 + kbx] = bxi->d; + } + + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI4_0; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_0) { + int i = i0 + i_offset * QI4_0 + k / blocks_per_tile_x_row; + if (need_check) { + i = min(i, i_max); + } + const block_q4_0* bxi = bx0 + i * blocks_per_row + kbxd; + x_dmf[i * (WARP_SIZE_GGUF / QI4_0) + i / QI4_0 + kbxd] = __half2float(bxi->d); + } +} + +static __device__ __forceinline__ float vec_dot_q4_0_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + (void)x_qh; + (void)x_sc; + + const int kyqs = k % (QI8_1 / 2) + QI8_1 * (k / (QI8_1 / 2)); + const float* x_dmf = (const float*)x_dm; + + int u[2 * VDR_Q4_0_Q8_1_MMQ]; + +#pragma unroll + for (int l = 0; l < VDR_Q4_0_Q8_1_MMQ; ++l) { + u[2 * l + 0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF]; + u[2 * l + 1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI4_0) % WARP_SIZE_GGUF]; + } + + return vec_dot_q4_0_q8_1_impl( + &x_ql[i * (WARP_SIZE_GGUF + 1) + k], + u, + x_dmf[i * (WARP_SIZE_GGUF / QI4_0) + i / QI4_0 + k / QI4_0], + y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (2 * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); +} + +static __device__ __forceinline__ float +vec_dot_q4_1_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { + const block_q4_1* bq4_1 = (const block_q4_1*)vbq; + + int v[VDR_Q4_1_Q8_1_MMVQ]; + int u[2 * VDR_Q4_1_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q4_1_Q8_1_MMVQ; ++i) { + v[i] = get_int_from_uint8_aligned(bq4_1->qs, iqs + i); + u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_1); + } + + return vec_dot_q4_1_q8_1_impl(v, u, bq4_1->dm, bq8_1->ds); +} + +template +static __device__ __forceinline__ void allocate_tiles_q4_1(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_qs[mmq_y * (WARP_SIZE_GGUF) + +mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI4_1) + mmq_y / QI4_1]; + *x_ql = tile_x_qs; + *x_dm = tile_x_dm; +} + +template +static __device__ __forceinline__ void load_tiles_q4_1( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI4_1; + const int kqsx = k % QI4_1; + + const block_q4_1* bx0 = (const block_q4_1*)vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + if (need_check) { + i = min(i, i_max); + } + const block_q4_1* bxi = bx0 + i * blocks_per_row + kbx; + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx); + } + + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI4_1; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_1) { + int i = i0 + i_offset * QI4_1 + k / blocks_per_tile_x_row; + if (need_check) { + i = min(i, i_max); + } + const block_q4_1* bxi = bx0 + i * blocks_per_row + kbxd; + x_dm[i * (WARP_SIZE_GGUF / QI4_1) + i / QI4_1 + kbxd] = bxi->dm; + } +} + +static __device__ __forceinline__ float vec_dot_q4_1_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + const int kyqs = k % (QI8_1 / 2) + QI8_1 * (k / (QI8_1 / 2)); + + int u[2 * VDR_Q4_1_Q8_1_MMQ]; + +#pragma unroll + for (int l = 0; l < VDR_Q4_1_Q8_1_MMQ; ++l) { + u[2 * l + 0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF]; + u[2 * l + 1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI4_1) % WARP_SIZE_GGUF]; + } + + return vec_dot_q4_1_q8_1_impl( + &x_ql[i * (WARP_SIZE_GGUF + 1) + k], + u, + x_dm[i * (WARP_SIZE_GGUF / QI4_1) + i / QI4_1 + k / QI4_1], + y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (2 * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); +} + +static __device__ __forceinline__ float +vec_dot_q5_0_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { + const block_q5_0* bq5_0 = (const block_q5_0*)vbq; + + int vl[VDR_Q5_0_Q8_1_MMVQ]; + int vh[VDR_Q5_0_Q8_1_MMVQ]; + int u[2 * VDR_Q5_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q5_0_Q8_1_MMVQ; ++i) { + vl[i] = get_int_from_uint8(bq5_0->qs, iqs + i); + vh[i] = get_int_from_uint8(bq5_0->qh, 0) >> (4 * (iqs + i)); + u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_0); + } + + return vec_dot_q5_0_q8_1_impl(vl, vh, u, __half2float(bq5_0->d), bq8_1->ds); +} + +template +static __device__ __forceinline__ void allocate_tiles_q5_0(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_ql[mmq_y * (2 * WARP_SIZE_GGUF) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI5_0) + mmq_y / QI5_0]; + + *x_ql = tile_x_ql; + *x_dm = (half2*)tile_x_d; +} + +template +static __device__ __forceinline__ void load_tiles_q5_0( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI5_0; + const int kqsx = k % QI5_0; + + const block_q5_0* bx0 = (const block_q5_0*)vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + const block_q5_0* bxi = bx0 + i * blocks_per_row + kbx; + const int ql = get_int_from_uint8(bxi->qs, kqsx); + const int qh = get_int_from_uint8(bxi->qh, 0) >> (4 * (k % QI5_0)); + + int qs0 = (ql >> 0) & 0x0F0F0F0F; + qs0 |= (qh << 4) & 0x00000010; // 0 -> 4 + qs0 |= (qh << 11) & 0x00001000; // 1 -> 12 + qs0 |= (qh << 18) & 0x00100000; // 2 -> 20 + qs0 |= (qh << 25) & 0x10000000; // 3 -> 28 + qs0 = __vsubss4(qs0, 0x10101010); // subtract 16 + + x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k + 0] = qs0; + + int qs1 = (ql >> 4) & 0x0F0F0F0F; + qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4 + qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12 + qs1 |= (qh << 2) & 0x00100000; // 18 -> 20 + qs1 |= (qh << 9) & 0x10000000; // 19 -> 28 + qs1 = __vsubss4(qs1, 0x10101010); // subtract 16 + + x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k + 1] = qs1; + } + + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI5_0; + const int kbxd = k % blocks_per_tile_x_row; + float* x_dmf = (float*)x_dm; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_0) { + int i = i0 + i_offset * QI5_0 + k / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_0* bxi = bx0 + i * blocks_per_row + kbxd; + x_dmf[i * (WARP_SIZE_GGUF / QI5_0) + i / QI5_0 + kbxd] = __half2float(bxi->d); + } +} + +static __device__ __forceinline__ float vec_dot_q5_0_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + const int kyqs = k % (QI8_1 / 2) + QI8_1 * (k / (QI8_1 / 2)); + const int index_bx = i * (WARP_SIZE_GGUF / QI5_0) + i / QI5_0 + k / QI5_0; + const float* x_dmf = (const float*)x_dm; + const float* y_df = (const float*)y_ds; + + int u[2 * VDR_Q5_0_Q8_1_MMQ]; + +#pragma unroll + for (int l = 0; l < VDR_Q5_0_Q8_1_MMQ; ++l) { + u[2 * l + 0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF]; + u[2 * l + 1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI5_0) % WARP_SIZE_GGUF]; + } + + return vec_dot_q8_0_q8_1_impl( + &x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k], + u, + x_dmf[index_bx], + y_df[j * (WARP_SIZE_GGUF / QI8_1) + (2 * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); +} + +static __device__ __forceinline__ float +vec_dot_q5_1_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { + const block_q5_1* bq5_1 = (const block_q5_1*)vbq; + + int vl[VDR_Q5_1_Q8_1_MMVQ]; + int vh[VDR_Q5_1_Q8_1_MMVQ]; + int u[2 * VDR_Q5_1_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q5_1_Q8_1_MMVQ; ++i) { + vl[i] = get_int_from_uint8_aligned(bq5_1->qs, iqs + i); + vh[i] = get_int_from_uint8_aligned(bq5_1->qh, 0) >> (4 * (iqs + i)); + u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_1); + } + + return vec_dot_q5_1_q8_1_impl(vl, vh, u, bq5_1->dm, bq8_1->ds); +} + +template +static __device__ __forceinline__ void allocate_tiles_q5_1(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_ql[mmq_y * (2 * WARP_SIZE_GGUF) + mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI5_1) + mmq_y / QI5_1]; + + *x_ql = tile_x_ql; + *x_dm = tile_x_dm; +} + +template +static __device__ __forceinline__ void load_tiles_q5_1( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI5_1; + const int kqsx = k % QI5_1; + + const block_q5_1* bx0 = (const block_q5_1*)vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_1* bxi = bx0 + i * blocks_per_row + kbx; + + const int ql = get_int_from_uint8_aligned(bxi->qs, kqsx); + const int qh = get_int_from_uint8_aligned(bxi->qh, 0) >> (4 * (k % QI5_1)); + + int qs0 = (ql >> 0) & 0x0F0F0F0F; + qs0 |= (qh << 4) & 0x00000010; // 0 -> 4 + qs0 |= (qh << 11) & 0x00001000; // 1 -> 12 + qs0 |= (qh << 18) & 0x00100000; // 2 -> 20 + qs0 |= (qh << 25) & 0x10000000; // 3 -> 28 + + x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k + 0] = qs0; + + int qs1 = (ql >> 4) & 0x0F0F0F0F; + qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4 + qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12 + qs1 |= (qh << 2) & 0x00100000; // 18 -> 20 + qs1 |= (qh << 9) & 0x10000000; // 19 -> 28 + + x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k + 1] = qs1; + } + + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI5_1; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_1) { + int i = i0 + i_offset * QI5_1 + k / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_1* bxi = bx0 + i * blocks_per_row + kbxd; + + x_dm[i * (WARP_SIZE_GGUF / QI5_1) + i / QI5_1 + kbxd] = bxi->dm; + } +} + +static __device__ __forceinline__ float vec_dot_q5_1_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + const int kyqs = k % (QI8_1 / 2) + QI8_1 * (k / (QI8_1 / 2)); + const int index_bx = i * (WARP_SIZE_GGUF / QI5_1) + +i / QI5_1 + k / QI5_1; + + int u[2 * VDR_Q5_1_Q8_1_MMQ]; + +#pragma unroll + for (int l = 0; l < VDR_Q5_1_Q8_1_MMQ; ++l) { + u[2 * l + 0] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l) % WARP_SIZE_GGUF]; + u[2 * l + 1] = y_qs[j * WARP_SIZE_GGUF + (kyqs + l + QI5_1) % WARP_SIZE_GGUF]; + } + + return vec_dot_q8_1_q8_1_impl( + &x_ql[i * (2 * WARP_SIZE_GGUF + 1) + 2 * k], + u, + x_dm[index_bx], + y_ds[j * (WARP_SIZE_GGUF / QI8_1) + (2 * k / QI8_1) % (WARP_SIZE_GGUF / QI8_1)]); +} + +static __device__ __forceinline__ float +vec_dot_q8_0_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { + const block_q8_0* bq8_0 = (const block_q8_0*)vbq; + + int v[VDR_Q8_0_Q8_1_MMVQ]; + int u[VDR_Q8_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q8_0_Q8_1_MMVQ; ++i) { + v[i] = get_int_from_int8(bq8_0->qs, iqs + i); + u[i] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + } + + return vec_dot_q8_0_q8_1_impl(v, u, __half2float(bq8_0->d), __low2float(bq8_1->ds)); +} + +template +static __device__ __forceinline__ void allocate_tiles_q8_0(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_qs[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE_GGUF / QI8_0) + mmq_y / QI8_0]; + + *x_ql = tile_x_qs; + *x_dm = (half2*)tile_x_d; +} + +template +static __device__ __forceinline__ void load_tiles_q8_0( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI8_0; + const int kqsx = k % QI8_0; + float* x_dmf = (float*)x_dm; + + const block_q8_0* bx0 = (const block_q8_0*)vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + const block_q8_0* bxi = bx0 + i * blocks_per_row + kbx; + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_int8(bxi->qs, kqsx); + } + + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI8_0; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI8_0) { + int i = i0 + i_offset * QI8_0 + k / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + const block_q8_0* bxi = bx0 + i * blocks_per_row + kbxd; + x_dmf[i * (WARP_SIZE_GGUF / QI8_0) + i / QI8_0 + kbxd] = __half2float(bxi->d); + } +} + +static __device__ __forceinline__ float vec_dot_q8_0_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + const float* x_dmf = (const float*)x_dm; + const float* y_df = (const float*)y_ds; + + return vec_dot_q8_0_q8_1_impl( + &x_ql[i * (WARP_SIZE_GGUF + 1) + k], + &y_qs[j * WARP_SIZE_GGUF + k], + x_dmf[i * (WARP_SIZE_GGUF / QI8_0) + i / QI8_0 + k / QI8_0], + y_df[j * (WARP_SIZE_GGUF / QI8_1) + k / QI8_1]); +} + +static __device__ __forceinline__ float +vec_dot_q2_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { + const block_q2_K* bq2_K = (const block_q2_K*)vbq; + + const int bq8_offset = QR2_K * (iqs / QI8_1); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1 / 2); + + const uint8_t* scales = bq2_K->scales + scale_offset; + + const int v = get_int_from_uint8_aligned(bq2_K->qs, iqs); + int u[QR2_K]; + float d8[QR2_K]; + +#pragma unroll + for (int i = 0; i < QR2_K; ++i) { + u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + i].ds); + } + + return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales, bq2_K->dm, d8); +} + +template +static __device__ __forceinline__ void allocate_tiles_q2_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI2_K) + mmq_y / QI2_K]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 4) + mmq_y / 4]; + + *x_ql = tile_x_ql; + *x_dm = tile_x_dm; + *x_sc = tile_x_sc; +} + +template +static __device__ __forceinline__ void load_tiles_q2_K( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI2_K; + const int kqsx = k % QI2_K; + + const block_q2_K* bx0 = (const block_q2_K*)vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + const block_q2_K* bxi = bx0 + i * blocks_per_row + kbx; + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx); + } + + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI2_K; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI2_K) { + int i = (i0 + i_offset * QI2_K + k / blocks_per_tile_x_row) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + const block_q2_K* bxi = bx0 + i * blocks_per_row + kbxd; + x_dm[i * (WARP_SIZE_GGUF / QI2_K) + i / QI2_K + kbxd] = bxi->dm; + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 4) { + int i = i0 + i_offset * 4 + k / (WARP_SIZE_GGUF / 4); + + if (need_check) { + i = min(i, i_max); + } + const block_q2_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 4)) / (QI2_K / 4); + x_sc[i * (WARP_SIZE_GGUF / 4) + i / 4 + k % (WARP_SIZE_GGUF / 4)] = + get_int_from_uint8_aligned(bxi->scales, k % (QI2_K / 4)); + } +} + +static __device__ __forceinline__ float vec_dot_q2_K_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + const int kbx = k / QI2_K; + const int ky = (k % QI2_K) * QR2_K; + const float* y_df = (const float*)y_ds; + + int v[QR2_K * VDR_Q2_K_Q8_1_MMQ]; + + const int kqsx = i * (WARP_SIZE_GGUF + 1) + kbx * QI2_K + (QI2_K / 2) * (ky / (2 * QI2_K)) + ky % (QI2_K / 2); + const int shift = 2 * ((ky % (2 * QI2_K)) / (QI2_K / 2)); + +#pragma unroll + for (int l = 0; l < QR2_K * VDR_Q2_K_Q8_1_MMQ; ++l) { + v[l] = (x_ql[kqsx + l] >> shift) & 0x03030303; + } + + const uint8_t* scales = ((const uint8_t*)&x_sc[i * (WARP_SIZE_GGUF / 4) + i / 4 + kbx * 4]) + ky / 4; + + const int index_y = j * WARP_SIZE_GGUF + (QR2_K * k) % WARP_SIZE_GGUF; + return vec_dot_q2_K_q8_1_impl_mmq( + v, &y_qs[index_y], scales, x_dm[i * (WARP_SIZE_GGUF / QI2_K) + i / QI2_K + kbx], y_df[index_y / QI8_1]); +} + +static __device__ __forceinline__ float +vec_dot_q3_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { + const block_q3_K* bq3_K = (const block_q3_K*)vbq; + + const int bq8_offset = QR3_K * (iqs / (QI3_K / 2)); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1 / 2); + + const float d = __half2float(bq3_K->d); + + const int vl = get_int_from_uint8(bq3_K->qs, iqs); + + // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted + const int vh = ~get_int_from_uint8(bq3_K->hmask, iqs % (QI3_K / 2)) >> bq8_offset; + + int u[QR3_K]; + float d8[QR3_K]; + +#pragma unroll + for (int i = 0; i < QR3_K; ++i) { + u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + i].ds); + } + + return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, bq3_K->scales, scale_offset, d, d8); +} + +template +static __device__ __forceinline__ void allocate_tiles_q3_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI3_K) + mmq_y / QI3_K]; + __shared__ int tile_x_qh[mmq_y * (WARP_SIZE_GGUF / 2) + mmq_y / 2]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 4) + mmq_y / 4]; + + *x_ql = tile_x_ql; + *x_dm = tile_x_dm; + *x_qh = tile_x_qh; + *x_sc = tile_x_sc; +} + +template +static __device__ __forceinline__ void load_tiles_q3_K( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI3_K; + const int kqsx = k % QI3_K; + + const block_q3_K* bx0 = (const block_q3_K*)vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + if (need_check) { + i = min(i, i_max); + } + const block_q3_K* bxi = bx0 + i * blocks_per_row + kbx; + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8(bxi->qs, kqsx); + } + + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI3_K; + const int kbxd = k % blocks_per_tile_x_row; + float* x_dmf = (float*)x_dm; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI3_K) { + int i = (i0 + i_offset * QI3_K + k / blocks_per_tile_x_row) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const block_q3_K* bxi = bx0 + i * blocks_per_row + kbxd; + x_dmf[i * (WARP_SIZE_GGUF / QI3_K) + i / QI3_K + kbxd] = __half2float(bxi->d); + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 2) { + int i = i0 + i_offset * 2 + k / (WARP_SIZE_GGUF / 2); + if (need_check) { + i = min(i, i_max); + } + const block_q3_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 2)) / (QI3_K / 2); + // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted + x_qh[i * (WARP_SIZE_GGUF / 2) + i / 2 + k % (WARP_SIZE_GGUF / 2)] = + ~get_int_from_uint8(bxi->hmask, k % (QI3_K / 2)); + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 4) { + int i = i0 + i_offset * 4 + k / (WARP_SIZE_GGUF / 4); + if (need_check) { + i = min(i, i_max); + } + const block_q3_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 4)) / (QI3_K / 4); + + const int ksc = k % (QI3_K / 4); + + const int ksc_low = ksc % (QI3_K / 8); + const int shift_low = 4 * (ksc / (QI3_K / 8)); + const int sc_low = (get_int_from_uint8(bxi->scales, ksc_low) >> shift_low) & 0x0F0F0F0F; + + const int ksc_high = QI3_K / 8; + const int shift_high = 2 * ksc; + const int sc_high = ((get_int_from_uint8(bxi->scales, ksc_high) >> shift_high) << 4) & 0x30303030; + + const int sc = __vsubss4(sc_low | sc_high, 0x20202020); + + x_sc[i * (WARP_SIZE_GGUF / 4) + i / 4 + k % (WARP_SIZE_GGUF / 4)] = sc; + } +} + +static __device__ __forceinline__ float vec_dot_q3_K_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + const int kbx = k / QI3_K; + const int ky = (k % QI3_K) * QR3_K; + const float* x_dmf = (const float*)x_dm; + const float* y_df = (const float*)y_ds; + + const int8_t* scales = ((const int8_t*)(x_sc + i * (WARP_SIZE_GGUF / 4) + i / 4 + kbx * 4)) + ky / 4; + + int v[QR3_K * VDR_Q3_K_Q8_1_MMQ]; + +#pragma unroll + for (int l = 0; l < QR3_K * VDR_Q3_K_Q8_1_MMQ; ++l) { + const int kqsx = i * (WARP_SIZE_GGUF + 1) + kbx * QI3_K + (QI3_K / 2) * (ky / (2 * QI3_K)) + ky % (QI3_K / 2); + const int shift = 2 * ((ky % 32) / 8); + const int vll = (x_ql[kqsx + l] >> shift) & 0x03030303; + + const int vh = x_qh[i * (WARP_SIZE_GGUF / 2) + i / 2 + kbx * (QI3_K / 2) + (ky + l) % 8] >> ((ky + l) / 8); + const int vlh = (vh << 2) & 0x04040404; + + v[l] = __vsubss4(vll, vlh); + } + + const int index_y = j * WARP_SIZE_GGUF + (k * QR3_K) % WARP_SIZE_GGUF; + return vec_dot_q3_K_q8_1_impl_mmq( + v, &y_qs[index_y], scales, x_dmf[i * (WARP_SIZE_GGUF / QI3_K) + i / QI3_K + kbx], y_df[index_y / QI8_1]); +} + +static __device__ __forceinline__ float +vec_dot_q4_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { + const block_q4_K* bq4_K = (const block_q4_K*)vbq; + + int v[2]; + int u[2 * QR4_K]; + float d8[QR4_K]; + + // iqs is in 0,2..30. bq8_offset = iqs/4 -> bq8_offset = 0, 2, 4, 6 + const int bq8_offset = QR4_K * ((iqs / 2) / (QI8_1 / 2)); + + // iqs = 0....3 -> bq8_offset = 0, want q4_offset = 0, 4, 8, 12 + // iqs = 4....7 -> bq8_offset = 2, want q4_offset = 32, 36, 40, 44 + // iqs = 8...11 -> bq8_offset = 4, want q4_offset = 64, 68, 72, 76 + // iqs = 12..15 -> bq8_offset = 6, want q4_offset = 96, 100, 104, 108 + + const int* q4 = (const int*)(bq4_K->qs + 16 * bq8_offset + 4 * ((iqs / 2) % 4)); + v[0] = q4[0]; + v[1] = q4[4]; + + const uint16_t* scales = (const uint16_t*)bq4_K->scales; + uint16_t aux[2]; + const int j = bq8_offset / 2; + if (j < 2) { + aux[0] = scales[j + 0] & 0x3f3f; + aux[1] = scales[j + 2] & 0x3f3f; + } else { + aux[0] = ((scales[j + 2] >> 0) & 0x0f0f) | ((scales[j - 2] & 0xc0c0) >> 2); + aux[1] = ((scales[j + 2] >> 4) & 0x0f0f) | ((scales[j - 0] & 0xc0c0) >> 2); + } + const uint8_t* sc = (const uint8_t*)aux; + const uint8_t* m = sc + 2; + + for (int i = 0; i < QR4_K; ++i) { + const block_q8_1* bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + + const int* q8 = (const int*)bq8i->qs + ((iqs / 2) % 4); + u[2 * i + 0] = q8[0]; + u[2 * i + 1] = q8[4]; + } + + return vec_dot_q4_K_q8_1_impl_vmmq(v, u, sc, m, bq4_K->dm, d8); +} + +template +static __device__ __forceinline__ void allocate_tiles_q4_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_ql[mmq_y * (WARP_SIZE_GGUF) + mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI4_K) + mmq_y / QI4_K]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 8) + mmq_y / 8]; + + *x_ql = tile_x_ql; + *x_dm = tile_x_dm; + *x_sc = tile_x_sc; +} + +template +static __device__ __forceinline__ void load_tiles_q4_K( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI4_K; // == 0 if QK_K == 256 + const int kqsx = k % QI4_K; // == k if QK_K == 256 + + const block_q4_K* bx0 = (const block_q4_K*)vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + const block_q4_K* bxi = bx0 + i * blocks_per_row + kbx; + x_ql[i * (WARP_SIZE_GGUF + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx); + } + + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI4_K; // == 1 if QK_K == 256 + const int kbxd = k % blocks_per_tile_x_row; // == 0 if QK_K == 256 + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_K) { + int i = (i0 + i_offset * QI4_K + k / blocks_per_tile_x_row) % mmq_y; + if (need_check) { + i = min(i, i_max); + } + const block_q4_K* bxi = bx0 + i * blocks_per_row + kbxd; + x_dm[i * (WARP_SIZE_GGUF / QI4_K) + i / QI4_K + kbxd] = bxi->dm; + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { + int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF / 8)) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 8)) / (QI4_K / 8); + + const int* scales = (const int*)bxi->scales; + + const int ksc = k % (WARP_SIZE_GGUF / 8); + // scale arrangement after the following two lines: sc0,...,sc3, sc4,...,sc7, m0,...,m3, m4,...,m8 + int scales8 = (scales[(ksc % 2) + (ksc != 0)] >> (4 * (ksc & (ksc / 2)))) & 0x0F0F0F0F; // lower 4 bits + scales8 |= (scales[ksc / 2] >> (2 * (ksc % 2))) & 0x30303030; // upper 2 bits + + x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + ksc] = scales8; + } +} + +static __device__ __forceinline__ float vec_dot_q4_K_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + (void)x_qh; + + const uint8_t* sc = ((const uint8_t*)&x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + k / 16]) + 2 * ((k % 16) / 8); + + const int index_y = j * WARP_SIZE_GGUF + (QR4_K * k) % WARP_SIZE_GGUF; + return vec_dot_q4_K_q8_1_impl_mmq( + &x_ql[i * (WARP_SIZE_GGUF + 1) + k], + &y_qs[index_y], + sc, + sc + 8, + x_dm[i * (WARP_SIZE_GGUF / QI4_K) + i / QI4_K], + &y_ds[index_y / QI8_1]); +} + +static __device__ __forceinline__ float +vec_dot_q5_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { + const block_q5_K* bq5_K = (const block_q5_K*)vbq; + + int vl[2]; + int vh[2]; + int u[2 * QR5_K]; + float d8[QR5_K]; + + const int bq8_offset = QR5_K * ((iqs / 2) / (QI8_1 / 2)); + const int* ql = (const int*)(bq5_K->qs + 16 * bq8_offset + 4 * ((iqs / 2) % 4)); + const int* qh = (const int*)(bq5_K->qh + 4 * ((iqs / 2) % 4)); + + vl[0] = ql[0]; + vl[1] = ql[4]; + + vh[0] = qh[0] >> bq8_offset; + vh[1] = qh[4] >> bq8_offset; + + const uint16_t* scales = (const uint16_t*)bq5_K->scales; + uint16_t aux[2]; + const int j = bq8_offset / 2; + if (j < 2) { + aux[0] = scales[j + 0] & 0x3f3f; + aux[1] = scales[j + 2] & 0x3f3f; + } else { + aux[0] = ((scales[j + 2] >> 0) & 0x0f0f) | ((scales[j - 2] & 0xc0c0) >> 2); + aux[1] = ((scales[j + 2] >> 4) & 0x0f0f) | ((scales[j - 0] & 0xc0c0) >> 2); + } + const uint8_t* sc = (const uint8_t*)aux; + const uint8_t* m = sc + 2; + +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const block_q8_1* bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + + const int* q8 = (const int*)bq8i->qs + ((iqs / 2) % 4); + u[2 * i + 0] = q8[0]; + u[2 * i + 1] = q8[4]; + } + + return vec_dot_q5_K_q8_1_impl_vmmq(vl, vh, u, sc, m, bq5_K->dm, d8); +} + +template +static __device__ __forceinline__ void allocate_tiles_q5_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_ql[mmq_y * (2 * WARP_SIZE_GGUF) + mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI5_K) + mmq_y / QI5_K]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 8) + mmq_y / 8]; + + *x_ql = tile_x_ql; + *x_dm = tile_x_dm; + *x_sc = tile_x_sc; +} + +template +static __device__ __forceinline__ void load_tiles_q5_K( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI5_K; // == 0 if QK_K == 256 + const int kqsx = k % QI5_K; // == k if QK_K == 256 + + const block_q5_K* bx0 = (const block_q5_K*)vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_K* bxi = bx0 + i * blocks_per_row + kbx; + const int ky = QR5_K * kqsx; + + const int ql = get_int_from_uint8_aligned(bxi->qs, kqsx); + const int ql0 = (ql >> 0) & 0x0F0F0F0F; + const int ql1 = (ql >> 4) & 0x0F0F0F0F; + + const int qh = get_int_from_uint8_aligned(bxi->qh, kqsx % (QI5_K / 4)); + const int qh0 = ((qh >> (2 * (kqsx / (QI5_K / 4)) + 0)) << 4) & 0x10101010; + const int qh1 = ((qh >> (2 * (kqsx / (QI5_K / 4)) + 1)) << 4) & 0x10101010; + + const int kq0 = ky - ky % (QI5_K / 2) + k % (QI5_K / 4) + 0; + const int kq1 = ky - ky % (QI5_K / 2) + k % (QI5_K / 4) + (QI5_K / 4); + + x_ql[i * (2 * WARP_SIZE_GGUF + 1) + kq0] = ql0 | qh0; + x_ql[i * (2 * WARP_SIZE_GGUF + 1) + kq1] = ql1 | qh1; + } + + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI5_K; // == 1 if QK_K == 256 + const int kbxd = k % blocks_per_tile_x_row; // == 0 if QK_K == 256 + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_K) { + int i = (i0 + i_offset * QI5_K + k / blocks_per_tile_x_row) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_K* bxi = bx0 + i * blocks_per_row + kbxd; + x_dm[i * (WARP_SIZE_GGUF / QI5_K) + i / QI5_K + kbxd] = bxi->dm; + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { + int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF / 8)) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 8)) / (QI5_K / 8); + + const int* scales = (const int*)bxi->scales; + + const int ksc = k % (WARP_SIZE_GGUF / 8); + + // scale arrangement after the following two lines: sc0,...,sc3, sc4,...,sc7, m0,...,m3, m4,...,m8 + int scales8 = (scales[(ksc % 2) + (ksc != 0)] >> (4 * (ksc & (ksc / 2)))) & 0x0F0F0F0F; // lower 4 bits + scales8 |= (scales[ksc / 2] >> (2 * (ksc % 2))) & 0x30303030; // upper 2 bits + + x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + ksc] = scales8; + } +} + +static __device__ __forceinline__ float vec_dot_q5_K_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + const uint8_t* sc = ((const uint8_t*)&x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + k / 16]) + 2 * ((k % 16) / 8); + + const int index_x = i * (QR5_K * WARP_SIZE_GGUF + 1) + QR5_K * k; + const int index_y = j * WARP_SIZE_GGUF + (QR5_K * k) % WARP_SIZE_GGUF; + return vec_dot_q5_K_q8_1_impl_mmq( + &x_ql[index_x], + &y_qs[index_y], + sc, + sc + 8, + x_dm[i * (WARP_SIZE_GGUF / QI5_K) + i / QI5_K], + &y_ds[index_y / QI8_1]); +} + +static __device__ __forceinline__ float +vec_dot_q6_K_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { + const block_q6_K* bq6_K = (const block_q6_K*)vbq; + + const int bq8_offset = 2 * QR6_K * (iqs / (QI6_K / 2)) + (iqs % (QI6_K / 2)) / (QI6_K / 4); + const int scale_offset = (QI6_K / 4) * (iqs / (QI6_K / 2)) + (iqs % (QI6_K / 2)) / (QI6_K / 8); + const int vh_shift = 2 * ((iqs % (QI6_K / 2)) / (QI6_K / 4)); + + const int vl = get_int_from_uint8(bq6_K->ql, iqs); + const int vh = get_int_from_uint8(bq6_K->qh, (QI6_K / 4) * (iqs / (QI6_K / 2)) + iqs % (QI6_K / 4)) >> vh_shift; + + const int8_t* scales = bq6_K->scales + scale_offset; + + int u[QR6_K]; + float d8[QR6_K]; + +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + 2 * i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + 2 * i].ds); + } + + return vec_dot_q6_K_q8_1_impl_mmvq(vl, vh, u, scales, __half2float(bq6_K->d), d8); +} + +template +static __device__ __forceinline__ void allocate_tiles_q6_K(int** x_ql, half2** x_dm, int** x_qh, int** x_sc) { + __shared__ int tile_x_ql[mmq_y * (2 * WARP_SIZE_GGUF) + mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE_GGUF / QI6_K) + mmq_y / QI6_K]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE_GGUF / 8) + mmq_y / 8]; + + *x_ql = tile_x_ql; + *x_dm = tile_x_dm; + *x_sc = tile_x_sc; +} + +template +static __device__ __forceinline__ void load_tiles_q6_K( + const void* __restrict__ vx, + int* __restrict__ x_ql, + half2* __restrict__ x_dm, + int* __restrict__ x_qh, + int* __restrict__ x_sc, + const int& i_offset, + const int& i_max, + const int& k, + const int& blocks_per_row) { + const int kbx = k / QI6_K; // == 0 if QK_K == 256 + const int kqsx = k % QI6_K; // == k if QK_K == 256 + + const block_q6_K* bx0 = (const block_q6_K*)vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + + const block_q6_K* bxi = bx0 + i * blocks_per_row + kbx; + const int ky = QR6_K * kqsx; + + const int ql = get_int_from_uint8(bxi->ql, kqsx); + const int ql0 = (ql >> 0) & 0x0F0F0F0F; + const int ql1 = (ql >> 4) & 0x0F0F0F0F; + + const int qh = get_int_from_uint8(bxi->qh, (QI6_K / 4) * (kqsx / (QI6_K / 2)) + kqsx % (QI6_K / 4)); + const int qh0 = ((qh >> (2 * ((kqsx % (QI6_K / 2)) / (QI6_K / 4)))) << 4) & 0x30303030; + const int qh1 = (qh >> (2 * ((kqsx % (QI6_K / 2)) / (QI6_K / 4)))) & 0x30303030; + + const int kq0 = ky - ky % QI6_K + k % (QI6_K / 2) + 0; + const int kq1 = ky - ky % QI6_K + k % (QI6_K / 2) + (QI6_K / 2); + + x_ql[i * (2 * WARP_SIZE_GGUF + 1) + kq0] = __vsubss4(ql0 | qh0, 0x20202020); + x_ql[i * (2 * WARP_SIZE_GGUF + 1) + kq1] = __vsubss4(ql1 | qh1, 0x20202020); + } + + const int blocks_per_tile_x_row = WARP_SIZE_GGUF / QI6_K; // == 1 if QK_K == 256 + const int kbxd = k % blocks_per_tile_x_row; // == 0 if QK_K == 256 + float* x_dmf = (float*)x_dm; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI6_K) { + int i = (i0 + i_offset * QI6_K + k / blocks_per_tile_x_row) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q6_K* bxi = bx0 + i * blocks_per_row + kbxd; + + x_dmf[i * (WARP_SIZE_GGUF / QI6_K) + i / QI6_K + kbxd] = __half2float(bxi->d); + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { + int i = (i0 + i_offset * 8 + k / (WARP_SIZE_GGUF / 8)) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q6_K* bxi = bx0 + i * blocks_per_row + (k % (WARP_SIZE_GGUF / 8)) / 4; + + x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + k % (WARP_SIZE_GGUF / 8)] = get_int_from_int8(bxi->scales, k % (QI6_K / 8)); + } +} + +static __device__ __forceinline__ float vec_dot_q6_K_q8_1_mul_mat( + const int* __restrict__ x_ql, + const half2* __restrict__ x_dm, + const int* __restrict__ x_qh, + const int* __restrict__ x_sc, + const int* __restrict__ y_qs, + const half2* __restrict__ y_ds, + const int& i, + const int& j, + const int& k) { + const float* x_dmf = (const float*)x_dm; + const float* y_df = (const float*)y_ds; + + const int8_t* sc = ((const int8_t*)&x_sc[i * (WARP_SIZE_GGUF / 8) + i / 8 + k / 8]); + + const int index_x = i * (QR6_K * WARP_SIZE_GGUF + 1) + QR6_K * k; + const int index_y = j * WARP_SIZE_GGUF + (QR6_K * k) % WARP_SIZE_GGUF; + return vec_dot_q6_K_q8_1_impl_mmq( + &x_ql[index_x], &y_qs[index_y], sc, x_dmf[i * (WARP_SIZE_GGUF / QI6_K) + i / QI6_K], &y_df[index_y / QI8_1]); +} + +static __device__ __forceinline__ float +vec_dot_iq2_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { + const block_iq2_xxs* bq2 = (const block_iq2_xxs*)vbq; + + const int ib32 = iqs; + const uint16_t* q2 = bq2->qs + 4 * ib32; + const uint8_t* aux8 = (const uint8_t*)q2; + const int8_t* q8 = bq8_1[ib32].qs; + uint32_t aux32 = q2[2] | (q2[3] << 16); + int sumi = 0; + for (int l = 0; l < 4; ++l) { + const uint8_t* grid = (const uint8_t*)(iq2xxs_grid + aux8[l]); + const uint8_t signs = ksigns_iq2xs[aux32 & 127]; + for (int j = 0; j < 8; ++j) { + sumi += q8[j] * grid[j] * (signs & kmask_iq2xs[j] ? -1 : 1); + } + q8 += 8; + aux32 >>= 7; + } + const float d = __half2float(bq2->d) * (0.5f + aux32) * __half2float(bq8_1[ib32].ds.x) * 0.25f; + return d * sumi; +} + +static __device__ __forceinline__ float +vec_dot_iq2_xs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { + const block_iq2_xs* bq2 = (const block_iq2_xs*)vbq; + + const int ib32 = iqs; + const uint16_t* q2 = bq2->qs + 4 * ib32; + const int8_t* q8 = bq8_1[ib32].qs; + const uint8_t ls1 = bq2->scales[ib32] & 0xf; + const uint8_t ls2 = bq2->scales[ib32] >> 4; + int sumi1 = 0; + for (int l = 0; l < 2; ++l) { + const uint8_t* grid = (const uint8_t*)(iq2xs_grid + (q2[l] & 511)); + const uint8_t signs = ksigns_iq2xs[q2[l] >> 9]; + for (int j = 0; j < 8; ++j) { + sumi1 += q8[j] * grid[j] * (signs & kmask_iq2xs[j] ? -1 : 1); + } + q8 += 8; + } + int sumi2 = 0; + for (int l = 2; l < 4; ++l) { + const uint8_t* grid = (const uint8_t*)(iq2xs_grid + (q2[l] & 511)); + const uint8_t signs = ksigns_iq2xs[q2[l] >> 9]; + for (int j = 0; j < 8; ++j) { + sumi2 += q8[j] * grid[j] * (signs & kmask_iq2xs[j] ? -1 : 1); + } + q8 += 8; + } + const float d = __half2float(bq2->d) * __half2float(bq8_1[ib32].ds.x) * 0.25f; + return d * ((0.5f + ls1) * sumi1 + (0.5f + ls2) * sumi2); +} + +static __device__ __forceinline__ float +vec_dot_iq2_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + const block_iq2_s* bq2 = (const block_iq2_s*)vbq; + + const int ib32 = iqs; + const int8_t* q8 = bq8_1[ib32].qs; + const uint8_t* signs = bq2->qs + QK_K / 8 + 4 * ib32; + const uint8_t ls1 = bq2->scales[ib32] & 0xf; + const uint8_t ls2 = bq2->scales[ib32] >> 4; + int sumi1 = 0; + for (int l = 0; l < 2; ++l) { + const uint32_t* grid = + (const uint32_t*)(iq2s_grid + (bq2->qs[4 * ib32 + l] | ((bq2->qh[ib32] << (8 - 2 * l)) & 0x300))); + const uint32_t signs0 = __vcmpeq4(((signs[l] & 0xf) * 0x01010101) & 0x08040201, 0x08040201); + const uint32_t signs1 = __vcmpeq4(((signs[l] >> 4) * 0x01010101) & 0x08040201, 0x08040201); + const int grid_l = __vsub4(grid[0] ^ signs0, signs0); + const int grid_h = __vsub4(grid[1] ^ signs1, signs1); + sumi1 = __dp4a(grid_l, *((const int*)q8 + 0), sumi1); + sumi1 = __dp4a(grid_h, *((const int*)q8 + 1), sumi1); + q8 += 8; + } + int sumi2 = 0; + for (int l = 2; l < 4; ++l) { + const uint32_t* grid = + (const uint32_t*)(iq2s_grid + (bq2->qs[4 * ib32 + l] | ((bq2->qh[ib32] << (8 - 2 * l)) & 0x300))); + const uint32_t signs0 = __vcmpeq4(((signs[l] & 0xf) * 0x01010101) & 0x08040201, 0x08040201); + const uint32_t signs1 = __vcmpeq4(((signs[l] >> 4) * 0x01010101) & 0x08040201, 0x08040201); + const int grid_l = __vsub4(grid[0] ^ signs0, signs0); + const int grid_h = __vsub4(grid[1] ^ signs1, signs1); + sumi2 = __dp4a(grid_l, *((const int*)q8 + 0), sumi2); + sumi2 = __dp4a(grid_h, *((const int*)q8 + 1), sumi2); + q8 += 8; + } + const float d = __half2float(bq2->d) * __low2float(bq8_1[ib32].ds) * 0.25f; + return d * ((0.5f + ls1) * sumi1 + (0.5f + ls2) * sumi2); +#endif +} + +static __device__ __forceinline__ float +vec_dot_iq3_xxs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + const block_iq3_xxs* bq2 = (const block_iq3_xxs*)vbq; + + const int ib32 = iqs; + const uint8_t* q3 = bq2->qs + 8 * ib32; + const uint16_t* gas = (const uint16_t*)(bq2->qs + QK_K / 4) + 2 * ib32; + const int8_t* q8 = bq8_1[ib32].qs; + uint32_t aux32 = gas[0] | (gas[1] << 16); + int sumi = 0; + for (int l = 0; l < 4; ++l) { + const uint32_t* grid1 = iq3xxs_grid + q3[2 * l + 0]; + const uint32_t* grid2 = iq3xxs_grid + q3[2 * l + 1]; + const uint32_t* signs = (const uint32_t*)(ksigns64 + (aux32 & 127)); + const int grid_l = __vsub4(grid1[0] ^ signs[0], signs[0]); + const int grid_h = __vsub4(grid2[0] ^ signs[1], signs[1]); + sumi = __dp4a(grid_l, *((int*)q8 + 0), sumi); + sumi = __dp4a(grid_h, *((int*)q8 + 1), sumi); + q8 += 8; + aux32 >>= 7; + } + const float d = __half2float(bq2->d) * (0.5f + aux32) * __low2float(bq8_1[ib32].ds) * 0.5f; + return d * sumi; +#endif +} + +static __device__ __forceinline__ float +vec_dot_iq3_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + const block_iq3_s* bq2 = (const block_iq3_s*)vbq; + + const int ib32 = iqs; + const uint8_t* qs = bq2->qs + 8 * ib32; + const int8_t* q8 = bq8_1[ib32].qs; + int sumi = 0; + for (int l = 0; l < 4; ++l) { + const uint32_t* grid1 = iq3xs_grid + (qs[2 * l + 0] | ((bq2->qh[ib32] << (8 - 2 * l)) & 256)); + const uint32_t* grid2 = iq3xs_grid + (qs[2 * l + 1] | ((bq2->qh[ib32] << (7 - 2 * l)) & 256)); + uint32_t signs0 = __vcmpeq4(((bq2->signs[4 * ib32 + l] & 0xf) * 0x01010101) & 0x08040201, 0x08040201); + uint32_t signs1 = __vcmpeq4(((bq2->signs[4 * ib32 + l] >> 4) * 0x01010101) & 0x08040201, 0x08040201); + const int grid_l = __vsub4(grid1[0] ^ signs0, signs0); + const int grid_h = __vsub4(grid2[0] ^ signs1, signs1); + sumi = __dp4a(grid_l, *((int*)q8 + 0), sumi); + sumi = __dp4a(grid_h, *((int*)q8 + 1), sumi); + q8 += 8; + } + const float d = __half2float(bq2->d) * (0.5f + ((bq2->scales[ib32 / 2] >> 4 * (ib32 % 2)) & 0xf)) * + __low2float(bq8_1[ib32].ds) * 0.5f; + return d * sumi; +#endif +} + +static __device__ __forceinline__ float +vec_dot_iq1_s_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + const block_iq1_s* bq1 = (const block_iq1_s*)vbq; + + const int qs_packed = get_int_b2(bq1->qs, iqs); + const uint8_t* qs = (const uint8_t*)&qs_packed; + + const int qh = bq1->qh[iqs]; + + int sumi = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int grid = iq1s_grid_gpu[qs[l0 / 2] | (((qh >> 3 * (l0 / 2)) & 0x07) << 8)]; + + const int grid0 = (grid >> 0) & 0x0F0F0F0F; + const int grid1 = (grid >> 4) & 0x0F0F0F0F; + + const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1); + + sumi = __dp4a(grid0, u0, sumi); + sumi = __dp4a(grid1, u1, sumi); + } + + const float d1q = __half2float(bq1->d) * (((qh >> 11) & 0x0E) + 1); + const float delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f * IQ1S_DELTA / 0x8000); + const float2 ds = __half22float2(bq8_1[iqs].ds); + return d1q * (ds.x * sumi + ds.y * delta); +#endif +} + +static __device__ __forceinline__ float +vec_dot_iq1_m_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + + const block_iq1_m* bq1 = (const block_iq1_m*)vbq; + + const int qs_packed = get_int_b4(bq1->qs, iqs); + const uint8_t* qs = (const uint8_t*)&qs_packed; + + int sumi[2] = {0}; + float sumf[2] = {0.0f}; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int qhl = bq1->qh[2 * iqs + l0 / 4] >> (4 * ((l0 / 2) % 2)); + + const int grid = iq1s_grid_gpu[qs[l0 / 2] | ((qhl & 0x07) << 8)]; + + const int grid0 = (grid >> 0) & 0x0F0F0F0F; + const int grid1 = (grid >> 4) & 0x0F0F0F0F; + + const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1); + + sumi[l0 / 4] = __dp4a(grid0, u0, sumi[l0 / 4]); + sumi[l0 / 4] = __dp4a(grid1, u1, sumi[l0 / 4]); + + const float delta = -1.0f + IQ1M_DELTA - (qhl & 0x08) * (2.0f * IQ1M_DELTA / 0x08); + int sumy = 0; + sumy = __dp4a(u0, 0x01010101, sumy); + sumy = __dp4a(u1, 0x01010101, sumy); + sumf[l0 / 4] += delta * sumy; + } + + const uint16_t* sc = (const uint16_t*)bq1->scales; + + iq1m_scale_t scale; + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00F0) | ((sc[2] >> 4) & 0x0F00) | (sc[3] & 0xF000); + const float d = __half2float(scale.f16) * __low2float(bq8_1[iqs].ds); + + const int tmp = sc[iqs / 2] >> (6 * (iqs % 2)); + const int sc0 = 2 * ((tmp >> 0) & 0x07) + 1; + const int sc1 = 2 * ((tmp >> 3) & 0x07) + 1; + return d * ((sumi[0] + sumf[0]) * sc0 + (sumi[1] + sumf[1]) * sc1); +#endif +} + +static __device__ __forceinline__ void +get_int_from_table_16(const uint32_t& q4, const uint8_t* values, int& val1, int& val2) { + uint32_t aux32; + const uint8_t* q8 = (const uint8_t*)&aux32; + aux32 = q4 & 0x0f0f0f0f; + uint16_t v1 = values[q8[0]] | (values[q8[1]] << 8); + uint16_t v2 = values[q8[2]] | (values[q8[3]] << 8); + val1 = v1 | (v2 << 16); + aux32 = (q4 >> 4) & 0x0f0f0f0f; + v1 = values[q8[0]] | (values[q8[1]] << 8); + v2 = values[q8[2]] | (values[q8[3]] << 8); + val2 = v1 | (v2 << 16); +} + +static __device__ __forceinline__ float +vec_dot_iq4_nl_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + + const block_iq4_nl* bq = (const block_iq4_nl*)vbq; + + const uint16_t* q4 = (const uint16_t*)bq->qs + 2 * iqs; + const int32_t* q8 = (const int32_t*)bq8_1->qs + iqs; + + const uint8_t* values = (const uint8_t*)kvalues_iq4nl; + + int v1, v2; + int sumi1 = 0, sumi2 = 0; + for (int l = 0; l < VDR_Q4_0_Q8_1_MMVQ; ++l) { + const uint32_t aux = q4[2 * l] | (q4[2 * l + 1] << 16); + get_int_from_table_16(aux, values, v1, v2); + sumi1 = __dp4a(v1, q8[l + 0], sumi1); + sumi2 = __dp4a(v2, q8[l + 4], sumi2); + } + const float d = __half2float(bq->d) * __low2float(bq8_1->ds); + return d * (sumi1 + sumi2); +#endif +} + +static __device__ __forceinline__ float +vec_dot_iq4_xs_q8_1(const void* __restrict__ vbq, const block_q8_1* __restrict__ bq8_1, const int& iqs) { +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 610 || defined USE_ROCM || defined USE_MUSA + const block_iq4_xs* bq4 = (const block_iq4_xs*)vbq; + const uint8_t* values = (const uint8_t*)kvalues_iq4nl; + + // iqs is 0...7 + const int ib32 = iqs; + const int32_t* q8 = (const int*)bq8_1[ib32].qs; + const uint32_t* q4 = (const uint32_t*)bq4->qs + 4 * ib32; + const int8_t ls = ((bq4->scales_l[ib32 / 2] >> 4 * (ib32 % 2)) & 0xf) | (((bq4->scales_h >> 2 * ib32) & 3) << 4); + const float d = __half2float(bq4->d) * (ls - 32) * __low2float(bq8_1[ib32].ds); + int v1, v2; + int sumi1 = 0, sumi2 = 0; + for (int j = 0; j < 4; ++j) { + get_int_from_table_16(q4[j], values, v1, v2); + sumi1 = __dp4a(v1, q8[j + 0], sumi1); + sumi2 = __dp4a(v2, q8[j + 4], sumi2); + } + return d * (sumi1 + sumi2); +#endif +} diff --git a/python/freetoken/kernel/gguf.py b/python/freetoken/kernel/gguf.py index f05b2def3..dce0c3819 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -22,6 +22,244 @@ _CSRC = pathlib.Path(__file__).parent / "csrc" / "gguf" +_GGML_QUANT_NAMES = { + 2: "Q4_0", 8: "Q8_0", 12: "Q4_K", 13: "Q5_K", 14: "Q6_K" +} +_GGUF_QUANT_TYPES = frozenset(_GGML_QUANT_NAMES) +_DISPATCH_COUNTS: dict[tuple, int] = {} +# Algorithm subset pinned to llama.cpp b10434 (commit 7e4c0a968). The local HIP ABI +# wrapper remains separate, so provenance does not imply full source replacement. +_GGUF_SOURCE_VERSION = "llama.cpp-7e4c0a968-q4q5q6q8-subset-v1" +_GGUF_ABI_VERSION = "moe-abi-v2" +_MMVQ_BS1_ABI_VERSION = "llama.cpp-b10434-mmvq-bs1-v1" +_CANDIDATE_MODULE_NAME = "freetoken_gguf_moe_gfx1100_v6" +_GGUF_MOE_ABI: dict[str, str] = { + "ggml_moe_a8_vec": _GGUF_ABI_VERSION, + "ggml_moe_a8_vec_strided": _GGUF_ABI_VERSION, +} + + +def register_gguf_moe_abi(callables: dict[str, str], *, abi_version: str) -> None: + """Register extension-bound MoE callables; metadata never advertises guesses.""" + if abi_version != _GGUF_ABI_VERSION: + raise ValueError(f"unsupported GGUF MoE ABI {abi_version!r}") + for name, version in callables.items(): + if version == _GGUF_ABI_VERSION: + _GGUF_MOE_ABI[name] = version + + +def mmvq_bs1_workspace_bytes(hidden: int, rows: int, channels: int) -> int: + """Exact caller-owned workspace formula for the b10434 single-token ABI.""" + hidden, rows, channels = int(hidden), int(rows), int(channels) + if min(hidden, rows, channels) <= 0: + raise ValueError("MMVQ b10434 shape must be positive") + if hidden % 32: + raise ValueError("MMVQ b10434 hidden dimension must be divisible by 32") + align = lambda value: (value + 255) // 256 * 256 + return align(hidden // 32 * 36) + align(channels * rows * 4) + + +def validate_mmvq_bs1_workspace(hidden: int, rows: int, channels: int, workspace_bytes: int) -> None: + required = mmvq_bs1_workspace_bytes(hidden, rows, channels) + if int(workspace_bytes) < required: + raise ValueError(f"MMVQ workspace too small: {workspace_bytes} < {required}") +_GGUF_BLOCK_SIZE = {2: 32, 8: 32, 12: 256, 13: 256, 14: 256} +_RDNA3_MOE_MMVQ_MAX = {2: 4, 8: 4, 12: 4, 13: 4, 14: 4} + + +def _runtime_backend() -> str: + if torch.version.hip is not None: + return "rocm" + if torch.version.cuda is not None: + return "cuda" + return "cpu" + + +def _runtime_arch() -> str | None: + forced = os.environ.get("FREETOKEN_KERNEL_CACHE_GFX") + if forced: + return forced + if _runtime_backend() == "rocm" and torch.cuda.is_available(): + try: + return torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName + except (AttributeError, RuntimeError): + pass + return None + + +def gguf_runtime_metadata() -> dict: + """Report GGUF JIT/runtime selection without compiling or mutating state.""" + backend = _runtime_backend() + arch = _runtime_arch() + flags = ["-O3"] + if backend == "rocm": + flags.extend(["USE_HIP=1", "USE_ROCM=1", f"offload-arch={arch or 'unknown'}"]) + elif backend == "cuda": + flags.append("expt-relaxed-constexpr") + return { + "backend": backend, + "arch": arch, + "device": torch.cuda.get_device_name(torch.cuda.current_device()) + if torch.cuda.is_available() + else None, + "torch": torch.__version__, + "rocm": torch.version.hip, + "compile_flags": flags, + "source_version": os.environ.get("FREETOKEN_GGUF_SOURCE_VERSION", _GGUF_SOURCE_VERSION), + "source": "gguf_kernel.cu", + } + + +def _record_dispatch(report: dict) -> None: + if os.environ.get("FREETOKEN_GGUF_DISPATCH_TRACE", "").lower() not in {"1", "true", "yes", "on"}: + return + key = tuple( + report[field] + for field in ("backend", "arch", "quant_type", "op", "implementation", "rows", "cols", "tokens") + ) + _DISPATCH_COUNTS[key] = _DISPATCH_COUNTS.get(key, 0) + 1 + + +def gguf_dispatch_report() -> list[dict]: + """Return aggregate dispatch observations collected when trace is enabled.""" + fields = ("backend", "arch", "quant_type", "op", "implementation", "rows", "cols", "tokens") + return [dict(zip(fields, key), calls=calls) for key, calls in _DISPATCH_COUNTS.items()] + + +def _arch_family(arch: str | None) -> str: + """Map GPU target names to b10434's kernel tuning families.""" + value = (arch or "").lower() + if value.startswith(("gfx11", "gfx12")): + return "rdna3" if value.startswith("gfx11") else "rdna4" + if value.startswith("gfx10"): + return "rdna2" + if value.startswith("gfx9"): + return "cdna" + if value.startswith("sm"): + return "nvidia" + return "generic" + + +def gguf_dispatch( + op: str, + quant_type: int, + rows: int, + cols: int, + tokens: int, + arch: str | None, + impl: str | None = None, +) -> dict: + """Resolve observable GGUF operation family without launching a kernel.""" + requested = (impl or os.environ.get("FREETOKEN_GGUF_MOE_IMPL", "legacy")).strip().lower() + if requested == "gfx1100": + requested = "rdna3_mmid" + if requested not in { + "auto", "legacy", "mmvq", "mmq", "rdna3_mmid", "rdna3_mmvdq", "grouped_mmq" + }: + raise ValueError(f"unsupported GGUF implementation {requested!r}") + # FREETOKEN_GGUF_MOE_IMPL is a MoE-only selector. Dense GGUF projections + # (including lm_head) must retain normal shape policy; otherwise forcing a + # MoE candidate silently sends dense layers through full dequantization. + if impl is None and op not in {"moe_decode", "moe", "moe_prefill", "grouped_prefill"}: + requested = "auto" + backend = _runtime_backend() + runtime = gguf_runtime_metadata() + report = { + "backend": backend, + "arch": arch, + "quant_type": _GGML_QUANT_NAMES.get(quant_type, f"unknown:{quant_type}"), + "op": op, + "rows": rows, + "cols": cols, + "tokens": tokens, + "implementation": "unsupported", + "requested": requested, + "compile_flags": runtime["compile_flags"], + "source_version": runtime["source_version"], + "abi_version": _GGUF_ABI_VERSION, + "library": "gguf-jit", + "stream": "current", + "id_space": "declared-by-caller", + "callable": None, + "reason": None, + } + if rows <= 0 or cols <= 0 or tokens <= 0: + report["reason"] = "non-positive shape" + elif quant_type not in _GGUF_QUANT_TYPES: + report["reason"] = "unsupported quantization type" + elif cols % _GGUF_BLOCK_SIZE[quant_type]: + report["reason"] = f"K dimension {cols} is not aligned to {_GGUF_BLOCK_SIZE[quant_type]}" + elif backend not in {"cuda", "rocm"}: + report["reason"] = "GPU backend unavailable" + elif backend == "rocm" and arch and arch.startswith("sm"): + report["reason"] = "NVIDIA architecture requested on ROCm" + elif backend == "cuda" and arch and arch.startswith("gfx"): + report["reason"] = "AMD architecture requested on CUDA" + elif op == "moe_prefill": + # b10434 uses the dedicated grouped MUL_MAT_ID path once the batch exceeds the + # architecture/type MMVQ window. The old local HIP grouped ABI has produced + # launch failures on gfx1100 for the Qwen3.6 prefill shape, so keep this opt-in + # until its exact model-shape A/B and error-free replay gate pass. The vector path + # is slower but proven and remains the fail-closed default. + family = _arch_family(arch) + limit = _RDNA3_MOE_MMVQ_MAX.get(quant_type, 8) if family == "rdna3" else 8 + grouped = os.environ.get("FREETOKEN_GGUF_GROUPED_PREFILL", "0").strip().lower() in { + "1", "true", "yes", "on" + } + report["implementation"] = "ggml_moe_a8" if tokens > limit and grouped else "ggml_moe_a8_vec" + report["callable"] = report["implementation"] + report["reason"] = ( + "multi-token grouped path" + if tokens > limit and grouped + else ("grouped path disabled after gfx1100 launch failure" if tokens > limit else None) + ) + elif op == "grouped_prefill": + report["implementation"] = "ggml_moe_a8" + report["callable"] = report["implementation"] + elif op in {"moe_decode", "moe"}: + candidate_callable = { + "rdna3_mmid": "ggml_moe_mmvq_id", + "rdna3_mmvdq": "ggml_moe_mmvdq_id", + "grouped_mmq": "ggml_moe_mmq_id_strided", + }.get(requested) + candidate_error = None + if requested in {"rdna3_mmid", "rdna3_mmvdq"} and backend == "rocm" and arch == "gfx1100": + try: + ensure_gguf_moe_candidate_ready() + except Exception as exc: + candidate_error = exc + report["reason"] = f"{requested} compile/self-test failed: {type(exc).__name__}" + if candidate_error is not None and requested not in {"auto", "legacy"}: + raise RuntimeError(report["reason"]) from candidate_error + # Candidate bindings are registered only after extension load. Until + # then explicit opt-in fails closed to the proven legacy callable. + shape_supported = requested != "rdna3_mmvdq" or quant_type in {12, 13, 14} + available = bool( + shape_supported + and candidate_callable + and _GGUF_MOE_ABI.get(candidate_callable) == _GGUF_ABI_VERSION + ) + if available: + report["implementation"] = requested + report["callable"] = candidate_callable + else: + report["implementation"] = "ggml_moe_a8_vec" + report["callable"] = "ggml_moe_a8_vec" + if requested not in {"auto", "legacy"}: + raise RuntimeError( + f"{requested} ABI callable unavailable; forced mode refuses legacy fallback" + ) + elif op in {"dense", "linear", "lm_head"}: + family = "mmvq" if tokens <= 8 else "mmq" + if requested not in {"auto", "legacy"} and requested != family: + report["reason"] = f"forced {requested} conflicts with shape policy {family}" + else: + report["implementation"] = "ggml_mul_mat_vec_a8" if family == "mmvq" else "ggml_mul_mat_a8" + else: + report["reason"] = "unsupported operation" + _record_dispatch(report) + return report + def _host_compiler() -> str | None: """A host compiler nvcc + libtorch headers accept. @@ -83,6 +321,24 @@ def _clear_stale_jit_lock(module_name: str) -> None: pass +def _rocm_jit_source(module_name: str, source: pathlib.Path) -> pathlib.Path: + """Prepare ROCm source in build dir without letting PyTorch rewrite repo files. + + PyTorch's HIP extension path runs hipify on CUDA sources and recursively rewrites + included headers. That can overwrite checked-in ``*_hip`` files when an absolute + repository source is supplied. A cache-local ``.cu`` copy keeps hipify's source + and generated output in the extension cache; ``-I`` resolves the tracked HIP + headers selected by the source. + """ + from torch.utils.cpp_extension import _get_build_directory + + build_dir = pathlib.Path(_get_build_directory(module_name, False)) + build_dir.mkdir(parents=True, exist_ok=True) + target = build_dir / source.name + target.write_text(source.read_text()) + return target + + @functools.cache def _module(): from torch.utils.cpp_extension import load @@ -96,11 +352,19 @@ def _module(): os.environ.setdefault("PYTORCH_ROCM_ARCH", gfx) extra_cuda_cflags = [ "-O3", f"--offload-arch={gfx}", "-DUSE_HIP=1", "-DUSE_ROCM=1", + "-DFREETOKEN_GGUF_SOURCE_VERSION=2", + f"-I{_CSRC}", ] os.environ.pop("CXX", None) os.environ.pop("CC", None) + source = _rocm_jit_source("freetoken_gguf_kernels", _CSRC / "gguf_kernel.cu") + include_paths = [] else: - extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"] + extra_cuda_cflags = [ + "-O3", + "--expt-relaxed-constexpr", + "-DFREETOKEN_GGUF_SOURCE_VERSION=2", + ] host_cxx = _host_compiler() if host_cxx is not None: # Point both nvcc's host pass (-ccbin) and torch's C++ compile (CXX) at a @@ -110,6 +374,8 @@ def _module(): extra_cuda_cflags += ["-ccbin", cxx_path] os.environ["CXX"] = cxx_path os.environ["CC"] = _c_compiler_for(cxx_path) + source = _CSRC / "gguf_kernel.cu" + include_paths = [str(_CSRC)] # Rows-per-warp for the MMVQ/MoE-vec launches (Inc 4, .plans/rocm-perf-parity). # Overridable for tuning/AB; the JIT cache keys on the cflags, so a changed value @@ -125,13 +391,196 @@ def _module(): # 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. _clear_stale_jit_lock("freetoken_gguf_kernels") - return load( + module = load( name="freetoken_gguf_kernels", - sources=[str(_CSRC / "gguf_kernel.cu")], - extra_include_paths=[str(_CSRC)], + sources=[str(source)], + extra_include_paths=include_paths, + extra_cuda_cflags=extra_cuda_cflags, + verbose=False, + ) + register_gguf_moe_abi( + { + name: _GGUF_ABI_VERSION + for name in ( + "ggml_moe_a8_vec", + "ggml_moe_a8_vec_strided", + "ggml_moe_a8_vec_workspace", + "ggml_moe_a8_vec_strided_workspace", + ) + if hasattr(module, name) + }, + abi_version=_GGUF_ABI_VERSION, + ) + return module + + +_candidate_ready = False + + +def _gguf_moe_impl() -> str: + impl = os.environ.get("FREETOKEN_GGUF_MOE_IMPL", "legacy").strip().lower() + if impl not in { + "auto", "legacy", "gfx1100", "rdna3_mmid", "rdna3_mmvdq", "grouped_mmq" + }: + raise ValueError( + f"FREETOKEN_GGUF_MOE_IMPL={impl!r}: expected auto, legacy, rdna3_mmid, " + "rdna3_mmvdq, grouped_mmq, or gfx1100" + ) + return impl + + +def _gfx1100_supported() -> bool: + if torch.version.hip is None or not torch.cuda.is_available(): + return False + try: + return torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName == "gfx1100" + except AttributeError: + return False + + +@functools.cache +def _candidate_module(): + from torch.utils.cpp_extension import load + + gfx = os.getenv("FREETOKEN_KERNEL_CACHE_GFX", "gfx1100") + if gfx != "gfx1100": + raise RuntimeError(f"gfx1100 candidate requires FREETOKEN_KERNEL_CACHE_GFX=gfx1100, got {gfx!r}") + os.environ.setdefault("PYTORCH_ROCM_ARCH", gfx) + os.environ.pop("CXX", None) + os.environ.pop("CC", None) + extra_cuda_cflags = [ + "-O3", + f"--offload-arch={gfx}", + "-DUSE_HIP=1", + "-DUSE_ROCM=1", + "-DFREETOKEN_GGUF_MOE_SOURCE_VERSION=2", + "-DFREETOKEN_GGUF_SOURCE_VERSION=2", + f"-I{_CSRC}", + ] + mmv_y = os.getenv("FREETOKEN_GGUF_MMV_Y", "1").strip() + if not mmv_y.isdigit() or int(mmv_y) not in (1, 2, 4, 8): + raise ValueError( + f"FREETOKEN_GGUF_MMV_Y={mmv_y!r}: expected one of 1, 2, 4, 8" + ) + extra_cuda_cflags.append(f"-DGGML_CUDA_MMV_Y={mmv_y}") + _clear_stale_jit_lock(_CANDIDATE_MODULE_NAME) + source = _rocm_jit_source( + _CANDIDATE_MODULE_NAME, _CSRC / "gguf_moe_gfx1100.cu" + ) + abi_source = _rocm_jit_source( + _CANDIDATE_MODULE_NAME, _CSRC / "gguf_b10434_kernel.cu" + ) + extra_cuda_cflags.append("-DFREETOKEN_GGUF_NO_PYBIND=1") + module = load( + name=_CANDIDATE_MODULE_NAME, + sources=[str(source), str(abi_source)], + extra_include_paths=[], extra_cuda_cflags=extra_cuda_cflags, verbose=False, ) + register_gguf_moe_abi( + { + name: _GGUF_ABI_VERSION + for name in ( + "ggml_moe_mmvq_id", "ggml_moe_mmvq_id_workspace", + "ggml_moe_mmvdq_id", "ggml_moe_mmvdq_id_workspace", + "ggml_moe_gate_up_swiglu_id", "ggml_moe_gate_up_swiglu_id_workspace", + "mmvq_bs1", "mmvq_bs1_workspace_bytes", + ) + if hasattr(module, name) + }, + abi_version=_GGUF_ABI_VERSION, + ) + return module + + +def _candidate_self_test(module) -> None: + """Exercise both candidate quant paths and synchronize before graph capture.""" + device = torch.device("cuda") + generator = torch.Generator(device="cpu").manual_seed(1100) + ids = torch.arange(8, dtype=torch.int32).reshape(1, 8).to(device) + x = torch.randn(1, 256, generator=generator, dtype=torch.bfloat16).to(device) + q4 = torch.zeros((8, 16, 144), dtype=torch.uint8, device=device) + q4[..., :4] = torch.tensor([128, 63, 128, 63], dtype=torch.uint8, device=device) + q4[..., 4:16] = torch.randint(1, 64, (8, 16, 12), generator=generator, dtype=torch.uint8).to(device) + q4[..., 16:] = torch.randint(0, 255, (8, 16, 128), generator=generator, dtype=torch.uint8).to(device) + q4_out = module.ggml_moe_a8_vec_gfx1100(x, q4, ids, 8, 12, 16, 1) + inter = torch.randn(8, 128, generator=generator, dtype=torch.bfloat16).to(device) + q8_blocks = torch.zeros((8, 16, 4, 34), dtype=torch.uint8, device=device) + q8_blocks[..., :2] = torch.tensor([128, 63], dtype=torch.uint8, device=device) + q8_blocks[..., 2:] = torch.randint( + 0, 255, (8, 16, 4, 32), generator=generator, dtype=torch.uint8 + ).to(device) + q8 = q8_blocks.reshape(8, 16, 136) + q8_out = module.ggml_moe_a8_vec_gfx1100(inter, q8, ids, 1, 8, 16, 8) + route_ids = ids.reshape(-1, 1) + id_out = torch.empty((8, 16), dtype=inter.dtype, device=device) + id_qx = torch.empty((8, 144), dtype=torch.int32, device=device) + id_out = module.ggml_moe_mmvq_id_workspace( + inter, q8, route_ids, 1, 8, 16, 8, + int(q8.stride(0)), int(q8.stride(1)), "slot", id_out, id_qx + ) + mmvdq_out = module.ggml_moe_mmvdq_id( + x, q4, ids, 8, 12, 16, 1, + int(q4.stride(0)), int(q4.stride(1)), "slot" + ) + fused_out = module.ggml_moe_gate_up_swiglu_id_workspace( + x, q4, ids, 8, 8, 1, + int(q4.stride(0)), int(q4.stride(1)), "slot", + torch.empty((8, 8), dtype=x.dtype, device=device), + torch.empty((1, 144), dtype=torch.int32, device=device), + ) + # Exercise the pinned b10434 caller-owned ABI separately from the model-facing + # BF16 candidate output. This catches module binding and workspace slicing before + # graph capture, without allocating inside the measured model path. + abi_x = torch.randn(1, 512, generator=generator, dtype=torch.bfloat16).to(device) + abi_q4 = torch.zeros((8, 16, 288), dtype=torch.uint8, device=device) + abi_q4[..., :4] = torch.tensor([128, 63, 128, 63], dtype=torch.uint8, device=device) + abi_q4[..., 4:] = 1 + abi_workspace = torch.empty( + mmvq_bs1_workspace_bytes(512, 16, 8), dtype=torch.uint8, device=device + ) + abi_out = torch.empty((8, 16), dtype=torch.float32, device=device) + abi_result = module.mmvq_bs1(abi_x, abi_q4, abi_out, abi_workspace, 12, 16, 8, ids) + torch.cuda.synchronize(device) + if ( + not torch.isfinite(q4_out).all() + or not torch.isfinite(q8_out).all() + or not torch.isfinite(id_out).all() + or not torch.isfinite(mmvdq_out).all() + or not torch.isfinite(fused_out).all() + or not torch.isfinite(abi_result).all() + ): + raise RuntimeError("gfx1100 candidate self-test produced non-finite output") + + +def ensure_gguf_moe_candidate_ready() -> bool: + """Compile and validate forced candidate before graph capture.""" + global _candidate_ready + impl = _gguf_moe_impl() + # Candidate remains forced-only until a later promotion is backed by measured evidence. + if impl not in {"gfx1100", "rdna3_mmid", "rdna3_mmvdq"}: + return False + if _candidate_ready: + return True + try: + module = _candidate_module() + _candidate_self_test(module) + _candidate_ready = True + return True + except Exception as exc: + raise RuntimeError("forced gfx1100 GGUF MoE candidate failed compile/self-test") from exc + + +def _moe_module(): + impl = _gguf_moe_impl() + if impl not in {"gfx1100", "rdna3_mmid", "rdna3_mmvdq"}: + return None + if not _gfx1100_supported(): + raise RuntimeError("forced gfx1100 GGUF MoE candidate requires ROCm gfx1100") + if not ensure_gguf_moe_candidate_ready(): + return None + return _candidate_module() # ---- thin typed wrappers (signatures mirror sgl_kernel.quantization.gguf) ---- @@ -184,9 +633,174 @@ def ggml_moe_a8_vec( quant_type: int, row: int, tokens: int, + output: torch.Tensor | None = None, ) -> 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) + candidate = _moe_module() + if output is not None: + if candidate is not None: + raise ValueError("reusable output is unsupported by gfx1100 candidate") + return _module().ggml_moe_a8_vec( + x, weight, topk_ids, top_k, quant_type, row, tokens, output + ) + if candidate is not None: + if quant_type in (8, 12): + return candidate.ggml_moe_a8_vec_gfx1100( + x, weight, topk_ids, top_k, quant_type, row, tokens + ) + # The candidate's ID-aware ABI owns native Q5_K/Q6_K rows and explicit + # strides; never send those lanes through its older compact Q4/Q8 helper. + return ggml_moe_mmvq_id( + x, weight, topk_ids, top_k, quant_type, row, tokens, + int(weight.stride(0)), int(weight.stride(1)), "slot", + ) + return _module().ggml_moe_a8_vec( + x, weight, topk_ids, top_k, quant_type, row, tokens + ) + + +def ggml_moe_a8_vec_strided( + x: torch.Tensor, + weight: torch.Tensor, + topk_ids: torch.Tensor, + top_k: int, + quant_type: int, + row: int, + tokens: int, + expert_stride_bytes: int, + row_stride_bytes: int, + output: torch.Tensor | None = None, +) -> torch.Tensor: + """Native Q5_K/Q6_K MoE GEMV over rows padded to a uniform Q6_K stride.""" + if output is None: + return _module().ggml_moe_a8_vec_strided( + x, weight, topk_ids, top_k, quant_type, row, tokens, + expert_stride_bytes, row_stride_bytes + ) + return _module().ggml_moe_a8_vec_strided( + x, weight, topk_ids, top_k, quant_type, row, tokens, + expert_stride_bytes, row_stride_bytes, output + ) + + +def ggml_moe_mmvq_id( + x: torch.Tensor, + weight: torch.Tensor, + topk_ids: torch.Tensor, + top_k: int, + quant_type: int, + row: int, + tokens: int, + expert_stride_bytes: int, + row_stride_bytes: int, + id_space: str, + output: torch.Tensor | None = None, + quant_x: torch.Tensor | None = None, +) -> torch.Tensor: + """Opt-in gfx1100 ID-aware MMVQ over raw or cache-slot IDs.""" + candidate = _moe_module() + if candidate is None or not hasattr(candidate, "ggml_moe_mmvq_id"): + raise RuntimeError("RDNA3 ID-aware GGUF MoE ABI is unavailable") + args = ( + x, weight, topk_ids, top_k, quant_type, row, tokens, + expert_stride_bytes, row_stride_bytes, id_space, + ) + if output is not None and quant_x is not None: + return candidate.ggml_moe_mmvq_id_workspace(*args, output, quant_x) + if output is not None or quant_x is not None: + raise ValueError("ggml_moe_mmvq_id requires output and quant_x together") + return candidate.ggml_moe_mmvq_id(*args) + + +def ggml_moe_mmvdq_id( + x: torch.Tensor, + weight: torch.Tensor, + topk_ids: torch.Tensor, + top_k: int, + quant_type: int, + row: int, + tokens: int, + expert_stride_bytes: int, + row_stride_bytes: int, + id_space: str, + output: torch.Tensor | None = None, +) -> torch.Tensor: + """Opt-in gfx1100 direct-float MMVDQ over native Q4_K/Q5_K/Q6_K rows.""" + candidate = _moe_module() + if candidate is None or not hasattr(candidate, "ggml_moe_mmvdq_id"): + raise RuntimeError("RDNA3 direct-float GGUF MoE ABI is unavailable") + args = ( + x, weight, topk_ids, top_k, quant_type, row, tokens, + expert_stride_bytes, row_stride_bytes, id_space, + ) + if output is None: + return candidate.ggml_moe_mmvdq_id(*args) + return candidate.ggml_moe_mmvdq_id_workspace(*args, output) + + +def ggml_moe_gate_up_swiglu_id( + x: torch.Tensor, + weight: torch.Tensor, + topk_ids: torch.Tensor, + top_k: int, + nrows: int, + tokens: int, + expert_stride_bytes: int, + row_stride_bytes: int, + id_space: str, + output: torch.Tensor | None = None, + quant_x: torch.Tensor | None = None, +) -> torch.Tensor: + """Opt-in gfx1100 fused Q4_K gate/up plus SwiGLU decode operation.""" + candidate = _moe_module() + if candidate is None or not hasattr(candidate, "ggml_moe_gate_up_swiglu_id"): + raise RuntimeError("RDNA3 fused gate/up GGUF MoE ABI is unavailable") + args = ( + x, weight, topk_ids, top_k, nrows, tokens, + expert_stride_bytes, row_stride_bytes, id_space, + ) + if output is not None and quant_x is not None: + return candidate.ggml_moe_gate_up_swiglu_id_workspace(*args, output, quant_x) + if output is not None or quant_x is not None: + raise ValueError("ggml_moe_gate_up_swiglu_id requires output and quant_x together") + return candidate.ggml_moe_gate_up_swiglu_id(*args) + + +def ggml_moe_a8_vec_workspace( + x: torch.Tensor, + weight: torch.Tensor, + topk_ids: torch.Tensor, + top_k: int, + quant_type: int, + row: int, + tokens: int, + output: torch.Tensor, + quant_x: torch.Tensor, +) -> torch.Tensor: + """Legacy MMVQ using caller-owned output and Q8_1 scratch tensors.""" + return _module().ggml_moe_a8_vec_workspace( + x, weight, topk_ids, top_k, quant_type, row, tokens, output, quant_x + ) + + +def ggml_moe_a8_vec_strided_workspace( + x: torch.Tensor, + weight: torch.Tensor, + topk_ids: torch.Tensor, + top_k: int, + quant_type: int, + row: int, + tokens: int, + expert_stride_bytes: int, + row_stride_bytes: int, + output: torch.Tensor, + quant_x: torch.Tensor, +) -> torch.Tensor: + """Strided native MMVQ using caller-owned output and Q8_1 scratch.""" + return _module().ggml_moe_a8_vec_strided_workspace( + x, weight, topk_ids, top_k, quant_type, row, tokens, + expert_stride_bytes, row_stride_bytes, output, quant_x + ) def ggml_moe_get_block_size(quant_type: int) -> int: @@ -194,10 +808,23 @@ def ggml_moe_get_block_size(quant_type: int) -> int: __all__ = [ + "mmvq_bs1_workspace_bytes", + "validate_mmvq_bs1_workspace", + "gguf_runtime_metadata", + "gguf_dispatch", + "gguf_dispatch_report", + "register_gguf_moe_abi", "ggml_dequantize", "ggml_mul_mat_vec_a8", "ggml_mul_mat_a8", "ggml_moe_a8", "ggml_moe_a8_vec", + "ggml_moe_a8_vec_strided", + "ggml_moe_a8_vec_workspace", + "ggml_moe_a8_vec_strided_workspace", + "ggml_moe_mmvq_id", + "ggml_moe_mmvdq_id", + "ggml_moe_gate_up_swiglu_id", "ggml_moe_get_block_size", + "ensure_gguf_moe_candidate_ready", ] diff --git a/python/freetoken/kernel/moe_impl.py b/python/freetoken/kernel/moe_impl.py index ef569fa2e..c5881296b 100644 --- a/python/freetoken/kernel/moe_impl.py +++ b/python/freetoken/kernel/moe_impl.py @@ -197,6 +197,32 @@ def moe_sum_reduce_triton(input: torch.Tensor, output: torch.Tensor) -> None: ) +def moe_weighted_sum_reduce_triton( + input: torch.Tensor, weights: torch.Tensor, output: torch.Tensor +) -> None: + """Graph-safe route weighting plus deterministic reduction in one launch.""" + import triton + + from .triton.fused_moe import moe_weighted_sum_reduce_kernel + + assert input.is_contiguous() and output.is_contiguous() + assert weights.ndim == 2 and weights.stride(1) == 1 + token_num, topk_num, hidden_dim = input.shape + assert weights.shape == (token_num, topk_num) + assert output.shape == (token_num, hidden_dim) + block_m = 1 if token_num <= 16 else 2 + block_dim = min(triton.next_power_of_2(hidden_dim), 1024) + grid = (triton.cdiv(token_num, block_m), triton.cdiv(hidden_dim, block_dim)) + moe_weighted_sum_reduce_kernel[grid]( + input, + input.stride(0), input.stride(1), input.stride(2), + weights, weights.stride(0), weights.stride(1), + output, output.stride(0), output.stride(1), + token_num, topk_num, hidden_dim, + BLOCK_M=block_m, BLOCK_DIM=block_dim, NUM_STAGE=4, + ) + + def mxfp4_fused_moe_kernel_t_triton( A: torch.Tensor, B_blocks_t: torch.Tensor, # transposed layout [E, K//2, N] (uint8, N innermost) diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index c2358d84f..a3f9f5102 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -947,3 +947,136 @@ def paged_attention( num_stages=2, ) return o + + +@triton.jit +def _q8_paged_attention_kernel( + q_ptr, k_ptr, v_ptr, ks_ptr, vs_ptr, o_ptr, + indptr_ptr, indices_ptr, q_to_req_ptr, q_pos_ptr, + sm_scale, + stride_qt, stride_qh, + stride_kslot, stride_kh, + stride_vslot, stride_vh, + stride_ksslot, stride_ksh, + stride_vsslot, stride_vsh, + stride_ot, stride_oh, + GROUP: tl.constexpr, D: tl.constexpr, BLOCK_D: tl.constexpr, + BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr, +): + """Paged causal attention that dequantizes Q8_0 rows in registers.""" + q_tok = tl.program_id(0) + q_head = tl.program_id(1) + kv_head = q_head // GROUP + req = tl.load(q_to_req_ptr + q_tok) + kv_start = tl.load(indptr_ptr + req) + kv_end = tl.load(indptr_ptr + req + 1) + kv_len = kv_end - kv_start + q_pos = tl.load(q_pos_ptr + q_tok) + offs_d = tl.arange(0, BLOCK_D) + mask_d = offs_d < D + block_d = offs_d // 32 + q = tl.load(q_ptr + q_tok * stride_qt + q_head * stride_qh + offs_d, mask=mask_d, other=0.0).to(tl.float32) + m_i = -float("inf") + l_i = 0.0 + acc = tl.zeros((BLOCK_D,), dtype=tl.float32) + for start in range(0, kv_len, BLOCK_N): + offs_n = start + tl.arange(0, BLOCK_N) + mask_n = (offs_n < kv_len) & (offs_n <= q_pos) + if SLIDING_WINDOW > 0: + mask_n = mask_n & ((offs_n + SLIDING_WINDOW) > q_pos) + if not tl.max(mask_n.to(tl.int32), axis=0) == 0: + slots = tl.load(indices_ptr + kv_start + offs_n, mask=offs_n < kv_len, other=0) + kq = tl.load( + k_ptr + slots[:, None] * stride_kslot + kv_head * stride_kh + offs_d[None, :], + mask=(offs_n[:, None] < kv_len) & mask_d[None, :], other=0, + ).to(tl.float32) + kscale = tl.load( + ks_ptr + slots[:, None] * stride_ksslot + kv_head * stride_ksh + block_d[None, :], + mask=(offs_n[:, None] < kv_len) & mask_d[None, :], other=0.0, + ) + scores = tl.sum(q[None, :] * (kq * kscale), axis=1) * sm_scale + scores = tl.where(mask_n, scores, -float("inf")) + row_max = tl.max(scores, axis=0) + m_new = tl.maximum(row_max, m_i) + alpha = tl.exp(m_i - m_new) + p = tl.exp(scores - m_new) + vq = tl.load( + v_ptr + slots[:, None] * stride_vslot + kv_head * stride_vh + offs_d[None, :], + mask=(offs_n[:, None] < kv_len) & mask_d[None, :], other=0, + ).to(tl.float32) + vscale = tl.load( + vs_ptr + slots[:, None] * stride_vsslot + kv_head * stride_vsh + block_d[None, :], + mask=(offs_n[:, None] < kv_len) & mask_d[None, :], other=0.0, + ) + acc = acc * alpha + tl.sum(p[:, None] * (vq * vscale), axis=0) + l_i = l_i * alpha + tl.sum(p, axis=0) + m_i = m_new + out = tl.where(l_i == 0.0, 0.0, acc / l_i) + tl.store(o_ptr + q_tok * stride_ot + q_head * stride_oh + offs_d, out.to(o_ptr.dtype.element_ty), mask=mask_d) + + +def q8_paged_attention( + q: torch.Tensor, + k_payload: torch.Tensor, + v_payload: torch.Tensor, + k_scales: torch.Tensor, + v_scales: torch.Tensor, + indptr: torch.Tensor, + indices: torch.Tensor, + q_to_req: torch.Tensor, + q_positions: torch.Tensor, + sm_scale: float, + sliding_window: int | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Q8_0 paged attention; storage stays packed and dequantization is per tile.""" + if q.ndim != 3 or k_payload.ndim != 3 or v_payload.ndim != 3: + raise ValueError("q8 attention expects q and packed K/V shaped [tokens, heads, dim]") + if k_payload.dtype != torch.int8 or v_payload.dtype != torch.int8: + raise TypeError("q8 attention payloads must be int8") + if k_scales.dtype != torch.float16 or v_scales.dtype != torch.float16: + raise TypeError("q8 attention scales must be float16") + batch_tokens, q_heads, head_dim = q.shape + kv_heads = k_payload.shape[1] + if q_heads % kv_heads or k_payload.shape[-1] != head_dim: + raise ValueError("q8 attention head geometry mismatch") + result = out if out is not None else torch.empty_like(q) + if not q.is_cuda: + result.zero_() + for token in range(batch_tokens): + req = int(q_to_req[token]) + begin, end = int(indptr[req]), int(indptr[req + 1]) + end = min(end, begin + int(q_positions[token]) + 1) + begin_window = max(begin, end - (sliding_window or end - begin)) + slots = indices[begin_window:end].to(torch.long) + if not slots.numel(): + continue + k = k_payload.index_select(0, slots).float() + v = v_payload.index_select(0, slots).float() + k = k * k_scales.index_select(0, slots).float().repeat_interleave(32, dim=-1) + v = v * v_scales.index_select(0, slots).float().repeat_interleave(32, dim=-1) + scores = torch.stack( + [q[token, head].float().matmul(k[:, head % kv_heads].transpose(0, 1)) + for head in range(q_heads)] + ) * sm_scale + probs = torch.softmax(scores, dim=-1) + for head in range(q_heads): + result[token, head] = probs[head].matmul(v[:, head % kv_heads]) + return result + block_d = triton.next_power_of_2(head_dim) + _q8_paged_attention_kernel[(batch_tokens, q_heads)]( + q, k_payload, v_payload, k_scales, v_scales, result, + indptr, indices, q_to_req, q_positions, sm_scale, + q.stride(0), q.stride(1), + k_payload.stride(0), k_payload.stride(1), + v_payload.stride(0), v_payload.stride(1), + k_scales.stride(0), k_scales.stride(1), + v_scales.stride(0), v_scales.stride(1), + result.stride(0), result.stride(1), + GROUP=q_heads // kv_heads, D=head_dim, BLOCK_D=block_d, BLOCK_N=32, + SLIDING_WINDOW=sliding_window or 0, num_warps=4, num_stages=2, + ) + return result + + +decode_q8_paged_attention = q8_paged_attention diff --git a/python/freetoken/kernel/triton/fused_moe.py b/python/freetoken/kernel/triton/fused_moe.py index 812e07e9e..bd62717c7 100644 --- a/python/freetoken/kernel/triton/fused_moe.py +++ b/python/freetoken/kernel/triton/fused_moe.py @@ -153,6 +153,52 @@ def moe_sum_reduce_kernel( ) +@triton.jit +def moe_weighted_sum_reduce_kernel( + input_ptr, + input_stride_0, + input_stride_1, + input_stride_2, + weights_ptr, + weights_stride_0, + weights_stride_1, + output_ptr, + output_stride_0, + output_stride_1, + token_num: int, + topk_num: int, + hidden_dim: int, + BLOCK_M: tl.constexpr, + BLOCK_DIM: tl.constexpr, + NUM_STAGE: tl.constexpr, +): + """Apply route weights and reduce in one fixed-order kernel.""" + token_block_id = tl.program_id(0) + dim_block_id = tl.program_id(1) + token_start = token_block_id * BLOCK_M + token_end = min((token_block_id + 1) * BLOCK_M, token_num) + dim_start = dim_block_id * BLOCK_DIM + dim_end = min((dim_block_id + 1) * BLOCK_DIM, hidden_dim) + offs_dim = dim_start + tl.arange(0, BLOCK_DIM) + for token_index in range(token_start, token_end): + accumulator = tl.zeros((BLOCK_DIM,), dtype=tl.float32) + input_t_ptr = input_ptr + token_index * input_stride_0 + offs_dim + for route in tl.range(0, topk_num, num_stages=NUM_STAGE): + value = tl.load( + input_t_ptr + route * input_stride_1, + mask=offs_dim < dim_end, + other=0.0, + ) + weight = tl.load(weights_ptr + token_index * weights_stride_0 + route * weights_stride_1) + accumulator += value.to(tl.float32) * weight.to(tl.float32) + store_t_ptr = output_ptr + token_index * output_stride_0 + offs_dim + tl.store( + store_t_ptr, + accumulator.to(output_ptr.dtype.element_ty), + mask=offs_dim < dim_end, + ) + + @triton.jit def fused_moe_kernel( # Pointers to matrices diff --git a/python/freetoken/kernel/triton/q8_kv.py b/python/freetoken/kernel/triton/q8_kv.py new file mode 100644 index 000000000..64cfe4f71 --- /dev/null +++ b/python/freetoken/kernel/triton/q8_kv.py @@ -0,0 +1,151 @@ +"""Graph-safe Q8_0 KV row store and its scalar reference. + +Rows are quantized independently: one K/V head at one token slot. No row crosses a +head, page, layer, or K/V slab. The scale is FP16 and the payload is signed INT8, +matching llama.cpp b10434's ``quantize_row_q8_0_ref`` contract. +""" + +from __future__ import annotations + +import torch + + +Q8_BLOCK = 32 + + +def _round_half_away_from_zero(value: torch.Tensor) -> torch.Tensor: + """C ``roundf`` semantics, unlike torch.round's ties-to-even behavior.""" + return torch.where(value >= 0, torch.floor(value + 0.5), torch.ceil(value - 0.5)) + + +def quantize_row_q8_0_ref(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(payload[int8], scales[float16])`` for rows in ``[..., D]``.""" + if x.ndim < 1 or x.shape[-1] % Q8_BLOCK: + raise ValueError(f"q8_0 requires last dimension divisible by 32, got {tuple(x.shape)}") + source = x.to(torch.float32).reshape(-1, x.shape[-1]) + blocks = source.reshape(source.shape[0], -1, Q8_BLOCK) + amax = blocks.abs().amax(dim=-1) + scales = amax / 127.0 + safe_scales = torch.where(scales == 0, torch.ones_like(scales), scales) + quant = _round_half_away_from_zero(blocks / safe_scales.unsqueeze(-1)) + quant = quant.clamp(-127, 127).to(torch.int8) + quant[amax == 0] = 0 + return quant.reshape_as(source), scales.to(torch.float16) + + +def validate_unique_destinations(indices: torch.Tensor) -> None: + """Reject racy physical destinations before graph capture.""" + flat = indices.reshape(-1) + if flat.ndim != 1 or flat.numel() == 0: + raise ValueError("q8_0 store indices must be a non-empty 1-D tensor") + if not flat.dtype in (torch.int32, torch.int64): + raise TypeError(f"q8_0 store indices must be int32/int64, got {flat.dtype}") + if torch.unique(flat).numel() != flat.numel(): + raise ValueError("q8_0 store has duplicate physical (page,offset) destinations") + + +def _store_reference( + payload: torch.Tensor, + scales: torch.Tensor, + indices: torch.Tensor, + values: torch.Tensor, +) -> None: + quant, row_scales = quantize_row_q8_0_ref(values) + payload.index_copy_(0, indices.to(torch.long), quant.reshape(-1, values.shape[-2], values.shape[-1])) + scales.index_copy_(0, indices.to(torch.long), row_scales.reshape(-1, values.shape[-2], values.shape[-1] // Q8_BLOCK)) + + +def _store_triton( + payload: torch.Tensor, + scales: torch.Tensor, + indices: torch.Tensor, + values: torch.Tensor, +) -> None: + import triton + import triton.language as tl + + @triton.jit + def kernel( + x_ptr, out_ptr, scale_ptr, idx_ptr, + sx0, sx1, sx2, so0, so1, so2, ss0, ss1, ss2, + BLOCKS: tl.constexpr, BLOCK_SIZE: tl.constexpr, + ): + row = tl.program_id(0) + block = tl.program_id(1) + heads = tl.num_programs(0) + token = row // heads + head = row - token * heads + slot = tl.load(idx_ptr + token) + offs = block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + x = tl.load(x_ptr + token * sx0 + head * sx1 + offs * sx2).to(tl.float32) + amax = tl.max(tl.abs(x), axis=0) + d = amax / 127.0 + safe_d = tl.where(d == 0, 1.0, d) + scaled = x / safe_d + q = tl.where(scaled >= 0, tl.floor(scaled + 0.5), tl.ceil(scaled - 0.5)) + q = tl.minimum(tl.maximum(q, -127.0), 127.0) + q = tl.where(amax == 0, 0.0, q) + tl.store(out_ptr + slot * so0 + head * so1 + offs * so2, q.to(tl.int8)) + tl.store(scale_ptr + slot * ss0 + head * ss1 + block * ss2, d.to(tl.float16)) + + n, heads, dim = values.shape + kernel[(n * heads, triton.cdiv(dim, Q8_BLOCK))]( + values, payload, scales, indices, + values.stride(0), values.stride(1), values.stride(2), + payload.stride(0), payload.stride(1), payload.stride(2), + scales.stride(0), scales.stride(1), scales.stride(2), + BLOCKS=dim // Q8_BLOCK, + BLOCK_SIZE=Q8_BLOCK, + num_warps=1, + ) + + +def store_q8_cache( + *, + k_payload: torch.Tensor, + v_payload: torch.Tensor, + k_scales: torch.Tensor, + v_scales: torch.Tensor, + indices: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, +) -> None: + """Quantize K and V directly into stable payload/scale page buffers.""" + if k.shape != v.shape or k.ndim != 3: + raise ValueError(f"q8_0 K/V must be matching [tokens, heads, dim], got {k.shape}/{v.shape}") + if not k.is_floating_point() or not v.is_floating_point(): + raise TypeError("q8_0 store inputs must be floating-point K/V tensors") + if k.shape[-1] % Q8_BLOCK: + raise ValueError(f"q8_0 head_dim must be divisible by 32, got {k.shape[-1]}") + if indices.numel() != k.shape[0]: + raise ValueError(f"q8_0 indices length {indices.numel()} != token count {k.shape[0]}") + # ``torch.unique`` is an eager validation barrier and cannot run inside HIP graph + # capture. Batch metadata validates destinations before capture; eager calls retain + # the loud duplicate check here. + capturing = False + if indices.is_cuda: + try: + capturing = bool(torch.cuda.is_current_stream_capturing()) + except (AttributeError, RuntimeError): + capturing = False + if not capturing: + validate_unique_destinations(indices) + for payload, scales in ((k_payload, k_scales), (v_payload, v_scales)): + if payload.dtype != torch.int8 or scales.dtype != torch.float16: + raise TypeError("q8_0 cache buffers must be int8 payload and float16 scales") + if payload.ndim != 3 or scales.shape != (*payload.shape[:2], payload.shape[-1] // Q8_BLOCK): + raise ValueError("q8_0 cache payload/scale geometry mismatch") + if not k.is_cuda: + _store_reference(k_payload, k_scales, indices, k) + _store_reference(v_payload, v_scales, indices, v) + return + _store_triton(k_payload, k_scales, indices, k) + _store_triton(v_payload, v_scales, indices, v) + + +__all__ = [ + "Q8_BLOCK", + "quantize_row_q8_0_ref", + "store_q8_cache", + "validate_unique_destinations", +] diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index c6c0f1bb9..291f94be4 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -12,6 +12,8 @@ BaseCacheHandle, BaseKVCachePool, BasePrefixCache, + KVStorageDescriptor, + QuantizedKVView, MatchResult, SizeInfo, ) @@ -111,6 +113,7 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt device=device, dtype=dtype, num_req_slots=config.max_running_req + 1, # + 1 for the dummy request row + storage_type=getattr(config, "kv_storage_type", None), ) @@ -122,6 +125,7 @@ def create_kvcache_pool( device: torch.device, num_swa_tokens: int | None = None, num_req_slots: int | None = None, + storage_type=None, ) -> BaseKVCachePool: if model_config.has_swa_attention: from .hybrid_swa_pool import HybridSWAKVCache @@ -238,6 +242,7 @@ def create_kvcache_pool( device=device, dtype=dtype, layer_ids=layer_ids, + storage_type=storage_type, ) diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py index 95669e8c8..9629f5ca6 100644 --- a/python/freetoken/kvcache/base.py +++ b/python/freetoken/kvcache/base.py @@ -10,6 +10,116 @@ logger = init_logger(__name__) +def _kv_storage_type(): + """Resolve enum lazily; importing engine.config at module load cycles through kvcache.""" + from freetoken.engine.config import KVStorageType + + return KVStorageType + + +@dataclass(frozen=True) +class KVStorageDescriptor: + """Immutable physical layout contract shared by allocation and accounting.""" + + storage_type: KVStorageType + block_size: int = 32 + payload_dtype: torch.dtype = torch.bfloat16 + scale_dtype: torch.dtype | None = None + contract_id: str = "native-16" + + def __post_init__(self) -> None: + kv_type = _kv_storage_type() + storage_type = kv_type.parse(self.storage_type) + object.__setattr__(self, "storage_type", storage_type) + if storage_type == kv_type.Q8_0: + if self.block_size != 32: + raise ValueError("q8_0 KV storage requires block_size=32") + object.__setattr__(self, "payload_dtype", torch.int8) + object.__setattr__(self, "scale_dtype", torch.float16) + object.__setattr__(self, "contract_id", "llama.cpp-b10434-q8_0-row-v1") + elif self.scale_dtype is not None: + raise ValueError(f"scale dtype is only valid for q8_0, got {storage_type.value}") + + @property + def is_quantized(self) -> bool: + return self.storage_type == _kv_storage_type().Q8_0 + + @property + def bytes_per_block(self) -> int: + if self.is_quantized: + return self.block_size + 2 + return self.block_size * self.payload_dtype.itemsize + + def validate_head_dim(self, head_dim: int) -> None: + if self.is_quantized and head_dim % self.block_size: + raise ValueError( + f"q8_0 KV storage requires head_dim divisible by {self.block_size}; " + f"got {head_dim}" + ) + + def row_bytes(self, head_dim: int) -> int: + self.validate_head_dim(head_dim) + if self.is_quantized: + return head_dim + 2 * (head_dim // self.block_size) + return head_dim * self.payload_dtype.itemsize + + def bytes_per_token( + self, *, num_layers: int, num_kv_heads: int, head_dim: int, slabs: int = 2 + ) -> int: + return slabs * num_layers * num_kv_heads * self.row_bytes(head_dim) + + +@dataclass(frozen=True) +class QuantizedKVView: + payload: torch.Tensor + scales: torch.Tensor + descriptor: KVStorageDescriptor + + +def kv_storage_descriptor(config, *, head_dim: int | None = None) -> KVStorageDescriptor: + """Resolve one descriptor without allocating a pool.""" + raw_type = getattr(config, "kv_storage_type", None) + if raw_type is None: + dtype = getattr(config, "dtype", torch.bfloat16) + kv_type = _kv_storage_type() + storage_type = kv_type.FP16 if dtype == torch.float16 else kv_type.BF16 + return KVStorageDescriptor(storage_type, payload_dtype=dtype) + kv_type = _kv_storage_type() + storage_type = kv_type.parse(raw_type) + if storage_type == kv_type.BF16: + return KVStorageDescriptor(storage_type, payload_dtype=torch.bfloat16) + if storage_type == kv_type.FP16: + return KVStorageDescriptor(storage_type, payload_dtype=torch.float16) + descriptor = KVStorageDescriptor(storage_type) + if head_dim is not None: + descriptor.validate_head_dim(int(head_dim)) + return descriptor + + +def validate_kv_storage_config(config) -> KVStorageDescriptor: + """Reject unsupported storage before model/cache allocation.""" + descriptor = kv_storage_descriptor(config) + if descriptor.storage_type != _kv_storage_type().Q8_0: + return descriptor + specs = tuple(getattr(config.model_config, "kv_cache_group_specs")()) + from freetoken.attention import AttnType + + unsupported = [ + spec.name for spec in specs + if getattr(spec, "attn_type", AttnType.FULL) is not AttnType.FULL + or spec.mla or spec.index_head_dim or spec.num_index_layers + ] + if unsupported: + raise ValueError( + "q8_0 KV storage currently supports plain full-attention MHA groups only; " + f"unsupported groups: {', '.join(unsupported)}" + ) + for spec in specs: + if spec.num_layers: + descriptor.validate_head_dim(int(spec.head_dim)) + return descriptor + + class CacheRebuildRejected(Exception): """A runtime cache rebuild was rejected BEFORE any destructive free (e.g. the requested geometry does not fit). The old caches are intact and serving continues -- @@ -25,13 +135,24 @@ def spec_kv_bytes_per_token(spec, config) -> int: ``index_ratio`` > 1 (QSA) stores one index key per token group, not per token; that slab's ring and scratch rows are fixed-size and priced in QSAKVCache.kv_cost instead.""" - per_token = ( - (1 if spec.mla else 2) # MLA latent groups store one slab (V aliases K) - * spec.head_dim - * div_even(spec.num_kv_heads, config.tp_info.size, allow_replicate=True) - * config.dtype.itemsize - * spec.num_layers + descriptor = kv_storage_descriptor(config, head_dim=spec.head_dim) + if descriptor.is_quantized and (spec.mla or spec.index_head_dim or spec.num_index_layers): + raise ValueError(f"q8_0 KV storage does not support group {spec.name!r}") + per_token = descriptor.bytes_per_token( + num_layers=spec.num_layers, + num_kv_heads=div_even(spec.num_kv_heads, config.tp_info.size, allow_replicate=True), + head_dim=spec.head_dim, + slabs=1 if spec.mla else 2, ) + if not descriptor.is_quantized: + # Preserve DSA/QSA index accounting for native 16-bit pools. + per_token = ( + (1 if spec.mla else 2) + * spec.head_dim + * div_even(spec.num_kv_heads, config.tp_info.size, allow_replicate=True) + * config.dtype.itemsize + * spec.num_layers + ) return per_token + spec.index_head_dim * spec.num_index_layers * 2 // spec.index_ratio diff --git a/python/freetoken/kvcache/cache_status.py b/python/freetoken/kvcache/cache_status.py index 10169d651..03ac40550 100644 --- a/python/freetoken/kvcache/cache_status.py +++ b/python/freetoken/kvcache/cache_status.py @@ -194,6 +194,23 @@ def compute_cache_status_meta(engine: "Engine") -> Dict[str, Any]: meta["free_vram_bytes"] = _pool_budget_free_vram_bytes(engine) meta["floors"] = compute_cache_floors(engine) meta["pools"] = compute_cache_pools(engine) + try: + runtime = engine.graph_runner.runtime_telemetry() + runtime["effective_moe_backend"] = getattr(engine.config, "moe_backend", None) + kv = dict(getattr(engine, "kv_storage_metadata", None) or {}) + if kv: + runtime["kv_storage"] = kv + runtime["kv_type"] = kv.get("storage_type") + runtime["memory_phases"] = [ + phase.as_dict() for phase in getattr(engine, "memory_phases", ()) + ] + if runtime.get("resident_gguf") and runtime.get("effective_moe_backend") == "fused": + runtime["execution_class"] = "resident_fused" + elif runtime.get("effective_moe_backend") in {"offload", "hybrid"}: + runtime["execution_class"] = "offload" + meta["execution"] = runtime + except Exception: # noqa: BLE001 -- telemetry must never block readiness + meta["execution"] = {} # Current window/full reuse ratio (the tunable knob), for DSV4 and radix-SWA; 0.0 otherwise. cfg = engine.config has_swa_ratio = cfg is not None and _supports_swa_ratio(cfg) diff --git a/python/freetoken/kvcache/mha_pool.py b/python/freetoken/kvcache/mha_pool.py index 8ed280b96..cde006ab1 100644 --- a/python/freetoken/kvcache/mha_pool.py +++ b/python/freetoken/kvcache/mha_pool.py @@ -6,7 +6,7 @@ from freetoken.distributed import get_tp_info from freetoken.utils import div_even -from .base import BaseKVCachePool +from .base import BaseKVCachePool, KVStorageDescriptor, kv_storage_descriptor class MHAKVCache(BaseKVCachePool): @@ -32,6 +32,7 @@ def __init__( dtype: torch.dtype, device: torch.device, layer_ids: Sequence[int] | None = None, + storage_type=None, ) -> None: tp_info = get_tp_info() local_kv_heads = div_even(num_kv_heads, tp_info.size, allow_replicate=True) @@ -47,13 +48,32 @@ def __init__( raise ValueError(f"KV layer id {global_id} outside [0, {num_layers})") layer_map[global_id] = dense self._layer_map = layer_map - self._kv_buffer = torch.empty( - (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim), - device=device, - dtype=dtype, - ) - self._k_buffer = self._kv_buffer[0] - self._v_buffer = self._kv_buffer[1] + if storage_type is None: + descriptor = KVStorageDescriptor( + "bf16" if dtype == torch.bfloat16 else "fp16", payload_dtype=dtype + ) + self._storage_dtype = dtype + else: + descriptor = kv_storage_descriptor(type("KVConfig", (), {"kv_storage_type": storage_type})(), head_dim=head_dim) + self._storage_dtype = descriptor.payload_dtype + self._descriptor = descriptor + self._generation = 0 + shape = (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim) + if descriptor.is_quantized: + descriptor.validate_head_dim(head_dim) + self._kv_buffer = None + self._k_buffer = torch.empty(shape, device=device, dtype=torch.int8)[0] + self._v_buffer = torch.empty(shape, device=device, dtype=torch.int8)[1] + scale_shape = (*shape[:4], local_kv_heads, head_dim // descriptor.block_size) + scales = torch.empty(scale_shape, device=device, dtype=descriptor.scale_dtype) + self._k_scales = scales[0] + self._v_scales = scales[1] + self._zero_dummy_page() + else: + self._kv_buffer = torch.empty(shape, device=device, dtype=self._storage_dtype) + self._k_buffer = self._kv_buffer[0] + self._v_buffer = self._kv_buffer[1] + self._k_scales = self._v_scales = None self._device = device self._storage_shape = (num_pages * page_size, local_kv_heads, head_dim) @@ -64,8 +84,8 @@ def rebuild(self, num_pages: int) -> None: existing buffer; only the page count changes. Views and ``_storage_shape`` are refreshed. Object identity is preserved so cached backend references stay valid. """ - _, num_storage_layers, _old_pages, page_size, local_kv_heads, head_dim = self._kv_buffer.shape - dtype = self._kv_buffer.dtype + old = self._k_buffer + num_storage_layers, _old_pages, page_size, local_kv_heads, head_dim = old.shape device = self._device self._k_buffer = None self._v_buffer = None @@ -73,14 +93,31 @@ def rebuild(self, num_pages: int) -> None: if device.type == "cuda": torch.cuda.synchronize(device) torch.cuda.empty_cache() - self._kv_buffer = torch.empty( - (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim), - device=device, - dtype=dtype, - ) - self._k_buffer = self._kv_buffer[0] - self._v_buffer = self._kv_buffer[1] + shape = (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim) + if self._descriptor.is_quantized: + self._k_buffer = torch.empty(shape, device=device, dtype=torch.int8)[0] + self._v_buffer = torch.empty(shape, device=device, dtype=torch.int8)[1] + scale_shape = (*shape[:4], local_kv_heads, head_dim // self._descriptor.block_size) + scales = torch.empty(scale_shape, device=device, dtype=self._descriptor.scale_dtype) + self._k_scales = scales[0] + self._v_scales = scales[1] + self._zero_dummy_page() + else: + self._kv_buffer = torch.empty(shape, device=device, dtype=self._storage_dtype) + self._k_buffer = self._kv_buffer[0] + self._v_buffer = self._kv_buffer[1] + self._k_scales = self._v_scales = None self._storage_shape = (num_pages * page_size, local_kv_heads, head_dim) + self._generation += 1 + + def _zero_dummy_page(self) -> None: + """Initialize reserved dummy page, including Q8 payload and scales.""" + if not self._descriptor.is_quantized: + return + self._k_buffer[:, -1].zero_() + self._v_buffer[:, -1].zero_() + self._k_scales[:, -1].zero_() + self._v_scales[:, -1].zero_() @classmethod def kv_cost(cls, config) -> tuple[int, int, int, int]: @@ -99,9 +136,12 @@ def rebuild_from_config( self.rebuild(num_pages + 1) # +1 for the dummy page (matches create_kvcache_pool) def unit_bytes(self) -> tuple[int, int]: - buf = self._kv_buffer - tokens = int(buf.shape[2]) * int(buf.shape[3]) - return int(buf.numel() * buf.element_size()) // tokens, 0 + tokens = int(self._k_buffer.shape[1]) * int(self._k_buffer.shape[2]) + total = self._k_buffer.numel() * self._k_buffer.element_size() + if self._descriptor.is_quantized: + total += self._k_scales.numel() * self._k_scales.element_size() + total *= 2 + return total // tokens, 0 def _dense(self, layer_id: int) -> int: if self._layer_map is None: @@ -124,6 +164,31 @@ def store_kv( out_loc: torch.Tensor, layer_id: int, ) -> None: + if self._descriptor.is_quantized: + from freetoken.kernel.triton.q8_kv import store_q8_cache + + # Qwen3.5 projection path keeps KV heads flattened as [tokens, + # heads * head_dim]. Q8 storage quantizes one row per head, so + # restore its explicit row geometry at cache boundary. reshape is + # view-only for normal contiguous projection output. + kv_heads, head_dim = self._storage_shape[1:] + flat_width = kv_heads * head_dim + if k.ndim == 2 and v.ndim == 2 and k.shape[1] == flat_width and v.shape[1] == flat_width: + k = k.reshape(-1, kv_heads, head_dim) + v = v.reshape(-1, kv_heads, head_dim) + + dense = self._dense(layer_id) + store_q8_cache( + k_payload=self._k_buffer[dense].view(self._storage_shape), + v_payload=self._v_buffer[dense].view(self._storage_shape), + k_scales=self._k_scales[dense].view(-1, self._storage_shape[1], self._storage_shape[2] // 32), + v_scales=self._v_scales[dense].view(-1, self._storage_shape[1], self._storage_shape[2] // 32), + indices=out_loc, + k=k, + v=v, + ) + return + from freetoken.kernel import store_cache dense = self._dense(layer_id) @@ -141,8 +206,44 @@ def device(self) -> torch.device: @property def dtype(self) -> torch.dtype: - return self._kv_buffer.dtype + return self._descriptor.payload_dtype @property def num_layers(self) -> int: return self._num_layers + + @property + def storage_descriptor(self) -> KVStorageDescriptor: + return self._descriptor + + @property + def pointer_generation(self) -> int: + return self._generation + + @property + def is_quantized(self) -> bool: + return self._descriptor.is_quantized + + def k_cache_view(self, index: int): + if not self.is_quantized: + return self.k_cache(index) + from .base import QuantizedKVView + + dense = self._dense(index) + return QuantizedKVView( + self._k_buffer[dense].view(self._storage_shape), + self._k_scales[dense].view(-1, self._storage_shape[1], self._storage_shape[2] // 32), + self._descriptor, + ) + + def v_cache_view(self, index: int): + if not self.is_quantized: + return self.v_cache(index) + from .base import QuantizedKVView + + dense = self._dense(index) + return QuantizedKVView( + self._v_buffer[dense].view(self._storage_shape), + self._v_scales[dense].view(-1, self._storage_shape[1], self._storage_shape[2] // 32), + self._descriptor, + ) diff --git a/python/freetoken/layers/gguf.py b/python/freetoken/layers/gguf.py index ac49b1a5b..61e5d70b3 100644 --- a/python/freetoken/layers/gguf.py +++ b/python/freetoken/layers/gguf.py @@ -21,7 +21,9 @@ GGML_F16, GGML_F32, GGML_NAME, + GGML_Q4_K, GGML_Q4_0, + GGML_Q5_K, GGML_Q6_K, GGML_Q8_0, row_bytes, @@ -32,30 +34,47 @@ # ggml type groups for kernel dispatch (subset we build kernels for). _UNQUANTIZED = {GGML_F32, GGML_F16, GGML_BF16} # standard + k-quants: both an MMVQ (small-batch GEMV) and MMQ (large-batch) kernel exist. -_MMVQ = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K} -_MMQ = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K} -_DEQUANT = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K} +_MMVQ = {GGML_Q4_0, GGML_Q4_K, GGML_Q5_K, GGML_Q8_0, GGML_Q6_K} +_MMQ = {GGML_Q4_0, GGML_Q4_K, GGML_Q5_K, GGML_Q8_0, GGML_Q6_K} +_DEQUANT = {GGML_Q4_0, GGML_Q4_K, GGML_Q5_K, GGML_Q8_0, GGML_Q6_K} # Below this token count, the MMVQ GEMV kernel wins (matches vLLM's heuristic). _MMVQ_SAFE = 6 -def fused_mul_mat_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int) -> torch.Tensor: +def fused_mul_mat_gguf( + x: torch.Tensor, qweight: torch.Tensor, qweight_type: int, operation: str = "dense" +) -> torch.Tensor: """y = x @ dequant(qweight).T, dispatched by batch size and quant type.""" from freetoken.kernel.gguf import ( ggml_dequantize, ggml_mul_mat_a8, ggml_mul_mat_vec_a8, + gguf_dispatch, + gguf_runtime_metadata, ) out_features = qweight.shape[0] if x.shape[0] == 0: return x.new_empty((0, out_features)) + if qweight_type in BLOCK_SHAPE: + block, type_size = BLOCK_SHAPE[qweight_type] + in_features = qweight.shape[1] // type_size * block + else: + in_features = x.shape[1] + dispatch = gguf_dispatch( + operation, + qweight_type, + out_features, + in_features, + x.shape[0], + gguf_runtime_metadata().get("arch"), + ) if qweight_type in _UNQUANTIZED: return x @ qweight.T - if x.shape[0] <= _MMVQ_SAFE and qweight_type in _MMVQ: + if dispatch["implementation"] == "ggml_mul_mat_vec_a8" and qweight_type in _MMVQ: return ggml_mul_mat_vec_a8(qweight, x, qweight_type, out_features) - if qweight_type in _MMQ: + if dispatch["implementation"] == "ggml_mul_mat_a8" and qweight_type in _MMQ: return ggml_mul_mat_a8(qweight, x, qweight_type, out_features) if qweight_type in _DEQUANT: block, type_size = BLOCK_SHAPE[qweight_type] @@ -74,19 +93,31 @@ def __init__( out_features: int, quant_type: int, has_bias: bool = False, + operation: str = "dense", + quant_role: str = "dense", ): self.in_features = in_features self.out_features = out_features self._quant_type = quant_type + self._operation = operation + self.quant_role = quant_role self.qweight = torch.empty(out_features, row_bytes(in_features, quant_type), dtype=torch.uint8) self.bias = torch.empty(out_features) if has_bias else None def forward(self, x: torch.Tensor) -> torch.Tensor: - out = fused_mul_mat_gguf(x, self.qweight, self._quant_type) + out = fused_mul_mat_gguf(x, self.qweight, self._quant_type, self._operation) if self.bias is not None: out = out + self.bias return out + @property + def quant_metadata(self) -> dict[str, object]: + return { + "role": self.quant_role, + "quant_type": self._quant_type, + "operation": self._operation, + } + class GGUFEmbedding(BaseOP): """Vocab embedding stored as a native GGUF block-quantized table. @@ -101,10 +132,12 @@ def __init__( embedding_dim: int, quant_type: int, embed_scale: float | None = None, + quant_role: str = "embedding", ): self.num_embeddings = num_embeddings self.embedding_dim = embedding_dim self._quant_type = quant_type + self.quant_role = quant_role self.qweight = torch.empty( num_embeddings, row_bytes(embedding_dim, quant_type), dtype=torch.uint8 ) @@ -124,5 +157,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y = y * self._embed_scale_t return y + @property + def quant_metadata(self) -> dict[str, object]: + return {"role": self.quant_role, "quant_type": self._quant_type} + __all__ = ["GGUFLinear", "GGUFEmbedding", "fused_mul_mat_gguf"] diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index ac01eed90..7421f3c15 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -38,6 +38,7 @@ def __init__( apply_router_weight_on_input: bool = False, allocate_experts: bool = True, weight_format: str = "bf16", + gguf_down_quant_type: int | None = None, ): super().__init__() @@ -53,6 +54,19 @@ def __init__( self.activation = activation self.apply_router_weight_on_input = apply_router_weight_on_input self.weight_format = weight_format + self.gguf_down_quant_type = gguf_down_quant_type + if weight_format == "gguf": + from freetoken.moe.fused_gguf import MoeDecodeWork + + self._gguf_workspaces = { + "moe_decode": MoeDecodeWork("moe_decode"), + "moe_prefill": MoeDecodeWork("moe_prefill"), + } + # Compatibility alias for callers that only know the old scratch name. + self._gguf_workspace = self._gguf_workspaces["moe_decode"] + else: + self._gguf_workspaces = {} + self._gguf_workspace = None intermediate_size_per_partition = div_even(intermediate_size, tp_size) if allocate_experts: self._alloc_resident_experts(intermediate_size_per_partition) @@ -77,6 +91,24 @@ def _alloc_resident_experts(self, intermediate_size_per_partition: int) -> None: self.down_proj = torch.empty(n, h, i, dtype=FP8) self.down_scale_inv = torch.empty(n, h // blk, i // blk, dtype=torch.bfloat16) return + if self.weight_format == "gguf": + from freetoken.models.gguf.dequant import ( + GGML_Q4_K, GGML_Q5_K, GGML_Q6_K, row_bytes, + ) + + if self.gguf_down_quant_type not in (GGML_Q5_K, GGML_Q6_K): + raise ValueError( + "resident GGUF expert down type must be Q5_K/Q6_K, got " + f"{self.gguf_down_quant_type!r}" + ) + n, i, h = self.num_experts, self.intermediate_size, self.hidden_size + self.gate_up_proj = torch.empty( + n, 2 * i, row_bytes(h, GGML_Q4_K), dtype=torch.uint8 + ) + self.down_proj = torch.empty( + n, h, row_bytes(i, self.gguf_down_quant_type), dtype=torch.uint8 + ) + return assert self.weight_format == "bf16", ( f"no resident expert allocation for weight_format {self.weight_format!r}" ) @@ -139,6 +171,22 @@ def _resident_gemm( hidden_states, self.gate_up_proj, self.gate_up_scale_inv, self.down_proj, self.down_scale_inv, topk_weights, topk_ids, ) + if self.weight_format == "gguf": + from freetoken.moe.fused_gguf import fused_experts_gguf_native + + return fused_experts_gguf_native( + hidden_states, + self.gate_up_proj, + self.down_proj, + topk_weights, + topk_ids, + self.activation, + down_quant_type=self.gguf_down_quant_type, + is_prefill=get_global_ctx().batch.is_prefill, + workspace=self._gguf_workspaces[ + "moe_prefill" if get_global_ctx().batch.is_prefill else "moe_decode" + ], + ) assert self.weight_format == "bf16", ( f"no resident expert kernel for weight_format {self.weight_format!r}" ) @@ -209,6 +257,7 @@ def __init__( renormalize: bool = True, activation: str = "silu", apply_router_weight_on_input: bool = False, + gguf_down_quant_type: int | None = None, ): super().__init__( num_experts=num_experts, @@ -221,7 +270,16 @@ def __init__( allocate_experts=False, ) self.layer_id = layer_id + self.gguf_down_quant_type = gguf_down_quant_type self.offload_cache: OffloadMoeCache | None = None + if gguf_down_quant_type is not None: + from freetoken.moe.fused_gguf import MoeDecodeWork + + self._gguf_workspaces = { + "moe_decode": MoeDecodeWork("moe_decode"), + "moe_prefill": MoeDecodeWork("moe_prefill"), + } + self._gguf_workspace = self._gguf_workspaces["moe_decode"] def forward( self, @@ -535,15 +593,54 @@ def _expert_gemm( return fused_experts_gguf_q4_0( hidden_states, gate_up, down, topk_weights, topk_ids, self.activation ) - if fmt == "gguf": - # Native GGUF qwen3.5-moe experts: gate_up stays Q4_K, down is stored as Q8_0 - # (re-quantized at load; a uniform format the cache can hold). Dequant-in-kernel - # grouped GEMV (MMVQ) over the streamed packed banks. + if fmt in ("gguf", "gguf_native"): + # GGUF offload keeps Q4_K gate/up. Legacy ``gguf`` converts routed + # down rows to Q8_0; ``gguf_native`` retains per-layer Q5_K/Q6_K + # rows in a Q6_K-stride cache and selects the exact kernel here. from freetoken.moe.fused_gguf import fused_experts_gguf + from freetoken.kernel.gguf import gguf_dispatch, gguf_runtime_metadata gate_up, down = views + down_type = 8 if fmt == "gguf" else self.gguf_down_quant_type + if down_type not in (8, 13, 14): + raise ValueError(f"invalid GGUF offload down type {down_type!r}") + phase = "moe_prefill" if is_prefill else "moe_decode" + arch = gguf_runtime_metadata().get("arch") + dispatch_metadata = { + "gate_up": gguf_dispatch( + phase, + 12, + gate_up.shape[1], + hidden_states.shape[1], + hidden_states.shape[0], + arch, + ), + "down": gguf_dispatch( + phase, + down_type, + down.shape[1], + gate_up.shape[1] // 2, + hidden_states.shape[0] * topk_ids.shape[1], + arch, + ), + } + work = self._gguf_workspaces.get( + "moe_prefill" if is_prefill else "moe_decode" + ) return fused_experts_gguf( - hidden_states, gate_up, down, topk_weights, topk_ids, self.activation + hidden_states, + gate_up, + down, + topk_weights, + topk_ids, + self.activation, + is_prefill=is_prefill, + dispatch_metadata=dispatch_metadata, + down_quant_type=down_type, + down_stride_bytes=(int(down.stride(0)) if fmt == "gguf_native" else None), + down_row_stride_bytes=(int(down.stride(1)) if fmt == "gguf_native" else None), + work=work, + id_space="raw" if is_prefill else "slot", ) if fmt == "mxfp4_triton": # gpt-oss MXFP4 experts (biased, clamped swiglu): transposed split-K GEMV @@ -612,6 +709,7 @@ def make_moe_layer( intermediate_size: int | None = None, resident_cls: type[MoELayer] | None = None, offload_cls: "type[OffloadMoELayer] | None" = None, + gguf_down_quant_type: int | None = None, extra_attrs: dict | None = None, ) -> MoELayer: """Build the experts layer for ``config.moe_backend`` -- the one construction @@ -642,8 +740,11 @@ def make_moe_layer( if offload: assert layer_id is not None, "offload MoE backends need the layer_id" kwargs["layer_id"] = layer_id + if gguf_down_quant_type is not None: + kwargs["gguf_down_quant_type"] = gguf_down_quant_type else: kwargs["weight_format"] = weight_format + kwargs["gguf_down_quant_type"] = gguf_down_quant_type layer = layer_cls(**kwargs) for name, value in (extra_attrs or {}).items(): setattr(layer, name, value) diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index 229cce812..6bbc870be 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -297,6 +297,9 @@ class ModelConfig: has_attn_bias: bool = False has_router_bias: bool = False moe_weight_format: str | None = None + # GGUF routed-expert down projection type per decoder layer. Populated by GGUF model + # adapters before resident construction; empty for every non-GGUF model. + gguf_down_quant_types: Tuple[int, ...] = () swiglu_limit: float | None = None hidden_act_alpha: float = 1.702 # Full DeepseekV4Args payload for the DSV4-specific machinery (MLA sparse attention, diff --git a/python/freetoken/models/gguf/dequant.py b/python/freetoken/models/gguf/dequant.py index 05648f9e4..22ddb6b8c 100644 --- a/python/freetoken/models/gguf/dequant.py +++ b/python/freetoken/models/gguf/dequant.py @@ -1,5 +1,4 @@ -"""GGML block-quant dequantization in pure torch (the formats this repo's GGUF -checkpoints use: Q4_0, Q6_K, plus trivial F32/F16/BF16). +"""GGML block-quant dequantization in pure torch. This is the *reference / CPU* path, NOT the engine's hot path: GGUF weights stay packed and are dequantized inside the borrowed ggml CUDA kernels (see @@ -36,6 +35,7 @@ GGML_Q4_0: (32, 18), GGML_Q8_0: (32, 34), GGML_Q4_K: (256, 144), + GGML_Q5_K: (256, 176), GGML_Q6_K: (256, 210), } @@ -84,6 +84,42 @@ def dequant_q4_0(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: return ((q - 8.0) * d).reshape(-1).to(out_dtype) +def _scale_min_k4(scales: torch.Tensor, index: int) -> tuple[torch.Tensor, torch.Tensor]: + if index < 4: + return scales[:, index] & 63, scales[:, index + 4] & 63 + return ( + (scales[:, index + 4] & 0xF) | ((scales[:, index - 4] >> 6) << 4), + (scales[:, index + 4] >> 4) | ((scales[:, index] >> 6) << 4), + ) + + +def dequant_q4_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + """Q4_K: 256-element super-block with eight 32-element scale/min groups.""" + raw = raw.reshape(-1, 144) + n = raw.shape[0] + dm = raw[:, 0:4].contiguous().view(torch.float16).to(torch.float32) + dall, dmin = dm[:, 0], dm[:, 1] + scales = raw[:, 4:16] + qs = raw[:, 16:144] + y = torch.empty((n, 256), dtype=torch.float32, device=raw.device) + for il in range(4): + s0, m0 = _scale_min_k4(scales, 2 * il) + s1, m1 = _scale_min_k4(scales, 2 * il + 1) + q = qs[:, 32 * il:32 * il + 32].to(torch.float32) + lo = 64 * il + y[:, lo:lo + 32] = q.remainder(16) * (dall * s0).unsqueeze(1) - (dmin * m0).unsqueeze(1) + y[:, lo + 32:lo + 64] = torch.div(q, 16, rounding_mode="floor") * (dall * s1).unsqueeze(1) - (dmin * m1).unsqueeze(1) + return y.reshape(-1).to(out_dtype) + + +def dequant_q8_0(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + """Q8_0: per-32-element block = fp16 scale followed by 32 signed quants.""" + raw = raw.reshape(-1, 34) + d = _f16_scales(raw, 0, 2) + q = raw[:, 2:34].contiguous().view(torch.int8).to(torch.float32) + return (q * d).reshape(-1).to(out_dtype) + + def dequant_q6_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: """Q6_K: 256-elem super-block = 128B low nibbles + 64B high 2-bits + 16 int8 sub-scales + fp16 ``d``. Direct vectorization of ggml's two-half loop.""" @@ -182,8 +218,10 @@ def _sm(j): _DEQUANT = { GGML_Q4_0: dequant_q4_0, + GGML_Q4_K: dequant_q4_k, GGML_Q5_K: dequant_q5_k, GGML_Q6_K: dequant_q6_k, + GGML_Q8_0: dequant_q8_0, } @@ -210,11 +248,14 @@ def dequantize(raw: torch.Tensor, ggml_type: int, out_dtype: torch.dtype) -> tor "GGML_Q4_0", "GGML_Q8_0", "GGML_Q4_K", + "GGML_Q5_K", "GGML_Q6_K", "GGML_NAME", "BLOCK_SHAPE", "row_bytes", "dequant_q4_0", + "dequant_q4_k", + "dequant_q8_0", "dequant_q5_k", "dequant_q6_k", "quantize_q8_0", diff --git a/python/freetoken/models/qwen3_5_moe/__init__.py b/python/freetoken/models/qwen3_5_moe/__init__.py index cae7dfd34..c32f4ffbf 100644 --- a/python/freetoken/models/qwen3_5_moe/__init__.py +++ b/python/freetoken/models/qwen3_5_moe/__init__.py @@ -4,6 +4,7 @@ is_gguf_model, iter_gguf_weights, load_gguf_expert_sources, + load_gguf_expert_sources_native, parse_gguf_config, ) from .model import Qwen3_5MoEForCausalLM @@ -28,4 +29,5 @@ "convert_qwen35moe_to_gguf", "is_gguf_model", "load_gguf_expert_sources", + "load_gguf_expert_sources_native", ] diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index e6ffb04cf..3b6eb4a3f 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -13,7 +13,10 @@ down, the token embedding and the lm_head) stay in their native packed block layout (Q8_0 projections, Q6_K head) and are yielded as ``.qweight`` (uint8); tiny F32 tensors (norms, router, GDN b/a) dequantize to bf16; GDN conv/A_log/dt_bias stay fp32. Routed experts -(Q4_K gate/up, Q5_K/Q6_K down) go to the offload cache (``load_gguf_expert_sources``). +(Q4_K gate/up, Q5_K/Q6_K down) stay packed. Resident mode loads them into model +buffers; plain GPU offload uses ``load_gguf_expert_sources_native`` and retains each +down type with a padded Q6_K row stride; CPU/hybrid/converter paths use +``load_gguf_expert_sources`` and its legacy uniform Q8_0 down banks. """ from __future__ import annotations @@ -31,6 +34,8 @@ from freetoken.models.gguf.dequant import ( GGML_F32, GGML_Q4_K, + GGML_Q5_K, + GGML_Q6_K, GGML_Q8_0, dequantize, quantize_q8_0, @@ -54,6 +59,8 @@ def _require_tp1(what: str) -> None: def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: m = shim.metadata + down_quant_types = _gguf_down_quant_types(shim.model_path) + def g(key: str): val = m.get(f"qwen35moe.{key}") if val is None: @@ -127,6 +134,7 @@ def g(key: str): dense_quant="gguf", lm_head_quant="gguf", moe_weight_format="gguf", + gguf_down_quant_types=down_quant_types, ) @@ -134,6 +142,36 @@ def is_gguf_model(config: ModelConfig) -> bool: return getattr(config, "moe_weight_format", None) == "gguf" +def _gguf_down_quant_types(model_path: str) -> tuple[int, ...]: + """Read routed down types from tensor headers without touching tensor payloads.""" + from freetoken.models.gguf.reader import _reader + + try: + reader = _reader(model_path) + except (OSError, ValueError, ImportError): + return () + types: dict[int, int] = {} + for tensor in reader.tensors: + if not tensor.name.startswith("blk.") or not tensor.name.endswith("ffn_down_exps.weight"): + continue + layer = int(tensor.name.split(".")[1]) + quant_type = int(tensor.tensor_type) + if quant_type not in (GGML_Q5_K, GGML_Q6_K): + raise ValueError( + f"{tensor.name}: resident GGUF supports Q5_K/Q6_K down experts, " + f"got ggml type {quant_type}" + ) + prior = types.setdefault(layer, quant_type) + if prior != quant_type: + raise ValueError(f"{tensor.name}: expert down quant type changed within layer") + if not types: + return () + last = max(types) + if set(types) != set(range(last + 1)): + raise ValueError(f"GGUF expert down types missing layers: expected 0..{last}, got {sorted(types)}") + return tuple(types[i] for i in range(last + 1)) + + # -------------------------------------------------------------------------------------- # Model layer swap: dense bf16 Linear -> native GGUF-quant ops. # -------------------------------------------------------------------------------------- @@ -146,7 +184,8 @@ def convert_qwen35moe_to_gguf(model, config: ModelConfig) -> None: embedding (Q8_0) and the (untied) lm_head (Q6_K), full-attention qkv/o (Q8_0), GDN in_proj_qkvz + out_proj (Q8_0; in_proj_ba stays dense bf16), and the shared-expert gate_up/down (Q8_0). Left dense bf16/fp32 (F32 in the GGUF): the norms, the two - routers, and the GDN conv1d/A_log/dt_bias. Routed experts stay on the offload cache. + routers, and the GDN conv1d/A_log/dt_bias. Routed experts are native-resident when + selected, or remain in host banks for offload. """ from freetoken.layers.gguf import GGUFEmbedding, GGUFLinear from freetoken.models.gguf.dequant import GGML_Q6_K, GGML_Q8_0 @@ -156,9 +195,11 @@ def convert_qwen35moe_to_gguf(model, config: ModelConfig) -> None: num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, quant_type=GGML_Q8_0, + quant_role="token_embedding", ) model.lm_head = GGUFLinear( - config.hidden_size, config.vocab_size, GGML_Q6_K, has_bias=False + config.hidden_size, config.vocab_size, GGML_Q6_K, has_bias=False, + operation="lm_head", quant_role="lm_head", ) shared_I = config.shared_expert_intermediate_size @@ -166,10 +207,11 @@ def convert_qwen35moe_to_gguf(model, config: ModelConfig) -> None: if layer._is_linear: g = layer.linear_attn g.in_proj_qkvz = GGUFLinear( - config.hidden_size, g.conv_dim + g.value_dim, GGML_Q8_0, has_bias=False + config.hidden_size, g.conv_dim + g.value_dim, GGML_Q8_0, + has_bias=False, quant_role="gdn_qkvz" ) g.out_proj = GGUFLinear( - g.value_dim, config.hidden_size, GGML_Q8_0, has_bias=False + g.value_dim, config.hidden_size, GGML_Q8_0, has_bias=False, quant_role="gdn_out" ) else: attn = layer.self_attn @@ -177,17 +219,20 @@ def convert_qwen35moe_to_gguf(model, config: ModelConfig) -> None: config.hidden_size, attn.num_q * attn.head_dim * 2 + 2 * attn.kv_attn_dim, GGML_Q8_0, - has_bias=False, + has_bias=False, quant_role="attention_qkv", ) attn.o_proj = GGUFLinear( - attn.qo_attn_dim, config.hidden_size, GGML_Q8_0, has_bias=False + attn.qo_attn_dim, config.hidden_size, GGML_Q8_0, has_bias=False, + quant_role="attention_output" ) m = layer.mlp m.shared_expert.gate_up_proj = GGUFLinear( - config.hidden_size, 2 * shared_I, GGML_Q8_0, has_bias=False + config.hidden_size, 2 * shared_I, GGML_Q8_0, has_bias=False, + quant_role="shared_gate_up" ) m.shared_expert.down_proj = GGUFLinear( - shared_I, config.hidden_size, GGML_Q8_0, has_bias=False + shared_I, config.hidden_size, GGML_Q8_0, has_bias=False, + quant_role="shared_down" ) @@ -289,21 +334,18 @@ def iter_gguf_weights( include_moe_experts: bool, include_non_moe: bool, ) -> Iterator[tuple[str, torch.Tensor]]: - """Yield (param_name, tensor) for every non-expert qwen3_5_moe param. + """Yield packed state tensors for Qwen3.5-MoE. Quantized projections stay packed and are yielded as ``.qweight`` (uint8); the F32 norms/router/GDN b,a dequantize to bf16; conv1d/A_log/dt_bias stay fp32. Full-attention q/k/v -> ``self_attn.qkv_proj.qweight``, GDN qkv|z -> ``linear_attn.in_proj_qkvz.qweight`` (Q8_0, concat along the output dim), GDN b|a -> ``linear_attn.in_proj_ba.weight`` (dense - bf16). Routed experts are skipped (offload cache). + bf16). Routed experts are skipped for offload and yielded as native packed banks for + resident mode. """ from freetoken.models.gguf.reader import iter_gguf_tensors from freetoken.utils import cached_load_hf_config - assert not include_moe_experts, ( - "qwen3.5-moe GGUF stores experts as Q4_K/Q5_K/Q6_K and only supports the offload " - "backend; experts are loaded into the offload cache via load_gguf_expert_sources()." - ) assert include_non_moe _require_tp1("weight loading") @@ -324,6 +366,7 @@ def iter_gguf_weights( qkvz_buf: dict[int, dict[str, torch.Tensor]] = {} ba_buf: dict[int, dict[str, torch.Tensor]] = {} shexp_buf: dict[int, dict[str, torch.Tensor]] = {} + expert_buf: dict[int, dict[str, tuple[torch.Tensor, int]]] = {} for t in iter_gguf_tensors(model_path): name = t.name @@ -340,8 +383,38 @@ def iter_gguf_weights( continue if not name.startswith("blk."): continue - if any(name.endswith(sfx) for sfx in _EXPERT_SUFFIXES): - continue # routed experts -> offload banks + expert_suffix = next((sfx for sfx in _EXPERT_SUFFIXES if name.endswith(sfx)), None) + if expert_suffix is not None: + if not include_moe_experts: + continue # routed experts -> offload banks + layer = int(name.split(".")[1]) + kind = { + "ffn_gate_exps.weight": "gate", + "ffn_up_exps.weight": "up", + "ffn_down_exps.weight": "down", + }[expert_suffix] + expected_type = (GGML_Q4_K,) if kind in ("gate", "up") else (GGML_Q5_K, GGML_Q6_K) + if t.ggml_type not in expected_type: + raise ValueError( + f"{name}: expected {expected_type}, got GGML type {t.ggml_type}" + ) + E = config.num_experts + I = config.moe_intermediate_size + packed = t.packed().reshape(E, I, t.row_bytes) + layer_buf = expert_buf.setdefault(layer, {}) + layer_buf[kind] = (packed, t.ggml_type) + if {"gate", "up", "down"} <= set(layer_buf): + gate, gate_type = layer_buf["gate"] + up, up_type = layer_buf["up"] + down, down_type = layer_buf["down"] + if gate_type != GGML_Q4_K or up_type != GGML_Q4_K: + raise AssertionError(f"layer {layer}: gate/up expert type is not Q4_K") + yield f"model.layers.{layer}.mlp.experts.gate_up_proj", torch.cat( + [gate, up], dim=1 + ) + yield f"model.layers.{layer}.mlp.experts.down_proj", down + del expert_buf[layer] + continue layer = int(name.split(".")[1]) suffix = name.split(".", 2)[2] @@ -430,6 +503,8 @@ def iter_gguf_weights( assert not qkvz_buf, f"incomplete GDN qkvz groups: {sorted(qkvz_buf)}" assert not ba_buf, f"incomplete GDN ba groups: {sorted(ba_buf)}" assert not shexp_buf, f"incomplete shared-expert gate/up: {sorted(shexp_buf)}" + if include_moe_experts: + assert not expert_buf, f"incomplete routed expert groups: {sorted(expert_buf)}" # -------------------------------------------------------------------------------------- @@ -499,10 +574,79 @@ def _load(sink) -> None: return banks +def load_gguf_expert_sources_native( + model_path: str, config: ModelConfig, *, layer_sink=None +) -> dict[str, list[torch.Tensor]]: + """Load Qwen routed experts without Q8 conversion for GPU offload. + + Q5_K and Q6_K have different packed row widths. The cache needs one shape, + so every down row uses a Q6_K-sized stride; Q5_K retains its 176-byte block + prefix and zero tail. ``ggml_moe_a8_vec_strided`` consumes the exact source + type and skips that tail. Converter callers stay on ``load_gguf_expert_sources`` + because its FTW schema intentionally remains the legacy uniform Q8 layout. + """ + from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, alloc_layer_banks + from freetoken.models.gguf.reader import iter_gguf_tensors + + _require_tp1("native expert banks") + L, E = config.num_layers, config.num_experts + H, I = config.hidden_size, config.moe_intermediate_size + h_rb = row_bytes(H, GGML_Q4_K) + q5_rb = row_bytes(I, GGML_Q5_K) + q6_rb = row_bytes(I, GGML_Q6_K) + specs = { + "gate_up": ((E, 2 * I, h_rb), torch.uint8), + "down": ((E, H, q6_rb), torch.uint8), + } + hb = alloc_layer_banks(specs, L) + banks = {name: [b.tensor for b in hb[name]] for name in specs} + + def _load(sink) -> None: + tracker = LayerCompletionTracker(2, hb, sink) if sink is not None else None + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + continue + layer = int(t.name.split(".")[1]) + if t.name.endswith("ffn_gate_exps.weight"): + banks["gate_up"][layer][:, :I].copy_( + t.packed().reshape(E, I, h_rb) + ) + elif t.name.endswith("ffn_up_exps.weight"): + banks["gate_up"][layer][:, I:].copy_( + t.packed().reshape(E, I, h_rb) + ) + elif t.name.endswith("ffn_down_exps.weight"): + if t.ggml_type == GGML_Q5_K: + packed = t.packed().reshape(E, H, q5_rb) + banks["down"][layer][:, :, :q5_rb].copy_(packed) + elif t.ggml_type == GGML_Q6_K: + banks["down"][layer].copy_( + t.packed().reshape(E, H, q6_rb) + ) + else: + raise ValueError( + f"native GGUF down layer {layer} has unsupported type {t.ggml_type}" + ) + else: + continue + if tracker is not None: + tracker.note(layer) + + if layer_sink is not None: + _load(layer_sink) + elif torch.cuda.is_available(): + with PinPipeline() as pins: + _load(pins) + else: + _load(None) + return banks + + __all__ = [ "parse_gguf_config", "iter_gguf_weights", "convert_qwen35moe_to_gguf", "is_gguf_model", "load_gguf_expert_sources", + "load_gguf_expert_sources_native", ] diff --git a/python/freetoken/models/qwen3_5_moe/moe.py b/python/freetoken/models/qwen3_5_moe/moe.py index b5ab1cf3d..0c07244d2 100644 --- a/python/freetoken/models/qwen3_5_moe/moe.py +++ b/python/freetoken/models/qwen3_5_moe/moe.py @@ -66,13 +66,22 @@ class Qwen3_5MoE(BaseOP): def __init__(self, config: ModelConfig, layer_id: int | None = None): weight_format = ( - "fp8_block" if getattr(config, "expert_quant", "none") == "fp8_block" else "bf16" + "gguf" + if getattr(config, "moe_weight_format", None) == "gguf" + else ( + "fp8_block" + if getattr(config, "expert_quant", "none") == "fp8_block" + else "bf16" + ) ) + down_types = getattr(config, "gguf_down_quant_types", ()) + down_type = down_types[layer_id] if weight_format == "gguf" and layer_id is not None else None self.experts = make_moe_layer( config, layer_id=layer_id, renormalize=config.norm_topk_prob, weight_format=weight_format, + gguf_down_quant_type=down_type, ) self.gate = LinearReplicated(config.hidden_size, config.num_experts, has_bias=False) self.shared_expert = _SharedExpert( diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 636d451ee..3bd1ce31f 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -356,6 +356,18 @@ def load_gguf_moe_expert_sources( return loader(model_path, model_config, layer_sink=layer_sink) +def load_gguf_moe_expert_sources_native( + model_path: str, + model_config, + *, + layer_sink=None, +) -> dict: + """Load packed Qwen GGUF expert banks with native Q5_K/Q6_K down rows.""" + _config, spec = _spec_for_model_path(model_path) + loader = _load_attr(spec.module, "load_gguf_expert_sources_native") + return loader(model_path, model_config, layer_sink=layer_sink) + + def _num_moe_layers(config) -> int: value = getattr(config, "num_moe_layers", None) if value is not None: diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index be1fd9668..c67049da1 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -260,11 +260,18 @@ def _gguf_banks(model_path, model_config, device, dtype, dummy, parallel=False, ) if dummy: raise NotImplementedError("gguf expert banks have no dummy path; load the real GGUF") - from freetoken.models.weight import load_gguf_moe_expert_sources + from freetoken.models.weight import ( + load_gguf_moe_expert_sources, + load_gguf_moe_expert_sources_native, + ) - sources = load_gguf_moe_expert_sources(model_path, model_config, layer_sink=layer_sink) + native = decode_target == "gpu" and layer_sink is None + loader = load_gguf_moe_expert_sources_native if native else load_gguf_moe_expert_sources + sources = loader(model_path, model_config, layer_sink=layer_sink) return ExpertBanks( - "gguf", {name: sources[name] for name in _BANK_SCHEMAS["gguf"]}, streamed=layer_sink is not None + "gguf_native" if native else "gguf", + {name: sources[name] for name in _BANK_SCHEMAS["gguf_native" if native else "gguf"]}, + streamed=layer_sink is not None, ) diff --git a/python/freetoken/moe/fused_gguf.py b/python/freetoken/moe/fused_gguf.py index db1e386a0..cf866a9f3 100644 --- a/python/freetoken/moe/fused_gguf.py +++ b/python/freetoken/moe/fused_gguf.py @@ -1,16 +1,18 @@ -"""Grouped expert GEMM over native GGUF Q4_K gate/up + Q8_0 down banks. +"""Expert GEMV over native GGUF Q4_K gate/up + K-quant down banks. Ports vLLM/sglang's ``_fused_moe_gguf`` MMVQ path onto FreeToken's offload-cache interface: experts are streamed to the GPU as packed block bytes and dequantized -*inside* ``ggml_moe_a8_vec`` -- no bf16 expert copy is materialized. ``gate_up`` stays -native Q4_K; ``down`` is stored as Q8_0 (re-quantized at load from the GGUF's -Q5_K/Q6_K -- 8-bit, >= the source precision, so no quality loss) because the offload -cache needs a single uniform per-bank format. We use the MMVQ (vector) kernel for both -prefill and decode, mirroring ``fused_experts_gguf_q4_0``. +*inside* GGUF kernels -- no bf16 expert copy is materialized. ``gate_up`` stays native +Q4_K; legacy CPU/hybrid offload ``down`` is Q8_0, while plain GPU offload and +resident mode retain each layer's native Q5_K/Q6_K type. Decode uses MMVQ; larger +prefill uses aligned grouped MMQ where its shape contract is proven. """ from __future__ import annotations +import os +from dataclasses import dataclass, field + import torch from freetoken.layers.activation import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul @@ -19,6 +21,247 @@ _ACT = {"silu": silu_and_mul, "gelu": gelu_and_mul, "gelu_tanh": gelu_tanh_and_mul} +@dataclass +class MoeDecodeWork: + """Fixed-shape GGUF MoE ABI and reusable decode scratch. + + ``id_space`` is part of call state: ``raw`` IDs address resident expert + banks, ``slot`` IDs address offload-cache rows. Kernels never infer one + space from tensor provenance. Buffers grow only during warmup; graph replay + sees stable addresses after the requested shape has been provisioned. + """ + + phase: str + buffers: dict[str, torch.Tensor] = field(default_factory=dict) + id_space: str | None = None + quant_type: int | None = None + down_quant_type: int | None = None + gate_expert_stride_bytes: int | None = None + gate_row_stride_bytes: int | None = None + down_expert_stride_bytes: int | None = None + down_row_stride_bytes: int | None = None + + def bind( + self, + hidden_states: torch.Tensor, + gate_up_q: torch.Tensor, + down_q: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + id_space: str, + down_quant_type: int, + ) -> None: + if self.phase not in {"moe_decode", "moe_prefill"}: + raise ValueError(f"invalid GGUF MoE phase {self.phase!r}") + if id_space not in {"raw", "slot"}: + raise ValueError(f"invalid GGUF expert ID space {id_space!r}") + if hidden_states.ndim != 2 or gate_up_q.ndim != 3 or down_q.ndim != 3: + raise ValueError("GGUF MoE expects hidden [T,H] and 3D packed banks") + if topk_ids.ndim != 2 or topk_weights.shape != topk_ids.shape: + raise ValueError("top-k IDs and weights must have equal [T,K] shape") + if topk_ids.dtype != torch.int32: + raise ValueError(f"GGUF MoE IDs must be int32, got {topk_ids.dtype}") + if gate_up_q.dtype != torch.uint8 or down_q.dtype != torch.uint8: + raise ValueError("GGUF MoE banks must use packed uint8 storage") + devices = {str(t.device) for t in (hidden_states, gate_up_q, down_q, topk_ids)} + if len(devices) != 1: + raise ValueError(f"GGUF MoE tensors must share device, got {sorted(devices)}") + if hidden_states.shape[0] != topk_ids.shape[0]: + raise ValueError("GGUF MoE hidden/token and route shapes disagree") + if topk_weights.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError(f"unsupported top-k weight dtype {topk_weights.dtype}") + if gate_up_q.stride(2) != 1 or down_q.stride(2) != 1: + raise ValueError("packed GGUF bank innermost stride must be one byte") + self.id_space = id_space + self.quant_type = 12 + self.down_quant_type = int(down_quant_type) + self.gate_expert_stride_bytes = int(gate_up_q.stride(0)) + self.gate_row_stride_bytes = int(gate_up_q.stride(1)) + self.down_expert_stride_bytes = int(down_q.stride(0)) + self.down_row_stride_bytes = int(down_q.stride(1)) + + def reserve(self, name: str, shape: tuple[int, ...], dtype: torch.dtype, device) -> torch.Tensor: + current = self.buffers.get(name) + if current is None or tuple(current.shape) != tuple(shape) or current.dtype != dtype or current.device != device: + current = torch.empty(shape, dtype=dtype, device=device) + self.buffers[name] = current + return current + + def tensor(self, name: str) -> torch.Tensor | None: + return self.buffers.get(name) + + +def _work_buffer(work: MoeDecodeWork | dict[str, torch.Tensor] | None, name: str): + if isinstance(work, MoeDecodeWork): + return work.tensor(name) + return work.get(name) if work is not None else None + + +def _reserve_work_buffer( + work: MoeDecodeWork | dict[str, torch.Tensor] | None, + name: str, + shape: tuple[int, ...], + dtype: torch.dtype, + device, +): + if isinstance(work, MoeDecodeWork): + return work.reserve(name, shape, dtype, device) + if work is None: + return None + value = work.get(name) + if value is None or tuple(value.shape) != tuple(shape) or value.dtype != dtype or value.device != device: + value = torch.empty(shape, dtype=dtype, device=device) + work[name] = value + return value + + +def _reduce_routes( + routes: torch.Tensor, + output: torch.Tensor | None, + weights: torch.Tensor | None = None, +) -> torch.Tensor: + """Weight/reduce [T,K,H] routes with one graph-safe launch when on GPU.""" + if weights is not None and weights.shape != routes.shape[:2]: + raise ValueError("route weights must match the first two route dimensions") + if not routes.is_cuda: + value = routes if weights is None else routes * weights.to(routes.dtype).unsqueeze(-1) + reduced = value.sum(dim=1) + if output is not None: + output.copy_(reduced) + return output + return reduced + if output is None: + output = torch.empty( + (routes.shape[0], routes.shape[2]), dtype=routes.dtype, device=routes.device + ) + if weights is None: + from freetoken.kernel import moe_sum_reduce_triton + + moe_sum_reduce_triton(routes, output) + else: + from freetoken.kernel import moe_weighted_sum_reduce_triton + + moe_weighted_sum_reduce_triton(routes, weights, output) + return output + + +def _gguf_moe_matmul( + x: torch.Tensor, + weights: torch.Tensor, + topk_ids: torch.Tensor, + quant_type: int, + row: int, + dispatch: dict | None, + output: torch.Tensor | None = None, + weight_stride_bytes: int | None = None, + weight_row_stride_bytes: int | None = None, + work: MoeDecodeWork | dict[str, torch.Tensor] | None = None, + quant_x: torch.Tensor | None = None, +) -> torch.Tensor: + """Run vector MMVQ or b10434-style aligned grouped MMQ, preserving route order.""" + from freetoken.kernel.gguf import ( + ggml_moe_a8, + ggml_moe_a8_vec, + ggml_moe_a8_vec_strided, + ggml_moe_get_block_size, + ) + + if work is not None and isinstance(work, MoeDecodeWork): + if work.id_space not in {"raw", "slot"}: + raise ValueError("MoeDecodeWork must be bound before GGUF dispatch") + + if dispatch is not None and dispatch.get("implementation") == "rdna3_mmvdq": + from freetoken.kernel.gguf import ggml_moe_mmvdq_id + + if quant_x is not None: + raise ValueError("rdna3_mmvdq must not receive Q8_1 scratch") + if weight_stride_bytes is None: + weight_stride_bytes = int(weights.stride(0)) + if weight_row_stride_bytes is None: + weight_row_stride_bytes = int(weights.stride(1)) + direct_id_space = work.id_space if isinstance(work, MoeDecodeWork) else "slot" + args = ( + x, weights, topk_ids, int(topk_ids.shape[1]), quant_type, row, x.shape[0], + weight_stride_bytes, weight_row_stride_bytes, direct_id_space, + ) + return ggml_moe_mmvdq_id(*args, output=output) + + if dispatch is not None and dispatch.get("implementation") == "rdna3_mmid": + from freetoken.kernel.gguf import ggml_moe_mmvq_id + + if weight_stride_bytes is None: + weight_stride_bytes = int(weights.stride(0)) + if weight_row_stride_bytes is None: + weight_row_stride_bytes = int(weights.stride(1)) + if work is not None and isinstance(work, MoeDecodeWork): + candidate_id_space = work.id_space + else: + candidate_id_space = "slot" + args = ( + x, weights, topk_ids, int(topk_ids.shape[1]), quant_type, row, x.shape[0], + weight_stride_bytes, weight_row_stride_bytes, candidate_id_space, + ) + if output is not None and quant_x is not None: + return ggml_moe_mmvq_id(*args, output=output, quant_x=quant_x) + if output is not None or quant_x is not None: + raise ValueError("rdna3_mmid requires output and quant_x together") + return ggml_moe_mmvq_id(*args) + + if ( + dispatch is not None + and dispatch.get("implementation") == "ggml_moe_a8" + and weight_stride_bytes is not None + ): + # Grouped strided ABI is not available yet. Keep selection visible and + # fall back explicitly; do not silently relabel strided MMVQ as grouped. + dispatch["reason"] = "grouped strided ABI unavailable; vector fallback" + + # Native mixed-Q5_K/Q6_K cache rows use a Q6_K-sized expert stride. + if weight_stride_bytes is not None: + if weight_row_stride_bytes is None: + weight_row_stride_bytes = int(weights.stride(1)) + args = (x, weights, topk_ids, int(topk_ids.shape[1]), quant_type, row, x.shape[0]) + if output is not None and quant_x is not None: + from freetoken.kernel.gguf import ggml_moe_a8_vec_strided_workspace + + return ggml_moe_a8_vec_strided_workspace( + *args, weight_stride_bytes, weight_row_stride_bytes, output, quant_x + ) + return ( + ggml_moe_a8_vec_strided(*args, weight_stride_bytes, weight_row_stride_bytes) + if output is None + else ggml_moe_a8_vec_strided( + *args, weight_stride_bytes, weight_row_stride_bytes, output + ) + ) + if dispatch is None or dispatch.get("implementation") != "ggml_moe_a8": + args = (x, weights, topk_ids, int(topk_ids.shape[1]), quant_type, row, x.shape[0]) + if output is not None and quant_x is not None: + from freetoken.kernel.gguf import ggml_moe_a8_vec_workspace + + return ggml_moe_a8_vec_workspace(*args, output, quant_x) + return ggml_moe_a8_vec(*args) if output is None else ggml_moe_a8_vec(*args, output) + + from freetoken.moe.fused import moe_align_block_size + + block_size = ggml_moe_get_block_size(quant_type) + sorted_ids, expert_ids, tokens_post_padded = moe_align_block_size( + topk_ids, block_size, weights.shape[0] + ) + return ggml_moe_a8( + x, + weights, + sorted_ids, + expert_ids, + tokens_post_padded, + quant_type, + row, + int(topk_ids.shape[1]), + x.shape[0], + ) + + def fused_experts_gguf( hidden_states: torch.Tensor, gate_up_q: torch.Tensor, # [num_slots, 2I, row_bytes(H, Q4_K)] uint8 @@ -26,9 +269,15 @@ def fused_experts_gguf( topk_weights: torch.Tensor, topk_ids: torch.Tensor, activation: str, + *, + is_prefill: bool = False, + dispatch_metadata: dict | None = None, + down_quant_type: int = GGML_Q8_0, + down_stride_bytes: int | None = None, + down_row_stride_bytes: int | None = None, + work: MoeDecodeWork | None = None, + id_space: str = "slot", ) -> torch.Tensor: - from freetoken.kernel.gguf import ggml_moe_a8_vec - act_fn = _ACT.get(activation) if act_fn is None: raise ValueError(f"unsupported MoE activation {activation!r}") @@ -37,22 +286,180 @@ def fused_experts_gguf( n2 = gate_up_q.shape[1] # 2 * intermediate h = down_q.shape[1] # hidden top_k = topk_ids.shape[1] + gate_dispatch = (dispatch_metadata or {}).get("gate_up") + down_dispatch = (dispatch_metadata or {}).get("down") + # Fused gate/up is an opt-in candidate optimization. Its BF16 activation + # and SiLU rounding differ from proven legacy MMVQ; keep correctness path + # on separate ID-aware MMVQ until model-level greedy parity is proven. + fused_gate_up = ( + gate_dispatch is not None + and gate_dispatch.get("implementation") == "rdna3_mmid" + and activation == "silu" + and os.environ.get("FREETOKEN_GGUF_FUSED_GATE_UP", "0").strip().lower() + in {"1", "true", "yes", "on"} + ) - # "moe_gate_up" / "moe_down" record_function labels = the profiler-segmented - # expert-GEMM halves of the fused MoE forward (Inc 2, .plans/rocm-perf-parity). - with torch.profiler.record_function("moe_gate_up"): - gate_up = ggml_moe_a8_vec( - hidden_states, gate_up_q, topk_ids, top_k, int(GGML_Q4_K), n2, num_tokens + if work is not None: + work.bind( + hidden_states, gate_up_q, down_q, topk_weights, topk_ids, + id_space=id_space, down_quant_type=down_quant_type, + ) + gate_out = None + if not fused_gate_up: + gate_out = work.reserve( + "gate_up", (num_tokens * top_k, n2), hidden_states.dtype, hidden_states.device + ) + inter = work.reserve( + "inter", (num_tokens * top_k, n2 // 2), hidden_states.dtype, hidden_states.device + ) + down_out = work.reserve( + "down", (num_tokens * top_k, h), hidden_states.dtype, hidden_states.device + ) + result = work.reserve( + "output", (num_tokens, h), hidden_states.dtype, hidden_states.device ) - inter = act_fn(gate_up) + gate_quant_x = None + if gate_dispatch is None or gate_dispatch.get("implementation") != "rdna3_mmvdq": + gate_quant_x = work.reserve( + "quant_x_gate", + (num_tokens, ((hidden_states.shape[1] + 511) // 512) * 144), + torch.int32, + hidden_states.device, + ) + down_quant_x = None + if down_dispatch is None or down_dispatch.get("implementation") != "rdna3_mmvdq": + down_quant_x = work.reserve( + "quant_x_down", + (num_tokens * top_k, ((n2 // 2 + 511) // 512) * 144), + torch.int32, + hidden_states.device, + ) + else: + gate_out = inter = down_out = result = gate_quant_x = down_quant_x = None + + # "moe_gate_up" / "moe_down" labels segment both vector decode and grouped prefill. + if fused_gate_up: + from freetoken.kernel.gguf import ggml_moe_gate_up_swiglu_id + + with torch.profiler.record_function("moe_gate_up_swiglu"): + inter = ggml_moe_gate_up_swiglu_id( + hidden_states, + gate_up_q, + topk_ids, + top_k, + n2 // 2, + num_tokens, + int(gate_up_q.stride(0)), + int(gate_up_q.stride(1)), + work.id_space if isinstance(work, MoeDecodeWork) else id_space, + output=inter, + quant_x=gate_quant_x, + ) + else: + with torch.profiler.record_function("moe_gate_up"): + gate_up = _gguf_moe_matmul( + hidden_states, gate_up_q, topk_ids, int(GGML_Q4_K), n2, gate_dispatch, + output=gate_out, quant_x=gate_quant_x, work=work, + ) + with torch.profiler.record_function("moe_activation"): + inter = act_fn(gate_up, out=inter) with torch.profiler.record_function("moe_down"): - out = ggml_moe_a8_vec( - inter, down_q, topk_ids, 1, int(GGML_Q8_0), h, num_tokens * top_k + route_ids = topk_ids.reshape(-1, 1) + out = _gguf_moe_matmul( + inter, down_q, route_ids, int(down_quant_type), h, down_dispatch, + weight_stride_bytes=down_stride_bytes, + weight_row_stride_bytes=down_row_stride_bytes, + output=down_out, + quant_x=down_quant_x, + work=work, ) - out = out.reshape(num_tokens, top_k, h) * topk_weights.reshape(num_tokens, top_k, 1).to( - out.dtype + out = out.reshape(num_tokens, top_k, h) + return _reduce_routes(out, result, topk_weights.reshape(num_tokens, top_k)) + + +def fused_experts_gguf_native( + hidden_states: torch.Tensor, + gate_up_q: torch.Tensor, + down_q: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: str, + *, + down_quant_type: int | None, + is_prefill: bool = False, + workspace: dict[str, torch.Tensor] | None = None, +) -> torch.Tensor: + """Resident native expert path; no cache slot remap or Q8 conversion.""" + from freetoken.kernel.gguf import gguf_dispatch, gguf_runtime_metadata + from freetoken.models.gguf.dequant import GGML_Q5_K, GGML_Q6_K + + if down_quant_type not in (GGML_Q5_K, GGML_Q6_K): + raise ValueError(f"unsupported resident GGUF down type {down_quant_type!r}") + act_fn = _ACT.get(activation) + if act_fn is None: + raise ValueError(f"unsupported MoE activation {activation!r}") + num_tokens = hidden_states.shape[0] + top_k = topk_ids.shape[1] + phase = "moe_prefill" if is_prefill else "moe_decode" + arch = gguf_runtime_metadata().get("arch") + gate_dispatch = gguf_dispatch( + phase, GGML_Q4_K, gate_up_q.shape[1], hidden_states.shape[1], num_tokens, arch + ) + down_dispatch = gguf_dispatch( + phase, down_quant_type, down_q.shape[1], gate_up_q.shape[1] // 2, + num_tokens * top_k, arch ) - return out.sum(dim=1) + if isinstance(workspace, MoeDecodeWork): + return fused_experts_gguf( + hidden_states, + gate_up_q, + down_q, + topk_weights, + topk_ids, + activation, + is_prefill=is_prefill, + dispatch_metadata={"gate_up": gate_dispatch, "down": down_dispatch}, + down_quant_type=down_quant_type, + down_stride_bytes=int(down_q.stride(0)), + down_row_stride_bytes=int(down_q.stride(1)), + work=workspace, + id_space="raw", + ) + with torch.profiler.record_function("moe_gate_up_native"): + gate_out = None + if workspace is not None: + gate_out = workspace.get("gate_up") + shape = (num_tokens * top_k, gate_up_q.shape[1]) + if gate_out is None or tuple(gate_out.shape) != shape or gate_out.dtype != hidden_states.dtype: + gate_out = torch.empty(shape, dtype=hidden_states.dtype, device=hidden_states.device) + workspace["gate_up"] = gate_out + gate_up = _gguf_moe_matmul( + hidden_states, gate_up_q, topk_ids, GGML_Q4_K, + gate_up_q.shape[1], gate_dispatch, gate_out + ) + with torch.profiler.record_function("moe_activation_native"): + inter = None + if workspace is not None: + inter = workspace.get("inter") + shape = (num_tokens * top_k, gate_up_q.shape[1] // 2) + if inter is None or tuple(inter.shape) != shape or inter.dtype != hidden_states.dtype: + inter = torch.empty(shape, dtype=hidden_states.dtype, device=hidden_states.device) + workspace["inter"] = inter + inter = act_fn(gate_up, out=inter) + with torch.profiler.record_function("moe_down_native"): + down_out = None + if workspace is not None: + down_out = workspace.get("down") + shape = (num_tokens * top_k, down_q.shape[1]) + if down_out is None or tuple(down_out.shape) != shape or down_out.dtype != hidden_states.dtype: + down_out = torch.empty(shape, dtype=hidden_states.dtype, device=hidden_states.device) + workspace["down"] = down_out + out = _gguf_moe_matmul( + inter, down_q, topk_ids.reshape(-1, 1), down_quant_type, + down_q.shape[1], down_dispatch, down_out + ) + out = out.reshape(num_tokens, top_k, down_q.shape[1]) + return _reduce_routes(out, None, topk_weights.reshape(num_tokens, top_k)) -__all__ = ["fused_experts_gguf"] +__all__ = ["MoeDecodeWork", "fused_experts_gguf", "fused_experts_gguf_native"] diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 8c08969b4..ba50906c8 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -49,6 +49,10 @@ # down re-quantized to Q8_0 [L*E, H, row_bytes(I,Q8_0)] (a uniform format the cache # can hold; the source GGUF down is Q5_K/Q6_K, both >= ... 8-bit >= source precision). "gguf": ("gate_up", "down"), + # Native Qwen GGUF offload: gate_up Q4_K plus Q5_K/Q6_K down rows stored in + # one Q6_K-sized row stride. The loader pads only Q5_K row tails; kernels + # receive the exact type and explicit expert stride. + "gguf_native": ("gate_up", "down"), # native ModelOpt rows for the Triton inline-dequant kernels: packed e2m1 codes + # fp8-e4m3 per-16 block scales + per-output-row fp16 globals (w1/w3 carry distinct # globals, and folding them into the e4m3 block scales would underflow) @@ -99,6 +103,9 @@ def fp8_block_scale_pad(rows: int, cols: int) -> int: "q4_0": lambda H, I: 2 * I * (H // 32) * 18 + H * (I // 32) * 18, # gate_up Q4_K row_bytes(H, Q4_K)=H//256*144; down Q8_0 row_bytes(I, Q8_0)=I//32*34. "gguf": lambda H, I: 2 * I * (H // 256) * 144 + H * (I // 32) * 34, + # Native down rows use largest supported K-quant row (Q6_K) as uniform + # cache stride; Q5_K data keeps its native packed prefix. + "gguf_native": lambda H, I: 2 * I * (H // 256) * 144 + H * (I // 256) * 210, "nvfp4": lambda H, I: 2 * I * (H // 2 + H // 16 + 2) + H * (I // 2 + I // 16 + 2), "mxfp4": lambda H, I: 2 * I * (H // 2 + H // 32 + 2) + H * (I // 2 + I // 32 + 2), "ds_fp4": lambda H, I: 2 * I * (H // 2 + H // 32) + H * (I // 2 + I // 32), @@ -123,7 +130,8 @@ class OffloadMoeCache: # coalesced runs). Requires prefill_overlap, cache_size > 2 * num_experts and # the fused copy plan; silently falls back to the full-layer copy otherwise. prefill_hit_d2d: bool = False - # "bf16" (default, dense expert weights) or one of the NVFP4 bank layouts: + # "bf16" (default, dense expert weights), native GGUF layouts ("gguf" legacy Q8_0 + # down and "gguf_native" mixed Q5_K/Q6_K down), or one of the NVFP4 bank layouts: # "nvfp4" (native ModelOpt rows, FreeToken Triton kernels), "nvfp4_marlin" # (Marlin-tiled, vLLM W4A16 GEMM, sm_80-99) or "nvfp4_b12x" (flashinfer SM12x # W4A16); or "mxfp4_triton" (gpt-oss transposed split-K GEMV decode + _t grouped diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index bd35d373e..81e26b374 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -25,6 +25,7 @@ load_tokenizer, load_toolcall_anchor_id, ) +from freetoken.utils.step_profiler import profiler_phase, step_profiler from .cache import CacheManager from .config import SchedulerConfig @@ -228,26 +229,26 @@ def overlap_loop(self, last_data: ForwardData | None) -> ForwardData | None: # table_idx can have its freshly copied prompt clobbered by the prior occupant's # still-pending output write -- corrupting tokens (e.g. dropping an image # placeholder, which the multimodal merge then rejects). - self.stream.wait_stream(self.engine.stream) - forward_input = self._schedule_next_batch() - ongoing_data = None - if forward_input is not None: - with self.engine_stream_ctx: # run the batch in the engine's stream - self.engine.stream.wait_stream(self.stream) - # COW-restore GDN snapshots for prefix hits ON THE ENGINE STREAM, after the - # cross-stream wait and before the forward reads the live slot (program order - # vs the prior batch's snapshot writes). Doing this on self.stream would race. - self._restore_linear_states(forward_input.batch) - ongoing_data = (forward_input, self._forward(forward_input)) - - # The drain issues GPU-visible writes to state the batch just launched still reads: the - # page-table re-point and, for the paged-SWA pools, the full->swa (DSV4: full->window) - # sentinel scatter. DSV4 stages the page table at replay time and translates - # full_to_window INSIDE the captured graph, so an unordered drain can redirect an - # in-flight forward. copy_done only covers batch N; order against N+1 explicitly. - self.stream.wait_stream(self.engine.stream) - self._process_last_data(last_data) - self._flush_abort_acks() + with step_profiler(): + self.stream.wait_stream(self.engine.stream) + forward_input = self._schedule_next_batch() + ongoing_data = None + if forward_input is not None: + with self.engine_stream_ctx: # run the batch in the engine's stream + self.engine.stream.wait_stream(self.stream) + # COW-restore GDN snapshots for prefix hits ON THE ENGINE STREAM, after the + # cross-stream wait and before the forward reads the live slot (program order + # vs the prior batch's snapshot writes). Doing this on self.stream would race. + self._restore_linear_states(forward_input.batch) + ongoing_data = (forward_input, self._forward(forward_input)) + + # The drain issues GPU-visible writes to state the batch just launched still reads: + # the page-table re-point and, for paged-SWA pools, the full->swa sentinel scatter. + # DSV4 stages the page table at replay time and translates full_to_window INSIDE the + # captured graph, so an unordered drain can redirect an in-flight forward. + self.stream.wait_stream(self.engine.stream) + self._process_last_data(last_data) + self._flush_abort_acks() return ongoing_data def normal_loop(self) -> None: @@ -267,15 +268,16 @@ def normal_loop(self) -> None: ): self._execute_pending_rebuild() - forward_input = self._schedule_next_batch() - ongoing_data = None - if forward_input is not None: - # already inside engine_stream_ctx (run_forever); restore on the engine stream - self._restore_linear_states(forward_input.batch) - ongoing_data = (forward_input, self._forward(forward_input)) + with step_profiler(): + forward_input = self._schedule_next_batch() + ongoing_data = None + if forward_input is not None: + # already inside engine_stream_ctx (run_forever); restore on the engine stream + self._restore_linear_states(forward_input.batch) + ongoing_data = (forward_input, self._forward(forward_input)) - self._process_last_data(ongoing_data) - self._flush_abort_acks() + self._process_last_data(ongoing_data) + self._flush_abort_acks() @torch.inference_mode() def run_forever(self) -> NoReturn: @@ -300,99 +302,112 @@ def shutdown(self) -> None: self.engine.shutdown() def _process_last_data(self, last_data: ForwardData | None) -> None: + with profiler_phase("token_d2h_event_drain"): + # Call helper through class so lightweight unbound-method test stubs do not + # need to carry a bound helper attribute. + Scheduler._process_last_data_impl(self, last_data) + + def _process_last_data_impl(self, last_data: ForwardData | None) -> None: if last_data is None: return batch, (_, next_tokens_cpu, copy_done) = last_data[0].batch, last_data[1] - copy_done.synchronize() - reply: List[DetokenizeMsg] = [] - new_finished_reqs: Set[Req] = set() - with self.cache_manager.lazy_free_region(): - for i, req in enumerate(batch.reqs): - if isinstance(req, ChunkedReq): - # Don't cache intermediate chunks; the full prompt is cached once when the - # final chunk is processed. Caching here snapshots a handle the next chunk - # already copied (overlap), so cache_req double-frees the prior chunk. + with profiler_phase("token_copy_wait"): + # D2H observation runs on TokenStaging.copy_stream. Query first so already-complete + # copies do not block CPU; synchronize remains exact fallback for slow copies. + query = getattr(copy_done, "query", None) + if query is None or not query(): + copy_done.synchronize() + with profiler_phase("decode_result_accounting"): + reply: List[DetokenizeMsg] = [] + new_finished_reqs: Set[Req] = set() + with self.cache_manager.lazy_free_region(): + for i, req in enumerate(batch.reqs): + if isinstance(req, ChunkedReq): + # Don't cache intermediate chunks; the full prompt is cached once when the + # final chunk is processed. Caching here snapshots a handle the next chunk + # already copied (overlap), so cache_req double-frees the prior chunk. + if req.aborted: + # Aborted mid-chunked-prefill while this chunk was in flight: the abort + # popped the pending continuation (no next chunk launches), and this + # drain point frees the chunk's pages/slots exactly once. + self._free_req_resources(req) + continue if req.aborted: - # Aborted mid-chunked-prefill while this chunk was in flight: the abort - # popped the pending continuation (no next chunk launches), and this - # drain point frees the chunk's pages/slots exactly once. + # Aborted while this final-chunk prefill / decode step was in flight: free + # here (the forward is drained) and finish the request. No DetokenizeMsg -- + # the abort ack flushed after this method stays the uid's terminal reply. + self.decode_manager.remove_req(req) self._free_req_resources(req) - continue - if req.aborted: - # Aborted while this final-chunk prefill / decode step was in flight: free - # here (the forward is drained) and finish the request. No DetokenizeMsg -- - # the abort ack flushed after this method stays the uid's terminal reply. - self.decode_manager.remove_req(req) - self._free_req_resources(req) - new_finished_reqs.add(req) - continue - if req in self.finished_reqs: - # Overlap scheduling launched one more decode step for a request that - # already terminated (filter_reqs keeps it while output budget remains, - # and the next batch is scheduled before this drain runs). Its resources - # are freed below/already; shipping this token would append past the - # client's terminal reply. - continue - next_token = next_tokens_cpu[i] - req.append_host(next_token.unsqueeze(0)) - next_token = int(next_token.item()) - # EOS / stop-string -> "stop", output budget exhausted -> "length"; - # EOS and stop strings win over length. - hit_length = not req.can_decode - hit_eos = ( - not req.sampling_params.ignore_eos and next_token in self.eos_token_ids - ) - matched_stop = ( - self._match_stop_str(req) - if not hit_eos and req.sampling_params.stop_strs - else None - ) - finished = hit_length or hit_eos or matched_stop is not None - finish_reason = ( - ("stop" if (hit_eos or matched_stop is not None) else "length") - if finished - else None - ) - if ( - next_token == self.toolcall_anchor_id - and req.toolcall_anchor_len is None - and not finished - ): - req.toolcall_anchor_len = req.input_ids.numel() - reply.append( - DetokenizeMsg( - uid=req.uid, - next_token=next_token, - finished=finished, - finish_reason=finish_reason, - matched_stop=matched_stop, - stop_strs=req.sampling_params.stop_strs or None, + new_finished_reqs.add(req) + continue + if req in self.finished_reqs: + # Overlap scheduling launched one more decode step for a request that + # already terminated (filter_reqs keeps it while output budget remains, + # and the next batch is scheduled before this drain runs). Its resources + # are freed below/already; shipping this token would append past the + # client's terminal reply. + continue + next_token = next_tokens_cpu[i] + req.append_host(next_token.unsqueeze(0)) + next_token = int(next_token.item()) + # EOS / stop-string -> "stop", output budget exhausted -> "length"; + # EOS and stop strings win over length. + hit_length = not req.can_decode + hit_eos = ( + not req.sampling_params.ignore_eos and next_token in self.eos_token_ids + ) + matched_stop = ( + self._match_stop_str(req) + if not hit_eos and req.sampling_params.stop_strs + else None + ) + finished = hit_length or hit_eos or matched_stop is not None + finish_reason = ( + ("stop" if (hit_eos or matched_stop is not None) else "length") + if finished + else None + ) + if ( + next_token == self.toolcall_anchor_id + and req.toolcall_anchor_len is None + and not finished + ): + req.toolcall_anchor_len = req.input_ids.numel() + reply.append( + DetokenizeMsg( + uid=req.uid, + next_token=next_token, + finished=finished, + finish_reason=finish_reason, + matched_stop=matched_stop, + stop_strs=req.sampling_params.stop_strs or None, + ) ) - ) - # NOTE: overlap scheduling may make the request freed twice, skip second free - if finished and req not in self.finished_reqs: - self.decode_manager.remove_req(req) - self._free_req_resources(req) - new_finished_reqs.add(req) - elif batch.is_prefill and req.table_idx != -1: - # for prefill, non-chunk req, cache the prefix. - # Polymorphic: the DSV4 naive manager keeps the request's slots (no-op); - # the generic manager inserts the prefix into its radix/naive cache. - # table_idx == -1 is defense-in-depth: aborts mark in-flight requests - # instead of freeing them (handled above), so a freed request should - # never reach this commit -- but if a future path frees one early, skip - # rather than re-read the freed page-table row (and on hybrid, deref the - # None'd GDN ping-pong slots). - self.cache_manager.cache_req(req, finished=False) + # NOTE: overlap scheduling may make the request freed twice, skip second free + if finished and req not in self.finished_reqs: + self.decode_manager.remove_req(req) + self._free_req_resources(req) + new_finished_reqs.add(req) + elif batch.is_prefill and req.table_idx != -1: + # for prefill, non-chunk req, cache the prefix. + # Polymorphic: the DSV4 naive manager keeps the request's slots (no-op); + # the generic manager inserts the prefix into its radix/naive cache. + # table_idx == -1 is defense-in-depth: aborts mark in-flight requests + # instead of freeing them (handled above), so a freed request should + # never reach this commit -- but if a future path frees one early, skip + # rather than re-read the freed page-table row (and on hybrid, deref the + # None'd GDN ping-pong slots). + self.cache_manager.cache_req(req, finished=False) self.finished_reqs = new_finished_reqs # Stamp each reply with the post-batch KV page occupancy so the frontend (shell # status bar) can show live KV usage without a separate query. - used, total = self._kv_usage_pages() - mamba_slots = self._mamba_slot_usage() - swa_tokens = self._swa_token_usage() + with profiler_phase("decode_usage_snapshot"): + used, total = self._kv_usage_pages() + mamba_slots = self._mamba_slot_usage() + swa_tokens = self._swa_token_usage() if reply: mem = self._gpu_mem_bytes() mamba_used, mamba_total = mamba_slots or (0, 0) @@ -405,18 +420,26 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: m.swa_used_tokens = swa_used m.swa_total_tokens = swa_total m.gpu_mem_bytes = mem - self.status_reporter.report_batch( - batch, - running_reqs=len(self.decode_manager.running_reqs), - queue_reqs=len(self.prefill_manager.pending_list), - kv_used_pages=used, - kv_total_pages=total, - page_size=self.config.page_size, - mamba_slots=mamba_slots, - swa_tokens=swa_tokens, - moe_stats=self._moe_stats_snapshot(), - ) - self.send_result(reply) + with profiler_phase("decode_status_report"): + stats_due = getattr(self.status_reporter, "decode_stats_due", None) + moe_stats = ( + self._moe_stats_snapshot() + if stats_due is None or stats_due(batch) + else None + ) + self.status_reporter.report_batch( + batch, + running_reqs=len(self.decode_manager.running_reqs), + queue_reqs=len(self.prefill_manager.pending_list), + kv_used_pages=used, + kv_total_pages=total, + page_size=self.config.page_size, + mamba_slots=mamba_slots, + swa_tokens=swa_tokens, + moe_stats=moe_stats, + ) + with profiler_phase("decode_result_send"): + self.send_result(reply) def _moe_stats_snapshot(self) -> dict | None: """Per-window MoE cache hit/miss stats for the decode log line, or None when @@ -773,6 +796,10 @@ def _log_cache_geometry(self, event: str) -> None: logger.warning(f"could not log cache geometry: {e!r}") def _prepare_batch(self, batch: Batch) -> ForwardInput: + with profiler_phase("scheduler_prepare"): + return self._prepare_batch_impl(batch) + + def _prepare_batch_impl(self, batch: Batch) -> ForwardInput: self.engine.graph_runner.pad_batch(batch) self._forward_iter += 1 if batch.is_decode: @@ -793,10 +820,11 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput: self.cache_manager.allocate_paged(batch.reqs) if batch.is_prefill: self._gather_multimodal(batch) - batch.positions = _make_positions(batch, self.device) - input_mapping = _make_input_tuple(batch, self.device) - write_mapping = _make_write_tuple(batch, self.device) - batch.out_loc = self.engine.page_table[input_mapping] + with profiler_phase("token_pool_gather"): + batch.positions = _make_positions(batch, self.device) + input_mapping = _make_input_tuple(batch, self.device) + write_mapping = _make_write_tuple(batch, self.device) + batch.out_loc = self.engine.page_table[input_mapping] if self.engine.linear_state_pool is not None: if batch.is_decode: # GPU GDN-state slot (one per padded request) for the decode gather/scatter; @@ -823,7 +851,8 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput: # This batch's padded per-row page-table rows. Backends that snapshot the table for # a captured replay (DSV4) read them in prepare_metadata / prepare_for_replay. batch.active_table_idx = input_mapping[0].view(-1) - self.engine.attn_backend.prepare_metadata(batch) + with profiler_phase("attention_metadata"): + self.engine.attn_backend.prepare_metadata(batch) return ForwardInput( batch=batch, sample_args=self.engine.sampler.prepare(batch), diff --git a/python/freetoken/scheduler/status.py b/python/freetoken/scheduler/status.py index c7be5e8f1..7be608b2e 100644 --- a/python/freetoken/scheduler/status.py +++ b/python/freetoken/scheduler/status.py @@ -59,6 +59,13 @@ def report_batch( moe_stats=moe_stats, ) + def decode_stats_due(self, batch: Batch) -> bool: + """Whether next decode report will log and may safely read device counters.""" + return ( + batch.is_decode + and (self._decode_forward_count + 1) % self.decode_log_interval == 0 + ) + def _report_prefill( self, batch: Batch, diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 3e2acc854..a1069feef 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -176,6 +176,8 @@ class FrontendManager: 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) + # Execution mode and graph/resident telemetry delivered with readiness metadata. + execution: Dict[str, Any] = field(default_factory=dict) # 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 @@ -791,6 +793,9 @@ def cache_geometry(state: Any) -> dict: # Engine's exact total cache VRAM budget (all pools), from the ("meta", …) ack. 0 when # unknown (pre-budget engine) — the desktop then reverse-derives it from the limits. "cache_budget_bytes": int(getattr(state, "cache_budget_bytes", 0) or 0), + # Engine-owned observed storage/graph/execution metadata. Requested CLI + # values stay separate; benchmark gates consume this server observation. + "execution": dict(getattr(state, "execution", None) or {}), "reasoning": reasoning, } # Per-pool slider bounds, sized against the cache budget the rebuild fit-check actually @@ -1013,6 +1018,7 @@ def _on_meta(meta: dict) -> 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.execution = dict(meta.pop("execution", 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 cb1662bd4..d07ecf16c 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -7,6 +7,7 @@ import torch from freetoken.distributed import DistributedInfo +from freetoken.engine.config import KVStorageType from freetoken.scheduler import SchedulerConfig from freetoken.utils import init_logger @@ -383,6 +384,16 @@ def _infer_reasoning_parser(model_path: str) -> str | None: " the first one is used for prefill and the second one for decode.", ) + parser.add_argument( + "--kv-storage", + "--kv-type", + dest="kv_storage_type", + type=lambda value: KVStorageType.parse(value), + choices=list(KVStorageType), + default=ServerArgs.kv_storage_type, + help="Physical KV format: bf16, fp16, or opt-in q8_0.", + ) + parser.add_argument( "--model-source", type=str, @@ -732,9 +743,9 @@ def _infer_reasoning_parser(model_path: str) -> str | None: kwargs["tp_info"] = DistributedInfo(0, kwargs["tensor_parallel_size"]) del kwargs["tensor_parallel_size"] - # ROCm (AMD) has no NVIDIA-native NVFP4/Marlin path. Reject NVIDIA-only NVFP4 backends - # at parse time with a clean error, and warn for the resident fused MoE backend (the - # offload/cpu/hybrid family is the supported AMD path). + # ROCm (AMD) has no NVIDIA-native NVFP4/Marlin path. Reject those backends at parse + # time with a clean error. GGUF resident fused MoE is resolved later, after metadata + # preflight, so it must remain selectable here. from freetoken.utils.arch import is_rocm if is_rocm(): @@ -744,13 +755,6 @@ def _infer_reasoning_parser(model_path: str) -> str | None: f"--nvfp4-backend {nvfp4} is NVIDIA-only and unavailable on ROCm/AMD; " f"use --nvfp4-backend triton (inline-dequant) or auto." ) - if kwargs.get("moe_backend") == "fused": - logger = init_logger(__name__) - logger.warning( - "--moe-backend fused relies on NVIDIA-native fused GEMM; on ROCm/AMD the " - "supported family is offload/hybrid/cpu (the triton/offload path)." - ) - result = ServerArgs(**kwargs) logger = init_logger(__name__) logger.info(f"Parsed arguments:\n{result}") diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py index 9fcb44856..c88c4ff5c 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -156,20 +156,18 @@ def start_subprocess() -> "BackendHandle": from .supervisor import BackendHandle mp.set_start_method("spawn", force=True) - # Graph-capture variant env (Inc 6, .plans/rocm-perf-parity): the gate probes - # rocBLAS/prewarm variants in subprocesses and may require e.g. - # TORCH_BLAS_PREFER_HIPBLASLT=0 for capture to be viable. Apply it to THIS - # (supervisor) process before the workers are spawned so the spawned engine - # worker inherits it from process start — before any torch import/GEMM; a late - # write at capture time may no-op depending on torch's BLAS-preference caching. - from freetoken.utils.graph_gate import graph_capture_env + # Resolve one BLAS policy before workers spawn. Explicit FREETOKEN_ROCM_BLAS + # wins over graph-gate auto output; worker inherits env before torch import/GEMM. + # A late write at capture time may no-op after PyTorch caches BLAS preference. + from freetoken.utils.graph_gate import graph_capture_env, rocm_blas_report gate_env = graph_capture_env() if gate_env: os.environ.update(gate_env) logger.info( - f"graph-capture gate env applied to workers: {sorted(gate_env)}" + f"ROCm BLAS/graph env applied to workers: {sorted(gate_env)}" ) + logger.info(f"BLAS policy before worker spawn: {rocm_blas_report()}") detach = server_args.shell_mode # see _detach_process_group world_size = server_args.tp_info.size diff --git a/python/freetoken/server/stats.py b/python/freetoken/server/stats.py index 76c6c308a..f148c1f27 100644 --- a/python/freetoken/server/stats.py +++ b/python/freetoken/server/stats.py @@ -161,6 +161,7 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: "swa": swa, "vram_bytes": tr.vram_bytes, "gpus": list(getattr(state, "gpus", None) or []), + "execution": dict(getattr(state, "execution", None) or {}), "throughput": { "decode_tps": round(tr.decode_tps(), 1), "prefill_tps": round(tr.prefill_tps(), 1), diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py index d138ac1a6..fd289efa3 100644 --- a/python/freetoken/tokenizer/detokenize.py +++ b/python/freetoken/tokenizer/detokenize.py @@ -91,6 +91,14 @@ def discard(self, uid: int) -> None: self.decode_map.pop(uid, None) def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: + try: + import torch + except ImportError: + return self._detokenize(msgs) + with torch.profiler.record_function("detokenize"): + return self._detokenize(msgs) + + def _detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: read_ids: List[List[int]] = [] surr_ids: List[List[int]] = [] for msg in msgs: diff --git a/python/freetoken/utils/graph_gate.py b/python/freetoken/utils/graph_gate.py index b521e921b..f1f8e14a6 100644 --- a/python/freetoken/utils/graph_gate.py +++ b/python/freetoken/utils/graph_gate.py @@ -16,6 +16,116 @@ from functools import lru_cache _CACHE_FILE = "freetoken_graph_gate.json" +_ROCM_BLAS_VALUES = frozenset({"auto", "hipblas", "hipblaslt", "rocblas"}) + + +def _rocm_blas_request(value: str | None = None) -> str: + requested = (value if value is not None else os.environ.get("FREETOKEN_ROCM_BLAS", "auto")) + requested = requested.strip().lower() + if requested not in _ROCM_BLAS_VALUES: + choices = "auto, hipblas, hipblaslt, rocblas" + raise ValueError(f"FREETOKEN_ROCM_BLAS={requested!r}: expected {choices}") + return requested + + +def _is_rocm() -> bool: + try: + import torch + + return torch.version.hip is not None + except Exception: + return False + + +def _blas_env(requested: str) -> dict[str, str]: + if requested in {"hipblas", "rocblas"}: + return {"TORCH_BLAS_PREFER_HIPBLASLT": "0"} + if requested == "hipblaslt": + return {"TORCH_BLAS_PREFER_HIPBLASLT": "1"} + return {} + + +def _effective_blas() -> tuple[str | None, str]: + """Return normalized PyTorch BLAS backend plus verification detail.""" + try: + import torch + + if not _is_rocm(): + return "not-applicable", "not-applicable" + api = getattr(torch.backends.cuda, "preferred_blas_library", None) + if api is None: + return None, "preferred_blas_library unavailable" + raw = api() + name = getattr(raw, "name", str(raw)).lower() + if "lt" in name: + return "hipblaslt", "reported by torch.backends.cuda.preferred_blas_library" + if "cublas" in name or "hipblas" in name: + return "hipblas", "reported by torch.backends.cuda.preferred_blas_library" + return name, "unsupported backend reported by PyTorch" + except Exception as exc: + return None, f"backend report failed: {type(exc).__name__}: {exc}" + + +def rocm_blas_report(requested: str | None = None, *, gate: dict | None = None) -> dict: + """Describe requested/effective ROCm BLAS policy without changing process state.""" + requested = _rocm_blas_request(requested) + rocm = _is_rocm() + gate = gate if gate is not None else (run_graph_gate() if rocm and requested == "auto" else None) + env = _blas_env(requested) + source = "explicit" if requested != "auto" else "inherited" + if requested == "auto" and gate and gate.get("ok") and gate.get("env"): + env = dict(gate["env"]) + source = "graph_gate" + effective, detail = _effective_blas() + expected = {"hipblas": "hipblas", "rocblas": "hipblas", "hipblaslt": "hipblaslt"}.get(requested) + if not rocm: + verification = "not-applicable" + elif effective is None: + verification = "unverified" + elif expected is not None and effective != expected: + verification = "mismatch" + else: + verification = "verified" + try: + import torch + + torch_version = torch.__version__ + rocm_version = torch.version.hip + device = torch.cuda.get_device_name(torch.cuda.current_device()) if torch.cuda.is_available() else None + except Exception: + torch_version = None + rocm_version = None + device = None + return { + "requested": requested, + "effective": effective, + "env": env, + "source": source, + "torch_version": torch_version, + "rocm_version": rocm_version, + "device": device, + "verification": verification, + "detail": detail, + } + + +def resolve_rocm_blas_env(*, gate: dict | None = None) -> dict[str, str]: + """Resolve one BLAS env map; explicit policy overrides graph-gate auto output.""" + requested = _rocm_blas_request() + if not _is_rocm(): + return {} + if requested != "auto": + try: + import torch + + if getattr(torch.backends.cuda, "preferred_blas_library", None) is None: + raise RuntimeError("torch.backends.cuda.preferred_blas_library unavailable") + except Exception as exc: + raise RuntimeError(f"explicit ROCm BLAS policy unavailable: {exc}") from exc + return _blas_env(requested) + if gate is None: + gate = run_graph_gate() + return dict(gate.get("env") or {}) if gate.get("ok") else {} def _cache_dir() -> str: @@ -144,9 +254,15 @@ def graph_capture_env() -> dict[str, str]: capture time, where a late env write may no-op. """ try: + requested = _rocm_blas_request() + if requested != "auto" and _is_rocm(): + return resolve_rocm_blas_env() result = run_graph_gate() - env = result.get("env") or {} - return env if result.get("ok") else {} + return resolve_rocm_blas_env(gate=result) + except ValueError: + raise + except RuntimeError: + raise except Exception: return {} diff --git a/python/freetoken/utils/step_profiler.py b/python/freetoken/utils/step_profiler.py index 693d0944f..20a302bc8 100644 --- a/python/freetoken/utils/step_profiler.py +++ b/python/freetoken/utils/step_profiler.py @@ -1,4 +1,4 @@ -"""Env-gated per-step torch.profiler wrapper (stage-level decode breakdowns). +"""Env-gated per-step profiler and optional ROCm marker wrapper. ``FREETOKEN_TORCH_PROFILE="::"``, unset/empty = no-op (the only hot-path cost is one cached check). Skip the first ``warm`` call(s) of the wrapped @@ -8,7 +8,11 @@ - a chrome trace to ```` and - a top-kernels table to ``-kernels.log``. -This is the Inc 2 instrument of .plans/rocm-perf-parity: NVTX is a no-op on ROCm, +``FREETOKEN_ROCTX_MARKERS=1`` emits low-overhead ROCTX ranges without torch +profiling. A non-boolean value also retains marker records in the profiler's +``-markers.json`` sidecar. + +This is the Inc 1 instrument of .plans/rocm-parity-next: NVTX is a no-op on ROCm, so the trace segments by the explicit ``torch.profiler.record_function`` range names (moe_router, moe_gate_up, moe_down, attn, Sampler) that appear as table row names on both backends. @@ -16,7 +20,13 @@ from __future__ import annotations +import atexit import os +import ctypes +import ctypes.util +import json +import statistics +import time _NO_SPEC = object() # "not parsed yet" sentinel @@ -44,12 +54,28 @@ def _read_spec() -> tuple[int, int, str] | None: class _State: - __slots__ = ("spec", "calls", "done") + __slots__ = ( + "spec", + "calls", + "done", + "prof", + "markers_enabled", + "marker_output", + "markers", + "roctx", + "roctx_error", + ) def __init__(self) -> None: self.spec: tuple[int, int, str] | None = _NO_SPEC # type: ignore[assignment] self.calls = 0 self.done = False + self.prof = None + self.markers_enabled: bool | None = None + self.marker_output: str | None = None + self.markers: list[dict] = [] + self.roctx = None + self.roctx_error = None _state = _State() @@ -70,55 +96,183 @@ def __exit__(self, *exc): class _ProfilerCtx: """One wrapped step inside the profiled window; exports on the last step's exit.""" - __slots__ = ("prof", "final") + __slots__ = ("prof", "final", "step", "phase", "stream_id") - def __init__(self, prof, final: bool) -> None: + def __init__(self, prof, final: bool, step: int, phase: str, stream_id: str | None) -> None: self.prof = prof self.final = final + self.step = step + self.phase = phase + self.stream_id = stream_id if final: _state.done = True def __enter__(self): - self.prof.__enter__() + _mark(self.step, self.phase, "begin", self.stream_id) + return self def __exit__(self, exc_type=None, exc=None, tb=None) -> bool: # noqa: ANN001 - self.prof.__exit__(exc_type, exc, tb) + _mark(self.step, self.phase, "end", self.stream_id) if self.final: + if self.prof is not None: + self.prof.__exit__(exc_type, exc, tb) warm, steps, out = _state.spec or (0, 0, "") - _export(self.prof, out, steps) + if self.prof is not None: + _export(self.prof, out, steps) + return False + + +class _PhaseCtx: + """Low-overhead phase range nested inside a scheduler step.""" + + __slots__ = ("phase", "step", "stream_id", "record") + + def __init__(self, phase: str, step: int, stream_id: str | None) -> None: + self.phase = phase + self.step = step + self.stream_id = stream_id + self.record = None + + def __enter__(self): + _mark(self.step, self.phase, "begin", self.stream_id) + # record_function is useful in torch traces even when ROCTX is unavailable. Import + # lazily so disabled profiling has no torch dependency in this helper. + try: + import torch.profiler + + self.record = torch.profiler.record_function(self.phase) + self.record.__enter__() + except Exception: + self.record = None + return self + + def __exit__(self, exc_type=None, exc=None, tb=None) -> bool: # noqa: ANN001 + if self.record is not None: + self.record.__exit__(exc_type, exc, tb) + _mark(self.step, self.phase, "end", self.stream_id) return False _NULL = _NullCtx() -def step_profiler(): +def _configure_markers() -> None: + raw = os.environ.get("FREETOKEN_ROCTX_MARKERS", "").strip() + _state.markers_enabled = bool(raw and raw.lower() not in {"0", "false", "off", "no"}) + if _state.markers_enabled and raw.lower() not in {"1", "true", "yes", "on"}: + _state.marker_output = raw + if not _state.markers_enabled: + return + if _state.marker_output: + # Launch-mode rocprof has no torch-profiler exit hook. Persist sidecar + # ranges at process exit so GPU trace and CPU phase ledger share a key. + atexit.register(_export_marker_sidecar) + library = ctypes.util.find_library("roctx64") or "libroctx64.so" + try: + _state.roctx = ctypes.CDLL(library) + _state.roctx.roctxRangePushA.argtypes = [ctypes.c_char_p] + _state.roctx.roctxRangePushA.restype = ctypes.c_int + _state.roctx.roctxRangePop.argtypes = [] + _state.roctx.roctxRangePop.restype = ctypes.c_int + except (AttributeError, OSError) as exc: + _state.roctx = None + _state.roctx_error = f"{type(exc).__name__}: {exc}" + + +def _current_stream_id() -> str | None: + try: + import torch + + stream = torch.cuda.current_stream() + value = getattr(stream, "cuda_stream", None) + return str(value) if value is not None else None + except Exception: + return None + + +def _mark(step: int, phase: str, event: str, stream_id: str | None) -> None: + if not _state.markers_enabled: + return + monotonic_ns = time.monotonic_ns() + marker = { + "step": step, + "phase": phase, + "event": event, + "monotonic_ns": monotonic_ns, + "stream_id": stream_id, + } + if _state.marker_output: + _state.markers.append(marker) + if _state.roctx is not None: + label = f"freetoken step={step} phase={phase}" + try: + if event == "begin": + _state.roctx.roctxRangePushA(label.encode()) + else: + _state.roctx.roctxRangePop() + except Exception as exc: # marker failure must never stop serving + _state.roctx_error = f"{type(exc).__name__}: {exc}" + + +def step_profiler(phase: str = "scheduler", stream_id: str | None = None): """Wrap one scheduler step. No-op unless FREETOKEN_TORCH_PROFILE is set (parsed on the first call); the window covers `steps` calls after the first `warm`.""" + if _state.markers_enabled is None: + _configure_markers() if _state.done: return _NULL if _state.spec is _NO_SPEC: # type: ignore[comparison-overlap] spec = _read_spec() if spec is None: - _state.done = True - return _NULL + if not _state.markers_enabled: + _state.done = True + return _NULL _state.spec = spec - warm, steps, _ = _state.spec # type: ignore[misc] _state.calls += 1 + marker_step = _state.calls + marker_stream = stream_id if stream_id is not None else _current_stream_id() + if _state.spec is None: + return _ProfilerCtx(None, final=False, step=marker_step, phase=phase, stream_id=marker_stream) + warm, steps, _ = _state.spec if _state.calls <= warm: + if _state.markers_enabled: + return _ProfilerCtx(None, final=False, step=marker_step, phase=phase, stream_id=marker_stream) return _NULL import torch.profiler - prof = torch.profiler.profile( - activities=[ - torch.profiler.ProfilerActivity.CPU, - torch.profiler.ProfilerActivity.CUDA, - ], - record_shapes=False, - with_stack=False, + if _state.prof is None: + _state.prof = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=False, + with_stack=False, + acc_events=True, + ) + _state.prof.__enter__() + return _ProfilerCtx( + _state.prof, + final=_state.calls >= warm + steps, + step=marker_step, + phase=phase, + stream_id=marker_stream, ) - return _ProfilerCtx(prof, final=_state.calls >= warm + steps) + + +def profiler_phase(phase: str, stream_id: str | None = None): + """Mark graph/sampler subranges without making missing profiler libs fatal.""" + if _state.markers_enabled is None: + _configure_markers() + if not _state.markers_enabled and _state.spec is _NO_SPEC: + # Keep record_function available for an explicitly active torch profiler, but avoid + # importing torch on the normal serving path when no instrumentation is configured. + if not os.environ.get("FREETOKEN_TORCH_PROFILE"): + return _NULL + marker_step = _state.calls or 0 + marker_stream = stream_id if stream_id is not None else _current_stream_id() + return _PhaseCtx(phase, marker_step, marker_stream) def _export(prof, out: str, steps: int) -> None: # pragma: no cover - heavy @@ -131,4 +285,119 @@ def _export(prof, out: str, steps: int) -> None: # pragma: no cover - heavy with open(table_path, "w") as f: f.write(f"torch.profiler key_averages over {steps} profiled step(s)\n") f.write(table) - f.write("\n") \ No newline at end of file + f.write("\n") + _export_marker_sidecar(out) + + +def _marker_path(out: str) -> str: + return out.rsplit(".", 1)[0] + "-markers.json" + + +def _export_marker_sidecar(out: str | None = None) -> None: + path = out or _state.marker_output + if not path or not _state.markers: + return + marker_path = _marker_path(path) + os.makedirs(os.path.dirname(marker_path) or ".", exist_ok=True) + with open(marker_path, "w") as f: + json.dump( + { + "schema": "freetoken-roctx-markers-v1", + "roctx_available": _state.roctx is not None, + "roctx_error": _state.roctx_error, + "markers": _state.markers, + }, + f, + sort_keys=True, + indent=2, + ) + + +def summarize_markers(markers: list[dict]) -> dict: + """Summarize paired ROCTX/sidecar ranges without summing overlapping stages. + + Durations are host monotonic-clock intervals. ``critical_step`` reports the + outer scheduler envelope; phase totals remain attribution data and are never + presented as additive wall time. + """ + open_ranges: dict[tuple[object, object, object], list[int]] = {} + phase_durations: dict[str, list[int]] = {} + step_durations: list[int] = [] + errors: list[str] = [] + for marker in markers: + if not isinstance(marker, dict): + errors.append("marker is not an object") + continue + try: + key = (marker["step"], marker["phase"], marker.get("stream_id")) + stamp = int(marker["monotonic_ns"]) + event = marker["event"] + phase = str(marker["phase"]) + except (KeyError, TypeError, ValueError): + errors.append("malformed marker") + continue + if event == "begin": + open_ranges.setdefault(key, []).append(stamp) + elif event == "end": + starts = open_ranges.get(key) + if not starts: + errors.append(f"unmatched end for step={key[0]} phase={phase}") + continue + duration = stamp - starts.pop() + if duration < 0: + errors.append(f"negative duration for step={key[0]} phase={phase}") + continue + phase_durations.setdefault(phase, []).append(duration) + if phase == "scheduler": + step_durations.append(duration) + else: + errors.append(f"unknown marker event {event!r}") + for step, phase, _stream in open_ranges: + if open_ranges[(step, phase, _stream)]: + errors.append(f"unmatched begin for step={step} phase={phase}") + + def stats(values: list[int]) -> dict[str, int | float]: + ordered = sorted(values) + return { + "count": len(ordered), + "total_ns": sum(ordered), + "median_ns": statistics.median(ordered), + "p95_ns": ordered[min(len(ordered) - 1, int(len(ordered) * 0.95))], + } + + return { + "schema": "freetoken-stage-summary-v1", + "complete": not errors, + "errors": errors, + "phases": {phase: stats(values) for phase, values in sorted(phase_durations.items())}, + "critical_step": stats(step_durations) if step_durations else None, + "note": "phase totals overlap; critical_step is the wall-clock envelope", + } + + +def summarize_marker_file(path: str) -> dict: + """Load a profiler marker sidecar and return a machine-readable stage summary.""" + with open(path) as marker_file: + payload = json.load(marker_file) + markers = payload.get("markers") if isinstance(payload, dict) else None + summary = summarize_markers(markers if isinstance(markers, list) else []) + summary["source"] = path + if isinstance(payload, dict): + summary["roctx_available"] = payload.get("roctx_available") + summary["roctx_error"] = payload.get("roctx_error") + return summary + + +if __name__ == "__main__": # pragma: no cover - command-line artifact helper + import argparse + + parser = argparse.ArgumentParser(description="Summarize FreeToken profiler markers") + parser.add_argument("markers") + parser.add_argument("--out") + args = parser.parse_args() + summary = json.dumps(summarize_marker_file(args.markers), indent=2, sort_keys=True) + if args.out: + with open(args.out, "w") as summary_file: + summary_file.write(summary + "\n") + else: + print(summary) diff --git a/scripts/build-llama-reference.sh b/scripts/build-llama-reference.sh new file mode 100755 index 000000000..c2fcecad1 --- /dev/null +++ b/scripts/build-llama-reference.sh @@ -0,0 +1,94 @@ +#!/usr/bin/env bash +# Build pinned llama.cpp reference binaries outside FreeToken. +# +# Required source state: llama.cpp commit 7e4c0a968 (b10434). The script never +# changes source HEAD; it fails before configure when source identity differs. +# +# Env: +# LLAMA_CPP_SRC local llama.cpp checkout (required) +# LLAMA_CPP_COMMIT exact commit (default 7e4c0a968) +# LLAMA_BUILD_ROOT output root (default /tmp/llama-cpp-b10434) +# LLAMA_BACKENDS comma list: hip,vulkan (default hip,vulkan) + +set -euo pipefail + +EXPECTED_COMMIT="${LLAMA_CPP_COMMIT:-7e4c0a968}" +LLAMA_CPP_SRC="${LLAMA_CPP_SRC:-}" +LLAMA_BUILD_ROOT="${LLAMA_BUILD_ROOT:-/tmp/llama-cpp-b10434}" +LLAMA_BACKENDS="${LLAMA_BACKENDS:-hip,vulkan}" +MANIFEST="$LLAMA_BUILD_ROOT/provenance.txt" + +die() { echo "ERROR: $*" >&2; exit 1; } + +[ -n "$LLAMA_CPP_SRC" ] || die "set LLAMA_CPP_SRC to a local llama.cpp checkout" +[ -d "$LLAMA_CPP_SRC/.git" ] || die "LLAMA_CPP_SRC is not a git checkout: $LLAMA_CPP_SRC" +command -v cmake >/dev/null 2>&1 || die "cmake not found" + +HEAD="$(git -C "$LLAMA_CPP_SRC" rev-parse HEAD 2>/dev/null)" || die "cannot read llama.cpp HEAD" +case "$HEAD" in + "$EXPECTED_COMMIT"|"$EXPECTED_COMMIT"*) ;; + *) die "llama.cpp HEAD=$HEAD does not match required $EXPECTED_COMMIT" ;; +esac + +case ",${LLAMA_BACKENDS}," in + *,hip,*) + command -v hipcc >/dev/null 2>&1 || die "hipcc not found for HIP reference" + ;; +esac +case ",${LLAMA_BACKENDS}," in + *,hip,*|*,vulkan,*) ;; + *) die "LLAMA_BACKENDS must contain hip and/or vulkan" ;; +esac + +mkdir -p "$LLAMA_BUILD_ROOT" +{ + echo "source=$LLAMA_CPP_SRC" + echo "source_head=$HEAD" + echo "required_commit=$EXPECTED_COMMIT" + echo "backends=$LLAMA_BACKENDS" + echo "cmake=$(cmake --version | head -1)" + echo "hipcc=$(command -v hipcc 2>/dev/null || echo unavailable)" + echo "hipcc_version=$(hipcc --version 2>/dev/null | head -1 || echo unavailable)" + echo "git_status=$(git -C "$LLAMA_CPP_SRC" status --short)" +} > "$MANIFEST" + +build_backend() { + local backend="$1" + local build_dir="$LLAMA_BUILD_ROOT/$backend" + local -a cmake_args=( + -S "$LLAMA_CPP_SRC" + -B "$build_dir" + -DCMAKE_BUILD_TYPE=Release + ) + case "$backend" in + hip) + cmake_args+=( + -DGGML_HIP=ON + -DCMAKE_HIP_ARCHITECTURES=gfx1100 + ) + ;; + vulkan) + cmake_args+=( -DGGML_VULKAN=ON ) + ;; + *) die "unsupported backend: $backend" ;; + esac + printf 'configure:' >> "$MANIFEST" + printf ' %q' cmake "${cmake_args[@]}" + printf '\n' >> "$MANIFEST" + cmake "${cmake_args[@]}" + cmake --build "$build_dir" --target llama-server -- -j"${CMAKE_BUILD_PARALLEL_LEVEL:-2}" + local binary="$build_dir/bin/llama-server" + [ -x "$binary" ] || binary="$build_dir/llama-server" + [ -x "$binary" ] || die "llama-server missing after $backend build: $build_dir" + printf 'binary_%s=%s\n' "$backend" "$binary" >> "$MANIFEST" + printf 'version_%s=%s\n' "$backend" "$("$binary" --version 2>&1 | head -1)" >> "$MANIFEST" + echo "built $backend: $binary" +} + +IFS=',' read -r -a backends <<< "$LLAMA_BACKENDS" +for backend in "${backends[@]}"; do + [ -n "$backend" ] || continue + build_backend "$backend" +done + +echo "manifest: $MANIFEST" diff --git a/scripts/profile-rocm-decode.sh b/scripts/profile-rocm-decode.sh index dd3262afe..ff19ae5b5 100755 --- a/scripts/profile-rocm-decode.sh +++ b/scripts/profile-rocm-decode.sh @@ -1,30 +1,64 @@ #!/usr/bin/env bash # profile-rocm-decode.sh # -# Inc 2 of .plans/rocm-perf-parity: stage-time breakdown of decode on gfx1100 for +# Inc 2 of .plans/qwen-moe-speed: stage-time breakdown of decode on gfx1100 for # the Qwen3.6-35B-A3B GGUF (matching the P217 baseline conditions). # -# Runs ft serve with FREETOKEN_TORCH_PROFILE=:: (the env reaches -# the server because serve-qwen-moe.sh execs $PY from this shell), sends one -# AIME-style request, waits for the exported trace, then stops the server. Env: +# Runs ft serve with torch-profiler or low-overhead rocprofv3 marker mode, sends the +# same pinned AIME fixture/request shape as bench_decode_moe.py, then stops server. +# Env: +# PROFILE_MODE torch|rocprofv3 (default torch) +# PROFILE_CAPTURE_START launch|attach for rocprofv3 (default launch) # PROFILE_OUT chrome-trace base path (default /tmp/ft-rocm-profile/chrome.json) # PROFILE_WARM decode steps skipped before the profiled window (default 40) # PROFILE_STEPS steps inside the profiled window (default 40) +# PROFILE_DECODE exact max_tokens for the request (default 512) +# AIME_JSONL local immutable AIME JSONL; required +# AIME_PROBLEM 0-based fixture row (default 0) # FT_MODEL model path (default the 7900 XTX box's Qwen3.6 GGUF) # FT_PORT server port (default 1920) +# ROCPROF_BIN rocprofv3 path (default rocprofv3) +# ROCPROF_TARGET_PID attach target override; default frontend PID +# PROFILE_ATTACH_MS rocprofv3 attach duration (default 30000) set -euo pipefail REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" PY="${PY:-$REPO/.venv-rocm/bin/python}" +PROFILE_MODE="${PROFILE_MODE:-torch}" +PROFILE_CAPTURE_START="${PROFILE_CAPTURE_START:-launch}" PROFILE_OUT="${PROFILE_OUT:-/tmp/ft-rocm-profile/chrome.json}" PROFILE_WARM="${PROFILE_WARM:-40}" PROFILE_STEPS="${PROFILE_STEPS:-40}" +PROFILE_DECODE="${PROFILE_DECODE:-512}" FT_MODEL="${FT_MODEL:-/home/smk/models/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf}" FT_PORT="${FT_PORT:-1920}" +AIME_JSONL="${AIME_JSONL:-${FREETOKEN_AIME25_JSONL:-}}" +AIME_PROBLEM="${AIME_PROBLEM:-0}" +FT_ATTN="${FT_ATTN:-triton}" +FT_MOE_CACHE="${FT_MOE_CACHE:-auto}" +FT_MEMORY_RATIO="${FT_MEMORY_RATIO:-0.9}" +FT_KV_TOKENS="${FT_KV_TOKENS:-9216}" +FT_MAX_OUTPUT="${FT_MAX_OUTPUT:-$PROFILE_DECODE}" +ROCPROF_BIN="${ROCPROF_BIN:-rocprofv3}" +PROFILE_ATTACH_MS="${PROFILE_ATTACH_MS:-30000}" PORT="$FT_PORT" -export FT_MODEL FT_PORT +ROCPROF_LOG="${ROCPROF_LOG:-${PROFILE_OUT%.json}-rocprofv3.log}" +export FT_MODEL FT_PORT AIME_JSONL AIME_PROBLEM FT_ATTN FT_MOE_CACHE FT_MEMORY_RATIO FT_KV_TOKENS FT_MAX_OUTPUT PROFILE_DECODE PROFILE_MODE + +case "$PROFILE_MODE" in + torch|rocprofv3) ;; + *) echo "ERROR: PROFILE_MODE must be torch or rocprofv3" >&2; exit 1 ;; +esac +case "$PROFILE_CAPTURE_START" in + launch|attach) ;; + *) echo "ERROR: PROFILE_CAPTURE_START must be launch or attach" >&2; exit 1 ;; +esac +if [ "$PROFILE_MODE" = torch ] && [ "$PROFILE_CAPTURE_START" != launch ]; then + echo "ERROR: PROFILE_CAPTURE_START only applies to PROFILE_MODE=rocprofv3" >&2 + exit 1 +fi if grep -q $'\r' "${BASH_SOURCE[0]}"; then echo "ERROR: CRLF line endings in profile-rocm-decode.sh" >&2 @@ -32,41 +66,101 @@ if grep -q $'\r' "${BASH_SOURCE[0]}"; then fi [ -x "$PY" ] || { echo "ERROR: python not found: $PY" >&2; exit 1; } [ -f "$FT_MODEL" ] || { echo "ERROR: model not found: $FT_MODEL" >&2; exit 1; } +[ -n "$AIME_JSONL" ] && [ -f "$AIME_JSONL" ] || { + echo "ERROR: AIME_JSONL must point to an immutable local JSONL fixture" >&2 + exit 1 +} OUT_DIR="$(dirname "$PROFILE_OUT")" mkdir -p "$OUT_DIR" rm -f "$PROFILE_OUT" "${PROFILE_OUT%.json}-kernels.log" +TRACE_DIR="${PROFILE_OUT%.json}-rocprofv3" +ENV_MANIFEST="${PROFILE_OUT%.json}-env.txt" +MARKER_SIDECAR="${PROFILE_OUT%.json}-markers.json" +rm -rf "$TRACE_DIR" +{ + echo "profile_mode=$PROFILE_MODE" + echo "profile_capture_start=$PROFILE_CAPTURE_START" + echo "profile_out=$PROFILE_OUT" + echo "profile_warm=$PROFILE_WARM" + echo "profile_steps=$PROFILE_STEPS" + echo "profile_decode=$PROFILE_DECODE" + echo "model=$FT_MODEL" + echo "model_sha256=$(sha256sum "$FT_MODEL" | awk '{print $1}')" + echo "aime=$AIME_JSONL" + echo "aime_sha256=$(sha256sum "$AIME_JSONL" | awk '{print $1}')" + echo "git_revision=$(git -C "$REPO" rev-parse HEAD 2>/dev/null || echo unknown)" + echo "python=$PY" + echo "rocprofv3=$ROCPROF_BIN" + echo "rocprofv3_version=$($ROCPROF_BIN --version 2>&1 | head -1 || echo unavailable)" + echo "hipcc_version=$(hipcc --version 2>/dev/null | head -1 || echo unavailable)" + env | sort +} > "$ENV_MANIFEST" +echo "Profile mode: $PROFILE_MODE" echo "Profile env: FREETOKEN_TORCH_PROFILE=${PROFILE_WARM}:${PROFILE_STEPS}:${PROFILE_OUT}" # serve-qwen-moe.sh refuses a double start; clear any previous instance first. ./scripts/serve-qwen-moe.sh stop >/dev/null 2>&1 || true # serve-qwen-moe.sh refuses a double start and returns once ready (~3-4 min load). -FREETOKEN_TORCH_PROFILE="${PROFILE_WARM}:${PROFILE_STEPS}:${PROFILE_OUT}" \ - ./scripts/serve-qwen-moe.sh +if [ "$PROFILE_MODE" = torch ]; then + FREETOKEN_TORCH_PROFILE="${PROFILE_WARM}:${PROFILE_STEPS}:${PROFILE_OUT}" \ + FREETOKEN_ROCTX_MARKERS="$PROFILE_OUT" ./scripts/serve-qwen-moe.sh +else + command -v "$ROCPROF_BIN" >/dev/null 2>&1 || { + echo "ERROR: rocprofv3 not found: $ROCPROF_BIN" >&2 + exit 1 + } + mkdir -p "$TRACE_DIR" + if [ "$PROFILE_CAPTURE_START" = launch ]; then + FREETOKEN_ROCTX_MARKERS="$PROFILE_OUT" \ + FT_ROCPROF_BIN="$ROCPROF_BIN" \ + FT_ROCPROF_TRACE_DIR="$TRACE_DIR" \ + FT_ROCPROF_LOG="$ROCPROF_LOG" \ + ./scripts/serve-qwen-moe.sh + ROCPROF_PID="$(pgrep -f 'rocprofv3.*freetoken-rocprof' | head -1 || true)" + echo "rocprofv3 launch profiler_pid=${ROCPROF_PID:-unknown}" + else + FREETOKEN_ROCTX_MARKERS="$PROFILE_OUT" ./scripts/serve-qwen-moe.sh + TARGET_PID="${ROCPROF_TARGET_PID:-$(pgrep -f 'freetoke[n].cli serve' | head -1 || true)}" + [ -n "$TARGET_PID" ] || { + echo "ERROR: could not find FreeToken server PID for rocprofv3 attach" >&2 + exit 1 + } + "$ROCPROF_BIN" --runtime-trace --marker-trace --kernel-trace \ + --memory-copy-trace --memory-allocation-trace -d "$TRACE_DIR" \ + --attach "$TARGET_PID" --attach-duration-msec "$PROFILE_ATTACH_MS" \ + > "$ROCPROF_LOG" 2>&1 & + ROCPROF_PID=$! + echo "rocprofv3 attach pid=$TARGET_PID profiler_pid=$ROCPROF_PID" + fi +fi -# Send the profiled request (max_tokens > PROFILE_STEPS, so the engine stays -# decoding through the whole profiled window; the trace exports at its last step). +# Send exact benchmark request shape (max_tokens > PROFILE_STEPS, so the engine +# stays decoding through the whole profiled window; trace exports at its last step). FREETOKEN_TORCH_PROFILE_MAX="${PROFILE_STEPS}" \ -"${PY}" - <<'PYEOF' -import json, os, sys, urllib.request +PYTHONPATH="$REPO/python:$REPO/benchmarks" "${PY}" - <<'PYEOF' +import json, os, urllib.request + +from bench_decode_moe import load_problem_details, resolve_sampling port = os.environ["FT_PORT"] -max_tokens = ( - int(os.environ["PROFILE_WARM"]) + int(os.environ["PROFILE_STEPS"]) + 80 -) -prompt = ( - "Every morning Aya goes for a 9 kilometer walk, stops at a coffee shop, then " - "walks back home. She walks at 4 km/h, and the coffee shop detour adds 15 " - "minutes. On a day when she walks at t km/h and the detour still costs her 2 " - "hours total, what is t? Answer with just the number." +max_tokens = int(os.environ["PROFILE_DECODE"]) + 1 +prompt, _, _ = load_problem_details( + os.environ["AIME_JSONL"], + int(os.environ["AIME_PROBLEM"]), ) +sampling, _ = resolve_sampling(os.environ["FT_MODEL"], greedy=False) body = json.dumps({ "model": os.path.basename(os.environ["FT_MODEL"]), "messages": [{"role": "user", "content": prompt}], "max_tokens": max_tokens, - "stream": False, + "ignore_eos": True, + "stream_options": {"include_usage": True}, + "chat_template_kwargs": {"enable_thinking": True}, + "stream": True, + **sampling, }).encode() req = urllib.request.Request( f"http://127.0.0.1:{port}/v1/chat/completions", @@ -74,7 +168,17 @@ req = urllib.request.Request( headers={"Content-Type": "application/json"}, ) with urllib.request.urlopen(req, timeout=1200) as r: - data = json.loads(r.read()) + data = {} + for raw in r: + line = raw.strip() + if not line.startswith(b"data:"): + continue + payload = line[len(b"data:"):].strip() + if payload == b"[DONE]": + break + chunk = json.loads(payload) + if chunk.get("usage"): + data["usage"] = chunk["usage"] print("request completed; completion_tokens:", data.get("usage", {}).get("completion_tokens")) PYEOF @@ -87,11 +191,35 @@ done ./scripts/serve-qwen-moe.sh stop || true +if [ "$PROFILE_MODE" = rocprofv3 ] && [ -n "${ROCPROF_PID:-}" ]; then + wait "$ROCPROF_PID" || true +fi + if [ -f "$PROFILE_OUT" ]; then echo "Trace: $PROFILE_OUT" echo "Kernels: ${PROFILE_OUT%.json}-kernels.log" + if [ -f "$MARKER_SIDECAR" ]; then + SUMMARY_OUT="${PROFILE_OUT%.json}-summary.json" + PYTHONPATH="$REPO/python" "$PY" -m freetoken.utils.step_profiler \ + "$MARKER_SIDECAR" --out "$SUMMARY_OUT" + echo "Stages: $SUMMARY_OUT" + fi else - echo "ERROR: trace not exported to $PROFILE_OUT; check /tmp/serve_qwen_moe.log" >&2 - tail -20 /tmp/serve_qwen_moe.log >&2 || true - exit 1 -fi \ No newline at end of file + if [ "$PROFILE_MODE" = rocprofv3 ] && find "$TRACE_DIR" -type f -print -quit 2>/dev/null | grep -q .; then + echo "Trace: $TRACE_DIR" + echo "Env: $ENV_MANIFEST" + echo "Markers: rocprofv3 runtime/marker/kernel trace" + if [ -f "$MARKER_SIDECAR" ]; then + SUMMARY_OUT="${PROFILE_OUT%.json}-summary.json" + PYTHONPATH="$REPO/python" "$PY" -m freetoken.utils.step_profiler \ + "$MARKER_SIDECAR" --out "$SUMMARY_OUT" + echo "Stages: $SUMMARY_OUT" + else + echo "Stages: unavailable (ROCTX marker sidecar missing)" + fi + else + echo "ERROR: trace not exported to $PROFILE_OUT; check /tmp/serve_qwen_moe.log" >&2 + tail -20 /tmp/serve_qwen_moe.log >&2 || true + exit 1 + fi +fi diff --git a/scripts/serve-qwen-moe.sh b/scripts/serve-qwen-moe.sh index b59df731a..ea357da5b 100755 --- a/scripts/serve-qwen-moe.sh +++ b/scripts/serve-qwen-moe.sh @@ -6,7 +6,8 @@ # CPU/offload cache) and the triton attention backend by default. # # Native context window: 262144 (256K) tokens (max_position_embeddings in the model). -# Graph capture is settled as a failure on ROCm, so this runs eager kernel-launch decode. +# Decode graph capture is probed at startup. Current gfx1100 passes the gate and uses +# replay; unsupported/failed capture falls back to eager kernel-launch decode. # # VS CODE NOTES: # - The server command is built as a bash array and launched on ONE physical @@ -29,6 +30,28 @@ set -euo pipefail REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" + +# --- .env loader: repo-local runtime defaults (FT_MODEL, FT_PORT, FT_KV_TOKENS, +# FT_SERVED_MODEL, ...) ------------------------------------------------------------ +# Contract (see .env header): values there are DEFAULTS — a variable already set in +# the shell wins, so one-off `FT_X=... ./serve-qwen-moe.sh` overrides still hold. +# Runs every launch; editing .env affects the next launch only. +if [ -f "$REPO/.env" ]; then + while IFS= read -r line; do + line="${line%%#*}" # strip comments + line="${line#"${line%%[![:space:]]*}"}" # ltrim + if [ -z "$line" ]; then continue; fi + key="${line%%=*}" + case "$key" in *[![:space:][:alnum:]_]*) continue ;; esac + val="${line#*=}" + # trim leading spaces from the value + val="${val#"${val%%[![:space:]]*}"}" + if [ -n "$key" ] && [ -z "${!key+x}" ]; then + export "${key}=${val}" + fi + done < "$REPO/.env" +fi + PY="${PY:-$REPO/.venv-rocm/bin/python}" MODEL="${FT_MODEL:-}" @@ -72,12 +95,16 @@ KV_TOKENS="${FT_KV_TOKENS:-131072}" # other clients inherit this. MAX_OUTPUT="${FT_MAX_OUTPUT:-65536}" MOE_CACHE="${FT_MOE_CACHE:-auto}" +KV_TYPE="${FT_KV_TYPE:-bf16}" # bf16 | fp16 | q8_0; q8_0 is opt-in and full-MHA only LOG="${FT_LOG:-/tmp/serve_qwen_moe.log}" +ROCPROF_BIN="${FT_ROCPROF_BIN:-}" +ROCPROF_TRACE_DIR="${FT_ROCPROF_TRACE_DIR:-}" +ROCPROF_LOG="${FT_ROCPROF_LOG:-/tmp/freetoken-rocprofv3.log}" # Advertised model id in /v1/models (--served-model-name). Copilot-fork custom # endpoints validate the configured model id against this list, so set it to the # exact id the client config declares (e.g. FT_SERVED_MODEL=qwen3.6). Empty = # server default (the GGUF filename). -SERVED_MODEL="${FT_SERVED_MODEL:-}" +SERVED_MODEL="${FT_SERVED_MODEL:-qwen3.6}" die() { echo "ERROR: $*" >&2; exit 1; } @@ -117,6 +144,7 @@ start() { $( [ -n "$CPU_LAYERS" ] && echo "--moe-cpu-layers" "$CPU_LAYERS" ) "--attention-backend" "$ATNN" "--num-tokens" "$KV_TOKENS" + "--kv-type" "$KV_TYPE" "--memory-ratio" "$MEMORY_RATIO" "--max-prefill-length" "$PREFILL_CHUNK" "--max-output-tokens" "$MAX_OUTPUT" @@ -141,15 +169,33 @@ start() { echo " attn : $ATNN" echo " moe : $MOE_BACKEND" echo " kv : $KV_TOKENS tokens (gpu moe cache: ${MOE_CACHE}${MOE_CACHE:+ }$( [ "$MOE_CACHE" = auto ] && echo "kv-reserve $KV_TOKENS" || echo slots))" + echo " kv type : $KV_TYPE" echo " python : $PY" echo " log : $LOG" + local -a EXEC_PREFIX=() + if [ -n "$ROCPROF_BIN" ]; then + [ -d "$ROCPROF_TRACE_DIR" ] || die "rocprof trace directory not found: $ROCPROF_TRACE_DIR" + [ -x "$(command -v "$ROCPROF_BIN" 2>/dev/null || true)" ] || die "rocprofiler not found: $ROCPROF_BIN" + EXEC_PREFIX=( + "$ROCPROF_BIN" + --runtime-trace + --marker-trace + --kernel-trace + --memory-copy-trace + --memory-allocation-trace + -d "$ROCPROF_TRACE_DIR" + -- + ) + echo " profiler: ${EXEC_PREFIX[*]}" + fi + # setsid: give the server its OWN session/process group. nohup alone only # ignores SIGHUP — a caller that dies (e.g. a VS Code task cancelled, an agent # tool timeout) still takes down the whole process group with SIGKILL/SIGTERM, # which silently killed the server mid model-load once already. cd "$REPO" - PYTHONPATH="$REPO/python" setsid nohup "$PY" -m freetoken.cli serve "${SERVE_ARGS[@]}" >"$LOG" 2>&1 "$LOG" 2>&1 /dev/null || true echo "pid=$pid — waiting for readiness (model load takes ~3-4 min)..." diff --git a/tests/benchmarks/__init__.py b/tests/benchmarks/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/benchmarks/test_paired_stats.py b/tests/benchmarks/test_paired_stats.py new file mode 100644 index 000000000..4af459c37 --- /dev/null +++ b/tests/benchmarks/test_paired_stats.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +BENCHMARKS = Path(__file__).resolve().parents[2] / "benchmarks" +if str(BENCHMARKS) in sys.path: + sys.path.remove(str(BENCHMARKS)) +sys.path.insert(0, str(BENCHMARKS)) + +from lib.paired_stats import bootstrap_median_interval, paired_summary + + +def test_paired_summary_reports_positive_candidate_recovery(): + result = paired_summary([10, 12, 11], [8, 9, 10]) + assert result["pairs"] == 3 + assert result["median_recovery_us"] == 2 + assert result["recovery_p02_5_us"] <= 2 <= result["recovery_p97_5_us"] + assert result["candidate_speedup_pct"] > 0 + + +def test_paired_summary_rejects_unpaired_inputs(): + with pytest.raises(ValueError, match="equal length"): + paired_summary([1], [1, 2]) + + +def test_empty_bootstrap_is_explicit(): + assert bootstrap_median_interval([]) == (None, None) diff --git a/tests/benchmarks/test_replay_manifest.py b/tests/benchmarks/test_replay_manifest.py new file mode 100644 index 000000000..c2e4bfbfe --- /dev/null +++ b/tests/benchmarks/test_replay_manifest.py @@ -0,0 +1,267 @@ +"""Inc 0 harness tests: replay manifest, gate lanes, and identity invalidation. + +Covers the pure benchmark-harness logic (no GPU): the replay-manifest schema +contract, route hashing, warmup-aware step summaries, disjoint lane +classification in the promotion gate, and candidate-fallback invalidation. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +BENCHMARKS = Path(__file__).resolve().parents[2] / "benchmarks" +if str(BENCHMARKS) not in sys.path: + sys.path.insert(0, str(BENCHMARKS)) + +import bench_decode_replay as replay # noqa: E402 +import check_decode_gate as gate # noqa: E402 + + +# --------------------------------------------------------------------------- +# Replay manifest schema +# --------------------------------------------------------------------------- + + +def valid_manifest(**overrides) -> dict: + manifest = { + "schema": "freetoken-replay-manifest-v1", + "lane": "teacher_forced_replay", + "prompt_ids": [11, 12, 13], + "continuation_ids": [70, 80, 90, 100], + "warmup_tokens": 1, + "measured_tokens": 4, + "model_sha256": "0" * 64, + "fixture_sha256": "2" * 64, + "tokenizer_sha256": "3" * 64, + "route_top_k": 8, + "golden": {"ids_sha256": "a" * 64}, + } + return {**manifest, **overrides} + + +def test_manifest_accepts_valid_contract(): + assert replay.validate_manifest(valid_manifest()) == [] + + +def test_manifest_rejects_schema_lane_and_identity_drift(): + problems = replay.validate_manifest({**valid_manifest(), "schema": "other-schema"}) + assert any("schema" in problem for problem in problems) + + problems = replay.validate_manifest( + {**valid_manifest(), "lane": "sampled_absolute"} + ) + assert any("lane" in problem for problem in problems) + + problems = replay.validate_manifest({**valid_manifest(), "measured_tokens": 5}) + assert any("measured_tokens" in problem for problem in problems) + + problems = replay.validate_manifest({**valid_manifest(), "model_sha256": "short"}) + assert any("model_sha256" in problem for problem in problems) + + problems = replay.validate_manifest({**valid_manifest(), "route_top_k": None}) + assert any("route_top_k" in problem for problem in problems) + + problems = replay.validate_manifest({**valid_manifest(), "golden": {"ids_sha256": None}}) + assert any("golden" in problem for problem in problems) + + problems = replay.validate_manifest({**valid_manifest(), "continuation_ids": [1]}) + assert any("continuation_ids" in problem for problem in problems) + + problems = replay.validate_manifest({**valid_manifest(), "prompt_ids": [1, -1]}) + assert any("prompt_ids" in problem for problem in problems) + + problems = replay.validate_manifest({**valid_manifest(), "model_sha256": "z" * 64}) + assert any("model_sha256" in problem for problem in problems) + + problems = replay.validate_manifest( + {**valid_manifest(), "prompt_text": "prompt", "prompt_text_sha256": "0" * 64} + ) + assert any("prompt_text_sha256" in problem for problem in problems) + + +def test_load_manifest_rejects_invalid_file(tmp_path): + path = tmp_path / "manifest.json" + path.write_text(json.dumps({**valid_manifest(), "schema": "wrong"})) + with pytest.raises(ValueError, match="rejected"): + replay.load_manifest(str(path)) + + good = valid_manifest() + path.write_text(json.dumps(good)) + assert replay.load_manifest(str(path)) == good + + +def test_route_hash_is_stable_and_disjoint(): + assert replay.route_hash([1, 2, 3]) == replay.route_hash([1, 2, 3]) + assert replay.route_hash([1]) != replay.route_hash([2]) + assert replay.route_hash(list(range(8))) != replay.route_hash(list(range(7, -1, -1))) + + +def test_route_digest_preserves_token_and_layer_order(): + one = replay._route_digest( + [ + {"layer": 2, "start": 0, "hashes": ["b", "c"]}, + {"layer": 1, "start": 0, "hashes": ["a", "d"]}, + ] + ) + two = replay._route_digest( + [ + {"layer": 1, "start": 0, "hashes": ["a", "d"]}, + {"layer": 2, "start": 0, "hashes": ["b", "c"]}, + ] + ) + assert one == two + assert one["0"] != one["1"] + + +def test_prompt_text_uses_manifest_and_checks_hash(): + prompt = "hello" + manifest = {"prompt_text": prompt, "prompt_text_sha256": replay._sha256_bytes(prompt.encode())} + assert replay.prompt_text_of(manifest) == prompt + with pytest.raises(ValueError, match="prompt_text"): + replay.prompt_text_of({"prompt_text": prompt, "prompt_text_sha256": "0" * 64}) + + +def test_tokens_list_accepts_llama_tokenize_response(): + assert replay.tokens_list({"tokens": [1, 2, 3]}) == [1, 2, 3] + assert replay.tokens_list({"tokens": "bad"}) is None + + +def test_summarize_steps_excludes_warmup_keeps_raw(): + steps = [1.0, 2.0, 3.0, 100.0] + stats = replay.summarize_steps(steps, warmup_steps=1) + assert stats["steps"] == 3 + assert stats["warmup_steps"] == 1 + assert stats["raw_steps_ms"] == steps + assert stats["ms_per_token_median"] == 3.0 + assert stats["ms_per_token_min"] == 2.0 + empty = replay.summarize_steps([], 0) + assert empty["steps"] == 0 and empty["ms_per_token_median"] is None + + +# --------------------------------------------------------------------------- +# Gate lane classification +# --------------------------------------------------------------------------- + + +def _sampled_row(**overrides) -> dict: + row = { + "schema": "freetoken-base-decode-v2", + "status": "accepted", + "lane": "sampled_absolute", + "sampling": {"temperature": 1.0}, + "decode_tok_s": 60.0, + "metadata": {"lane": "sampled_absolute", "mtp": "off"}, + } + return {**row, **overrides} + + +def _greedy_row() -> dict: + return { + "schema": "freetoken-base-decode-v2", + "status": "accepted", + "lane": "greedy_correctness", + "sampling": {"temperature": 0.0}, + "metadata": {"lane": "greedy_correctness", "mtp": "off"}, + } + + +def _replay_row() -> dict: + return { + "schema": "freetoken-replay-v1", + "status": "accepted", + "lane": "teacher_forced_replay", + } + + +def test_lane_classification_legacy_and_explicit(): + assert gate._row_lane(_sampled_row()) == gate.LANE_SAMPLED + assert gate._row_lane(_greedy_row()) == gate.LANE_GREEDY + assert gate._row_lane(_replay_row()) == gate.LANE_REPLAY + legacy = { + "schema": "freetoken-base-decode-v2", + "sampling": {"temperature": 1.0}, + "metadata": {"mtp": "off"}, + } + assert gate._row_lane(legacy) == gate.LANE_SAMPLED + legacy["sampling"] = {"temperature": 0.0} + assert gate._row_lane(legacy) == gate.LANE_GREEDY + assert gate._row_lane({"schema": "unrelated-v9"}) is None + + +def test_lane_mixing_is_rejected(): + assert gate._lane_reasons([_sampled_row(), _sampled_row()]) == [] + mixed = gate._lane_reasons([_sampled_row(), _replay_row()]) + assert any("mix disjoint lanes" in reason for reason in mixed) + unknown = gate._lane_reasons([{"schema": "unrelated-v9", "status": "accepted"}]) + assert any("unknown/mismatched lane" in reason for reason in unknown) + + +def test_fallback_evidence_rejects_gate_row(): + row = _sampled_row() + row["metadata"] = { + **row["metadata"], + "kernel_observed": {"fallback_markers": 2, "source": "server_log"}, + } + result = gate.evaluate_gate( + [row], reference_rows=[], probe=None, min_runs=1, directional_threshold=75.0 + ) + assert any("candidate-eligible fallback" in reason for reason in result["reasons"]) + + +def test_zero_fallback_keeps_row_eligible(): + row = _sampled_row() + row["metadata"] = { + **row["metadata"], + "kernel_observed": {"fallback_markers": 0, "source": "server_log"}, + } + result = gate.evaluate_gate( + [row], reference_rows=[], probe=None, min_runs=1, directional_threshold=75.0 + ) + assert not any("fallback" in reason for reason in result["reasons"]) + + +def _complete_replay(rate: float, *, repeat: int, route: dict | None = None) -> dict: + return { + "schema": "freetoken-replay-v1", + "lane": "teacher_forced_replay", + "status": "accepted", + "runtime": "freetoken", + "repeat": repeat, + "forced": True, + "ids_match": True, + "route_hash_status": "matched", + "route_digest": route or {"0": "route"}, + "model_sha256": "a" * 64, + "fixture_sha256": "b" * 64, + "tokenizer_sha256": "c" * 64, + "manifest_ids_sha256": "d" * 64, + "context": 9216, + "batch": 512, + "ubatch": 512, + "kv_type": "q8_0", + "mtp": "off", + "speculative": False, + "decode_batch_size": 1, + "decode_tok_s": rate, + } + + +def test_gate_b_is_reported_separately_for_matched_replay(): + candidate = [_complete_replay(82.0, repeat=0)] + reference = [_complete_replay(81.0, repeat=0)] + result = gate.evaluate_gate(candidate, reference, min_runs=1) + assert result["gate_b"]["gate"] is True + assert result["gates"] == {"gate_a": False, "gate_b": True} + assert result["gate"] is False # Gate A has no sampled lane. + + +def test_gate_b_rejects_route_mismatch(): + candidate = [_complete_replay(82.0, repeat=0)] + reference = [_complete_replay(81.0, repeat=0, route={"0": "other"})] + result = gate.evaluate_gate(candidate, reference, min_runs=1) + assert result["gate_b"]["gate"] is False + assert any("route hash mismatch" in reason for reason in result["gate_b"]["reasons"]) diff --git a/tests/benchmarks/test_rocm_trace.py b/tests/benchmarks/test_rocm_trace.py new file mode 100644 index 000000000..6f5ae0e10 --- /dev/null +++ b/tests/benchmarks/test_rocm_trace.py @@ -0,0 +1,78 @@ +import importlib.util +import sys +from pathlib import Path + +_TRACE_PATH = Path(__file__).parents[2] / "benchmarks" / "lib" / "rocm_trace.py" +_SPEC = importlib.util.spec_from_file_location("_freetoken_rocm_trace", _TRACE_PATH) +assert _SPEC and _SPEC.loader +_MODULE = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _MODULE +_SPEC.loader.exec_module(_MODULE) + +from _freetoken_rocm_trace import ( + calibrate_clocks, + correlate_events, + intersection, + token_ledger, + union, + warm_offload_summary, +) + + +def test_clock_calibration_rejects_large_residual(): + good = calibrate_clocks([(0, 100), (1000, 1100), (2000, 2100)]) + assert good.accepted + bad = calibrate_clocks([(0, 0), (1000, 100000), (2000, 0)]) + assert not bad.accepted + + +def test_union_and_intersection_do_not_double_count_overlap(): + assert union([(0, 10), (5, 20), (30, 40)]) == [(0, 20), (30, 40)] + assert intersection([(0, 10), (15, 30)], [(5, 20)]) == [(5, 10), (15, 20)] + + +def test_token_ledger_clips_async_events_and_separates_host_wait(): + rows = token_ledger( + [{"token": 3, "start_ns": 100, "end_ns": 300}], + [ + {"kind": "kernel", "start_ns": 50, "end_ns": 180}, + {"kind": "copy", "start_ns": 160, "end_ns": 260}, + ], + [ + {"category": "host_active", "start_ns": 120, "end_ns": 220}, + {"category": "host_wait", "start_ns": 220, "end_ns": 280}, + ], + ) + assert rows[0]["gpu_ns"] == 160 + assert rows[0]["host_active_ns"] == 100 + assert rows[0]["host_wait_ns"] == 60 + assert rows[0]["unattributed_ns"] == 20 + assert rows[0]["host_active_only_ns"] == 0 + + +def test_correlate_events_preserves_unmatched_identity(): + rows = correlate_events( + [{"correlation_id": 4, "stream_id": "s0"}], + [{"correlation_id": 4, "start_ns": 1, "end_ns": 2}], + [{"correlation_id": 9, "start_ns": 3, "end_ns": 4}], + ) + assert rows[0]["correlation_matched"] is True + assert rows[1]["correlation_matched"] is False + + +def test_warm_offload_summary_keeps_missing_bytes_explicit(): + summary = warm_offload_summary([ + {"kind": "kernel", "name": "ensure_experts", "start_ns": 0, "end_ns": 5}, + {"kind": "kernel", "name": "fast_index_copy_multi_jit", "start_ns": 5, "end_ns": 9}, + ]) + assert summary["status"] == "measured" + assert summary["all_hit"] is False + assert summary["copy_missing_bytes"] is None + + +def test_warm_offload_summary_reports_all_hit_zero_copy_count(): + summary = warm_offload_summary([ + {"kind": "kernel", "name": "ensure_experts", "start_ns": 0, "end_ns": 5}, + ]) + assert summary["all_hit"] is True + assert summary["copy_missing_count"] == 0 diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index 9ac2a4f4c..c2b4e08c5 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -218,7 +218,9 @@ class Cfg: cuda_graph_bs = [1, 2] max_seq_len = 1024 page_size = 1 - attention_backend = "fi" + # Keep fixture portable: ``fi`` is NVIDIA-only and config validation must + # reject it on the ROCm lane before this arithmetic test runs. + attention_backend = "triton" nvfp4_backend = "auto" num_page_override = None num_token_override = 5000 @@ -350,7 +352,7 @@ def _offload_engine_config(**overrides): model_path="/tmp/freetoken-test-model", tp_info=DistributedInfo(rank=0, size=1), dtype=torch.bfloat16, - attention_backend="fi", + attention_backend="triton", **overrides, ) object.__setattr__( diff --git a/tests/engine/test_gguf_resident.py b/tests/engine/test_gguf_resident.py new file mode 100644 index 000000000..3af2160e0 --- /dev/null +++ b/tests/engine/test_gguf_resident.py @@ -0,0 +1,12 @@ +from freetoken.engine.resident_budget import ResidentBudget + + +def test_observed_phase_can_raise_static_resident_requirement(): + budget = ResidentBudget( + free_bytes=100, total_vram_bytes=100, + packed_model_bytes=10, kv_bytes=10, gdn_state_bytes=0, + page_table_bytes=0, graph_reserve_bytes=0, peak_load_scratch_bytes=0, + safety_bytes=0, + phases=(), + ) + assert budget.required_bytes == 20 diff --git a/tests/engine/test_graph_blas_policy.py b/tests/engine/test_graph_blas_policy.py new file mode 100644 index 000000000..a09ec3c04 --- /dev/null +++ b/tests/engine/test_graph_blas_policy.py @@ -0,0 +1,114 @@ +"""Pure policy tests for ROCm BLAS selection and graph-gate precedence.""" + +from __future__ import annotations + +import pytest + +import freetoken.utils.graph_gate as gate + + +def test_graph_runner_detects_nested_resident_format_without_module_assumptions(): + from freetoken.engine.graph import _has_weight_format + + class Node: + pass + + root, child = Node(), Node() + root.child = child + root.loop = root + child.weight_format = "gguf" + assert _has_weight_format(root, "gguf") + assert not _has_weight_format(root, "fp8_block") + + +def test_graph_runner_runtime_telemetry_is_read_only_metadata(): + from freetoken.engine.graph import GraphRunner + + runner = object.__new__(GraphRunner) + runner.graph_telemetry = { + "expert_storage": "resident_gguf", + "expert_fetches": 0, + "expert_remaps": 0, + } + runner.resident_gguf = True + runner.graph_map = {1: object(), 4: object()} + runner.sampler_graph_map = {1: object()} + assert runner.runtime_telemetry() == { + "expert_storage": "resident_gguf", + "expert_fetches": 0, + "expert_remaps": 0, + "resident_gguf": True, + "graph_batches": [1, 4], + "sampler_graph_batches": [1], + } + + +@pytest.mark.parametrize("value", ["auto", "hipblas", "hipblaslt", "rocblas"]) +def test_rocm_blas_request_accepts_supported_values(value): + assert gate._rocm_blas_request(value) == value + + +def test_rocm_blas_request_rejects_unknown_value(): + with pytest.raises(ValueError, match="expected auto, hipblas, hipblaslt, rocblas"): + gate._rocm_blas_request("cublas") + + +@pytest.mark.parametrize( + ("value", "expected"), + [("hipblas", {"TORCH_BLAS_PREFER_HIPBLASLT": "0"}), + ("rocblas", {"TORCH_BLAS_PREFER_HIPBLASLT": "0"}), + ("hipblaslt", {"TORCH_BLAS_PREFER_HIPBLASLT": "1"}), + ("auto", {})], +) +def test_blas_env_aliases(value, expected): + assert gate._blas_env(value) == expected + + +def test_explicit_policy_overrides_graph_gate(monkeypatch): + monkeypatch.setenv("FREETOKEN_ROCM_BLAS", "hipblaslt") + monkeypatch.setattr(gate, "_is_rocm", lambda: True) + monkeypatch.setattr(gate, "run_graph_gate", lambda: pytest.fail("graph gate must not run")) + assert gate.resolve_rocm_blas_env() == {"TORCH_BLAS_PREFER_HIPBLASLT": "1"} + + +def test_auto_policy_uses_passing_graph_gate(monkeypatch): + monkeypatch.setenv("FREETOKEN_ROCM_BLAS", "auto") + monkeypatch.setattr(gate, "_is_rocm", lambda: True) + result = {"ok": True, "env": {"TORCH_BLAS_PREFER_HIPBLASLT": "0"}} + monkeypatch.setattr(gate, "run_graph_gate", lambda: result) + assert gate.resolve_rocm_blas_env() == result["env"] + + +def test_explicit_policy_fails_when_effective_api_unavailable(monkeypatch): + monkeypatch.setenv("FREETOKEN_ROCM_BLAS", "hipblas") + monkeypatch.setattr(gate, "_is_rocm", lambda: True) + monkeypatch.setattr(gate, "torch", None, raising=False) + class Backends: + cuda = object() + class Torch: + backends = Backends() + monkeypatch.setitem(__import__("sys").modules, "torch", Torch()) + with pytest.raises(RuntimeError, match="preferred_blas_library unavailable"): + gate.resolve_rocm_blas_env() + + +def test_report_normalizes_rocm_alias(monkeypatch): + monkeypatch.setenv("FREETOKEN_ROCM_BLAS", "rocblas") + monkeypatch.setattr(gate, "_is_rocm", lambda: True) + monkeypatch.setattr( + gate, + "_effective_blas", + lambda: ("hipblas", "reported"), + ) + report = gate.rocm_blas_report(gate={}) + assert report["requested"] == "rocblas" + assert report["effective"] == "hipblas" + assert report["verification"] == "verified" + + +def test_graph_capture_env_invalid_policy_is_not_swallowed(monkeypatch): + gate.graph_capture_env.cache_clear() + monkeypatch.setenv("FREETOKEN_ROCM_BLAS", "bad") + with pytest.raises(ValueError): + gate.graph_capture_env() + gate.graph_capture_env.cache_clear() diff --git a/tests/engine/test_graph_capture_env.py b/tests/engine/test_graph_capture_env.py new file mode 100644 index 000000000..32639c5e6 --- /dev/null +++ b/tests/engine/test_graph_capture_env.py @@ -0,0 +1,24 @@ +"""Graph worker environment resolution tests.""" + +from __future__ import annotations + +import freetoken.utils.graph_gate as gate + + +def test_cpu_graph_env_is_noop(monkeypatch): + gate.graph_capture_env.cache_clear() + monkeypatch.setenv("FREETOKEN_ROCM_BLAS", "hipblaslt") + monkeypatch.setattr(gate, "_is_rocm", lambda: False) + monkeypatch.setattr(gate, "run_graph_gate", lambda: (_ for _ in ()).throw(AssertionError())) + assert gate.graph_capture_env() == {} + gate.graph_capture_env.cache_clear() + + +def test_auto_graph_env_is_single_gate_result(monkeypatch): + gate.graph_capture_env.cache_clear() + monkeypatch.setenv("FREETOKEN_ROCM_BLAS", "auto") + monkeypatch.setattr(gate, "_is_rocm", lambda: True) + expected = {"ok": True, "env": {"TORCH_BLAS_PREFER_HIPBLASLT": "0"}} + monkeypatch.setattr(gate, "run_graph_gate", lambda: expected) + assert gate.graph_capture_env() == expected["env"] + gate.graph_capture_env.cache_clear() diff --git a/tests/engine/test_resident_budget.py b/tests/engine/test_resident_budget.py new file mode 100644 index 000000000..98a821151 --- /dev/null +++ b/tests/engine/test_resident_budget.py @@ -0,0 +1,28 @@ +import pytest + +from freetoken.engine.resident_budget import phase_memory, required_phase_bytes + + +def test_phase_uses_driver_high_water_without_double_counting_non_torch_bytes(): + phase = phase_memory( + "capture", + start_free_bytes=900, + end_free_bytes=700, + allocator_peak_allocated_bytes=100, + allocator_peak_reserved_bytes=150, + minimum_driver_free_bytes=650, + total_driver_bytes=1000, + ) + assert phase.driver_used_high_water_bytes == 350 + assert phase.non_torch_bytes == 200 + assert phase.required_bytes == 350 + assert required_phase_bytes([phase], safety_bytes=50) == 400 + + +def test_phase_rejects_impossible_driver_counter(): + with pytest.raises(ValueError, match="within total"): + phase_memory( + "load", start_free_bytes=1, end_free_bytes=1, + allocator_peak_allocated_bytes=0, allocator_peak_reserved_bytes=0, + minimum_driver_free_bytes=2, total_driver_bytes=1, + ) diff --git a/tests/engine/test_sample.py b/tests/engine/test_sample.py new file mode 100644 index 000000000..497cc5153 --- /dev/null +++ b/tests/engine/test_sample.py @@ -0,0 +1,44 @@ +"""Sampler semantic and preparation-cache gates.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.core import SamplingParams +from freetoken.engine.sample import Sampler, apply_penalties + + +def test_penalties_ignore_prompt_and_count_generated_tokens(): + logits = torch.zeros(1, 8) + req = SimpleNamespace( + prompt_len=3, + input_ids=torch.tensor([1, 2, 3, 4, 4, 6], dtype=torch.int32), + sampling_params=SamplingParams(presence_penalty=1.0, frequency_penalty=0.5), + ) + apply_penalties(logits, [req]) + assert logits[0, 1].item() == 0.0 # prompt token is excluded + assert logits[0, 4].item() == -2.0 # presence + two generated occurrences + assert logits[0, 6].item() == -1.5 + + +def test_greedy_without_penalty_has_no_sampling_tensors(): + sampler = Sampler(torch.device("cpu"), vocab_size=8) + req = SimpleNamespace(sampling_params=SamplingParams()) + args = sampler.prepare(SimpleNamespace(reqs=[req])) + assert args.temperatures is None + assert args.top_k is None and args.top_p is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA/ROCm device") +def test_sampling_args_reused_for_stable_request_semantics(): + sampler = Sampler(torch.device("cuda"), vocab_size=8) + params = SamplingParams(temperature=0.7, top_k=4, top_p=0.95) + req = SimpleNamespace(sampling_params=params) + first = sampler.prepare(SimpleNamespace(reqs=[req])) + second = sampler.prepare(SimpleNamespace(reqs=[req])) + assert first is second + assert first.temperatures is not None + assert first.top_k is not None and first.top_p is not None diff --git a/tests/engine/test_sample_capture.py b/tests/engine/test_sample_capture.py new file mode 100644 index 000000000..533938137 --- /dev/null +++ b/tests/engine/test_sample_capture.py @@ -0,0 +1,74 @@ +"""Capture-safe sampler contract tests.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.core import SamplingParams +from freetoken.engine.engine import DeviceTokenChain +from freetoken.engine.sample import BatchSamplingArgs, Sampler + + +def test_sample_into_reuses_int32_output(): + sampler = Sampler(torch.device("cpu"), vocab_size=4) + logits = torch.tensor([[1.0, 4.0, 2.0, 3.0], [8.0, 2.0, 1.0, 0.0]]) + out = torch.full((2,), -1, dtype=torch.int32) + scratch = torch.empty(2, dtype=torch.int64) + result = sampler.sample_into(logits, BatchSamplingArgs(temperatures=None), None, out, scratch) + assert result.data_ptr() == out.data_ptr() + assert result.tolist() == [1, 0] + + +def test_sample_into_device_preserves_stable_output_address(): + sampler = Sampler(torch.device("cpu"), vocab_size=4) + logits = torch.tensor([[1.0, 4.0, 2.0, 3.0]]) + out = torch.full((1,), -1, dtype=torch.int32) + scratch = torch.empty(1, dtype=torch.int64) + result = sampler.sample_into_device( + logits, BatchSamplingArgs(temperatures=None), None, out, scratch + ) + assert result.data_ptr() == out.data_ptr() + assert result.tolist() == [1] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA/ROCm device") +def test_device_token_chain_keeps_async_observation_sources_distinct(): + device = torch.device("cuda") + stream = torch.cuda.Stream(device=device) + chain = DeviceTokenChain(device, capacity=1) + with torch.cuda.stream(stream): + first = chain.publish(torch.tensor([17], device=device), stream) + _, cpu0, event0 = chain.stage(first, stream) + second = chain.publish(torch.tensor([23], device=device), stream) + _, cpu1, event1 = chain.stage(second, stream) + event0.synchronize() + event1.synchronize() + assert first.data_ptr() != second.data_ptr() + assert cpu0.item() == 17 and cpu1.item() == 23 + + +@pytest.mark.parametrize( + "args", + [ + BatchSamplingArgs(torch.ones(1)), + BatchSamplingArgs(None, apply_penalties=True), + ], +) +def test_sample_into_rejects_dynamic_modes(args): + sampler = Sampler(torch.device("cpu"), vocab_size=4) + with pytest.raises(ValueError, match="greedy sampling without penalties"): + sampler.sample_into(torch.zeros(1, 4), args, None, torch.empty(1, dtype=torch.int32)) + + +def test_sampler_cpu_greedy_penalty_path_has_no_nvtx_dependency(): + sampler = Sampler(torch.device("cpu"), vocab_size=4) + req = SimpleNamespace( + prompt_len=0, + input_ids=torch.tensor([1], dtype=torch.int32), + sampling_params=SamplingParams(presence_penalty=1.0), + ) + args = BatchSamplingArgs(None, apply_penalties=True) + assert sampler.sample(torch.zeros(1, 4), args, SimpleNamespace(reqs=[req])).item() == 0 diff --git a/tests/kernels/test_attention_q8.py b/tests/kernels/test_attention_q8.py new file mode 100644 index 000000000..cb525aced --- /dev/null +++ b/tests/kernels/test_attention_q8.py @@ -0,0 +1,68 @@ +import torch + +from freetoken.kernel.triton.attention import q8_paged_attention +from freetoken.kernel.triton.q8_kv import quantize_row_q8_0_ref + + +def test_q8_paged_attention_matches_dequantized_reference(): + torch.manual_seed(2) + values_k = torch.randn(3, 2, 32) + values_v = torch.randn(3, 2, 32) + kq, ks = quantize_row_q8_0_ref(values_k) + vq, vs = quantize_row_q8_0_ref(values_v) + kq, vq = kq.reshape(3, 2, 32), vq.reshape(3, 2, 32) + ks, vs = ks.reshape(3, 2, 1), vs.reshape(3, 2, 1) + q = torch.randn(1, 4, 32) + out = q8_paged_attention( + q, + kq, + vq, + ks, + vs, + torch.tensor([0, 3], dtype=torch.int32), + torch.tensor([0, 1, 2], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + torch.tensor([2], dtype=torch.int32), + 32**-0.5, + ) + k = kq.float() * ks.float().repeat_interleave(32, dim=-1) + v = vq.float() * vs.float().repeat_interleave(32, dim=-1) + scores = torch.stack([ + q[0, h].float().matmul(k[:, h % 2].transpose(0, 1)) for h in range(4) + ]) * 32**-0.5 + ref = torch.stack([ + torch.softmax(scores[h], dim=-1).matmul(v[:, h % 2]) for h in range(4) + ]) + assert torch.allclose(out[0].float(), ref, rtol=1e-2, atol=1e-2) + + +def test_q8_paged_attention_prefill_is_causal_and_handles_page_slots(): + torch.manual_seed(7) + values_k = torch.randn(4, 1, 32) + values_v = torch.randn(4, 1, 32) + kq, ks = quantize_row_q8_0_ref(values_k) + vq, vs = quantize_row_q8_0_ref(values_v) + kq, vq = kq.reshape(4, 1, 32), vq.reshape(4, 1, 32) + ks, vs = ks.reshape(4, 1, 1), vs.reshape(4, 1, 1) + q = torch.randn(4, 2, 32) + out = q8_paged_attention( + q, kq, vq, ks, vs, + torch.tensor([0, 4], dtype=torch.int32), + # Non-monotonic physical slots model a page-table transition. + torch.tensor([3, 0, 2, 1], dtype=torch.int32), + torch.zeros(4, dtype=torch.int32), + torch.arange(4, dtype=torch.int32), + 32**-0.5, + ) + logical_k = kq[[3, 0, 2, 1]].float() * ks[[3, 0, 2, 1]].float().repeat_interleave(32, dim=-1) + logical_v = vq[[3, 0, 2, 1]].float() * vs[[3, 0, 2, 1]].float().repeat_interleave(32, dim=-1) + for token in range(4): + scores = torch.stack([ + q[token, head].float().matmul(logical_k[: token + 1, head % 1].transpose(0, 1)) + for head in range(2) + ]) * 32**-0.5 + expected = torch.stack([ + torch.softmax(scores[head], dim=-1).matmul(logical_v[: token + 1, 0]) + for head in range(2) + ]) + assert torch.allclose(out[token].float(), expected, rtol=1e-2, atol=1e-2) diff --git a/tests/kernels/test_gguf_dispatch.py b/tests/kernels/test_gguf_dispatch.py new file mode 100644 index 000000000..aee60505e --- /dev/null +++ b/tests/kernels/test_gguf_dispatch.py @@ -0,0 +1,122 @@ +"""Pure GGUF dispatch matrix; no kernel build or GPU required.""" + +from __future__ import annotations + +import pytest + +import freetoken.kernel.gguf as gguf + + +def _backend(monkeypatch, *, rocm: bool): + monkeypatch.setattr(gguf.torch.version, "hip", "7.2" if rocm else None) + monkeypatch.setattr(gguf.torch.version, "cuda", None if rocm else "13.0") + + +@pytest.mark.parametrize( + ("tokens", "implementation"), + [(1, "ggml_mul_mat_vec_a8"), (6, "ggml_mul_mat_vec_a8"), (8, "ggml_mul_mat_vec_a8")], +) +def test_dense_shape_policy(monkeypatch, tokens, implementation): + _backend(monkeypatch, rocm=True) + report = gguf.gguf_dispatch("dense", 8, 2048, 2048, tokens, "gfx1100") + assert report["implementation"] == implementation + assert report["quant_type"] == "Q8_0" + assert report["rows"] == 2048 + assert report["cols"] == 2048 + + +@pytest.mark.parametrize( + ("op", "implementation", "reason"), + [ + ("moe_decode", "ggml_moe_a8_vec", None), + ("moe_prefill", "ggml_moe_a8_vec", None), + ("grouped_prefill", "ggml_moe_a8", None), + ], +) +def test_moe_phase_is_explicit(monkeypatch, op, implementation, reason): + _backend(monkeypatch, rocm=True) + report = gguf.gguf_dispatch(op, 12, 4096, 2048, 4, "gfx1100") + assert report["implementation"] == implementation + assert report["reason"] == reason + + +def test_rdna3_prefill_switches_to_grouped_mmq(monkeypatch): + _backend(monkeypatch, rocm=True) + monkeypatch.setenv("FREETOKEN_GGUF_GROUPED_PREFILL", "1") + report = gguf.gguf_dispatch("moe_prefill", 12, 4096, 2048, 5, "gfx1100") + assert report["implementation"] == "ggml_moe_a8" + assert report["reason"] == "multi-token grouped path" + + +def test_rdna3_prefill_fails_closed_to_vector_by_default(monkeypatch): + _backend(monkeypatch, rocm=True) + monkeypatch.delenv("FREETOKEN_GGUF_GROUPED_PREFILL", raising=False) + report = gguf.gguf_dispatch("moe_prefill", 12, 4096, 2048, 5, "gfx1100") + assert report["implementation"] == "ggml_moe_a8_vec" + assert report["reason"] == "grouped path disabled after gfx1100 launch failure" + + +def test_quant_k_alignment_is_fail_closed(monkeypatch): + _backend(monkeypatch, rocm=True) + report = gguf.gguf_dispatch("dense", 13, 4096, 255, 1, "gfx1100") + assert report["implementation"] == "unsupported" + assert report["reason"] == "K dimension 255 is not aligned to 256" + + +def test_q4_k_and_cuda_matrix(monkeypatch): + _backend(monkeypatch, rocm=False) + report = gguf.gguf_dispatch("dense", 12, 4096, 2048, 1, "sm90") + assert report["backend"] == "cuda" + assert report["implementation"] == "ggml_mul_mat_vec_a8" + assert report["quant_type"] == "Q4_K" + + +def test_arch_mismatch_is_unsupported(monkeypatch): + _backend(monkeypatch, rocm=True) + report = gguf.gguf_dispatch("dense", 8, 32, 32, 1, "sm90") + assert report["implementation"] == "unsupported" + assert report["reason"] == "NVIDIA architecture requested on ROCm" + + +def test_nvidia_only_forced_implementation_fails_loudly(monkeypatch): + _backend(monkeypatch, rocm=True) + with pytest.raises(ValueError, match="unsupported GGUF implementation"): + gguf.gguf_dispatch("dense", 8, 16, 16, 1, "gfx1100", impl="marlin") + + +def test_dispatch_trace_aggregates_calls(monkeypatch): + _backend(monkeypatch, rocm=True) + monkeypatch.setenv("FREETOKEN_GGUF_DISPATCH_TRACE", "1") + gguf._DISPATCH_COUNTS.clear() + gguf.gguf_dispatch("dense", 8, 32, 32, 1, "gfx1100") + rows = gguf.gguf_dispatch_report() + assert rows[-1]["implementation"] == "ggml_mul_mat_vec_a8" + assert rows[-1]["calls"] == 1 + + +def test_rocm_jit_source_stays_cache_local(monkeypatch, tmp_path): + from torch.utils import cpp_extension + + build_root = tmp_path / "torch_extensions" + source = tmp_path / "gguf_kernel.cu" + source.write_text("// source\n") + monkeypatch.setattr( + cpp_extension, + "_get_build_directory", + lambda name, verbose: str(build_root / name), + ) + + target = gguf._rocm_jit_source("gguf_probe", source) + + assert target == build_root / "gguf_probe" / "gguf_kernel.cu" + assert target.read_text() == source.read_text() + + +def test_b10434_bs1_workspace_is_aligned_and_fail_closed(): + from freetoken.kernel.gguf import mmvq_bs1_workspace_bytes, validate_mmvq_bs1_workspace + + required = mmvq_bs1_workspace_bytes(2048, 512, 8) + assert required % 256 == 0 + validate_mmvq_bs1_workspace(2048, 512, 8, required) + with pytest.raises(ValueError, match="too small"): + validate_mmvq_bs1_workspace(2048, 512, 8, required - 1) diff --git a/tests/kernels/test_gguf_linear.py b/tests/kernels/test_gguf_linear.py new file mode 100644 index 000000000..b373b47d0 --- /dev/null +++ b/tests/kernels/test_gguf_linear.py @@ -0,0 +1,105 @@ +"""Independent finite/reference gate for native GGUF dense GEMV.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel.gguf import ggml_mul_mat_vec_a8 +from freetoken.layers.gguf import fused_mul_mat_gguf +from freetoken.models.gguf.dequant import ( + GGML_Q4_K, + GGML_Q5_K, + GGML_Q6_K, + GGML_Q8_0, + dequantize, + dequant_q6_k, + row_bytes, +) + + +@pytest.mark.parametrize("quant_type", [GGML_Q4_K, GGML_Q5_K, GGML_Q6_K, GGML_Q8_0]) +def test_native_k_quant_reference_shape(quant_type): + raw = torch.zeros(row_bytes(256, quant_type), dtype=torch.uint8) + decoded = dequantize(raw, quant_type, torch.float32) + assert decoded.shape == (256,) + assert torch.equal(decoded, torch.zeros(256)) + + +@pytest.mark.parametrize("operation", ["dense", "lm_head"]) +def test_dense_policy_covers_lm_head_operation(monkeypatch, operation): + import freetoken.kernel.gguf as gguf + + monkeypatch.setattr(gguf, "gguf_runtime_metadata", lambda: {"arch": "gfx1100"}) + monkeypatch.setattr(gguf, "_runtime_backend", lambda: "rocm") + seen = {} + + def dispatch(op, quant_type, rows, cols, tokens, arch): + seen.update(op=op, quant_type=quant_type, rows=rows, cols=cols, tokens=tokens, arch=arch) + return {"implementation": "unsupported"} + + monkeypatch.setattr(gguf, "gguf_dispatch", dispatch) + monkeypatch.setattr(gguf, "ggml_dequantize", lambda *args: torch.zeros((args[2], args[3]))) + x = torch.zeros((1, 256), dtype=torch.float32) + qweight = torch.zeros((7, 210), dtype=torch.uint8) + out = fused_mul_mat_gguf(x, qweight, GGML_Q6_K, operation) + assert out.shape == (1, 7) + assert seen == { + "op": operation, "quant_type": GGML_Q6_K, "rows": 7, "cols": 256, + "tokens": 1, "arch": "gfx1100", + } + + +ROCM_DEVICE = pytest.mark.skipif( + torch.version.hip is None or not torch.cuda.is_available(), + reason="needs CUDA/ROCm device", +) + + +def _q8_weight(rows: int, cols: int) -> torch.Tensor: + blocks = torch.zeros((rows, cols // 32, 34), dtype=torch.uint8, device="cuda") + blocks[..., :2] = torch.tensor([128, 63], dtype=torch.uint8, device="cuda") + blocks[..., 2:] = torch.randint(0, 255, blocks[..., 2:].shape, dtype=torch.uint8, device="cuda") + return blocks.reshape(rows, row_bytes(cols, GGML_Q8_0)) + + +def _q6_weight(rows: int, cols: int) -> torch.Tensor: + blocks = torch.zeros((rows, cols // 256, 210), dtype=torch.uint8, device="cuda") + blocks[..., 192:208] = torch.randint( + 1, 255, blocks[..., 192:208].shape, dtype=torch.uint8, device="cuda" + ) + blocks[..., 208:210] = torch.tensor([0, 60], dtype=torch.uint8, device="cuda") + return blocks.reshape(rows, row_bytes(cols, GGML_Q6_K)) + + +def _q8_reference(weight: torch.Tensor) -> torch.Tensor: + blocks = weight.reshape(weight.shape[0], -1, 34) + scale = blocks[..., :2].contiguous().view(torch.float16).float() + quant = blocks[..., 2:].contiguous().view(torch.int8).float() + return (scale * quant).reshape(weight.shape[0], -1) + + +def _q8_1_activation_reference(x: torch.Tensor) -> torch.Tensor: + """Mirror gguf_kernel.cu quantize_q8_1, including per-32-block scales.""" + padded = ((x.shape[1] + 511) // 512) * 512 + padded_x = torch.nn.functional.pad(x.float(), (0, padded - x.shape[1])) + grouped = padded_x.reshape(x.shape[0], -1, 32) + d = grouped.abs().amax(dim=-1, keepdim=True) / 127.0 + d = d.clamp_min(1e-9).to(torch.float16).float() + return (torch.round(grouped / d) * d).reshape_as(padded_x)[:, : x.shape[1]] + + +@ROCM_DEVICE +@pytest.mark.parametrize("quant_type", [GGML_Q8_0, GGML_Q6_K]) +def test_gguf_linear_matches_independent_dequant(quant_type): + rows, cols = (7, 256) if quant_type == GGML_Q8_0 else (5, 256) + weight = _q8_weight(rows, cols) if quant_type == GGML_Q8_0 else _q6_weight(rows, cols) + x = torch.randn(1, cols, generator=torch.Generator(device="cpu").manual_seed(2026), dtype=torch.bfloat16) + x = x.to("cuda") + output = ggml_mul_mat_vec_a8(weight, x, quant_type, rows).float() + dense = _q8_reference(weight) if quant_type == GGML_Q8_0 else dequant_q6_k(weight, torch.float32).reshape(rows, cols) + reference = _q8_1_activation_reference(x) @ dense.T + assert torch.isfinite(output).all() + # Q6_K GEMV accumulates packed sub-blocks in device precision/order; keep + # this reference gate aligned with the MoE packed-kernel gate. + torch.testing.assert_close(output, reference, rtol=5e-2, atol=0.5) diff --git a/tests/kernels/test_gguf_moe.py b/tests/kernels/test_gguf_moe.py new file mode 100644 index 000000000..b348709d2 --- /dev/null +++ b/tests/kernels/test_gguf_moe.py @@ -0,0 +1,271 @@ +"""Reference and finite-output gate for the gfx1100 GGUF MoE candidate.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel.gguf import ( + ggml_dequantize, + ggml_moe_a8_vec, + ggml_moe_a8_vec_strided, + ggml_moe_gate_up_swiglu_id, + ggml_moe_mmvdq_id, + ggml_moe_mmvq_id, +) +from freetoken.models.gguf.dequant import GGML_Q5_K, GGML_Q6_K +from freetoken.moe.fused_gguf import MoeDecodeWork + + +def test_weighted_route_reduce_is_fixed_order_on_cpu(): + from freetoken.moe.fused_gguf import _reduce_routes + + routes = torch.arange(12, dtype=torch.float32).reshape(2, 2, 3) + weights = torch.tensor([[0.25, 0.75], [0.6, 0.4]]) + output = _reduce_routes(routes, torch.empty(2, 3), weights) + torch.testing.assert_close(output, (routes * weights[..., None]).sum(dim=1)) + + +ROCM_GFX1100 = pytest.mark.skipif( + torch.version.hip is None + or not torch.cuda.is_available() + or getattr(torch.cuda.get_device_properties(0), "gcnArchName", "") != "gfx1100", + reason="needs ROCm gfx1100", +) + + +def test_moe_decode_work_has_explicit_id_space_and_reuses_buffers(): + work = MoeDecodeWork("moe_decode") + hidden = torch.zeros((1, 256), dtype=torch.bfloat16) + gate_up = torch.empty((4, 16, 144), dtype=torch.uint8) + down = torch.empty((4, 8, 210), dtype=torch.uint8) + ids = torch.tensor([[0, 3]], dtype=torch.int32) + weights = torch.tensor([[0.25, 0.75]], dtype=torch.float32) + + work.bind(hidden, gate_up, down, weights, ids, id_space="slot", down_quant_type=GGML_Q6_K) + first = work.reserve("output", (1, 8), torch.bfloat16, torch.device("cpu")) + second = work.reserve("output", (1, 8), torch.bfloat16, torch.device("cpu")) + + assert work.id_space == "slot" + assert work.gate_expert_stride_bytes == gate_up.stride(0) + assert work.down_row_stride_bytes == down.stride(1) + assert first.data_ptr() == second.data_ptr() + + +def test_moe_decode_work_rejects_raw_route_dtype_mismatch(): + work = MoeDecodeWork("moe_decode") + tensors = ( + torch.zeros((1, 256), dtype=torch.bfloat16), + torch.empty((4, 16, 144), dtype=torch.uint8), + torch.empty((4, 8, 210), dtype=torch.uint8), + torch.ones((1, 2), dtype=torch.float32), + torch.ones((1, 2), dtype=torch.int64), + ) + with pytest.raises(ValueError, match="IDs must be int32"): + work.bind(*tensors, id_space="raw", down_quant_type=GGML_Q6_K) + + +def test_native_mixed_down_types_preserve_dispatch_and_routes(monkeypatch): + import freetoken.moe.fused_gguf as fused + + calls = [] + + def fake_dispatch(phase, quant_type, rows, cols, tokens, arch): + calls.append((phase, quant_type, rows, cols, tokens, arch)) + return {"implementation": "ggml_moe_a8_vec"} + + def fake_matmul(x, weights, ids, quant_type, row, dispatch, output=None): + calls.append(("matmul", quant_type, tuple(x.shape), tuple(ids.shape), row, output)) + if output is not None: + output.zero_() + return output + return torch.zeros((ids.shape[0] * ids.shape[1], row), dtype=x.dtype) + + monkeypatch.setattr(fused, "_gguf_moe_matmul", fake_matmul) + def fake_act(x, out=None): + value = x[..., : x.shape[-1] // 2] + if out is not None: + out.copy_(value) + return out + return value + + monkeypatch.setattr(fused, "_ACT", {"test": fake_act}) + monkeypatch.setattr("freetoken.kernel.gguf.gguf_runtime_metadata", lambda: {"arch": "gfx1100"}) + monkeypatch.setattr("freetoken.kernel.gguf.gguf_dispatch", fake_dispatch) + + hidden = torch.ones((2, 4), dtype=torch.bfloat16) + gate_up = torch.empty((3, 8, 1), dtype=torch.uint8) + topk_weights = torch.tensor([[0.25, 0.75], [0.6, 0.4]], dtype=torch.float32) + topk_ids = torch.tensor([[2, 0], [1, 2]], dtype=torch.int32) + for down_type in (GGML_Q5_K, GGML_Q6_K): + calls.clear() + workspace = {} + out = fused.fused_experts_gguf_native( + hidden, gate_up, torch.empty((3, 4, 1), dtype=torch.uint8), + topk_weights, topk_ids, "test", down_quant_type=down_type, + workspace=workspace, + ) + assert out.shape == (2, 4) + assert torch.isfinite(out).all() + assert topk_ids.tolist() == [[2, 0], [1, 2]] + dispatch_types = [row[1] for row in calls if row[0] == "moe_decode"] + assert dispatch_types == [12, down_type] + assert calls[-1][-1] is workspace["down"] + + +def test_grouped_helper_preserves_aligned_route_contract(monkeypatch): + import freetoken.moe.fused_gguf as fused + + calls = {} + + def fake_align(ids, block_size, experts): + calls["align"] = (ids.clone(), block_size, experts) + return (torch.tensor([0, 1, 4, 4], dtype=torch.int32), + torch.tensor([1, 2], dtype=torch.int32), + torch.tensor([4], dtype=torch.int32)) + + def fake_vec(*args): + calls["vec"] = True + return torch.empty((2, 4)) + + def fake_grouped(x, weights, sorted_ids, expert_ids, padded, quant_type, row, top_k, tokens): + calls["grouped"] = (tuple(x.shape), tuple(weights.shape), tuple(sorted_ids.shape), + tuple(expert_ids.shape), int(padded.item()), quant_type, row, + top_k, tokens) + return torch.zeros((tokens * top_k, row), dtype=x.dtype) + + monkeypatch.setattr("freetoken.moe.fused.moe_align_block_size", fake_align) + monkeypatch.setattr("freetoken.kernel.gguf.ggml_moe_get_block_size", lambda _qt: 32) + monkeypatch.setattr("freetoken.kernel.gguf.ggml_moe_a8_vec", fake_vec) + monkeypatch.setattr("freetoken.kernel.gguf.ggml_moe_a8", fake_grouped) + out = fused._gguf_moe_matmul( + torch.zeros((2, 4)), torch.zeros((3, 4, 1), dtype=torch.uint8), + torch.tensor([[1, 2], [0, 1]], dtype=torch.int32), 12, 4, + {"implementation": "ggml_moe_a8"}, + ) + assert out.shape == (4, 4) + assert calls["align"][0].tolist() == [[1, 2], [0, 1]] + assert calls["grouped"][-3:] == (4, 2, 2) + + +@ROCM_GFX1100 +@pytest.mark.slow +def test_native_kquant_stride_matches_compact_rows(): + """Q5_K rows padded to the Q6_K cache stride retain native GEMV output.""" + device = torch.device("cuda") + for quant_type, row_bytes in ((GGML_Q5_K, 176), (GGML_Q6_K, 210)): + x = torch.randn((2, 256), dtype=torch.bfloat16, device=device) + compact = torch.zeros((3, 2, row_bytes), dtype=torch.uint8, device=device) + if quant_type == GGML_Q5_K: + compact[..., :4] = torch.tensor([128, 63, 128, 63], dtype=torch.uint8, device=device) + compact[..., 4:16] = torch.randint(1, 64, (3, 2, 12), dtype=torch.uint8, device=device) + compact[..., 16:] = torch.randint(0, 255, (3, 2, row_bytes - 16), dtype=torch.uint8, device=device) + else: + compact[..., :208] = torch.randint(0, 255, (3, 2, 208), dtype=torch.uint8, device=device) + compact[..., 208:] = torch.tensor([128, 63], dtype=torch.uint8, device=device) + padded = torch.zeros((3, 2, 210), dtype=torch.uint8, device=device) + padded[..., :row_bytes].copy_(compact) + ids = torch.tensor([[2, 0], [1, 2]], dtype=torch.int32, device=device) + expected = ggml_moe_a8_vec(x, compact, ids, 2, quant_type, 2, 2) + actual = ggml_moe_a8_vec_strided( + x, padded, ids, 2, quant_type, 2, 2, + int(padded.stride(0)), int(padded.stride(1)), + ) + torch.cuda.synchronize(device) + torch.testing.assert_close(actual, expected, rtol=5e-2, atol=5e-2) + + +@ROCM_GFX1100 +@pytest.mark.slow +def test_gfx1100_moe_matches_legacy(monkeypatch): + device = torch.device("cuda") + generator = torch.Generator(device="cpu").manual_seed(2026) + ids = torch.arange(8, dtype=torch.int32, device=device).reshape(1, 8) + + x_gate = torch.randn(1, 256, generator=generator, dtype=torch.bfloat16, device="cpu").to(device) + q4 = torch.zeros((8, 16, 144), dtype=torch.uint8, device=device) + q4[..., :4] = torch.tensor([128, 63, 128, 63], dtype=torch.uint8, device=device) + q4[..., 4:16] = torch.randint(1, 64, (8, 16, 12), generator=generator, dtype=torch.uint8).to(device) + q4[..., 16:] = torch.randint(0, 255, (8, 16, 128), generator=generator, dtype=torch.uint8).to(device) + monkeypatch.setenv("FREETOKEN_GGUF_MOE_IMPL", "legacy") + legacy_gate = ggml_moe_a8_vec(x_gate, q4, ids, 8, 12, 16, 1) + monkeypatch.setenv("FREETOKEN_GGUF_MOE_IMPL", "gfx1100") + candidate_gate = ggml_moe_a8_vec(x_gate, q4, ids, 8, 12, 16, 1) + + x_down = torch.randn(8, 128, generator=generator, dtype=torch.bfloat16, device="cpu").to(device) + q8_blocks = torch.zeros((8, 16, 4, 34), dtype=torch.uint8, device=device) + q8_blocks[..., :2] = torch.tensor([128, 63], dtype=torch.uint8, device=device) + q8_blocks[..., 2:] = torch.randint( + 0, 255, (8, 16, 4, 32), generator=generator, dtype=torch.uint8 + ).to(device) + q8 = q8_blocks.reshape(8, 16, 136) + monkeypatch.setenv("FREETOKEN_GGUF_MOE_IMPL", "legacy") + legacy_down = ggml_moe_a8_vec(x_down, q8, ids, 1, 8, 16, 8) + monkeypatch.setenv("FREETOKEN_GGUF_MOE_IMPL", "gfx1100") + candidate_down = ggml_moe_a8_vec(x_down, q8, ids, 1, 8, 16, 8) + torch.cuda.synchronize(device) + + assert torch.isfinite(candidate_gate).all() + assert torch.isfinite(candidate_down).all() + torch.testing.assert_close(candidate_gate, legacy_gate, rtol=5e-2, atol=5e-2) + torch.testing.assert_close(candidate_down, legacy_down, rtol=5e-2, atol=5e-2) + + +@ROCM_GFX1100 +@pytest.mark.slow +def test_rdna3_mmvdq_and_fused_gate_up_match_dequant_reference(monkeypatch): + device = torch.device("cuda") + generator = torch.Generator(device="cpu").manual_seed(2027) + hidden, intermediate, top_k, experts = 256, 256, 2, 8 + ids = torch.tensor([[1, 7]], dtype=torch.int32, device=device) + x = torch.randn((1, hidden), generator=generator, dtype=torch.float32, device="cpu").to(device) + gate = torch.randint( + 0, 255, (experts, 2 * intermediate, 144), generator=generator, dtype=torch.uint8 + ).to(device) + gate[..., :4] = torch.tensor([128, 63, 128, 63], dtype=torch.uint8, device=device) + gate[..., 4:16] = torch.randint( + 1, 4, gate[..., 4:16].shape, generator=generator, dtype=torch.uint8 + ).to(device) + monkeypatch.setenv("FREETOKEN_GGUF_MOE_IMPL", "rdna3_mmvdq") + mmvdq = ggml_moe_mmvdq_id( + x, gate, ids, top_k, 12, 2 * intermediate, 1, + int(gate.stride(0)), int(gate.stride(1)), "raw", + ) + monkeypatch.setenv("FREETOKEN_GGUF_MOE_IMPL", "rdna3_mmid") + mmvq = ggml_moe_mmvq_id( + x, gate, ids, top_k, 12, 2 * intermediate, 1, + int(gate.stride(0)), int(gate.stride(1)), "raw", + ) + fused = ggml_moe_gate_up_swiglu_id( + x, gate, ids, top_k, intermediate, 1, + int(gate.stride(0)), int(gate.stride(1)), "raw", + ) + reference = torch.cat( + [x @ ggml_dequantize(gate[expert], 12, 2 * intermediate, hidden, torch.float32).t() + for expert in ids[0].tolist()], + dim=0, + ) + mmvq_fused = ( + mmvq[:, :intermediate] / + (1.0 + torch.exp(-mmvq[:, :intermediate])) * mmvq[:, intermediate:] + ) + torch.cuda.synchronize(device) + torch.testing.assert_close(mmvdq, reference, rtol=3e-4, atol=3e-3) + torch.testing.assert_close(fused, mmvq_fused, rtol=3e-4, atol=3e-3) + + q6 = torch.randint(0, 255, (experts, hidden, 210), generator=generator, dtype=torch.uint8).to(device) + q6[..., -2:] = torch.tensor([128, 63], dtype=torch.uint8, device=device) + down_x = torch.randn((top_k, intermediate), generator=generator, dtype=torch.float32, device="cpu").to(device) + route_ids = ids.reshape(-1, 1) + monkeypatch.setenv("FREETOKEN_GGUF_MOE_IMPL", "rdna3_mmvdq") + down = ggml_moe_mmvdq_id( + down_x, q6, route_ids, 1, GGML_Q6_K, hidden, top_k, + int(q6.stride(0)), int(q6.stride(1)), "raw", + ) + down_reference = torch.cat( + [down_x[i:i + 1] @ ggml_dequantize(q6[expert], GGML_Q6_K, hidden, intermediate, torch.float32).t() + for i, expert in enumerate(ids[0].tolist())], + dim=0, + ) + torch.cuda.synchronize(device) + torch.testing.assert_close(down, down_reference, rtol=3e-4, atol=3e-3) diff --git a/tests/kvcache/test_kv_storage.py b/tests/kvcache/test_kv_storage.py new file mode 100644 index 000000000..880def685 --- /dev/null +++ b/tests/kvcache/test_kv_storage.py @@ -0,0 +1,40 @@ +import pytest +import torch + +from freetoken.engine.config import EngineConfig, KVStorageType +from freetoken.kvcache.base import KVStorageDescriptor, spec_kv_bytes_per_token + + +def test_q8_descriptor_matches_b10434_row_contract(): + desc = KVStorageDescriptor(KVStorageType.Q8_0) + assert desc.bytes_per_block == 34 + assert desc.row_bytes(2048) == 2048 + 2 * 64 + assert desc.bytes_per_token(num_layers=1, num_kv_heads=4, head_dim=2048) == 2 * 4 * 2176 + with pytest.raises(ValueError, match="divisible"): + desc.row_bytes(33) + + +def test_q8_spec_price_counts_local_heads_and_both_slabs(): + class TP: + size = 2 + + class Spec: + name = "full" + mla = False + index_head_dim = 0 + num_index_layers = 0 + index_ratio = 1 + head_dim = 64 + num_kv_heads = 8 + num_layers = 3 + + class Config: + kv_storage_type = KVStorageType.Q8_0 + tp_info = TP() + + assert spec_kv_bytes_per_token(Spec(), Config()) == 2 * 3 * 4 * (64 + 4) + + +def test_engine_config_normalizes_kv_storage_type_without_touching_gpu(): + cfg = EngineConfig(model_path="dummy", tp_info=type("TP", (), {"rank": 0, "size": 1})(), dtype=torch.bfloat16, kv_storage_type="q8_0") + assert cfg.kv_storage_type is KVStorageType.Q8_0 diff --git a/tests/kvcache/test_mha_q8.py b/tests/kvcache/test_mha_q8.py new file mode 100644 index 000000000..61adda94c --- /dev/null +++ b/tests/kvcache/test_mha_q8.py @@ -0,0 +1,79 @@ +import pytest +import torch + +from freetoken.kvcache.mha_pool import MHAKVCache +from freetoken.kernel.triton.q8_kv import quantize_row_q8_0_ref, store_q8_cache + + +def test_q8_store_matches_reference_and_keeps_rows_independent(): + values = torch.tensor( + [[[0.5, -0.5, 1.0, -1.0] * 8], [[2.0, 0.0, -2.0, 0.0] * 8]], dtype=torch.float32 + ) + payload = torch.empty((4, 1, 32), dtype=torch.int8) + scales = torch.empty((4, 1, 1), dtype=torch.float16) + store_q8_cache( + k_payload=payload, + v_payload=payload.clone(), + k_scales=scales, + v_scales=scales.clone(), + indices=torch.tensor([2, 0], dtype=torch.int32), + k=values, + v=values, + ) + expected_payload, expected_scales = quantize_row_q8_0_ref(values) + assert torch.equal( + payload.index_select(0, torch.tensor([2, 0])).reshape(2, 32), expected_payload + ) + assert torch.equal( + scales.index_select(0, torch.tensor([2, 0])), expected_scales.reshape(2, 1, 1) + ) + + +def test_q8_duplicate_destination_rejected_before_store(): + with pytest.raises(ValueError, match="duplicate"): + store_q8_cache( + k_payload=torch.empty((2, 1, 32), dtype=torch.int8), + v_payload=torch.empty((2, 1, 32), dtype=torch.int8), + k_scales=torch.empty((2, 1, 1), dtype=torch.float16), + v_scales=torch.empty((2, 1, 1), dtype=torch.float16), + indices=torch.tensor([0, 0], dtype=torch.int32), + k=torch.ones((2, 1, 32)), + v=torch.ones((2, 1, 32)), + ) + + +def test_q8_mha_pool_reports_packed_unit_bytes_and_generation(): + from freetoken.distributed import set_tp_info, try_get_tp_info + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + pool = MHAKVCache( + num_kv_heads=2, + num_layers=3, + head_dim=32, + num_pages=2, + page_size=4, + dtype=torch.bfloat16, + device=torch.device("cpu"), + storage_type="q8_0", + ) + assert pool.is_quantized + assert pool.unit_bytes() == (2 * 3 * 2 * 34, 0) + before = pool.pointer_generation + pool.rebuild(3) + assert pool.pointer_generation == before + 1 + assert pool.k_cache_view(0).payload.dtype is torch.int8 + assert pool.k_cache_view(0).scales.dtype is torch.float16 + + +def test_q8_zero_rows_initialize_payload_and_scales(): + payload = torch.full((1, 1, 32), 9, dtype=torch.int8) + scales = torch.full((1, 1, 1), 3, dtype=torch.float16) + store_q8_cache( + k_payload=payload, v_payload=payload.clone(), + k_scales=scales, v_scales=scales.clone(), + indices=torch.tensor([0], dtype=torch.int32), + k=torch.zeros((1, 1, 32)), v=torch.zeros((1, 1, 32)), + ) + assert torch.count_nonzero(payload) == 0 + assert torch.count_nonzero(scales) == 0 diff --git a/tests/models/test_qwen35moe_gguf_deint.py b/tests/models/test_qwen35moe_gguf_deint.py index 4154fe27e..176233ff8 100644 --- a/tests/models/test_qwen35moe_gguf_deint.py +++ b/tests/models/test_qwen35moe_gguf_deint.py @@ -7,12 +7,20 @@ """ import torch +import pytest from freetoken.models.qwen3_5_moe.gguf import ( _gdn_head_perm, _deint_dense_rows, _deint_q8_cols, _deint_q8_rows, + _gguf_down_quant_types, +) +from freetoken.models.gguf.dequant import ( + GGML_Q5_K, + GGML_Q6_K, + dequantize, + quantize_q8_0, ) @@ -69,3 +77,57 @@ def test_deint_q8_cols_recovers_contiguous(): inter[:, perm[h] * bbp:(perm[h] + 1) * bbp] = x[:, h * bbp:(h + 1) * bbp] rec = _deint_q8_cols(inter, nv, blocks_per_head) assert torch.allclose(rec, x, atol=1e-6) + + +def _dequant_q8_0_reference(raw: torch.Tensor) -> torch.Tensor: + """Independent Q8_0 unpack used to check the cache's re-quantized bytes.""" + blocks = raw.reshape(-1, 34) + scales = blocks[:, :2].contiguous().view(torch.float16).float() + values = blocks[:, 2:].contiguous().view(torch.int8).float() + return (values * scales).reshape(raw.shape[0], -1) + + +@pytest.mark.parametrize( + ("ggml_type", "row_bytes"), + [(GGML_Q5_K, 176), (GGML_Q6_K, 210)], +) +def test_k_quant_source_requantizes_to_q8_cache_within_tolerance(ggml_type, row_bytes): + """FreeToken's Q8_0 expert cache preserves Q5_K/Q6_K source values closely.""" + generator = torch.Generator().manual_seed(100 + ggml_type) + packed = torch.randint(0, 256, (4, row_bytes), dtype=torch.uint8, generator=generator) + scales = torch.tensor([0.0, -0.5, 0.75, 1.25], dtype=torch.float16) + if ggml_type == GGML_Q5_K: + # Q5_K stores (dall, dmin) in its first two fp16 values. + packed[:, :4] = torch.stack((scales, scales.abs() / 2), dim=1).view(torch.uint8) + else: + packed[:, 208:210] = scales.view(torch.uint8).reshape(4, 2) + + source = dequantize(packed, ggml_type, torch.float32).view(4, 256) + cache_bytes = quantize_q8_0(source) + cached = _dequant_q8_0_reference(cache_bytes) + assert torch.isfinite(source).all() + assert torch.isfinite(cached).all() + assert torch.equal(cached[0], torch.zeros_like(cached[0])) + relative_rmse = (cached - source).square().mean().sqrt() / source.square().mean().sqrt() + assert relative_rmse < 0.01 + + +def test_down_quant_type_reader_preserves_mixed_layer_types(monkeypatch): + from freetoken.models.gguf import reader + + class Tensor: + def __init__(self, name, quant_type): + self.name = name + self.tensor_type = quant_type + + class FakeReader: + tensors = [ + Tensor("blk.0.ffn_down_exps.weight", GGML_Q5_K), + Tensor("blk.1.ffn_down_exps.weight", GGML_Q6_K), + Tensor("blk.2.ffn_down_exps.weight", GGML_Q5_K), + ] + + monkeypatch.setattr(reader, "_reader", lambda _path: FakeReader()) + assert _gguf_down_quant_types("synthetic.gguf") == ( + GGML_Q5_K, GGML_Q6_K, GGML_Q5_K + ) diff --git a/tests/models/test_qwen35moe_moe.py b/tests/models/test_qwen35moe_moe.py new file mode 100644 index 000000000..062a02779 --- /dev/null +++ b/tests/models/test_qwen35moe_moe.py @@ -0,0 +1,113 @@ +"""Resident Qwen3.5 MoE construction and metadata-only budget tests.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.distributed import set_tp_info, try_get_tp_info +from freetoken.models.config import ModelConfig, RotaryConfig +from freetoken.models.gguf.dequant import GGML_Q5_K, GGML_Q6_K +from freetoken.models.qwen3_5_moe.moe import Qwen3_5MoE + + +@pytest.fixture(scope="module", autouse=True) +def _single_rank_tp(): + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + +def _config(*, expert_quant="none", moe_weight_format=None, down_types=()): + return ModelConfig( + num_layers=1, + num_qo_heads=4, + num_kv_heads=4, + head_dim=64, + hidden_size=256, + vocab_size=320, + intermediate_size=0, + rms_norm_eps=1e-6, + rotary_config=RotaryConfig(64, 64, 1024, 10000.0, None), + hidden_act="silu", + tie_word_embeddings=False, + num_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=256, + norm_topk_prob=True, + model_type="qwen3_5_moe", + architectures=["Qwen35moeForCausalLM"], + moe_enabled=True, + expert_quant=expert_quant, + moe_weight_format=moe_weight_format, + gguf_down_quant_types=down_types, + shared_expert_intermediate_size=256, + ) + + +@pytest.mark.parametrize("down_type", [GGML_Q5_K, GGML_Q6_K]) +def test_gguf_resident_construction_uses_native_packed_shapes(down_type): + config = _config(moe_weight_format="gguf", expert_quant="gguf", down_types=(down_type,)) + with torch.device("meta"): + moe = Qwen3_5MoE(config, layer_id=0) + experts = moe.experts + assert experts.weight_format == "gguf" + assert experts.gate_up_proj.shape == (2, 512, 144) + assert experts.down_proj.shape[-1] == (176 if down_type == GGML_Q5_K else 210) + assert experts.gate_up_proj.dtype == torch.uint8 + assert experts.down_proj.dtype == torch.uint8 + + +def test_bf16_resident_construction_unchanged(): + config = _config() + with torch.device("meta"): + moe = Qwen3_5MoE(config, layer_id=0) + assert moe.experts.weight_format == "bf16" + assert moe.experts.gate_up_proj.shape == (2, 512, 256) + assert moe.experts.down_proj.shape == (2, 256, 256) + + +def test_fp8_resident_construction_unchanged(): + config = _config(expert_quant="fp8_block") + with torch.device("meta"): + moe = Qwen3_5MoE(config, layer_id=0) + assert moe.experts.weight_format == "fp8_block" + assert moe.experts.gate_up_proj.shape == (2, 512, 256) + + +def test_resident_budget_is_allocation_free(monkeypatch): + from freetoken.engine import resident_budget + + class Pool: + @classmethod + def kv_cost(cls, _config): + return 1024, 2048, 16, 0 + + model_config = SimpleNamespace( + moe_weight_format="gguf", + linear_attention_group=lambda: None, + ) + config = SimpleNamespace( + model_path="synthetic.gguf", + model_config=model_config, + max_seq_len=64, + page_size=16, + max_running_req=1, + cache_type="naive", + tp_info=SimpleNamespace(size=1), + dtype=torch.bfloat16, + ) + monkeypatch.setattr(resident_budget, "_gguf_payload_bytes", lambda _path: (10000, 700)) + monkeypatch.setattr(resident_budget, "_total_vram_bytes", lambda _free: 8 * (1 << 30)) + monkeypatch.setattr("freetoken.kvcache.resolve_pool_class", lambda _mc: Pool) + monkeypatch.setattr("freetoken.kvcache.linear_state_pool.state_pool_bytes", lambda *_a, **_k: 3000) + + budget = resident_budget.estimate_gguf_resident_budget(config.model_path, config, 2 * (1 << 30)) + + assert budget.packed_model_bytes == 10000 + assert budget.kv_bytes == 6144 + assert budget.gdn_state_bytes == 3000 + assert budget.graph_reserve_bytes == 768 * (1 << 20) + assert budget.peak_load_scratch_bytes == 512 * (1 << 20) + assert budget.safety_bytes == 1_500 * (1 << 20) + assert budget.required_bytes > budget.free_bytes + assert budget.fits is False diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 477d215d8..453bc6a14 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -415,7 +415,9 @@ def test_adjust_config_converts_moe_cache_rate_to_cache_size(): model_path="/tmp/freetoken-test-model", tp_info=DistributedInfo(rank=0, size=1), dtype=torch.float16, - attention_backend="fi", + # ``fi`` is NVIDIA-only; this test exercises MoE config arithmetic on both + # CUDA and ROCm, so use portable attention backend. + attention_backend="triton", moe_cache_rate=0.3, ) object.__setattr__( diff --git a/tests/scheduler/test_decode_handoff.py b/tests/scheduler/test_decode_handoff.py new file mode 100644 index 000000000..1332adfb0 --- /dev/null +++ b/tests/scheduler/test_decode_handoff.py @@ -0,0 +1,43 @@ +"""One-step overlap token staging/event-order gates.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.engine.engine import TokenStaging +from freetoken.scheduler.status import SchedulerStatusReporter + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA/ROCm device") +def test_token_staging_alternates_slots_and_preserves_cpu_values(): + device = torch.device("cuda") + stream = torch.cuda.Stream(device=device) + staging = TokenStaging(device, capacity=2) + with torch.cuda.stream(stream): + gpu0, cpu0, event0 = staging.stage(torch.tensor([11], device=device), stream) + gpu1, cpu1, event1 = staging.stage(torch.tensor([22], device=device), stream) + stream.synchronize() + event0.synchronize() + event1.synchronize() + assert gpu0.data_ptr() != 0 and gpu1.data_ptr() != 0 + assert cpu0.data_ptr() != cpu1.data_ptr() + assert event0 is not event1 + assert cpu0.item() == 11 and cpu1.item() == 22 + + # First event is drained before slot zero is reused, matching scheduler overlap order. + event0.synchronize() + with torch.cuda.stream(stream): + _, cpu2, event2 = staging.stage(torch.tensor([33], device=device), stream) + stream.synchronize() + assert event2 is event0 + event2.synchronize() + assert cpu2.item() == 33 + + +def test_decode_stats_only_read_on_configured_log_interval(): + reporter = SchedulerStatusReporter(log=lambda _: None, decode_log_interval=2) + batch = type("DecodeBatch", (), {"is_decode": True})() + assert not reporter.decode_stats_due(batch) + reporter._decode_forward_count = 1 + assert reporter.decode_stats_due(batch) diff --git a/tests/utils/test_decode_benchmark_metadata.py b/tests/utils/test_decode_benchmark_metadata.py new file mode 100644 index 000000000..534bebc5b --- /dev/null +++ b/tests/utils/test_decode_benchmark_metadata.py @@ -0,0 +1,138 @@ +"""Pure provenance and rejection gates for ROCm decode benchmark adapters.""" + +from __future__ import annotations + +import hashlib +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +BENCHMARKS = Path(__file__).resolve().parents[2] / "benchmarks" +if str(BENCHMARKS) not in sys.path: + sys.path.insert(0, str(BENCHMARKS)) + +from bench_decode_moe import acceptance_status, execution_metadata, model_fingerprint # noqa: E402 +from bench_decode_ollama import ollama_blob_identity # noqa: E402 +from bench_llama_cpp_hip import _parse_cli_timing # noqa: E402 + + +def _args(**overrides): + values = { + "decode": 512, + "problem": 0, + "attention_backend": "triton", + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_file_model_identity_is_full_sha256(tmp_path): + model = tmp_path / "model.gguf" + payload = b"stable model bytes\n" + model.write_bytes(payload) + expected = hashlib.sha256(payload).hexdigest() + + identity = model_fingerprint(str(model)) + + assert identity["sha256"] == expected + assert identity["identity"] == "full-file-sha256" + assert identity["size_bytes"] == len(payload) + + +def test_directory_identity_ignores_mtime_but_tracks_content(tmp_path): + model = tmp_path / "checkpoint" + model.mkdir() + shard = model / "weights.bin" + shard.write_bytes(b"one") + first = model_fingerprint(str(model)) + shard.touch() + assert model_fingerprint(str(model))["sha256"] == first["sha256"] + shard.write_bytes(b"two") + assert model_fingerprint(str(model))["sha256"] != first["sha256"] + + +@pytest.mark.parametrize( + ("completion", "events", "text", "graph_state", "accepted"), + [ + (512, 4, "answer", "replay", True), + (511, 4, "answer", "replay", False), + (512, 1, "answer", "replay", False), + (512, 4, "", "replay", False), + (512, 4, "answer", "unknown", False), + ], +) +def test_acceptance_status_rejects_non_comparable_rows( + completion, events, text, graph_state, accepted +): + result = acceptance_status( + args=_args(), + result={"stamps": [0.0] * events, "text": text}, + graph={"state": graph_state}, + usage={"completion_tokens": completion}, + model={"sha256": "model-sha"}, + ) + + assert result["accepted"] is accepted + assert result["status"] == ("accepted" if accepted else "rejected") + + +def test_ollama_manifest_digest_is_not_gguf_identity(tmp_path): + model = tmp_path / "model.gguf" + model.write_bytes(b"same") + + unproven = ollama_blob_identity("sha256:manifest") + verified = ollama_blob_identity( + "sha256:manifest", + ollama_gguf=str(model), + reference_gguf=str(model), + ) + + assert unproven["same_blob"] is False + assert unproven["status"] == "unproven" + assert verified["same_blob"] is True + assert verified["status"] == "verified" + assert verified["manifest_digest_kind"] == "ollama-model-manifest" + + +def test_ollama_gguf_mismatch_is_explicit(tmp_path): + ollama_model = tmp_path / "ollama.gguf" + reference_model = tmp_path / "reference.gguf" + ollama_model.write_bytes(b"ollama") + reference_model.write_bytes(b"freetoken") + + identity = ollama_blob_identity( + "sha256:manifest", + ollama_gguf=str(ollama_model), + reference_gguf=str(reference_model), + ) + + assert identity["same_blob"] is False + assert identity["status"] == "mismatch" + + +def test_llama_cli_timing_parser_requires_eval_count(): + output = "eval time = 6400.0 ms / 512 runs (12.5 ms per token, 80.0 tokens per second)" + assert _parse_cli_timing(output) == (512, 80.0) + assert _parse_cli_timing("no timings") is None + + +def test_execution_metadata_reads_observed_cache_geometry(): + observed = { + "effective_moe_backend": "fused", + "expert_storage": "resident_gguf", + "resident_gguf": True, + "expert_fetches": 0, + "expert_remaps": 0, + } + value = execution_metadata( + args=_args(), + backend="fused", + graph={"state": "replay", "gate": "pass"}, + cache_status={"geometry": {"execution": observed}}, + ) + assert value["effective_moe_backend"] == "fused" + assert value["expert_storage"] == "resident_gguf" + assert value["graph_state"] == "replay" diff --git a/tests/utils/test_decode_gate.py b/tests/utils/test_decode_gate.py new file mode 100644 index 000000000..c368cd47f --- /dev/null +++ b/tests/utils/test_decode_gate.py @@ -0,0 +1,210 @@ +"""Pure final decode gate tests.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +BENCHMARKS = Path(__file__).resolve().parents[2] / "benchmarks" +if str(BENCHMARKS) not in sys.path: + sys.path.insert(0, str(BENCHMARKS)) + +from check_decode_gate import evaluate_gate # noqa: E402 + + +MODEL_SHA = "a" * 64 +PROMPT_SHA = "b" * 64 +FIXTURE_SHA = "c" * 64 +SAMPLING = {"temperature": 1.0, "top_p": 0.95, "top_k": 64} +PROBE = { + "schema": "qwen-moe-base-probe-v2", + "model": {"sha256": MODEL_SHA}, + "prompt_sha256": PROMPT_SHA, + "mtp": "off", + "speculative": False, + "eager": {"finite_logits": True, "decode_rows": 8}, + "graph": {"finite_logits": True, "decode_rows": 8}, + "comparison": {"token_ids_equal": True}, +} + + +def _execution() -> dict: + return { + "effective_moe_backend": "fused", + "expert_storage": "resident_gguf", + "resident_gguf": True, + "expert_fetches": 0, + "expert_remaps": 0, + "attention_backend": "triton", + "graph_state": "replay", + "graph_gate": "pass", + "decode_batch_size": 1, + "mtp": "off", + "speculative": False, + } + + +def _free(value: float, *, status: str = "accepted", repeat: int | None = None) -> dict: + return { + "status": status, + "metadata": {"model_sha256": MODEL_SHA, "prompt_sha256": PROMPT_SHA, "sampling": SAMPLING, "mtp": "off"}, + "sampling": SAMPLING, + "context": 9216, + "batch": 512, + "ubatch": 512, + "kv_type": "q8_0", + "fixture_sha256": FIXTURE_SHA, + "completion_tokens": 512, + "decode_requested": 512, + "decode_tok_s": value, + "repeat": repeat, + "execution": _execution(), + "acceptance": {"accepted": status == "accepted"}, + } + + +def _ollama(value: float, *, repeat: int | None = None) -> dict: + return { + "schema": "ollama-base-decode-v1", + "status": "accepted", + "reference_identity": { + "status": "verified", + "same_blob": True, + "reference_gguf": {"sha256": MODEL_SHA}, + }, + "prompt_sha256": PROMPT_SHA, + "context": 9216, + "batch": 512, + "ubatch": 512, + "kv_type": "q8_0", + "fixture_sha256": FIXTURE_SHA, + "mtp": "off", + "speculative": False, + "acceptance": {"accepted": True, "checks": {"mtp_off": True}}, + "options": SAMPLING, + "client_arrival_tok_s": value, + "repeat": repeat, + } + + +def test_gate_passes_only_with_verified_reference_and_confidence(): + free = [_free(82.0, repeat=i) for i in range(10)] + reference = [_ollama(81.0, repeat=i) for i in range(10)] + result = evaluate_gate(free, reference, PROBE, min_runs=10) + assert result["gate"] is True + assert result["reference_identity"]["kind"] == "ollama-client-arrival" + assert result["runs"]["p02_5_bootstrap_tok_s"] == 82.0 + + +def test_directional_threshold_never_promotes_unproven_reference(): + result = evaluate_gate([_free(90.0) for _ in range(10)], [], min_runs=10) + assert result["gate"] is False + assert result["threshold_source"] == "directional-ollama-unproven" + assert any("matched reference unavailable" in reason for reason in result["reasons"]) + + +@pytest.mark.parametrize("value", [69.9, 75.0]) +def test_gate_rejects_low_run_or_confidence(value): + result = evaluate_gate( + [_free(value, repeat=i) for i in range(10)], + [_ollama(60.0, repeat=i) for i in range(10)], + min_runs=10, + ) + assert result["gate"] is False + assert any("below 70" in reason or "p02.5" in reason for reason in result["reasons"]) + + +def test_rejected_rows_cannot_be_hidden_from_gate(): + result = evaluate_gate( + [_free(90.0, repeat=0), _free(90.0, status="rejected", repeat=1)], + [_ollama(1.0, repeat=0)], + min_runs=1, + ) + assert result["gate"] is False + assert result["rejected"] == 1 + + +def test_empty_input_is_machine_readable_rejection(): + result = evaluate_gate([], [], min_runs=1) + assert result["gate"] is False + assert result["runs"]["median_tok_s"] is None + + +def test_malformed_accepted_row_is_rejected_without_gate_crash(): + row = _free(90.0) + row["model_fingerprint"] = None + row["metadata"] = [] + row["acceptance"] = None + result = evaluate_gate([row], [], min_runs=1) + assert result["gate"] is False + assert "FreeToken rows do not carry one full model SHA-256" in result["reasons"] + assert "accepted rows contain exact-completion or acceptance mismatch" in result["reasons"] + + +def test_non_sha_model_identity_cannot_pass_as_full_identity(): + row = _free(90.0) + row["metadata"]["model_sha256"] = "head-tail-sha256:abc" + result = evaluate_gate([row for _ in range(10)], [], min_runs=10) + assert result["gate"] is False + assert "FreeToken rows do not carry one full model SHA-256" in result["reasons"] + + +def test_missing_execution_evidence_cannot_promote(): + rows = [_free(90.0, repeat=i) for i in range(10)] + for row in rows: + del row["execution"] + result = evaluate_gate(rows, [], min_runs=10) + assert result["gate"] is False + assert any("missing execution-mode evidence" in reason for reason in result["reasons"]) + + +def test_missing_finite_logit_probe_cannot_promote_verified_reference(): + result = evaluate_gate( + [_free(90.0, repeat=i) for i in range(10)], + [_ollama(81.0, repeat=i) for i in range(10)], + min_runs=10, + ) + assert result["gate"] is False + assert "finite-logit/parity probe unavailable" in result["reasons"] + + +def test_probe_identity_and_parity_are_required(): + probe = dict(PROBE) + probe["prompt_sha256"] = "c" * 64 + probe["comparison"] = {"token_ids_equal": False} + result = evaluate_gate( + [_free(90.0, repeat=i) for i in range(10)], + [_ollama(81.0, repeat=i) for i in range(10)], + probe, + min_runs=10, + ) + assert result["gate"] is False + assert "finite-logit probe prompt SHA-256 does not match FreeToken rows" in result["reasons"] + assert "eager/graph greedy token parity failed or is unavailable" in result["reasons"] + + +@pytest.mark.parametrize("field", ["context", "batch", "ubatch", "kv_type", "fixture_sha256"]) +def test_gate_rejects_comparator_mismatch(field): + rows = [_free(90.0, repeat=i) for i in range(10)] + wrong = { + "context": 8192, + "batch": 1, + "ubatch": 1, + "kv_type": "bf16", + "fixture_sha256": "d" * 64, + } + rows[0][field] = wrong[field] + result = evaluate_gate(rows, [], PROBE, min_runs=10) + assert result["gate"] is False + needle = "fixture" if field == "fixture_sha256" else field + assert any(needle in reason for reason in result["reasons"]) + + +def test_gate_rejects_missing_fixture_identity(): + rows = [_free(90.0, repeat=i) for i in range(10)] + del rows[0]["fixture_sha256"] + result = evaluate_gate(rows, [], PROBE, min_runs=10) + assert result["gate"] is False + assert any("fixture SHA-256" in reason for reason in result["reasons"]) diff --git a/tests/utils/test_step_profiler.py b/tests/utils/test_step_profiler.py new file mode 100644 index 000000000..f84c36d63 --- /dev/null +++ b/tests/utils/test_step_profiler.py @@ -0,0 +1,25 @@ +"""Pure marker pairing and stage-summary tests.""" + +from freetoken.utils.step_profiler import summarize_markers + + +def test_summarize_markers_keeps_outer_step_separate_from_overlapping_phases(): + markers = [ + {"step": 1, "phase": "scheduler", "event": "begin", "monotonic_ns": 0}, + {"step": 1, "phase": "attention_metadata", "event": "begin", "monotonic_ns": 10}, + {"step": 1, "phase": "attention_metadata", "event": "end", "monotonic_ns": 40}, + {"step": 1, "phase": "scheduler", "event": "end", "monotonic_ns": 100}, + ] + summary = summarize_markers(markers) + assert summary["complete"] is True + assert summary["critical_step"]["median_ns"] == 100 + assert summary["phases"]["attention_metadata"]["median_ns"] == 30 + assert summary["note"].startswith("phase totals overlap") + + +def test_summarize_markers_rejects_unmatched_ranges(): + summary = summarize_markers([ + {"step": 1, "phase": "scheduler", "event": "end", "monotonic_ns": 1}, + ]) + assert summary["complete"] is False + assert "unmatched end" in summary["errors"][0]