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 new file mode 100644 index 000000000..64894ed07 --- /dev/null +++ b/.github/workflows/unit-nvidia.yml @@ -0,0 +1,103 @@ +# 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: 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." diff --git a/.github/workflows/unit-rocm.yml b/.github/workflows/unit-rocm.yml new file mode 100644 index 000000000..a4a379dd1 --- /dev/null +++ b/.github/workflows/unit-rocm.yml @@ -0,0 +1,213 @@ +# 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 + 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 + +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: 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 \ + > /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 }} + 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 + + 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/.gitignore b/.gitignore index bf804e075..edc08af25 100644 --- a/.gitignore +++ b/.gitignore @@ -227,3 +227,13 @@ benchmarks/cross_framework # local e2e/bench artifacts (harnesses may run with repo cwd) /results/ + +# torch.utils.cpp_extension ROCm hipify build artifacts (preprocessed .cu -> .hip, +# and includes rewritten to *_hip.cuh) written next to the GGUF sources. +python/freetoken/kernel/csrc/gguf/*.hip +python/freetoken/kernel/csrc/gguf/*_hip.cuh +python/freetoken/kernel/csrc/gguf/ggml-common_hip.h + +# local session junk +nohup.out +.plans/ diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 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/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_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 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 new file mode 100644 index 000000000..6edbea21c --- /dev/null +++ b/benchmarks/bench_gguf_moe_kernels.py @@ -0,0 +1,318 @@ +"""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 +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_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 + + +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( + "--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 + 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. + 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(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). + 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() + 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") + 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" + + 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_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): + 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() + 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()) 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 new file mode 100644 index 000000000..423826dd4 --- /dev/null +++ b/docs/install-amd.md @@ -0,0 +1,207 @@ +# AMD GPU (ROCm) support + +FreeToken targets Linux + NVIDIA CUDA by default. AMD (ROCm) is a supported, tested +configuration with a **single-GPU** milestone: correct functional path first, performance +recovered via HIP ports where safe. This page covers installing and running on RX 7000. + +> Status: **experimental.** The default and best-tested path remains CUDA. AMD brings up a +> correct functional path (Triton attention + offload/CPU MoE + portable quant) and is +> recovering performance via the HIP kernel ports. + +## Requirements + +| Component | Requirement | +| --- | --- | +| OS | Linux x86_64 (Windows WDDM pinned-memory is a known edge, not supported yet) | +| GPU | AMD RX 7000 (RDNA 3, `gfx1100`); RX 9000 (`gfx1201`) is future work | +| ROCm | ROCm toolkit with `hipcc` (`/opt/rocm/bin/hipcc` or on `PATH`) | +| torch | ROCm build, e.g. `torch==2.5.1+rocm6.2` | + +The build refuses to mix toolchains: it will **not** silently fall back to `nvcc`/`libcudart` +when only the ROCm toolkit is present, and vice versa. + +## Install + +```bash +# ROCm torch (PyTorch official ROCm wheels) -- must satisfy the repo's torch>=2.11,<2.12 +# build pin, so use the rocm7.2 index (rocm6.2 only carries torch up to 2.5.1). +pip install --index-url https://download.pytorch.org/whl/rocm7.2 \ + "torch==2.11.0+rocm7.2" torchvision triton-rocm==3.6.0 + +# FreeToken with the ROCm extra (builds the native extensions with hipcc) +uv pip install -e ".[rocm]" --no-build-isolation +``` + +`pip install ".[rocm]"` pulls ROCm-compatible `torch`/`triton`; the NVIDIA-only `[accel]` +packages (`flashinfer`, `sgl-kernel`, `triton_kernels`, Marlin) are **not** installed on AMD +and their backends are rejected with a clean error if requested. + +## Verified feature matrix + +| Feature | On AMD | Notes | +| --- | --- | --- | +| Attention | `--attention-backend triton` | flashinfer/fa/trtllm are NVIDIA-only and rejected | +| MoE | `--moe-backend 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) | | + +## 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` → 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). + +## 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; `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 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 historical 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** | + +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 + `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). 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..2153c7897 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ classifiers = [ "Intended Audience :: Science/Research", "Operating System :: POSIX :: Linux", "Environment :: GPU :: NVIDIA CUDA", + "Environment :: GPU :: AMD ROCm", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", @@ -78,6 +79,13 @@ fi = ["flashinfer-python[cu13]>=0.6,<0.7"] # renamed from sgl-kernel at 0.4; still imports as `sgl_kernel`, so never co-install both sgl = ["sglang-kernel==0.4.5"] accel = ["freetoken[fi,sgl]"] +# ROCm (AMD) install: the NVIDIA-only fi/sgl/Marlin packages are NOT pulled in. torch must +# come from the ROCm wheel index (e.g. `pip install torch==2.11.0+rocm7.2` from +# https://download.pytorch.org/whl/rocm7.2) so the native extensions build against the HIP +# runtime; this extra pins the rest. See docs/install-amd.md. +rocm = [ + "triton==3.6.0; platform_system == 'Linux'", +] # NOTE: the Marlin W4A16 NVFP4 expert-GEMM path (sm_80-99) borrows vLLM's AOT wheel # (vllm>=0.14,<0.15), which pins transformers>=4.56,<5 and so is INCOMPATIBLE with the # core transformers>=5.5 requirement. It is therefore not a lockable extra and is left diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index 8a410ffe4..28a6f3c8a 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") @@ -92,6 +95,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..c19a9fcf9 --- /dev/null +++ b/python/freetoken/attention/torch.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, List + +import os +import torch + +from freetoken.core import Batch, get_global_ctx + +from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata + +if TYPE_CHECKING: + from freetoken.models import ModelConfig + + +@dataclass +class TorchMetadata(BaseAttnMetadata): + """Minimal contiguous-gather metadata for the pure-torch backend. + + ``indices`` maps every logical KV position (across all padded requests, in + ``seqlens_k`` order) to its physical paged-cache slot, exactly like the triton + backend's gather. The torch backend reads the SAME paged cache as triton, so a + triton-vs-torch logit difference isolates the attention *compute* from the + cache addressing. + """ + + indices: torch.Tensor + seqlens_q: List[int] + seqlens_k: List[int] + cached_lens: List[int] + is_decode: bool + cu_seqlens_q: torch.Tensor + + def get_last_indices(self, bs: int) -> torch.Tensor: + return self.cu_seqlens_q[1 : 1 + bs] - 1 + + +class TorchAttentionBackend(BaseAttnBackend): + """Pure-PyTorch full-attention backend (no Triton/CUDA kernels). + + Serves as a numerically-explicit ground truth for debugging the hybrid + qwen35moe model. It stores K/V into the same paged MHAKVCache as + ``TritonAttentionBackend``, gathers the request's full K/V history via the + identical ``indices`` page gather, and computes GQA softmax attention with + PyTorch ops so every intermediate is auditable. + + Intended for correctness debugging / backend A-B comparison, not production + serving. Registered as the ``"torch"`` attention backend (``AttnType.FULL``). + """ + + def __init__(self, config: ModelConfig): + self.config = config + self.kvcache = get_global_ctx().kv_cache + self.device = self.kvcache.device + self.num_q_heads = int(getattr(config, "num_qo_heads", 1)) + self.num_kv_heads = int(getattr(config, "num_kv_heads", 1)) + self.head_dim = int(getattr(config, "head_dim", 1)) + # Prefer the full-attention group spec head_dim (authoritative for kv heads). + specs = getattr(config, "kv_cache_group_specs", lambda: ())() + for spec in specs: + name = getattr(spec, "attn_type", None) + if name is not None and str(name) == "AttnType.FULL": + self.head_dim = int(getattr(spec, "head_dim", self.head_dim)) + self.num_kv_heads = int(getattr(spec, "num_kv_heads", self.num_kv_heads)) + break + # Debugging: contiguous (per-request) cache instead of the paged pool, to + # isolate cache addressing from the attention compute. + self._contig: dict[tuple[int, int], list] = {} + self._use_contig = os.environ.get("FT_DEBUG_CONTIG_CACHE") == "1" + + def _build_metadata(self, batch: Batch) -> TorchMetadata: + ctx = get_global_ctx() + page_table = ctx.page_table + reqs = batch.padded_reqs + seqlens_q = [req.extend_len for req in reqs] + seqlens_k = [req.device_len for req in reqs] + cached_lens = [req.cached_len for req in reqs] + is_decode = max(seqlens_q) == 1 + indices = torch.cat([page_table[req.table_idx, : req.device_len] for req in reqs]) + if is_decode: + cu_seqlens_q = torch.arange(0, len(reqs) + 1, dtype=torch.int32, device=self.device) + else: + cu_seqlens_q = torch.tensor( + [0] + seqlens_q, dtype=torch.int32, device=self.device + ).cumsum_(0) + return TorchMetadata( + indices=indices, + seqlens_q=seqlens_q, + seqlens_k=seqlens_k, + cached_lens=cached_lens, + is_decode=is_decode, + cu_seqlens_q=cu_seqlens_q, + ) + + def prepare_metadata(self, batch: Batch) -> None: + batch.attn_metadata = self._build_metadata(batch) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer_id: int, + batch: Batch, + attn_spec: AttentionSpec | None = None, + ) -> torch.Tensor: + if self._use_contig: + return self._forward_contig(q, k, v, layer_id, batch, attn_spec) + self.kvcache.store_kv(k, v, batch.out_loc, layer_id) + + k_raw = self.kvcache.k_cache(layer_id) + v_raw = self.kvcache.v_cache(layer_id) + kv_heads, head_dim = k_raw.shape[-2], k_raw.shape[-1] + assert head_dim == q.shape[-1], f"head_dim {head_dim} != {q.shape[-1]}" + k_cache = k_raw.view(-1, kv_heads, head_dim) + v_cache = v_raw.view(-1, kv_heads, head_dim) + + metadata = batch.attn_metadata + assert isinstance(metadata, TorchMetadata) + k_all = k_cache[metadata.indices] # [total_kv, kv_heads, head_dim] + v_all = v_cache[metadata.indices] # [total_kv, kv_heads, head_dim] + + spec = attn_spec or AttentionSpec() + scale = spec.sm_scale if spec.sm_scale is not None else head_dim ** -0.5 + group = self.num_q_heads // kv_heads + + num_q_tokens = q.shape[0] + out = torch.empty((num_q_tokens, self.num_q_heads, head_dim), dtype=q.dtype, device=q.device) + q_off = 0 + k_off = 0 + for lq, lk, cached in zip( + metadata.seqlens_q, metadata.seqlens_k, metadata.cached_lens + ): + qs = q[q_off : q_off + lq] # [lq, num_q, head_dim] + ks = k_all[k_off : k_off + lk] # [lk, kv_heads, head_dim] + vs = v_all[k_off : k_off + lk] + if group > 1: + ks = ks.repeat_interleave(group, dim=1) # [lk, num_q, head_dim] + vs = vs.repeat_interleave(group, dim=1) + # [num_q, lq, lk] + scores = torch.einsum("qhd,khd->hqk", qs.float(), ks.float()) * scale + # causal: query i (global cached+i) attends key col j <= cached+i + if lq > 1 or lk > lq: + rows = torch.arange(lq, device=scores.device) + cols = torch.arange(lk, device=scores.device) + masked = (cols[None, :] > (cached + rows)[:, None]).to(scores.device) + scores = scores.masked_fill(masked[None, :, :], float("-inf")) + probs = torch.softmax(scores, dim=-1) + o = torch.einsum("hqk,khd->qhd", probs, vs.float()).to(q.dtype) # [lq, num_q, head_dim] + out[q_off : q_off + lq] = o + q_off += lq + k_off += lk + + return out + + def _forward_contig(self, q, k, v, layer_id, batch, attn_spec=None): + """Contiguous (non-paged) attention: K/V accumulate per (layer, request) in + a Python list keyed by logical position, so cache addressing is trivially + correct. Isolates the paged-cache addressing from the attention compute.""" + metadata = self._build_metadata(batch) + kv_heads = self.num_kv_heads + head_dim = self.head_dim + spec = attn_spec or AttentionSpec() + scale = spec.sm_scale if spec.sm_scale is not None else head_dim ** -0.5 + group = self.num_q_heads // kv_heads + # Store this forward's K/V rows per request (append in global position order). + q_off = 0 + for i, lq in enumerate(metadata.seqlens_q): + uid = batch.padded_reqs[i].uid + buf = self._contig.setdefault((layer_id, uid), {"k": [], "v": []}) + kseg = k[q_off : q_off + lq].view(lq, kv_heads, head_dim) + vseg = v[q_off : q_off + lq].view(lq, kv_heads, head_dim) + for t in range(lq): + buf["k"].append(kseg[t]) + buf["v"].append(vseg[t]) + q_off += lq + # compute attention from the contiguous cache + q_off = 0 + out = torch.empty( + (q.shape[0], self.num_q_heads, head_dim), dtype=q.dtype, device=q.device + ) + for i, lq in enumerate(metadata.seqlens_q): + uid = batch.padded_reqs[i].uid + buf = self._contig[(layer_id, uid)] + ks = torch.stack(buf["k"]) # [acc, kv_heads, head_dim] + vs = torch.stack(buf["v"]) + acc = ks.shape[0] + cached = acc - lq + qs = q[q_off : q_off + lq] + if group > 1: + ks = ks.repeat_interleave(group, dim=1) + vs = vs.repeat_interleave(group, dim=1) + scores = torch.einsum("qhd,khd->hqk", qs.float(), ks.float()) * scale + rows = torch.arange(lq, device=scores.device) + cols = torch.arange(acc, device=scores.device) + masked = (cols[None, :] > (cached + rows)[:, None]).to(scores.device) + scores = scores.masked_fill(masked[None, :, :], float("-inf")) + probs = torch.softmax(scores, dim=-1) + o = torch.einsum("hqk,khd->qhd", probs, vs.float()).to(q.dtype) + out[q_off : q_off + lq] = o + q_off += lq + return out + + def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + # ROCm/graph capture is disabled for this debugging backend; no-op. + return None + + def prepare_for_capture(self, batch: Batch) -> None: + return None + + def prepare_for_replay(self, batch: Batch) -> None: + return None + + +__all__ = ["TorchAttentionBackend", "TorchMetadata"] diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py index 9eed1e1d2..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] @@ -167,7 +187,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, @@ -238,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/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/__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 bcbe6bcf2..9d5e93a35 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" @@ -83,6 +105,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 b5a6fa3b0..f624a5f31 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -16,13 +16,14 @@ 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 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 @@ -51,6 +75,11 @@ def _flashinfer_available() -> bool: def _sgl_flash_attn_available() -> bool: + from freetoken.utils.arch import is_rocm + + # sgl_kernel is NVIDIA-only; never select it on ROCm even if a stray copy is importable. + if is_rocm(): + return False try: from sgl_kernel.flash_attn import flash_attn_with_kvcache # noqa: F401 except Exception as exc: @@ -102,9 +131,14 @@ def _backend_parts_serve(name: str, required: frozenset[AttnType]) -> bool: def _backend_requirements_met(name: str) -> bool: + infos = [attention_backend_info(part) for part in name.split(",")] + # On ROCm (AMD) only the portable backends exist: flashinfer/sgl/trtllm (and anything + # sm_100-gated) are NVIDIA-only, so short-circuit before probing them at all. + if is_rocm(): + return all(not i.requires_flashinfer and not i.requires_sgl_kernel + and not i.requires_sm100 for i in infos) # flashinfer first across ALL parts: the sgl probe logs a "falls back to fi" warning, # which would mislead when the candidate is about to fail on flashinfer anyway. - infos = [attention_backend_info(part) for part in name.split(",")] if any(i.requires_flashinfer for i in infos) and not _flashinfer_available(): return False if any(i.requires_sgl_kernel for i in infos) and not _sgl_flash_attn_available(): @@ -197,6 +231,15 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att # explicit --attention-backend choices. for part in backend_parts: info = attention_backend_info(part) + + # 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 " + 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 " @@ -290,6 +333,87 @@ class ForwardOutput(NamedTuple): 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() @@ -299,12 +423,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. @@ -322,6 +458,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 @@ -343,6 +480,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] @@ -355,6 +493,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() @@ -386,6 +531,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( @@ -396,6 +545,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)}") @@ -414,6 +570,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, @@ -426,10 +583,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: @@ -758,6 +958,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 @@ -919,13 +1123,24 @@ 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 - use_graph = self.graph_runner.can_use_cuda_graph(batch) + 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 profiler_phase("model_forward_setup"): + use_graph = self.graph_runner.can_use_cuda_graph(batch) with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, use_graph): - logits = self.graph_runner.replay(batch) if use_graph else self.model.forward() + 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. @@ -935,10 +1150,30 @@ 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_cpu = next_tokens_gpu.to("cpu", non_blocking=True) - copy_done_event = torch.cuda.Event() - copy_done_event.record(self.stream) + 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() @@ -1030,6 +1265,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 @@ -1229,6 +1467,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) @@ -1332,6 +1571,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 @@ -1470,7 +1714,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 4f2025025..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,9 +214,35 @@ 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, 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. + detail = run_graph_gate().get("detail", "") + logger.info_rank0( + "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: 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) @@ -145,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( @@ -181,21 +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: - 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: + 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 @@ -213,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 01d14b1aa..0519c2fae 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -1,13 +1,13 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, List import torch 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: @@ -54,12 +84,30 @@ 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] - 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) + 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] @@ -70,11 +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(temperatures, top_k=top_k, top_p=top_p) + 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) -> torch.Tensor: - with torch.cuda.nvtx.range("Sampler"): + def sample( + self, logits: torch.Tensor, args: BatchSamplingArgs, batch: Batch + ) -> torch.Tensor: + 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/_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/aot_models.py b/python/freetoken/kernel/aot_models.py index 0ce249de8..b6af1b0f7 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -94,6 +94,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 @@ -182,6 +187,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/backend.py b/python/freetoken/kernel/backend.py index fc42c49fd..8e685bbc1 100644 --- a/python/freetoken/kernel/backend.py +++ b/python/freetoken/kernel/backend.py @@ -10,6 +10,8 @@ import functools import importlib.util +from freetoken.utils.arch import is_rocm + def _importable(name: str) -> bool: # find_spec normally returns None when a package is absent, but it can raise @@ -21,16 +23,48 @@ 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") +@functools.cache +def is_triton_kernels_installed() -> bool: + """OpenAI's ``triton_kernels`` (the fused MoE router used by ``moe.fused.fused_topk``). + + Distinct from the ``triton`` runtime we always depend on: it ships with the Triton + source tree and has no Windows wheel. It is also not one of the six ops + ``freetoken.kernel.triton`` reimplements, so its call-site carries its own fallback. + """ + if is_rocm(): + return False + return _importable("triton_kernels") + + @functools.cache def driver_cuda_version() -> int | None: """Max CUDA version the installed NVIDIA driver supports (``13000`` == CUDA 13.0), 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 8349123bd..93b4a1d00 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 @@ -622,7 +644,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 @@ -1968,13 +1990,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/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/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..3f2eaa52b 100644 --- a/python/freetoken/kernel/csrc/gguf/ggml-common.h +++ b/python/freetoken/kernel/csrc/gguf/ggml-common.h @@ -8,7 +8,18 @@ #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 +// 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 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 d88960d5f..d04c26460 100644 --- a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu +++ b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu @@ -1,14 +1,48 @@ -// 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 +// itself is not directly includable on ROCm (missing a generated header). +#include +#include +#include +#include +using cudaStream_t = hipStream_t; +#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" +#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" @@ -16,6 +50,7 @@ #include "mmq.cuh" #include "moe.cuh" #include "moe_vec.cuh" +#endif // clang-format off // Q8 gemv @@ -77,11 +112,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 +134,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 +232,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 +379,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", [&] { @@ -545,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 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(); + 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) { @@ -809,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: @@ -843,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/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..1e063374c 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -10,6 +10,16 @@ #include #include +#if defined(USE_HIP) +#include +// HIP has no __grid_constant__ (a CUDA read-only-constant optimization). Define it +// empty so `const __grid_constant__ Params params` compiles as a plain by-value +// parameter, which is correct (just without the CUDA constant-cache hint). +#ifndef __grid_constant__ +#define __grid_constant__ +#endif +#endif + namespace device { inline constexpr auto kWarpThreads = 32u; @@ -42,16 +52,23 @@ __always_inline __device__ auto offset(const T *ptr, U... offset) -> const namespace PDL { +// Programmatic Dependent Launch is a CUDA-only optimization (griddepcontrol). +// HIP has no equivalent; the wait/launch are no-ops there. PDL is optional in the +// kernels (use_pdl defaults to false), so dropping it is purely a perf change. template __always_inline __device__ void wait() { +#if !defined(USE_HIP) if constexpr (kUsePDL) { asm volatile("griddepcontrol.wait;" ::: "memory"); } +#endif } template __always_inline __device__ void launch() { +#if !defined(USE_HIP) if constexpr (kUsePDL) { asm volatile("griddepcontrol.launch_dependents;" :::); } +#endif } } // namespace PDL @@ -60,12 +77,27 @@ template __always_inline __device__ void launch() { namespace host { +#if defined(USE_HIP) +inline auto +HIP_CHECK(::hipError_t error, + std::source_location location = std::source_location::current()) + -> void { + if (error != ::hipSuccess) [[unlikely]] { + ::host::panic(location, "HIP error: ", ::hipGetErrorString(error)); + } +} + +inline auto +HIP_CHECK(std::source_location location = std::source_location::current()) + -> void { + return HIP_CHECK(::hipGetLastError(), location); +} +#else inline auto CUDA_CHECK(::cudaError_t error, std::source_location location = std::source_location::current()) -> void { - if (error != ::cudaSuccess) { - [[unlikely]]; + if (error != ::cudaSuccess) [[unlikely]] { ::host::panic(location, "CUDA error: ", ::cudaGetErrorString(error)); } } @@ -75,7 +107,21 @@ CUDA_CHECK(std::source_location location = std::source_location::current()) -> void { return CUDA_CHECK(::cudaGetLastError(), location); } +#endif +#if defined(USE_HIP) +template inline void set_smem_once(std::size_t smem_size) { + static const auto last_smem_size = [&] { + HIP_CHECK(::hipFuncSetAttribute( + F, ::hipFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + return smem_size; + }(); + RuntimeCheck( + smem_size <= last_smem_size, + "Dynamic shared memory size exceeds the previously set maximum size: ", + last_smem_size, " bytes"); +} +#else template inline void set_smem_once(std::size_t smem_size) { static const auto last_smem_size = [&] { CUDA_CHECK(::cudaFuncSetAttribute( @@ -87,7 +133,52 @@ template inline void set_smem_once(std::size_t smem_size) { "Dynamic shared memory size exceeds the previously set maximum size: ", last_smem_size, " bytes"); } +#endif + +#if defined(USE_HIP) +struct LaunchKernel { +public: + explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, DLDevice device, + std::size_t dynamic_shared_mem_bytes = 0) noexcept + : m_grid(grid_dim), m_block(block_dim), + m_stream(resolve_device(device)), m_smem(dynamic_shared_mem_bytes) {} + + explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, hipStream_t stream, + std::size_t dynamic_shared_mem_bytes = 0) noexcept + : m_grid(grid_dim), m_block(block_dim), m_stream(stream), + m_smem(dynamic_shared_mem_bytes) {} + + static auto resolve_device(DLDevice device) -> hipStream_t { + return static_cast( + ::TVMFFIEnvGetStream(device.device_type, device.device_id)); + } + + LaunchKernel(const LaunchKernel &) = delete; + LaunchKernel &operator=(const LaunchKernel &) = delete; + + template + auto operator()(T &&kernel, Args &&...args) const -> void { + // hipLaunchKernel takes a void** args array (pointers to each argument value), + // unlike cudaLaunchKernelEx's variadic form. The array is consumed at launch. + void *arg_array[sizeof...(Args)] = { + const_cast(static_cast(&args))...}; + HIP_CHECK(::hipLaunchKernel(reinterpret_cast(kernel), m_grid, + m_block, arg_array, m_smem, m_stream)); + } + auto with_attr(bool /*use_pdl*/) -> LaunchKernel & { + // HIP has no programmatic dependent launch / launch attributes; PDL is a + // CUDA-only optimization and is dropped here. + return *this; + } + +private: + dim3 m_grid; + dim3 m_block; + hipStream_t m_stream; + std::size_t m_smem; +}; +#else struct LaunchKernel { public: explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, DLDevice device, @@ -140,5 +231,6 @@ private: cudaLaunchConfig_t m_config; cudaLaunchAttribute m_attr_cache; }; +#endif } // namespace host diff --git a/python/freetoken/kernel/csrc/jit/batch_memcpy.cuh b/python/freetoken/kernel/csrc/jit/batch_memcpy.cuh index 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..dce0c3819 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -16,11 +16,250 @@ import os import pathlib import shutil +import time import torch _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. @@ -47,30 +286,301 @@ 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). 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( + 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 + lock.unlink() + except Exception: # noqa: BLE001 - hygiene must never break the build path + 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 - extra_cuda_cflags = ["-O3", "--expt-relaxed-constexpr"] - host_cxx = _host_compiler() - if host_cxx is not None: - # Point both nvcc's host pass (-ccbin) and torch's C++ compile (CXX) at a - # libtorch/nvcc-compatible compiler. Force (not setdefault): the system - # default (CXX unset -> g++) can be a gcc too new for the torch headers. - cxx_path = shutil.which(host_cxx) or host_cxx - extra_cuda_cflags += ["-ccbin", cxx_path] - os.environ["CXX"] = cxx_path - os.environ["CC"] = _c_compiler_for(cxx_path) + if torch.version.hip is not None: + # ROCm: hipcc (torch.utils.cpp_extension picks it up), pass the HIP defines so + # the kernels compile their HIP branches; drop the CUDA-only -ccbin/flag logic. + # Explicit --offload-arch (plus PYTORCH_ROCM_ARCH) prevents torch from auto- + # emitting ~14 gfx arches, which would multiply build time per arch. + gfx = os.getenv("FREETOKEN_KERNEL_CACHE_GFX", "gfx1100") + os.environ.setdefault("PYTORCH_ROCM_ARCH", gfx) + extra_cuda_cflags = [ + "-O3", f"--offload-arch={gfx}", "-DUSE_HIP=1", "-DUSE_ROCM=1", + "-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", + "-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 + # 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) + 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 + # 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. - return load( + _clear_stale_jit_lock("freetoken_gguf_kernels") + 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=True, + 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) ---- @@ -123,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: @@ -133,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/activation.py b/python/freetoken/kernel/triton/activation.py index 0e4d28c8f..505bb417e 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/python/freetoken/kernel/triton/activation.py @@ -50,20 +50,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 @@ -140,10 +135,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/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/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/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/kernel/utils.py b/python/freetoken/kernel/utils.py index 7a0164b59..df7583003 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -40,6 +40,32 @@ def _rank(a: str) -> int: cc = max(arch_list, key=_rank).rstrip("a").replace(".", "") flags = flags + [f"-gencode=arch=compute_{cc},code=compute_{cc}"] return flags + + +def _rocm_cflags(extra: List[str]) -> List[str]: + """HIP/ROCm flags for a kernel build. Adds the ``USE_HIP``/``USE_ROCM`` defines so the + shared ``.cu``/``.cuh`` sources compile their HIP branches (llama.cpp-style). When + ``TVM_FFI_ROCM_ARCH_LIST`` (e.g. "gfx1100") is set (AOT cache build), we pin + ``--offload-arch``; otherwise hipcc targets the local GPU.""" + flags = DEFAULT_CFLAGS + ["-DUSE_HIP=1", "-DUSE_ROCM=1"] + list(extra) + arch_list = os.getenv("TVM_FFI_ROCM_ARCH_LIST", "").split() + if arch_list: + flags = flags + [f"--offload-arch={arch_list[0]}"] + return flags + + +def _arch_flags(extra: List[str]) -> List[str]: + """GPU kernel build flags for the active backend: HIP flags on ROCm torch, else CUDA.""" + try: + from freetoken.kernel._toolchain import is_rocm_torch + + if is_rocm_torch(): + return _rocm_cflags(extra) + except Exception: + pass + return _cuda_cflags(extra) + + CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool] @@ -80,10 +106,24 @@ def _version_parts(version: str) -> Tuple[str, List[str]]: return base, local.split(".") if local else [] -def _build_stamps(segments: List[str]) -> set[str]: +def _build_stamps(local_segments: List[str]) -> List[str]: """The `g` commit-stamp tokens of a local version segment list - (stamped by scripts/build-release-wheels.sh).""" - return {s for s in segments if re.fullmatch(r"g[0-9a-f]{7,40}", s)} + (``["cu130", "g3f01615"]`` -> ``["g3f01615"]``). + + 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]: + """The backend-tag tokens of a local segment list (``cu130`` or ``rocm``). A cache + wheel and runtime wheel must carry the SAME backend tag -- a ``+rocm`` cache is + meaningless to a ``+cu130`` runtime (the fatbin is gfx SASS vs sm SASS) and vice + versa. Returns the tags (usually one, e.g. ``["cu130"]`` or ``["rocm"]``).""" + return [s for s in local_segments if s.startswith("cu") or s.startswith("rocm")] def _kernel_cache_version_ok(cache_version: str, runtime_version: str) -> bool: @@ -93,14 +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/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index 41b82f1d5..11aaba862 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -12,6 +12,8 @@ BaseCacheHandle, BaseKVCachePool, BasePrefixCache, + KVStorageDescriptor, + QuantizedKVView, MatchResult, SizeInfo, ) @@ -117,6 +119,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), ) @@ -128,6 +131,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 @@ -255,6 +259,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 d68d8ded5..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}" ) @@ -171,12 +219,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) ) @@ -205,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, @@ -217,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, @@ -531,6 +593,55 @@ def _expert_gemm( return fused_experts_gguf_q4_0( hidden_states, gate_up, down, topk_weights, topk_ids, self.activation ) + 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, + 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 # decode + grouped `_t` prefill. The swiglu scalars live on the layer @@ -598,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 @@ -628,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 8ce69d654..e0ced8a72 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -308,6 +308,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/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..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 @@ -23,6 +22,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 +34,8 @@ GGML_BF16: (1, 2), 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), } @@ -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", } @@ -79,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.""" @@ -115,9 +156,72 @@ 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_Q4_K: dequant_q4_k, + GGML_Q5_K: dequant_q5_k, GGML_Q6_K: dequant_q6_k, + GGML_Q8_0: dequant_q8_0, } @@ -143,11 +247,17 @@ 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_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", "dequantize", ] diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index 6d5481c17..83d8be840 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -10,10 +10,84 @@ from typing import Any +from freetoken.utils import init_logger + from .reader import gguf_architecture, load_gguf_metadata +logger = init_logger(__name__) + # GGUF architecture -> transformers GGUF tokenizer-converter key. -_TOKENIZER_ARCH = {"gemma4": "gemma4_text"} +_TOKENIZER_ARCH = { + "gemma4": "gemma4_text", + "qwen35moe": "qwen3_moe", +} + + +# ggml token_type enum values that mark non-mergeable control strings. +# 1=NORMAL 2=UNKNOWN 3=CONTROL 4=USER_DEFINED 5=UNUSED 6=BYTE +_GGML_SPECIAL_TYPES = (2, 3, 4) + + +def _register_control_tokens(tokenizer, tokens: list[str], types: list[int]) -> None: + """Force every CONTROL/USER_DEFINED/UNKNOWN vocab entry to tokenize atomically. + + The GGUF->fast-tokenizer converter relies on BPE merges for special strings: + ``<|im_start|>`` happens to be a merge-table entry and survives, but ```` + is NOT (it never appears in merged training text), so it silently splits into + plain pieces (''). A model fed those garbage ids answers with + gibberish and EOS -- observed as "thought then returned empty" on Qwen3.6 GGUFs. + Registering each string as an added token makes the AddedVocabulary extract it + before BPE; because the string already exists in the base vocab, the existing id + is reused and the vocab never grows (transformers' own ``add_tokens`` wrapper + no-ops for in-vocab strings, so this goes through the backend directly). + """ + missing = [ + name + for i, (name, ty) in enumerate(zip(tokens, types)) + if int(ty) in _GGML_SPECIAL_TYPES + # Skip BYTE-fallback and unused; skip anything already atomic at its id. + and tokenizer.encode(name, add_special_tokens=False) != [i] + ] + if missing: + tokenizer.backend_tokenizer.add_tokens(missing) + logger.info( + "registered %d unmerged control tokens for atomic encoding (e.g. %s)", + len(missing), + ", ".join(repr(t) for t in missing[:6]), + ) + + +def _resolve_chat_template(meta: dict[str, Any], model_path: str) -> str | None: + """Chat template for GGUF checkpoints: explicit mirrors win, then metadata. + + Priority: a ``chat_template.jinja`` dropped NEXT TO the .gguf file, then + ``FT_CHAT_TEMPLATE_REPO`` ( on the HF Hub, fetched via huggingface_hub), + then the template embedded in the GGUF's ``tokenizer.chat_template`` metadata. + GGUF packagers (llama.cpp/unsloth) sometimes ship modified variants of the + official template — placing the official file beside the checkpoint overrides + it without repacking. The embedded one is the last resort, never wrong-by-default. + """ + import os + + sidecar = os.path.join(os.path.dirname(model_path), "chat_template.jinja") + if os.path.isfile(sidecar): + logger.info("using chat template sidecar %s", sidecar) + with open(sidecar, encoding="utf-8") as fh: + return fh.read() + repo = os.environ.get("FT_CHAT_TEMPLATE_REPO") + if repo: + try: + from huggingface_hub import hf_hub_download + + path = hf_hub_download(repo_id=repo, filename="chat_template.jinja") + with open(path, encoding="utf-8") as fh: + return fh.read() + except Exception as exc: # noqa: BLE001 — offline/bad repo is not fatal + logger.warning("FT_CHAT_TEMPLATE_REPO=%s fetch failed: %s", repo, exc) + embedded = meta.get("tokenizer.chat_template") + if isinstance(embedded, str) and embedded.strip(): + return embedded + return None def load_gguf_tokenizer(model_path: str): @@ -46,7 +120,11 @@ def tok_for(id_key: str, default: str) -> str: unk_token=tok_for("unknown_token_id", ""), pad_token=tok_for("padding_token_id", ""), ) - chat_template = meta.get("tokenizer.chat_template") + # GGUFs are not required to carry per-token types; absent means all-normal. + types = meta.get("tokenizer.ggml.token_type") or [] + if types: + _register_control_tokens(tokenizer, tokens, types) + chat_template = _resolve_chat_template(meta, str(model_path)) if chat_template: tokenizer.chat_template = chat_template return tokenizer diff --git a/python/freetoken/models/qwen3_5_moe/__init__.py b/python/freetoken/models/qwen3_5_moe/__init__.py index 98936e9f2..c32f4ffbf 100644 --- a/python/freetoken/models/qwen3_5_moe/__init__.py +++ b/python/freetoken/models/qwen3_5_moe/__init__.py @@ -1,4 +1,12 @@ from .config import parse_config +from .gguf import ( + convert_qwen35moe_to_gguf, + is_gguf_model, + iter_gguf_weights, + load_gguf_expert_sources, + load_gguf_expert_sources_native, + parse_gguf_config, +) from .model import Qwen3_5MoEForCausalLM from .weight import ( iter_weights, @@ -16,4 +24,10 @@ "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", + "load_gguf_expert_sources_native", ] diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index 2e7320051..9c208e608 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -76,10 +76,14 @@ def __init__( # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM). self._block_fp8 = expert_quant == "fp8_block" self._pertensor_fp8 = attn_quant == "fp8_pertensor" - self._fp8 = self._block_fp8 or self._pertensor_fp8 + self._gguf = expert_quant == "gguf" + # "Split" projection layout (qkv|z + ba as two GEMMs): used by the fp8 paths + # and by GGUF (native-quant qkv|z + dense bf16 ba). The fused 4-way in_proj is + # only for the plain bf16 case. + self._split_proj = self._block_fp8 or self._pertensor_fp8 or self._gguf self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] - if self._fp8: + if self._block_fp8 or self._pertensor_fp8: ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged self.in_proj_qkvz = ColMerged( hidden_size, [self.conv_dim, self.value_dim], has_bias=False @@ -87,6 +91,15 @@ def __init__( self.in_proj_ba = LinearColParallelMerged( hidden_size, [num_v_heads, num_v_heads], has_bias=False ) + elif self._gguf: + # GGUF: qkv|z are native-quant (Q8_0, swapped to GGUFLinear after build), b|a + # stay dense bf16. Same split as the fp8 path (matches the GGUF tensor layout). + self.in_proj_qkvz = LinearColParallelMerged( + hidden_size, [self.conv_dim, self.value_dim], has_bias=False + ) + self.in_proj_ba = LinearColParallelMerged( + hidden_size, [num_v_heads, num_v_heads], has_bias=False + ) else: # Fused input projection (one GEMM instead of four): qkv | z | b | a. self.in_proj = LinearColParallelMerged(hidden_size, self._in_proj_split, has_bias=False) @@ -161,7 +174,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: fla = build_fla_metadata(batch, hidden_states.device) batch.fla_metadata = fla - if self._fp8: + if self._split_proj: qkvz = self.in_proj_qkvz.forward(hidden_states) conv_in, z = torch.split(qkvz, [self.conv_dim, self.value_dim], dim=-1) ba = self.in_proj_ba.forward(hidden_states) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py new file mode 100644 index 000000000..3b6eb4a3f --- /dev/null +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -0,0 +1,652 @@ +"""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) 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 + +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_Q5_K, + GGML_Q6_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 + + down_quant_types = _gguf_down_quant_types(shim.model_path) + + 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", + gguf_down_quant_types=down_quant_types, + ) + + +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. +# -------------------------------------------------------------------------------------- + + +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 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 + + inner = model.model + inner.embed_tokens = GGUFEmbedding( + 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, + operation="lm_head", quant_role="lm_head", + ) + 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, quant_role="gdn_qkvz" + ) + g.out_proj = GGUFLinear( + g.value_dim, config.hidden_size, GGML_Q8_0, has_bias=False, quant_role="gdn_out" + ) + 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, quant_role="attention_qkv", + ) + attn.o_proj = GGUFLinear( + 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, + quant_role="shared_gate_up" + ) + m.shared_expert.down_proj = GGUFLinear( + shared_I, config.hidden_size, GGML_Q8_0, has_bias=False, + quant_role="shared_down" + ) + + +# -------------------------------------------------------------------------------------- +# 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 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 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 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]] = {} + expert_buf: dict[int, dict[str, tuple[torch.Tensor, int]]] = {} + + 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 + 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] + 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)}" + if include_moe_experts: + assert not expert_buf, f"incomplete routed expert groups: {sorted(expert_buf)}" + + +# -------------------------------------------------------------------------------------- +# Routed-expert host banks for the offload cache. +# +# The GGUF stores gate/up as Q4_K on every layer, but ``down`` as Q5_K on 37 layers and +# Q6_K on 3 -- heterogeneous row widths the offload cache cannot hold in one uniform bank +# (``set_bank_sources`` requires every layer to share a shape, and ``ggml_moe_a8_vec`` +# derives the row stride from the quant type). We keep ``gate_up`` native Q4_K and +# re-quantize the ``down`` experts to Q8_0 (8-bit, >= Q5_K/Q6_K precision, so no quality +# loss; a uniform per-bank format that fits the cache machinery). +# -------------------------------------------------------------------------------------- + + +def load_gguf_expert_sources( + model_path: str, config: ModelConfig, *, layer_sink=None +) -> dict[str, list[torch.Tensor]]: + """Per-layer host banks of the routed experts: ``gate_up`` native Q4_K + ``[E, 2I, row_bytes(H, Q4_K)]`` and ``down`` Q8_0 ``[E, H, row_bytes(I, Q8_0)]``. + + ``ffn_{gate,up}_exps`` are each ``[E, I, row_bytes(H, Q4_K)]`` packed and are fused + along the intermediate dim into ``gate_up``; ``ffn_down_exps`` (Q5_K/Q6_K) is + dequantized and re-quantized to Q8_0. Whole layers complete in two writes + (gate_up + down). ``layer_sink=None`` (serving): pin each layer's banks as they + complete via an internal :class:`PinPipeline`. + """ + from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, alloc_layer_banks + from freetoken.models.gguf.reader import iter_gguf_tensors + + _require_tp1("expert banks") + L, E = config.num_layers, config.num_experts + H, I = config.hidden_size, config.moe_intermediate_size + h_rb = row_bytes(H, GGML_Q4_K) + i_rb = row_bytes(I, GGML_Q8_0) + specs = { + "gate_up": ((E, 2 * I, h_rb), torch.uint8), + "down": ((E, H, i_rb), torch.uint8), + } + hb = alloc_layer_banks(specs, L) + banks = {name: [b.tensor for b in hb[name]] for name in specs} + + def _load(sink) -> None: + tracker = LayerCompletionTracker(2, hb, sink) if sink is not None else None + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + continue + layer = int(t.name.split(".")[1]) + if t.name.endswith("ffn_gate_exps.weight"): + banks["gate_up"][layer][:, :I].copy_(t.packed().reshape(E, I, h_rb)) + elif t.name.endswith("ffn_up_exps.weight"): + banks["gate_up"][layer][:, I:] = t.packed().reshape(E, I, h_rb) + elif t.name.endswith("ffn_down_exps.weight"): + flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.float32) + banks["down"][layer].copy_(quantize_q8_0(flat.reshape(E, H, I))) + else: + continue + if tracker is not None: + tracker.note(layer) + + if layer_sink is not None: + _load(layer_sink) + elif torch.cuda.is_available(): + with PinPipeline() as pins: + _load(pins) + else: + _load(None) + return banks + + +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/model.py b/python/freetoken/models/qwen3_5_moe/model.py index eba7fd24f..df0d41f0b 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -109,8 +109,28 @@ def __init__(self, config: ModelConfig): ) super().__init__() + # GGUF checkpoints carry native block-quantized weights: swap the dense + # projections + embedding for GGUF-quant ops (experts stay on the offload cache). + from .gguf import convert_qwen35moe_to_gguf, is_gguf_model + + if is_gguf_model(config): + convert_qwen35moe_to_gguf(self, config) + def forward(self) -> torch.Tensor: - output = self.model.forward(get_global_ctx().batch.input_ids) + ctx = get_global_ctx() + batch = ctx.batch + output = self.model.forward(batch.input_ids) + if batch.is_prefill: + # GGUFLinear (unlike ParallelLMHead / Nvfp4LMHead) does not gather each + # request's final row itself, but the engine contract is [batch.size, vocab] + # (the logits to sample after each request's prompt). Without this gather + # the sampler reads the FIRST prompt position's logits: the first generated + # token comes from position 0, and -- worse -- differs between a fresh + # prefill and a radix-cache continuation of the same prompt (the restored + # continuation has a different first row), which made greedy output flip + # between server states / cache states. + indices = batch.attn_metadata.get_last_indices(batch.size) + output = output[indices] return self.lm_head.forward(output) diff --git a/python/freetoken/models/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/register.py b/python/freetoken/models/register.py index b8c17322a..a339f901f 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", + ), # Qwen3.8-Flash-Next (model_type qwen4_exp): multimodal wrapper config (text tower in # text_config, weights under model.language_model.); served text-only. 36 GDN + 12 QSA # compressed-sparse attention layers on 4 hyper-connection residual streams, a PLE diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 6a34f3b9b..3bd1ce31f 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -342,6 +342,32 @@ 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 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 8b6116ba8..c67049da1 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -252,6 +252,29 @@ def _q4_0_banks(model_path, model_config, device, dtype, dummy, parallel=False, ) +def _gguf_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks: + if parallel: + raise NotImplementedError( + "parallel reader not implemented for gguf: GGUF is a single packed file " + "(not safetensors), so the common reader doesn't apply." + ) + if dummy: + raise NotImplementedError("gguf expert banks have no dummy path; load the real GGUF") + from freetoken.models.weight import ( + load_gguf_moe_expert_sources, + load_gguf_moe_expert_sources_native, + ) + + 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_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, + ) + + 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 +324,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.py b/python/freetoken/moe/fused.py index ecc45eae4..02ce3bc9f 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -6,7 +6,11 @@ import torch from freetoken.moe import BaseMoeBackend -from freetoken.utils import div_ceil +from freetoken.utils import div_ceil, init_logger + +logger = init_logger(__name__) + +_warned_torch_topk = False def _torch_fused_topk( @@ -40,9 +44,66 @@ def fused_topk( ) -> Tuple[torch.Tensor, torch.Tensor]: assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" - from freetoken.kernel.triton.moe_router import fused_topk_softmax - - return fused_topk_softmax(gating_output, topk, renormalize, num_token_non_padded) + 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. + if not is_triton_kernels_installed(): + global _warned_torch_topk + if not _warned_torch_topk: + _warned_torch_topk = True + # Once, not per call: this runs every MoE forward. On Linux a missing + # triton_kernels used to fail fast with ImportError; keep the misconfiguration + # visible without giving up the fallback that Windows needs. + 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. 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) + + if topk & (topk - 1): + # triton_kernels.topk builds tl.arange(0, k), which must be a power of 2; a + # top-10 router (qwen4_exp) takes the equivalent vendored triton router instead. + from freetoken.kernel.triton.moe_router import fused_topk_softmax + + return fused_topk_softmax(gating_output, topk, renormalize, num_token_non_padded) + + from triton_kernels.topk import topk as triton_kernels_topk + + logits = gating_output.float() + softmax_first = not renormalize + if softmax_first: + logits = torch.softmax(logits, dim=-1) + sparse_topk = triton_kernels_topk( + logits, + topk, + apply_softmax=not softmax_first, + ) + if hasattr(sparse_topk, "vals"): + topk_weights = sparse_topk.vals + topk_ids = sparse_topk.indx + else: + topk_weights, topk_ids = sparse_topk[:2] + topk_ids = topk_ids.to(torch.int32) + if num_token_non_padded is not None: + indices = torch.arange(0, topk_ids.shape[0], device=topk_ids.device) + topk_ids[indices >= num_token_non_padded, :] = -1 + return topk_weights, topk_ids def moe_align_block_size( diff --git a/python/freetoken/moe/fused_gguf.py b/python/freetoken/moe/fused_gguf.py new file mode 100644 index 000000000..cf866a9f3 --- /dev/null +++ b/python/freetoken/moe/fused_gguf.py @@ -0,0 +1,465 @@ +"""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* 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 +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} + + +@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 + down_q: torch.Tensor, # [num_slots, H, row_bytes(I, Q8_0)] uint8 + 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: + 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_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"} + ) + + 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 + ) + 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"): + 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) + 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 + ) + 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__ = ["MoeDecodeWork", "fused_experts_gguf", "fused_experts_gguf_native"] 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..a81bf9e90 --- /dev/null +++ b/python/freetoken/moe/nvfp4_to_mxfp4.py @@ -0,0 +1,227 @@ +"""NVFP4 -> MXFP4 (gpt-oss/FreeToken ``mxfp4_triton``) weight converter. + +ModelOpt NVFP4 (the format stored in FreeToken's native ``nvfp4`` banks) packs e2m1 +codes with a *fp8-e4m3* per-16 block scale and a per-output-row fp16 *global* scale. +MXFP4 (FreeToken's ``mxfp4_triton`` banks) packs e2m1 codes with an *e8m0* per-32 +block scale and a per-block bias, and is the native format of the gpt-oss family -- +the AMD-supported quant matrix alongside BF16/GGUF. + +This module converts a checkpoint's NVFP4 expert weights to MXFP4 *on load* (once, +cached per model), so a checkpoint that only ships NVFP4 can still run on AMD via the +portable MXFP4 path (the converter runs on the host, not on the GPU). The two formats +share the e2m1 code-packing (2 codes per byte, low nibble first), so only the scale +granularity (16 -> 32) and scale format (e4m3 -> e8m0) change. + +Layouts handled here (per projection/expert, ``N`` = output rows, ``K`` = input cols): + +* input (NVFP4, native ModelOpt rows): ``packed [N, K//2]`` uint8, ``scale [N, K//16]`` + fp8-e4m3, ``global [N]`` fp16. +* output (MXFP4, ``mxfp4_triton`` bank layout): ``blocks_t [N, K//2]`` uint8, + ``scales_t [N, K//32]`` uint8 e8m0. (FreeToken's MXFP4 stores the projection + transposed, N innermost, matching gpt-oss -- the converter emits that shape.) + +All numeric work is done in numpy so the core is unit-testable without a GPU/torch +runtime; the public entry converts torch tensors to/from numpy on the host. +""" + +from __future__ import annotations + +import math +from typing import Sequence + +try: + import numpy as _np +except ImportError: # pragma: no cover - numpy is a hard dep + _np = None + +__all__ = [ + "convert_nvfp4_to_mxfp4", + "dequantize_nvfp4_block", + "fp4_e2m1_table", + "e8m0_scale_and_codes", +] + +# --------------------------------------------------------------------------- +# e2m1 (fp4) code -> value table, and e8m0 (block scale) encode. +# +# e2m1: 1 sign + 2 exponent + 1 mantissa. With exponent bias 1 the finite set is: +# code value code value +# 0 0.0 8 -0.0 +# 1 0.5 9 -0.5 +# 2 1.0 10 -1.0 +# 3 1.5 11 -1.5 +# 4 2.0 12 -2.0 +# 5 3.0 13 -3.0 +# 6 4.0 14 -4.0 +# 7 6.0 15 -6.0 +# --------------------------------------------------------------------------- +_FP4_CODES = ( + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, +) +_FP4_TABLE = _np.asarray(_FP4_CODES, dtype=_np.float32) +# Magnitudes sorted ascending for the nearest-code search. +_FP4_SORT = _np.asarray(sorted(abs(v) for v in _FP4_CODES[1:8]), dtype=_np.float32) + + +def fp4_e2m1_table() -> Sequence[float]: + """The 16 e2m1 values keyed by 4-bit code (index == code).""" + return list(_FP4_CODES) + + +def e8m0_scale_and_codes(values: _np.ndarray, block: int = 32) -> tuple[_np.ndarray, _np.ndarray]: + """Return ``(scale_codes, fp4_codes)`` for ``values`` shaped ``[..., block]``: an + e8m0 ``uint8`` scale per block (the smallest power-of-2 scale covering the block + max-abs, in the MX ``2**(v-127)`` encoding) and the requantized 4-bit codes. + + ``scale_codes`` has shape ``values.shape[:-1]``; ``fp4_codes`` matches ``values``. + """ + v = values.reshape(-1, block) + amax = _np.max(_np.abs(v), axis=-1) + # e2m1 max positive magnitude is 6.0; choose the smallest power-of-2 scale so the + # block max-abs maps near the top of the e2m1 range (best precision): + # s = 2^ceil(log2(max_abs / 6.0)), stored as e8m0 code v with 2**(v-127) == s. + amax_safe = _np.maximum(amax, 1e-38) + exp = _np.ceil(_np.log2(amax_safe / 6.0)).astype(_np.float32) + exp = _np.where(amax == 0.0, 0.0, exp) + scale_codes = (127.0 + exp).astype(_np.uint8) # v-127 == exp + scale_v = (2.0 ** exp).astype(_np.float32) + # Requantize values in the block by its scale, then nearest-e2m1-code. + q = v / scale_v[:, None] + codes = _fp4_quantize(q).astype(_np.uint8) + return scale_codes.reshape(values.shape[:-1]), codes.reshape(values.shape) + + +def _fp4_quantize(values: _np.ndarray) -> _np.ndarray: + a = _np.abs(values) + diff = _np.abs(a[..., None] - _FP4_SORT) + idx = diff.argmin(axis=-1) + mag = _FP4_SORT[idx] + out = _np.where(mag < 0.25, 0, (idx + 1).astype(_np.uint8)) + out = _np.where(_np.signbit(values), out | 8, out) + return out + + +# --------------------------------------------------------------------------- +# NVFP4 -> MXFP4 +# --------------------------------------------------------------------------- + + +def dequantize_nvfp4_block( + packed: _np.ndarray, + scale: _np.ndarray, + global_scale: _np.ndarray, + *, + block: int = 16, +) -> _np.ndarray: + """Dequantize one native NVFP4 projection back to fp32. + + ``packed [N, K//2]`` uint8 (e2m1 pairs), ``scale [N, K//16]`` fp32 (already + converted from fp8-e4m3), ``global_scale [N]`` fp32. Returns ``[N, K]`` fp32. + """ + N, K2 = packed.shape + K = K2 * 2 + lo = (packed & 0x0F).astype(_np.uint8) + hi = (packed >> 4).astype(_np.uint8) + codes = _np.stack([lo, hi], axis=-1).reshape(N, K) # [N, K] + vals = _FP4_TABLE[codes.astype(_np.int64)] # [N, K] + # Per-16 block scale broadcast over K. + bs = _np.repeat(scale, block, axis=-1) # [N, K] + return (vals * bs).astype(_np.float32) * global_scale[:, None].astype(_np.float32) + + +def _pack_codes(codes: _np.ndarray) -> _np.ndarray: + """Pack ``[N, K]`` uint8 4-bit codes -> ``[N, K//2]`` uint8 (low nibble first).""" + N, K = codes.shape + even = codes[..., 0::2] + odd = codes[..., 1::2] + return (even | (odd << 4)).astype(_np.uint8) + + +def convert_nvfp4_to_mxfp4( + packed, + scale, + global_scale, + *, + axis: int = -1, + block: int = 32, +): + """Convert one projection's native NVFP4 expert weights to the MXFP4 layout. + + Args: + packed: ``[..., K//2]`` uint8 e2m1 pairs (low nibble = first code). + scale: ``[..., K//16]`` fp8-e4m3 block scale (fp32/fp16 input accepted). + global_scale: ``[...]`` per-output-row fp16 global scale. + axis: the K (contraction) axis along which blocks are grouped. + + Returns ``(mxfp4_packed [..., K//2] uint8, mxfp4_scales [..., K//32] uint8 e8m0)`` + matching the ``mxfp4_triton`` per-expert bank shape (K innermost). + """ + np = _np + packed = np.asarray(packed) + scale = np.asarray(scale, dtype=np.float32) + global_scale = np.asarray(global_scale, dtype=np.float32) + + if axis not in (-1, packed.ndim - 1): + raise NotImplementedError("converter requires the K axis to be innermost") + + # Dequantize NVFP4 to fp32, move K to the last axis. + K2 = packed.shape[-1] + K = K2 * 2 + assert K % block == 0, ( + f"NVFP4 K={K} is not a multiple of the MXFP4 block ({block}); " + "cannot convert without truncating weights" + ) + codes = np.stack([packed & 0x0F, (packed >> 4)], axis=-1).reshape( + *packed.shape[:-1], K + ) + vals = _FP4_TABLE[codes.astype(np.int64)] + bs = np.repeat(scale, 16, axis=-1) + f32 = (vals * bs).astype(np.float32) * global_scale[..., None].astype(np.float32) + + # Requantize to per-`block` e8m0 + e2m1. + flat = f32.reshape(-1, K) + n_blocks = K // block + flat_b = flat.reshape(-1, block) + scale_codes, mxfp4_codes = e8m0_scale_and_codes(flat_b, block=block) + out_codes = mxfp4_codes.reshape(flat.shape[0], n_blocks * block) + + mxfp4_packed = _pack_codes(out_codes) # [..., K//2] + mxfp4_scales = scale_codes.reshape(*packed.shape[:-1], n_blocks) + return mxfp4_packed, mxfp4_scales + + +# --------------------------------------------------------------------------- +# torch-tensor entrypoint +# --------------------------------------------------------------------------- + + +def _to_numpy(t: object) -> _np.ndarray: + if _np is not None and isinstance(t, _np.ndarray): + return t + try: + return t.detach().cpu().numpy() + except Exception as exc: # pragma: no cover + raise TypeError( + f"convert_nvfp4_to_mxfp4 expects torch tensors or numpy arrays, got {type(t)!r}" + ) from exc + + +def convert_torch_nvfp4_to_mxfp4( + packed, + scale, + global_scale, + *, + block: int = 32, +): + """Torch-tensor variant returning ``(mxfp4_packed, mxfp4_scales)`` torch tensors on + the same device as ``packed`` (host converter: input tensors are pulled to CPU and + the results copied back). Used at load time, cached per model.""" + import torch + + device = packed.device + dtype = packed.dtype + p, s = convert_nvfp4_to_mxfp4( + _to_numpy(packed), _to_numpy(scale), _to_numpy(global_scale), block=block + ) + return torch.from_numpy(p).to(device=device, dtype=dtype), torch.from_numpy(s).to(device=device) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index e1f20dd2f..4255728bb 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -45,6 +45,14 @@ # 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 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) @@ -90,6 +98,11 @@ + (H // 128) * fp8_block_scale_pad(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, + # 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), @@ -114,7 +127,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 48923e3b0..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,17 +420,38 @@ 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, - ) - 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 + 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 @@ -760,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: @@ -780,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; @@ -810,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 ca706c506..7be608b2e 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,8 +56,16 @@ def report_batch( page_size=page_size, mamba_slots=mamba_slots, swa_tokens=swa_tokens, + 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, @@ -101,6 +110,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 +131,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 +139,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/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 5b4db587d..70fb41593 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 @@ -152,6 +153,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" @@ -377,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, @@ -612,6 +629,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", @@ -725,6 +753,18 @@ 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 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(): + 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." + ) result = ServerArgs(**kwargs) logger = init_logger(__name__) logger.info(f"Parsed arguments:\n{result}") diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index 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/launch.py b/python/freetoken/server/launch.py index acef55175..c88c4ff5c 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -156,6 +156,18 @@ def start_subprocess() -> "BackendHandle": from .supervisor import BackendHandle mp.set_start_method("spawn", force=True) + # 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"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/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/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/__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..f1f8e14a6 --- /dev/null +++ b/python/freetoken/utils/graph_gate.py @@ -0,0 +1,364 @@ +"""HIP/CUDA graph-capture parity probe. + +Whether ``torch.cuda.graph`` graph capture works on the target GPU is the single +highest-informational-risk assumption for AMD (ROCm) support. This module probes it +once and records a PASS/FAIL + device result that the engine reads when deciding +whether to use CUDA-graph decode. On CUDA it is expected to PASS; on ROCm it may +fail on some consumer cards, in which case decode must use the kernel-launch path. + +The result is cached to disk under the user cache dir so it survives across runs, +and keyed by device kind + device name so a change of GPU invalidates it. +""" +from __future__ import annotations + +import json +import os +from functools import lru_cache + +_CACHE_FILE = "freetoken_graph_gate.json" +_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: + 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(variants: tuple[tuple[str, dict[str, str], str], ...] | None = None) -> dict: + """Run the capture probes on the current device across the variants. + + 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 + 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}", + } + + # 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 + + 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, + ) + 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: + requested = _rocm_blas_request() + if requested != "auto" and _is_rocm(): + return resolve_rocm_blas_env() + result = run_graph_gate() + return resolve_rocm_blas_env(gate=result) + except ValueError: + raise + except RuntimeError: + raise + except Exception: + 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() + 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": 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: + data = json.load(f) + 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: + pass + return None + + +def run_graph_gate() -> dict: + """Run (or reuse the cached result for the current device of) the capture probe.""" + cached = _load_cached() + if cached is not None: + return cached + result = probe_graph_capture() + try: + with open(_cache_path(), "w") as f: + json.dump(result, f) + except Exception: + pass + return result + + +@lru_cache(maxsize=1) +def graph_capture_status() -> str: + """Cached graph-capture status: ``"pass"``, ``"fail"``, or ``"unknown"`` (no device / + probe unavailable). The graph runner reads this to pick HIP-graph vs kernel-launch decode.""" + try: + result = run_graph_gate() + if result["ok"]: + return "pass" + if result.get("device_kind") and result["device_kind"] != "cpu": + return "fail" + return "unknown" + except Exception: + return "unknown" diff --git a/python/freetoken/utils/step_profiler.py b/python/freetoken/utils/step_profiler.py new file mode 100644 index 000000000..20a302bc8 --- /dev/null +++ b/python/freetoken/utils/step_profiler.py @@ -0,0 +1,403 @@ +"""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 +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``. + +``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. +""" + +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 + + +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", + "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() + + +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", "step", "phase", "stream_id") + + 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): + _mark(self.step, self.phase, "begin", self.stream_id) + return self + + def __exit__(self, exc_type=None, exc=None, tb=None) -> bool: # noqa: ANN001 + _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, "") + 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 _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: + if not _state.markers_enabled: + _state.done = True + return _NULL + _state.spec = spec + _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 + + 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, + ) + + +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 + 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") + _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/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/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/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..1ab821e7a --- /dev/null +++ b/scripts/mirror-hf-configs.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +# mirror-hf-configs.sh +# +# Mirror a model's official Hugging Face config files next to a local GGUF so +# FreeToken uses them instead of whatever the GGUF packager embedded: +# +# ./mirror-hf-configs.sh [gguf-file] +# +# ./mirror-hf-configs.sh Qwen/Qwen3.6-35B-A3B /path/to/model.gguf +# +# With no gguf argument, uses $FT_MODEL if set; otherwise errors. The GGUF must be +# passed explicitly (or via FT_MODEL) -- no implicit model-directory search. +# +# Files fetched (when they exist upstream): +# chat_template.jinja — read by FreeToken's GGUF loader (overrides embedded) +# generation_config.json — default sampling params +# tokenizer_config.json — reference only (FreeToken builds its tokenizer from GGUF) +# +# To revert, delete the mirrored files — the GGUF metadata is never modified. + +set -euo pipefail + +REPO_ID="${1:?usage: mirror-hf-configs.sh [gguf-file]}" +shift || true + +if [ "$#" -ge 1 ]; then + GGUF="$1" +else + GGUF="${FT_MODEL:-}" + [ -n "$GGUF" ] || { echo "ERROR: no gguf file given; pass one explicitly or set FT_MODEL" >&2; exit 1; } +fi +[ -f "$GGUF" ] || { echo "ERROR: not a file: $GGUF" >&2; exit 1; } + +DIR="$(dirname "$GGUF")" +echo "repo : $REPO_ID" +echo "into : $DIR" + +PY="${PY:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)/.venv-rocm/bin/python}" + +for f in chat_template.jinja generation_config.json tokenizer_config.json; do + if "$PY" - "$REPO_ID" "$f" "$DIR" <<'EOF' +import sys +from huggingface_hub import hf_hub_download +repo, fname, dest_dir = sys.argv[1:4] +try: + p = hf_hub_download(repo_id=repo, filename=fname) +except Exception: + sys.exit(1) +import shutil +shutil.copyfile(p, f"{dest_dir}/{fname}") +print(f"{fname}: OK") +EOF + then + : + else + echo "$f: not present upstream, skipped" + fi +done + +echo "done — restart the server to pick them up." diff --git a/scripts/profile-rocm-decode.sh b/scripts/profile-rocm-decode.sh new file mode 100755 index 000000000..ff19ae5b5 --- /dev/null +++ b/scripts/profile-rocm-decode.sh @@ -0,0 +1,225 @@ +#!/usr/bin/env bash +# profile-rocm-decode.sh +# +# 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 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" +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 + 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; } +[ -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). +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 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}" \ +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_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, + "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", + data=body, + headers={"Content-Type": "application/json"}, +) +with urllib.request.urlopen(req, timeout=1200) as r: + 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 + +# 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 [ "$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 + 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 new file mode 100755 index 000000000..ea357da5b --- /dev/null +++ b/scripts/serve-qwen-moe.sh @@ -0,0 +1,265 @@ +#!/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). +# 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 +# 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)" + +# --- .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:-}" +HOST="${FT_HOST:-127.0.0.1}" +PORT="${FT_PORT:-1920}" +ATNN="${FT_ATTN:-triton}" # triton | torch +MOE_BACKEND="${FT_MOE:-offload}" # offload required for K-quant experts +# radix = cross-request GDN-state prefix reuse (default); naive = no prefix reuse. +# Suspect for long-context degradation: radix GDN-state reuse corrupting later requests. +CACHE_TYPE="${FT_CACHE_TYPE:-radix}" +# MoE cache hit/miss stats in the decode log line (--moe-collect-stats). +MOE_STATS="${FT_MOE_STATS:-1}" +# Disable the two-buffer prefill MoE overlap (diagnostic: race check). +PREFILL_OVERLAP="${FT_PREFILL_OVERLAP:-1}" +# MoE layers computed on the CPU executor (diagnostic: '0' = all-GPU, no CPU path). +CPU_LAYERS="${FT_CPU_LAYERS:-}" +# --num-tokens sizes the KV cache in tokens (hybrid arch: only 10/40 layers are full +# attention, 2 kv heads x 256 dim -> ~20 KiB/token bf16; 128k ~= 2.7 GiB). +# The GPU expert-slot cache is sized by FT_MOE_CACHE: +# auto -> --moe-cache-auto: the engine derives slot bytes from the real expert +# tensors and fills all free VRAM AFTER reserving kv-reserve-tokens for KV. +# --kv-reserve-tokens MUST equal --num-tokens here: with an explicit +# --num-tokens the engine skips auto's own KV-half plan (num_page_override +# is set), so without a matching reservation greedy expert fill would eat +# VRAM the pinned KV still needs -> late CUDA OOM. +# -> fixed --moe-cache-size N slots (legacy behavior). +# Headroom: the engine may use FT_MEMORY_RATIO of free VRAM for weights+KV+experts +# combined (default here 0.80, upstream default 0.9). The remainder absorbs prefill +# transients -- the MoE overlap double-buffer alone needs ~3.8 GiB at this model's +# batch shapes; 0.9 left only ~2 GiB and OOM'd mid-prefill under VS Code payloads. +MEMORY_RATIO="${FT_MEMORY_RATIO:-0.80}" +# --max-prefill-length caps chunked-prefill chunk size (engine default 8192). The lm_head +# materializes logits for EVERY chunk token: an 8192-token chunk spikes ~3.8 GiB +# transiently -- enough to OOM on VS Code-sized prompts even with healthy headroom. +# 4096 halves the spike; long prompts just prefill in more chunks. +PREFILL_CHUNK="${FT_PREFILL_CHUNK:-4096}" +KV_TOKENS="${FT_KV_TOKENS:-131072}" +# Default max output tokens for requests that omit max_tokens. The reasoning model +# sometimes needs more room to finish its reasoning before answering; the engine +# default is 32k. Copilot sends its own max_tokens, which we cannot override, but +# other clients inherit this. +MAX_OUTPUT="${FT_MAX_OUTPUT:-65536}" +MOE_CACHE="${FT_MOE_CACHE:-auto}" +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:-qwen3.6}" + +die() { echo "ERROR: $*" >&2; exit 1; } + +# Map FT_MOE_CACHE to argv. Only the auto path carries --kv-reserve-tokens: +# with a fixed size the CLI already validates fit against the pinned KV. +# NOTE: validation must run in THIS shell (no process substitution), or die() +# would only exit a subshell and the launch would continue unvalidated. + +# Guard against the exact failure mode VS Code caused before: if this file ever +# gets saved with CRLF endings, every argument silently grows a trailing \r. +if grep -q $'\r' "${BASH_SOURCE[0]}"; then + die "CRLF line endings detected in $(basename "${BASH_SOURCE[0]}"). Run: sed -i 's/\r$//' ${BASH_SOURCE[0]}" +fi + +[ -x "$PY" ] || die "python not found: $PY (set PY=/path/to/venv-python)" + +server_pids() { + pgrep -f "freetoke[n].cli serve" || true +} + +start() { + [ -n "$MODEL" ] || die "no model configured: set FT_MODEL=/path/to/model.gguf" + if [ -n "$(server_pids)" ]; then + echo "A server is already running (pid $(server_pids | tr '\n' ' ')). Use 'stop' first." + exit 1 + fi + [ -f "$MODEL" ] || die "model not found: $MODEL" + + # One arg per array element; expanded once, single line, no continuations. + local -a SERVE_ARGS=( + "--model" "$MODEL" + $( [ -n "$SERVED_MODEL" ] && echo "--served-model-name" "$SERVED_MODEL" ) + "--moe-backend" "$MOE_BACKEND" + "--cache-type" "$CACHE_TYPE" + $( [ "$MOE_STATS" = 1 ] && echo "--moe-collect-stats" ) + $( [ "$PREFILL_OVERLAP" = 0 ] && echo "--disable-moe-prefill-overlap" ) + $( [ -n "$CPU_LAYERS" ] && echo "--moe-cpu-layers" "$CPU_LAYERS" ) + "--attention-backend" "$ATNN" + "--num-tokens" "$KV_TOKENS" + "--kv-type" "$KV_TYPE" + "--memory-ratio" "$MEMORY_RATIO" + "--max-prefill-length" "$PREFILL_CHUNK" + "--max-output-tokens" "$MAX_OUTPUT" + "--host" "$HOST" + "--port" "$PORT" + ) + case "$MOE_CACHE" in + auto) + SERVE_ARGS+=("--moe-cache-auto" "--kv-reserve-tokens" "$KV_TOKENS") + ;; + ''|*[!0-9]*) + die "FT_MOE_CACHE='$MOE_CACHE' is invalid: use 'auto' or a slot count" + ;; + *) + SERVE_ARGS+=("--moe-cache-size" "$MOE_CACHE") + ;; + esac + + echo "Launching FreeToken server" + echo " model : $MODEL" + echo " listen : $HOST:$PORT" + echo " attn : $ATNN" + echo " moe : $MOE_BACKEND" + echo " kv : $KV_TOKENS tokens (gpu moe cache: ${MOE_CACHE}${MOE_CACHE:+ }$( [ "$MOE_CACHE" = auto ] && echo "kv-reserve $KV_TOKENS" || echo slots))" + echo " 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 "${EXEC_PREFIX[@]}" "$PY" -m freetoken.cli serve "${SERVE_ARGS[@]}" >"$LOG" 2>&1 /dev/null || true + echo "pid=$pid — waiting for readiness (model load takes ~3-4 min)..." +} + +wait_ready() { + for _ in $(seq 1 90); do + if grep -q "API server is ready" "$LOG" 2>/dev/null; then + echo "READY on $HOST:$PORT" + return 0 + fi + if [ -z "$(server_pids)" ]; then + echo "server exited; last log lines:" >&2 + tail -20 "$LOG" >&2 || true + return 1 + fi + sleep 5 + done + echo "timed out waiting for readiness; see $LOG" >&2 + return 1 +} + +stop() { + local pids + pids="$(server_pids)" + if [ -z "$pids" ]; then + echo "no server running" + else + pkill -9 -f "freetoke[n].cli serve" 2>/dev/null || true + echo "stopped (was pid $pids)" + fi + pkill -9 -f "multiprocessing.spawn" 2>/dev/null || true + # free the distributed worker port (server_port+1) + local p pid + for p in "$PORT" "$((PORT + 1))"; do + pid="$(ss -ltnp 2>/dev/null | grep ":$p" | grep -oP 'pid=\K[0-9]+' | head -1 || true)" + if [ -n "$pid" ]; then + kill -9 "$pid" 2>/dev/null || true + fi + done +} + +status() { + local pids + pids="$(server_pids)" + if [ -n "$pids" ]; then + echo "RUNNING (pid $(echo "$pids" | tr '\n' ' ')) on $HOST:$PORT" + else + echo "NOT RUNNING" + fi + # Surface the auto-resolved expert-cache split so users don't have to read + # engine log lines; only present when FT_MOE_CACHE=auto booted the server. + if [ -f "$LOG" ]; then + local resolved + resolved="$(grep -o -- '--moe-cache-auto resolved moe_cache_size=[0-9]* num_pages=[0-9]*' "$LOG" | tail -1)" + [ -n "$resolved" ] && echo "moe cache: ${resolved//--moe-cache-auto resolved /}" + fi + [ -f "$LOG" ] && echo "--- last 5 log lines ($LOG) ---" && tail -5 "$LOG" +} + +case "${1:-start}" in + start) start && wait_ready ;; + stop) stop ;; + status) status ;; + log) exec tail -f "$LOG" ;; + *) echo "usage: $0 [start|stop|status|log]" >&2; exit 1 ;; +esac diff --git a/setup.py b/setup.py index d8e705ab6..744624398 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,7 @@ from __future__ import annotations import importlib.util +import os from pathlib import Path import sys @@ -12,15 +13,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 " @@ -30,12 +48,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=[ @@ -44,13 +93,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. @@ -59,10 +109,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"], ), # --ple-backend disk row store; Linux-only until the TableFile/BatchReader seams grow Windows bodies *([ diff --git a/tests/attention/test_torch_backend.py b/tests/attention/test_torch_backend.py new file mode 100644 index 000000000..3a132777d --- /dev/null +++ b/tests/attention/test_torch_backend.py @@ -0,0 +1,62 @@ +"""The debug ``"torch"`` attention backend. + +Verifies (a) the backend is registered and selectable, and (b) its pure-PyTorch +GQA attention math (with causal masking and the per-head output gate) matches +``torch.nn.functional.scaled_dot_product_attention`` on a synthetic case with the +qwen35moe full-attention geometry (num_q=16, num_kv=2, head_dim=256, gate). + +This is the ground-truth backend used to A/B the production triton path. +""" + +import torch + + +def _attention_math(q, ks, vs, lk, scale, group, gate): + """Mirror of TorchAttentionBackend's per-request attention (fp32 intermediates).""" + ks = ks.repeat_interleave(group, dim=1).float() + vs = vs.repeat_interleave(group, dim=1).float() + scores = torch.einsum("qhd,khd->hqk", q.float(), ks) * scale + lq = q.shape[0] + rows = torch.arange(lq) + cols = torch.arange(lk) + masked = (cols[None, :] > (lk - lq + rows)[:, None]).to(scores.device) + scores = scores.masked_fill(masked[None], float("-inf")) + probs = torch.softmax(scores, dim=-1) + o = torch.einsum("hqk,khd->qhd", probs, vs) # [lq, num_q, hd] + lq_, nq, hd_ = o.shape + o = o.reshape(lq_, nq * hd_) * torch.sigmoid(gate.float()) # [lq, num_q*hd] + return o.reshape(lq_, nq, hd_) + + +def test_torch_backend_registered(): + from freetoken.attention import SUPPORTED_ATTENTION_BACKENDS, attention_backend_info, AttnType + + assert "torch" in SUPPORTED_ATTENTION_BACKENDS.supported_names() + info = attention_backend_info("torch") + assert AttnType.FULL in info.supported_types + assert info.hybrid_linear_ok # must coexist with GDN layers + + +def test_torch_attention_matches_sdpa_prefill(): + num_q, num_kv, hd, group = 16, 2, 256, 8 + T = 4 + scale = hd ** -0.5 + q = torch.randn(T, num_q, hd) + k = torch.randn(T, num_kv, hd) + v = torch.randn(T, num_kv, hd) + gate = torch.randn(T, num_q * hd) + + out = _attention_math(q, k, v, T, scale, group, gate) + # reference: sdpa on GQA-expanded heads, then gate + no o_proj here + ke = k.repeat_interleave(group, dim=1) + ve = v.repeat_interleave(group, dim=1) + ref = torch.nn.functional.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0), + ke.transpose(0, 1).unsqueeze(0), + ve.transpose(0, 1).unsqueeze(0), + is_causal=True, + )[0].transpose(0, 1) + ref = ref.reshape(T, num_q * hd) * torch.sigmoid(gate) + ref = ref.reshape(T, num_q, hd) + + assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3) diff --git a/tests/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_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 ef7dc7ff8..cb52e3e17 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/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/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_backend_rocm.py b/tests/kernels/test_backend_rocm.py new file mode 100644 index 000000000..6b62e89e5 --- /dev/null +++ b/tests/kernels/test_backend_rocm.py @@ -0,0 +1,77 @@ +"""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(): + # 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 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(): + 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_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/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/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_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 diff --git a/tests/models/test_qwen35moe_gguf_deint.py b/tests/models/test_qwen35moe_gguf_deint.py new file mode 100644 index 000000000..176233ff8 --- /dev/null +++ b/tests/models/test_qwen35moe_gguf_deint.py @@ -0,0 +1,133 @@ +"""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 +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, +) + + +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) + + +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_fused_moe.py b/tests/moe/test_fused_moe.py index cafe77844..3db2f6431 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) diff --git a/tests/moe/test_nvfp4_backends_rocm.py b/tests/moe/test_nvfp4_backends_rocm.py new file mode 100644 index 000000000..cc71e4727 --- /dev/null +++ b/tests/moe/test_nvfp4_backends_rocm.py @@ -0,0 +1,107 @@ +"""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(): + # 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 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(): + 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/moe/test_offload.py b/tests/moe/test_offload.py index 422ca8675..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__( @@ -447,6 +449,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, 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_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" 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]